Enforce offline_store during feast apply, rename entity_dataset_name …

…to dataset Signed-off-by: Tsotne Tabidze <tsotne@tecton.ai>
feast-dev · feast-ci-bot · May 19, 2021 · May 10, 2021 · May 18, 2021 · May 19, 2021
commit 100f88df511670086aa7d21475fbed03b96ae806
@@ -60,3 +60,10 @@ def __init__(self, extras_type: str, nested_error: str):
             + f"You may need run {Style.BRIGHT + Fore.GREEN}pip install 'feast[{extras_type}]'{Style.RESET_ALL}"
         )
         super().__init__(message)
+
+
+class FeastOfflineStoreUnsupportedDataSource(Exception):
+    def __init__(self, offline_store_name: str, data_source_name: str):
+        super().__init__(
+            f"Offline Store '{offline_store_name}' does not support data source '{data_source_name}'"
+        )
@@ -5,15 +5,14 @@
 
 import mmh3
 import pandas
-import pyarrow
 from tqdm import tqdm
 
 from feast import FeatureTable, utils
 from feast.entity import Entity
 from feast.errors import FeastProviderLoginError
 from feast.feature_view import FeatureView
 from feast.infra.key_encoding_utils import serialize_entity_key
-from feast.infra.offline_stores.helpers import get_offline_store_from_sources
+from feast.infra.offline_stores.helpers import get_offline_store_from_config
 from feast.infra.provider import (
     Provider,
     RetrievalJob,
@@ -28,7 +27,7 @@
 
 try:
     from google.auth.exceptions import DefaultCredentialsError
-    from google.cloud import bigquery, datastore
+    from google.cloud import datastore
 except ImportError as e:
     from feast.errors import FeastExtrasDependencyImportError
 
@@ -40,11 +39,14 @@ class GcpProvider(Provider):
 
     def __init__(self, config: RepoConfig):
         assert isinstance(config.online_store, DatastoreOnlineStoreConfig)
+        assert config.offline_store is not None
         if config and config.online_store and config.online_store.project_id:
             self._gcp_project_id = config.online_store.project_id
         else:
             self._gcp_project_id = None
 
+        self.offline_store = get_offline_store_from_config(config.offline_store)
+
     def _initialize_client(self):
         try:
             if self._gcp_project_id is not None:
@@ -168,8 +170,7 @@ def materialize_single_feature_view(
         start_date = utils.make_tzaware(start_date)
         end_date = utils.make_tzaware(end_date)
 
-        offline_store = get_offline_store_from_sources([feature_view.input])
-        table = offline_store.pull_latest_from_table_or_query(
+        table = self.offline_store.pull_latest_from_table_or_query(
             data_source=feature_view.input,
             join_key_columns=join_key_columns,
             feature_name_columns=feature_name_columns,
@@ -193,25 +194,16 @@ def materialize_single_feature_view(
         feature_view.materialization_intervals.append((start_date, end_date))
         registry.apply_feature_view(feature_view, project)
 
-    @staticmethod
-    def _pull_query(query: str) -> pyarrow.Table:
-        client = bigquery.Client()
-        query_job = client.query(query)
-        return query_job.to_arrow()
-
-    @staticmethod
     def get_historical_features(
+        self,
         config: RepoConfig,
         feature_views: List[FeatureView],
         feature_refs: List[str],
         entity_df: Union[pandas.DataFrame, str],
         registry: Registry,
         project: str,
     ) -> RetrievalJob:
-        offline_store = get_offline_store_from_sources(
-            [feature_view.input for feature_view in feature_views]
-        )
-        job = offline_store.get_historical_features(
+        job = self.offline_store.get_historical_features(
             config=config,
             feature_views=feature_views,
             feature_refs=feature_refs,

@@ -12,7 +12,7 @@
 from feast.entity import Entity
 from feast.feature_view import FeatureView
 from feast.infra.key_encoding_utils import serialize_entity_key
-from feast.infra.offline_stores.helpers import get_offline_store_from_sources
+from feast.infra.offline_stores.helpers import get_offline_store_from_config
 from feast.infra.provider import (
     Provider,
     RetrievalJob,
@@ -30,16 +30,15 @@ class LocalProvider(Provider):
     _db_path: Path
 
     def __init__(self, config: RepoConfig, repo_path: Path):
-
         assert config is not None
-        assert config.online_store is not None
-        local_online_store_config = config.online_store
-        assert isinstance(local_online_store_config, SqliteOnlineStoreConfig)
-        local_path = Path(local_online_store_config.path)
+        assert isinstance(config.online_store, SqliteOnlineStoreConfig)
+        assert config.offline_store is not None
+        local_path = Path(config.online_store.path)
         if local_path.is_absolute():
             self._db_path = local_path
         else:
             self._db_path = repo_path.joinpath(local_path)
+        self.offline_store = get_offline_store_from_config(config.offline_store)
 
     def _get_conn(self):
         Path(self._db_path).parent.mkdir(exist_ok=True)
@@ -184,8 +183,7 @@ def materialize_single_feature_view(
         start_date = utils.make_tzaware(start_date)
         end_date = utils.make_tzaware(end_date)
 
-        offline_store = get_offline_store_from_sources([feature_view.input])
-        table = offline_store.pull_latest_from_table_or_query(
+        table = self.offline_store.pull_latest_from_table_or_query(
             data_source=feature_view.input,
             join_key_columns=join_key_columns,
             feature_name_columns=feature_name_columns,
@@ -209,19 +207,16 @@ def materialize_single_feature_view(
         feature_view.materialization_intervals.append((start_date, end_date))
         registry.apply_feature_view(feature_view, project)
 
-    @staticmethod
     def get_historical_features(
+        self,
         config: RepoConfig,
         feature_views: List[FeatureView],
         feature_refs: List[str],
         entity_df: Union[pd.DataFrame, str],
         registry: Registry,
         project: str,
     ) -> RetrievalJob:
-        offline_store = get_offline_store_from_sources(
-            [feature_view.input for feature_view in feature_views]
-        )
-        return offline_store.get_historical_features(
+        return self.offline_store.get_historical_features(
             config=config,
             feature_views=feature_views,
             feature_refs=feature_refs,

@@ -102,7 +102,7 @@ def get_historical_features(
 
             table_id = _upload_entity_df_into_bigquery(
                 config.project,
-                getattr(config.offline_store, "entity_dataset_name", "feast"),
+                getattr(config.offline_store, "dataset", "feast"),
                 entity_df,
                 client,
             )
@@ -201,13 +201,11 @@ class FeatureViewQueryContext:
     entity_selections: List[str]
 
 
-def _upload_entity_df_into_bigquery(
-    project, entity_dataset_name, entity_df, client
-) -> str:
+def _upload_entity_df_into_bigquery(project, dataset_name, entity_df, client) -> str:
     """Uploads a Pandas entity dataframe into a BigQuery table and returns a reference to the resulting table"""
 
     # First create the BigQuery dataset if it doesn't exist
-    dataset = bigquery.Dataset(f"{client.project}.{entity_dataset_name}")
+    dataset = bigquery.Dataset(f"{client.project}.{dataset_name}")
     dataset.location = "US"
     client.create_dataset(
         dataset, exists_ok=True
@@ -218,9 +216,7 @@ def _upload_entity_df_into_bigquery(
 
     # Upload the dataframe into BigQuery, creating a temporary table
     job_config = bigquery.LoadJobConfig()
-    table_id = (
-        f"{client.project}.{entity_dataset_name}.entity_df_{project}_{int(time.time())}"
-    )
+    table_id = f"{client.project}.{dataset_name}.entity_df_{project}_{int(time.time())}"
     job = client.load_table_from_dataframe(entity_df, table_id, job_config=job_config,)
     job.result()
 

@@ -1,28 +1,41 @@
-from typing import List
-
 from feast.data_source import BigQuerySource, DataSource, FileSource
+from feast.errors import FeastOfflineStoreUnsupportedDataSource
 from feast.infra.offline_stores.offline_store import OfflineStore
+from feast.repo_config import (
+    BigqueryOfflineStoreConfig,
+    FileOfflineStoreConfig,
+    OfflineStoreConfig,
+)
 
 
-def get_offline_store_from_sources(sources: List[DataSource]) -> OfflineStore:
-    """Detect which offline store should be used for retrieving historical features"""
-
-    source_types = [type(source) for source in sources]
+def get_offline_store_from_config(
+    offline_store_config: OfflineStoreConfig,
+) -> OfflineStore:
+    """Get the offline store from offline store config"""
 
-    # Retrieve features from ParquetOfflineStore
-    if all(source == FileSource for source in source_types):
+    if isinstance(offline_store_config, FileOfflineStoreConfig):
         from feast.infra.offline_stores.file import FileOfflineStore
 
         return FileOfflineStore()
-
-    # Retrieve features from BigQueryOfflineStore
-    if all(source == BigQuerySource for source in source_types):
+    elif isinstance(offline_store_config, BigqueryOfflineStoreConfig):
         from feast.infra.offline_stores.bigquery import BigQueryOfflineStore
 
         return BigQueryOfflineStore()
 
-    # Could not map inputs to an OfflineStore implementation
-    raise NotImplementedError(
-        "Unsupported combination of feature view input source types. Please ensure that all source types are "
-        "consistent and available in the same offline store."
+    raise ValueError(f"Unsupported offline store config '{offline_store_config}'")
+
+
+def assert_offline_store_supports_data_source(
+    offline_store_config: OfflineStoreConfig, data_source: DataSource
+):
+    if (
+        isinstance(offline_store_config, FileOfflineStoreConfig)
+        and isinstance(data_source, FileSource)
+    ) or (
+        isinstance(offline_store_config, BigqueryOfflineStoreConfig)
+        and isinstance(data_source, BigQuerySource)
+    ):
+        return
+    raise FeastOfflineStoreUnsupportedDataSource(
+        offline_store_config.type, data_source.__class__.__name__
     )
@@ -107,9 +107,9 @@ def materialize_single_feature_view(
     ) -> None:
         pass
 
-    @staticmethod
     @abc.abstractmethod
     def get_historical_features(
+        self,
         config: RepoConfig,
         feature_views: List[FeatureView],
         feature_refs: List[str],

@@ -50,8 +50,8 @@ class BigqueryOfflineStoreConfig(FeastBaseModel):
     type: Literal["bigquery"] = "bigquery"
     """ Offline store type selector"""
 
-    entity_dataset_name: Optional[StrictStr] = "feast"
-    """ (optional) Bigquery Dataset name for uploaded entity dataframes for point-in-time join """
+    dataset: Optional[StrictStr] = "feast"
+    """ (optional) Bigquery Dataset name for temporary tables """
 
 
 OfflineStoreConfig = Union[FileOfflineStoreConfig, BigqueryOfflineStoreConfig]

@@ -11,6 +11,7 @@
 
 from feast import Entity, FeatureTable
 from feast.feature_view import FeatureView
+from feast.infra.offline_stores.helpers import assert_offline_store_supports_data_source
 from feast.infra.provider import get_provider
 from feast.names import adjectives, animals
 from feast.registry import Registry
@@ -130,6 +131,14 @@ def apply_total(repo_config: RepoConfig, repo_path: Path):
     for t in repo.feature_views:
         repo_table_names.add(t.name)
 
+    data_sources = [t.input for t in repo.feature_views]
+
+    # Make sure the data source used by this feature view is supported by
+    for data_source in data_sources:
+        assert_offline_store_supports_data_source(
+            repo_config.offline_store, data_source
+        )
+
     tables_to_delete = []
     for registry_table in registry.list_feature_tables(project=project):
         if registry_table.name not in repo_table_names:

@@ -40,7 +40,7 @@ def run_with_output(self, args: List[str], cwd: Path) -> Tuple[int, bytes]:
             return e.returncode, e.output
 
     @contextmanager
-    def local_repo(self, example_repo_py: str):
+    def local_repo(self, example_repo_py: str, offline_store: str):
         """
         Convenience method to set up all the boilerplate for a local feature repo.
         """
@@ -63,6 +63,8 @@ def local_repo(self, example_repo_py: str):
             provider: local
             online_store:
                 path: {data_path / "online_store.db"}
+            offline_store:
+                type: {offline_store}
             """
                 )
             )

@@ -31,6 +31,8 @@ def test_workflow() -> None:
         provider: local
         online_store:
             path: {data_path / "online_store.db"}
+        offline_store:
+            type: bigquery
         """
             )
         )
@@ -96,6 +98,8 @@ def test_non_local_feature_repo() -> None:
         provider: local
         online_store:
             path: data/online_store.db
+        offline_store:
+            type: bigquery
         """
             )
         )

@@ -76,7 +76,8 @@ def test_e2e_local() -> None:
         with runner.local_repo(
             get_example_repo("example_feature_repo_2.py").replace(
                 "%PARQUET_PATH%", driver_stats_path
-            )
+            ),
+            "file",
         ) as store:
 
             assert store.repo_path is not None

@@ -407,7 +407,7 @@ def test_historical_features_from_bigquery_sources(
                     ),
                     provider="gcp",
                     offline_store=BigqueryOfflineStoreConfig(
-                        type="bigquery", entity_dataset_name="foo",
+                        type="bigquery", dataset="foo"
                     ),
                 )
             )

@@ -17,7 +17,9 @@ def test_online() -> None:
     Test reading from the online store in local mode.
     """
     runner = CliRunner()
-    with runner.local_repo(get_example_repo("example_feature_repo_1.py")) as store:
+    with runner.local_repo(
+        get_example_repo("example_feature_repo_1.py"), "bigquery"
+    ) as store:
         # Write some data to two tables
 
         driver_locations_fv = store.get_feature_view(name="driver_locations")

@@ -12,7 +12,9 @@ def test_partial() -> None:
     """
 
     runner = CliRunner()
-    with runner.local_repo(get_example_repo("example_feature_repo_1.py")) as store:
+    with runner.local_repo(
+        get_example_repo("example_feature_repo_1.py"), "bigquery"
+    ) as store:
 
         driver_locations_source = BigQuerySource(
             table_ref="rh_prod.ride_hailing_co.drivers",