Repository: oom-ai/oomstore
Pull request #826 (2021-12-27 08:04:15)

Jinghan/use `FeatureFullName` for API Join
This PR refactors API `Join`: use `FeatureFullName` instead of `FeatureName`.

related to #811 


TODO:

- [ ] Deserialize `Join` output data via value_type, unit test does not work


Patch:
diff --git a/internal/database/offline/bigquery/join.go b/internal/database/offline/bigquery/join.go
index cd308c6c..9b863764 100644
--- a/internal/database/offline/bigquery/join.go
+++ b/internal/database/offline/bigquery/join.go
@@ -3,6 +3,7 @@ package bigquery
 import (
 	"context"
 	"fmt"
+	"strings"
 
 	"github.com/oom-ai/oomstore/internal/database/offline/sqlutil"
 
@@ -55,7 +56,8 @@ func bigqueryQueryResults(ctx context.Context, dbOpt dbutil.DBOpt, query string,
 			}
 			record := make([]interface{}, 0, len(recordMap))
 			for _, h := range header {
-				record = append(record, recordMap[h])
+				column := strings.Split(h, ".")
+				record = append(record, recordMap[column[len(column)-1]])
 			}
 			data <- record
 		}
diff --git a/internal/database/offline/sqlutil/join.go b/internal/database/offline/sqlutil/join.go
index 63674f98..6dc1954d 100644
--- a/internal/database/offline/sqlutil/join.go
+++ b/internal/database/offline/sqlutil/join.go
@@ -170,7 +170,7 @@ func ReadJoinedTable(ctx context.Context, dbOpt dbutil.DBOpt, opt ReadJoinedTabl
 	for _, name := range tableNames {
 		for _, f := range opt.FeatureMap[name] {
 			fields = append(fields, fmt.Sprintf("%s.%s", qt(name), qt(f.Name)))
-			header = append(header, f.Name)
+			header = append(header, f.FullName)
 		}
 	}
 
diff --git a/oomagent/server.go b/oomagent/server.go
index dba1861f..fabeb76f 100644
--- a/oomagent/server.go
+++ b/oomagent/server.go
@@ -206,9 +206,9 @@ func (s *server) ChannelJoin(stream codegen.OomAgent_ChannelJoinServer) error {
 	// This goroutine runs the join operation, and send whatever joined as the response
 	go func() {
 		joinResult, err := s.oomstore.ChannelJoin(context.Background(), types.ChannelJoinOpt{
-			FeatureNames: firstReq.FeatureNames,
-			EntityRows:   entityRows,
-			ValueNames:   firstReq.ValueNames,
+			FeatureFullNames: firstReq.FeatureNames,
+			EntityRows:       entityRows,
+			ValueNames:       firstReq.ValueNames,
 		})
 		if err != nil {
 			globalErr = err
@@ -276,9 +276,9 @@ func (s *server) ChannelJoin(stream codegen.OomAgent_ChannelJoinServer) error {
 
 func (s *server) Join(ctx context.Context, req *codegen.JoinRequest) (*codegen.JoinResponse, error) {
 	err := s.oomstore.Join(ctx, types.JoinOpt{
-		FeatureNames:   req.FeatureNames,
-		InputFilePath:  req.InputFilePath,
-		OutputFilePath: req.OutputFilePath,
+		FeatureFullNames: req.FeatureNames,
+		InputFilePath:    req.InputFilePath,
+		OutputFilePath:   req.OutputFilePath,
 	})
 	if err != nil {
 		return &codegen.JoinResponse{
diff --git a/oomcli/cmd/join_helper.go b/oomcli/cmd/join_helper.go
index 61cb983f..2df91d94 100644
--- a/oomcli/cmd/join_helper.go
+++ b/oomcli/cmd/join_helper.go
@@ -24,9 +24,9 @@ func join(ctx context.Context, store *oomstore.OomStore, opt JoinOpt, output str
 	}
 
 	joinResult, err := store.ChannelJoin(ctx, types.ChannelJoinOpt{
-		FeatureNames: opt.FeatureNames,
-		EntityRows:   entityRows,
-		ValueNames:   header[2:],
+		FeatureFullNames: opt.FeatureNames,
+		EntityRows:       entityRows,
+		ValueNames:       header[2:],
 	})
 	if err != nil {
 		return err
diff --git a/pkg/oomstore/join.go b/pkg/oomstore/join.go
index a578e1ab..17514a35 100644
--- a/pkg/oomstore/join.go
+++ b/pkg/oomstore/join.go
@@ -16,11 +16,11 @@ import (
 	"github.com/spf13/cast"
 )
 
-// Get point-in-time correct feature values for each entity row.
+// ChannelJoin gets point-in-time correct feature values for each entity row.
 // Currently, this API only supports batch features.
 func (s *OomStore) ChannelJoin(ctx context.Context, opt types.ChannelJoinOpt) (*types.JoinResult, error) {
 	features, err := s.metadata.ListFeature(ctx, metadata.ListFeatureOpt{
-		FeatureFullNames: &opt.FeatureNames,
+		FeatureFullNames: &opt.FeatureFullNames,
 	})
 	if err != nil {
 		return nil, err
@@ -63,7 +63,7 @@ func (s *OomStore) ChannelJoin(ctx context.Context, opt types.ChannelJoinOpt) (*
 	})
 }
 
-// Get point-in-time correct feature values for each entity row.
+// Join gets point-in-time correct feature values for each entity row.
 // The method is similar to Join, except that both input and output are files on disk.
 // Input File should contain header, the first two columns of Input File should be
 // entity_key, unix_milli, then followed by other real-time feature values.
@@ -74,9 +74,9 @@ func (s *OomStore) Join(ctx context.Context, opt types.JoinOpt) error {
 	}
 
 	joinResult, err := s.ChannelJoin(ctx, types.ChannelJoinOpt{
-		FeatureNames: opt.FeatureNames,
-		EntityRows:   entityRows,
-		ValueNames:   header[2:],
+		FeatureFullNames: opt.FeatureFullNames,
+		EntityRows:       entityRows,
+		ValueNames:       header[2:],
 	})
 	if err != nil {
 		return err
diff --git a/pkg/oomstore/types/options.go b/pkg/oomstore/types/options.go
index c638d76e..ce89df74 100644
--- a/pkg/oomstore/types/options.go
+++ b/pkg/oomstore/types/options.go
@@ -55,15 +55,15 @@ type OnlineMultiGetOpt struct {
 }
 
 type ChannelJoinOpt struct {
-	FeatureNames []string
-	EntityRows   <-chan EntityRow
-	ValueNames   []string
+	FeatureFullNames []string
+	EntityRows       <-chan EntityRow
+	ValueNames       []string
 }
 
 type JoinOpt struct {
-	FeatureNames   []string
-	InputFilePath  string
-	OutputFilePath string
+	FeatureFullNames []string
+	InputFilePath    string
+	OutputFilePath   string
 }
 
 type UpdateEntityOpt struct {


Test patch:
diff --git a/internal/database/offline/test_impl/join.go b/internal/database/offline/test_impl/join.go
index 20b13425..eb137327 100644
--- a/internal/database/offline/test_impl/join.go
+++ b/internal/database/offline/test_impl/join.go
@@ -110,7 +110,7 @@ func TestJoin(t *testing.T, prepareStore PrepareStoreFn, destoryStore DestroySto
 			} else {
 				expectedValues := extractValues(tc.expected.Data)
 				actualValues := extractValues(actual.Data)
-				assert.ObjectsAreEqual(tc.expected.Header, actual.Header)
+				assert.ElementsMatch(t, tc.expected.Header, actual.Header)
 				assert.ObjectsAreEqual(expectedValues, actualValues)
 			}
 		})
@@ -120,16 +120,19 @@ func TestJoin(t *testing.T, prepareStore PrepareStoreFn, destoryStore DestroySto
 func prepareFeatures(oneGroup bool) (types.FeatureList, map[string]types.FeatureList) {
 	price := &types.Feature{
 		Name:      "price",
+		FullName:  "device_basic.price",
 		ValueType: types.Int64,
 		GroupID:   1,
 	}
 	model := &types.Feature{
 		Name:      "model",
+		FullName:  "device_basic.model",
 		ValueType: types.String,
 		GroupID:   1,
 	}
 	isActive := &types.Feature{
 		Name:      "is_active",
+		FullName:  "device_advanced.is_active",
 		ValueType: types.Bool,
 		GroupID:   2,
 	}
@@ -225,46 +228,46 @@ func prepareEntityRows(isEmpty bool, withValue bool) <-chan types.EntityRow {
 }
 
 func prepareResult(oneGroup bool, withValue bool) *types.JoinResult {
-	header := []string{"entity_key", "unix_milli", "model", "price", "is_active"}
+	header := []string{"entity_key", "unix_milli", "device_basic.model", "device_basic.price", "device_advanced.is_active"}
 	if withValue {
-		header = []string{"entity_key", "unix_milli", "value_1", "value_2", "model", "price", "is_active"}
+		header = []string{"entity_key", "unix_milli", "value_1", "value_2", "device_basic.model", "device_basic.price", "device_advanced.is_active"}
 	}
 	values := []map[string]interface{}{
 		{
-			"entity_key": "1234",
-			"unix_milli": int64(10),
-			"value_1":    1,
-			"value_2":    2,
-			"model":      "xiaomi",
-			"price":      int64(100),
-			"is_active":  true,
+			"entity_key":                "1234",
+			"unix_milli":                int64(10),
+			"value_1":                   "1",
+			"value_2":                   "2",
+			"device_basic.model":        "xiaomi",
+			"device_basic.price":        int64(100),
+			"device_advanced.is_active": true,
 		},
 		{
-			"entity_key": "1234",
-			"unix_milli": int64(20),
-			"value_1":    3,
-			"value_2":    4,
-			"model":      "galaxy",
-			"price":      int64(300),
-			"is_active":  true,
+			"entity_key":                "1234",
+			"unix_milli":                int64(20),
+			"value_1":                   "3",
+			"value_2":                   "4",
+			"device_basic.model":        "galaxy",
+			"device_basic.price":        int64(300),
+			"device_advanced.is_active": true,
 		},
 		{
-			"entity_key": "1235",
-			"unix_milli": int64(5),
-			"value_1":    5,
-			"value_2":    6,
-			"model":      "apple",
-			"price":      int64(200),
-			"is_active":  false,
+			"entity_key":                "1235",
+			"unix_milli":                int64(5),
+			"value_1":                   "5",
+			"value_2":                   "6",
+			"device_basic.model":        "apple",
+			"device_basic.price":        int64(200),
+			"device_advanced.is_active": false,
 		},
 		{
-			"entity_key": "1235",
-			"unix_milli": int64(15),
-			"value_1":    7,
-			"value_2":    8,
-			"model":      "oneplus",
-			"price":      int64(240),
-			"is_active":  false,
+			"entity_key":                "1235",
+			"unix_milli":                int64(14),
+			"value_1":                   "7",
+			"value_2":                   "8",
+			"device_basic.model":        "apple",
+			"device_basic.price":        int64(200),
+			"device_advanced.is_active": false,
 		},
 	}
 	if oneGroup {
diff --git a/oomcli/test/test_join.sh b/oomcli/test/test_join.sh
index 9ad177ff..ec79c7a0 100755
--- a/oomcli/test/test_join.sh
+++ b/oomcli/test/test_join.sh
@@ -23,14 +23,14 @@ EOF
 
 case='oomcli join historical-feature'
 expected="
-entity_key,unix_milli,price,model
+entity_key,unix_milli,phone.price,phone.model
 1,$t1,,
 2,$t1,,
 1,$t2,3999,xiaomi-mix3
 2,$t2,5299,huawei-p40
 "
 actual=$(oomcli join \
-    --feature model,price \
+    --feature phone.model,phone.price \
     --input-file entity_rows.csv \
     --output csv
 )
@@ -51,7 +51,7 @@ EOF
 
 case='oomcli join historical-feature with real-time feature values'
 expected="
-entity_key,unix_milli,value_1,value_2,price,model
+entity_key,unix_milli,value_1,value_2,phone.price,phone.model
 1,$t1,1,2,,
 2,$t1,3,4,,
 1,$t2,5,6,3999,xiaomi-mix3
@@ -59,7 +59,7 @@ entity_key,unix_milli,value_1,value_2,price,model
 "
 
 actual=$(oomcli join \
-    --feature model,price \
+    --feature phone.model,phone.price \
     --input-file entity_rows_with_values.csv \
     --output csv
 )
diff --git a/pkg/oomstore/join_test.go b/pkg/oomstore/join_test.go
index d2194f74..f0dc8083 100644
--- a/pkg/oomstore/join_test.go
+++ b/pkg/oomstore/join_test.go
@@ -57,8 +57,8 @@ func TestChannelJoin(t *testing.T) {
 		{
 			description: "no valid features, return nil",
 			opt: types.ChannelJoinOpt{
-				FeatureNames: streamFeatures.Names(),
-				EntityRows:   entityRows,
+				FeatureFullNames: streamFeatures.FullNames(),
+				EntityRows:       entityRows,
 			},
 			features:      streamFeatures,
 			expectedError: nil,
@@ -67,8 +67,8 @@ func TestChannelJoin(t *testing.T) {
 		{
 			description: "inconsistent features, return nil",
 			opt: types.ChannelJoinOpt{
-				FeatureNames: inconsistentFeatures.Names(),
-				EntityRows:   entityRows,
+				FeatureFullNames: inconsistentFeatures.FullNames(),
+				EntityRows:       entityRows,
 			},
 			features:      inconsistentFeatures,
 			expectedError: fmt.Errorf("inconsistent entity type: %v", map[string]string{"device": "price", "user": "age"}),
@@ -77,8 +77,8 @@ func TestChannelJoin(t *testing.T) {
 		{
 			description: "nil joined, return nil",
 			opt: types.ChannelJoinOpt{
-				FeatureNames: consistentFeatures.Names(),
-				EntityRows:   entityRows,
+				FeatureFullNames: consistentFeatures.FullNames(),
+				EntityRows:       entityRows,
 			},
 			entity:   &entity,
 			features: consistentFeatures,
@@ -93,8 +93,8 @@ func TestChannelJoin(t *testing.T) {
 		{
 			description: "consistent entity type, succeed",
 			opt: types.ChannelJoinOpt{
-				FeatureNames: consistentFeatures.Names(),
-				EntityRows:   entityRows,
+				FeatureFullNames: consistentFeatures.FullNames(),
+				EntityRows:       entityRows,
 			},
 			entity:   &entity,
 			features: consistentFeatures,
@@ -111,7 +111,7 @@ func TestChannelJoin(t *testing.T) {
 	for _, tc := range testCases {
 		t.Run(tc.description, func(t *testing.T) {
 			metadataStore.EXPECT().Refresh().Return(nil).AnyTimes()
-			metadataStore.EXPECT().ListFeature(gomock.Any(), metadata.ListFeatureOpt{FeatureFullNames: &tc.opt.FeatureNames}).Return(tc.features, nil)
+			metadataStore.EXPECT().ListFeature(gomock.Any(), metadata.ListFeatureOpt{FeatureFullNames: &tc.opt.FeatureFullNames}).Return(tc.features, nil)
 			if tc.entity != nil {
 				for _, featureList := range tc.featureMap {
 					metadataStore.EXPECT().ListRevision(gomock.Any(), &featureList[0].GroupID).Return(revisions, nil).AnyTimes()