|
42 | 42 | from feast.protos.feast.types.Value_pb2 import RepeatedValue as RepeatedValueProto |
43 | 43 | from feast.protos.feast.types.Value_pb2 import Value as ValueProto |
44 | 44 | from feast.type_map import python_values_to_proto_values |
45 | | -from feast.types import from_feast_to_pyarrow_type |
| 45 | +from feast.types import ComplexFeastType, PrimitiveFeastType, from_feast_to_pyarrow_type |
46 | 46 | from feast.value_type import ValueType |
47 | 47 | from feast.version import get_version |
48 | 48 |
|
@@ -552,13 +552,27 @@ def _augment_response_with_on_demand_transforms( |
552 | 552 | selected_subset = [f for f in transformed_columns if f in _feature_refs] |
553 | 553 |
|
554 | 554 | proto_values = [] |
| 555 | + schema_dict = {k.name: k.dtype for k in odfv.schema} |
555 | 556 | for selected_feature in selected_subset: |
556 | 557 | feature_vector = transformed_features[selected_feature] |
| 558 | + selected_feature_type = schema_dict.get(selected_feature, None) |
| 559 | + feature_type: ValueType = ValueType.UNKNOWN |
| 560 | + if selected_feature_type is not None: |
| 561 | + if isinstance( |
| 562 | + selected_feature_type, (ComplexFeastType, PrimitiveFeastType) |
| 563 | + ): |
| 564 | + feature_type = selected_feature_type.to_value_type() |
| 565 | + elif not isinstance(selected_feature_type, ValueType): |
| 566 | + raise TypeError( |
| 567 | + f"Unexpected type for feature_type: {type(feature_type)}" |
| 568 | + ) |
| 569 | + |
557 | 570 | proto_values.append( |
558 | | - python_values_to_proto_values(feature_vector, ValueType.UNKNOWN) |
559 | | - if odfv.mode == "python" |
560 | | - else python_values_to_proto_values( |
561 | | - feature_vector.to_numpy(), ValueType.UNKNOWN |
| 571 | + python_values_to_proto_values( |
| 572 | + feature_vector |
| 573 | + if odfv.mode == "python" |
| 574 | + else feature_vector.to_numpy(), |
| 575 | + feature_type, |
562 | 576 | ) |
563 | 577 | ) |
564 | 578 |
|
|
0 commit comments