Repository: oom-ai/oomstore
Pull request #963 (2022-01-13 07:57:02)

Jinghan/implement offline method `Export`
This PR re-implements offline method `Export`, which supports exporting multi-group, multi-type (batch / streaming) features.

related to #962 

TODO

- [ ] implement `Export` for bigquery


Patch:
diff --git a/internal/database/offline/bigquery/export.go b/internal/database/offline/bigquery/export.go
index 7bb57e2e..f10fa480 100644
--- a/internal/database/offline/bigquery/export.go
+++ b/internal/database/offline/bigquery/export.go
@@ -17,20 +17,20 @@ import (
 	"github.com/oom-ai/oomstore/pkg/oomstore/types"
 )
 
-func (db *DB) Export(ctx context.Context, opt offline.ExportOpt) (<-chan types.ExportRecord, <-chan error) {
+func (db *DB) ExportOneGroup(ctx context.Context, opt offline.ExportOneGroupOpt) (<-chan types.ExportRecord, <-chan error) {
 	dbOpt := dbutil.DBOpt{
 		Backend:    types.BackendBigQuery,
 		BigQueryDB: db.Client,
 		DatasetID:  &db.datasetID,
 	}
-	doExportOpt := sqlutil.DoExportOpt{
-		ExportOpt:    opt,
-		QueryResults: bigqueryQueryExportResults,
+	doExportOpt := sqlutil.DoExportOneGroupOpt{
+		ExportOneGroupOpt: opt,
+		QueryResults:      bigqueryQueryExportResults,
 	}
-	return sqlutil.DoExport(ctx, dbOpt, doExportOpt)
+	return sqlutil.DoExportOneGroup(ctx, dbOpt, doExportOpt)
 }
 
-func bigqueryQueryExportResults(ctx context.Context, dbOpt dbutil.DBOpt, opt offline.ExportOpt, query string, args []interface{}) (<-chan types.ExportRecord, <-chan error) {
+func bigqueryQueryExportResults(ctx context.Context, dbOpt dbutil.DBOpt, opt offline.ExportOneGroupOpt, query string, args []interface{}) (<-chan types.ExportRecord, <-chan error) {
 	stream := make(chan types.ExportRecord)
 	errs := make(chan error, 1) // at most 1 error
 	for _, arg := range args {
@@ -66,3 +66,8 @@ func bigqueryQueryExportResults(ctx context.Context, dbOpt dbutil.DBOpt, opt off
 
 	return stream, errs
 }
+
+func (db *DB) Export(ctx context.Context, opt offline.ExportOpt) (<-chan types.ExportRecord, <-chan error) {
+	//TODO implement me
+	panic("implement me")
+}
diff --git a/internal/database/offline/mock_offline/store.go b/internal/database/offline/mock_offline/store.go
index cfc256fe..580d286d 100644
--- a/internal/database/offline/mock_offline/store.go
+++ b/internal/database/offline/mock_offline/store.go
@@ -79,6 +79,21 @@ func (mr *MockStoreMockRecorder) Export(ctx, opt interface{}) *gomock.Call {
 	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Export", reflect.TypeOf((*MockStore)(nil).Export), ctx, opt)
 }
 
+// ExportOneGroup mocks base method.
+func (m *MockStore) ExportOneGroup(ctx context.Context, opt offline.ExportOneGroupOpt) (<-chan types.ExportRecord, <-chan error) {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "ExportOneGroup", ctx, opt)
+	ret0, _ := ret[0].(<-chan types.ExportRecord)
+	ret1, _ := ret[1].(<-chan error)
+	return ret0, ret1
+}
+
+// ExportOneGroup indicates an expected call of ExportOneGroup.
+func (mr *MockStoreMockRecorder) ExportOneGroup(ctx, opt interface{}) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ExportOneGroup", reflect.TypeOf((*MockStore)(nil).ExportOneGroup), ctx, opt)
+}
+
 // Import mocks base method.
 func (m *MockStore) Import(ctx context.Context, opt offline.ImportOpt) (int64, error) {
 	m.ctrl.T.Helper()
diff --git a/internal/database/offline/mysql/store.go b/internal/database/offline/mysql/store.go
index 485243da..ffb69792 100644
--- a/internal/database/offline/mysql/store.go
+++ b/internal/database/offline/mysql/store.go
@@ -37,6 +37,10 @@ func (db *DB) Import(ctx context.Context, opt offline.ImportOpt) (int64, error)
 	return sqlutil.Import(ctx, db.DB, opt, dbutil.LoadDataFromSource(Backend, MySQLBatchSize), Backend)
 }
 
+func (db *DB) ExportOneGroup(ctx context.Context, opt offline.ExportOneGroupOpt) (<-chan types.ExportRecord, <-chan error) {
+	return sqlutil.ExportOneGroup(ctx, db.DB, opt, Backend)
+}
+
 func (db *DB) Export(ctx context.Context, opt offline.ExportOpt) (<-chan types.ExportRecord, <-chan error) {
 	return sqlutil.Export(ctx, db.DB, opt, Backend)
 }
diff --git a/internal/database/offline/postgres/store.go b/internal/database/offline/postgres/store.go
index d78e4d05..6ae9ffb9 100644
--- a/internal/database/offline/postgres/store.go
+++ b/internal/database/offline/postgres/store.go
@@ -34,8 +34,13 @@ func (db *DB) Import(ctx context.Context, opt offline.ImportOpt) (int64, error)
 	return sqlutil.Import(ctx, db.DB, opt, loadDataFromSource, Backend)
 }
 
+func (db *DB) ExportOneGroup(ctx context.Context, opt offline.ExportOneGroupOpt) (<-chan types.ExportRecord, <-chan error) {
+	return sqlutil.ExportOneGroup(ctx, db.DB, opt, Backend)
+}
+
 func (db *DB) Export(ctx context.Context, opt offline.ExportOpt) (<-chan types.ExportRecord, <-chan error) {
 	return sqlutil.Export(ctx, db.DB, opt, Backend)
+
 }
 
 func (db *DB) Join(ctx context.Context, opt offline.JoinOpt) (*types.JoinResult, error) {
diff --git a/internal/database/offline/redshift/store.go b/internal/database/offline/redshift/store.go
index 9f136f8a..66f3f3b6 100644
--- a/internal/database/offline/redshift/store.go
+++ b/internal/database/offline/redshift/store.go
@@ -37,6 +37,10 @@ func (db *DB) Import(ctx context.Context, opt offline.ImportOpt) (int64, error)
 	return sqlutil.Import(ctx, db.DB, opt, dbutil.LoadDataFromSource(Backend, RedshiftBatchSize), Backend)
 }
 
+func (db *DB) ExportOneGroup(ctx context.Context, opt offline.ExportOneGroupOpt) (<-chan types.ExportRecord, <-chan error) {
+	return sqlutil.ExportOneGroup(ctx, db.DB, opt, Backend)
+}
+
 func (db *DB) Export(ctx context.Context, opt offline.ExportOpt) (<-chan types.ExportRecord, <-chan error) {
 	return sqlutil.Export(ctx, db.DB, opt, Backend)
 }
diff --git a/internal/database/offline/snowflake/store.go b/internal/database/offline/snowflake/store.go
index c7c284c9..b89f054c 100644
--- a/internal/database/offline/snowflake/store.go
+++ b/internal/database/offline/snowflake/store.go
@@ -51,6 +51,10 @@ func (db *DB) Import(ctx context.Context, opt offline.ImportOpt) (int64, error)
 	return sqlutil.Import(ctx, db.DB, opt, dbutil.LoadDataFromSource(Backend, SnowflakeBatchSize), Backend)
 }
 
+func (db *DB) ExportOneGroup(ctx context.Context, opt offline.ExportOneGroupOpt) (<-chan types.ExportRecord, <-chan error) {
+	return sqlutil.ExportOneGroup(ctx, db.DB, opt, Backend)
+}
+
 func (db *DB) Export(ctx context.Context, opt offline.ExportOpt) (<-chan types.ExportRecord, <-chan error) {
 	return sqlutil.Export(ctx, db.DB, opt, Backend)
 }
diff --git a/internal/database/offline/sqlite/store.go b/internal/database/offline/sqlite/store.go
index 9b24a83a..6d55a7fa 100644
--- a/internal/database/offline/sqlite/store.go
+++ b/internal/database/offline/sqlite/store.go
@@ -37,6 +37,10 @@ func (db *DB) Import(ctx context.Context, opt offline.ImportOpt) (int64, error)
 	return sqlutil.Import(ctx, db.DB, opt, dbutil.LoadDataFromSource(Backend, SQLiteBatchSize), Backend)
 }
 
