diff --git a/src/zenml/integrations/feast/__init__.py b/src/zenml/integrations/feast/__init__.py index 2746f93316b..efa1b672947 100644 --- a/src/zenml/integrations/feast/__init__.py +++ b/src/zenml/integrations/feast/__init__.py @@ -31,7 +31,7 @@ class FeastIntegration(Integration): NAME = FEAST # click is added to keep the feast click version in sync with ZenML's click - REQUIREMENTS = ["feast", "click>=8.0.1,<8.1.4"] + REQUIREMENTS = ["feast>=0.12.0", "click>=8.0.1,<8.1.4"] REQUIREMENTS_IGNORED_ON_UNINSTALL = ["click", "pandas"] @classmethod diff --git a/src/zenml/integrations/feast/feature_stores/feast_feature_store.py b/src/zenml/integrations/feast/feature_stores/feast_feature_store.py index 4213d540d36..ffd4e84340f 100644 --- a/src/zenml/integrations/feast/feature_stores/feast_feature_store.py +++ b/src/zenml/integrations/feast/feature_stores/feast_feature_store.py @@ -16,7 +16,7 @@ from typing import Any, Dict, List, Union, cast import pandas as pd -from feast import FeatureStore # type: ignore +from feast import FeatureService, FeatureStore # type: ignore from feast.infra.registry.base_registry import BaseRegistry # type: ignore from zenml.feature_stores.base_feature_store import BaseFeatureStore @@ -43,14 +43,14 @@ def config(self) -> FeastFeatureStoreConfig: def get_historical_features( self, entity_df: Union[pd.DataFrame, str], - features: List[str], + features: Union[List[str], FeatureService], full_feature_names: bool = False, ) -> pd.DataFrame: """Returns the historical features for training or batch scoring. Args: entity_df: The entity DataFrame or entity name. - features: The features to retrieve. + features: The features to retrieve or a FeatureService. full_feature_names: Whether to return the full feature names. Raise: @@ -70,14 +70,14 @@ def get_historical_features( def get_online_features( self, entity_rows: List[Dict[str, Any]], - features: List[str], + features: Union[List[str], FeatureService], full_feature_names: bool = False, ) -> Dict[str, Any]: """Returns the latest online feature data. Args: entity_rows: The entity rows to retrieve. - features: The features to retrieve. + features: The features to retrieve or a FeatureService. full_feature_names: Whether to return the full feature names. Raise: @@ -118,17 +118,21 @@ def get_entities(self) -> List[str]: fs = FeatureStore(repo_path=self.config.feast_repo) return [ds.name for ds in fs.list_entities()] - def get_feature_services(self) -> List[str]: - """Returns the feature service names. + def get_feature_services(self) -> List[FeatureService]: + """Returns the feature services. Raise: ConnectionError: If the online component (Redis) is not available. Returns: - The feature service names. + The feature services. """ fs = FeatureStore(repo_path=self.config.feast_repo) - return [ds.name for ds in fs.list_feature_services()] + feature_services: List[FeatureService] = list( + fs.list_feature_services() + ) + + return feature_services def get_feature_views(self) -> List[str]: """Returns the feature view names.