Skip to content

Commit 753dee5

Browse files
JwredeJwrede
authored
feat(bigquery): Support DATE-type event timestamp columns (#6362)
* feat(bigquery): Support DATE-type event timestamp columns When the event_timestamp column in BigQuery is a DATE type, the generated SQL wraps comparison values in TIMESTAMP(), causing a type mismatch error. This adds a timestamp_field_type parameter to BigQuerySource that, when set to "DATE", generates DATE() comparisons instead. Closes #2530 (part 2) Signed-off-by: Jonathan Wrede <wrede.jonathan00@gmail.com> * fix(bigquery): Use protobuf 4.25.x compatible generated code The proto files were regenerated with protobuf 6.31.1 / grpcio-tools 1.80.0, which imports runtime_version -- a module that does not exist in protobuf 4.25.x used by the project. Revert generated code to 4.25.1 format while keeping the new timestamp_field_type field. Signed-off-by: Jonathan Wrede <wrede.jonathan00@gmail.com> * fix(bigquery): Add Literal type annotation for cast_style Mypy infers str from the ternary expression; annotate with the exact Literal union so the call to get_timestamp_filter_sql passes type checking. Signed-off-by: Jonathan Wrede <wrede.jonathan00@gmail.com> * fix: Make timestamp_field_type default to None in FeatureViewQueryContext Callers that do not use DATE-typed timestamp fields (e.g. Spark offline store tests) should not be forced to pass timestamp_field_type. Adding a default keeps the new field backward-compatible. Signed-off-by: Jonathan Wrede <wrede.jonathan00@gmail.com> * fix: Keep timestamp_field_type required in FeatureViewQueryContext A default value on timestamp_field_type breaks the SparkFeatureViewQueryContext subclass because its non-default fields (min_date_partition, max_date_partition) would follow a field with a default. Instead, keep it required and update the Spark test to pass it. Signed-off-by: Jonathan Wrede <wrede.jonathan00@gmail.com> * fix: regenerate protos matching upstream mypy-protobuf style Reset all non-DataSource generated files to match master. Only DataSource_pb2.py and DataSource_pb2.pyi contain our timestamp_field_type additions (field 28). The .pyi stub is hand-edited to match the existing import style used on master. Signed-off-by: Jonathan Wrede <wrede.jonathan00@gmail.com> --------- Signed-off-by: Jonathan Wrede <wrede.jonathan00@gmail.com>
1 parent d86b13d commit 753dee5

9 files changed

Lines changed: 208 additions & 45 deletions

File tree

‎protos/feast/core/DataSource.proto‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -29,7 +29,7 @@ import "feast/types/Value.proto";
2929
import "feast/core/Feature.proto";
3030

3131
// Defines a Data Source that can be used source Feature data
32-
// Next available id: 28
32+
// Next available id: 29
3333
message DataSource {
3434
// Field indexes should *not* be reused. Not sure if fields 6-10 were used previously or not,
3535
// but they are going to be reserved for backwards compatibility.
@@ -81,6 +81,10 @@ message DataSource {
8181
// Must specify creation timestamp column name
8282
string created_timestamp_column = 5;
8383

84+
// (Optional) Type of the timestamp_field column ("TIMESTAMP" or "DATE").
85+
// When set to "DATE", SQL generation uses date-only comparisons.
86+
string timestamp_field_type = 28;
87+
8488
// This is an internal field that is represents the python class for the data source object a proto object represents.
8589
// This should be set by feast, and not by users.
8690
// The field is used primarily by custom data sources and is mandatory for them to set. Feast may set it for

‎sdk/python/feast/data_source.py‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -205,6 +205,7 @@ class DataSource(ABC):
205205
tags: Dict[str, str]
206206
owner: str
207207
date_partition_column: str
208+
timestamp_field_type: str
208209
created_timestamp: Optional[datetime]
209210
last_updated_timestamp: Optional[datetime]
210211

@@ -219,6 +220,7 @@ def __init__(
219220
tags: Optional[Dict[str, str]] = None,
220221
owner: Optional[str] = "",
221222
date_partition_column: Optional[str] = None,
223+
timestamp_field_type: Optional[str] = None,
222224
):
223225
"""
224226
Creates a DataSource object.
@@ -237,6 +239,9 @@ def __init__(
237239
owner (optional): The owner of the data source, typically the email of the primary
238240
maintainer.
239241
date_partition_column (optional): Timestamp column used for partitioning. Not supported by all stores
242+
timestamp_field_type (optional): Type of the timestamp_field column.
243+
Defaults to "TIMESTAMP". Set to "DATE" when the event timestamp column
244+
is a DATE type, so SQL generation uses date-only comparisons.
240245
"""
241246
self.name = name
242247
self.timestamp_field = timestamp_field or ""
@@ -257,6 +262,7 @@ def __init__(
257262
self.date_partition_column = (
258263
date_partition_column if date_partition_column else ""
259264
)
265+
self.timestamp_field_type = timestamp_field_type if timestamp_field_type else ""
260266
now = _utc_now()
261267
self.created_timestamp = now
262268
self.last_updated_timestamp = now
@@ -280,6 +286,7 @@ def __eq__(self, other):
280286
or self.created_timestamp_column != other.created_timestamp_column
281287
or self.field_mapping != other.field_mapping
282288
or self.date_partition_column != other.date_partition_column
289+
or self.timestamp_field_type != other.timestamp_field_type
283290
or self.description != other.description
284291
or self.tags != other.tags
285292
or self.owner != other.owner

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

Lines changed: 19 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -161,13 +161,18 @@ def pull_latest_from_table_or_query(
161161
project=project_id,
162162
location=config.offline_store.location,
163163
)
164+
cast_style: Literal["date_func", "timestamp_func"] = (
165+
"date_func"
166+
if data_source.timestamp_field_type == "DATE"
167+
else "timestamp_func"
168+
)
164169
timestamp_filter = get_timestamp_filter_sql(
165170
start_date,
166171
end_date,
167172
timestamp_field,
168173
date_partition_column=data_source.date_partition_column,
169174
quote_fields=False,
170-
cast_style="timestamp_func",
175+
cast_style=cast_style,
171176
)
172177
query = f"""
173178
SELECT
@@ -220,13 +225,18 @@ def pull_all_from_table_or_query(
220225
+ BigQueryOfflineStore._escape_query_columns(feature_name_columns)
221226
+ timestamp_fields
222227
)
228+
cast_style: Literal["date_func", "timestamp_func"] = (
229+
"date_func"
230+
if data_source.timestamp_field_type == "DATE"
231+
else "timestamp_func"
232+
)
223233
timestamp_filter = get_timestamp_filter_sql(
224234
start_date,
225235
end_date,
226236
timestamp_field,
227237
date_partition_column=data_source.date_partition_column,
228238
quote_fields=False,
229-
cast_style="timestamp_func",
239+
cast_style=cast_style,
230240
)
231241
query = f"""
232242
SELECT {field_string}
@@ -938,10 +948,17 @@ def arrow_schema_to_bq_schema(arrow_schema: pyarrow.Schema) -> List[SchemaField]
938948
{% if loop.last %}{% else %}, {% endif %}
939949
{% endfor %}
940950
FROM {{ featureview.table_subquery }}
951+
{% if featureview.timestamp_field_type == "DATE" %}
952+
WHERE {{ featureview.timestamp_field }} <= DATE('{{ featureview.max_event_timestamp[:10] }}')
953+
{% if featureview.ttl == 0 %}{% else %}
954+
AND {{ featureview.timestamp_field }} >= DATE('{{ featureview.min_event_timestamp[:10] }}')
955+
{% endif %}
956+
{% else %}
941957
WHERE {{ featureview.timestamp_field }} <= '{{ featureview.max_event_timestamp }}'
942958
{% if featureview.ttl == 0 %}{% else %}
943959
AND {{ featureview.timestamp_field }} >= '{{ featureview.min_event_timestamp }}'
944960
{% endif %}
961+
{% endif %}
945962
{% if featureview.date_partition_column %}
946963
AND {{ featureview.date_partition_column | backticks }} <= '{{ featureview.max_event_timestamp[:10] }}'
947964
{% if featureview.min_event_timestamp %}

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

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -35,6 +35,7 @@ def __init__(
3535
created_timestamp_column: Optional[str] = "",
3636
field_mapping: Optional[Dict[str, str]] = None,
3737
date_partition_column: Optional[str] = None,
38+
timestamp_field_type: Optional[str] = None,
3839
query: Optional[str] = None,
3940
description: Optional[str] = "",
4041
tags: Optional[Dict[str, str]] = None,
@@ -54,6 +55,9 @@ def __init__(
5455
field_mapping (optional): A dictionary mapping of column names in this data source to feature names in a feature table
5556
or view. Only used for feature columns, not entities or timestamp columns.
5657
date_partition_column (optional): Timestamp column used for partitioning.
58+
timestamp_field_type (optional): Type of the timestamp_field column.
59+
Set to "DATE" when the event timestamp column is a DATE type,
60+
so SQL generation uses date-only comparisons instead of TIMESTAMP().
5761
query (optional): The query to be executed to obtain the features. When both 'table'
5862
and 'query' are provided, 'query' takes priority for reads.
5963
description (optional): A human-readable description.
@@ -81,6 +85,7 @@ def __init__(
8185
created_timestamp_column=created_timestamp_column,
8286
field_mapping=field_mapping,
8387
date_partition_column=date_partition_column,
88+
timestamp_field_type=timestamp_field_type,
8489
description=description,
8590
tags=tags,
8691
owner=owner,
@@ -121,6 +126,7 @@ def from_proto(data_source: DataSourceProto):
121126
timestamp_field=data_source.timestamp_field,
122127
created_timestamp_column=data_source.created_timestamp_column,
123128
date_partition_column=data_source.date_partition_column,
129+
timestamp_field_type=data_source.timestamp_field_type or None,
124130
query=data_source.bigquery_options.query,
125131
description=data_source.description,
126132
tags=dict(data_source.tags),
@@ -139,6 +145,7 @@ def _to_proto_impl(self) -> DataSourceProto:
139145
timestamp_field=self.timestamp_field,
140146
created_timestamp_column=self.created_timestamp_column,
141147
date_partition_column=self.date_partition_column,
148+
timestamp_field_type=self.timestamp_field_type,
142149
)
143150

144151
return data_source_proto

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

Lines changed: 12 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -98,6 +98,7 @@ class FeatureViewQueryContext:
9898
date_partition_column: Optional[
9999
str
100100
] # this attribute is added because partition pruning affects Athena's query performance.
101+
timestamp_field_type: Optional[str]
101102

102103

103104
def get_feature_view_query_context(
@@ -160,6 +161,10 @@ def get_feature_view_query_context(
160161
feature_view.batch_source.date_partition_column,
161162
)
162163

164+
timestamp_field_type = getattr(
165+
feature_view.batch_source, "timestamp_field_type", ""
166+
)
167+
163168
max_event_timestamp = to_naive_utc(entity_df_timestamp_range[1]).isoformat()
164169
min_event_timestamp = None
165170
if feature_view.ttl:
@@ -181,6 +186,7 @@ def get_feature_view_query_context(
181186
min_event_timestamp=min_event_timestamp,
182187
max_event_timestamp=max_event_timestamp,
183188
date_partition_column=date_partition_column,
189+
timestamp_field_type=timestamp_field_type or None,
184190
)
185191
query_context.append(context)
186192

@@ -340,7 +346,7 @@ def get_timestamp_filter_sql(
340346
date_partition_column: Optional[str] = None,
341347
tz: Optional[timezone] = None,
342348
cast_style: Literal[
343-
"timestamp", "timestamp_func", "timestamptz", "raw"
349+
"timestamp", "timestamp_func", "timestamptz", "raw", "date_func"
344350
] = "timestamp",
345351
date_time_separator: str = "T",
346352
quote_fields: bool = True,
@@ -355,10 +361,11 @@ def get_timestamp_filter_sql(
355361
date_partition_column: optional partition column (for pruning)
356362
tz: optional timezone for datetime inputs
357363
cast_style: one of:
358-
- "timestamp": TIMESTAMP '...' → Common Sql engine Snowflake, Redshift etc.
364+
- "timestamp": TIMESTAMP '...' → Common Sql engine Snowflake, Redshift etc.
359365
- "timestamp_func": TIMESTAMP('...') → BigQuery, Couchbase etc.
360366
- "timestamptz": '...'::timestamptz → PostgreSQL
361367
- "raw": '...' → no cast, string only
368+
- "date_func": DATE('...') → BigQuery DATE columns
362369
date_time_separator: separator for datetime strings (default is "T")
363370
(e.g. "2023-10-01T00:00:00" or "2023-10-01 00:00:00")
364371
quote_fields: whether to quote the timestamp and partition column names
@@ -384,6 +391,9 @@ def format_casted_ts(val: Union[str, datetime]) -> str:
384391
return f"TIMESTAMP '{val_str}'"
385392
elif cast_style == "timestamp_func":
386393
return f"TIMESTAMP('{val_str}')"
394+
elif cast_style == "date_func":
395+
date_str = val_str[:10] if len(val_str) >= 10 else val_str
396+
return f"DATE('{date_str}')"
387397
elif cast_style == "timestamptz":
388398
return f"'{val_str}'::{cast_style}"
389399
else:

0 commit comments

Comments
 (0)