enh: reflect connection time zone in Datetime for DuckDB, support convert_time_zone and replace_time_zone (#2590)

Author: MarcoGorelli

narwhals/_duckdb/dataframe.py

@@ -9,6 +9,7 @@
 from duckdb import FunctionExpression, StarExpression
 
 from narwhals._duckdb.utils import (
+    DeferredTimeZone,
     col,
     evaluate_exprs,
     generate_partition_by_sql,
@@ -34,6 +35,7 @@
     import pandas as pd
     import pyarrow as pa
     from duckdb import Expression
+    from duckdb.typing import DuckDBPyType
     from typing_extensions import Self, TypeIs
 
     from narwhals._compliant.typing import CompliantDataFrameAny
@@ -70,7 +72,7 @@ def __init__(
         self._native_frame: duckdb.DuckDBPyRelation = df
         self._version = version
         self._backend_version = backend_version
-        self._cached_schema: dict[str, DType] | None = None
+        self._cached_native_schema: dict[str, DuckDBPyType] | None = None
         self._cached_columns: list[str] | None = None
         validate_backend_version(self._implementation, self._backend_version)
 
@@ -212,23 +214,25 @@ def filter(self, predicate: DuckDBExpr) -> Self:
 
     @property
     def schema(self) -> dict[str, DType]:
-        if self._cached_schema is None:
-            # Note: prefer `self._cached_schema` over `functools.cached_property`
+        if self._cached_native_schema is None:
+            # Note: prefer `self._cached_native_schema` over `functools.cached_property`
             # due to Python3.13 failures.
-            self._cached_schema = {
-                column_name: native_to_narwhals_dtype(duckdb_dtype, self._version)
-                for column_name, duckdb_dtype in zip(
-                    self.native.columns, self.native.types
-                )
-            }
-        return self._cached_schema
+            self._cached_native_schema = dict(zip(self.columns, self.native.types))
+
+        deferred_time_zone = DeferredTimeZone(self.native)
+        return {
+            column_name: native_to_narwhals_dtype(
+                duckdb_dtype, self._version, deferred_time_zone
+            )
+            for column_name, duckdb_dtype in zip(self.native.columns, self.native.types)
+        }
 
     @property
     def columns(self) -> list[str]:
         if self._cached_columns is None:
             self._cached_columns = (
                 list(self.schema)
-                if self._cached_schema is not None
+                if self._cached_native_schema is not None
                 else self.native.columns
             )
         return self._cached_columns

narwhals/_duckdb/expr_dt.py

@@ -1,16 +1,17 @@
 from __future__ import annotations
 
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Sequence
 
 from duckdb import FunctionExpression
 
-from narwhals._duckdb.utils import UNITS_DICT, lit
+from narwhals._duckdb.utils import UNITS_DICT, fetch_rel_time_zone, lit
 from narwhals._duration import parse_interval_string
 from narwhals.utils import not_implemented
 
 if TYPE_CHECKING:
     from duckdb import Expression
 
+    from narwhals._duckdb.dataframe import DuckDBLazyFrame
     from narwhals._duckdb.expr import DuckDBExpr
 
 
@@ -124,13 +125,36 @@ def _truncate(expr: Expression) -> Expression:
 
         return self._compliant_expr._with_callable(_truncate)
 
+    def _no_op_time_zone(self, time_zone: str) -> DuckDBExpr:
+        def func(df: DuckDBLazyFrame) -> Sequence[Expression]:
+            native_series_list = self._compliant_expr(df)
+            conn_time_zone = fetch_rel_time_zone(df.native)
+            if conn_time_zone != time_zone:
+                msg = (
+                    "DuckDB stores the time zone in the connection, rather than in the "
+                    f"data type, so changing the timezone to anything other than {conn_time_zone} "
+                    " (the current connection time zone) is not supported."
+                )
+                raise NotImplementedError(msg)
+            return native_series_list
+
+        return self._compliant_expr.__class__(
+            func,
+            evaluate_output_names=self._compliant_expr._evaluate_output_names,
+            alias_output_names=self._compliant_expr._alias_output_names,
+            backend_version=self._compliant_expr._backend_version,
+            version=self._compliant_expr._version,
+        )
+
+    def convert_time_zone(self, time_zone: str) -> DuckDBExpr:
+        return self._no_op_time_zone(time_zone)
+
     def replace_time_zone(self, time_zone: str | None) -> DuckDBExpr:
         if time_zone is None:
             return self._compliant_expr._with_callable(
                 lambda _input: _input.cast("timestamp")
             )
-        else:  # pragma: no cover
-            msg = "`replace_time_zone` with non-null `time_zone` not yet implemented for duckdb"
-            raise NotImplementedError(msg)
+        else:
+            return self._no_op_time_zone(time_zone)
 
     total_nanoseconds = not_implemented()

narwhals/_duckdb/series.py

@@ -2,7 +2,7 @@
 
 from typing import TYPE_CHECKING
 
-from narwhals._duckdb.utils import native_to_narwhals_dtype
+from narwhals._duckdb.utils import DeferredTimeZone, native_to_narwhals_dtype
 from narwhals.dependencies import get_duckdb
 
 if TYPE_CHECKING:
@@ -28,7 +28,11 @@ def __native_namespace__(self) -> ModuleType:
 
     @property
     def dtype(self) -> DType:
-        return native_to_narwhals_dtype(self._native_series.types[0], self._version)
+        return native_to_narwhals_dtype(
+            self._native_series.types[0],
+            self._version,
+            DeferredTimeZone(self._native_series),
+        )
 
     def __getattr__(self, attr: str) -> Never:
         msg = (  # pragma: no cover

narwhals/_duckdb/utils.py

@@ -8,7 +8,7 @@
 from narwhals.utils import Version, isinstance_or_issubclass
 
 if TYPE_CHECKING:
-    from duckdb import Expression
+    from duckdb import DuckDBPyRelation, Expression
     from duckdb.typing import DuckDBPyType
 
     from narwhals._duckdb.dataframe import DuckDBLazyFrame
@@ -95,21 +95,56 @@ def evaluate_exprs(
     return native_results
 
 
-def native_to_narwhals_dtype(duckdb_dtype: DuckDBPyType, version: Version) -> DType:
+class DeferredTimeZone:
+    """Object which gets passed between `native_to_narwhals_dtype` calls.
+
+    DuckDB stores the time zone in the connection, rather than in the dtypes, so
+    this ensures that when calculating the schema of a dataframe with multiple
+    timezone-aware columns, that the connection's time zone is only fetched once.
+
+    Note: we cannot make the time zone a cached `DuckDBLazyFrame` property because
+    the time zone can be modified after `DuckDBLazyFrame` creation:
+
+    ```python
+    df = nw.from_native(rel)
+    print(df.collect_schema())
+    rel.query("set timezone = 'Asia/Kolkata'")
+    print(df.collect_schema())  # should change to reflect new time zone
+    ```
+    """
+
+    _cached_time_zone: str | None = None
+
+    def __init__(self, rel: DuckDBPyRelation) -> None:
+        self._rel = rel
+
+    @property
+    def time_zone(self) -> str:
+        """Fetch relation time zone (if it wasn't calculated already)."""
+        if self._cached_time_zone is None:
+            self._cached_time_zone = fetch_rel_time_zone(self._rel)
+        return self._cached_time_zone
+
+
+def native_to_narwhals_dtype(
+    duckdb_dtype: DuckDBPyType, version: Version, deferred_time_zone: DeferredTimeZone
+) -> DType:
     duckdb_dtype_id = duckdb_dtype.id
     dtypes = version.dtypes
 
     # Handle nested data types first
     if duckdb_dtype_id == "list":
-        return dtypes.List(native_to_narwhals_dtype(duckdb_dtype.child, version=version))
+        return dtypes.List(
+            native_to_narwhals_dtype(duckdb_dtype.child, version, deferred_time_zone)
+        )
 
     if duckdb_dtype_id == "struct":
         children = duckdb_dtype.children
         return dtypes.Struct(
             [
                 dtypes.Field(
                     name=child[0],
-                    dtype=native_to_narwhals_dtype(child[1], version=version),
+                    dtype=native_to_narwhals_dtype(child[1], version, deferred_time_zone),
                 )
                 for child in children
             ]
@@ -123,7 +158,7 @@ def native_to_narwhals_dtype(duckdb_dtype: DuckDBPyType, version: Version) -> DT
             child, size = child[1].children
             shape.insert(0, size[1])
 
-        inner = native_to_narwhals_dtype(child[1], version=version)
+        inner = native_to_narwhals_dtype(child[1], version, deferred_time_zone)
         return dtypes.Array(inner=inner, shape=tuple(shape))
 
     if duckdb_dtype_id == "enum":
@@ -132,9 +167,20 @@ def native_to_narwhals_dtype(duckdb_dtype: DuckDBPyType, version: Version) -> DT
         categories = duckdb_dtype.children[0][1]
         return dtypes.Enum(categories=categories)
 
+    if duckdb_dtype_id == "timestamp with time zone":
+        return dtypes.Datetime(time_zone=deferred_time_zone.time_zone)
+
     return _non_nested_native_to_narwhals_dtype(duckdb_dtype_id, version)
 
 
+def fetch_rel_time_zone(rel: duckdb.DuckDBPyRelation) -> str:
+    result = rel.query(
+        "duckdb_settings()", "select value from duckdb_settings() where name = 'TimeZone'"
+    ).fetchone()
+    assert result is not None  # noqa: S101
+    return result[0]
+
+
 @lru_cache(maxsize=16)
 def _non_nested_native_to_narwhals_dtype(duckdb_dtype_id: str, version: Version) -> DType:
     dtypes = version.dtypes
@@ -154,9 +200,6 @@ def _non_nested_native_to_narwhals_dtype(duckdb_dtype_id: str, version: Version)
         "varchar": dtypes.String(),
         "date": dtypes.Date(),
         "timestamp": dtypes.Datetime(),
-        # TODO(marco): is UTC correct, or should we be getting the connection timezone?
-        # https://github.com/narwhals-dev/narwhals/issues/2165
-        "timestamp with time zone": dtypes.Datetime(time_zone="UTC"),
         "boolean": dtypes.Boolean(),
         "interval": dtypes.Duration(),
         "decimal": dtypes.Decimal(),

tests/conftest.py

@@ -136,6 +136,8 @@ def duckdb_lazy_constructor(obj: Data) -> duckdb.DuckDBPyRelation:
     import duckdb
     import polars as pl
 
+    duckdb.sql("""set timezone = 'UTC'""")
+
     _df = pl.LazyFrame(obj)
     return duckdb.table("_df")
 

tests/dtypes_test.py

@@ -457,3 +457,28 @@ def test_enum_repr() -> None:
 def test_enum_hash() -> None:
     assert nw.Enum(["a", "b"]) in {nw.Enum(["a", "b"])}
     assert nw.Enum(["a", "b"]) not in {nw.Enum(["a", "b", "c"])}
+
+
+def test_datetime_w_tz_duckdb() -> None:
+    pytest.importorskip("duckdb")
+    pytest.importorskip("zoneinfo")
+    import duckdb
+
+    duckdb.sql("""set timezone = 'Europe/Amsterdam'""")
+    df = nw.from_native(
+        duckdb.sql("""select * from values (timestamptz '2020-01-01')df(a)""")
+    )
+    result = df.collect_schema()
+    assert result["a"] == nw.Datetime("us", "Europe/Amsterdam")
+    duckdb.sql("""set timezone = 'Asia/Kathmandu'""")
+    result = df.collect_schema()
+    assert result["a"] == nw.Datetime("us", "Asia/Kathmandu")
+
+    df = nw.from_native(
+        duckdb.sql(
+            """select * from values (timestamptz '2020-01-01', [[timestamptz '2020-01-02']])df(a,b)"""
+        )
+    )
+    result = df.collect_schema()
+    assert result["a"] == nw.Datetime("us", "Asia/Kathmandu")
+    assert result["b"] == nw.List(nw.List(nw.Datetime("us", "Asia/Kathmandu")))

tests/expr_and_series/dt/convert_time_zone_test.py

@@ -134,3 +134,22 @@ def test_convert_time_zone_to_none_series(constructor_eager: ConstructorEager) -
     df = nw.from_native(constructor_eager(data))
     with pytest.raises(TypeError, match="Target `time_zone` cannot be `None`"):
         df["a"].dt.convert_time_zone(None)  # type: ignore[arg-type]
+
+
+def test_convert_time_zone_to_connection_tz_duckdb() -> None:
+    pytest.importorskip("duckdb")
+    pytest.importorskip("zoneinfo")
+    import duckdb
+    from zoneinfo import ZoneInfo
+
+    duckdb.sql("set timezone = 'Asia/Kolkata'")
+    rel = duckdb.sql("""select * from values (timestamptz '2020-01-01') df(a)""")
+    result = nw.from_native(rel).with_columns(
+        nw.col("a").dt.convert_time_zone("Asia/Kolkata")
+    )
+    expected = {"a": [datetime(2020, 1, 1, tzinfo=ZoneInfo("Asia/Kolkata"))]}
+    assert_equal_data(result, expected)
+    with pytest.raises(NotImplementedError):
+        result = nw.from_native(rel).with_columns(
+            nw.col("a").dt.convert_time_zone("Asia/Kathmandu")
+        )

tests/expr_and_series/dt/replace_time_zone_test.py

@@ -28,7 +28,8 @@ def test_replace_time_zone(
         or ("pyarrow_table" in str(constructor) and PYARROW_VERSION < (12,))
     ):
         pytest.skip()
-    if any(x in str(constructor) for x in ("cudf", "duckdb", "pyspark", "ibis")):
+
+    if any(x in str(constructor) for x in ("cudf", "pyspark", "ibis", "duckdb")):
         request.applymarker(pytest.mark.xfail)
     data = {
         "a": [
@@ -129,3 +130,22 @@ def test_replace_time_zone_none_series(constructor_eager: ConstructorEager) -> N
     result_str = result.select(df["a"].dt.to_string("%Y-%m-%dT%H:%M"))
     expected = {"a": ["2020-01-01T00:00", "2020-01-02T00:00"]}
     assert_equal_data(result_str, expected)
+
+
+def test_replace_time_zone_to_connection_tz_duckdb() -> None:
+    pytest.importorskip("duckdb")
+    pytest.importorskip("zoneinfo")
+    import duckdb
+    from zoneinfo import ZoneInfo
+
+    duckdb.sql("set timezone = 'Asia/Kolkata'")
+    rel = duckdb.sql("""select * from values (timestamptz '2020-01-01') df(a)""")
+    result = nw.from_native(rel).with_columns(
+        nw.col("a").dt.replace_time_zone("Asia/Kolkata")
+    )
+    expected = {"a": [datetime(2020, 1, 1, tzinfo=ZoneInfo("Asia/Kolkata"))]}
+    assert_equal_data(result, expected)
+    with pytest.raises(NotImplementedError):
+        result = nw.from_native(rel).with_columns(
+            nw.col("a").dt.replace_time_zone("Asia/Kathmandu")
+        )