Skip to content

Commit

Permalink
Add missing FeatureStore methods (#1423)
Browse files Browse the repository at this point in the history
* Add test feature store

Signed-off-by: Jacob Klegar <jacob@tecton.ai>

* Finish moving test

Signed-off-by: Jacob Klegar <jacob@tecton.ai>

* lint

Signed-off-by: Jacob Klegar <jacob@tecton.ai>

* Address comments

Signed-off-by: Jacob Klegar <jacob@tecton.ai>

* Use tempfile for online store path

Signed-off-by: Jacob Klegar <jacob@tecton.ai>
  • Loading branch information
jklegar authored Mar 31, 2021
1 parent 0892cd1 commit dd81448
Show file tree
Hide file tree
Showing 4 changed files with 303 additions and 143 deletions.
53 changes: 53 additions & 0 deletions sdk/python/feast/feature_store.py
Original file line number Diff line number Diff line change
Expand Up @@ -82,6 +82,59 @@ def _get_provider(self) -> Provider:
def _get_registry(self) -> Registry:
return Registry(self.config.metadata_store)

def list_entities(self) -> List[Entity]:
"""
Retrieve a list of entities from the registry
Returns:
List of entities
"""
return self._get_registry().list_entities(self.project)

def list_feature_views(self) -> List[FeatureView]:
"""
Retrieve a list of feature views from the registry
Returns:
List of feature views
"""
return self._get_registry().list_feature_views(self.project)

def get_entity(self, name: str) -> Entity:
"""
Retrieves an entity.
Args:
name: Name of entity
Returns:
Returns either the specified entity, or raises an exception if
none is found
"""
return self._get_registry().get_entity(name, self.project)

def get_feature_view(self, name: str) -> FeatureView:
"""
Retrieves a feature view.
Args:
name: Name of feature view
Returns:
Returns either the specified feature view, or raises an exception if
none is found
"""
return self._get_registry().get_feature_view(name, self.project)

def delete_feature_view(self, name: str):
"""
Deletes a feature view or raises an exception if not found.
Args:
name: Name of feature view
"""
return self._get_registry().delete_feature_view(name, self.project)

def apply(self, objects: List[Union[FeatureView, Entity]]):
"""Register objects to metadata store and update related infrastructure.
Expand Down
2 changes: 1 addition & 1 deletion sdk/python/feast/registry.py
Original file line number Diff line number Diff line change
Expand Up @@ -200,7 +200,7 @@ def list_feature_views(self, project: str) -> List[FeatureView]:
project: Filter feature tables based on project name
Returns:
List of feature tables
List of feature views
"""
registry_proto = self._registry_store.get_registry()
feature_views = []
Expand Down
144 changes: 2 additions & 142 deletions sdk/python/tests/test_client.py
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,6 @@
import socket
from concurrent import futures
from datetime import datetime, timedelta
from tempfile import mkstemp
from typing import Tuple
from unittest import mock

Expand Down Expand Up @@ -135,24 +134,6 @@ def secure_mock_client_with_auth(self):
client._serving_url = SERVING_URL
return client

@pytest.fixture
def client_with_object_registry(self):
fd, path = mkstemp()
return Client(registry_path=path,)

@pytest.fixture
def client_with_object_registry_gcs(self):
# clear any old registry left there
from google.cloud import storage

storage_client = storage.Client()
bucket = storage_client.bucket("feast-registry-test")
blob = bucket.blob("test")
if blob.exists():
blob.delete()

return Client(registry_path="gs://feast-registry-test/test",)

@pytest.fixture
def server_credentials(self):
private_key = pkgutil.get_data(__name__, _PRIVATE_KEY_RESOURCE_PATH)
Expand Down Expand Up @@ -389,12 +370,7 @@ def test_get_historical_features(self, mocked_client, mocker):
assert 1 == 1

@pytest.mark.parametrize(
"test_client",
[
lazy_fixture("client"),
lazy_fixture("secure_client"),
lazy_fixture("client_with_object_registry"),
],
"test_client", [lazy_fixture("client"), lazy_fixture("secure_client")],
)
def test_apply_entity_success(self, test_client):

Expand All @@ -420,50 +396,8 @@ def test_apply_entity_success(self, test_client):
and entity.labels["team"] == "matchmaking"
)

@pytest.mark.integration
@pytest.mark.parametrize(
"test_client", [lazy_fixture("client_with_object_registry_gcs")],
)
def test_apply_entity_integration(self, test_client):

