Skip to content

Commit f0c8521

Browse files
author
Tsotne Tabidze
authored
Add offline_store config (#1552)
* Add offline_store config Signed-off-by: Tsotne Tabidze <tsotne@tecton.ai> * Enforce offline_store during feast apply, rename entity_dataset_name to dataset Signed-off-by: Tsotne Tabidze <tsotne@tecton.ai> * Remove ugly getattr since it's unnecessary anymore Signed-off-by: Tsotne Tabidze <tsotne@tecton.ai> * Rename Bigquery to BigQuery Signed-off-by: Tsotne Tabidze <tsotne@tecton.ai>
1 parent 3017dd4 commit f0c8521

14 files changed

Lines changed: 206 additions & 84 deletions

‎sdk/python/feast/errors.py‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -60,3 +60,10 @@ def __init__(self, extras_type: str, nested_error: str):
6060
+ f"You may need run {Style.BRIGHT + Fore.GREEN}pip install 'feast[{extras_type}]'{Style.RESET_ALL}"
6161
)
6262
super().__init__(message)
63+
64+
65+
class FeastOfflineStoreUnsupportedDataSource(Exception):
66+
def __init__(self, offline_store_name: str, data_source_name: str):
67+
super().__init__(
68+
f"Offline Store '{offline_store_name}' does not support data source '{data_source_name}'"
69+
)

‎sdk/python/feast/infra/gcp.py‎

Lines changed: 8 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -5,15 +5,14 @@
55

66
import mmh3
77
import pandas
8-
import pyarrow
98
from tqdm import tqdm
109

