Skip to content

Commit 251eee4

Browse files
adchiaadchia
authored
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>
1 parent fdb5f21 commit 251eee4

12 files changed

+259
-211
lines changed

‎sdk/python/feast/__init__.py

+2
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@
1515
from .feature_view import FeatureView
1616
from .on_demand_feature_view import OnDemandFeatureView
1717
from .repo_config import RepoConfig
18+
from .request_feature_view import RequestFeatureView
1819
from .value_type import ValueType
1920

2021
logging.basicConfig(
@@ -44,5 +45,6 @@
4445
"BigQuerySource",
4546
"FileSource",
4647
"RedshiftSource",
48+
"RequestFeatureView",
4749
"SnowflakeSource",
4850
]

‎sdk/python/feast/diff/registry_diff.py

+24-14
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,10 @@
11
from dataclasses import dataclass
2-
from typing import Any, Dict, Generic, Iterable, List, Set, Tuple, TypeVar
2+
from typing import Any, Dict, Iterable, List, Optional, Set, Tuple, TypeVar, cast
33

44
from feast.base_feature_view import BaseFeatureView
55
from feast.diff.property_diff import PropertyDiff, TransitionType
66
from feast.entity import Entity
7+
from feast.feast_object import FeastObject
78
from feast.feature_service import FeatureService
89
from feast.feature_view import DUMMY_ENTITY_NAME
910
from feast.protos.feast.core.Entity_pb2 import Entity as EntityProto
@@ -20,15 +21,13 @@
2021
from feast.registry import FEAST_OBJECT_TYPES, FeastObjectType, Registry
2122
from feast.repo_contents import RepoContents
2223

23-
FeastObject = TypeVar("FeastObject", Entity, BaseFeatureView, FeatureService)
24-
2524

2625
@dataclass
27-
class FeastObjectDiff(Generic[FeastObject]):
26+
class FeastObjectDiff:
2827
name: str
2928
feast_object_type: FeastObjectType
30-
current_feast_object: FeastObject
31-
new_feast_object: FeastObject
29+
current_feast_object: Optional[FeastObject]
30+
new_feast_object: Optional[FeastObject]
3231
feast_object_property_diffs: List[PropertyDiff]
3332
transition_type: TransitionType
3433

