Skip to content

Commit 77ed779

Browse files
JwredeJwrede
authored
fix(bigquery): Prefer query over table in get_table_query_string (#6360)
* fix(bigquery): prefer query over table in get_table_query_string When both `table` and `query` are set on a BigQuerySource, `get_table_query_string()` now returns the query (wrapped in parens) instead of the table reference. This allows PushSource users to provide a custom read query (e.g. for deduplication) while keeping `table` for offline writes via `offline_write_batch()`. Also applies the same priority inversion to `get_table_column_names_and_types()` so schema inference matches the actual read path. Closes #6200 Signed-off-by: Jonathan Wrede <wrede.jonathan00@gmail.com> * docs(bigquery): update docstring to reflect both table and query support The validation only enforces at least one of table/query, not exactly one. Update the docstring to document the supported behavior when both are set. Signed-off-by: Jonathan Wrede <wrede.jonathan00@gmail.com> --------- Signed-off-by: Jonathan Wrede <wrede.jonathan00@gmail.com>
1 parent 153a279 commit 77ed779

2 files changed

Lines changed: 61 additions & 13 deletions

File tree

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

Lines changed: 13 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -47,15 +47,15 @@ def __init__(
4747
case the table must be specified.
4848
timestamp_field (optional): Event timestamp field used for point in time
4949
joins of feature values.
50-
table (optional): BigQuery table where the features are stored. Exactly one of 'table'
51-
and 'query' must be specified.
52-
table (optional): The BigQuery table where features can be found.
50+
table (optional): BigQuery table where the features are stored. At least one of 'table'
51+
and 'query' must be specified. When both are set, 'query' is used for reads and
52+
'table' is used as the write destination.
5353
created_timestamp_column (optional): Timestamp column when row was created, used for deduplicating rows.
5454
field_mapping (optional): A dictionary mapping of column names in this data source to feature names in a feature table
5555
or view. Only used for feature columns, not entities or timestamp columns.
5656
date_partition_column (optional): Timestamp column used for partitioning.
57-
query (optional): The query to be executed to obtain the features. Exactly one of 'table'
58-
and 'query' must be specified.
57+
query (optional): The query to be executed to obtain the features. When both 'table'
58+
and 'query' are provided, 'query' takes priority for reads.
5959
description (optional): A human-readable description.
6060
tags (optional): A dictionary of key-value pairs to store arbitrary metadata.
6161
owner (optional): The owner of the bigquery source, typically the email of the primary
@@ -156,10 +156,10 @@ def validate(self, config: RepoConfig):
156156

157157
def get_table_query_string(self) -> str:
158158
"""Returns a string that can directly be used to reference this table in SQL"""
159-
if self.table:
160-
return f"`{self.table}`"
161-
else:
159+
if self.query:
162160
return f"({self.query})"
161+
else:
162+
return f"`{self.table}`"
163163

164164
@staticmethod
165165
def source_datatype_to_feast_value_type() -> Callable[[str], ValueType]:
@@ -185,14 +185,14 @@ def get_table_column_names_and_types(
185185
location=config.offline_store.location,
186186
client_info=http_client_info.ClientInfo(user_agent=get_user_agent()),
187187
)
188-
if self.table:
189-
schema = client.get_table(self.table).schema
190-
if not isinstance(schema[0], bigquery.schema.SchemaField):
191-
raise TypeError("Could not parse BigQuery table schema.")
192-
else:
188+
if self.query:
193189
bq_columns_query = f"SELECT * FROM ({self.query}) LIMIT 0"
194190
query_res = client.query(bq_columns_query).result()
195191
schema = query_res.schema
192+
else:
193+
schema = client.get_table(self.table).schema
194+
if not isinstance(schema[0], bigquery.schema.SchemaField):
195+
raise TypeError("Could not parse BigQuery table schema.")
196196

197197
name_type_pairs: List[Tuple[str, str]] = []
198198
for field in schema:

‎sdk/python/tests/unit/infra/offline_stores/test_bigquery.py‎

Lines changed: 48 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -152,3 +152,51 @@ def test_pull_all_from_table_or_query_partition_pruning(mock_get_bigquery_client
152152
)
153153
assert "partition_date >= '2021-01-01'" in actual_query
154154
assert "partition_date <= '2021-01-02'" in actual_query
155+
156+
157+
class TestBigQuerySourceGetTableQueryString:
158+
def test_table_only(self):
159+
source = BigQuerySource(
160+
name="test",
161+
table="project.dataset.table",
162+
timestamp_field="ts",
163+
)
164+
assert source.get_table_query_string() == "`project.dataset.table`"
165+
166+
def test_query_only(self):
167+
source = BigQuerySource(
168+
name="test",
169+
query="SELECT * FROM `project.dataset.table` WHERE active = TRUE",
170+
timestamp_field="ts",
171+
)
172+
assert (
173+
source.get_table_query_string()
174+
== "(SELECT * FROM `project.dataset.table` WHERE active = TRUE)"
175+
)
176+
177+
def test_both_table_and_query_prefers_query(self):
178+
"""When both table and query are set, query takes priority for reads."""
179+
query = (
180+
"SELECT * FROM `project.dataset.table`"
181+
" QUALIFY ROW_NUMBER() OVER (PARTITION BY entity_id, event_time) = 1"
182+
)
183+
source = BigQuerySource(
184+
name="test",
185+
table="project.dataset.table",
186+
query=query,
187+
timestamp_field="ts",
188+
)
189+
result = source.get_table_query_string()
190+
assert result.startswith("(")
191+
assert "QUALIFY" in result
192+
assert result != "`project.dataset.table`"
193+
194+
def test_table_property_unaffected_by_query_priority(self):
195+
"""The .table property is still accessible for write paths."""
196+
source = BigQuerySource(
197+
name="test",
198+
table="project.dataset.write_target",
199+
query="SELECT * FROM `project.dataset.write_target` WHERE deduped",
200+
timestamp_field="ts",
201+
)
202+
assert source.table == "project.dataset.write_target"

0 commit comments

Comments
 (0)