1110
from feast import FeatureTable, utils
1211
from feast.entity import Entity
1312
from feast.errors import FeastProviderLoginError
1413
from feast.feature_view import FeatureView
1514
from feast.infra.key_encoding_utils import serialize_entity_key
16-
from feast.infra.offline_stores.helpers import get_offline_store_from_sources
15+
from feast.infra.offline_stores.helpers import get_offline_store_from_config
1716
from feast.infra.provider import (
1817
Provider,
1918
RetrievalJob,
@@ -28,7 +27,7 @@
2827

2928
try:
3029
from google.auth.exceptions import DefaultCredentialsError
31-
from google.cloud import bigquery, datastore
30+
from google.cloud import datastore
3231
except ImportError as e:
3332
from feast.errors import FeastExtrasDependencyImportError
3433

@@ -40,11 +39,14 @@ class GcpProvider(Provider):
4039

4140
def __init__(self, config: RepoConfig):
4241
assert isinstance(config.online_store, DatastoreOnlineStoreConfig)
42+
assert config.offline_store is not None
4343
if config and config.online_store and config.online_store.project_id:
4444
self._gcp_project_id = config.online_store.project_id
4545
else:
4646
self._gcp_project_id = None
4747

48+
self.offline_store = get_offline_store_from_config(config.offline_store)
49+
4850
def _initialize_client(self):
4951
try:
5052
if self._gcp_project_id is not None:
@@ -168,8 +170,7 @@ def materialize_single_feature_view(
168170
start_date = utils.make_tzaware(start_date)
169171
end_date = utils.make_tzaware(end_date)
170172

171-
offline_store = get_offline_store_from_sources([feature_view.input])
172-
table = offline_store.pull_latest_from_table_or_query(
173+
table = self.offline_store.pull_latest_from_table_or_query(
173174
data_source=feature_view.input,
174175
join_key_columns=join_key_columns,
175176
feature_name_columns=feature_name_columns,
@@ -193,25 +194,16 @@ def materialize_single_feature_view(
193194
feature_view.materialization_intervals.append((start_date, end_date))
194195
registry.apply_feature_view(feature_view, project)
195196

196-
@staticmethod
197-
def _pull_query(query: str) -> pyarrow.Table:
198-
client = bigquery.Client()
199-
query_job = client.query(query)
200-
return query_job.to_arrow()
201-
202-
@staticmethod
203197
def get_historical_features(
198+
self,
204199
config: RepoConfig,
205200
feature_views: List[FeatureView],
206201
feature_refs: List[str],
207202
entity_df: Union[pandas.DataFrame, str],
208203
registry: Registry,
209204
project: str,
210205
) -> RetrievalJob:
211-
offline_store = get_offline_store_from_sources(
212-
[feature_view.input for feature_view in feature_views]
213-
)
214-
job = offline_store.get_historical_features(
206+
job = self.offline_store.get_historical_features(
215207
config=config,
216208
feature_views=feature_views,
217209
feature_refs=feature_refs,

‎sdk/python/feast/infra/local.py‎

Lines changed: 8 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,7 @@
1212
from feast.entity import Entity
1313
from feast.feature_view import FeatureView
1414
from feast.infra.key_encoding_utils import serialize_entity_key
15-
from feast.infra.offline_stores.helpers import get_offline_store_from_sources
15+
from feast.infra.offline_stores.helpers import get_offline_store_from_config
1616
from feast.infra.provider import (
1717
Provider,
1818
RetrievalJob,
@@ -30,16 +30,15 @@ class LocalProvider(Provider):
3030
_db_path: Path
3131

3232
def __init__(self, config: RepoConfig, repo_path: Path):
33-
3433
assert config is not None
35-
assert config.online_store is not None
36-
local_online_store_config = config.online_store
37-
assert isinstance(local_online_store_config, SqliteOnlineStoreConfig)
38-
local_path = Path(local_online_store_config.path)
34+
assert isinstance(config.online_store, SqliteOnlineStoreConfig)
35+
assert config.offline_store is not None
36+
local_path = Path(config.online_store.path)
3937
if local_path.is_absolute():
4038
self._db_path = local_path
4139
else:
4240
self._db_path = repo_path.joinpath(local_path)
41+
self.offline_store = get_offline_store_from_config(config.offline_store)
4342

4443
def _get_conn(self):
4544
Path(self._db_path).parent.mkdir(exist_ok=True)
@@ -184,8 +183,7 @@ def materialize_single_feature_view(
184183
start_date = utils.make_tzaware(start_date)
185184
end_date = utils.make_tzaware(end_date)
186185

187-
offline_store = get_offline_store_from_sources([feature_view.input])
188-
table = offline_store.pull_latest_from_table_or_query(
186+
table = self.offline_store.pull_latest_from_table_or_query(
189187
data_source=feature_view.input,
190188
join_key_columns=join_key_columns,
191189
feature_name_columns=feature_name_columns,
@@ -209,19 +207,16 @@ def materialize_single_feature_view(
209207
feature_view.materialization_intervals.append((start_date, end_date))
210208
registry.apply_feature_view(feature_view, project)
211209

212-
@staticmethod
213210
def get_historical_features(
211+
self,
214212
config: RepoConfig,
215213
feature_views: List[FeatureView],
216214
feature_refs: List[str],
217215
entity_df: Union[pd.DataFrame, str],
218216
registry: Registry,
219217
project: str,
220218
) -> RetrievalJob:
221-
offline_store = get_offline_store_from_sources(
222-
[feature_view.input for feature_view in feature_views]
223-
)
224-
return offline_store.get_historical_features(
219+
return self.offline_store.get_historical_features(
225220
config=config,
226221
feature_views=feature_views,
227222
feature_refs=feature_refs,

‎sdk/python/feast/infra/offline_stores/bigquery.py‎

Lines changed: 7 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,7 @@
1717
_get_requested_feature_views_to_features_dict,
1818
)
1919
from feast.registry import Registry
20-
from feast.repo_config import RepoConfig
20+
from feast.repo_config import BigQueryOfflineStoreConfig, RepoConfig
2121

2222
try:
2323
from google.auth.exceptions import DefaultCredentialsError
@@ -100,8 +100,10 @@ def get_historical_features(
100100
entity_df
101101
)
102102

103+
assert isinstance(config.offline_store, BigQueryOfflineStoreConfig)
104+
103105
table_id = _upload_entity_df_into_bigquery(
104-
config.project, entity_df, client
106+
config.project, config.offline_store.dataset, entity_df, client,
105107
)
106108
entity_df_sql_table = f"`{table_id}`"
107109
else:
@@ -198,11 +200,11 @@ class FeatureViewQueryContext:
198200
entity_selections: List[str]
199201

200202

201-
def _upload_entity_df_into_bigquery(project, entity_df, client) -> str:
203+
def _upload_entity_df_into_bigquery(project, dataset_name, entity_df, client) -> str:
202204
"""Uploads a Pandas entity dataframe into a BigQuery table and returns a reference to the resulting table"""
203205

204206
# First create the BigQuery dataset if it doesn't exist
205-
dataset = bigquery.Dataset(f"{client.project}.feast_{project}")
207+
dataset = bigquery.Dataset(f"{client.project}.{dataset_name}")
206208
dataset.location = "US"
207209
client.create_dataset(
208210
dataset, exists_ok=True
@@ -213,7 +215,7 @@ def _upload_entity_df_into_bigquery(project, entity_df, client) -> str:
213215

214216
# Upload the dataframe into BigQuery, creating a temporary table
215217
job_config = bigquery.LoadJobConfig()
216-
table_id = f"{client.project}.feast_{project}.entity_df_{int(time.time())}"
218+
table_id = f"{client.project}.{dataset_name}.entity_df_{project}_{int(time.time())}"
217219
job = client.load_table_from_dataframe(entity_df, table_id, job_config=job_config,)
218220
job.result()
219221

‎sdk/python/feast/infra/offline_stores/helpers.py‎

Lines changed: 28 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -1,28 +1,41 @@
1-
from typing import List
2-
31
from feast.data_source import BigQuerySource, DataSource, FileSource
2+
from feast.errors import FeastOfflineStoreUnsupportedDataSource
43
from feast.infra.offline_stores.offline_store import OfflineStore
4+
from feast.repo_config import (
5+
BigQueryOfflineStoreConfig,
6+
FileOfflineStoreConfig,
7+
OfflineStoreConfig,
8+
)
59

610

7-
def get_offline_store_from_sources(sources: List[DataSource]) -> OfflineStore:
8-
"""Detect which offline store should be used for retrieving historical features"""
9-
10-
source_types = [type(source) for source in sources]
11+
def get_offline_store_from_config(
12+
offline_store_config: OfflineStoreConfig,
13+
) -> OfflineStore:
14+
"""Get the offline store from offline store config"""
1115

12-
# Retrieve features from ParquetOfflineStore
13-
if all(source == FileSource for source in source_types):
16+
if isinstance(offline_store_config, FileOfflineStoreConfig):
1417
from feast.infra.offline_stores.file import FileOfflineStore
1518

1619
return FileOfflineStore()
17-
18-
# Retrieve features from BigQueryOfflineStore
19-
if all(source == BigQuerySource for source in source_types):
20+
elif isinstance(offline_store_config, BigQueryOfflineStoreConfig):
2021
from feast.infra.offline_stores.bigquery import BigQueryOfflineStore
2122

2223
return BigQueryOfflineStore()
2324

24-
# Could not map inputs to an OfflineStore implementation
25-
raise NotImplementedError(
26-
"Unsupported combination of feature view input source types. Please ensure that all source types are "
27-
"consistent and available in the same offline store."
25+
raise ValueError(f"Unsupported offline store config '{offline_store_config}'")
26+
27+
28+
def assert_offline_store_supports_data_source(
29+
offline_store_config: OfflineStoreConfig, data_source: DataSource
30+
):
31+
if (
32+
isinstance(offline_store_config, FileOfflineStoreConfig)
33+
and isinstance(data_source, FileSource)
34+
) or (
35+
isinstance(offline_store_config, BigQueryOfflineStoreConfig)
36+
and isinstance(data_source, BigQuerySource)
37+
):
38+
return
39+
raise FeastOfflineStoreUnsupportedDataSource(
40+
offline_store_config.type, data_source.__class__.__name__
2841
)

‎sdk/python/feast/infra/provider.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -107,9 +107,9 @@ def materialize_single_feature_view(
107107
) -> None:
108108
pass
109109

110-
@staticmethod
111110
@abc.abstractmethod
112111
def get_historical_features(
112+
self,
113113
config: RepoConfig,
114114
feature_views: List[FeatureView],
115115
feature_refs: List[str],

0 commit comments

Comments
 (0)