Skip to content

Commit a55497b

Browse files
achalsTsotne Tabidze
achals
authored and
Tsotne Tabidze
committed
Add entity column validations when getting historical features from bigquery (#1614)
* Add entity column validations when getting historical features from bigquery Signed-off-by: Achal Shah <achals@gmail.com> * make format Signed-off-by: Achal Shah <achals@gmail.com> * Remove wrong file Signed-off-by: Achal Shah <achals@gmail.com> * Add tests Signed-off-by: Achal Shah <achals@gmail.com>
1 parent 1d392bd commit a55497b

3 files changed

Lines changed: 102 additions & 2 deletions

File tree

‎sdk/python/feast/errors.py‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -67,3 +67,11 @@ def __init__(self, offline_store_name: str, data_source_name: str):
6767
super().__init__(
6868
f"Offline Store '{offline_store_name}' does not support data source '{data_source_name}'"
6969
)
70+
71+
72+
class FeastEntityDFMissingColumnsError(Exception):
73+
def __init__(self, expected, missing):
74+
super().__init__(
75+
f"The entity dataframe you have provided must contain columns {expected}, "
76+
f"but {missing} were missing."
77+
)

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

Lines changed: 52 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,12 +1,13 @@
11
import time
22
from dataclasses import asdict, dataclass
33
from datetime import datetime, timedelta
4-
from typing import List, Optional, Union
4+
from typing import List, Optional, Set, Union
55

66
import pandas
77
import pyarrow
88
from jinja2 import BaseLoader, Environment
99

10+
from feast import errors
1011
from feast.data_source import BigQuerySource, DataSource
1112
from feast.errors import FeastProviderLoginError
1213
from feast.feature_view import FeatureView
@@ -87,13 +88,18 @@ def get_historical_features(
8788

8889
client = _get_bigquery_client()
8990

91+
expected_join_keys = _get_join_keys(project, feature_views, registry)
92+
9093
if type(entity_df) is str:
9194
entity_df_job = client.query(entity_df)
9295
entity_df_result = entity_df_job.result() # also starts job
9396

9497
entity_df_event_timestamp_col = _infer_event_timestamp_from_bigquery_query(
9598
entity_df_result
9699
)
100+
_assert_expected_columns_in_bigquery(
101+
expected_join_keys, entity_df_event_timestamp_col, entity_df_result
102+
)
97103

98104
entity_df_sql_table = f"`{entity_df_job.destination.project}.{entity_df_job.destination.dataset_id}.{entity_df_job.destination.table_id}`"
99105
elif isinstance(entity_df, pandas.DataFrame):
@@ -103,6 +109,10 @@ def get_historical_features(
103109

104110
assert isinstance(config.offline_store, BigQueryOfflineStoreConfig)
105111

112+
_assert_expected_columns_in_dataframe(
113+
expected_join_keys, entity_df_event_timestamp_col, entity_df
114+
)
115+
106116
table_id = _upload_entity_df_into_bigquery(
107117
config.project, config.offline_store.dataset, entity_df, client
108118
)
@@ -132,6 +142,47 @@ def get_historical_features(
132142
return job
133143

134144

145+
def _assert_expected_columns_in_dataframe(
146+
join_keys: Set[str], entity_df_event_timestamp_col: str, entity_df: pandas.DataFrame
147+
):
148+
entity_df_columns = set(entity_df.columns.values)
149+
expected_columns = join_keys.copy()
150+
expected_columns.add(entity_df_event_timestamp_col)
151+
152+
missing_keys = expected_columns - entity_df_columns
153+
154+
if len(missing_keys) != 0:
155+
raise errors.FeastEntityDFMissingColumnsError(expected_columns, missing_keys)
156+
157+
158+
def _assert_expected_columns_in_bigquery(
159+
join_keys: Set[str], entity_df_event_timestamp_col: str, entity_df_result
160+
):
161+
entity_columns = set()
162+
for schema_field in entity_df_result.schema:
163+
entity_columns.add(schema_field.name)
164+
165+
expected_columns = join_keys.copy()
166+
expected_columns.add(entity_df_event_timestamp_col)
167+
168+
missing_keys = expected_columns - entity_columns
169+
170+
if len(missing_keys) != 0:
171+
raise errors.FeastEntityDFMissingColumnsError(expected_columns, missing_keys)
172+
173+
174+
def _get_join_keys(
175+
project: str, feature_views: List[FeatureView], registry: Registry
176+
) -> Set[str]:
177+
join_keys = set()
178+
for feature_view in feature_views:
179+
entities = feature_view.entities
180+
for entity_name in entities:
181+
entity = registry.get_entity(entity_name, project)
182+
join_keys.add(entity.join_key)
183+
return join_keys
184+
185+
135186
def _infer_event_timestamp_from_bigquery_query(entity_df_result) -> str:
136187
if any(
137188
schema_field.name == DEFAULT_ENTITY_DF_EVENT_TIMESTAMP_COL

‎sdk/python/tests/test_historical_retrieval.py‎

Lines changed: 42 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,7 @@
1414
from pytz import utc
1515

1616
import feast.driver_test_data as driver_data
17-
from feast import utils
17+
from feast import errors, utils
1818
from feast.data_source import BigQuerySource, FileSource
1919
from feast.entity import Entity
2020
from feast.feature import Feature
@@ -450,6 +450,30 @@ def test_historical_features_from_bigquery_sources(
450450
check_dtype=False,
451451
)
452452

453+
timestamp_column = (
454+
"e_ts"
455+
if infer_event_timestamp_col
456+
else DEFAULT_ENTITY_DF_EVENT_TIMESTAMP_COL
457+
)
458+
459+
entity_df_query_with_invalid_join_key = (
460+
f"select order_id, driver_id, customer_id as customer, "
461+
f"order_is_success, {timestamp_column}, FROM {gcp_project}.{table_id}"
462+
)
463+
# Rename the join key; this should now raise an error.
464+
assertpy.assert_that(store.get_historical_features).raises(
465+
errors.FeastEntityDFMissingColumnsError
466+
).when_called_with(
467+
entity_df=entity_df_query_with_invalid_join_key,
468+
feature_refs=[
469+
"driver_stats:conv_rate",
470+
"driver_stats:avg_daily_trips",
471+
"customer_profile:current_balance",
472+
"customer_profile:avg_passenger_count",
473+
"customer_profile:lifetime_trip_count",
474+
],
475+
)
476+
453477
job_from_df = store.get_historical_features(
454478
entity_df=orders_df,
455479
feature_refs=[
@@ -461,6 +485,23 @@ def test_historical_features_from_bigquery_sources(
461485
],
462486
)
463487

488+
# Rename the join key; this should now raise an error.
489+
orders_df_with_invalid_join_key = orders_df.rename(
490+
{"customer_id": "customer"}, axis="columns"
491+
)
492+
assertpy.assert_that(store.get_historical_features).raises(
493+
errors.FeastEntityDFMissingColumnsError
494+
).when_called_with(
495+
entity_df=orders_df_with_invalid_join_key,
496+
feature_refs=[
497+
"driver_stats:conv_rate",
498+
"driver_stats:avg_daily_trips",
499+
"customer_profile:current_balance",
500+
"customer_profile:avg_passenger_count",
501+
"customer_profile:lifetime_trip_count",
502+
],
503+
)
504+
464505
# Make sure that custom dataset name is being used from the offline_store config
465506
if provider_type == "gcp_custom_offline_config":
466507
assertpy.assert_that(job_from_df.query).contains("foo.entity_df")

0 commit comments

Comments
 (0)