Repository: narwhals-dev/narwhals
Pull request #380 (2024-07-02 08:44:07)

allow slicing by sequence
<!--
# 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.



MarcoGorelli: @EdAbati is this what you needed for scikit-learn?

Patch:
diff --git a/narwhals/_arrow/dataframe.py b/narwhals/_arrow/dataframe.py
index 9cf65e72..c3f2c0fb 100644
--- a/narwhals/_arrow/dataframe.py
+++ b/narwhals/_arrow/dataframe.py
@@ -2,10 +2,12 @@ from __future__ import annotations
 
 from typing import TYPE_CHECKING
 from typing import Any
+from typing import Sequence
 from typing import overload
 
 from narwhals._arrow.utils import translate_dtype
 from narwhals._pandas_like.utils import evaluate_into_exprs
+from narwhals.dependencies import get_numpy
 from narwhals.dependencies import get_pyarrow
 
 if TYPE_CHECKING:
@@ -65,17 +67,22 @@ class ArrowDataFrame:
             return ArrowSeries(self._dataframe[item], name=item)
 
         elif isinstance(item, slice):
-            from narwhals._arrow.dataframe import ArrowDataFrame
-
             if item.step is not None and item.step != 1:
                 msg = "Slicing with step is not supported on PyArrow tables"
                 raise NotImplementedError(msg)
             start = item.start or 0
             stop = item.stop or len(self._dataframe)
-            return ArrowDataFrame(
+            return self._from_dataframe(
                 self._dataframe.slice(item.start, stop - start),
             )
 
+        elif isinstance(item, Sequence) or (
+            (np := get_numpy()) is not None
+            and isinstance(item, np.ndarray)
+            and item.ndim == 1
+        ):
+            return self._from_dataframe(self._dataframe.take(item))
+
         else:  # pragma: no cover
             msg = f"Expected str or slice, got: {type(item)}"
             raise TypeError(msg)
diff --git a/narwhals/_arrow/expr.py b/narwhals/_arrow/expr.py
index 80d0740d..9af2b37a 100644
--- a/narwhals/_arrow/expr.py
+++ b/narwhals/_arrow/expr.py
@@ -68,6 +68,12 @@ class ArrowExpr:
     def cum_sum(self) -> Self:
         return reuse_series_implementation(self, "cum_sum")  # type: ignore[type-var]
 
+    def any(self) -> Self:
+        return reuse_series_implementation(self, "any", returns_scalar=True)  # type: ignore[type-var]
+
+    def all(self) -> Self:
+        return reuse_series_implementation(self, "all", returns_scalar=True)  # type: ignore[type-var]
+
     @property
     def dt(self) -> ArrowExprDateTimeNamespace:
         return ArrowExprDateTimeNamespace(self)
diff --git a/narwhals/_arrow/namespace.py b/narwhals/_arrow/namespace.py
index bc7182b7..e2bfa9da 100644
--- a/narwhals/_arrow/namespace.py
+++ b/narwhals/_arrow/namespace.py
@@ -1,10 +1,12 @@
 from __future__ import annotations
 
 from typing import TYPE_CHECKING
+from typing import Any
 from typing import Iterable
 
 from narwhals import dtypes
 from narwhals._arrow.expr import ArrowExpr
+from narwhals._arrow.series import ArrowSeries
 from narwhals.utils import flatten
 
 if TYPE_CHECKING:
@@ -12,7 +14,6 @@ if TYPE_CHECKING:
 
     from narwhals._arrow.dataframe import ArrowDataFrame
     from narwhals._arrow.expr import ArrowExpr
-    from narwhals._arrow.series import ArrowSeries
 
 
 class ArrowNamespace:
@@ -66,6 +67,12 @@ class ArrowNamespace:
             output_names=None,
         )
 
+    def _create_series_from_scalar(self, value: Any, series: ArrowSeries) -> ArrowSeries:
+        return ArrowSeries.from_iterable(
+            [value],
+            name=series.name,
+        )
+
     # --- not in spec ---
     def __init__(self) -> None: ...
 
diff --git a/narwhals/_arrow/series.py b/narwhals/_arrow/series.py
index 0044e3da..4fdafb97 100644
--- a/narwhals/_arrow/series.py
+++ b/narwhals/_arrow/series.py
@@ -2,8 +2,10 @@ from __future__ import annotations
 
 from typing import TYPE_CHECKING
 from typing import Any
