Fix Go

Signed-off-by: Felix Wang <[email protected]>
feast-dev · May 2, 2022 · 7dda0b4 · 7dda0b4
1 parent 523e021
commit 7dda0b4
Show file tree

Hide file tree

Showing 7 changed files with 48 additions and 31 deletions.
diff --git a/go/embedded/online_features.go b/go/embedded/online_features.go
@@ -88,7 +88,7 @@ func (s *OnlineFeatureService) GetEntityTypesMap(featureRefs []string) (map[stri
 		}
 		for _, entityName := range view.Entities {
 			entity := entitiesByName[entityName]
-			joinKeyTypes[entity.JoinKey] = int32(entity.ValueType.Number())
+			joinKeyTypes[entity.JoinKey] = int32(view.GetEntityType(entity.JoinKey).Number())
 		}
 	}
 
@@ -117,7 +117,7 @@ func (s *OnlineFeatureService) GetEntityTypesMapByFeatureService(featureServiceN
 		}
 		for _, entityName := range view.Entities {
 			entity := entitiesByName[entityName]
-			joinKeyTypes[entity.JoinKey] = int32(entity.ValueType.Number())
+			joinKeyTypes[entity.JoinKey] = int32(view.GetEntityType(entity.JoinKey).Number())
 		}
 	}
 

