Skip to content

Commit 7bff5ed

Browse files
felixwang9817felixwang9817
authored
Implement feature_store._apply_diffs to handle registry and infra diffs (#2238)
* Implement feature_store._apply_diffs to handle registry and infra diffs Signed-off-by: Felix Wang <wangfelix98@gmail.com> * Add logging Signed-off-by: Felix Wang <wangfelix98@gmail.com> * Factor out validation and inference logic Signed-off-by: Felix Wang <wangfelix98@gmail.com> * Clean up FeastObjectType enum Signed-off-by: Felix Wang <wangfelix98@gmail.com> * Move logic into enum Signed-off-by: Felix Wang <wangfelix98@gmail.com> * Add TODO Signed-off-by: Felix Wang <wangfelix98@gmail.com> * Add positional argument name Signed-off-by: Felix Wang <wangfelix98@gmail.com>
1 parent 46c4722 commit 7bff5ed

File tree

5 files changed

+289
-149
lines changed

5 files changed

+289
-149
lines changed

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

Lines changed: 87 additions & 44 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@
55
from feast.diff.property_diff import PropertyDiff, TransitionType
66
from feast.entity import Entity
77
from feast.feature_service import FeatureService
8+
from feast.feature_view import DUMMY_ENTITY_NAME
89
from feast.protos.feast.core.Entity_pb2 import Entity as EntityProto
910
from feast.protos.feast.core.FeatureService_pb2 import (
1011
FeatureService as FeatureServiceProto,
@@ -16,26 +17,16 @@
1617
from feast.protos.feast.core.RequestFeatureView_pb2 import (
1718
RequestFeatureView as RequestFeatureViewProto,
1819
)
19-
from feast.registry import FeastObjectType, Registry
20+
from feast.registry import FEAST_OBJECT_TYPES, FeastObjectType, Registry
2021
from feast.repo_contents import RepoContents
2122

22-
FEAST_OBJECT_TYPE_TO_STR = {
23-
FeastObjectType.ENTITY: "entity",
24-
FeastObjectType.FEATURE_VIEW: "feature view",
25-
FeastObjectType.ON_DEMAND_FEATURE_VIEW: "on demand feature view",
26-
FeastObjectType.REQUEST_FEATURE_VIEW: "request feature view",
27-
FeastObjectType.FEATURE_SERVICE: "feature service",
28-
}
29-
30-
FEAST_OBJECT_TYPES = FEAST_OBJECT_TYPE_TO_STR.keys()
31-
3223
Fco = TypeVar("Fco", Entity, BaseFeatureView, FeatureService)
3324

3425

3526
@dataclass
3627
class FcoDiff(Generic[Fco]):
3728
name: str
38-
fco_type: str
29+
fco_type: FeastObjectType
3930
current_fco: Fco
4031
new_fco: Fco
4132
fco_property_diffs: List[PropertyDiff]
@@ -52,6 +43,28 @@ def __init__(self):
5243
def add_fco_diff(self, fco_diff: FcoDiff):
5344
self.fco_diffs.append(fco_diff)
5445

46+
def to_string(self):
47+
from colorama import Fore, Style
48+
49+
log_string = ""
50+
51+
message_action_map = {
52+
TransitionType.CREATE: ("Created", Fore.GREEN),
53+
TransitionType.DELETE: ("Deleted", Fore.RED),
54+
TransitionType.UNCHANGED: ("Unchanged", Fore.LIGHTBLUE_EX),
55+
TransitionType.UPDATE: ("Updated", Fore.YELLOW),
56+
}
57+
for fco_diff in self.fco_diffs:
58+
if fco_diff.name == DUMMY_ENTITY_NAME:
59+
continue
60+
action, color = message_action_map[fco_diff.transition_type]
61+
log_string += f"{action} {fco_diff.fco_type.value} {Style.BRIGHT + color}{fco_diff.name}{Style.RESET_ALL}\n"
62+
if fco_diff.transition_type == TransitionType.UPDATE:
63+
for _p in fco_diff.fco_property_diffs:
64+
log_string += f"\t{_p.property_name}: {Style.BRIGHT + color}{_p.val_existing}{Style.RESET_ALL} -> {Style.BRIGHT + Fore.LIGHTGREEN_EX}{_p.val_declared}{Style.RESET_ALL}\n"
65+
66+
return log_string
67+
5568

5669
def tag_objects_for_keep_delete_update_add(
5770
existing_objs: Iterable[Fco], desired_objs: Iterable[Fco]
@@ -93,7 +106,9 @@ def tag_proto_objects_for_keep_delete_add(
93106
FIELDS_TO_IGNORE = {"project"}
94107

95108

96-
def diff_registry_objects(current: Fco, new: Fco, object_type: str) -> FcoDiff:
109+
def diff_registry_objects(
110+
current: Fco, new: Fco, object_type: FeastObjectType
111+
) -> FcoDiff:
97112
current_proto = current.to_proto()
98113
new_proto = new.to_proto()
99114
assert current_proto.DESCRIPTOR.full_name == new_proto.DESCRIPTOR.full_name
@@ -145,30 +160,12 @@ def extract_objects_for_keep_delete_update_add(
145160
objs_to_update = {}
146161
objs_to_add = {}
147162

148-
registry_object_type_to_objects: Dict[FeastObjectType, List[Any]]
149-
registry_object_type_to_objects = {
150-
FeastObjectType.ENTITY: registry.list_entities(project=current_project),
151-
FeastObjectType.FEATURE_VIEW: registry.list_feature_views(
152-
project=current_project
153-
),
154-
FeastObjectType.ON_DEMAND_FEATURE_VIEW: registry.list_on_demand_feature_views(
155-
project=current_project
156-
),
157-
FeastObjectType.REQUEST_FEATURE_VIEW: registry.list_request_feature_views(
158-
project=current_project
159-
),
160-
FeastObjectType.FEATURE_SERVICE: registry.list_feature_services(
161-
project=current_project
162-
),
163-
}
164-
registry_object_type_to_repo_contents: Dict[FeastObjectType, Set[Any]]
165-
registry_object_type_to_repo_contents = {
166-
FeastObjectType.ENTITY: desired_repo_contents.entities,
167-
FeastObjectType.FEATURE_VIEW: desired_repo_contents.feature_views,
168-
FeastObjectType.ON_DEMAND_FEATURE_VIEW: desired_repo_contents.on_demand_feature_views,
169-
FeastObjectType.REQUEST_FEATURE_VIEW: desired_repo_contents.request_feature_views,
170-
FeastObjectType.FEATURE_SERVICE: desired_repo_contents.feature_services,
171-
}
163+
registry_object_type_to_objects: Dict[
164+
FeastObjectType, List[Any]
165+
] = FeastObjectType.get_objects_from_registry(registry, current_project)
166+
registry_object_type_to_repo_contents: Dict[
167+
FeastObjectType, Set[Any]
168+
] = FeastObjectType.get_objects_from_repo_contents(desired_repo_contents)
172169

173170
for object_type in FEAST_OBJECT_TYPES:
174171
(
@@ -221,7 +218,7 @@ def diff_between(
221218
diff.add_fco_diff(
222219
FcoDiff(
223220
name=e.name,
224-
fco_type=FEAST_OBJECT_TYPE_TO_STR[object_type],
221+
fco_type=object_type,
225222
current_fco=None,
226223
new_fco=e,
227224
fco_property_diffs=[],
@@ -232,7 +229,7 @@ def diff_between(
232229
diff.add_fco_diff(
233230
FcoDiff(
234231
name=e.name,
235-
fco_type=FEAST_OBJECT_TYPE_TO_STR[object_type],
232+
fco_type=object_type,
236233
current_fco=e,
237234
new_fco=None,
238235
fco_property_diffs=[],
@@ -241,10 +238,56 @@ def diff_between(
241238
)
242239
for e in objects_to_update:
243240
current_obj = [_e for _e in objects_to_keep if _e.name == e.name][0]
244-
diff.add_fco_diff(
245-
diff_registry_objects(
246-
current_obj, e, FEAST_OBJECT_TYPE_TO_STR[object_type]
247-
)
248-
)
241+
diff.add_fco_diff(diff_registry_objects(current_obj, e, object_type))
249242

250243
return diff
244+
245+
246+
def apply_diff_to_registry(
247+
registry: Registry, registry_diff: RegistryDiff, project: str, commit: bool = True
248+
):
249+
"""
250+
Applies the given diff to the given Feast project in the registry.
251+
252+
Args:
253+
registry: The registry to be updated.
254+
registry_diff: The diff to apply.
255+
project: Feast project to be updated.
256+
commit: Whether the change should be persisted immediately
257+
"""
258+
for fco_diff in registry_diff.fco_diffs:
259+
# There is no need to delete the FCO on an update, since applying the new FCO
260+
# will automatically delete the existing FCO.
261+
if fco_diff.transition_type == TransitionType.DELETE:
262+
if fco_diff.fco_type == FeastObjectType.ENTITY:
263+
registry.delete_entity(fco_diff.current_fco.name, project, commit=False)
264+
elif fco_diff.fco_type == FeastObjectType.FEATURE_SERVICE:
265+
registry.delete_feature_service(
266+
fco_diff.current_fco.name, project, commit=False
267+
)
268+
elif fco_diff.fco_type in [
269+
FeastObjectType.FEATURE_VIEW,
270+
FeastObjectType.ON_DEMAND_FEATURE_VIEW,
271+
FeastObjectType.REQUEST_FEATURE_VIEW,
272+
]:
273+
registry.delete_feature_view(
274+
fco_diff.current_fco.name, project, commit=False,
275+
)
276+
277+
if fco_diff.transition_type in [
278+
TransitionType.CREATE,
279+
TransitionType.UPDATE,
280+
]:
281+
if fco_diff.fco_type == FeastObjectType.ENTITY:
282+
registry.apply_entity(fco_diff.new_fco, project, commit=False)
283+
elif fco_diff.fco_type == FeastObjectType.FEATURE_SERVICE:
284+
registry.apply_feature_service(fco_diff.new_fco, project, commit=False)
285+
elif fco_diff.fco_type in [
286+
FeastObjectType.FEATURE_VIEW,
287+
FeastObjectType.ON_DEMAND_FEATURE_VIEW,
288+
FeastObjectType.REQUEST_FEATURE_VIEW,
289+
]:
290+
registry.apply_feature_view(fco_diff.new_fco, project, commit=False)
291+
292+
if commit:
293+
registry.commit()

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

Lines changed: 18 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -60,7 +60,24 @@ def update(self):
6060
infra_object.update()
6161

6262
def to_string(self):
63-
pass
63+
from colorama import Fore, Style
64+
65+
log_string = ""
66+
67+
message_action_map = {
68+
TransitionType.CREATE: ("Created", Fore.GREEN),
69+
TransitionType.DELETE: ("Deleted", Fore.RED),
70+
TransitionType.UNCHANGED: ("Unchanged", Fore.LIGHTBLUE_EX),
71+
TransitionType.UPDATE: ("Updated", Fore.YELLOW),
72+
}
73+
for infra_object_diff in self.infra_object_diffs:
74+
action, color = message_action_map[infra_object_diff.transition_type]
75+
log_string += f"{action} {infra_object_diff.infra_object_type} {Style.BRIGHT + color}{infra_object_diff.name}{Style.RESET_ALL}\n"
76+
if infra_object_diff.transition_type == TransitionType.UPDATE:
77+
for _p in infra_object_diff.infra_object_property_diffs:
78+
log_string += f"\t{_p.property_name}: {Style.BRIGHT + color}{_p.val_existing}{Style.RESET_ALL} -> {Style.BRIGHT + Fore.LIGHTGREEN_EX}{_p.val_declared}{Style.RESET_ALL}\n"
79+
80+
return log_string
6481

6582

6683
def tag_infra_proto_objects_for_keep_delete_add(

0 commit comments

Comments
 (0)