82
82
from feast .repo_contents import RepoContents
83
83
from feast .request_feature_view import RequestFeatureView
84
84
from feast .saved_dataset import SavedDataset , SavedDatasetStorage
85
- from feast .type_map import python_values_to_proto_values
85
+ from feast .type_map import (
86
+ feast_value_type_to_python_type ,
87
+ python_values_to_proto_values ,
88
+ )
86
89
from feast .usage import log_exceptions , log_exceptions_and_usage , set_usage_attribute
87
90
from feast .value_type import ValueType
88
91
from feast .version import get_version
@@ -135,6 +138,7 @@ def __init__(
135
138
self ._registry = Registry (registry_config , repo_path = self .repo_path )
136
139
self ._registry ._initialize_registry ()
137
140
self ._provider = get_provider (self .config , self .repo_path )
141
+ self ._go_server = None
138
142
139
143
@log_exceptions
140
144
def version (self ) -> str :
@@ -1284,7 +1288,29 @@ def get_online_features(
1284
1288
except KeyError as e :
1285
1289
raise ValueError ("All entity_rows must have the same keys." ) from e
1286
1290
1287
- # If Go feature server is enabled, send request to it instead of going through a regular Python logic
1291
+ return self ._get_online_features (
1292
+ features = features ,
1293
+ entity_values = columnar ,
1294
+ full_feature_names = full_feature_names ,
1295
+ native_entity_values = True ,
1296
+ )
1297
+
1298
+ def _get_online_features (
1299
+ self ,
1300
+ features : Union [List [str ], FeatureService ],
1301
+ entity_values : Mapping [
1302
+ str , Union [Sequence [Any ], Sequence [Value ], RepeatedValue ]
1303
+ ],
1304
+ full_feature_names : bool = False ,
1305
+ native_entity_values : bool = True ,
1306
+ ):
1307
+ # Extract Sequence from RepeatedValue Protobuf.
1308
+ entity_value_lists : Dict [str , Union [List [Any ], List [Value ]]] = {
1309
+ k : list (v ) if isinstance (v , Sequence ) else list (v .val )
1310
+ for k , v in entity_values .items ()
1311
+ }
1312
+
1313
+ # If Go feature server is enabled, send request to it instead of going through regular Python logic
1288
1314
if self .config .go_feature_server :
1289
1315
from feast .embedded_go .online_features_service import (
1290
1316
EmbeddedOnlineFeatureServer ,
@@ -1296,32 +1322,31 @@ def get_online_features(
1296
1322
str (self .repo_path .absolute ()), self .config , self
1297
1323
)
1298
1324
1325
+ entity_native_values : Dict [str , List [Any ]]
1326
+ if not native_entity_values :
1327
+ # Convert proto types to native types since Go feature server currently
1328
+ # only handles native types.
1329
+ # TODO(felixwang9817): Remove this logic once native types are supported.
1330
+ entity_native_values = {
1331
+ k : [
1332
+ feast_value_type_to_python_type (proto_value )
1333
+ for proto_value in v
1334
+ ]
1335
+ for k , v in entity_value_lists .items ()
1336
+ }
1337
+ else :
1338
+ entity_native_values = entity_value_lists
1339
+
1299
1340
return self ._go_server .get_online_features (
1300
1341
features_refs = features if isinstance (features , list ) else [],
1301
1342
feature_service = features
1302
1343
if isinstance (features , FeatureService )
1303
1344
else None ,
1304
- entities = columnar ,
1345
+ entities = entity_native_values ,
1305
1346
request_data = {}, # TODO: add request data parameter to public API
1306
1347
full_feature_names = full_feature_names ,
1307
1348
)
1308
1349
1309
- return self ._get_online_features (
1310
- features = features ,
1311
- entity_values = columnar ,
1312
- full_feature_names = full_feature_names ,
1313
- native_entity_values = True ,
1314
- )
1315
-
1316
- def _get_online_features (
1317
- self ,
1318
- features : Union [List [str ], FeatureService ],
1319
- entity_values : Mapping [
1320
- str , Union [Sequence [Any ], Sequence [Value ], RepeatedValue ]
1321
- ],
1322
- full_feature_names : bool = False ,
1323
- native_entity_values : bool = True ,
1324
- ):
1325
1350
_feature_refs = self ._get_features (features , allow_cache = True )
1326
1351
(
1327
1352
requested_feature_views ,
@@ -1344,12 +1369,6 @@ def _get_online_features(
1344
1369
join_keys_set ,
1345
1370
) = self ._get_entity_maps (requested_feature_views )
1346
1371
1347
- # Extract Sequence from RepeatedValue Protobuf.
1348
- entity_value_lists : Dict [str , Union [List [Any ], List [Value ]]] = {
1349
- k : list (v ) if isinstance (v , Sequence ) else list (v .val )
1350
- for k , v in entity_values .items ()
1351
- }
1352
-
1353
1372
entity_proto_values : Dict [str , List [Value ]]
1354
1373
if native_entity_values :
1355
1374
# Convert values to Protobuf once.
0 commit comments