+func (db *DB) ExportOneGroup(ctx context.Context, opt offline.ExportOneGroupOpt) (<-chan types.ExportRecord, <-chan error) {
+	return sqlutil.ExportOneGroup(ctx, db.DB, opt, Backend)
+}
+
 func (db *DB) Export(ctx context.Context, opt offline.ExportOpt) (<-chan types.ExportRecord, <-chan error) {
 	return sqlutil.Export(ctx, db.DB, opt, Backend)
 }
diff --git a/internal/database/offline/sqlutil/export.go b/internal/database/offline/sqlutil/export.go
index 61894db5..7d5b44af 100644
--- a/internal/database/offline/sqlutil/export.go
+++ b/internal/database/offline/sqlutil/export.go
@@ -3,6 +3,7 @@ package sqlutil
 import (
 	"context"
 	"fmt"
+	"sort"
 
 	"github.com/jmoiron/sqlx"
 	"github.com/spf13/cast"
@@ -13,13 +14,47 @@ import (
 	"github.com/oom-ai/oomstore/pkg/oomstore/types"
 )
 
-type QueryExportResults func(ctx context.Context, dbOpt dbutil.DBOpt, opt offline.ExportOpt, query string, args []interface{}) (<-chan types.ExportRecord, <-chan error)
+type QueryExportResults func(ctx context.Context, dbOpt dbutil.DBOpt, opt offline.ExportOneGroupOpt, query string, args []interface{}) (<-chan types.ExportRecord, <-chan error)
+
+type DoExportOneGroupOpt struct {
+	offline.ExportOneGroupOpt
+	QueryResults QueryExportResults
+}
 
 type DoExportOpt struct {
 	offline.ExportOpt
 	QueryResults QueryExportResults
 }
 
+func ExportOneGroup(ctx context.Context, db *sqlx.DB, opt offline.ExportOneGroupOpt, backend types.BackendType) (<-chan types.ExportRecord, <-chan error) {
+	dbOpt := dbutil.DBOpt{
+		Backend: backend,
+		SqlxDB:  db,
+	}
+	doJoinOpt := DoExportOneGroupOpt{
+		ExportOneGroupOpt: opt,
+		QueryResults:      sqlxQueryExportResults,
+	}
+	return DoExportOneGroup(ctx, dbOpt, doJoinOpt)
+}
+
+func DoExportOneGroup(ctx context.Context, dbOpt dbutil.DBOpt, opt DoExportOneGroupOpt) (<-chan types.ExportRecord, <-chan error) {
+	var (
+		emptyStream = make(chan types.ExportRecord)
+		errs        = make(chan error, 1) // at most 1 error
+	)
+	defer close(emptyStream)
+	defer close(errs)
+
+	query, args, err := buildExportOneGroupQuery(dbOpt, opt.ExportOneGroupOpt)
+	if err != nil {
+		errs <- errdefs.WithStack(err)
+		return emptyStream, errs
+	}
+
+	return opt.QueryResults(ctx, dbOpt, opt.ExportOneGroupOpt, query, args)
+}
+
 func Export(ctx context.Context, db *sqlx.DB, opt offline.ExportOpt, backend types.BackendType) (<-chan types.ExportRecord, <-chan error) {
 	dbOpt := dbutil.DBOpt{
 		Backend: backend,
@@ -40,16 +75,119 @@ func DoExport(ctx context.Context, dbOpt dbutil.DBOpt, opt DoExportOpt) (<-chan
 	defer close(emptyStream)
 	defer close(errs)
 
-	query, args, err := buildExportQuery(dbOpt, opt.ExportOpt)
+	// Step 0: prepare variables
+	var (
+		snapshotTables = make([]string, 0, len(opt.SnapshotTables))
+		cdcTables      = make([]string, 0, len(opt.CdcTables))
+		groupIDs       = make([]int, 0, len(opt.Features))
+	)
+	for groupID := range opt.Features {
+		groupIDs = append(groupIDs, groupID)
+	}
+	sort.Slice(groupIDs, func(i, j int) bool {
+		return groupIDs[i] < groupIDs[j]
+	})
+	for _, groupID := range groupIDs {
+		snapshotTables = append(snapshotTables, opt.SnapshotTables[groupID])
+		if _, ok := opt.CdcTables[groupID]; ok {
+			cdcTables = append(cdcTables, opt.CdcTables[groupID])
+		}
+	}
+
+	// Step 1: prepare export_entity table, which contains all entity keys from source tables
+	tableName, err := prepareEntityTable(ctx, dbOpt, opt.ExportOpt, snapshotTables, cdcTables)
+	if err != nil {
+		errs <- errdefs.WithStack(err)
+		return emptyStream, errs
+	}
+
+	// Step 2: join export_entity table, snapshot tables and cdc tables
+	qt := dbutil.QuoteFn(dbOpt.Backend)
+	var (
+		fields      []string
+		featureList types.FeatureList
+	)
+	for _, groupID := range groupIDs {
+		features := opt.Features[groupID]
+		if features[0].Group.Category == types.CategoryBatch {
+			for _, f := range features {
+				fields = append(fields, fmt.Sprintf("%s.%s AS %s", qt(opt.SnapshotTables[groupID]), qt(f.Name), qt(f.FullName)))
+				featureList = append(featureList, f)
+			}
+		} else {
+			for _, f := range features {
+				cdc := fmt.Sprintf("%s.%s", opt.CdcTables[groupID]+"_0", qt(f.Name))
+				snapshot := fmt.Sprintf("%s.%s", qt(opt.SnapshotTables[groupID]), qt(f.Name))
+				fields = append(fields, fmt.Sprintf("(CASE WHEN %s IS NULL THEN %s ELSE %s END) AS %s", cdc, snapshot, cdc, qt(f.FullName)))
+				featureList = append(featureList, f)
+			}
+		}
+	}
+	query, err := buildExportQuery(exportQueryParams{
+		EntityTableName: tableName,
+		EntityName:      opt.EntityName,
+		UnixMilli:       "unix_milli",
+		SnapshotTables:  snapshotTables,
+		CdcTables:       cdcTables,
+		Fields:          fields,
+		Backend:         dbOpt.Backend,
+	})
 	if err != nil {
 		errs <- errdefs.WithStack(err)
 		return emptyStream, errs
 	}
+	args := make([]interface{}, 0, len(opt.CdcTables)*2)
+	for i := 0; i < len(opt.CdcTables)*2; i++ {
+		args = append(args, opt.UnixMilli)
+	}
+	return queryExportResults(ctx, dbOpt, query, args, featureList)
+}
+
+func queryExportResults(ctx context.Context, dbOpt dbutil.DBOpt, query string, args []interface{}, featureList types.FeatureList) (<-chan types.ExportRecord, <-chan error) {
+	stream := make(chan types.ExportRecord)
+	errs := make(chan error, 1) // at most 1 error
 
-	return opt.QueryResults(ctx, dbOpt, opt.ExportOpt, query, args)
+	go func() {
+		defer close(stream)
+		defer close(errs)
+		stmt, err := dbOpt.SqlxDB.Preparex(dbOpt.SqlxDB.Rebind(query))
+		if err != nil {
+			errs <- errdefs.WithStack(err)
+			return
+		}
+		defer stmt.Close()
+		rows, err := stmt.Queryx(args...)
+		if err != nil {
+			errs <- errdefs.WithStack(err)
+			return
+		}
+		defer rows.Close()
+		for rows.Next() {
+			record, err := rows.SliceScan()
+			if err != nil {
+				errs <- errdefs.Errorf("failed at rows.SliceScan, err=%v", err)
+				return
+			}
+			record[0] = cast.ToString(record[0])
+			for i, f := range featureList {
+				if record[i+1] == nil {
+					continue
+				}
+				deserializedValue, err := dbutil.DeserializeByValueType(record[i+1], f.ValueType, dbOpt.Backend)
+				if err != nil {
+					errs <- err
+					return
+				}
+				record[i+1] = deserializedValue
+			}
+			stream <- record
+		}
+	}()
+
+	return stream, errs
 }
 
-func sqlxQueryExportResults(ctx context.Context, dbOpt dbutil.DBOpt, opt offline.ExportOpt, query string, args []interface{}) (<-chan types.ExportRecord, <-chan error) {
+func sqlxQueryExportResults(ctx context.Context, dbOpt dbutil.DBOpt, opt offline.ExportOneGroupOpt, query string, args []interface{}) (<-chan types.ExportRecord, <-chan error) {
 	stream := make(chan types.ExportRecord)
 	errs := make(chan error, 1) // at most 1 error
 
@@ -93,7 +231,7 @@ func sqlxQueryExportResults(ctx context.Context, dbOpt dbutil.DBOpt, opt offline
 	return stream, errs
 }
 
-func buildExportQuery(dbOpt dbutil.DBOpt, opt offline.ExportOpt) (string, []interface{}, error) {
+func buildExportOneGroupQuery(dbOpt dbutil.DBOpt, opt offline.ExportOneGroupOpt) (string, []interface{}, error) {
 	if opt.CdcTable == nil && opt.UnixMilli == nil {
 		return buildExportBatchQuery(dbOpt, opt), nil, nil
 	}
@@ -103,7 +241,7 @@ func buildExportQuery(dbOpt dbutil.DBOpt, opt offline.ExportOpt) (string, []inte
 	return "", nil, fmt.Errorf("invalid option %+v", opt)
 }
 
-func buildExportBatchQuery(dbOpt dbutil.DBOpt, opt offline.ExportOpt) string {
+func buildExportBatchQuery(dbOpt dbutil.DBOpt, opt offline.ExportOneGroupOpt) string {
 	fields := append([]string{opt.EntityName}, opt.Features.Names()...)
 	qt := dbutil.QuoteFn(dbOpt.Backend)
 
@@ -118,7 +256,7 @@ func buildExportBatchQuery(dbOpt dbutil.DBOpt, opt offline.ExportOpt) string {
 	return query
 }
 
-func buildExportStreamQuery(dbOpt dbutil.DBOpt, opt offline.ExportOpt) (string, []interface{}, error) {
+func buildExportStreamQuery(dbOpt dbutil.DBOpt, opt offline.ExportOneGroupOpt) (string, []interface{}, error) {
 	query, err := buildAggregateQuery(aggregateQueryParams{
 		EntityName:            opt.EntityName,
 		UnixMilli:             "unix_milli",
diff --git a/internal/database/offline/sqlutil/export_helper.go b/internal/database/offline/sqlutil/export_helper.go
index 8e06dfa6..22b51866 100644
--- a/internal/database/offline/sqlutil/export_helper.go
+++ b/internal/database/offline/sqlutil/export_helper.go
@@ -2,15 +2,57 @@ package sqlutil
 
 import (
 	"bytes"
+	"context"
 	"fmt"
 	"strings"
 	"text/template"
 
+	"github.com/oom-ai/oomstore/pkg/errdefs"
+
+	"github.com/oom-ai/oomstore/internal/database/offline"
+
 	"github.com/oom-ai/oomstore/internal/database/dbutil"
 	"github.com/oom-ai/oomstore/pkg/oomstore/types"
 	"github.com/pkg/errors"
 )
 
+const UNION_ENTITY_QUERY = `
+INSERT INTO {{ qt .TableName }}
+{{ snapshot .SnapshotTables }}
+{{ cdc .CdcTables }}
+`
+
+const EXPORT_QUERY = `
+SELECT
+	e.{{ qt .EntityName }},
+	{{ fieldJoin .Fields }}
+FROM {{ qt .EntityTableName }} AS e
+{{ range $table := .SnapshotTables }}
+LEFT JOIN {{ qt $table }}
+ON e.{{ qt $.EntityName }} = {{ qt $table }}.{{ qt $.EntityName }}
+{{end}}
+{{ range $table := .CdcTables }}
+LEFT JOIN
+(
+	SELECT
+		{{ $table }}_1.*
+	FROM {{ qt $table }} AS {{ $table }}_1
+	JOIN
+	(SELECT
+		{{ qt $.EntityName }},
+		MAX({{ qt $.UnixMilli }}) AS {{ qt $.UnixMilli }}
+	FROM {{ qt $table }}
+	WHERE {{ qt $table }}.{{ qt $.UnixMilli }} <= ?
+	GROUP BY {{ qt $.EntityName }}
+	) AS {{ $table }}_2
+	ON {{ $table }}_1.{{ qt $.EntityName }} = {{ $table }}_2.{{ qt $.EntityName }} AND {{ $table }}_1.{{ qt $.UnixMilli }} = {{ $table }}_2.{{ qt $.UnixMilli }}
+	WHERE {{ $table }}_1.{{ qt $.UnixMilli }} <= ?
+) AS {{ $table }}_0
+ON e.{{ qt $.EntityName }} = {{ $table }}_0.{{ qt $.EntityName }}
+{{end}}
+
+`
+
 const AGGREGATE_QUERY = `
 SELECT
 	l.{{ qt .EntityName }},
@@ -79,10 +121,7 @@ func buildAggregateQuery(params aggregateQueryParams) (string, error) {
 	}
 	qt := dbutil.QuoteFn(params.Backend)
 	t := template.Must(template.New("snapshot").Funcs(template.FuncMap{
-		"qt": qt,
-		"columnJoin": func(columns []string) string {
-			return qt(columns...)
-		},
+		"qt": dbutil.QuoteFn(params.Backend),
 		"featureValue": func(features []string) string {
 			values := make([]string, 0, len(features))
 			for _, f := range features {
@@ -98,3 +137,115 @@ func buildAggregateQuery(params aggregateQueryParams) (string, error) {
 	}
 	return buf.String(), nil
 }
+
+type unionEntityQueryParams struct {
+	TableName      string
+	EntityName     string
+	SnapshotTables []string
+	CdcTables      []string
+	UnixMilli      int64
+	Backend        types.BackendType
+}
+
+func buildUnionEntityQuery(params unionEntityQueryParams) (string, []interface{}, error) {
+	qt := dbutil.QuoteFn(params.Backend)
+	var args []interface{}
+	t := template.Must(template.New("union_entity").Funcs(template.FuncMap{
+		"qt": qt,
+		"snapshot": func(tables []string) string {
+			query := make([]string, 0, len(tables))
+			for _, t := range tables {
+				query = append(query, fmt.Sprintf("SELECT %s FROM %s", qt(params.EntityName), qt(t)))
+			}
+			return strings.Join(query, "UNION \n\t")
+		},
+		"cdc": func(tables []string) string {
+			if len(tables) == 0 {
+				return ""
+			}
+			query := make([]string, 0, len(tables))
+			for _, t := range tables {
+				query = append(query, fmt.Sprintf("SELECT %s FROM %s WHERE %s <= ?", qt(params.EntityName), qt(t), qt("unix_milli")))
+				args = append(args, params.UnixMilli)
+			}
+			return fmt.Sprintf("%s%s", "UNION \n\t", strings.Join(query, "UNION \n\t"))
+		},
+	}).Parse(UNION_ENTITY_QUERY))
+
+	buf := bytes.NewBuffer(nil)
+	if err := t.Execute(buf, params); err != nil {
+		return "", nil, errors.WithStack(err)
+	}
+	return buf.String(), args, nil
+}
+
+type exportQueryParams struct {
+	EntityTableName string
+	EntityName      string
+	UnixMilli       string
+	SnapshotTables  []string
+	CdcTables       []string
+	Fields          []string
+	Backend         types.BackendType
+}
+
+func buildExportQuery(params exportQueryParams) (string, error) {
+	t := template.Must(template.New("export").Funcs(template.FuncMap{
+		"qt": dbutil.QuoteFn(params.Backend),
+		"fieldJoin": func(fields []string) string {
+			return strings.Join(fields, ",\n\t")
+		},
+	}).Parse(EXPORT_QUERY))
+
+	buf := bytes.NewBuffer(nil)
+	if err := t.Execute(buf, params); err != nil {
+		return "", errors.WithStack(err)
+	}
+	return buf.String(), nil
+}
+
+func prepareEntityTable(ctx context.Context, dbOpt dbutil.DBOpt, opt offline.ExportOpt, snapshotTables, cdcTables []string) (string, error) {
+	// Step 1: create table export_entity
+	tableName := dbutil.TempTable("export_entity")
+	qtTableName, columnDefs, err := prepareTableSchema(dbOpt, prepareTableSchemaParams{
+		tableName:    tableName,
+		entityName:   opt.EntityName,
+		hasUnixMilli: false,
+	})
+	if err != nil {
+		return "", err
+	}
+	schema := fmt.Sprintf(`
+		CREATE TABLE %s (
+			%s
+		);
+	`, qtTableName, strings.Join(columnDefs, ",\n"))
+	if err = dbOpt.ExecContext(ctx, schema, nil); err != nil {
+		return "", err
+	}
+
+	// Step 2: aggregate all entity keys
+	query, args, err := buildUnionEntityQuery(unionEntityQueryParams{
+		TableName:      tableName,
+		EntityName:     opt.EntityName,
+		SnapshotTables: snapshotTables,
+		CdcTables:      cdcTables,
+		UnixMilli:      opt.UnixMilli,
+		Backend:        dbOpt.Backend,
+	})
+	if err != nil {
+		return "", errdefs.WithStack(err)
+	}
+	if err = dbOpt.ExecContext(ctx, query, args); err != nil {
+		return "", errdefs.WithStack(err)
+	}
+
+	// Step 3: create index on table entity_rows
+	if supportIndex(dbOpt.Backend) {
+		index := fmt.Sprintf(`CREATE UNIQUE INDEX idx_%s ON %s (%s)`, tableName, tableName, opt.EntityName)
+		if err = dbOpt.ExecContext(ctx, index, nil); err != nil {
+			return "", err
+		}
+	}
+	return tableName, nil
+}
diff --git a/internal/database/offline/sqlutil/join_helper.go b/internal/database/offline/sqlutil/join_helper.go
index d80fe188..f9daad57 100644
--- a/internal/database/offline/sqlutil/join_helper.go
+++ b/internal/database/offline/sqlutil/join_helper.go
@@ -79,7 +79,12 @@ func PrepareJoinedTable(
 ) (string, error) {
 	// Step 1: create table joined_
 	tableName := dbutil.TempTable(fmt.Sprintf("joined_%s", groupName))
-	qtTableName, columnDefs, err := prepareTableSchema(dbOpt, entity, tableName, valueNames)
+	qtTableName, columnDefs, err := prepareTableSchema(dbOpt, prepareTableSchemaParams{
+		tableName:    tableName,
+		entityName:   "entity_key",
+		valueNames:   valueNames,
+		hasUnixMilli: true,
+	})
 	if err != nil {
 		return "", err
 	}
@@ -124,7 +129,12 @@ func PrepareEntityRowsTable(ctx context.Context,
 ) (string, error) {
 	// Step 1: create table entity_rows
 	tableName := dbutil.TempTable("entity_rows")
-	qtTableName, columnDefs, err := prepareTableSchema(dbOpt, entity, tableName, valueNames)
+	qtTableName, columnDefs, err := prepareTableSchema(dbOpt, prepareTableSchemaParams{
+		tableName:    tableName,
+		entityName:   "entity_key",
+		valueNames:   valueNames,
+		hasUnixMilli: true,
+	})
 	if err != nil {
 		return "", err
 	}
@@ -154,7 +164,14 @@ func PrepareEntityRowsTable(ctx context.Context,
 	return tableName, nil
 }
 
-func prepareTableSchema(dbOpt dbutil.DBOpt, entity types.Entity, tableName string, valueNames []string) (string, []string, error) {
+type prepareTableSchemaParams struct {
+	tableName    string
+	entityName   string
+	valueNames   []string
+	hasUnixMilli bool
+}
+
+func prepareTableSchema(dbOpt dbutil.DBOpt, params prepareTableSchemaParams) (string, []string, error) {
 	columnFormat, err := dbutil.GetColumnFormat(dbOpt.Backend)
 	if err != nil {
 		return "", nil, err
@@ -167,26 +184,28 @@ func prepareTableSchema(dbOpt dbutil.DBOpt, entity types.Entity, tableName strin
 	case types.BackendBigQuery:
 		entityType = "STRING"
 		valueType = "STRING"
-		qtTableName = fmt.Sprintf("%s.%s", *dbOpt.DatasetID, qt(tableName))
+		qtTableName = fmt.Sprintf("%s.%s", *dbOpt.DatasetID, qt(params.tableName))
 	case types.BackendMySQL:
 		entityType = "VARCHAR(255)"
 		valueType = "TEXT"
-		qtTableName = qt(tableName)
+		qtTableName = qt(params.tableName)
 	case types.BackendSnowflake:
 		entityType = "TEXT"
 		valueType = "TEXT"
-		qtTableName = fmt.Sprintf("PUBLIC.%s", qt(tableName))
+		qtTableName = fmt.Sprintf("PUBLIC.%s", qt(params.tableName))
 	default:
 		entityType = "TEXT"
 		valueType = "TEXT"
-		qtTableName = qt(tableName)
+		qtTableName = qt(params.tableName)
 	}
 
 	columnDefs := []string{
-		fmt.Sprintf(`%s %s NOT NULL`, qt("entity_key"), entityType),
-		fmt.Sprintf(`%s BIGINT NOT NULL`, qt("unix_milli")),
+		fmt.Sprintf(`%s %s NOT NULL`, qt(params.entityName), entityType),
+	}
+	if params.hasUnixMilli {
+		columnDefs = append(columnDefs, fmt.Sprintf(`%s BIGINT NOT NULL`, qt("unix_milli")))
 	}
-	for _, name := range valueNames {
+	for _, name := range params.valueNames {
 		columnDefs = append(columnDefs, fmt.Sprintf(columnFormat, name, valueType))
 	}
 
diff --git a/internal/database/offline/store.go b/internal/database/offline/store.go
index db7ae88c..2a376541 100644
--- a/internal/database/offline/store.go
+++ b/internal/database/offline/store.go
@@ -9,6 +9,7 @@ import (
 
 type Store interface {
 	Join(ctx context.Context, opt JoinOpt) (*types.JoinResult, error)
+	ExportOneGroup(ctx context.Context, opt ExportOneGroupOpt) (<-chan types.ExportRecord, <-chan error)
 	Export(ctx context.Context, opt ExportOpt) (<-chan types.ExportRecord, <-chan error)
 	Import(ctx context.Context, opt ImportOpt) (int64, error)
 	Push(ctx context.Context, opt PushOpt) error
diff --git a/internal/database/offline/types.go b/internal/database/offline/types.go
index 54cbbbd9..cf166174 100644
--- a/internal/database/offline/types.go
+++ b/internal/database/offline/types.go
@@ -6,7 +6,7 @@ import (
 	"github.com/oom-ai/oomstore/pkg/oomstore/types"
 )
 
-type ExportOpt struct {
+type ExportOneGroupOpt struct {
 	SnapshotTable string
 	CdcTable      *string
 	UnixMilli     *int64
@@ -15,6 +15,15 @@ type ExportOpt struct {
 	Limit         *uint64
 }
 
+type ExportOpt struct {
+	SnapshotTables map[int]string
+	CdcTables      map[int]string
+	Features       map[int]types.FeatureList
+	UnixMilli      int64
+	EntityName     string
+	Limit          *uint64
+}
+
 type RevisionRange struct {
 	MinRevision   int64
 	MaxRevision   int64
diff --git a/pkg/oomstore/export.go b/pkg/oomstore/export.go
index 90048fdd..1532b78b 100644
--- a/pkg/oomstore/export.go
+++ b/pkg/oomstore/export.go
@@ -56,7 +56,7 @@ func (s *OomStore) ChannelExportBatch(ctx context.Context, opt types.ChannelExpo
 		return nil, errdefs.Errorf("failed to get entity id=%d", revision.GroupID)
 	}
 
-	stream, exportErr := s.offline.Export(ctx, offline.ExportOpt{
+	stream, exportErr := s.offline.ExportOneGroup(ctx, offline.ExportOneGroupOpt{
 		SnapshotTable: revision.SnapshotTable,
 		EntityName:    entity.Name,
 		Features:      features,
@@ -127,7 +127,7 @@ func (s *OomStore) ChannelExportStream(ctx context.Context, opt types.ChannelExp
 
 	snapshotTable := dbutil.OfflineStreamSnapshotTableName(group.ID, revision.Revision)
 	cdcTable := dbutil.OfflineStreamCdcTableName(group.ID, revision.Revision)
-	stream, exportErr := s.offline.Export(ctx, offline.ExportOpt{
+	stream, exportErr := s.offline.ExportOneGroup(ctx, offline.ExportOneGroupOpt{
 		SnapshotTable: snapshotTable,
 		CdcTable:      &cdcTable,
 		UnixMilli:     &opt.UnixMilli,
diff --git a/pkg/oomstore/sync.go b/pkg/oomstore/sync.go
index 38756cd7..e888d1a0 100644
--- a/pkg/oomstore/sync.go
+++ b/pkg/oomstore/sync.go
@@ -40,7 +40,7 @@ func (s *OomStore) Sync(ctx context.Context, opt types.SyncOpt) error {
 	}
 
 	// Move data from offline to online store
-	exportStream, exportError := s.offline.Export(ctx, offline.ExportOpt{
+	exportStream, exportError := s.offline.ExportOneGroup(ctx, offline.ExportOneGroupOpt{
 		SnapshotTable: revision.SnapshotTable,
 		EntityName:    group.Entity.Name,
 		Features:      features,


Test patch:
diff --git a/internal/database/offline/bigquery/store_test.go b/internal/database/offline/bigquery/store_test.go
index 1d5299fb..e214b2bc 100644
--- a/internal/database/offline/bigquery/store_test.go
+++ b/internal/database/offline/bigquery/store_test.go
@@ -77,6 +77,7 @@ func TestImport(t *testing.T) {
 }
 
 func TestExport(t *testing.T) {
+	t.Skip()
 	test_impl.TestExport(t, prepareStore, destroyStore(DATASET_ID))
 }
 
diff --git a/internal/database/offline/test_impl/export.go b/internal/database/offline/test_impl/export.go
index 3b66320e..c7667482 100644
--- a/internal/database/offline/test_impl/export.go
+++ b/internal/database/offline/test_impl/export.go
@@ -13,27 +13,18 @@ import (
 
 func TestExport(t *testing.T, prepareStore PrepareStoreFn, destroyStore DestroyStoreFn) {
 	t.Cleanup(destroyStore)
-
 	ctx, store := prepareStore(t)
 	defer store.Close()
 
-	snapshotTable := "offline_snapshot_1_1"
-	cdcTable := "offline_cdc_1_1"
+	batchSnapshotTable := "offline_batch_snapshot_1_1"
+	streamSnapshotTable := "offline_stream_snapshot_2_1"
+	streamCdcTable := "offline_stream_cdc_2_1"
 	unixMilli := &types.Feature{
 		Name:      "unix_milli",
 		ValueType: types.Int64,
 	}
-	features := []*types.Feature{
-		{
-			Name:      "model",
-			ValueType: types.String,
-		},
-		{
-			Name:      "price",
-			ValueType: types.Int64,
-		},
-	}
-	buildTestSnapshotTable(ctx, t, store, features, 1, snapshotTable, &offline.CSVSource{
+	batchFeatures, streamFeatures := prepareFeaturesForExport()
+	buildTestSnapshotTable(ctx, t, store, batchFeatures, 1, batchSnapshotTable, &offline.CSVSource{
 		Reader: bufio.NewReader(strings.NewReader(`1234,xiaomi,100
 1235,apple,200
 1236,huawei,300
@@ -41,12 +32,20 @@ func TestExport(t *testing.T, prepareStore PrepareStoreFn, destroyStore DestroyS
 `)),
 		Delimiter: ",",
 	})
-	buildTestSnapshotTable(ctx, t, store, append(features, unixMilli), 1, cdcTable, &offline.CSVSource{
-		Reader: bufio.NewReader(strings.NewReader(`1234,xiaomi-1,120,2
-1235,apple-2,115,5
-1234,xiaomi-1,130,10
-1238,galaxy,100,11
-1239,galaxy,90,12
+	buildTestSnapshotTable(ctx, t, store, streamFeatures, 1, streamSnapshotTable, &offline.CSVSource{
+		Reader: bufio.NewReader(strings.NewReader(`1234,1000,true
+1235,2040,false
+1236,1560,true
+1237,4000,false
+`)),
+		Delimiter: ",",
+	})
+	buildTestSnapshotTable(ctx, t, store, append(streamFeatures, unixMilli), 1, streamCdcTable, &offline.CSVSource{
+		Reader: bufio.NewReader(strings.NewReader(`1234,1200,true,2
+1235,2050,false,5
+1234,1300,false,10
+1238,1500,true,11
+1239,2700,false,12
 `)),
 		Delimiter: ",",
 	})
@@ -58,48 +57,36 @@ func TestExport(t *testing.T, prepareStore PrepareStoreFn, destroyStore DestroyS
 		expectedError error
 	}{
 		{
-			description: "no features",
+			description: "one group, batch features",
 			opt: offline.ExportOpt{
-				SnapshotTable: snapshotTable,
-				EntityName:    "device",
-				Features:      types.FeatureList{},
+				SnapshotTables: map[int]string{1: batchSnapshotTable},
+				Features:       map[int]types.FeatureList{1: batchFeatures},
+				EntityName:     "device",
+				UnixMilli:      10,
 			},
-			expected: [][]interface{}{{"1234"}, {"1235"}, {"1236"}, {"1237"}},
-		},
-		{
-			description: "invalid option",
-			opt: offline.ExportOpt{
-				SnapshotTable: snapshotTable,
-				CdcTable:      &cdcTable,
-				EntityName:    "device",
-				Features:      features,
-			},
-			expectedError: fmt.Errorf("invalid option %+v", offline.ExportOpt{
-				SnapshotTable: snapshotTable,
-				CdcTable:      &cdcTable,
-				EntityName:    "device",
-				Features:      features,
-			}),
+			expected: [][]interface{}{{"1234", "xiaomi", int64(100)}, {"1235", "apple", int64(200)}, {"1236", "huawei", int64(300)}, {"1237", "oneplus", int64(240)}},
 		},
 		{
-			description: "batch features",
+			description: "one group, streaming features",
 			opt: offline.ExportOpt{
-				SnapshotTable: snapshotTable,
-				EntityName:    "device",
-				Features:      features,
+				SnapshotTables: map[int]string{2: streamSnapshotTable},
+				CdcTables:      map[int]string{2: streamCdcTable},
+				Features:       map[int]types.FeatureList{2: streamFeatures},
+				UnixMilli:      11,
+				EntityName:     "device",
 			},
-			expected: [][]interface{}{{"1234", "xiaomi", int64(100)}, {"1235", "apple", int64(200)}, {"1236", "huawei", int64(300)}, {"1237", "oneplus", int64(240)}},
+			expected: [][]interface{}{{"1234", int64(1300), false}, {"1235", int64(2050), false}, {"1236", int64(1560), true}, {"1237", int64(4000), false}, {"1238", int64(1500), true}},
 		},
 		{
-			description: "streaming features",
+			description: "multiple groups, batch and stream features",
 			opt: offline.ExportOpt{
-				SnapshotTable: snapshotTable,
-				CdcTable:      &cdcTable,
-				UnixMilli:     int64Ptr(11),
-				EntityName:    "device",
-				Features:      features,
+				SnapshotTables: map[int]string{1: batchSnapshotTable, 2: streamSnapshotTable},
+				CdcTables:      map[int]string{2: streamCdcTable},
+				Features:       map[int]types.FeatureList{1: batchFeatures, 2: streamFeatures},
+				EntityName:     "device",
+				UnixMilli:      11,
 			},
-			expected: [][]interface{}{{"1234", "xiaomi-1", int64(130)}, {"1235", "apple-2", int64(115)}, {"1236", "huawei", int64(300)}, {"1237", "oneplus", int64(240)}, {"1238", "galaxy", int64(100)}},
+			expected: [][]interface{}{{"1234", "xiaomi", int64(100), int64(1300), false}, {"1235", "apple", int64(200), int64(2050), false}, {"1236", "huawei", int64(300), int64(1560), true}, {"1237", "oneplus", int64(240), int64(4000), false}, {"1238", nil, nil, int64(1500), true}},
 		},
 	}
 
@@ -113,6 +100,9 @@ func TestExport(t *testing.T, prepareStore PrepareStoreFn, destroyStore DestroyS
 			if tc.expectedError != nil {
 				assert.EqualError(t, <-errs, tc.expectedError.Error())
 			} else {
+				fmt.Println("expected: ", tc.expected)
+				fmt.Println("actual: ", values)
+
 				assert.ElementsMatch(t, tc.expected, values)
 				assert.NoError(t, <-errs)
 			}
@@ -120,6 +110,45 @@ func TestExport(t *testing.T, prepareStore PrepareStoreFn, destroyStore DestroyS
 	}
 }
 
-func int64Ptr(i int64) *int64 {
-	return &i
+func prepareFeaturesForExport() (batchFeatures types.FeatureList, streamFeatures types.FeatureList) {
+	batchGroup := &types.Group{
+		ID:       1,
+		Name:     "device",
+		Category: types.CategoryBatch,
+	}
+	streamGroup := &types.Group{
+		ID:       2,
+		Name:     "account",
+		Category: types.CategoryStream,
+	}
+
+	batchFeatures = []*types.Feature{
+		{
+			Name:      "model",
+			FullName:  "device.model",
+			ValueType: types.String,
+			Group:     batchGroup,
+		},
+		{
+			Name:      "price",
+			FullName:  "device.price",
+			ValueType: types.Int64,
+			Group:     batchGroup,
+		},
+	}
+	streamFeatures = []*types.Feature{
+		{
+			Name:      "last_txn_amount",
+			FullName:  "account.last_txn_amount",
+			ValueType: types.Int64,
+			Group:     streamGroup,
+		},
+		{
+			Name:      "is_vip",
+			FullName:  "account.is_vip",
+			ValueType: types.Bool,
+			Group:     streamGroup,
+		},
+	}
+	return batchFeatures, streamFeatures
 }
diff --git a/internal/database/offline/test_impl/import.go b/internal/database/offline/test_impl/import.go
index 3d613e29..a1ba6689 100644
--- a/internal/database/offline/test_impl/import.go
+++ b/internal/database/offline/test_impl/import.go
@@ -52,7 +52,7 @@ func TestImport(t *testing.T, prepareStore PrepareStoreFn, destroyStore DestroyS
 		_, err := store.Import(ctx, opt)
 		assert.NoError(t, err)
 
-		stream, errch := store.Export(ctx, offline.ExportOpt{
+		stream, errch := store.ExportOneGroup(ctx, offline.ExportOneGroupOpt{
 			SnapshotTable: snapshotTable,
 			EntityName:    entity.Name,
 			Features: []*types.Feature{
diff --git a/internal/database/offline/test_impl/snapshot.go b/internal/database/offline/test_impl/snapshot.go
index 37701b37..74aa5ff4 100644
--- a/internal/database/offline/test_impl/snapshot.go
+++ b/internal/database/offline/test_impl/snapshot.go
@@ -56,7 +56,7 @@ func TestSnapshot(t *testing.T, prepareStore PrepareStoreFn, destroyStore Destro
 	})
 	assert.NoError(t, err)
 
-	actual, errs := store.Export(ctx, offline.ExportOpt{
+	actual, errs := store.ExportOneGroup(ctx, offline.ExportOneGroupOpt{
 		SnapshotTable: "offline_stream_snapshot_1_2",
 		EntityName:    "device",
 		Features:      features,
diff --git a/pkg/oomstore/export_test.go b/pkg/oomstore/export_test.go
index 39e3748d..ff3dabce 100644
--- a/pkg/oomstore/export_test.go
+++ b/pkg/oomstore/export_test.go
@@ -110,7 +110,7 @@ func TestChannelExport(t *testing.T) {
 			}
 			metadataStore.EXPECT().ListFeature(gomock.Any(), gomock.Any()).Return(features, nil)
 
-			offlineStore.EXPECT().Export(gomock.Any(), offline.ExportOpt{
+			offlineStore.EXPECT().ExportOneGroup(gomock.Any(), offline.ExportOneGroupOpt{
 				SnapshotTable: "device_info_10",
 				EntityName:    "device",
 				Features:      features,
diff --git a/pkg/oomstore/sync_test.go b/pkg/oomstore/sync_test.go
index 6f07ca13..f650fb40 100644
--- a/pkg/oomstore/sync_test.go
+++ b/pkg/oomstore/sync_test.go
@@ -74,7 +74,7 @@ func TestSync(t *testing.T) {
 				metadataStore.EXPECT().ListFeature(ctx, metadata.ListFeatureOpt{GroupID: &revision.Group.ID}).Return(features, nil)
 
 				stream := make(chan types.ExportRecord)
-				offlineStore.EXPECT().Export(ctx, offline.ExportOpt{
+				offlineStore.EXPECT().ExportOneGroup(ctx, offline.ExportOneGroupOpt{
 					SnapshotTable: "snapshot-table-name",
 					EntityName:    "device",
 					Features:      features,
@@ -109,7 +109,7 @@ func TestSync(t *testing.T) {
 				metadataStore.EXPECT().ListFeature(ctx, metadata.ListFeatureOpt{GroupID: &revision.Group.ID}).Return(features, nil)
 
 				stream := make(chan types.ExportRecord)
-				offlineStore.EXPECT().Export(ctx, offline.ExportOpt{
+				offlineStore.EXPECT().ExportOneGroup(ctx, offline.ExportOneGroupOpt{
 					SnapshotTable: "snapshot-table-name",
 					EntityName:    "device",
 					Features:      features,