From 6127039cbbb35dd16cfeee3ceffb0bab2c49a737 Mon Sep 17 00:00:00 2001 From: breken-ai <312387581+breken-ai@users.noreply.github.com> Date: Fri, 25 Sep 2026 18:52:07 -0700 Subject: [PATCH] Respect time unit and time zone in polars datetime pushdown The polars lazy-frame filter pushdown treated every Datetime literal as microseconds and emitted a naive TIMESTAMP. For TIMESTAMP_MS / TIMESTAMP_S columns (polars Datetime('ms')) the literal landed in January 1970, so filters like pl.col('t') > datetime(...) returned the wrong rows. For TIMESTAMPTZ columns the UTC instant was cast to a naive TIMESTAMP and then re-read in the session TimeZone, shifting it by the UTC offset. Convert milliseconds and microseconds exactly, emit TIMESTAMPTZ when the literal carries a time zone, and let polars apply nanosecond filters. --- duckdb/polars_io.py | 25 ++++++++++++++++++++----- tests/fast/arrow/test_polars.py | 31 +++++++++++++++++++++++++++++++ 2 files changed, 51 insertions(+), 5 deletions(-) diff --git a/duckdb/polars_io.py b/duckdb/polars_io.py index c5cad457..c2d017a9 100644 --- a/duckdb/polars_io.py +++ b/duckdb/polars_io.py @@ -290,12 +290,27 @@ def _pl_tree_to_sql(tree: _ExpressionTree) -> str: ) return str(Decimal(decimal_value[0]) / Decimal(10 ** decimal_value[-1])) - # Datetime with microseconds since epoch + # Datetime: [value since epoch, time unit, time zone] if dtype.startswith("{'Datetime'") or dtype == "Datetime": - micros = value["Datetime"] - assert isinstance(micros, list), f"A {dtype} should be a one member list but got {type(micros)}" - dt_timestamp = datetime.datetime.fromtimestamp(micros[0] / 1_000_000, tz=datetime.timezone.utc) - return f"'{dt_timestamp!s}'::TIMESTAMP" + datetime_value = value["Datetime"] + assert isinstance(datetime_value, list), f"A {dtype} should be a list but got {type(datetime_value)}" + epoch_value = datetime_value[0] + time_unit = datetime_value[1] if len(datetime_value) > 1 else "Microseconds" + time_zone = datetime_value[2] if len(datetime_value) > 2 else None + assert isinstance(epoch_value, int), f"A {dtype} value should be an int but got {type(epoch_value)}" + if time_unit == "Milliseconds": + micros = epoch_value * 1_000 + elif time_unit == "Microseconds": + micros = epoch_value + else: + # A TIMESTAMP literal cannot hold nanoseconds, let polars apply the filter + msg = f"Unsupported datetime time unit {time_unit!r}" + raise NotImplementedError(msg) + dt_timestamp = datetime.datetime(1970, 1, 1) + datetime.timedelta(microseconds=micros) + if time_zone is None: + return f"'{dt_timestamp!s}'::TIMESTAMP" + # The value is UTC; compare as an instant so the session TimeZone does not shift it + return f"'{dt_timestamp!s}+00'::TIMESTAMPTZ" # Match simple numeric/boolean types if dtype in ( diff --git a/tests/fast/arrow/test_polars.py b/tests/fast/arrow/test_polars.py index f3d7e072..bfbcfabf 100644 --- a/tests/fast/arrow/test_polars.py +++ b/tests/fast/arrow/test_polars.py @@ -1,5 +1,6 @@ import datetime import json +from zoneinfo import ZoneInfo import pytest from packaging.version import parse as parse_version @@ -438,6 +439,36 @@ def test_polars_lazy_pushdown_timestamp(self, duckdb_cursor): lazy_df.filter((pl.col("a") == ts_2020) | (pl.col("b") == ts_2008)).select(pl.len()).collect().item() == 2 ) + @pytest.mark.parametrize("data_type", ["TIMESTAMP_S", "TIMESTAMP_MS", "TIMESTAMP", "TIMESTAMP_NS"]) + def test_polars_lazy_pushdown_timestamp_units(self, data_type, duckdb_cursor): + duckdb_cursor.execute(f"CREATE TABLE test_timestamp_units (a {data_type})") + duckdb_cursor.execute( + """ + INSERT INTO test_timestamp_units VALUES + ('2008-01-01 00:00:01'), ('2010-01-01 10:00:01'), ('2020-03-01 10:00:01'), (NULL) + """ + ) + lazy_df = duckdb_cursor.table("test_timestamp_units").pl(lazy=True) + ts_2010 = datetime.datetime(2010, 1, 1, 10, 0, 1) + + assert lazy_df.filter(pl.col("a") == ts_2010).select(pl.len()).collect().item() == 1 + assert lazy_df.filter(pl.col("a") > ts_2010).select(pl.len()).collect().item() == 1 + assert lazy_df.filter(pl.col("a") >= ts_2010).select(pl.len()).collect().item() == 2 + assert lazy_df.filter(pl.col("a") < ts_2010).select(pl.len()).collect().item() == 1 + + def test_polars_lazy_pushdown_timestamptz(self, duckdb_cursor): + duckdb_cursor.execute("SET TimeZone = 'America/New_York'") + duckdb_cursor.execute("CREATE TABLE test_timestamptz (a TIMESTAMPTZ)") + duckdb_cursor.execute( + "INSERT INTO test_timestamptz VALUES ('2024-01-01 10:00:00+00'), ('2024-01-01 14:00:00+00'), (NULL)" + ) + lazy_df = duckdb_cursor.table("test_timestamptz").pl(lazy=True) + # 12:00 UTC + noon_utc = datetime.datetime(2024, 1, 1, 7, 0, tzinfo=ZoneInfo("America/New_York")) + + assert lazy_df.filter(pl.col("a") < noon_utc).select(pl.len()).collect().item() == 1 + assert lazy_df.filter(pl.col("a") > noon_utc).select(pl.len()).collect().item() == 1 + @pytest.mark.skipif(pl_pre_1_35_0, reason="Polars < 1.36.0 expressions on dates produce casts in predicates") def test_polars_predicate_to_expression_post_1_36_0(self): ts_2008 = datetime.datetime(2008, 1, 1, 0, 0, 1)