Repository: narwhals-dev/narwhals
Pull request #847 (2024-08-22 14:01:00)

feat: Support more indexing: boolean lists to DataFrame.filter and Series.filter, add DataFrame.row, 
A bunch more things which would be useful to scikit-learn

<!--
# Thanks for contributing a pull request! 
## Please make sure you see our contribution guidelines: https://github.com/narwhals-dev/narwhals/blob/main/CONTRIBUTING.md
-->

## What type of PR is this? (check all applicable)

- [ ] 💾 Refactor
- [ ] ✨ Feature
- [ ] 🐛 Bug Fix
- [ ] 🔧 Optimization
- [ ] 📝 Documentation
- [ ] ✅ Test
- [ ] 🐳 Other

## Related issues 

- Related issue # 
- Closes #

## Checklist

- [ ] Code follows style guide (ruff)
- [ ] Tests added 
- [ ] Documented the changes

## If you have comments or can explain your changes, please do so below.




Patch:
diff --git a/.github/release-drafter.yml b/.github/release-drafter.yml
index 7b29b1cb..8f4c9881 100644
--- a/.github/release-drafter.yml
+++ b/.github/release-drafter.yml
@@ -9,31 +9,31 @@ autolabeler:
   - label: breaking
     title:
       # Example: feat!: ...
-      - '/^([B|b]uild|[C|c]hore|[CI|ci]|[D|d]epr|[D|d]oc|DOC|[F|f]eat|[F|f]ix|[P|p]erf|[R|r]efactor|[R|r]elease|[T|t]est)\! /'
+      - '/^([Bb]uild|[Cc]hore|CI|ci|[Dd]epr|[Dd]oc|DOC|[Ff]eat|[Ff]ix|[Pp]erf|[Rr]efactor|[Rr]elease|[Tt]est)\! /'
   - label: build
     title:
-      - '/^([B|b]uild)/'
+      - '/^([Bb]uild)/'
   - label: internal
     title:
-      - '/^(chore|ci|refactor|test|template|bench|Chore|CI|Refactor|Test|Template|Bench)'
+      - '/^(chore|ci|refactor|test|template|bench|Chore|CI|Refactor|Test|Template|Bench)/'
   - label: deprecation
     title:
-      - '/^([D|d]epr)/'
+      - '/^([Dd]epr)/'
   - label: documentation
     title:
-      - '/^([D|d]oc|DOC)/'
+      - '/^([Dd]oc|DOC)/'
   - label: enhancement
     title:
       - '/^(feat|enh|Feat|ENH|Enh)/'
   - label: fix
     title:
-      - '/^([F|f]ix)/'
+      - '/^([Ff]ix)/'
   - label: performance
     title:
-      - '/^([P|p]erf)/'
+      - '/^([Pp]erf)/'
   - label: release
     title:
-      - '/^([R|r]elease)/'
+      - '/^([Rr]elease)/'
       
 version-resolver:
   major:
diff --git a/docs/api-reference/dataframe.md b/docs/api-reference/dataframe.md
index 08ad0f7c..c144b4af 100644
--- a/docs/api-reference/dataframe.md
+++ b/docs/api-reference/dataframe.md
@@ -26,6 +26,7 @@
         - null_count
         - pipe
         - rename
+        - row
         - rows
         - schema
         - select
diff --git a/docs/installation.md b/docs/installation.md
index e1035e4f..3e36e339 100644
--- a/docs/installation.md
+++ b/docs/installation.md
@@ -11,6 +11,6 @@ Then, if you start the Python REPL and see the following:
 ```python
 >>> import narwhals
 >>> narwhals.__version__
-'1.5.2'
+'1.5.3'
 ```
 then installation worked correctly!
diff --git a/narwhals/__init__.py b/narwhals/__init__.py
index a23bd7b7..716ffeb5 100644
--- a/narwhals/__init__.py
+++ b/narwhals/__init__.py
@@ -22,11 +22,11 @@ from narwhals.dtypes import UInt32
 from narwhals.dtypes import UInt64
 from narwhals.dtypes import Unknown
 from narwhals.expr import Expr
-from narwhals.expr import all
+from narwhals.expr import all_ as all
 from narwhals.expr import all_horizontal
 from narwhals.expr import any_horizontal
 from narwhals.expr import col
