Skip to content

fix: read all source columns in BigQuery pulls when the feature list is empty #7048

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
Mrun25 wants to merge 1 commit into
base: master
Choose a base branch
from
Open

fix: read all source columns in BigQuery pulls when the feature list is empty #7048

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
36 changes: 25 additions & 11 deletions sdk/python/feast/infra/offline_stores/bigquery.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -166,11 +166,19 @@ def pull_latest_from_table_or_query(
if created_timestamp_column:
timestamps.append(created_timestamp_column)
timestamp_desc_string = " DESC, ".join(timestamps) + " DESC"
field_string = ", ".join(
BigQueryOfflineStore._escape_query_columns(join_key_columns)
+ BigQueryOfflineStore._escape_query_columns(feature_name_columns)
+ timestamps
)
# An empty feature_name_columns signals "read all source columns"
# (BatchFeatureView with a pandas/python UDF); see
# FeatureBuilder.get_column_info.
if not feature_name_columns:
inner_field_string = "*"
# The outer query must not expose the ROW_NUMBER helper column.
field_string = "* EXCEPT(_feast_row)"
else:
inner_field_string = field_string = ", ".join(
BigQueryOfflineStore._escape_query_columns(join_key_columns)
+ BigQueryOfflineStore._escape_query_columns(feature_name_columns)
+ timestamps
)
project_id = (
config.offline_store.billing_project_id or config.offline_store.project_id
)
Expand All @@ -196,7 +204,7 @@ def pull_latest_from_table_or_query(
{field_string}
{f", {repr(DUMMY_ENTITY_VAL)} AS {DUMMY_ENTITY_ID}" if not join_key_columns else ""}
FROM (
SELECT {field_string},
SELECT {inner_field_string},
ROW_NUMBER() OVER({partition_by_join_key_string} ORDER BY {timestamp_desc_string}) AS _feast_row
FROM {from_expression}
WHERE {timestamp_filter}
Expand Down Expand Up @@ -237,11 +245,17 @@ def pull_all_from_table_or_query(
timestamp_fields = [timestamp_field]
if created_timestamp_column:
timestamp_fields.append(created_timestamp_column)
field_string = ", ".join(
BigQueryOfflineStore._escape_query_columns(join_key_columns)
+ BigQueryOfflineStore._escape_query_columns(feature_name_columns)
+ timestamp_fields
)
# An empty feature_name_columns signals "read all source columns"
# (BatchFeatureView with a pandas/python UDF); see
# FeatureBuilder.get_column_info.
if not feature_name_columns:
field_string = "*"
else:
field_string = ", ".join(
BigQueryOfflineStore._escape_query_columns(join_key_columns)
+ BigQueryOfflineStore._escape_query_columns(feature_name_columns)
+ timestamp_fields
)
cast_style: Literal["date_func", "timestamp_func"] = (
"date_func"
if data_source.timestamp_field_type == "DATE"
Expand Down
63 changes: 63 additions & 0 deletions sdk/python/tests/unit/infra/offline_stores/test_bigquery.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -154,6 +154,69 @@ def test_pull_all_from_table_or_query_partition_pruning(mock_get_bigquery_client
assert "partition_date <= '2021-01-02'" in actual_query


@patch("feast.infra.offline_stores.bigquery._get_bigquery_client")
def test_pull_all_empty_feature_columns_selects_all(mock_get_bigquery_client):
"""BatchFeatureView pandas/python UDF mode passes an empty feature list to
signal "read all source columns" (FeatureBuilder.get_column_info)."""
mock_get_bigquery_client.return_value = Mock()
test_repo_config = RepoConfig(
registry="gs://ml-test/repo/registry.db",
project="test",
provider="gcp",
online_store=SqliteOnlineStoreConfig(type="sqlite"),
offline_store=BigQueryOfflineStoreConfig(type="bigquery", dataset="feast"),
)
test_data_source = BigQuerySource(
table="project:dataset.table",
timestamp_field="event_timestamp",
)
retrieval_job = BigQueryOfflineStore.pull_all_from_table_or_query(
config=test_repo_config,
data_source=test_data_source,
join_key_columns=["driver_id"],
feature_name_columns=[],
timestamp_field="event_timestamp",
start_date=datetime(2021, 1, 1, tzinfo=timezone.utc),
end_date=datetime(2021, 1, 2, tzinfo=timezone.utc),
)
actual_query = retrieval_job.to_sql()
normalized = " ".join(actual_query.split())
assert "SELECT * FROM" in normalized
assert "`driver_id`" not in normalized.split("FROM")[0]


@patch("feast.infra.offline_stores.bigquery._get_bigquery_client")
def test_pull_latest_empty_feature_columns_selects_all(mock_get_bigquery_client):
"""BatchFeatureView pandas/python UDF mode passes an empty feature list to
signal "read all source columns" (FeatureBuilder.get_column_info)."""
mock_get_bigquery_client.return_value = Mock()
test_repo_config = RepoConfig(
registry="gs://ml-test/repo/registry.db",
project="test",
provider="gcp",
online_store=SqliteOnlineStoreConfig(type="sqlite"),
offline_store=BigQueryOfflineStoreConfig(type="bigquery", dataset="feast"),
)
test_data_source = BigQuerySource(
table="project:dataset.table",
timestamp_field="event_timestamp",
)
retrieval_job = BigQueryOfflineStore.pull_latest_from_table_or_query(
config=test_repo_config,
data_source=test_data_source,
join_key_columns=["driver_id"],
feature_name_columns=[],
timestamp_field="event_timestamp",
created_timestamp_column=None,
start_date=datetime(2021, 1, 1, tzinfo=timezone.utc),
end_date=datetime(2021, 1, 2, tzinfo=timezone.utc),
)
actual_query = retrieval_job.to_sql()
normalized = " ".join(actual_query.split())
assert "SELECT * EXCEPT(_feast_row)" in normalized
assert "SELECT *, ROW_NUMBER()" in normalized


@patch("feast.infra.offline_stores.bigquery._get_bigquery_client")
def test_pull_latest_date_type_timestamp_field(mock_get_bigquery_client):
mock_get_bigquery_client.return_value = Mock()
Expand Down