+from typing import Iterable
 
 from narwhals._arrow.utils import translate_dtype
+from narwhals._pandas_like.utils import native_series_from_iterable
 from narwhals.dependencies import get_pyarrow_compute
 
 if TYPE_CHECKING:
@@ -29,6 +31,15 @@ class ArrowSeries:
             name=self._name,
         )
 
+    @classmethod
+    def from_iterable(cls: type[Self], data: Iterable[Any], name: str) -> Self:
+        return cls(
+            native_series_from_iterable(
+                data, name=name, index=None, implementation="arrow"
+            ),
+            name=name,
+        )
+
     def __len__(self) -> int:
         return len(self._series)
 
@@ -62,6 +73,14 @@ class ArrowSeries:
         pc = get_pyarrow_compute()
         return self._from_series(pc.cumulative_sum(self._series))
 
+    def any(self) -> bool:
+        pc = get_pyarrow_compute()
+        return pc.any(self._series)  # type: ignore[no-any-return]
+
+    def all(self) -> bool:
+        pc = get_pyarrow_compute()
+        return pc.all(self._series)  # type: ignore[no-any-return]
+
     @property
     def shape(self) -> tuple[int]:
         return (len(self._series),)
diff --git a/narwhals/_pandas_like/dataframe.py b/narwhals/_pandas_like/dataframe.py
index c58f08ce..dab477b0 100644
--- a/narwhals/_pandas_like/dataframe.py
+++ b/narwhals/_pandas_like/dataframe.py
@@ -6,6 +6,7 @@ from typing import Any
 from typing import Iterable
 from typing import Iterator
 from typing import Literal
+from typing import Sequence
 from typing import overload
 
 from narwhals._pandas_like.expr import PandasExpr
@@ -17,13 +18,12 @@ from narwhals._pandas_like.utils import validate_dataframe_comparand
 from narwhals._pandas_like.utils import validate_indices
 from narwhals.dependencies import get_cudf
 from narwhals.dependencies import get_modin
+from narwhals.dependencies import get_numpy
 from narwhals.dependencies import get_pandas
 from narwhals.utils import flatten
 from narwhals.utils import parse_version
 
 if TYPE_CHECKING:
-    from collections.abc import Sequence
-
     from typing_extensions import Self
 
     from narwhals._pandas_like.group_by import PandasGroupBy
@@ -101,12 +101,18 @@ class PandasDataFrame:
                 implementation=self._implementation,
             )
 
-        elif isinstance(item, slice):
+        elif isinstance(item, (slice, Sequence)):
             from narwhals._pandas_like.dataframe import PandasDataFrame
 
             return PandasDataFrame(
                 self._dataframe.iloc[item], implementation=self._implementation
             )
+        elif (
+            (np := get_numpy()) is not None
+            and isinstance(item, np.ndarray)
+            and item.ndim == 1
+        ):
+            return self._from_dataframe(self._dataframe.iloc[item])
 
         else:  # pragma: no cover
             msg = f"Expected str or slice, got: {type(item)}"
diff --git a/narwhals/_pandas_like/utils.py b/narwhals/_pandas_like/utils.py
index 1ff880e5..dc1d6893 100644
--- a/narwhals/_pandas_like/utils.py
+++ b/narwhals/_pandas_like/utils.py
@@ -10,6 +10,7 @@ from narwhals.dependencies import get_cudf
 from narwhals.dependencies import get_modin
 from narwhals.dependencies import get_numpy
 from narwhals.dependencies import get_pandas
+from narwhals.dependencies import get_pyarrow
 from narwhals.utils import flatten
 from narwhals.utils import isinstance_or_issubclass
 from narwhals.utils import parse_version
