fix: return Int64 from null_count on SQL backends, and 0 on empty frames

src/narwhals/_duckdb/expr.py

@@ -229,9 +229,6 @@ def func(expr: Expression) -> Expression:
     def len(self) -> Self:
         return self._with_callable(lambda _expr: F("count"))
 
-    def null_count(self) -> Self:
-        return self._with_callable(lambda expr: F("sum", expr.isnull().cast("int")))
-
     def is_nan(self) -> Self:
         return self._with_elementwise(lambda expr: F("isnan", expr))
 

src/narwhals/_ibis/expr.py

@@ -264,9 +264,6 @@ def func(df: IbisLazyFrame) -> Sequence[ir.IntegerScalar]:
             version=self._version,
         )
 
-    def null_count(self) -> Self:
-        return self._with_callable(lambda expr: expr.isnull().sum())
-
     def is_nan(self) -> Self:
         def func(expr: ir.FloatingValue) -> ir.Value:
             otherwise = expr.isnan() if is_floating(expr.type()) else False

src/narwhals/_spark_like/expr.py

@@ -294,12 +294,6 @@ def _median(expr: Column) -> Column:
 
         return self._with_callable(_median)
 
-    def null_count(self) -> Self:
-        def _null_count(expr: Column) -> Column:
-            return self._F.count_if(self._F.isnull(expr))
-
-        return self._with_callable(_null_count)
-
     def is_finite(self) -> Self:
         def _is_finite(expr: Column) -> Column:
             # A value is finite if it's not NaN, and not infinite, while NULLs should be

src/narwhals/_sql/expr.py

@@ -514,6 +514,25 @@ def min(self) -> Self:
     def count(self) -> Self:
         return self._with_callable(lambda expr: self._function("count", expr))
 
+    def null_count(self) -> Self:
+        dtypes = self._version.dtypes
+        is_null = self.is_null().cast(dtypes.Int32())
+
+        def f(expr: NativeExprT) -> NativeExprT:
+            return self._coalesce(self._function("sum", expr), self._lit(0))
+
+        def window_f(
+            df: SQLLazyFrameT, inputs: WindowInputs[NativeExprT]
+        ) -> Sequence[NativeExprT]:
+            # Partitions are never empty, so unlike `sum` there is no NULL to coalesce.
+            return [
+                self._window_expression(self._function("sum", expr), inputs.partition_by)
+                for expr in is_null(df)
+            ]
+
+        # DuckDB (also behind SQLFrame) widens integer `sum` and `count_if` to HUGEINT.
+        return is_null._with_callable(f, window_f).cast(dtypes.Int64())
+
     def sum(self) -> Self:
         def f(expr: NativeExprT) -> NativeExprT:
             return self._coalesce(self._function("sum", expr), self._lit(0))

tests/expr_and_series/null_count_test.py

@@ -1,16 +1,67 @@
 from __future__ import annotations
 
+from typing import Any
+
+import pytest
+
 import narwhals as nw
 from tests.utils import Constructor, ConstructorEager, assert_equal_data
 
-data = {"a": [1.0, None, None, 3.0], "b": [1.0, None, 4.0, 5.0]}
+data = {"a": [1.0, None, None, 3.0], "b": [1.0, None, 4.0, 5.0], "g": [1, 1, 1, 2]}
+
+
+def assert_integer_dtype(result: nw.LazyFrame[Any], *names: str) -> None:
+    schema, collected = result.collect_schema(), result.collect().schema
+    for name in names:
+        assert schema[name].is_integer()
+        assert collected[name] == schema[name]
 
 
-def test_null_count_expr(constructor: Constructor) -> None:
-    df = nw.from_native(constructor(data))
+def test_null_count_select(constructor: Constructor) -> None:
+    df = nw.from_native(constructor(data)).lazy()
     result = df.select(nw.col("a", "b").null_count())
-    expected = {"a": [2], "b": [1]}
+    assert_equal_data(result, {"a": [2], "b": [1]})
+    assert_integer_dtype(result, "a", "b")
+
+    result = df.with_columns(nw.col("a", "b").null_count()).sort("g")
+    expected = {"a": [2, 2, 2, 2], "b": [1, 1, 1, 1], "g": [1, 1, 1, 2]}
     assert_equal_data(result, expected)
+    assert_integer_dtype(result, "a", "b")
+
+
+@pytest.mark.filterwarnings("ignore:Found complex group-by:UserWarning")
+def test_null_count_group_by(
+    constructor: Constructor, request: pytest.FixtureRequest
+) -> None:
+    if any(x in str(constructor) for x in ("pyarrow_table", "dask")):
+        reason = "`null_count` is not a supported group-by aggregation on these backends"
+        request.applymarker(pytest.mark.xfail(reason=reason))
+    df = nw.from_native(constructor(data)).lazy()
+    result = df.group_by("g").agg(nw.col("a", "b").null_count()).sort("g")
+    assert_equal_data(result, {"g": [1, 2], "a": [2, 0], "b": [1, 0]})
+    assert_integer_dtype(result, "a", "b")
+
+
+def test_null_count_over(
+    constructor: Constructor, request: pytest.FixtureRequest
+) -> None:
+    if any(
+        x in str(constructor)
+        for x in ("pandas", "modin", "cudf", "dask", "pyarrow_table")
+    ):
+        reason = "`null_count().over()` is not supported on these backends"
+        request.applymarker(pytest.mark.xfail(reason=reason))
+    df = nw.from_native(constructor(data)).lazy()
+    result = df.with_columns(nw.col("a", "b").null_count().over("g")).sort("g")
+    expected = {"a": [2, 2, 2, 0], "b": [1, 1, 1, 0], "g": [1, 1, 1, 2]}
+    assert_equal_data(result, expected)
+    assert_integer_dtype(result, "a", "b")
+
+
+def test_null_count_empty(constructor: Constructor) -> None:
+    df = nw.from_native(constructor(data)).filter(nw.col("b") > 100)
+    result = df.select(nw.col("a", "b").null_count())
+    assert_equal_data(result, {"a": [0], "b": [0]})
 
 
 def test_null_count_series(constructor_eager: ConstructorEager) -> None: