Fix OnDemandFeatureView type inference for array types

Signed-off-by: Alex Mirrington <alex.mirrington@rokt.com>
feast-dev · Jun 24, 2024 · 6b7afb9 · 6b7afb9
1 parent de5b0eb
commit 6b7afb9
Show file tree

Hide file tree

Showing 3 changed files with 9 additions and 3 deletions.
diff --git a/sdk/python/feast/transformation/pandas_transformation.py b/sdk/python/feast/transformation/pandas_transformation.py
@@ -44,7 +44,9 @@ def infer_features(self, random_input: dict[str, list[Any]]) -> list[Field]:
             Field(
                 name=f,
                 dtype=from_value_type(
-                    python_type_to_feast_value_type(f, type_name=str(dt))
+                    python_type_to_feast_value_type(
+                        f, value=output_df[f].tolist()[0], type_name=str(dt)
+                    )
                 ),
             )
             for f, dt in zip(output_df.columns, output_df.dtypes)

diff --git a/sdk/python/feast/transformation/python_transformation.py b/sdk/python/feast/transformation/python_transformation.py
@@ -44,7 +44,9 @@ def infer_features(self, random_input: dict[str, list[Any]]) -> list[Field]:
             Field(
                 name=f,
                 dtype=from_value_type(
-                    python_type_to_feast_value_type(f, type_name=type(dt[0]).__name__)
+                    python_type_to_feast_value_type(
+                        f, value=dt[0], type_name=type(dt[0]).__name__
+                    )
                 ),
             )
             for f, dt in output_dict.items()

diff --git a/sdk/python/feast/transformation/substrait_transformation.py b/sdk/python/feast/transformation/substrait_transformation.py
@@ -64,7 +64,9 @@ def infer_features(self, random_input: dict[str, list[Any]]) -> list[Field]:
             Field(
                 name=f,
                 dtype=from_value_type(
-                    python_type_to_feast_value_type(f, type_name=str(dt))
+                    python_type_to_feast_value_type(
+                        f, value=output_df[f].tolist()[0], type_name=str(dt)
+                    )
                 ),
             )
             for f, dt in zip(output_df.columns, output_df.dtypes)