feat: Improve DataFrame.__getitem__ consistency

narwhals/_arrow/dataframe.py

@@ -16,19 +16,17 @@
 
 from narwhals._arrow.series import ArrowSeries
 from narwhals._arrow.utils import align_series_full_broadcast
-from narwhals._arrow.utils import convert_str_slice_to_int_slice
 from narwhals._arrow.utils import native_to_narwhals_dtype
-from narwhals._arrow.utils import select_rows
 from narwhals._compliant import EagerDataFrame
 from narwhals._expression_parsing import ExprKind
-from narwhals.dependencies import is_numpy_array_1d
+from narwhals.dependencies import is_numpy_array
 from narwhals.exceptions import ShapeError
 from narwhals.utils import Implementation
 from narwhals.utils import Version
 from narwhals.utils import check_column_exists
 from narwhals.utils import check_column_names_are_unique
+from narwhals.utils import convert_str_slice_to_int_slice
 from narwhals.utils import generate_temporary_column_name
-from narwhals.utils import is_sequence_but_not_str
 from narwhals.utils import not_implemented
 from narwhals.utils import parse_columns_to_drop
 from narwhals.utils import parse_version
@@ -51,7 +49,6 @@
     from narwhals._arrow.group_by import ArrowGroupBy
     from narwhals._arrow.namespace import ArrowNamespace
     from narwhals._arrow.typing import ArrowChunkedArray
-    from narwhals._arrow.typing import Indices  # type: ignore[attr-defined]
     from narwhals._arrow.typing import Mask  # type: ignore[attr-defined]
     from narwhals._arrow.typing import Order  # type: ignore[attr-defined]
     from narwhals._translate import IntoArrowTable
@@ -60,10 +57,14 @@
     from narwhals.typing import CompliantDataFrame
     from narwhals.typing import CompliantLazyFrame
     from narwhals.typing import JoinStrategy
+    from narwhals.typing import SizedMultiIndexSelector
+    from narwhals.typing import SizedMultiNameSelector
     from narwhals.typing import SizeUnit
     from narwhals.typing import UniqueKeepStrategy
     from narwhals.typing import _1DArray
     from narwhals.typing import _2DArray
+    from narwhals.typing import _SliceIndex
+    from narwhals.typing import _SliceName
     from narwhals.utils import Version
     from narwhals.utils import _FullContext
 
@@ -80,8 +81,9 @@
     PromoteOptions: TypeAlias = Literal["none", "default", "permissive"]
 
 
