feat: Support drop_nulls(OneOrIterable[ColumnNameOrSelector])

Author: dangotbanned

narwhals/_plan/_parse.py

@@ -136,10 +136,10 @@ def parse_into_selector_ir(input: ColumnNameOrSelector | Expr, /) -> SelectorIR:
         from narwhals._plan import selectors as cs
 
         selector = cs.by_name(input)
-    elif is_expr(input):  # pragma: no cover
+    elif is_expr(input):
         selector = input.meta.as_selector()
-    else:  # pragma: no cover
-        msg = f"cannot turn {qualified_type_name(input)!r} into selector"
+    else:
+        msg = f"cannot turn {qualified_type_name(input)!r} into a selector"
         raise TypeError(msg)
     return selector._ir
 

narwhals/_plan/dataframe.py

@@ -114,9 +114,11 @@ def drop(self, *columns: str, strict: bool = True) -> Self:
         return self._with_compliant(self._compliant.drop(columns, strict=strict))
 
     def drop_nulls(
-        self, subset: str | Sequence[str] | None = None
-    ) -> Self:  # pragma: no cover
-        subset = [subset] if isinstance(subset, str) else subset
+        self, subset: OneOrIterable[ColumnNameOrSelector] | None = None
+    ) -> Self:
+        if subset is not None:
+            s_irs = _parse.parse_into_seq_of_selector_ir(subset)
+            subset = expand_selector_irs_names(s_irs, schema=self)
         return self._with_compliant(self._compliant.drop_nulls(subset))
 
     def rename(self, mapping: Mapping[str, str]) -> Self:

tests/plan/compliant_test.py

@@ -5,6 +5,7 @@
 import pytest
 
 from narwhals._plan import selectors as ncs
+from narwhals.exceptions import ColumnNotFoundError, InvalidOperationError
 
 pytest.importorskip("pyarrow")
 pytest.importorskip("numpy")
@@ -19,7 +20,9 @@
 if TYPE_CHECKING:
     from collections.abc import Sequence
 
+    from narwhals._plan.typing import ColumnNameOrSelector, OneOrIterable
     from narwhals.typing import PythonLiteral
+    from tests.conftest import Data
 
 
 @pytest.fixture
@@ -44,12 +47,19 @@ def data_small() -> dict[str, Any]:
 
 
 @pytest.fixture
-def data_smaller(data_small: dict[str, Any]) -> dict[str, Any]:
+def data_small_af(data_small: dict[str, Any]) -> dict[str, Any]:
     """Use only columns `"a"-"f"`."""
     keep = {"a", "b", "c", "d", "e", "f"}
     return {k: v for k, v in data_small.items() if k in keep}
 
 
+@pytest.fixture
+def data_small_dh(data_small: dict[str, Any]) -> dict[str, Any]:
+    """Use only columns `"d"-"h"`."""
+    keep = {"d", "e", "f", "g", "h"}
+    return {k: v for k, v in data_small.items() if k in keep}
+
+
 @pytest.fixture
 def data_indexed() -> dict[str, Any]:
     """Used in https://github.com/narwhals-dev/narwhals/pull/2528."""
@@ -472,9 +482,9 @@ def test_select(
 def test_with_columns(
     expr: nwp.Expr | Sequence[nwp.Expr],
     expected: dict[str, Any],
-    data_smaller: dict[str, Any],
+    data_small_af: dict[str, Any],
 ) -> None:
-    result = dataframe(data_smaller).with_columns(expr)
+    result = dataframe(data_small_af).with_columns(expr)
     assert_equal_data(result, expected)
 
 
@@ -518,6 +528,62 @@ def test_row_is_py_literal(
     assert result == polars_result
 
 
+def test_drop_nulls(data_small_dh: Data) -> None:
+    df = dataframe(data_small_dh)
+    expected: Data = {"d": [], "e": [], "f": [], "g": [], "h": []}
+    result = df.drop_nulls()
+    assert_equal_data(result, expected)
+
+
+def test_drop_nulls_invalid(data_small_dh: Data) -> None:
+    df = dataframe(data_small_dh)
+    with pytest.raises(TypeError, match=r"cannot turn.+int.+into a selector"):
+        df.drop_nulls(123)  # type: ignore[arg-type]
+    with pytest.raises(
+        InvalidOperationError, match=r"cannot turn.+col\('a'\).first\(\).+into a selector"
+    ):
+        df.drop_nulls(nwp.col("a").first())  # type: ignore[arg-type]
+
+    with pytest.raises(ColumnNotFoundError):
+        df.drop_nulls(["j", "k"])
+
+    with pytest.raises(ColumnNotFoundError):
+        df.drop_nulls(ncs.by_name("j", "k"))
+
+    with pytest.raises(ColumnNotFoundError):
+        df.drop_nulls(ncs.by_index(-999))
+
+
+DROP_ROW_1: Data = {
+    "d": [7, 8],
+    "e": [9, 7],
+    "f": [False, None],
+    "g": [None, False],
+    "h": [None, True],
+}
+KEEP_ROW_3: Data = {"d": [8], "e": [7], "f": [None], "g": [False], "h": [True]}
+
+
+@pytest.mark.parametrize(
+    ("subset", "expected"),
+    [
+        ("e", DROP_ROW_1),
+        (nwp.col("e"), DROP_ROW_1),
+        (ncs.by_index(1), DROP_ROW_1),
+        (ncs.integer(), DROP_ROW_1),
+        ([ncs.numeric() | ~ncs.boolean()], DROP_ROW_1),
+        (["g", "h"], KEEP_ROW_3),
+        ([ncs.by_name("g", "h"), "d"], KEEP_ROW_3),
+    ],
+)
+def test_drop_nulls_subset(
+    data_small_dh: Data, subset: OneOrIterable[ColumnNameOrSelector], expected: Data
+) -> None:
+    df = dataframe(data_small_dh)
+    result = df.drop_nulls(subset)
+    assert_equal_data(result, expected)
+
+
 if TYPE_CHECKING:
     from typing_extensions import assert_type