@@ -257,20 +256,25 @@ def apply_diff_to_registry(
257256
# will automatically delete the existing object.
258257
if feast_object_diff.transition_type == TransitionType.DELETE:
259258
if feast_object_diff.feast_object_type == FeastObjectType.ENTITY:
260-
registry.delete_entity(
261-
feast_object_diff.current_feast_object.name, project, commit=False
262-
)
259+
entity_obj = cast(Entity, feast_object_diff.current_feast_object)
260+
registry.delete_entity(entity_obj.name, project, commit=False)
263261
elif feast_object_diff.feast_object_type == FeastObjectType.FEATURE_SERVICE:
262+
feature_service_obj = cast(
263+
FeatureService, feast_object_diff.current_feast_object
264+
)
264265
registry.delete_feature_service(
265-
feast_object_diff.current_feast_object.name, project, commit=False
266+
feature_service_obj.name, project, commit=False
266267
)
267268
elif feast_object_diff.feast_object_type in [
268269
FeastObjectType.FEATURE_VIEW,
269270
FeastObjectType.ON_DEMAND_FEATURE_VIEW,
270271
FeastObjectType.REQUEST_FEATURE_VIEW,
271272
]:
273+
feature_view_obj = cast(
274+
BaseFeatureView, feast_object_diff.current_feast_object
275+
)
272276
registry.delete_feature_view(
273-
feast_object_diff.current_feast_object.name, project, commit=False,
277+
feature_view_obj.name, project, commit=False,
274278
)
275279

276280
if feast_object_diff.transition_type in [
@@ -279,19 +283,25 @@ def apply_diff_to_registry(
279283
]:
280284
if feast_object_diff.feast_object_type == FeastObjectType.ENTITY:
281285
registry.apply_entity(
282-
feast_object_diff.new_feast_object, project, commit=False
286+
cast(Entity, feast_object_diff.new_feast_object),
287+
project,
288+
commit=False,
283289
)
284290
elif feast_object_diff.feast_object_type == FeastObjectType.FEATURE_SERVICE:
285291
registry.apply_feature_service(
286-
feast_object_diff.new_feast_object, project, commit=False
292+
cast(FeatureService, feast_object_diff.new_feast_object),
293+
project,
294+
commit=False,
287295
)
288296
elif feast_object_diff.feast_object_type in [
289297
FeastObjectType.FEATURE_VIEW,
290298
FeastObjectType.ON_DEMAND_FEATURE_VIEW,
291299
FeastObjectType.REQUEST_FEATURE_VIEW,
292300
]:
293301
registry.apply_feature_view(
294-
feast_object_diff.new_feast_object, project, commit=False
302+
cast(BaseFeatureView, feast_object_diff.new_feast_object),
303+
project,
304+
commit=False,
295305
)
296306

297307
if commit:

‎sdk/python/feast/feast_object.py

+12
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,12 @@
1+
from typing import Union
2+
3+
from .entity import Entity
4+
from .feature_service import FeatureService
5+
from .feature_view import FeatureView
6+
from .on_demand_feature_view import OnDemandFeatureView
7+
from .request_feature_view import RequestFeatureView
8+
9+
# Convenience type representing all Feast objects
10+
FeastObject = Union[
11+
FeatureView, OnDemandFeatureView, RequestFeatureView, Entity, FeatureService
12+
]

‎sdk/python/feast/feature_store.py

+3-20
Original file line numberDiff line numberDiff line change
@@ -50,6 +50,7 @@
5050
RequestDataNotFoundInEntityDfException,
5151
RequestDataNotFoundInEntityRowsException,
5252
)
53+
from feast.feast_object import FeastObject
5354
from feast.feature_service import FeatureService
5455
from feast.feature_view import (
5556
DUMMY_ENTITY,
@@ -539,27 +540,9 @@ def apply(
539540
OnDemandFeatureView,
540541
RequestFeatureView,
541542
FeatureService,
542-
List[
543-
Union[
544-
FeatureView,
545-
OnDemandFeatureView,
546-
RequestFeatureView,
547-
Entity,
548-
FeatureService,
549-
]
550-
],
543+
List[FeastObject],
551544
],
552-
objects_to_delete: Optional[
553-
List[
554-
Union[
555-
FeatureView,
556-
OnDemandFeatureView,
557-
RequestFeatureView,
558-
Entity,
559-
FeatureService,
560-
]
561-
]
562-
] = None,
545+
objects_to_delete: Optional[List[FeastObject]] = None,
563546
partial: bool = True,
564547
):
565548
"""Register objects to metadata store and update related infrastructure.

‎sdk/python/tests/benchmarks/test_benchmark_universal_online_retrieval.py

+6-5
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,10 @@
11
import random
2+
from typing import List
23

34
import pytest
45

56
from feast import FeatureService
7+
from feast.feast_object import FeastObject
68
from tests.integration.feature_repos.repo_configuration import (
79
construct_universal_feature_views,
810
)
@@ -16,25 +18,24 @@
1618
@pytest.mark.benchmark
1719
@pytest.mark.integration
1820
def test_online_retrieval(environment, universal_data_sources, benchmark):
19-
2021
fs = environment.feature_store
2122
entities, datasets, data_sources = universal_data_sources
2223
feature_views = construct_universal_feature_views(data_sources)
2324

2425
feature_service = FeatureService(
2526
"convrate_plus100",
26-
features=[feature_views["driver"][["conv_rate"]], feature_views["driver_odfv"]],
27+
features=[feature_views.driver[["conv_rate"]], feature_views.driver_odfv],
2728
)
2829

29-
feast_objects = []
30+
feast_objects: List[FeastObject] = []
3031
feast_objects.extend(feature_views.values())
3132
feast_objects.extend([driver(), customer(), location(), feature_service])
3233
fs.apply(feast_objects)
3334
fs.materialize(environment.start_date, environment.end_date)
3435

35-
sample_drivers = random.sample(entities["driver"], 10)
36+
sample_drivers = random.sample(entities.driver_vals, 10)
3637

37-
sample_customers = random.sample(entities["customer"], 10)
38+
sample_customers = random.sample(entities.customer_vals, 10)
3839

3940
entity_rows = [
4041
{"driver": d, "customer_id": c, "val_to_add": 50}

‎sdk/python/tests/conftest.py

+4-13
Original file line numberDiff line numberDiff line change
@@ -32,10 +32,9 @@
3232
FULL_REPO_CONFIGS,
3333
REDIS_CONFIG,
3434
Environment,
35+
TestData,
3536
construct_test_environment,
36-
construct_universal_data_sources,
37-
construct_universal_datasets,
38-
construct_universal_entities,
37+
construct_universal_test_data,
3938
)
4039

4140
logger = logging.getLogger(__name__)
@@ -184,21 +183,13 @@ def cleanup():
184183

185184

186185
@pytest.fixture(scope="session")
187-
def universal_data_sources(request, environment):
188-
entities = construct_universal_entities()
189-
datasets = construct_universal_datasets(
190-
entities, environment.start_date, environment.end_date
191-
)
192-
datasources = construct_universal_data_sources(
193-
datasets, environment.data_source_creator
194-
)
195-
186+
def universal_data_sources(request, environment) -> TestData:
196187
def cleanup():
197188
# logger.info("Running cleanup in %s, Request: %s", worker_id, request.param)
198189
environment.data_source_creator.teardown()
199190

200191
request.addfinalizer(cleanup)
201-
return entities, datasets, datasources
192+
return construct_universal_test_data(environment)
202193

203194

204195
@pytest.fixture(scope="session")

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

+2-2
Original file line numberDiff line numberDiff line change
@@ -67,7 +67,7 @@ def test_historical_retrieval_with_validation(environment, universal_data_source
6767

6868
store.apply([driver(), customer(), location(), *feature_views.values()])
6969

70-
entity_df = datasets["entity"].drop(
70+
entity_df = datasets.entity_df.drop(
7171
columns=["order_id", "origin_id", "destination_id"]
7272
)
7373

@@ -101,7 +101,7 @@ def test_historical_retrieval_fails_on_validation(environment, universal_data_so
101101

102102
store.apply([driver(), customer(), location(), *feature_views.values()])
103103

104-
entity_df = datasets["entity"].drop(
104+
entity_df = datasets.entity_df.drop(
105105
columns=["order_id", "origin_id", "destination_id"]
106106
)
107107

0 commit comments

Comments
 (0)