Repository: oom-ai/oomstore
Pull request #995 (2022-01-17 07:34:39)

Jinghan/Add `unix_milli` column to snapshot table
This PR adds `unix_milli` column to streaming feature's snapshot table.

TODO:

- [ ] refactor offline method `CreateTable`: handle `unix_milli` issue separately


related to #994 


Patch:
diff --git a/internal/database/offline/sqlutil/snapshot.go b/internal/database/offline/sqlutil/snapshot.go
index f5c9ab90..7f70963b 100644
--- a/internal/database/offline/sqlutil/snapshot.go
+++ b/internal/database/offline/sqlutil/snapshot.go
@@ -15,9 +15,10 @@ import (
 )
 
 const SNAPSHOT_QUERY = `
-INSERT INTO {{ qt .CurrSnapshotTableName }} ({{ qt .EntityName }}, {{ columnJoin .FeatureNames }})
+INSERT INTO {{ qt .CurrSnapshotTableName }} ({{ qt .EntityName }}, {{ qt .UnixMilli }}, {{ columnJoin .FeatureNames }})
 SELECT
 	l.{{ qt .EntityName }},
+	{{ unixMilli .UnixMilli }},
 	{{ featureValue .FeatureNames }}
 FROM {{ qt .PrevSnapshotTableName }} AS l
 LEFT JOIN
@@ -38,6 +39,7 @@ ON l.{{ qt .EntityName }} = r.{{ qt .EntityName }}
 {{ .Union }}
 SELECT
 	r.{{ qt .EntityName }},
+	{{ unixMilli .UnixMilli }},
 	{{ featureValue .FeatureNames }}
 FROM
 (
@@ -69,7 +71,7 @@ func Snapshot(ctx context.Context, dbOpt dbutil.DBOpt, opt offline.SnapshotOpt)
 		currCdcTableName = fmt.Sprintf("%s.%s", *dbOpt.DatasetID, currCdcTableName)
 	}
 
-	schema := dbutil.BuildTableSchema(currSnapshotTableName, opt.Group.Entity, false, opt.Features, []string{opt.Group.Entity.Name}, dbOpt.Backend)
+	schema := dbutil.BuildTableSchema(currSnapshotTableName, opt.Group.Entity, true, opt.Features, []string{opt.Group.Entity.Name}, dbOpt.Backend)
 	if err := dbOpt.ExecContext(ctx, schema, nil); err != nil {
 		return err
 	}
@@ -116,6 +118,9 @@ func buildSnapshotQuery(params snapshotQueryParams) (string, error) {
 		"columnJoin": func(columns []string) string {
 			return qt(columns...)
 		},
+		"unixMilli": func(unixMilli string) string {
+			return fmt.Sprintf("(CASE WHEN r.%s IS NULL THEN l.%s ELSE r.%s END) AS %s", qt(params.FeatureNames[0]), qt(unixMilli), qt(unixMilli), unixMilli)
+		},
 		"featureValue": func(features []string) string {
 			values := make([]string, 0, len(features))
 			for _, f := range features {
diff --git a/pkg/oomstore/revision.go b/pkg/oomstore/revision.go
index 6cc57eaa..0f35ff4e 100644
--- a/pkg/oomstore/revision.go
+++ b/pkg/oomstore/revision.go
@@ -47,6 +47,7 @@ func (s *OomStore) createFirstSnapshotTable(ctx context.Context, revision *types
 		TableName: snapshotTable,
 		Entity:    revision.Group.Entity,
 		Features:  features,
+		IsCDC:     true, // In order to add `unix_milli` column, set IsCDC as true temporarily
 	}); err != nil {
 		return err
 	}


Test patch:
diff --git a/internal/database/offline/test_impl/join.go b/internal/database/offline/test_impl/join.go
index 933a86d7..c03c98f7 100644
--- a/internal/database/offline/test_impl/join.go
+++ b/internal/database/offline/test_impl/join.go
@@ -147,11 +147,13 @@ func prepareFeatures(oneGroup bool) (types.FeatureList, map[string]types.Feature
 	groupBasic := &types.Group{
 		ID:       1,
 		Name:     "device_basic",
+		Entity:   &types.Entity{Name: "device"},
 		Category: types.CategoryBatch,
 	}
 	groupAdvanced := &types.Group{
 		ID:       2,
 		Name:     "device_advanced",
+		Entity:   &types.Entity{Name: "device"},
 		Category: types.CategoryBatch,
 	}
 	price := &types.Feature{
diff --git a/internal/database/offline/test_impl/snapshot.go b/internal/database/offline/test_impl/snapshot.go
index 3ef7ef93..fe15e0b3 100644
--- a/internal/database/offline/test_impl/snapshot.go
+++ b/internal/database/offline/test_impl/snapshot.go
@@ -20,20 +20,18 @@ func TestSnapshot(t *testing.T, prepareStore PrepareStoreFn, destroyStore Destro
 	ctx, store := prepareStore(t)
 	defer store.Close()
 
-	group := &types.Group{
-		ID:     1,
-		Entity: &types.Entity{Name: "device"},
-	}
+	features, _ := prepareFeatures(true)
+	group := features[0].Group
 	unixMilli := &types.Feature{
 		Name:      "unix_milli",
 		ValueType: types.Int64,
+		Group:     group,
 	}
-	features, _ := prepareFeatures(true)
 
-	buildTestSnapshotTable(ctx, t, store, features, 1, "offline_stream_snapshot_1_1", &offline.CSVSource{
-		Reader: bufio.NewReader(strings.NewReader(`1234,xiaomi,100
-1235,apple,200
-1236,oneplus,155
+	buildTestSnapshotTable(ctx, t, store, append(features, unixMilli), 1, "offline_stream_snapshot_1_1", &offline.CSVSource{
+		Reader: bufio.NewReader(strings.NewReader(`1234,xiaomi,100,1
+1235,apple,200,5
+1236,oneplus,155,8
 `)),
 		Delimiter: ",",
 	})
@@ -59,7 +57,7 @@ func TestSnapshot(t *testing.T, prepareStore PrepareStoreFn, destroyStore Destro
 	result, err := store.Export(ctx, offline.ExportOpt{
 		SnapshotTables: map[int]string{1: "offline_stream_snapshot_1_2"},
 		EntityName:     "device",
-		Features:       map[int]types.FeatureList{1: features},
+		Features:       map[int]types.FeatureList{1: append(features, unixMilli)},
 	})
 	values := make([][]interface{}, 0)
 	for row := range result.Data {
@@ -69,10 +67,10 @@ func TestSnapshot(t *testing.T, prepareStore PrepareStoreFn, destroyStore Destro
 		return cast.ToInt64(values[i][0]) < cast.ToInt64(values[j][0])
 	})
 	expected := [][]interface{}{
-		{"1234", "xiaomi-1", int64(130)},
-		{"1235", "apple-1", int64(113)},
-		{"1236", "oneplus", int64(155)},
-		{"1237", "pixel", int64(200)},
+		{"1234", "xiaomi-1", int64(130), int64(10)},
+		{"1235", "apple-1", int64(113), int64(15)},
+		{"1236", "oneplus", int64(155), int64(8)},
+		{"1237", "pixel", int64(200), int64(11)},
 	}
 	assert.Equal(t, expected, values)
 	assert.NoError(t, result.CheckStreamError())