Skip to content

Commit

Permalink
fix error to search with key instead of name
Browse files Browse the repository at this point in the history
  • Loading branch information
msistla96 committed May 2, 2024
1 parent e093be7 commit 3d12dbf
Show file tree
Hide file tree
Showing 2 changed files with 17 additions and 19 deletions.
12 changes: 7 additions & 5 deletions go/internal/feast/featurestore.go
Original file line number Diff line number Diff line change
Expand Up @@ -291,13 +291,15 @@ func (fs *FeatureStore) ListEntities(hideDummyEntity bool) ([]*model.Entity, err
return entities, nil
}

func (fs *FeatureStore) ListEntity(entityName string) (*model.Entity, error) {
func (fs *FeatureStore) ListEntityByJoinKey(entitykey string) (*model.Entity, error) {

entity, err := fs.registry.GetEntity(fs.config.Project, entityName)
if err != nil {
return entity, err
entities, _ := fs.ListEntities(false)
for _, entity := range entities {
if entity.JoinKey == entitykey {
return entity, nil
}
}
return entity, nil
return nil, fmt.Errorf("entity with join key %s not found", entitykey)
}

func (fs *FeatureStore) ListOnDemandFeatureViews() ([]*model.OnDemandFeatureView, error) {
Expand Down
24 changes: 10 additions & 14 deletions go/internal/feast/server/http_server.go
Original file line number Diff line number Diff line change
Expand Up @@ -171,10 +171,10 @@ func logWithSpanContext(span tracer.Span) zerolog.Logger {

/*
*
This function is to typecast entity values unmarshalled from JSON to the actual entity type specified in the entity definition.
This is because of a bug when using an entity defined with ValueType_INT32
This function fixes the bug when the entity join key type is INT32.
It typecasts the entity values to int32, as unmarshalled JSON values only return int64.
*/
func typecastToEntityIntType(val *repeatedValue, entityType prototypes.ValueType_Enum) {
func typecastToInt32Type(val *repeatedValue, entityType prototypes.ValueType_Enum) {
if val.int64Val != nil {
if entityType == prototypes.ValueType_INT32 {
for _, v := range val.int64Val {
Expand All @@ -183,14 +183,6 @@ func typecastToEntityIntType(val *repeatedValue, entityType prototypes.ValueType
val.int64Val = nil
}
}
if val.int32Val != nil {
if entityType == prototypes.ValueType_INT64 {
for _, v := range val.int32Val {
val.int64Val = append(val.int64Val, int64(v))
}
val.int32Val = nil
}
}
}

func (s *httpServer) getOnlineFeatures(w http.ResponseWriter, r *http.Request) {
Expand Down Expand Up @@ -238,11 +230,15 @@ func (s *httpServer) getOnlineFeatures(w http.ResponseWriter, r *http.Request) {
}
entitiesProto := make(map[string]*prototypes.RepeatedValue)
for key, value := range request.Entities {
var entity, err = s.fs.ListEntity(key)
var entity, err = s.fs.ListEntityByJoinKey(key)
if err != nil {

logSpanContext.Error().Err(err).Msg("Error getting entity from registry")
writeJSONError(w, fmt.Errorf("Error getting entity from registry: %+v", err), http.StatusInternalServerError)
return
}
if value.int64Val != nil {
typecastToInt32Type(&value, entity.Dtype)
}
typecastToEntityIntType(&value, entity.Dtype)
entitiesProto[key] = value.ToProto()
}
requestContextProto := make(map[string]*prototypes.RepeatedValue)
Expand Down

0 comments on commit 3d12dbf

Please sign in to comment.