Respect time unit and time zone in polars datetime filter pushdown

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 (

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)