linted and updated

franciscojavierarceo · franciscojavierarceo · commit 3068c3b760ce · 2024-09-29T15:06:29.000-04:00
Signed-off-by: Francisco Javier Arceo &lt;farceo@redhat.com&gt;
diff --git a/sdk/python/feast/inference.py b/sdk/python/feast/inference.py
@@ -13,9 +13,9 @@
 from feast.infra.offline_stores.file_source import FileSource
 from feast.infra.offline_stores.redshift_source import RedshiftSource
 from feast.infra.offline_stores.snowflake_source import SnowflakeSource
+from feast.on_demand_feature_view import OnDemandFeatureView
 from feast.repo_config import RepoConfig
 from feast.stream_feature_view import StreamFeatureView
-from feast.on_demand_feature_view import OnDemandFeatureView
 from feast.types import String
 from feast.value_type import ValueType
 
@@ -257,6 +257,8 @@ def _infer_features_and_entities(
                 if field.name not in [feature.name for feature in fv.features]:
                     fv.features.append(field)
 
+    fv.entity_columns = entity_columns
+
 
 def _infer_on_demand_features_and_entities(
     fv: OnDemandFeatureView,
diff --git a/sdk/python/feast/on_demand_feature_view.py b/sdk/python/feast/on_demand_feature_view.py
@@ -3,7 +3,7 @@
 import inspect
 import warnings
 from types import FunctionType
-from typing import Any, Optional, Union, get_type_hints, List
+from typing import Any, List, Optional, Union, get_type_hints
 
 import dill
 import pandas as pd
diff --git a/sdk/python/tests/unit/local_feast_tests/test_local_feature_store.py b/sdk/python/tests/unit/local_feast_tests/test_local_feature_store.py
@@ -392,7 +392,9 @@ def test_apply_dummy_entity_and_feature_view_columns(test_feature_store):
 
     # Register Feature View
     test_feature_store.apply([e1, fv_no_entity, fv_with_entity])
-    fv_from_online_store = test_feature_store.get_feature_view("my_feature_view_no_entity")
+    fv_from_online_store = test_feature_store.get_feature_view(
+        "my_feature_view_no_entity"
+    )
 
     # Note that after the apply() the feature_view serializes the Dummy Entity ID
     assert fv_no_entity.entity_columns[0].name == DUMMY_ENTITY_ID
diff --git a/sdk/python/tests/unit/test_on_demand_python_transformation.py b/sdk/python/tests/unit/test_on_demand_python_transformation.py
@@ -380,15 +380,15 @@ def python_view(inputs: dict[str, Any]) -> dict[str, Any]:
             self.store.apply(
                 [driver, driver_stats_source, driver_stats_fv, python_view]
             )
-            self.store.write_to_online_store(
-                feature_view_name="driver_hourly_stats", df=driver_df
-            )
-
             fv_applied = self.store.get_feature_view("driver_hourly_stats")
             assert fv_applied.entities == [driver.name]
             # Note here that after apply() is called, the entity_columns are populated with the join_key
             assert fv_applied.entity_columns[0].name == driver.join_key
 
+            self.store.write_to_online_store(
+                feature_view_name="driver_hourly_stats", df=driver_df
+            )
+
     def test_python_transformation_returning_all_data_types(self):
         entity_rows = [
             {
