diff --git a/src/narwhals/_duckdb/expr.py b/src/narwhals/_duckdb/expr.py index 26c6524133..b95b567e1f 100644 --- a/src/narwhals/_duckdb/expr.py +++ b/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)) diff --git a/src/narwhals/_ibis/expr.py b/src/narwhals/_ibis/expr.py index 6178b8a618..48fcf0b85a 100644 --- a/src/narwhals/_ibis/expr.py +++ b/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 diff --git a/src/narwhals/_spark_like/expr.py b/src/narwhals/_spark_like/expr.py index 917e1aef23..aedc45921f 100644 --- a/src/narwhals/_spark_like/expr.py +++ b/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 diff --git a/src/narwhals/_sql/expr.py b/src/narwhals/_sql/expr.py index 0e9ece8241..a9e1aa6268 100644 --- a/src/narwhals/_sql/expr.py +++ b/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)) diff --git a/tests/expr_and_series/null_count_test.py b/tests/expr_and_series/null_count_test.py index 6efef0ada1..ee7c1003b2 100644 --- a/tests/expr_and_series/null_count_test.py +++ b/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: