From 41aaeebaa5908f44cda28a4410e2fca412f53e92 Mon Sep 17 00:00:00 2001
From: Francisco Arceo <arceofrancisco@gmail.com>
Date: Wed, 23 Oct 2024 04:28:38 -0400
Subject: [PATCH] fix: Fixing failure of protos during ODFV transformations for
 missing entities (#4667)

---
 sdk/python/feast/utils.py                     | 24 +++++++++++++----
 .../test_on_demand_python_transformation.py   | 26 +++++++++----------
 2 files changed, 32 insertions(+), 18 deletions(-)

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=[