refactor: Refactor universal test data to have more well defined types (

#2340)

* Refactor universal test suite to generate fixtures as dataclasses instead of loosely typed dictionaries

Signed-off-by: Danny Chiao <danny@tecton.ai>

* lint

Signed-off-by: Danny Chiao <danny@tecton.ai>

* lint

Signed-off-by: Danny Chiao <danny@tecton.ai>

* other lint

Signed-off-by: Danny Chiao <danny@tecton.ai>

* other lint

Signed-off-by: Danny Chiao <danny@tecton.ai>

* test fix

Signed-off-by: Danny Chiao <danny@tecton.ai>

* test fix

Signed-off-by: Danny Chiao <danny@tecton.ai>

sdk/python/feast/__init__.py

@@ -15,6 +15,7 @@
 from .feature_view import FeatureView
 from .on_demand_feature_view import OnDemandFeatureView
 from .repo_config import RepoConfig
+from .request_feature_view import RequestFeatureView
 from .value_type import ValueType
 
 logging.basicConfig(
@@ -44,5 +45,6 @@
     "BigQuerySource",
     "FileSource",
     "RedshiftSource",
+    "RequestFeatureView",
     "SnowflakeSource",
 ]

sdk/python/feast/diff/registry_diff.py

@@ -1,9 +1,10 @@
 from dataclasses import dataclass
-from typing import Any, Dict, Generic, Iterable, List, Set, Tuple, TypeVar
+from typing import Any, Dict, Iterable, List, Optional, Set, Tuple, TypeVar, cast
 
 from feast.base_feature_view import BaseFeatureView
 from feast.diff.property_diff import PropertyDiff, TransitionType
 from feast.entity import Entity
+from feast.feast_object import FeastObject
 from feast.feature_service import FeatureService
 from feast.feature_view import DUMMY_ENTITY_NAME
 from feast.protos.feast.core.Entity_pb2 import Entity as EntityProto
@@ -20,15 +21,13 @@
 from feast.registry import FEAST_OBJECT_TYPES, FeastObjectType, Registry
 from feast.repo_contents import RepoContents
 
-FeastObject = TypeVar("FeastObject", Entity, BaseFeatureView, FeatureService)
-
 
 @dataclass
-class FeastObjectDiff(Generic[FeastObject]):
+class FeastObjectDiff:
     name: str
     feast_object_type: FeastObjectType
-    current_feast_object: FeastObject
-    new_feast_object: FeastObject
+    current_feast_object: Optional[FeastObject]
+    new_feast_object: Optional[FeastObject]
     feast_object_property_diffs: List[PropertyDiff]
     transition_type: TransitionType
 
@@ -257,20 +256,25 @@ def apply_diff_to_registry(
         # will automatically delete the existing object.
         if feast_object_diff.transition_type == TransitionType.DELETE:
             if feast_object_diff.feast_object_type == FeastObjectType.ENTITY:
-                registry.delete_entity(
-                    feast_object_diff.current_feast_object.name, project, commit=False
-                )
+                entity_obj = cast(Entity, feast_object_diff.current_feast_object)
+                registry.delete_entity(entity_obj.name, project, commit=False)
             elif feast_object_diff.feast_object_type == FeastObjectType.FEATURE_SERVICE:
+                feature_service_obj = cast(
+                    FeatureService, feast_object_diff.current_feast_object
+                )
                 registry.delete_feature_service(
-                    feast_object_diff.current_feast_object.name, project, commit=False
+                    feature_service_obj.name, project, commit=False
                 )
             elif feast_object_diff.feast_object_type in [
                 FeastObjectType.FEATURE_VIEW,
                 FeastObjectType.ON_DEMAND_FEATURE_VIEW,
                 FeastObjectType.REQUEST_FEATURE_VIEW,
             ]:
+                feature_view_obj = cast(
+                    BaseFeatureView, feast_object_diff.current_feast_object
+                )
                 registry.delete_feature_view(
-                    feast_object_diff.current_feast_object.name, project, commit=False,
+                    feature_view_obj.name, project, commit=False,
                 )
 
         if feast_object_diff.transition_type in [
@@ -279,19 +283,25 @@ def apply_diff_to_registry(
         ]:
             if feast_object_diff.feast_object_type == FeastObjectType.ENTITY:
                 registry.apply_entity(
-                    feast_object_diff.new_feast_object, project, commit=False
+                    cast(Entity, feast_object_diff.new_feast_object),
+                    project,
+                    commit=False,
                 )
             elif feast_object_diff.feast_object_type == FeastObjectType.FEATURE_SERVICE:
                 registry.apply_feature_service(
-                    feast_object_diff.new_feast_object, project, commit=False
+                    cast(FeatureService, feast_object_diff.new_feast_object),
+                    project,
+                    commit=False,
                 )
             elif feast_object_diff.feast_object_type in [
                 FeastObjectType.FEATURE_VIEW,
                 FeastObjectType.ON_DEMAND_FEATURE_VIEW,
                 FeastObjectType.REQUEST_FEATURE_VIEW,
             ]:
                 registry.apply_feature_view(
-                    feast_object_diff.new_feast_object, project, commit=False
+                    cast(BaseFeatureView, feast_object_diff.new_feast_object),
+                    project,
+                    commit=False,
                 )
 
     if commit:

sdk/python/feast/feast_object.py

@@ -0,0 +1,12 @@
+from typing import Union
+
+from .entity import Entity
+from .feature_service import FeatureService
+from .feature_view import FeatureView
+from .on_demand_feature_view import OnDemandFeatureView
+from .request_feature_view import RequestFeatureView
+
+# Convenience type representing all Feast objects
+FeastObject = Union[
+    FeatureView, OnDemandFeatureView, RequestFeatureView, Entity, FeatureService
+]

sdk/python/feast/feature_store.py

@@ -50,6 +50,7 @@
     RequestDataNotFoundInEntityDfException,
     RequestDataNotFoundInEntityRowsException,
 )
+from feast.feast_object import FeastObject
 from feast.feature_service import FeatureService
 from feast.feature_view import (
     DUMMY_ENTITY,
@@ -539,27 +540,9 @@ def apply(
             OnDemandFeatureView,
             RequestFeatureView,
             FeatureService,
-            List[
-                Union[
-                    FeatureView,
-                    OnDemandFeatureView,
-                    RequestFeatureView,
-                    Entity,
-                    FeatureService,
-                ]
-            ],
+            List[FeastObject],
         ],
-        objects_to_delete: Optional[
-            List[
-                Union[
-                    FeatureView,
-                    OnDemandFeatureView,
-                    RequestFeatureView,
-                    Entity,
-                    FeatureService,
-                ]
-            ]
-        ] = None,
+        objects_to_delete: Optional[List[FeastObject]] = None,
         partial: bool = True,
     ):
         """Register objects to metadata store and update related infrastructure.

sdk/python/tests/benchmarks/test_benchmark_universal_online_retrieval.py

@@ -1,8 +1,10 @@
 import random
+from typing import List
 
 import pytest
 
 from feast import FeatureService
+from feast.feast_object import FeastObject
 from tests.integration.feature_repos.repo_configuration import (
     construct_universal_feature_views,
 )
@@ -16,25 +18,24 @@
 @pytest.mark.benchmark
 @pytest.mark.integration
 def test_online_retrieval(environment, universal_data_sources, benchmark):
-
     fs = environment.feature_store
     entities, datasets, data_sources = universal_data_sources
     feature_views = construct_universal_feature_views(data_sources)
 
     feature_service = FeatureService(
         "convrate_plus100",
-        features=[feature_views["driver"][["conv_rate"]], feature_views["driver_odfv"]],
+        features=[feature_views.driver[["conv_rate"]], feature_views.driver_odfv],
     )
 
-    feast_objects = []
+    feast_objects: List[FeastObject] = []
     feast_objects.extend(feature_views.values())
     feast_objects.extend([driver(), customer(), location(), feature_service])
     fs.apply(feast_objects)
     fs.materialize(environment.start_date, environment.end_date)
 
-    sample_drivers = random.sample(entities["driver"], 10)
+    sample_drivers = random.sample(entities.driver_vals, 10)
 
-    sample_customers = random.sample(entities["customer"], 10)
+    sample_customers = random.sample(entities.customer_vals, 10)
 
     entity_rows = [
         {"driver": d, "customer_id": c, "val_to_add": 50}

sdk/python/tests/conftest.py

@@ -32,10 +32,9 @@
     FULL_REPO_CONFIGS,
     REDIS_CONFIG,
     Environment,
+    TestData,
     construct_test_environment,
-    construct_universal_data_sources,
-    construct_universal_datasets,
-    construct_universal_entities,
+    construct_universal_test_data,
 )
 
 logger = logging.getLogger(__name__)
@@ -184,21 +183,13 @@ def cleanup():
 
 
 @pytest.fixture(scope="session")
-def universal_data_sources(request, environment):
-    entities = construct_universal_entities()
-    datasets = construct_universal_datasets(
-        entities, environment.start_date, environment.end_date
-    )
-    datasources = construct_universal_data_sources(
-        datasets, environment.data_source_creator
-    )
-
+def universal_data_sources(request, environment) -> TestData:
     def cleanup():
         # logger.info("Running cleanup in %s, Request: %s", worker_id, request.param)
         environment.data_source_creator.teardown()
 
     request.addfinalizer(cleanup)
-    return entities, datasets, datasources
+    return construct_universal_test_data(environment)
 
 
 @pytest.fixture(scope="session")

sdk/python/tests/integration/e2e/test_validation.py

@@ -67,7 +67,7 @@ def test_historical_retrieval_with_validation(environment, universal_data_source
 
     store.apply([driver(), customer(), location(), *feature_views.values()])
 
-    entity_df = datasets["entity"].drop(
+    entity_df = datasets.entity_df.drop(
         columns=["order_id", "origin_id", "destination_id"]
     )
 
@@ -101,7 +101,7 @@ def test_historical_retrieval_fails_on_validation(environment, universal_data_so
 
     store.apply([driver(), customer(), location(), *feature_views.values()])
 
-    entity_df = datasets["entity"].drop(
+    entity_df = datasets.entity_df.drop(
         columns=["order_id", "origin_id", "destination_id"]
     )