Skip to content

Commit 1b6bdf1

Browse files
achalsachals
committed
Use the same tag method
Signed-off-by: Achal Shah <achals@gmail.com>
1 parent 0d18cb1 commit 1b6bdf1

3 files changed

Lines changed: 26 additions & 75 deletions

File tree

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

Lines changed: 17 additions & 61 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,9 @@
11
from dataclasses import dataclass
22
from enum import Enum
3-
from typing import Any, List, Set, Tuple, Union
3+
from typing import Any, List, Set, Tuple, Union, TypeVar
4+
from unittest.mock import Base
45

6+
from feast.base_feature_view import BaseFeatureView
57
from feast.entity import Entity
68
from feast.feature_service import FeatureService
79
from feast.feature_table import FeatureTable
@@ -43,67 +45,21 @@ def add_fco_diff(self, fco_diff: FcoDiff):
4345
self.fco_diffs.append(fco_diff)
4446

4547

46-
def _tag_registry_entities_for_keep_delete(
47-
existing_entities: Set[Entity], desired_entities: Set[Entity]
48-
) -> Tuple[Set[Entity], Set[Entity], Set[Entity]]:
49-
existing_entity_names = {e.name for e in existing_entities}
50-
desired_entity_names = {e.name for e in desired_entities}
48+
T = TypeVar('T', Entity, BaseFeatureView, FeatureService, FeatureTable)
5149

52-
entities_to_add = {
53-
e for e in desired_entities if e.name not in existing_entity_names
54-
}
55-
entities_to_keep = {e for e in desired_entities if e.name in existing_entity_names}
56-
entities_to_delete = {
57-
e for e in existing_entities if e.name not in desired_entity_names
58-
}
5950

