chore: use less sql for duckdb, do extra validation (#2331)

* chore: use less sql for duckdb

* join_asof too

* coverage

* rename assert_type -> ensure_type

* docs: Link to why ignore

* simplify predicate

* walrus

* refactor: Be a little fancier in `ensure_type`

---------

Co-authored-by: dangotbanned <[email protected]>

Author: MarcoGorelli

narwhals/_duckdb/dataframe.py

@@ -265,18 +265,21 @@ def join(
             if self._backend_version < (1, 1, 4):
                 msg = f"DuckDB>=1.1.4 is required for cross-join, found version: {self._backend_version}"
                 raise NotImplementedError(msg)
-            rel = self.native.set_alias("lhs").cross(
-                other.native.set_alias("rhs")
-            )  # pragma: no cover
+            rel = self.native.set_alias("lhs").cross(other.native.set_alias("rhs"))
         else:
             # help mypy
             assert left_on is not None  # noqa: S101
             assert right_on is not None  # noqa: S101
-            condition = " and ".join(
-                f'lhs."{left}" = rhs."{right}"' for left, right in zip(left_on, right_on)
+            it = (
+                col(f'lhs."{left}"') == col(f'rhs."{right}"')
+                for left, right in zip(left_on, right_on)
             )
+            condition: duckdb.Expression = reduce(and_, it)
             rel = self.native.set_alias("lhs").join(
-                other.native.set_alias("rhs"), condition=condition, how=native_how
+                other.native.set_alias("rhs"),
+                # NOTE: Fixed in `--pre` https://github.com/duckdb/duckdb/pull/16933
+                condition=condition,  # type: ignore[arg-type, unused-ignore]
+                how=native_how,
             )
 
         if native_how in {"inner", "left", "cross", "outer"}:
@@ -310,21 +313,22 @@ def join_asof(
     ) -> Self:
         lhs = self.native
         rhs = other.native
-        conditions = []
+        conditions: list[duckdb.Expression] = []
         if by_left is not None and by_right is not None:
-            conditions += [
-                f'lhs."{left}" = rhs."{right}"' for left, right in zip(by_left, by_right)
-            ]
+            conditions.extend(
+                col(f'lhs."{left}"') == col(f'rhs."{right}"')
+                for left, right in zip(by_left, by_right)
+            )
         else:
             by_left = by_right = []
         if strategy == "backward":
-            conditions += [f'lhs."{left_on}" >= rhs."{right_on}"']
+            conditions.append(col(f'lhs."{left_on}"') >= col(f'rhs."{right_on}"'))
         elif strategy == "forward":
-            conditions += [f'lhs."{left_on}" <= rhs."{right_on}"']
+            conditions.append(col(f'lhs."{left_on}"') <= col(f'rhs."{right_on}"'))
         else:
             msg = "Only 'backward' and 'forward' strategies are currently supported for DuckDB"
             raise NotImplementedError(msg)
-        condition = " and ".join(conditions)
+        condition: duckdb.Expression = reduce(and_, conditions)
         select = ["lhs.*"]
         for name in rhs.columns:
             if name in lhs.columns and (
@@ -333,6 +337,8 @@ def join_asof(
                 select.append(f'rhs."{name}" as "{name}{suffix}"')
             elif right_on is None or name not in {right_on, *by_right}:
                 select.append(f'"{name}"')
+        # Replace with Python API call once
+        # https://github.com/duckdb/duckdb/discussions/16947 is addressed.
         query = f"""
             SELECT {",".join(select)}
             FROM lhs
@@ -350,8 +356,7 @@ def collect_schema(self: Self) -> dict[str, DType]:
     def unique(
         self: Self, subset: Sequence[str] | None, *, keep: Literal["any", "none"]
     ) -> Self:
-        subset_ = subset if keep == "any" else (subset or self.columns)
-        if subset_:
+        if subset_ := subset if keep == "any" else (subset or self.columns):
             # Sanitise input
             if any(x not in self.columns for x in subset_):
                 msg = f"Columns {set(subset_).difference(self.columns)} not found in {self.columns}."
@@ -366,13 +371,10 @@ def unique(
                         count(*) over ({partition_by_sql}) as "{count_name}"
                 from rel
                 """  # noqa: S608
-            if keep == "none":
-                keep_condition = col(count_name) == lit(1)
-            else:
-                keep_condition = col(idx_name) == lit(1)
+            name = count_name if keep == "none" else idx_name
             return self._with_native(
                 duckdb.sql(query)
-                .filter(keep_condition)
+                .filter(col(name) == lit(1))
                 .select(StarExpression(exclude=[count_name, idx_name]))
             )
         return self._with_native(self.native.unique(", ".join(self.columns)))
@@ -463,6 +465,8 @@ def unpivot(
 
         unpivot_on = ", ".join(f'"{name}"' for name in on_)
         rel = self.native  # noqa: F841
+        # Replace with Python API once
+        # https://github.com/duckdb/duckdb/discussions/16980 is addressed.
         query = f"""
             unpivot rel
             on {unpivot_on}

narwhals/_duckdb/expr.py

@@ -20,6 +20,7 @@
 from narwhals._duckdb.expr_struct import DuckDBExprStructNamespace
 from narwhals._duckdb.utils import WindowInputs
 from narwhals._duckdb.utils import col
+from narwhals._duckdb.utils import ensure_type
 from narwhals._duckdb.utils import generate_order_by_sql
 from narwhals._duckdb.utils import generate_partition_by_sql
 from narwhals._duckdb.utils import lit
@@ -109,6 +110,8 @@ def _rolling_window_func(
         min_samples: int,
         ddof: int | None = None,
     ) -> WindowFunction:
+        ensure_type(window_size, int, type(None))
+        ensure_type(min_samples, int)
         supported_funcs = ["sum", "mean", "std", "var"]
         if center:
             half = (window_size - 1) // 2
@@ -139,11 +142,10 @@ def func(window_inputs: WindowInputs) -> duckdb.Expression:
             else:  # pragma: no cover
                 msg = f"Only the following functions are supported: {supported_funcs}.\nGot: {func_name}."
                 raise ValueError(msg)
-            sql = (
-                f"case when count({window_inputs.expr}) over {window} >= {min_samples}"
-                f"then {func_}({window_inputs.expr}) over {window} end"
-            )
-            return SQLExpression(sql)  # type: ignore[no-any-return, unused-ignore]
+            condition_sql = f"count({window_inputs.expr}) over {window} >= {min_samples}"
+            condition = SQLExpression(condition_sql)
+            value = SQLExpression(f"{func_}({window_inputs.expr}) over {window}")
+            return when(condition, value)
 
         return func
 
@@ -540,6 +542,8 @@ def round(self: Self, decimals: int) -> Self:
         )
 
     def shift(self, n: int) -> Self:
+        ensure_type(n, int)
+
         def func(window_inputs: WindowInputs) -> duckdb.Expression:
             order_by_sql = generate_order_by_sql(*window_inputs.order_by, ascending=True)
             partition_by_sql = generate_partition_by_sql(*window_inputs.partition_by)
@@ -560,8 +564,8 @@ def func(window_inputs: WindowInputs) -> duckdb.Expression:
                 )
             else:
                 partition_by_sql = f"partition by {window_inputs.expr}"
-            sql = f"row_number() over({partition_by_sql} {order_by_sql}) == 1"
-            return SQLExpression(sql)  # type: ignore[no-any-return, unused-ignore]
+            sql = f"row_number() over({partition_by_sql} {order_by_sql})"
+            return SQLExpression(sql) == lit(1)  # type: ignore[no-any-return, unused-ignore]
 
         return self._with_window_function(func)
 
@@ -575,8 +579,8 @@ def func(window_inputs: WindowInputs) -> duckdb.Expression:
                 )
             else:
                 partition_by_sql = f"partition by {window_inputs.expr}"
-            sql = f"row_number() over({partition_by_sql} {order_by_sql}) == 1"
-            return SQLExpression(sql)  # type: ignore[no-any-return, unused-ignore]
+            sql = f"row_number() over({partition_by_sql} {order_by_sql})"
+            return SQLExpression(sql) == lit(1)  # type: ignore[no-any-return, unused-ignore]
 
         return self._with_window_function(func)
 

narwhals/_duckdb/utils.py

@@ -235,3 +235,12 @@ def generate_order_by_sql(*order_by: str, ascending: bool) -> str:
     else:
         by_sql = ", ".join([f'"{x}" desc nulls last' for x in order_by])
     return f"order by {by_sql}"
+
+
+def ensure_type(obj: Any, *valid_types: type[Any]) -> None:
+    # Use this for extra (possibly redundant) validation in places where we
+    # use SQLExpression, as an extra guard against unsafe inputs.
+    if not isinstance(obj, valid_types):  # pragma: no cover
+        tp_names = " | ".join(tp.__name__ for tp in valid_types)
+        msg = f"Expected {tp_names!r}, got: {type(obj).__name__!r}"
+        raise TypeError(msg)

narwhals/_spark_like/utils.py

@@ -43,7 +43,7 @@ def __init__(
 # NOTE: don't lru_cache this as `ModuleType` isn't hashable
 def native_to_narwhals_dtype(
     dtype: _NativeDType, version: Version, spark_types: ModuleType
-) -> DType:  # pragma: no cover
+) -> DType:
     dtypes = import_dtypes_module(version=version)
     if TYPE_CHECKING:
         native = sqlframe_types
@@ -69,13 +69,15 @@ def native_to_narwhals_dtype(
     if isinstance(dtype, native.DateType):
         return dtypes.Date()
     if isinstance(dtype, native.TimestampNTZType):
-        return dtypes.Datetime()
+        # TODO(marco): cover this
+        return dtypes.Datetime()  # pragma: no cover
     if isinstance(dtype, native.TimestampType):
         # TODO(marco): is UTC correct, or should we be getting the connection timezone?
         # https://github.com/narwhals-dev/narwhals/issues/2165
         return dtypes.Datetime(time_zone="UTC")
     if isinstance(dtype, native.DecimalType):
-        return dtypes.Decimal()
+        # TODO(marco): cover this
+        return dtypes.Decimal()  # pragma: no cover
     if isinstance(dtype, native.ArrayType):
         return dtypes.List(
             inner=native_to_narwhals_dtype(
@@ -96,7 +98,7 @@ def native_to_narwhals_dtype(
         )
     if isinstance(dtype, native.BinaryType):
         return dtypes.Binary()
-    return dtypes.Unknown()
+    return dtypes.Unknown()  # pragma: no cover
 
 
 def narwhals_to_native_dtype(