Skip to content

Commit

Permalink
Allow apply to take a single Entity or FeatureView (#1457)
Browse files Browse the repository at this point in the history
Signed-off-by: Jacob Klegar <jacob@tecton.ai>
  • Loading branch information
jklegar authored Apr 13, 2021
1 parent 74ff628 commit c6c9f34
Show file tree
Hide file tree
Showing 2 changed files with 8 additions and 4 deletions.
6 changes: 5 additions & 1 deletion sdk/python/feast/feature_store.py
Original file line number Diff line number Diff line change
Expand Up @@ -176,7 +176,9 @@ def delete_feature_view(self, name: str):

return self._registry.delete_feature_view(name, self.project)

def apply(self, objects: List[Union[FeatureView, Entity]]):
def apply(
self, objects: Union[Entity, FeatureView, List[Union[FeatureView, Entity]]]
):
"""Register objects to metadata store and update related infrastructure.
The apply method registers one or more definitions (e.g., Entity, FeatureView) and registers or updates these
Expand Down Expand Up @@ -210,6 +212,8 @@ def apply(self, objects: List[Union[FeatureView, Entity]]):
# TODO: Add locking
# TODO: Optimize by only making a single call (read/write)

if isinstance(objects, Entity) or isinstance(objects, FeatureView):
objects = [objects]
views_to_update = []
for ob in objects:
if isinstance(ob, FeatureView):
Expand Down
6 changes: 3 additions & 3 deletions sdk/python/tests/test_feature_store.py
Original file line number Diff line number Diff line change
Expand Up @@ -79,8 +79,8 @@ def test_apply_entity_success(test_feature_store):
labels={"team": "matchmaking"},
)

# Register Entity with Core
test_feature_store.apply([entity])
# Register Entity
test_feature_store.apply(entity)

entities = test_feature_store.list_entities()

Expand All @@ -107,7 +107,7 @@ def test_apply_entity_integration(test_feature_store):
labels={"team": "matchmaking"},
)

# Register Entity with Core
# Register Entity
test_feature_store.apply([entity])

entities = test_feature_store.list_entities()
Expand Down

0 comments on commit c6c9f34

Please sign in to comment.