entity = Entity(
name="driver_car_id",
description="Car driver id",
value_type=ValueType.STRING,
labels={"team": "matchmaking"},
)

# Register Entity with Core
test_client.apply(entity)

entities = test_client.list_entities()

entity = entities[0]
assert (
len(entities) == 1
and entity.name == "driver_car_id"
and entity.value_type == ValueType(ValueProto.ValueType.STRING)
and entity.description == "Car driver id"
and "team" in entity.labels
and entity.labels["team"] == "matchmaking"
)

entity = test_client.get_entity("driver_car_id")
assert (
entity.name == "driver_car_id"
and entity.value_type == ValueType(ValueProto.ValueType.STRING)
and entity.description == "Car driver id"
and "team" in entity.labels
and entity.labels["team"] == "matchmaking"
)

@pytest.mark.parametrize(
"test_client",
[
lazy_fixture("client"),
lazy_fixture("secure_client"),
lazy_fixture("client_with_object_registry"),
],
"test_client", [lazy_fixture("client"), lazy_fixture("secure_client")],
)
def test_apply_feature_table_success(self, test_client):

Expand Down Expand Up @@ -517,80 +451,6 @@ def test_apply_feature_table_success(self, test_client):
and feature_tables[0].entities[0] == "fs1-my-entity-1"
)

@pytest.mark.integration
@pytest.mark.parametrize(
"test_client", [lazy_fixture("client_with_object_registry_gcs")],
)
def test_apply_feature_table_integration(self, test_client):

# Create Feature Tables
batch_source = FileSource(
file_format=ParquetFormat(),
file_url="file://feast/*",
event_timestamp_column="ts_col",
created_timestamp_column="timestamp",
date_partition_column="date_partition_col",
)

stream_source = KafkaSource(
bootstrap_servers="localhost:9094",
message_format=ProtoFormat("class.path"),
topic="test_topic",
event_timestamp_column="ts_col",
)

ft1 = FeatureTable(
name="my-feature-table-1",
features=[
Feature(name="fs1-my-feature-1", dtype=ValueType.INT64),
Feature(name="fs1-my-feature-2", dtype=ValueType.STRING),
Feature(name="fs1-my-feature-3", dtype=ValueType.STRING_LIST),
Feature(name="fs1-my-feature-4", dtype=ValueType.BYTES_LIST),
],
entities=["fs1-my-entity-1"],
labels={"team": "matchmaking"},
batch_source=batch_source,
stream_source=stream_source,
)

# Register Feature Table with Core
test_client.apply(ft1)

feature_tables = test_client.list_feature_tables()

# List Feature Tables
assert (
len(feature_tables) == 1
and feature_tables[0].name == "my-feature-table-1"
and feature_tables[0].features[0].name == "fs1-my-feature-1"
and feature_tables[0].features[0].dtype == ValueType.INT64
and feature_tables[0].features[1].name == "fs1-my-feature-2"
and feature_tables[0].features[1].dtype == ValueType.STRING
and feature_tables[0].features[2].name == "fs1-my-feature-3"
and feature_tables[0].features[2].dtype == ValueType.STRING_LIST
and feature_tables[0].features[3].name == "fs1-my-feature-4"
and feature_tables[0].features[3].dtype == ValueType.BYTES_LIST
and feature_tables[0].entities[0] == "fs1-my-entity-1"
)

feature_table = test_client.get_feature_table("my-feature-table-1")
assert (
feature_table.name == "my-feature-table-1"
and feature_table.features[0].name == "fs1-my-feature-1"
and feature_table.features[0].dtype == ValueType.INT64
and feature_table.features[1].name == "fs1-my-feature-2"
and feature_table.features[1].dtype == ValueType.STRING
and feature_table.features[2].name == "fs1-my-feature-3"
and feature_table.features[2].dtype == ValueType.STRING_LIST
and feature_table.features[3].name == "fs1-my-feature-4"
and feature_table.features[3].dtype == ValueType.BYTES_LIST
and feature_table.entities[0] == "fs1-my-entity-1"
)

test_client.delete_feature_table("my-feature-table-1")
feature_tables = test_client.list_feature_tables()
assert len(feature_tables) == 0

@pytest.mark.parametrize(
"test_client", [lazy_fixture("client"), lazy_fixture("secure_client")]
)
Expand Down
Loading

0 comments on commit dd81448

Please sign in to comment.