chore: Support passing list of Entities to a feature view definition (#2528)

achals · web-flow · commit c5e0da2ab0c7 · 2022-04-12T15:44:42.000-07:00
Signed-off-by: Achal Shah &lt;achals@gmail.com&gt;
diff --git a/sdk/python/feast/feature_view.py b/sdk/python/feast/feature_view.py
@@ -97,7 +97,7 @@ def __init__(
         self,
         *args,
         name: Optional[str] = None,
-        entities: Optional[List[str]] = None,
+        entities: Optional[Union[List[Entity], List[str]]] = None,
         ttl: Optional[Union[Duration, timedelta]] = None,
         batch_source: Optional[DataSource] = None,
         stream_source: Optional[DataSource] = None,
@@ -167,7 +167,11 @@ def __init__(
             raise ValueError("feature view name needs to be specified")
 
         self.name = _name
-        self.entities = _entities if _entities else [DUMMY_ENTITY_NAME]
+        self.entities = (
+            [e.name if isinstance(e, Entity) else e for e in _entities]
+            if _entities
+            else [DUMMY_ENTITY_NAME]
+        )
 
         self._initialize_sources(_name, batch_source, stream_source, source)
 
diff --git a/sdk/python/tests/integration/feature_repos/universal/feature_views.py b/sdk/python/tests/integration/feature_repos/universal/feature_views.py
@@ -18,6 +18,7 @@
 )
 from feast.data_source import DataSource, RequestSource
 from feast.types import FeastType
+from tests.integration.feature_repos.universal.entities import location
 
 
 def driver_feature_view(
@@ -217,7 +218,7 @@ def create_order_feature_view(source, infer_features: bool = False):
 def create_location_stats_feature_view(source, infer_features: bool = False):
     location_stats_feature_view = FeatureView(
         name="location_stats",
-        entities=["location_id"],
+        entities=[location()],
         schema=None if infer_features else [Field(name="temperature", dtype=Int32)],
         source=source,
         ttl=timedelta(days=2),