-class ArrowDataFrame(EagerDataFrame["ArrowSeries", "ArrowExpr", "pa.Table"]):
-    # --- not in the spec ---
+class ArrowDataFrame(
+    EagerDataFrame["ArrowSeries", "ArrowExpr", "pa.Table", "pa.ChunkedArray[Any]"]
+):
     def __init__(
         self,
         native_dataframe: pa.Table,
@@ -248,118 +250,61 @@ def get_column(self, name: str) -> ArrowSeries:
     def __array__(self, dtype: Any, *, copy: bool | None) -> _2DArray:
         return self.native.__array__(dtype, copy=copy)
 
-    @overload
-    def __getitem__(  # type: ignore[overload-overlap, unused-ignore]
-        self, item: str | tuple[slice | Sequence[int] | _1DArray, int | str]
-    ) -> ArrowSeries: ...
-    @overload
-    def __getitem__(
-        self,
-        item: (
-            int
-            | slice
-            | Sequence[int]
-            | Sequence[str]
-            | _1DArray
-            | tuple[
-                slice | Sequence[int] | _1DArray, slice | Sequence[int] | Sequence[str]
-            ]
-        ),
-    ) -> Self: ...
-    def __getitem__(
-        self,
-        item: (
-            str
-            | int
-            | slice
-            | Sequence[int]
-            | Sequence[str]
-            | _1DArray
-            | tuple[slice | Sequence[int] | _1DArray, int | str]
-            | tuple[
-                slice | Sequence[int] | _1DArray, slice | Sequence[int] | Sequence[str]
-            ]
-        ),
-    ) -> ArrowSeries | Self:
-        if isinstance(item, tuple):
-            item = tuple(list(i) if is_sequence_but_not_str(i) else i for i in item)  # pyright: ignore[reportAssignmentType]
-
-        if isinstance(item, str):
-            return ArrowSeries.from_native(self.native[item], context=self, name=item)
-        elif (
-            isinstance(item, tuple)
-            and len(item) == 2
-            and is_sequence_but_not_str(item[1])
-            and not isinstance(item[0], str)
-        ):
-            if len(item[1]) == 0:
-                # Return empty dataframe
-                return self._with_native(self.native.slice(0, 0).select([]))
-            selected_rows = select_rows(self.native, item[0])
-            return self._with_native(selected_rows.select(cast("Indices", item[1])))
-
-        elif isinstance(item, tuple) and len(item) == 2:
-            if isinstance(item[1], slice):
-                columns = self.columns
-                indices = cast("Indices", item[0])
-                if item[1] == slice(None):
-                    if isinstance(item[0], Sequence) and len(item[0]) == 0:
-                        return self._with_native(self.native.slice(0, 0))
-                    return self._with_native(self.native.take(indices))
-                if isinstance(item[1].start, str) or isinstance(item[1].stop, str):
-                    start, stop, step = convert_str_slice_to_int_slice(item[1], columns)
-                    return self._with_native(
-                        self.native.take(indices).select(columns[start:stop:step])
-                    )
-                if isinstance(item[1].start, int) or isinstance(item[1].stop, int):
-                    return self._with_native(
-                        self.native.take(indices).select(
-                            columns[item[1].start : item[1].stop : item[1].step]
-                        )
-                    )
-                msg = f"Expected slice of integers or strings, got: {type(item[1])}"  # pragma: no cover
-                raise TypeError(msg)  # pragma: no cover
-
-            # PyArrow columns are always strings
-            col_name = (
-                item[1]
-                if isinstance(item[1], str)
-                else self.columns[cast("int", item[1])]
-            )
-            if isinstance(item[0], str):  # pragma: no cover
-                msg = "Can not slice with tuple with the first element as a str"
-                raise TypeError(msg)
-            if (isinstance(item[0], slice)) and (item[0] == slice(None)):
-                return ArrowSeries.from_native(
-                    self.native[col_name], context=self, name=col_name
-                )
-            selected_rows = select_rows(self.native, item[0])
-            return ArrowSeries.from_native(
-                selected_rows[col_name], context=self, name=col_name
-            )
+    def _gather(self, rows: SizedMultiIndexSelector[ArrowChunkedArray]) -> Self:
+        if len(rows) == 0:
+            return self._with_native(self.native.slice(0, 0))
+        if self._backend_version < (18,) and isinstance(rows, tuple):
+            rows = list(rows)
+        return self._with_native(self.native.take(rows))  # pyright: ignore[reportArgumentType]
+
+    def _gather_slice(self, rows: _SliceIndex | range) -> Self:
+        start = rows.start or 0
+        stop = rows.stop if rows.stop is not None else len(self.native)
+        if start < 0:
+            start = len(self.native) + start
+        if stop < 0:
+            stop = len(self.native) + stop
+        if rows.step is not None and rows.step != 1:
+            msg = "Slicing with step is not supported on PyArrow tables"
+            raise NotImplementedError(msg)
+        return self._with_native(self.native.slice(start, stop - start))
+
+    def _select_slice_name(self, columns: _SliceName) -> Self:
+        start, stop, step = convert_str_slice_to_int_slice(columns, self.columns)
+        return self._with_native(self.native.select(self.columns[start:stop:step]))
+
+    def _select_slice_index(self, columns: _SliceIndex | range) -> Self:
+        return self._with_native(
+            self.native.select(self.columns[columns.start : columns.stop : columns.step])
+        )
 
-        elif isinstance(item, slice):
-            if item.step is not None and item.step != 1:
-                msg = "Slicing with step is not supported on PyArrow tables"
-                raise NotImplementedError(msg)
-            columns = self.columns
-            if isinstance(item.start, str) or isinstance(item.stop, str):
-                start, stop, step = convert_str_slice_to_int_slice(item, columns)
-                return self._with_native(self.native.select(columns[start:stop:step]))
-            start = item.start or 0
-            stop = item.stop if item.stop is not None else len(self.native)
-            return self._with_native(self.native.slice(start, stop - start))
-
-        elif isinstance(item, Sequence) or is_numpy_array_1d(item):
-            if isinstance(item, Sequence) and len(item) > 0 and isinstance(item[0], str):
-                return self._with_native(self.native.select(cast("Indices", item)))
-            if isinstance(item, Sequence) and len(item) == 0:
-                return self._with_native(self.native.slice(0, 0))
-            return self._with_native(self.native.take(cast("Indices", item)))
+    def _select_multi_index(
+        self, columns: SizedMultiIndexSelector[ArrowChunkedArray]
+    ) -> Self:
+        selector: Sequence[int]
+        if isinstance(columns, pa.ChunkedArray):
+            # TODO @dangotbanned: Fix upstream with `pa.ChunkedArray.to_pylist(self) -> list[Any]:`
+            selector = cast("Sequence[int]", columns.to_pylist())
+        # TODO @dangotbanned: Fix upstream, it is actually much narrower
+        # **Doesn't accept `ndarray`**
+        elif is_numpy_array(columns):
+            selector = columns.tolist()
+        else:
+            selector = columns
+        return self._with_native(self.native.select(selector))
 
-        else:  # pragma: no cover
-            msg = f"Expected str or slice, got: {type(item)}"
-            raise TypeError(msg)
+    def _select_multi_name(
+        self, columns: SizedMultiNameSelector[ArrowChunkedArray]
+    ) -> Self:
+        selector: Sequence[str] | _1DArray
+        if isinstance(columns, pa.ChunkedArray):
+            # TODO @dangotbanned: Fix upstream with `pa.ChunkedArray.to_pylist(self) -> list[Any]:`
+            selector = cast("Sequence[str]", columns.to_pylist())
+        else:
+            selector = columns
+        # TODO @dangotbanned: Fix upstream `pa.Table.select` https://github.com/zen-xu/pyarrow-stubs/blob/f899bb35e10b36f7906a728e9f8acf3e0a1f9f64/pyarrow-stubs/__lib_pxi/table.pyi#L597
+        # NOTE: Investigate what `cython` actually checks
+        return self._with_native(self.native.select(selector))  # pyright: ignore[reportArgumentType]
 
     @property
     def schema(self) -> dict[str, DType]:

narwhals/_arrow/series.py

@@ -67,9 +67,11 @@
     from narwhals.typing import PythonLiteral
     from narwhals.typing import RankMethod
     from narwhals.typing import RollingInterpolationMethod
+    from narwhals.typing import SizedMultiIndexSelector
     from narwhals.typing import TemporalLiteral
     from narwhals.typing import _1DArray
     from narwhals.typing import _2DArray
+    from narwhals.typing import _SliceIndex
     from narwhals.utils import Version
     from narwhals.utils import _FullContext
 
@@ -406,20 +408,24 @@ def __native_namespace__(self) -> ModuleType:
     def name(self) -> str:
         return self._name
 
-    @overload
-    def __getitem__(self, idx: int) -> Any: ...
-
-    @overload
-    def __getitem__(self, idx: slice | Sequence[int] | ArrowChunkedArray) -> Self: ...
-
-    def __getitem__(
-        self, idx: int | slice | Sequence[int] | ArrowChunkedArray
-    ) -> Any | Self:
-        if isinstance(idx, int):
-            return maybe_extract_py_scalar(self.native[idx], return_py_scalar=True)
-        if isinstance(idx, (Sequence, pa.ChunkedArray)):
-            return self._with_native(self.native.take(idx))
-        return self._with_native(self.native[idx])
+    def _gather(self, rows: SizedMultiIndexSelector[ArrowChunkedArray]) -> Self:
+        if len(rows) == 0:
+            return self._with_native(self.native.slice(0, 0))
+        if self._backend_version < (18,) and isinstance(rows, tuple):
+            rows = list(rows)
+        return self._with_native(self.native.take(rows))
+
+    def _gather_slice(self, rows: _SliceIndex | range) -> Self:
+        start = rows.start or 0
+        stop = rows.stop if rows.stop is not None else len(self.native)
+        if start < 0:
+            start = len(self.native) + start
+        if stop < 0:
+            stop = len(self.native) + stop
+        if rows.step is not None and rows.step != 1:
+            msg = "Slicing with step is not supported on PyArrow tables"
+            raise NotImplementedError(msg)
+        return self._with_native(self.native.slice(start, stop - start))
 
     def scatter(self, indices: int | Sequence[int], values: Any) -> Self:
         import numpy as np  # ignore-banned-import
@@ -911,7 +917,7 @@ def rolling_sum(self, window_size: int, *, min_samples: int, center: bool) -> Se
         result = self._with_native(
             pc.if_else((count_in_window >= min_samples).native, rolling_sum.native, None)
         )
-        return result[offset:]
+        return result._gather_slice(slice(offset, None))
 
     def rolling_mean(self, window_size: int, *, min_samples: int, center: bool) -> Self:
         min_samples = min_samples if min_samples is not None else window_size
@@ -940,7 +946,7 @@ def rolling_mean(self, window_size: int, *, min_samples: int, center: bool) -> S
             )
             / count_in_window
         )