-from narwhals.expr import len
+from narwhals.expr import len_ as len
 from narwhals.expr import lit
 from narwhals.expr import max
 from narwhals.expr import mean
@@ -50,7 +50,7 @@ from narwhals.utils import maybe_align_index
 from narwhals.utils import maybe_convert_dtypes
 from narwhals.utils import maybe_set_index
 
-__version__ = "1.5.2"
+__version__ = "1.5.3"
 
 __all__ = [
     "selectors",
diff --git a/narwhals/_arrow/dataframe.py b/narwhals/_arrow/dataframe.py
index 65a155e0..20e50716 100644
--- a/narwhals/_arrow/dataframe.py
+++ b/narwhals/_arrow/dataframe.py
@@ -63,6 +63,9 @@ class ArrowDataFrame:
     def __len__(self) -> int:
         return len(self._native_frame)
 
+    def row(self, index: int) -> tuple[Any, ...]:
+        return tuple(col[index] for col in self._native_frame)
+
     def rows(
         self, *, named: bool = False
     ) -> list[tuple[Any, ...]] | list[dict[str, Any]]:
@@ -377,11 +380,18 @@ class ArrowDataFrame:
         self,
         *predicates: IntoArrowExpr,
     ) -> Self:
-        plx = self.__narwhals_namespace__()
-        expr = plx.all_horizontal(*predicates)
-        # Safety: all_horizontal's expression only returns a single column.
-        mask = expr._call(self)[0]
-        return self._from_native_frame(self._native_frame.filter(mask._native_series))
+        if (
+            len(predicates) == 1
+            and isinstance(predicates[0], list)
+            and all(isinstance(x, bool) for x in predicates[0])
+        ):
+            mask = predicates[0]
+        else:
+            plx = self.__narwhals_namespace__()
+            expr = plx.all_horizontal(*predicates)
+            # Safety: all_horizontal's expression only returns a single column.
+            mask = expr._call(self)[0]._native_series
+        return self._from_native_frame(self._native_frame.filter(mask))
 
     def null_count(self) -> Self:
         import pyarrow as pa  # ignore-banned-import()
diff --git a/narwhals/_arrow/series.py b/narwhals/_arrow/series.py
index fb15f3aa..39394924 100644
--- a/narwhals/_arrow/series.py
+++ b/narwhals/_arrow/series.py
@@ -233,7 +233,8 @@ class ArrowSeries:
         return len(self._native_series)
 
     def filter(self, other: Any) -> Self:
-        other = validate_column_comparand(other)
+        if not (isinstance(other, list) and all(isinstance(x, bool) for x in other)):
+            other = validate_column_comparand(other)
         return self._from_native_series(self._native_series.filter(other))
 
     def mean(self) -> int:
diff --git a/narwhals/_dask/dataframe.py b/narwhals/_dask/dataframe.py
index 35ef28bb..9774d6c8 100644
--- a/narwhals/_dask/dataframe.py
+++ b/narwhals/_dask/dataframe.py
@@ -72,12 +72,19 @@ class DaskLazyFrame:
         self,
         *predicates: DaskExpr,
     ) -> Self:
-        from narwhals._dask.namespace import DaskNamespace
+        if (
+            len(predicates) == 1
+            and isinstance(predicates[0], list)
+            and all(isinstance(x, bool) for x in predicates[0])
+        ):
+            mask = predicates[0]
+        else:
+            from narwhals._dask.namespace import DaskNamespace
 
-        plx = DaskNamespace(backend_version=self._backend_version)
-        expr = plx.all_horizontal(*predicates)
-        # Safety: all_horizontal's expression only returns a single column.
-        mask = expr._call(self)[0]
+            plx = DaskNamespace(backend_version=self._backend_version)
+            expr = plx.all_horizontal(*predicates)
+            # Safety: all_horizontal's expression only returns a single column.
+            mask = expr._call(self)[0]
         return self._from_native_frame(self._native_frame.loc[mask])
 
     def lazy(self) -> Self:
diff --git a/narwhals/_pandas_like/dataframe.py b/narwhals/_pandas_like/dataframe.py
index c6a51be2..193955cb 100644
--- a/narwhals/_pandas_like/dataframe.py
+++ b/narwhals/_pandas_like/dataframe.py
@@ -283,15 +283,25 @@ class PandasLikeDataFrame:
             )
         )
 
+    def row(self, row: int) -> tuple[Any, ...]:
+        return tuple(x for x in self._native_frame.iloc[row])
+
     def filter(
         self,
         *predicates: IntoPandasLikeExpr,
     ) -> Self:
         plx = self.__narwhals_namespace__()
-        expr = plx.all_horizontal(*predicates)
-        # Safety: all_horizontal's expression only returns a single column.
-        mask = expr._call(self)[0]
-        _mask = validate_dataframe_comparand(self._native_frame.index, mask)
+        if (
+            len(predicates) == 1
+            and isinstance(predicates[0], list)
+            and all(isinstance(x, bool) for x in predicates[0])
+        ):
+            _mask = predicates[0]
+        else:
+            expr = plx.all_horizontal(*predicates)
+            # Safety: all_horizontal's expression only returns a single column.
+            mask = expr._call(self)[0]
+            _mask = validate_dataframe_comparand(self._native_frame.index, mask)
         return self._from_native_frame(self._native_frame.loc[_mask])
 
     def with_columns(
diff --git a/narwhals/_pandas_like/expr.py b/narwhals/_pandas_like/expr.py
index 193b1786..44154453 100644
--- a/narwhals/_pandas_like/expr.py
+++ b/narwhals/_pandas_like/expr.py
@@ -227,8 +227,8 @@ class PandasLikeExpr:
 
     def filter(self, *predicates: Any) -> Self:
         plx = self.__narwhals_namespace__()
-        expr = plx.all_horizontal(*predicates)
-        return reuse_series_implementation(self, "filter", other=expr)
+        other = plx.all_horizontal(*predicates)
+        return reuse_series_implementation(self, "filter", other=other)
 
     def drop_nulls(self) -> Self:
         return reuse_series_implementation(self, "drop_nulls")
diff --git a/narwhals/_pandas_like/series.py b/narwhals/_pandas_like/series.py
index 90434ebd..9532c39c 100644
--- a/narwhals/_pandas_like/series.py
+++ b/narwhals/_pandas_like/series.py
@@ -229,7 +229,8 @@ class PandasLikeSeries:
 
     def filter(self, other: Any) -> PandasLikeSeries:
         ser = self._native_series
-        other = validate_column_comparand(self._native_series.index, other)
+        if not (isinstance(other, list) and all(isinstance(x, bool) for x in other)):
+            other = validate_column_comparand(self._native_series.index, other)
         return self._from_native_series(self._rename(ser.loc[other], ser.name))
 
     def __eq__(self, other: object) -> PandasLikeSeries:  # type: ignore[override]
diff --git a/narwhals/dataframe.py b/narwhals/dataframe.py
index b09f12cc..20514972 100644
--- a/narwhals/dataframe.py
+++ b/narwhals/dataframe.py
@@ -158,8 +158,13 @@ class BaseFrame(Generic[FrameT]):
             )
         )
 
-    def filter(self, *predicates: IntoExpr | Iterable[IntoExpr]) -> Self:
-        predicates, _ = self._flatten_and_extract(*predicates)
+    def filter(self, *predicates: IntoExpr | Iterable[IntoExpr] | list[bool]) -> Self:
+        if not (
+            len(predicates) == 1
+            and isinstance(predicates[0], list)
+            and all(isinstance(x, bool) for x in predicates[0])
+        ):
+            predicates, _ = self._flatten_and_extract(*predicates)
         return self._from_compliant_dataframe(
             self._compliant_frame.filter(*predicates),
         )
