chore!: use Python API more for DuckDB support (#2370)

Author: MarcoGorelli

.github/workflows/check_tpch_queries.yml

@@ -29,9 +29,8 @@ jobs:
           cache-dependency-glob: "pyproject.toml"
       - name: local-install
         run: |
-          # once https://github.com/duckdb/duckdb/issues/16445
-          # is addressed, try using `--pre`
           uv pip install -U -e ".[dask]" --group core-tests --system
+          uv pip install -U --pre duckdb --system
       - name: generate-data
         run: cd tpch && python generate_data.py
       - name: tpch-tests 

narwhals/_duckdb/dataframe.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import contextlib
 from functools import reduce
 from operator import and_
 from typing import TYPE_CHECKING
@@ -21,6 +22,7 @@
 from narwhals.exceptions import ColumnNotFoundError
 from narwhals.exceptions import InvalidOperationError
 from narwhals.typing import CompliantDataFrame
+from narwhals.typing import CompliantLazyFrame
 from narwhals.utils import Implementation
 from narwhals.utils import Version
 from narwhals.utils import generate_temporary_column_name
@@ -48,7 +50,8 @@
     from narwhals.typing import LazyUniqueKeepStrategy
     from narwhals.utils import _FullContext
 
-from narwhals.typing import CompliantLazyFrame
+with contextlib.suppress(ImportError):  # requires duckdb>=1.3.0
+    from duckdb import SQLExpression  # type: ignore[attr-defined, unused-ignore]
 
 
 class DuckDBLazyFrame(CompliantLazyFrame["DuckDBExpr", "duckdb.DuckDBPyRelation"]):
@@ -338,7 +341,7 @@ 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}"')
+                select.append(str(col(name)))
         # Replace with Python API call once
         # https://github.com/duckdb/duckdb/discussions/16947 is addressed.
         query = f"""
@@ -359,23 +362,28 @@ def unique(
         self, subset: Sequence[str] | None, *, keep: LazyUniqueKeepStrategy
     ) -> Self:
         if subset_ := subset if keep == "any" else (subset or self.columns):
+            if self._backend_version < (1, 3):
+                msg = (
+                    "At least version 1.3 of DuckDB is required for `unique` operation\n"
+                    "with `subset` specified."
+                )
+                raise NotImplementedError(msg)
             # 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}."
                 raise ColumnNotFoundError(msg)
             idx_name = generate_temporary_column_name(8, self.columns)
             count_name = generate_temporary_column_name(8, [*self.columns, idx_name])
             partition_by_sql = generate_partition_by_sql(*(subset_))
-            rel = self.native  # noqa: F841
-            query = f"""
-                select *,
-                        row_number() over ({partition_by_sql}) as "{idx_name}",
-                        count(*) over ({partition_by_sql}) as "{count_name}"
-                from rel
-                """  # noqa: S608
             name = count_name if keep == "none" else idx_name
+            idx_expr = SQLExpression(
+                f"{FunctionExpression('row_number')} over ({partition_by_sql})"
+            ).alias(idx_name)
+            count_expr = SQLExpression(
+                f"{FunctionExpression('count', StarExpression())} over ({partition_by_sql})"
+            ).alias(count_name)
             return self._with_native(
-                duckdb.sql(query)
+                self.native.select(StarExpression(), idx_expr, count_expr)
                 .filter(col(name) == lit(1))
                 .select(StarExpression(exclude=[count_name, idx_name]))
             )
@@ -465,7 +473,7 @@ def unpivot(
             msg = "`value_name` cannot be empty string for duckdb backend."
             raise NotImplementedError(msg)
 
-        unpivot_on = ", ".join(f'"{name}"' for name in on_)
+        unpivot_on = ", ".join(str(col(name)) for name in on_)
         rel = self.native  # noqa: F841
         # Replace with Python API once
         # https://github.com/duckdb/duckdb/discussions/16980 is addressed.

narwhals/_duckdb/expr.py

@@ -565,7 +565,7 @@ 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})"
+            sql = f"{FunctionExpression('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)
@@ -580,7 +580,7 @@ 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})"
+            sql = f"{FunctionExpression('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

@@ -225,15 +225,15 @@ def narwhals_to_native_dtype(dtype: DType | type[DType], version: Version) -> st
 def generate_partition_by_sql(*partition_by: str) -> str:
     if not partition_by:
         return ""
-    by_sql = ", ".join([f'"{x}"' for x in partition_by])
+    by_sql = ", ".join([f"{col(x)}" for x in partition_by])
     return f"partition by {by_sql}"
 
 
 def generate_order_by_sql(*order_by: str, ascending: bool) -> str:
     if ascending:
-        by_sql = ", ".join([f'"{x}" asc nulls first' for x in order_by])
+        by_sql = ", ".join([f"{col(x)} asc nulls first" for x in order_by])
     else:
-        by_sql = ", ".join([f'"{x}" desc nulls last' for x in order_by])
+        by_sql = ", ".join([f"{col(x)} desc nulls last" for x in order_by])
     return f"order by {by_sql}"
 
 

tests/frame/unique_test.py

@@ -8,6 +8,7 @@
 # becomes LazyFrame instead of DataFrame
 import narwhals as nw
 from narwhals.exceptions import ColumnNotFoundError
+from tests.utils import DUCKDB_VERSION
 from tests.utils import Constructor
 from tests.utils import ConstructorEager
 from tests.utils import assert_equal_data
@@ -36,6 +37,8 @@ def test_unique_eager(
 
 
 def test_unique_invalid_subset(constructor: Constructor) -> None:
+    if "duckdb" in str(constructor) and DUCKDB_VERSION < (1, 3):
+        pytest.skip()
     df_raw = constructor(data)
     df = nw.from_native(df_raw)
     with pytest.raises(ColumnNotFoundError):
@@ -56,6 +59,8 @@ def test_unique(
     keep: Literal["any", "none"],
     expected: dict[str, list[float]],
 ) -> None:
+    if "duckdb" in str(constructor) and DUCKDB_VERSION < (1, 3):
+        pytest.skip()
     df_raw = constructor(data)
     df = nw.from_native(df_raw)
     result = df.unique(subset, keep=keep).sort("z")
@@ -76,6 +81,8 @@ def test_unique_full_subset(
     keep: Literal["any", "none"],
     expected: dict[str, list[float]],
 ) -> None:
+    if "duckdb" in str(constructor) and DUCKDB_VERSION < (1, 3):
+        pytest.skip()
     data = {"a": [1, 1, 1, 2], "b": [3, 3, 4, 4]}
     df_raw = constructor(data)
     df = nw.from_native(df_raw)