diff --git a/go/internal/feast/featurestore.go b/go/internal/feast/featurestore.go
@@ -132,7 +132,7 @@ func (fs *FeatureStore) GetOnlineFeatures(
 	if entitylessCase {
 		dummyEntityColumn := &prototypes.RepeatedValue{Val: make([]*prototypes.Value, numRows)}
 		for index := 0; index < numRows; index++ {
-			dummyEntityColumn.Val[index] = &model.DUMMY_ENTITY
+			dummyEntityColumn.Val[index] = &model.DUMMY_ENTITY_VALUE
 		}
 		joinKeyToEntityValues[model.DUMMY_ENTITY_ID] = dummyEntityColumn
 	}

diff --git a/go/internal/feast/model/entity.go b/go/internal/feast/model/entity.go
@@ -2,18 +2,16 @@ package model
 
 import (
 	"github.com/feast-dev/feast/go/protos/feast/core"
-	"github.com/feast-dev/feast/go/protos/feast/types"
 )
 
 type Entity struct {
-	Name      string
-	ValueType types.ValueType_Enum
-	JoinKey   string
+	Name    string
+	JoinKey string
 }
 
 func NewEntityFromProto(proto *core.Entity) *Entity {
-	return &Entity{Name: proto.Spec.Name,
-		ValueType: proto.Spec.ValueType,
-		JoinKey:   proto.Spec.JoinKey,
+	return &Entity{
+		Name:    proto.Spec.Name,
+		JoinKey: proto.Spec.JoinKey,
 	}
 }
diff --git a/go/internal/feast/model/featureview.go b/go/internal/feast/model/featureview.go
@@ -13,12 +13,13 @@ const (
 	DUMMY_ENTITY_VAL  = ""
 )
 
-var DUMMY_ENTITY types.Value = types.Value{Val: &types.Value_StringVal{StringVal: DUMMY_ENTITY_VAL}}
+var DUMMY_ENTITY_VALUE types.Value = types.Value{Val: &types.Value_StringVal{StringVal: DUMMY_ENTITY_VAL}}
 
 type FeatureView struct {
-	Base     *BaseFeatureView
-	Ttl      *durationpb.Duration
-	Entities []string
+	Base          *BaseFeatureView
+	Ttl           *durationpb.Duration
+	Entities      []string
+	EntityColumns []*Feature
 }
 
 func NewFeatureViewFromProto(proto *core.FeatureView) *FeatureView {
@@ -30,23 +31,37 @@ func NewFeatureViewFromProto(proto *core.FeatureView) *FeatureView {
 	} else {
 		featureView.Entities = proto.Spec.Entities
 	}
+	entityColumns := make([]*Feature, len(proto.Spec.EntityColumns))
+	for i, entityColumn := range proto.Spec.EntityColumns {
+		entityColumns[i] = NewFeatureFromProto(entityColumn)
+	}
+	featureView.EntityColumns = entityColumns
 	return featureView
 }
 
-func (fs *FeatureView) NewFeatureViewFromBase(base *BaseFeatureView) *FeatureView {
-	ttl := durationpb.Duration{Seconds: fs.Ttl.Seconds, Nanos: fs.Ttl.Nanos}
+func (fv *FeatureView) NewFeatureViewFromBase(base *BaseFeatureView) *FeatureView {
+	ttl := durationpb.Duration{Seconds: fv.Ttl.Seconds, Nanos: fv.Ttl.Nanos}
 	featureView := &FeatureView{Base: base,
 		Ttl:      &ttl,
-		Entities: fs.Entities,
+		Entities: fv.Entities,
 	}
 	return featureView
 }
 
-func (fs *FeatureView) HasEntity(lookup string) bool {
-	for _, entityName := range fs.Entities {
+func (fv *FeatureView) HasEntity(lookup string) bool {
+	for _, entityName := range fv.Entities {
 		if entityName == lookup {
 			return true
 		}
 	}
 	return false
 }
+
+func (fv *FeatureView) GetEntityType(lookup string) types.ValueType_Enum {
+	for _, entityColumn := range fv.EntityColumns {
+		if entityColumn.Name == lookup {
+			return entityColumn.Dtype
+		}
+	}
+	return types.ValueType_INVALID
+}
diff --git a/go/internal/feast/server/logging/featureserviceschema.go b/go/internal/feast/server/logging/featureserviceschema.go
@@ -66,7 +66,7 @@ func generateSchema(featureService *model.FeatureService, entityMap map[string]*
 				}
 
 				joinKeysSet[joinKey] = nil
-				entityJoinKeyToType[joinKey] = entity.ValueType
+				entityJoinKeyToType[joinKey] = fv.GetEntityType(entity.JoinKey)
 			}
 		} else if odFv, ok := odFvMap[featureViewName]; ok {
 			for _, f := range featureProjection.Features {

diff --git a/go/internal/feast/server/logging/featureserviceschema_test.go b/go/internal/feast/server/logging/featureserviceschema_test.go
@@ -145,8 +145,12 @@ func InitializeFeatureRepoVariablesForTest() (*model.FeatureService, []*model.En
 		[]*model.Feature{f1, f2},
 		projection1,
 	)
-	featureView1 := test.CreateFeatureView(baseFeatureView1, nil, []string{"driver_id"})
-	entity1 := test.CreateNewEntity("driver_id", types.ValueType_INT64, "driver_id")
+	entity1 := test.CreateNewEntity("driver_id", "driver_id")
+	entitycolumn1 := test.CreateNewFeature(
+		"driver_id",
+		types.ValueType_INT64,
+	)
+	featureView1 := test.CreateFeatureView(baseFeatureView1, nil, []string{"driver_id"}, []*model.Feature{entitycolumn1})
 	f3 := test.CreateNewFeature(
 		"int32",
 		types.ValueType_INT32,
@@ -166,7 +170,7 @@ func InitializeFeatureRepoVariablesForTest() (*model.FeatureService, []*model.En
 		[]*model.Feature{f3, f4},
 		projection2,
 	)
-	featureView2 := test.CreateFeatureView(baseFeatureView2, nil, []string{"driver_id"})
+	featureView2 := test.CreateFeatureView(baseFeatureView2, nil, []string{"driver_id"}, []*model.Feature{entitycolumn1})
 
 	f5 := test.CreateNewFeature(
 		"odfv_f1",

diff --git a/go/internal/test/go_integration_test_utils.go b/go/internal/test/go_integration_test_utils.go
@@ -201,11 +201,10 @@ func CreateBaseFeatureView(name string, features []*model.Feature, projection *m
 	}
 }
 
-func CreateNewEntity(name string, valueType types.ValueType_Enum, joinKey string) *model.Entity {
+func CreateNewEntity(name string, joinKey string) *model.Entity {
 	return &model.Entity{
-		Name:      name,
-		ValueType: valueType,
-		JoinKey:   joinKey,
+		Name:    name,
+		JoinKey: joinKey,
 	}
 }
 
@@ -233,10 +232,11 @@ func CreateNewFeatureViewProjection(name string, nameAlias string, features []*m
 	}
 }
 
-func CreateFeatureView(base *model.BaseFeatureView, ttl *durationpb.Duration, entities []string) *model.FeatureView {
+func CreateFeatureView(base *model.BaseFeatureView, ttl *durationpb.Duration, entities []string, entityColumns []*model.Feature) *model.FeatureView {
 	return &model.FeatureView{
-		Base:     base,
-		Ttl:      ttl,
-		Entities: entities,
+		Base:          base,
+		Ttl:           ttl,
+		Entities:      entities,
+		EntityColumns: entityColumns,
 	}
 }