60-
return entities_to_keep, entities_to_delete, entities_to_add
61-
62-
63-
def _tag_registry_views_for_keep_delete(
64-
existing_views: Union[
65-
Set[FeatureView], Set[RequestFeatureView], Set[OnDemandFeatureView]
66-
],
67-
desired_views: Union[
68-
Set[FeatureView], Set[RequestFeatureView], Set[OnDemandFeatureView]
69-
],
70-
) -> Tuple[
71-
Set[Union[FeatureView, RequestFeatureView, OnDemandFeatureView]],
72-
Set[Union[FeatureView, RequestFeatureView, OnDemandFeatureView]],
73-
Set[Union[FeatureView, RequestFeatureView, OnDemandFeatureView]],
74-
]:
75-
existing_view_names = {v.name for v in existing_views}
76-
desired_view_names = {v.name for v in desired_views}
77-
78-
views_to_add = {v for v in desired_views if v.name not in existing_view_names}
79-
views_to_keep = {v for v in desired_views if v.name in existing_view_names}
80-
views_to_delete = {v for v in existing_views if v.name not in desired_view_names}
81-
return views_to_keep, views_to_delete, views_to_add
82-
83-
84-
def _tag_registry_tables_for_keep_delete(
85-
existing_tables: Set[FeatureTable], desired_tables: Set[FeatureTable]
86-
) -> Tuple[Set[FeatureTable], Set[FeatureTable], Set[FeatureTable]]:
87-
existing_table_names = {v.name for v in existing_tables}
88-
desired_table_names = {v.name for v in desired_tables}
89-
90-
tables_to_add = {t for t in desired_tables if t.name not in existing_table_names}
91-
tables_to_keep = {t for t in desired_tables if t.name in existing_table_names}
92-
tables_to_delete = {t for t in existing_tables if t.name not in desired_table_names}
93-
return tables_to_keep, tables_to_delete, tables_to_add
94-
95-
96-
def _tag_registry_services_for_keep_delete(
97-
existing_service: Set[FeatureService], desired_service: Set[FeatureService]
98-
) -> Tuple[Set[FeatureService], Set[FeatureService], Set[FeatureService]]:
99-
existing_service_names = {v.name for v in existing_service}
100-
desired_service_names = {v.name for v in desired_service}
101-
102-
services_to_add = {
103-
s for s in desired_service if s.name not in existing_service_names
51+
def tag_objects_for_keep_delete_add(
52+
existing_objs: Set[T], desired_objs: Set[T]
53+
) -> Tuple[Set[T], Set[T], Set[T]]:
54+
existing_obj_names = {e.name for e in existing_objs}
55+
desired_obj_names = {e.name for e in desired_objs}
56+
57+
objs_to_add = {
58+
e for e in desired_objs if e.name not in existing_obj_names
10459
}
105-
services_to_delete = {
106-
s for s in existing_service if s.name not in desired_service_names
60+
objs_to_keep = {e for e in desired_objs if e.name in existing_obj_names}
61+
objs_to_delete = {
62+
e for e in existing_objs if e.name not in desired_obj_names
10763
}
108-
services_to_keep = {s for s in desired_service if s.name in existing_service_names}
109-
return services_to_keep, services_to_delete, services_to_add
64+
65+
return objs_to_keep, objs_to_delete, objs_to_add

‎sdk/python/feast/registry.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -28,7 +28,7 @@
2828
FcoDiff,
2929
RegistryDiff,
3030
TransitionType,
31-
_tag_registry_entities_for_keep_delete,
31+
tag_objects_for_keep_delete_add,
3232
)
3333
from feast.entity import Entity
3434
from feast.errors import (
@@ -151,7 +151,7 @@ def diff_between(
151151
entities_to_keep,
152152
entities_to_delete,
153153
entities_to_add,
154-
) = _tag_registry_entities_for_keep_delete(
154+
) = tag_objects_for_keep_delete_add(
155155
set(current_registry.list_entities(project=project, allow_cache=True)),
156156
set(new_registry.list_entities(project=project, allow_cache=True)),
157157
)

‎sdk/python/feast/repo_operations.py‎

Lines changed: 7 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -12,12 +12,7 @@
1212
from click.exceptions import BadParameter
1313

1414
from feast.base_feature_view import BaseFeatureView
15-
from feast.diff.FcoDiff import (
16-
_tag_registry_entities_for_keep_delete,
17-
_tag_registry_services_for_keep_delete,
18-
_tag_registry_tables_for_keep_delete,
19-
_tag_registry_views_for_keep_delete,
20-
)
15+
from feast.diff.FcoDiff import tag_objects_for_keep_delete_add
2116
from feast.entity import Entity
2217
from feast.feature_service import FeatureService
2318
from feast.feature_store import FeatureStore
@@ -162,13 +157,13 @@ def apply_total(repo_config: RepoConfig, repo_path: Path, skip_source_validation
162157
entities_to_keep,
163158
entities_to_delete,
164159
entities_to_add,
165-
) = _tag_registry_entities_for_keep_delete(
160+
) = tag_objects_for_keep_delete_add(
166161
set(registry.list_entities(project=project)), repo.entities
167162
)
168163
# TODO(achals): This code path should be refactored to handle added & kept entities separately.
169164
entities_to_keep = entities_to_keep.union(entities_to_add)
170165

171-
views = _tag_registry_views_for_keep_delete(
166+
views = tag_objects_for_keep_delete_add(
172167
set(registry.list_feature_views(project=project)), repo.feature_views
173168
)
174169
views_to_keep, views_to_delete, views_to_add = (
@@ -177,7 +172,7 @@ def apply_total(repo_config: RepoConfig, repo_path: Path, skip_source_validation
177172
cast(Set[FeatureView], views[2]),
178173
)
179174

180-
request_views = _tag_registry_views_for_keep_delete(
175+
request_views = tag_objects_for_keep_delete_add(
181176
set(registry.list_request_feature_views(project=project)),
182177
repo.request_feature_views,
183178
)
@@ -201,7 +196,7 @@ def apply_total(repo_config: RepoConfig, repo_path: Path, skip_source_validation
201196
*request_views_to_delete,
202197
}
203198

204-
odfvs = _tag_registry_views_for_keep_delete(
199+
odfvs = tag_objects_for_keep_delete_add(
205200
set(registry.list_on_demand_feature_views(project=project)),
206201
repo.on_demand_feature_views,
207202
)
@@ -216,7 +211,7 @@ def apply_total(repo_config: RepoConfig, repo_path: Path, skip_source_validation
216211
tables_to_keep,
217212
tables_to_delete,
218213
tables_to_add,
219-
) = _tag_registry_tables_for_keep_delete(
214+
) = tag_objects_for_keep_delete_add(
220215
set(registry.list_feature_tables(project=project)), repo.feature_tables
221216
)
222217
tables_to_keep = tables_to_keep.union(tables_to_add)
@@ -225,7 +220,7 @@ def apply_total(repo_config: RepoConfig, repo_path: Path, skip_source_validation
225220
services_to_keep,
226221
services_to_delete,
227222
services_to_add,
228-
) = _tag_registry_services_for_keep_delete(
223+
) = tag_objects_for_keep_delete_add(
229224
set(registry.list_feature_services(project=project)), repo.feature_services
230225
)
231226
services_to_keep = services_to_keep.union(services_to_add)

0 commit comments

Comments
 (0)