Repository: oom-ai/oomstore
Pull request #938 (2022-01-11 09:10:28)

Jinghan/modify offline method `Export`
This PR modifies offline method `Export` to make it support streaming features.

TODO:

- [ ] modify bigquery's `Export` method

related to #802 


Patch:
diff --git a/internal/database/offline/sqlutil/export.go b/internal/database/offline/sqlutil/export.go
index 77d8f0ac..6f5e53f1 100644
--- a/internal/database/offline/sqlutil/export.go
+++ b/internal/database/offline/sqlutil/export.go
@@ -12,28 +12,30 @@ import (
 	"github.com/spf13/cast"
 )
 
-func Export(ctx context.Context, db *sqlx.DB, opt offline.ExportOpt, backendType types.BackendType) (<-chan types.ExportRecord, <-chan error) {
+func Export(ctx context.Context, db *sqlx.DB, opt offline.ExportOpt, backend types.BackendType) (<-chan types.ExportRecord, <-chan error) {
 	var (
-		fields = append([]string{opt.EntityName}, opt.Features.Names()...)
 		stream = make(chan types.ExportRecord)
 		errs   = make(chan error, 1) // at most 1 error
 	)
-	qt := dbutil.QuoteFn(backendType)
-	query := fmt.Sprintf("SELECT %s FROM %s", qt(fields...), qt(opt.SnapshotTable))
-	if opt.Limit != nil {
-		query += fmt.Sprintf(" LIMIT %d", *opt.Limit)
+
+	query, args, err := buildExportQuery(opt, backend)
+	if err != nil {
+		defer close(stream)
+		defer close(errs)
+		errs <- errors.WithStack(err)
+		return stream, errs
 	}
 
 	go func() {
 		defer close(stream)
 		defer close(errs)
-		stmt, err := db.Preparex(query)
+		stmt, err := db.Preparex(db.Rebind(query))
 		if err != nil {
 			errs <- errors.WithStack(err)
 			return
 		}
 		defer stmt.Close()
-		rows, err := stmt.Queryx()
+		rows, err := stmt.Queryx(args...)
 		if err != nil {
 			errs <- errors.WithStack(err)
 			return
@@ -50,7 +52,7 @@ func Export(ctx context.Context, db *sqlx.DB, opt offline.ExportOpt, backendType
 				if record[i+1] == nil {
 					continue
 				}
-				deserializedValue, err := dbutil.DeserializeByValueType(record[i+1], f.ValueType, backendType)
+				deserializedValue, err := dbutil.DeserializeByValueType(record[i+1], f.ValueType, backend)
 				if err != nil {
 					errs <- err
 					return
@@ -63,3 +65,42 @@ func Export(ctx context.Context, db *sqlx.DB, opt offline.ExportOpt, backendType
 
 	return stream, errs
 }
+
+func buildExportQuery(opt offline.ExportOpt, backend types.BackendType) (string, []interface{}, error) {
+	if opt.CdcTable == nil && opt.UnixMilli == nil {
+		return buildExportBatchQuery(opt, backend), nil, nil
+	}
+	if opt.CdcTable != nil && opt.UnixMilli != nil {
+		return buildExportStreamQuery(opt, backend)
+	}
+	return "", nil, fmt.Errorf("invalid option %+v", opt)
+}
+
+func buildExportBatchQuery(opt offline.ExportOpt, backend types.BackendType) string {
+	fields := append([]string{opt.EntityName}, opt.Features.Names()...)
+	qt := dbutil.QuoteFn(backend)
+
+	query := fmt.Sprintf("SELECT %s FROM %s", qt(fields...), qt(opt.SnapshotTable))
+	if opt.Limit != nil {
+		query += fmt.Sprintf(" LIMIT %d", *opt.Limit)
+	}
+	return query
+}
+
+func buildExportStreamQuery(opt offline.ExportOpt, backend types.BackendType) (string, []interface{}, error) {
+	query, err := buildAggregateQuery(aggregateQueryParams{
+		EntityName:            opt.EntityName,
+		UnixMilli:             "unix_milli",
+		FeatureNames:          opt.Features.Names(),
+		PrevSnapshotTableName: opt.SnapshotTable,
+		CurrCdcTableName:      *opt.CdcTable,
+		Backend:               backend,
+	})
+	if err != nil {
+		return "", nil, errors.WithStack(err)
+	}
+	if opt.Limit != nil {
+		query += fmt.Sprintf(" LIMIT %d", *opt.Limit)
+	}
+	return query, []interface{}{*opt.UnixMilli, *opt.UnixMilli, *opt.UnixMilli, *opt.UnixMilli}, nil
+}
diff --git a/internal/database/offline/sqlutil/export_helper.go b/internal/database/offline/sqlutil/export_helper.go
new file mode 100644
index 00000000..7467008c
--- /dev/null
+++ b/internal/database/offline/sqlutil/export_helper.go
@@ -0,0 +1,98 @@
+package sqlutil
+
+import (
+	"bytes"
+	"fmt"
+	"strings"
+	"text/template"
+
+	"github.com/oom-ai/oomstore/internal/database/dbutil"
+	"github.com/oom-ai/oomstore/pkg/oomstore/types"
+	"github.com/pkg/errors"
+)
+
+const AGGREGATE_QUERY = `
+SELECT
+	l.{{ qt .EntityName }},
+	{{ featureValue .FeatureNames }}
+FROM {{ qt .PrevSnapshotTableName }} AS l
+LEFT JOIN
+(
+	SELECT
+		t1.*
+	FROM
+		{{ qt .CurrCdcTableName }} AS t1
+	JOIN
+		(SELECT
+			{{ qt .EntityName }},
+			MAX({{ qt .UnixMilli }}) AS {{ qt .UnixMilli }}
+		FROM {{ qt .CurrCdcTableName }}
+		WHERE {{ qt .UnixMilli }} <= ?
+		GROUP BY {{ qt .EntityName }}) AS t2
+	ON t1.{{ qt .EntityName }} = t2.{{ qt .EntityName }} AND t1.{{ qt .UnixMilli }} = t2.{{ qt .UnixMilli }}
+	WHERE t1.{{ qt .UnixMilli }} <= ?
+) AS r
+ON l.{{ qt .EntityName }} = r.{{ qt .EntityName }}
+{{ .Union }}
+SELECT
+	r.{{ qt .EntityName }},
+	{{ featureValue .FeatureNames }}
+FROM
+(
+	SELECT
+		t1.*
+	FROM
+		{{ qt .CurrCdcTableName }} AS t1
+	JOIN
+		(SELECT
+			{{ qt .EntityName }},
+			MAX({{ qt .UnixMilli }}) AS {{ qt .UnixMilli }}
+		FROM {{ qt .CurrCdcTableName }}
+		WHERE {{ qt .UnixMilli }} <= ?
+		GROUP BY {{ qt .EntityName }}) AS t2
+	ON t1.{{ qt .EntityName }} = t2.{{ qt .EntityName }} AND t1.{{ qt .UnixMilli }} = t2.{{ qt .UnixMilli }}
+	WHERE t1.{{ qt .UnixMilli }} <= ?
+) AS r
+LEFT JOIN
+{{ qt .PrevSnapshotTableName }} AS l
+ON l.{{ qt .EntityName }} = r.{{ qt .EntityName }}
+`
+
+type aggregateQueryParams struct {
+	EntityName            string
+	UnixMilli             string
+	FeatureNames          []string
+	PrevSnapshotTableName string
+	CurrCdcTableName      string
+	Union                 string
+	Backend               types.BackendType
+	DatasetID             *string
+}
+
+func buildAggregateQuery(params aggregateQueryParams) (string, error) {
+	if params.Backend == types.BackendBigQuery {
+		params.Union = "UNION DISTINCT"
+	} else {
+		params.Union = "UNION"
+	}
+	qt := dbutil.QuoteFn(params.Backend)
+	t := template.Must(template.New("snapshot").Funcs(template.FuncMap{
+		"qt": qt,
+		"columnJoin": func(columns []string) string {
+			return qt(columns...)
+		},
+		"featureValue": func(features []string) string {
+			values := make([]string, 0, len(features))
+			for _, f := range features {
+				values = append(values, fmt.Sprintf("(CASE WHEN r.%s IS NULL THEN l.%s ELSE r.%s END) AS %s", qt(f), qt(f), qt(f), f))
+			}
+			return strings.Join(values, ",")
+		},
+	}).Parse(AGGREGATE_QUERY))
+
+	buf := bytes.NewBuffer(nil)
+	if err := t.Execute(buf, params); err != nil {
+		return "", errors.WithStack(err)
+	}
+	return buf.String(), nil
+}


Test patch:
diff --git a/internal/database/offline/bigquery/store_test.go b/internal/database/offline/bigquery/store_test.go
index 1d5299fb..22f21907 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("skip export test for bigquery until we support exporting stream features")
 	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 265acfab..3b66320e 100644
--- a/internal/database/offline/test_impl/export.go
+++ b/internal/database/offline/test_impl/export.go
@@ -2,6 +2,7 @@ package test_impl
 
 import (
 	"bufio"
+	"fmt"
 	"strings"
 	"testing"
 
@@ -16,7 +17,12 @@ func TestExport(t *testing.T, prepareStore PrepareStoreFn, destroyStore DestroyS
 	ctx, store := prepareStore(t)
 	defer store.Close()
 
-	snapshotTable := "offline_1_1"
+	snapshotTable := "offline_snapshot_1_1"
+	cdcTable := "offline_cdc_1_1"
+	unixMilli := &types.Feature{
+		Name:      "unix_milli",
+		ValueType: types.Int64,
+	}
 	features := []*types.Feature{
 		{
 			Name:      "model",
@@ -32,14 +38,24 @@ func TestExport(t *testing.T, prepareStore PrepareStoreFn, destroyStore DestroyS
 1235,apple,200
 1236,huawei,300
 1237,oneplus,240
+`)),
+		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
 `)),
 		Delimiter: ",",
 	})
 
 	testCases := []struct {
-		description string
-		opt         offline.ExportOpt
-		expected    [][]interface{}
+		description   string
+		opt           offline.ExportOpt
+		expected      [][]interface{}
+		expectedError error
 	}{
 		{
 			description: "no features",
@@ -51,7 +67,22 @@ func TestExport(t *testing.T, prepareStore PrepareStoreFn, destroyStore DestroyS
 			expected: [][]interface{}{{"1234"}, {"1235"}, {"1236"}, {"1237"}},
 		},
 		{
-			description: "valid features and valid entity rows",
+			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,
+			}),
+		},
+		{
+			description: "batch features",
 			opt: offline.ExportOpt{
 				SnapshotTable: snapshotTable,
 				EntityName:    "device",
@@ -59,6 +90,17 @@ func TestExport(t *testing.T, prepareStore PrepareStoreFn, destroyStore DestroyS
 			},
 			expected: [][]interface{}{{"1234", "xiaomi", int64(100)}, {"1235", "apple", int64(200)}, {"1236", "huawei", int64(300)}, {"1237", "oneplus", int64(240)}},
 		},
+		{
+			description: "streaming features",
+			opt: offline.ExportOpt{
+				SnapshotTable: snapshotTable,
+				CdcTable:      &cdcTable,
+				UnixMilli:     int64Ptr(11),
+				EntityName:    "device",
+				Features:      features,
+			},
+			expected: [][]interface{}{{"1234", "xiaomi-1", int64(130)}, {"1235", "apple-2", int64(115)}, {"1236", "huawei", int64(300)}, {"1237", "oneplus", int64(240)}, {"1238", "galaxy", int64(100)}},
+		},
 	}
 
 	for _, tc := range testCases {
@@ -68,8 +110,16 @@ func TestExport(t *testing.T, prepareStore PrepareStoreFn, destroyStore DestroyS
 			for row := range actual {
 				values = append(values, row)
 			}
-			assert.Equal(t, tc.expected, values)
-			assert.NoError(t, <-errs)
+			if tc.expectedError != nil {
+				assert.EqualError(t, <-errs, tc.expectedError.Error())
+			} else {
+				assert.ElementsMatch(t, tc.expected, values)
+				assert.NoError(t, <-errs)
+			}
 		})
 	}
 }
+
+func int64Ptr(i int64) *int64 {
+	return &i
+}