From ad65e0bf8217e5241ff2e52b570ec52159284711 Mon Sep 17 00:00:00 2001 From: NIKHIL Date: Wed, 4 Nov 2026 03:09:46 +0530 Subject: [PATCH] fix: Materialize zone-aware nanosecond timestamps on DuckDB 2 The ibis pull functions cast the timestamp column to a zone-less timestamp for the window filter. DuckDB 2 refuses that cast for TIMESTAMPTZ_NS, which pandas writes by default, and ibis reports the type as unknown. Keep the zone when casting. Signed-off-by: NIKHIL --- sdk/python/feast/infra/offline_stores/ibis.py | 21 +++++-- .../unit/infra/offline_stores/test_ibis.py | 56 ++++++++++++++++++- 2 files changed, 71 insertions(+), 6 deletions(-) diff --git a/sdk/python/feast/infra/offline_stores/ibis.py b/sdk/python/feast/infra/offline_stores/ibis.py index ccd694fb0fc..5f3bd686977 100644 --- a/sdk/python/feast/infra/offline_stores/ibis.py +++ b/sdk/python/feast/infra/offline_stores/ibis.py @@ -93,8 +93,8 @@ def pull_latest_from_table_or_query_ibis( table = table.filter( ibis.and_( - table[timestamp_field].cast("timestamp") >= ibis.literal(start_date), - table[timestamp_field].cast("timestamp") <= ibis.literal(end_date), + _comparable_timestamp(table[timestamp_field]) >= ibis.literal(start_date), + _comparable_timestamp(table[timestamp_field]) <= ibis.literal(end_date), ) ) @@ -117,6 +117,19 @@ def pull_latest_from_table_or_query_ibis( ) +def _comparable_timestamp(column): + """Cast a timestamp column for comparison with a window bound, keeping any zone. + + DuckDB 2 refuses to cast TIMESTAMPTZ_NS, which is what pandas writes by default, to a + TIMESTAMP without a zone, and ibis reports that column's type as unknown. + """ + column_type = column.type() + if column_type.is_unknown(): + return column.cast(dt.Timestamp(timezone="UTC")) + timezone = getattr(column_type, "timezone", None) + return column.cast(dt.Timestamp(timezone=timezone) if timezone else "timestamp") + + def _get_entity_df_event_timestamp_range( entity_df: pd.DataFrame, entity_df_event_timestamp_col: str ) -> Tuple[datetime, datetime]: @@ -345,10 +358,10 @@ def pull_all_from_table_or_query_ibis( table = table.filter( ibis.and_( - table[timestamp_field].cast("timestamp") >= ibis.literal(start_date) + _comparable_timestamp(table[timestamp_field]) >= ibis.literal(start_date) if start_date else ibis.literal(True), - table[timestamp_field].cast("timestamp") <= ibis.literal(end_date) + _comparable_timestamp(table[timestamp_field]) <= ibis.literal(end_date) if end_date else ibis.literal(True), ) diff --git a/sdk/python/tests/unit/infra/offline_stores/test_ibis.py b/sdk/python/tests/unit/infra/offline_stores/test_ibis.py index e74ecb887d1..20c10cd19e8 100644 --- a/sdk/python/tests/unit/infra/offline_stores/test_ibis.py +++ b/sdk/python/tests/unit/infra/offline_stores/test_ibis.py @@ -1,11 +1,18 @@ -from datetime import datetime, timedelta +from datetime import datetime, timedelta, timezone from typing import Dict, List, Tuple +from unittest.mock import MagicMock import ibis +import pandas as pd import pyarrow as pa import pyarrow.compute as pc +import pytest -from feast.infra.offline_stores.ibis import point_in_time_join +from feast.infra.offline_stores.ibis import ( + point_in_time_join, + pull_all_from_table_or_query_ibis, + pull_latest_from_table_or_query_ibis, +) def pa_datetime(year, month, day): @@ -226,3 +233,48 @@ def test_point_in_time_join_keeps_entity_rows_apart(): for row in actual.to_pylist() } assert rows == {"1": (1.0, 110.0), "11": (11.0, 10.0), "7": (7.0, None)} + + +@pytest.mark.parametrize( + "dtype", + ["datetime64[ns, UTC]", "datetime64[us, UTC]", "datetime64[ns]", "datetime64[us]"], +) +@pytest.mark.parametrize( + "pull", [pull_latest_from_table_or_query_ibis, pull_all_from_table_or_query_ibis] +) +def test_pull_filters_parquet_timestamps_of_any_precision(tmp_path, dtype, pull): + """Materialize pulls read what pandas writes: nanosecond and zone-aware included.""" + start = datetime(2026, 1, 1, tzinfo=timezone.utc) + path = tmp_path / "source.parquet" + pd.DataFrame( + { + "driver_id": [1, 2, 3, 4], + "event_timestamp": [start + timedelta(hours=h) for h in range(4)], + "trips": [10, 20, 30, 40], + } + ).astype( + {"event_timestamp": dtype if "UTC" in dtype else "datetime64[ns, UTC]"} + ).pipe( + lambda df: ( + df + if "UTC" in dtype + else df.assign( + event_timestamp=df.event_timestamp.dt.tz_localize(None).astype(dtype) + ) + ) + ).to_parquet(path) + + table = pull( + config=MagicMock(), + data_source=MagicMock(), + join_key_columns=["driver_id"], + feature_name_columns=["trips"], + timestamp_field="event_timestamp", + created_timestamp_column=None, + start_date=start + timedelta(hours=1), + end_date=start + timedelta(hours=2), + data_source_reader=lambda *_: ibis.read_parquet(path), + data_source_writer=MagicMock(), + ).to_arrow() + + assert sorted(table["driver_id"].to_pylist()) == [2, 3]