fix truncate for timezone-aware timestamps for DuckDB (#2577)

* fix `truncate` for timezone-aware timestamps for DuckDB

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: MarcoGorelli

narwhals/_duckdb/expr_dt.py

@@ -1,6 +1,5 @@
 from __future__ import annotations
 
-from datetime import datetime
 from typing import TYPE_CHECKING
 
 from duckdb import FunctionExpression
@@ -10,6 +9,8 @@
 from narwhals.utils import not_implemented
 
 if TYPE_CHECKING:
+    from duckdb import Expression
+
     from narwhals._duckdb.expr import DuckDBExpr
 
 
@@ -109,15 +110,19 @@ def total_microseconds(self) -> DuckDBExpr:
 
     def truncate(self, every: str) -> DuckDBExpr:
         multiple, unit = parse_interval_string(every)
+        if multiple != 1:
+            # https://github.com/duckdb/duckdb/issues/17554
+            msg = f"Only multiple 1 is currently supported for DuckDB.\nGot {multiple!s}."
+            raise ValueError(msg)
         if unit == "ns":
             msg = "Truncating to nanoseconds is not yet supported for DuckDB."
             raise NotImplementedError(msg)
-        every = f"{multiple!s} {UNITS_DICT[unit]}"
-        return self._compliant_expr._with_callable(
-            lambda expr: FunctionExpression(
-                "time_bucket", lit(every), expr, lit(datetime(1970, 1, 1))
-            )
-        )
+        format = lit(UNITS_DICT[unit])
+
+        def _truncate(expr: Expression) -> Expression:
+            return FunctionExpression("date_trunc", format, expr)
+
+        return self._compliant_expr._with_callable(_truncate)
 
     def replace_time_zone(self, time_zone: str | None) -> DuckDBExpr:
         if time_zone is None:

tests/expr_and_series/dt/truncate_test.py

@@ -108,7 +108,7 @@ def test_truncate_multiples(
     every: str,
     expected: list[datetime],
 ) -> None:
-    if any(x in str(constructor) for x in ("sqlframe", "cudf", "pyspark")):
+    if any(x in str(constructor) for x in ("sqlframe", "cudf", "pyspark", "duckdb")):
         # Reasons:
         # - sqlframe: https://github.com/eakmanrq/sqlframe/issues/383
         # - cudf: https://github.com/rapidsai/cudf/issues/18654
@@ -194,3 +194,20 @@ def test_pandas_numpy_nat() -> None:
     expected = {"a": [datetime(2020, 1, 1), None, datetime(2020, 1, 1)]}
     assert_equal_data(result, expected)
     assert result.item(1, 0) is pd.NaT
+
+
+def test_truncate_tz_aware_duckdb() -> None:
+    pytest.importorskip("duckdb")
+    pytest.importorskip("zoneinfo")
+    import duckdb
+    from zoneinfo import ZoneInfo
+
+    duckdb.sql("""set timezone = 'Europe/Amsterdam'""")
+    rel = duckdb.sql("""select * from values (timestamptz '2020-10-25') df(a)""")
+    result = nw.from_native(rel).with_columns(a_truncated=nw.col("a").dt.truncate("1mo"))
+    expected = {
+        "a": [datetime(2020, 10, 25, tzinfo=ZoneInfo("Europe/Amsterdam"))],
+        "a_truncated": [datetime(2020, 10, 1, tzinfo=ZoneInfo("Europe/Amsterdam"))],
+    }
+    assert_equal_data(result, expected)
+    duckdb.sql("""set timezone = 'UTC'""")

tests/utils.py

@@ -4,6 +4,7 @@
 import os
 import sys
 import warnings
+from datetime import date, datetime
 from pathlib import Path
 from typing import TYPE_CHECKING, Any, Callable, Iterator, Mapping, Sequence
 
@@ -116,8 +117,11 @@ def assert_equal_data(result: Any, expected: Mapping[str, Any]) -> None:
                 )
             elif pd.isna(lhs):
                 are_equivalent_values = pd.isna(rhs)
+            elif type(lhs) is date and type(rhs) is datetime:
+                are_equivalent_values = datetime(lhs.year, lhs.month, lhs.day) == rhs
             else:
                 are_equivalent_values = lhs == rhs
+
             assert are_equivalent_values, (
                 f"Mismatch at index {i}: {lhs} != {rhs}\nExpected: {expected}\nGot: {result}"
             )