@@ -599,6 +604,8 @@ class DataFrame(BaseFrame[FrameT]):
                 2
             ]
         """
+        if isinstance(item, int):
+            item = [item]
         if (
             isinstance(item, tuple)
             and len(item) == 2
@@ -693,6 +700,40 @@ class DataFrame(BaseFrame[FrameT]):
             }
         return self._compliant_frame.to_dict(as_series=as_series)  # type: ignore[no-any-return]
 
+    def row(self, index: int) -> tuple[Any, ...]:
+        """
+        Get values at given row.
+
+        !!!note
+            You should NEVER use this method to iterate over a DataFrame;
+            if you require row-iteration you should strongly prefer use of iter_rows() instead.
+
+        Arguments:
+            index: Row number.
+
+        Examples:
+            >>> import narwhals as nw
+            >>> import pandas as pd
+            >>> import polars as pl
+            >>> data = {"a": [1, 2, 3], "b": [4, 5, 6]}
+            >>> df_pd = pd.DataFrame(data)
+            >>> df_pl = pl.DataFrame(data)
+
+            Let's define a library-agnostic function to get the second row.
+
+            >>> @nw.narwhalify
+            ... def func(df):
+            ...     return df.row(1)
+
+            We can then pass pandas / Polars / any other supported library:
+
+            >>> func(df_pd)
+            (2, 5)
+            >>> func(df_pl)
+            (2, 5)
+        """
+        return self._compliant_frame.row(index)  # type: ignore[no-any-return]
+
     # inherited
     def pipe(self, function: Callable[[Any], Self], *args: Any, **kwargs: Any) -> Self:
         """
@@ -1483,14 +1524,15 @@ class DataFrame(BaseFrame[FrameT]):
         """
         return super().unique(subset, keep=keep, maintain_order=maintain_order)
 
-    def filter(self, *predicates: IntoExpr | Iterable[IntoExpr]) -> Self:
+    def filter(self, *predicates: IntoExpr | Iterable[IntoExpr] | list[bool]) -> Self:
         r"""
         Filter the rows in the DataFrame based on one or more predicate expressions.
 
         The original order of the remaining rows is preserved.
 
         Arguments:
-            predicates: Expression(s) that evaluates to a boolean Series.
+            *predicates: Expression(s) that evaluates to a boolean Series. Can
+                also be a (single!) boolean list.
 
         Examples:
             >>> import pandas as pd
@@ -2949,14 +2991,15 @@ class LazyFrame(BaseFrame[FrameT]):
         """
         return super().unique(subset, keep=keep, maintain_order=maintain_order)
 
-    def filter(self, *predicates: IntoExpr | Iterable[IntoExpr]) -> Self:
+    def filter(self, *predicates: IntoExpr | Iterable[IntoExpr] | list[bool]) -> Self:
         r"""
         Filter the rows in the LazyFrame based on a predicate expression.
 
         The original order of the remaining rows is preserved.
 
         Arguments:
-            *predicates: Expression that evaluates to a boolean Series.
+            *predicates: Expression that evaluates to a boolean Series. Can
+                also be a (single!) boolean list.
 
         Examples:
             >>> import pandas as pd
diff --git a/narwhals/expr.py b/narwhals/expr.py
index 74ead79e..a6b330bf 100644
--- a/narwhals/expr.py
+++ b/narwhals/expr.py
@@ -3657,7 +3657,8 @@ def col(*names: str | Iterable[str]) -> Expr:
     return Expr(func)
 
 
-def all() -> Expr:
+# Add underscore so it doesn't conflict with builtin `all`
+def all_() -> Expr:
     """
     Instantiate an expression representing all columns.
 
@@ -3696,7 +3697,8 @@ def all() -> Expr:
     return Expr(lambda plx: plx.all())
 
 
-def len() -> Expr:
+# Add underscore so it doesn't conflict with builtin `len`
+def len_() -> Expr:
     """
     Return the number of rows.
 
diff --git a/pyproject.toml b/pyproject.toml
index f49f65c9..83e3329a 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "narwhals"
-version = "1.5.2"
+version = "1.5.3"
 authors = [
   { name="Marco Gorelli", email="33491632+MarcoGorelli@users.noreply.github.com" },
 ]


Test patch:
diff --git a/tests/expr_and_series/filter_test.py b/tests/expr_and_series/filter_test.py
index ede0c49a..b55a0368 100644
--- a/tests/expr_and_series/filter_test.py
+++ b/tests/expr_and_series/filter_test.py
@@ -27,3 +27,6 @@ def test_filter_series(constructor_eager: Any) -> None:
     result = df.select(df["a"].filter((df["i"] < 2) & (df["c"] == 5)))
     expected = {"a": [0]}
     compare_dicts(result, expected)
