fix: use specialised duckdb functions for std_pop/std_samp/var_pop/var_samp (#2289)

Author: MarcoGorelli

narwhals/_duckdb/expr.py

@@ -372,22 +372,38 @@ def len(self: Self) -> Self:
         return self._from_call(lambda _input: FunctionExpression("count"))
 
     def std(self: Self, ddof: int) -> Self:
+        if ddof == 0:
+            return self._from_call(
+                lambda _input: FunctionExpression("stddev_pop", _input)
+            )
+        if ddof == 1:
+            return self._from_call(
+                lambda _input: FunctionExpression("stddev_samp", _input)
+            )
+
         def _std(_input: duckdb.Expression) -> duckdb.Expression:
             n_samples = FunctionExpression("count", _input)
-            # NOTE: Not implemented Error: Unable to transform python value of type '<class 'duckdb.duckdb.Expression'>' to DuckDB LogicalType
             return (
                 FunctionExpression("stddev_pop", _input)
                 * FunctionExpression("sqrt", n_samples)
-                / (FunctionExpression("sqrt", (n_samples - ddof)))  # type: ignore[operator]
+                / (FunctionExpression("sqrt", (n_samples - lit(ddof))))
             )
 
         return self._from_call(_std)
 
     def var(self: Self, ddof: int) -> Self:
+        if ddof == 0:
+            return self._from_call(lambda _input: FunctionExpression("var_pop", _input))
+        if ddof == 1:
+            return self._from_call(lambda _input: FunctionExpression("var_samp", _input))
+
         def _var(_input: duckdb.Expression) -> duckdb.Expression:
             n_samples = FunctionExpression("count", _input)
-            # NOTE: Not implemented Error: Unable to transform python value of type '<class 'duckdb.duckdb.Expression'>' to DuckDB LogicalType
-            return FunctionExpression("var_pop", _input) * n_samples / (n_samples - ddof)  # type: ignore[operator, no-any-return]
+            return (
+                FunctionExpression("var_pop", _input)
+                * n_samples
+                / (n_samples - lit(ddof))
+            )
 
         return self._from_call(_var)
 

tests/expr_and_series/arithmetic_test.py

@@ -338,3 +338,26 @@ def test_arithmetic_series_left_literal(
     df = nw.from_native(constructor_eager(data))
     result = df.select(getattr(lhs, attr)(nw.col("a")))
     assert_equal_data(result, {"literal": expected})
+
+
+def test_std_broadcating(constructor: Constructor) -> None:
+    # `std(ddof=2)` fails for duckdb here
+    df = nw.from_native(constructor({"a": [1, 2, 3]}))
+    result = df.with_columns(b=nw.col("a").std()).sort("a")
+    expected = {"a": [1, 2, 3], "b": [1.0, 1.0, 1.0]}
+    assert_equal_data(result, expected)
+    result = df.with_columns(b=nw.col("a").var()).sort("a")
+    expected = {"a": [1, 2, 3], "b": [1.0, 1.0, 1.0]}
+    assert_equal_data(result, expected)
+    result = df.with_columns(b=nw.col("a").std(ddof=0)).sort("a")
+    expected = {
+        "a": [1, 2, 3],
+        "b": [0.816496580927726, 0.816496580927726, 0.816496580927726],
+    }
+    assert_equal_data(result, expected)
+    result = df.with_columns(b=nw.col("a").var(ddof=0)).sort("a")
+    expected = {
+        "a": [1, 2, 3],
+        "b": [0.6666666666666666, 0.6666666666666666, 0.6666666666666666],
+    }
+    assert_equal_data(result, expected)

tests/expr_and_series/over_test.py

@@ -50,8 +50,6 @@ def test_over_single(constructor: Constructor) -> None:
 
 
 def test_over_std_var(request: pytest.FixtureRequest, constructor: Constructor) -> None:
-    if "duckdb" in str(constructor):
-        request.applymarker(pytest.mark.xfail)
     if "cudf" in str(constructor):
         # https://github.com/rapidsai/cudf/issues/18159
         request.applymarker(pytest.mark.xfail)

tests/frame/add_test.py

@@ -1,15 +1,11 @@
 from __future__ import annotations
 
-import pytest
-
 import narwhals.stable.v1 as nw
 from tests.utils import Constructor
 from tests.utils import assert_equal_data
 
 
-def test_add(constructor: Constructor, request: pytest.FixtureRequest) -> None:
-    if "duckdb" in str(constructor):
-        request.applymarker(pytest.mark.xfail)
+def test_add(constructor: Constructor) -> None:
     data = {"a": [1, 3, 2], "b": [4, 4, 6], "z": [7.0, 8.0, 9.0]}
     df = nw.from_native(constructor(data))
     result = df.with_columns(