diff --git a/sdk/python/feast/utils.py b/sdk/python/feast/utils.py index ec2da79782..32cd2f606c 100644 --- a/sdk/python/feast/utils.py +++ b/sdk/python/feast/utils.py @@ -42,7 +42,7 @@ from feast.protos.feast.types.Value_pb2 import RepeatedValue as RepeatedValueProto from feast.protos.feast.types.Value_pb2 import Value as ValueProto from feast.type_map import python_values_to_proto_values -from feast.types import from_feast_to_pyarrow_type +from feast.types import ComplexFeastType, PrimitiveFeastType, from_feast_to_pyarrow_type from feast.value_type import ValueType from feast.version import get_version @@ -552,13 +552,27 @@ def _augment_response_with_on_demand_transforms( selected_subset = [f for f in transformed_columns if f in _feature_refs] proto_values = [] + schema_dict = {k.name: k.dtype for k in odfv.schema} for selected_feature in selected_subset: feature_vector = transformed_features[selected_feature] + selected_feature_type = schema_dict.get(selected_feature, None) + feature_type: ValueType = ValueType.UNKNOWN + if selected_feature_type is not None: + if isinstance( + selected_feature_type, (ComplexFeastType, PrimitiveFeastType) + ): + feature_type = selected_feature_type.to_value_type() + elif not isinstance(selected_feature_type, ValueType): + raise TypeError( + f"Unexpected type for feature_type: {type(feature_type)}" + ) + proto_values.append( - python_values_to_proto_values(feature_vector, ValueType.UNKNOWN) - if odfv.mode == "python" - else python_values_to_proto_values( - feature_vector.to_numpy(), ValueType.UNKNOWN + python_values_to_proto_values( + feature_vector + if odfv.mode == "python" + else feature_vector.to_numpy(), + feature_type, ) ) diff --git a/sdk/python/tests/unit/test_on_demand_python_transformation.py b/sdk/python/tests/unit/test_on_demand_python_transformation.py index b7ddfb9e75..530bf1fa0a 100644 --- a/sdk/python/tests/unit/test_on_demand_python_transformation.py +++ b/sdk/python/tests/unit/test_on_demand_python_transformation.py @@ -609,19 +609,19 @@ def pandas_view(features_df: pd.DataFrame) -> pd.DataFrame: "val_to_add", "val_to_add_2", ] - with pytest.raises(TypeError): - _ = self.store.get_online_features( - entity_rows=[ - {"driver_id": 1234567890, "val_to_add": 0, "val_to_add_2": 1} - ], - features=[ - "driver_hourly_stats:conv_rate", - "driver_hourly_stats:acc_rate", - "driver_hourly_stats:avg_daily_trips", - "pandas_view:conv_rate_plus_val1", - "pandas_view:conv_rate_plus_val2", - ], - ) + resp_online_missing_entity = self.store.get_online_features( + entity_rows=[ + {"driver_id": 1234567890, "val_to_add": 0, "val_to_add_2": 1} + ], + features=[ + "driver_hourly_stats:conv_rate", + "driver_hourly_stats:acc_rate", + "driver_hourly_stats:avg_daily_trips", + "pandas_view:conv_rate_plus_val1", + "pandas_view:conv_rate_plus_val2", + ], + ) + assert resp_online_missing_entity is not None resp_online = self.store.get_online_features( entity_rows=[{"driver_id": 1001, "val_to_add": 0, "val_to_add_2": 1}], features=[