Merge remote-tracking branch 'upstream/main' into simp-pandas-group-by

Author: dangotbanned

narwhals/_duckdb/expr.py

@@ -12,6 +12,7 @@
 from narwhals._duckdb.expr_str import DuckDBExprStringNamespace
 from narwhals._duckdb.expr_struct import DuckDBExprStructNamespace
 from narwhals._duckdb.utils import (
+    DeferredTimeZone,
     F,
     col,
     lit,
@@ -678,11 +679,26 @@ def _fill_constant(expr: Expression, value: Any) -> Expression:
         return self._with_elementwise(_fill_constant, value=value)
 
     def cast(self, dtype: IntoDType) -> Self:
-        def func(expr: Expression) -> Expression:
-            native_dtype = narwhals_to_native_dtype(dtype, self._version)
-            return expr.cast(DuckDBPyType(native_dtype))
+        def func(df: DuckDBLazyFrame) -> list[Expression]:
+            tz = DeferredTimeZone(df.native)
+            native_dtype = narwhals_to_native_dtype(dtype, self._version, tz)
+            return [expr.cast(DuckDBPyType(native_dtype)) for expr in self(df)]
+
+        def window_f(df: DuckDBLazyFrame, inputs: DuckDBWindowInputs) -> list[Expression]:
+            tz = DeferredTimeZone(df.native)
+            native_dtype = narwhals_to_native_dtype(dtype, self._version, tz)
+            return [
+                expr.cast(DuckDBPyType(native_dtype))
+                for expr in self.window_function(df, inputs)
+            ]
 
-        return self._with_elementwise(func)
+        return self.__class__(
+            func,
+            window_f,
+            evaluate_output_names=self._evaluate_output_names,
+            alias_output_names=self._alias_output_names,
+            version=self._version,
+        )
 
     @requires.backend_version((1, 3))
     def is_unique(self) -> Self:

narwhals/_duckdb/expr_dt.py

@@ -29,71 +29,71 @@ class DuckDBExprDateTimeNamespace(
     LazyExprNamespace["DuckDBExpr"], DateTimeNamespace["DuckDBExpr"]
 ):
     def year(self) -> DuckDBExpr:
-        return self.compliant._with_callable(lambda expr: F("year", expr))
+        return self.compliant._with_elementwise(lambda expr: F("year", expr))
 
     def month(self) -> DuckDBExpr:
-        return self.compliant._with_callable(lambda expr: F("month", expr))
+        return self.compliant._with_elementwise(lambda expr: F("month", expr))
 
     def day(self) -> DuckDBExpr:
-        return self.compliant._with_callable(lambda expr: F("day", expr))
+        return self.compliant._with_elementwise(lambda expr: F("day", expr))
 
     def hour(self) -> DuckDBExpr:
-        return self.compliant._with_callable(lambda expr: F("hour", expr))
+        return self.compliant._with_elementwise(lambda expr: F("hour", expr))
 
     def minute(self) -> DuckDBExpr:
-        return self.compliant._with_callable(lambda expr: F("minute", expr))
+        return self.compliant._with_elementwise(lambda expr: F("minute", expr))
 
     def second(self) -> DuckDBExpr:
-        return self.compliant._with_callable(lambda expr: F("second", expr))
+        return self.compliant._with_elementwise(lambda expr: F("second", expr))
 
     def millisecond(self) -> DuckDBExpr:
-        return self.compliant._with_callable(
+        return self.compliant._with_elementwise(
             lambda expr: F("millisecond", expr) - F("second", expr) * lit(MS_PER_SECOND)
         )
 
     def microsecond(self) -> DuckDBExpr:
-        return self.compliant._with_callable(
+        return self.compliant._with_elementwise(
             lambda expr: F("microsecond", expr) - F("second", expr) * lit(US_PER_SECOND)
         )
 
     def nanosecond(self) -> DuckDBExpr:
-        return self.compliant._with_callable(
+        return self.compliant._with_elementwise(
             lambda expr: F("nanosecond", expr) - F("second", expr) * lit(NS_PER_SECOND)
         )
 
     def to_string(self, format: str) -> DuckDBExpr:
-        return self.compliant._with_callable(
+        return self.compliant._with_elementwise(
             lambda expr: F("strftime", expr, lit(format))
         )
 
     def weekday(self) -> DuckDBExpr:
-        return self.compliant._with_callable(lambda expr: F("isodow", expr))
+        return self.compliant._with_elementwise(lambda expr: F("isodow", expr))
 
     def ordinal_day(self) -> DuckDBExpr:
-        return self.compliant._with_callable(lambda expr: F("dayofyear", expr))
+        return self.compliant._with_elementwise(lambda expr: F("dayofyear", expr))
 
     def date(self) -> DuckDBExpr:
-        return self.compliant._with_callable(lambda expr: expr.cast("date"))
+        return self.compliant._with_elementwise(lambda expr: expr.cast("date"))
 
     def total_minutes(self) -> DuckDBExpr:
-        return self.compliant._with_callable(
+        return self.compliant._with_elementwise(
             lambda expr: F("datepart", lit("minute"), expr)
         )
 
     def total_seconds(self) -> DuckDBExpr:
-        return self.compliant._with_callable(
+        return self.compliant._with_elementwise(
             lambda expr: lit(SECONDS_PER_MINUTE) * F("datepart", lit("minute"), expr)
             + F("datepart", lit("second"), expr)
         )
 
     def total_milliseconds(self) -> DuckDBExpr:
-        return self.compliant._with_callable(
+        return self.compliant._with_elementwise(
             lambda expr: lit(MS_PER_MINUTE) * F("datepart", lit("minute"), expr)
             + F("datepart", lit("millisecond"), expr)
         )
 
     def total_microseconds(self) -> DuckDBExpr:
-        return self.compliant._with_callable(
+        return self.compliant._with_elementwise(
             lambda expr: lit(US_PER_MINUTE) * F("datepart", lit("minute"), expr)
             + F("datepart", lit("microsecond"), expr)
         )
@@ -112,7 +112,7 @@ def truncate(self, every: str) -> DuckDBExpr:
         def _truncate(expr: Expression) -> Expression:
             return F("date_trunc", format, expr)
 
-        return self.compliant._with_callable(_truncate)
+        return self.compliant._with_elementwise(_truncate)
 
     def _no_op_time_zone(self, time_zone: str) -> DuckDBExpr:
         def func(df: DuckDBLazyFrame) -> Sequence[Expression]:
@@ -139,7 +139,7 @@ def convert_time_zone(self, time_zone: str) -> DuckDBExpr:
 
     def replace_time_zone(self, time_zone: str | None) -> DuckDBExpr:
         if time_zone is None:
-            return self.compliant._with_callable(lambda _input: _input.cast("timestamp"))
+            return self.compliant._with_elementwise(lambda expr: expr.cast("timestamp"))
         else:
             return self._no_op_time_zone(time_zone)
 

narwhals/_duckdb/expr_list.py

@@ -14,4 +14,4 @@ class DuckDBExprListNamespace(
     LazyExprNamespace["DuckDBExpr"], ListNamespace["DuckDBExpr"]
 ):
     def len(self) -> DuckDBExpr:
-        return self.compliant._with_callable(lambda expr: F("len", expr))
+        return self.compliant._with_elementwise(lambda expr: F("len", expr))

narwhals/_duckdb/expr_str.py

@@ -17,12 +17,12 @@ class DuckDBExprStringNamespace(
     LazyExprNamespace["DuckDBExpr"], StringNamespace["DuckDBExpr"]
 ):
     def starts_with(self, prefix: str) -> DuckDBExpr:
-        return self.compliant._with_callable(
+        return self.compliant._with_elementwise(
             lambda expr: F("starts_with", expr, lit(prefix))
         )
 
     def ends_with(self, suffix: str) -> DuckDBExpr:
-        return self.compliant._with_callable(
+        return self.compliant._with_elementwise(
             lambda expr: F("ends_with", expr, lit(suffix))
         )
 
@@ -32,7 +32,7 @@ def func(expr: Expression) -> Expression:
                 return F("contains", expr, lit(pattern))
             return F("regexp_matches", expr, lit(pattern))
 
-        return self.compliant._with_callable(func)
+        return self.compliant._with_elementwise(func)
 
     def slice(self, offset: int, length: int | None) -> DuckDBExpr:
         def func(expr: Expression) -> Expression:
@@ -46,35 +46,37 @@ def func(expr: Expression) -> Expression:
                 F("length", expr) if length is None else lit(length) + offset_lit,
             )
 
-        return self.compliant._with_callable(func)
+        return self.compliant._with_elementwise(func)
 
     def split(self, by: str) -> DuckDBExpr:
-        return self.compliant._with_callable(lambda expr: F("str_split", expr, lit(by)))
+        return self.compliant._with_elementwise(
+            lambda expr: F("str_split", expr, lit(by))
+        )
 
     def len_chars(self) -> DuckDBExpr:
-        return self.compliant._with_callable(lambda expr: F("length", expr))
+        return self.compliant._with_elementwise(lambda expr: F("length", expr))
 
     def to_lowercase(self) -> DuckDBExpr:
-        return self.compliant._with_callable(lambda expr: F("lower", expr))
+        return self.compliant._with_elementwise(lambda expr: F("lower", expr))
 
     def to_uppercase(self) -> DuckDBExpr:
-        return self.compliant._with_callable(lambda expr: F("upper", expr))
+        return self.compliant._with_elementwise(lambda expr: F("upper", expr))
 
     def strip_chars(self, characters: str | None) -> DuckDBExpr:
         import string
 
-        return self.compliant._with_callable(
+        return self.compliant._with_elementwise(
             lambda expr: F(
                 "trim", expr, lit(string.whitespace if characters is None else characters)
             )
         )
 
     def replace_all(self, pattern: str, value: str, *, literal: bool) -> DuckDBExpr:
         if not literal:
-            return self.compliant._with_callable(
+            return self.compliant._with_elementwise(
                 lambda expr: F("regexp_replace", expr, lit(pattern), lit(value), lit("g"))
             )
-        return self.compliant._with_callable(
+        return self.compliant._with_elementwise(
             lambda expr: F("replace", expr, lit(pattern), lit(value))
         )
 
@@ -83,7 +85,7 @@ def to_datetime(self, format: str | None) -> DuckDBExpr:
             msg = "Cannot infer format with DuckDB backend, please specify `format` explicitly."
             raise NotImplementedError(msg)
 
-        return self.compliant._with_callable(
+        return self.compliant._with_elementwise(
             lambda expr: F("strptime", expr, lit(format))
         )
 
@@ -119,6 +121,8 @@ def func(expr: Expression) -> Expression:
                 .otherwise(expr)
             )
 
+        # can't use `_with_elementwise` due to `when` operator.
+        # TODO(unassigned): implement `window_func` like we do in `Expr.cast`
         return self.compliant._with_callable(func)
 
     replace = not_implemented()

narwhals/_duckdb/expr_struct.py

@@ -14,6 +14,6 @@ class DuckDBExprStructNamespace(
     LazyExprNamespace["DuckDBExpr"], StructNamespace["DuckDBExpr"]
 ):
     def field(self, name: str) -> DuckDBExpr:
-        return self.compliant._with_callable(
+        return self.compliant._with_elementwise(
             lambda expr: F("struct_extract", expr, lit(name))
         ).alias(name)

narwhals/_duckdb/namespace.py

@@ -13,7 +13,14 @@
 from narwhals._duckdb.dataframe import DuckDBLazyFrame
 from narwhals._duckdb.expr import DuckDBExpr
 from narwhals._duckdb.selectors import DuckDBSelectorNamespace
-from narwhals._duckdb.utils import F, concat_str, lit, narwhals_to_native_dtype, when
+from narwhals._duckdb.utils import (
+    DeferredTimeZone,
+    F,
+    concat_str,
+    lit,
+    narwhals_to_native_dtype,
+    when,
+)
 from narwhals._expression_parsing import (
     combine_alias_output_names,
     combine_evaluate_output_names,
@@ -142,13 +149,11 @@ def when(self, predicate: DuckDBExpr) -> DuckDBWhen:
         return DuckDBWhen.from_expr(predicate, context=self)
 
     def lit(self, value: NonNestedLiteral, dtype: IntoDType | None) -> DuckDBExpr:
-        def func(_df: DuckDBLazyFrame) -> list[Expression]:
+        def func(df: DuckDBLazyFrame) -> list[Expression]:
+            tz = DeferredTimeZone(df.native)
             if dtype is not None:
-                return [
-                    lit(value).cast(
-                        narwhals_to_native_dtype(dtype, version=self._version)  # type: ignore[arg-type]
-                    )
-                ]
+                target = narwhals_to_native_dtype(dtype, self._version, tz)
+                return [lit(value).cast(target)]  # type: ignore[arg-type]
             return [lit(value)]
 
         return self._expr(

narwhals/_duckdb/utils.py

@@ -30,6 +30,12 @@
     "us": "microsecond",
     "ns": "nanosecond",
 }
+UNIT_TO_TIMESTAMPS = {
+    "s": "TIMESTAMP_S",
+    "ms": "TIMESTAMP_MS",
+    "us": "TIMESTAMP",
+    "ns": "TIMESTAMP_NS",
+}
 
 col = duckdb.ColumnExpression
 """Alias for `duckdb.ColumnExpression`."""
@@ -182,7 +188,10 @@ def _non_nested_native_to_narwhals_dtype(duckdb_dtype_id: str, version: Version)
         "float": dtypes.Float32(),
         "varchar": dtypes.String(),
         "date": dtypes.Date(),
+        "timestamp_s": dtypes.Datetime("s"),
+        "timestamp_ms": dtypes.Datetime("ms"),
         "timestamp": dtypes.Datetime(),
+        "timestamp_ns": dtypes.Datetime("ns"),
         "boolean": dtypes.Boolean(),
         "interval": dtypes.Duration(),
         "decimal": dtypes.Decimal(),
@@ -191,7 +200,9 @@ def _non_nested_native_to_narwhals_dtype(duckdb_dtype_id: str, version: Version)
     }.get(duckdb_dtype_id, dtypes.Unknown())
 
 
-def narwhals_to_native_dtype(dtype: IntoDType, version: Version) -> str:  # noqa: C901, PLR0912, PLR0915
+def narwhals_to_native_dtype(  # noqa: PLR0912,PLR0915,C901
+    dtype: IntoDType, version: Version, deferred_time_zone: DeferredTimeZone
+) -> str:
     dtypes = version.dtypes
     if isinstance_or_issubclass(dtype, dtypes.Decimal):
         msg = "Casting to Decimal is not supported yet."
@@ -242,32 +253,40 @@ def narwhals_to_native_dtype(dtype: IntoDType, version: Version) -> str:  # noqa
         raise ValueError(msg)
 
     if isinstance_or_issubclass(dtype, dtypes.Datetime):
-        _time_unit = dtype.time_unit
-        _time_zone = dtype.time_zone
-        msg = "todo"
-        raise NotImplementedError(msg)
-    if isinstance_or_issubclass(dtype, dtypes.Duration):  # pragma: no cover
-        _time_unit = dtype.time_unit
-        msg = "todo"
-        raise NotImplementedError(msg)
-    if isinstance_or_issubclass(dtype, dtypes.Date):  # pragma: no cover
+        tu = dtype.time_unit
+        tz = dtype.time_zone
+        if not tz:
+            return UNIT_TO_TIMESTAMPS[tu]
+        if tu != "us":
+            msg = f"Only microsecond precision is supported for timezone-aware `Datetime` in DuckDB, got {tu} precision"
+            raise ValueError(msg)
+        if tz != (rel_tz := deferred_time_zone.time_zone):  # pragma: no cover
+            msg = f"Only the connection time zone {rel_tz} is supported, got: {tz}."
+            raise ValueError(msg)
+        # TODO(unassigned): cover once https://github.com/narwhals-dev/narwhals/issues/2742 addressed
+        return "TIMESTAMPTZ"  # pragma: no cover
+    if isinstance_or_issubclass(dtype, dtypes.Duration):
+        if (tu := dtype.time_unit) != "us":  # pragma: no cover
+            msg = f"Only microsecond-precision Duration is supported, got {tu} precision"
+        return "INTERVAL"
+    if isinstance_or_issubclass(dtype, dtypes.Date):
         return "DATE"
     if isinstance_or_issubclass(dtype, dtypes.List):
-        inner = narwhals_to_native_dtype(dtype.inner, version)
+        inner = narwhals_to_native_dtype(dtype.inner, version, deferred_time_zone)
         return f"{inner}[]"
-    if isinstance_or_issubclass(dtype, dtypes.Struct):  # pragma: no cover
+    if isinstance_or_issubclass(dtype, dtypes.Struct):
         inner = ", ".join(
-            f'"{field.name}" {narwhals_to_native_dtype(field.dtype, version)}'
+            f'"{field.name}" {narwhals_to_native_dtype(field.dtype, version, deferred_time_zone)}'
             for field in dtype.fields
         )
         return f"STRUCT({inner})"
-    if isinstance_or_issubclass(dtype, dtypes.Array):  # pragma: no cover
+    if isinstance_or_issubclass(dtype, dtypes.Array):
         shape = dtype.shape
         duckdb_shape_fmt = "".join(f"[{item}]" for item in shape)
         inner_dtype: Any = dtype
         for _ in shape:
             inner_dtype = inner_dtype.inner
-        duckdb_inner = narwhals_to_native_dtype(inner_dtype, version)
+        duckdb_inner = narwhals_to_native_dtype(inner_dtype, version, deferred_time_zone)
         return f"{duckdb_inner}{duckdb_shape_fmt}"
     msg = f"Unknown dtype: {dtype}"  # pragma: no cover
     raise AssertionError(msg)