-        return result[offset:]
+        return result._gather_slice(slice(offset, None))
 
     def rolling_var(
         self, window_size: int, *, min_samples: int, center: bool, ddof: int
@@ -983,7 +989,7 @@ def rolling_var(
             )
         ) / self._with_native(pc.max_element_wise((count_in_window - ddof).native, 0))
 
-        return result[offset:]
+        return result._gather_slice(slice(offset, None, None))
 
     def rolling_std(
         self, window_size: int, *, min_samples: int, center: bool, ddof: int

narwhals/_arrow/utils.py

@@ -7,7 +7,6 @@
 from typing import Iterator
 from typing import Sequence
 from typing import cast
-from typing import overload
 
 import pyarrow as pa
 import pyarrow.compute as pc
@@ -18,8 +17,6 @@
 from narwhals.utils import isinstance_or_issubclass
 
 if TYPE_CHECKING:
-    from typing import TypeVar
-
     from typing_extensions import TypeAlias
     from typing_extensions import TypeIs
 
@@ -33,15 +30,12 @@
     from narwhals._arrow.typing import ScalarAny
     from narwhals.dtypes import DType
     from narwhals.typing import PythonLiteral
-    from narwhals.typing import _AnyDArray
     from narwhals.utils import Version
 
     # NOTE: stubs don't allow for `ChunkedArray[StructArray]`
     # Intended to represent the `.chunks` property storing `list[pa.StructArray]`
     ChunkedArrayStructArray: TypeAlias = ArrowChunkedArray
 
-    _T = TypeVar("_T")
-
     def is_timestamp(t: Any) -> TypeIs[pa.TimestampType[Any, Any]]: ...
     def is_duration(t: Any) -> TypeIs[pa.DurationType[Any]]: ...
     def is_list(t: Any) -> TypeIs[pa.ListType[Any]]: ...
@@ -324,41 +318,6 @@ def cast_for_truediv(
     return arrow_array, pa_object
 
 
-@overload
-def convert_slice_to_nparray(num_rows: int, rows_slice: slice) -> _AnyDArray: ...
-@overload
-def convert_slice_to_nparray(num_rows: int, rows_slice: _T) -> _T: ...
-def convert_slice_to_nparray(num_rows: int, rows_slice: slice | _T) -> _AnyDArray | _T:
-    if isinstance(rows_slice, slice):
-        import numpy as np  # ignore-banned-import
-
-        return np.arange(num_rows)[rows_slice]
-    else:
-        return rows_slice
-
-
-def select_rows(
-    table: pa.Table, rows: slice | int | Sequence[int] | _AnyDArray
-) -> pa.Table:
-    if isinstance(rows, slice) and rows == slice(None):
-        selected_rows = table
-    elif isinstance(rows, Sequence) and not rows:
-        selected_rows = table.slice(0, 0)
-    else:
-        range_ = convert_slice_to_nparray(num_rows=len(table), rows_slice=rows)
-        selected_rows = table.take(cast("list[int]", range_))
-    return selected_rows
-
-
-def convert_str_slice_to_int_slice(
-    str_slice: slice, columns: list[str]
-) -> tuple[int | None, int | None, int | None]:
-    start = columns.index(str_slice.start) if str_slice.start is not None else None
-    stop = columns.index(str_slice.stop) + 1 if str_slice.stop is not None else None
-    step = str_slice.step
-    return (start, stop, step)
-
-
 # Regex for date, time, separator and timezone components
 DATE_RE = r"(?P<date>\d{1,4}[-/.]\d{1,2}[-/.]\d{1,4}|\d{8})"
 SEP_RE = r"(?P<sep>\s|T)"