fix: fix compile

oom-ai · Nov 12, 2021 · 072cfe5 · 072cfe5
1 parent 7c6d386
commit 072cfe5
Show file tree

Hide file tree

Showing 3 changed files with 6 additions and 19 deletions.
diff --git a/internal/database/online/types.go b/internal/database/online/types.go
@@ -1,6 +1,7 @@
 package online
 
 import (
+	"github.com/oom-ai/oomstore/pkg/oomstore/types"
 	"github.com/oom-ai/oomstore/pkg/oomstore/typesv2"
 )
 
@@ -21,6 +22,6 @@ type MultiGetOpt struct {
 type ImportOpt struct {
 	Revision    *typesv2.Revision
 	Entity      *typesv2.Entity
-	Stream      <-chan *typesv2.RawFeatureValueRecord
+	Stream      <-chan *types.RawFeatureValueRecord
 	FeatureList typesv2.FeatureList
 }
diff --git a/pkg/oomstore/online_feature_value.go b/pkg/oomstore/online_feature_value.go
@@ -33,7 +33,7 @@ func (s *OomStore) GetOnlineFeatureValues(ctx context.Context, opt types.GetOnli
 			continue
 		}
 		featureValues, err := s.online.Get(ctx, online.GetOpt{
-			EntityID:    entity.Name,
+			Entity:      entity,
 			RevisionID:  onlineRevisionId,
 			EntityKey:   opt.EntityKey,
 			FeatureList: features,
@@ -65,15 +65,15 @@ func (s *OomStore) MultiGetOnlineFeatureValues(ctx context.Context, opt types.Mu
 	featureMap := groupFeaturesByRevisionId(features)
 
 	// entity_key -> feature_name -> feature_value
-	featureValueMap, err := s.getFeatureValueMap(ctx, opt.EntityKeys, featureMap, entity.Name)
+	featureValueMap, err := s.getFeatureValueMap(ctx, opt.EntityKeys, featureMap, entity)
 	if err != nil {
 		return nil, err
 	}
 
 	return buildFeatureDataSet(featureValueMap, features.Names(), opt.EntityKeys)
 }
 
-func (s *OomStore) getFeatureValueMap(ctx context.Context, entityKeys []string, featureMap map[int32]typesv2.FeatureList, entityName string) (map[string]dbutil.RowMap, error) {
+func (s *OomStore) getFeatureValueMap(ctx context.Context, entityKeys []string, featureMap map[int32]typesv2.FeatureList, entity *typesv2.Entity) (map[string]dbutil.RowMap, error) {
 	// entity_key -> types.RecordMap
 	featureValueMap := make(map[string]dbutil.RowMap)
 
@@ -82,7 +82,7 @@ func (s *OomStore) getFeatureValueMap(ctx context.Context, entityKeys []string,
 			continue
 		}
 		featureValues, err := s.online.MultiGet(ctx, online.MultiGetOpt{
-			EntityName:  entityName,
+			Entity:      entity,
 			RevisionID:  onlineRevisionId,
 			EntityKeys:  entityKeys,
 			FeatureList: features,

diff --git a/pkg/oomstore/typesv2/raw_feature_value_record.go b/pkg/oomstore/typesv2/raw_feature_value_record.go