@@ -365,6 +366,9 @@ def native_series_from_iterable(
         mpd = get_modin()
 
         return mpd.Series(data, name=name, index=index)
+    if implementation == "arrow":
+        pa = get_pyarrow()
+        return pa.chunked_array([data])
     msg = f"Unknown implementation: {implementation}"  # pragma: no cover
     raise TypeError(msg)  # pragma: no cover
 
diff --git a/narwhals/dataframe.py b/narwhals/dataframe.py
index 9e770774..280c85e3 100644
--- a/narwhals/dataframe.py
+++ b/narwhals/dataframe.py
@@ -13,6 +13,7 @@ from narwhals._arrow.dataframe import ArrowDataFrame
 from narwhals._pandas_like.dataframe import PandasDataFrame
 from narwhals.dependencies import get_cudf
 from narwhals.dependencies import get_modin
+from narwhals.dependencies import get_numpy
 from narwhals.dependencies import get_pandas
 from narwhals.dependencies import get_polars
 from narwhals.dependencies import get_pyarrow
@@ -388,19 +389,26 @@ class DataFrame(BaseFrame):
         """
         return self._dataframe.shape  # type: ignore[no-any-return]
 
+    @overload
+    def __getitem__(self, item: Sequence[int]) -> Series: ...
+
     @overload
     def __getitem__(self, item: str) -> Series: ...
 
     @overload
     def __getitem__(self, item: slice) -> DataFrame: ...
 
-    def __getitem__(self, item: str | slice) -> Series | DataFrame:
+    def __getitem__(self, item: str | slice | Sequence[int]) -> Series | DataFrame:
         if isinstance(item, str):
             from narwhals.series import Series
 
             return Series(self._dataframe[item])
 
-        elif isinstance(item, (range, slice)):
+        elif isinstance(item, (Sequence, slice)) or (
+            (np := get_numpy()) is not None
+            and isinstance(item, np.ndarray)
+            and item.ndim == 1
+        ):
             return self._from_dataframe(self._dataframe[item])
 
         else:


Test patch:
diff --git a/tests/expr/any_all_test.py b/tests/expr/any_all_test.py
index 05aef972..c5bb88cb 100644
--- a/tests/expr/any_all_test.py
+++ b/tests/expr/any_all_test.py
@@ -1,12 +1,20 @@
 from typing import Any
 
+import pyarrow as pa
+import pytest
+
 import narwhals as nw
+from narwhals.utils import parse_version
 from tests.utils import compare_dicts
 
 
-def test_any_all(constructor: Any) -> None:
+def test_any_all(constructor_with_pyarrow: Any, request: Any) -> None:
+    if "table" in str(constructor_with_pyarrow) and parse_version(
+        pa.__version__
+    ) < parse_version("12.0.0"):  # pragma: no cover
+        request.applymarker(pytest.mark.xfail)
     df = nw.from_native(
-        constructor(
+        constructor_with_pyarrow(
             {
                 "a": [True, False, True],
                 "b": [True, True, True],
@@ -14,9 +22,9 @@ def test_any_all(constructor: Any) -> None:
             }
         )
     )
-    result = nw.to_native(df.select(nw.all().all()))
+    result = df.select(nw.all().all())
     expected = {"a": [False], "b": [True], "c": [False]}
     compare_dicts(result, expected)
-    result = nw.to_native(df.select(nw.all().any()))
+    result = df.select(nw.all().any())
     expected = {"a": [True], "b": [True], "c": [False]}
     compare_dicts(result, expected)
diff --git a/tests/frame/slice_test.py b/tests/frame/slice_test.py
index 5c47f287..83ed31fe 100644
--- a/tests/frame/slice_test.py
+++ b/tests/frame/slice_test.py
@@ -1,5 +1,7 @@
 from typing import Any
 
+import numpy as np
+import pandas as pd
 import polars as pl
 import pyarrow as pa
 import pytest
@@ -47,3 +49,23 @@ def test_slice_lazy_fails() -> None:
 def test_slice_int_fails(constructor_with_pyarrow: Any) -> None:
     with pytest.raises(TypeError, match="Expected str or slice, got: <class 'int'>"):
         _ = nw.from_native(constructor_with_pyarrow(data))[1]  # type: ignore[call-overload,index]
+
+
+def test_gather(constructor_with_pyarrow: Any) -> None:
+    df = nw.from_native(constructor_with_pyarrow(data), eager_only=True)
+    result = df[[0, 3, 1]]
+    expected = {
+        "a": [1.0, 4.0, 2.0],
+        "b": [11, 14, 12],
+    }
+    compare_dicts(result, expected)
+    result = df[np.array([0, 3, 1])]
+    compare_dicts(result, expected)
+
+
+def test_gather_pandas_index() -> None:
+    # check that we're slicing positionally, and not on the pandas index
+    df = pd.DataFrame({"a": [4, 1, 2], "b": [1, 4, 2]}, index=[2, 1, 3])
+    result = nw.from_native(df, eager_only=True)[[1, 2]]
+    expected = {"a": [1, 2], "b": [4, 2]}
+    compare_dicts(result, expected)