+    result_s = df["a"].filter([True, False, False, False, False])
+    expected = {"a": [0]}
+    compare_dicts({"a": result_s}, expected)
diff --git a/tests/frame/filter_test.py b/tests/frame/filter_test.py
index 9b7ed45d..a8d3144a 100644
--- a/tests/frame/filter_test.py
+++ b/tests/frame/filter_test.py
@@ -12,11 +12,9 @@ def test_filter(constructor: Any) -> None:
     compare_dicts(result, expected)
 
 
-def test_filter_series(constructor_eager: Any) -> None:
+def test_filter_with_boolean_list(constructor: Any) -> None:
     data = {"a": [1, 3, 2], "b": [4, 4, 6], "z": [7.0, 8, 9]}
-    df = nw.from_native(constructor_eager(data), eager_only=True).with_columns(
-        mask=nw.col("a") > 1
-    )
-    result = df.filter(df["mask"]).drop("mask")
+    df = nw.from_native(constructor(data))
+    result = df.filter([False, True, True])
     expected = {"a": [3, 2], "b": [4, 6], "z": [8.0, 9.0]}
     compare_dicts(result, expected)
diff --git a/tests/frame/get_column_test.py b/tests/frame/get_column_test.py
index 70f65562..58766ac3 100644
--- a/tests/frame/get_column_test.py
+++ b/tests/frame/get_column_test.py
@@ -24,6 +24,9 @@ def test_non_string_name() -> None:
     result = nw.from_native(df, eager_only=True).get_column(0)  # type: ignore[arg-type]
     compare_dicts({"a": result}, {"a": [1, 2]})
     assert result.name == 0  # type: ignore[comparison-overlap]
-    with pytest.raises(TypeError, match="Expected str or slice"):
-        # Check that getitem would have raised
-        nw.from_native(df, eager_only=True)[0]  # type: ignore[call-overload]
+
+
+def test_get_single_row() -> None:
+    df = pd.DataFrame({"a": [1, 2], "b": [3, 4]})
+    result = nw.from_native(df, eager_only=True)[0]  # type: ignore[call-overload]
+    compare_dicts(result, {"a": [1], "b": [3]})
diff --git a/tests/frame/row_test.py b/tests/frame/row_test.py
new file mode 100644
index 00000000..602c50f5
--- /dev/null
+++ b/tests/frame/row_test.py
@@ -0,0 +1,14 @@
+from typing import Any
+
+import narwhals.stable.v1 as nw
+
+
+def test_row_column(constructor_eager: Any) -> None:
+    data = {
+        "a": [1.0, 2.0, 3.0, 4.0, 5.0, 6.0],
+        "b": [11, 12, 13, 14, 15, 16],
+    }
+    result = nw.from_native(constructor_eager(data), eager_only=True).row(2)
+    if "pyarrow_table" in str(constructor_eager):
+        result = tuple(x.as_py() for x in result)
+    assert result == (3.0, 13)
diff --git a/tests/frame/slice_test.py b/tests/frame/slice_test.py
index 222717d1..eea94d44 100644
--- a/tests/frame/slice_test.py
+++ b/tests/frame/slice_test.py
@@ -49,9 +49,16 @@ def test_slice_lazy_fails() -> None:
         _ = nw.from_native(pl.LazyFrame(data))[1:]
 
 
-def test_slice_int_fails(constructor_eager: Any) -> None:
-    with pytest.raises(TypeError, match="Expected str or slice, got: <class 'int'>"):
-        _ = nw.from_native(constructor_eager(data))[1]  # type: ignore[call-overload,index]
+def test_slice_int(constructor_eager: Any) -> None:
+    result = nw.from_native(constructor_eager(data), eager_only=True)[1]  # type: ignore[call-overload]
+    compare_dicts(result, {"a": [2], "b": [12]})
+
+
+def test_slice_fails(constructor_eager: Any) -> None:
+    class Foo: ...
+
+    with pytest.raises(TypeError, match="Expected str or slice, got:"):
+        nw.from_native(constructor_eager(data), eager_only=True)[Foo()]  # type: ignore[call-overload]
 
 
 def test_gather(constructor_eager: Any) -> None: