Repository: narwhals-dev/narwhals
Pull request #622 (2024-07-25 17:15:53)

feat: allow slicing by tuple of sequences of ints
<!--
# 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 fancy taking a look, as you did quite some work on slicing?
MarcoGorelli: Thanks for your review 🙏 !

`_pandas_like` and `_arrow` aren't user-facing - it's better if we get the types there for internal correctness, but I'm not stressing out over them too much 😄 

Patch:
diff --git a/narwhals/_arrow/dataframe.py b/narwhals/_arrow/dataframe.py
index abf3ff7c..612f5773 100644
--- a/narwhals/_arrow/dataframe.py
+++ b/narwhals/_arrow/dataframe.py
@@ -123,6 +123,15 @@ class ArrowDataFrame:
                 name=item,
                 backend_version=self._backend_version,
             )
+        elif (
+            isinstance(item, tuple)
+            and len(item) == 2
+            and isinstance(item[1], (list, tuple))
+        ):
+            return self._from_native_dataframe(
+                self._native_dataframe.take(item[0]).select(item[1])
+            )
+
         elif isinstance(item, tuple) and len(item) == 2:
             from narwhals._arrow.series import ArrowSeries
 
diff --git a/narwhals/_pandas_like/dataframe.py b/narwhals/_pandas_like/dataframe.py
index ca3c5138..d2f7f2c7 100644
--- a/narwhals/_pandas_like/dataframe.py
+++ b/narwhals/_pandas_like/dataframe.py
@@ -122,11 +122,30 @@ class PandasLikeDataFrame:
                 backend_version=self._backend_version,
             )
 
+        elif (
+            isinstance(item, tuple)
+            and len(item) == 2
+            and isinstance(item[1], (tuple, list))
+        ):
+            if all(isinstance(x, int) for x in item[1]):
+                return self._from_native_dataframe(self._native_dataframe.iloc[item])
+            if all(isinstance(x, str) for x in item[1]):
+                item = (
+                    item[0],
+                    self._native_dataframe.columns.get_indexer(item[1]),
+                )
+                return self._from_native_dataframe(self._native_dataframe.iloc[item])
+            msg = (
+                f"Expected sequence str or int, got: {type(item[1])}"  # pragma: no cover
+            )
+            raise TypeError(msg)  # pragma: no cover
+
         elif isinstance(item, tuple) and len(item) == 2:
             from narwhals._pandas_like.series import PandasLikeSeries
 
             if isinstance(item[1], str):
-                native_series = self._native_dataframe.loc[item]
+                item = (item[0], self._native_dataframe.columns.get_loc(item[1]))
+                native_series = self._native_dataframe.iloc[item]
             elif isinstance(item[1], int):
                 native_series = self._native_dataframe.iloc[item]
             else:  # pragma: no cover
diff --git a/narwhals/dataframe.py b/narwhals/dataframe.py
index 629bdfca..0f4d8e94 100644
--- a/narwhals/dataframe.py
+++ b/narwhals/dataframe.py
@@ -449,7 +449,13 @@ class DataFrame(BaseFrame[FrameT]):
         )
 
     @overload
-    def __getitem__(self, item: tuple[Sequence[int], str | int]) -> Series: ...  # type: ignore[overload-overlap]
+    def __getitem__(self, item: tuple[Sequence[int], Sequence[int]]) -> Self: ...
+    @overload
+    def __getitem__(self, item: tuple[Sequence[int], str]) -> Series: ...  # type: ignore[overload-overlap]
+    @overload
+    def __getitem__(self, item: tuple[Sequence[int], Sequence[str]]) -> Self: ...
+    @overload
+    def __getitem__(self, item: tuple[Sequence[int], int]) -> Series: ...  # type: ignore[overload-overlap]
 
     @overload
     def __getitem__(self, item: Sequence[int]) -> Self: ...
@@ -461,7 +467,12 @@ class DataFrame(BaseFrame[FrameT]):
     def __getitem__(self, item: slice) -> Self: ...
 
     def __getitem__(
-        self, item: str | slice | Sequence[int] | tuple[Sequence[int], str | int]
+        self,
+        item: str
+        | slice
+        | Sequence[int]
+        | tuple[Sequence[int], str | int]
+        | tuple[Sequence[int], Sequence[int] | Sequence[str]],
     ) -> Series | Self:
         """
         Extract column or slice of DataFrame.
@@ -472,9 +483,10 @@ class DataFrame(BaseFrame[FrameT]):
                 - str: extract column
                 - slice or Sequence of integers: slice rows from dataframe.
                 - tuple of Sequence of integers and str or int: slice rows and extract column at the same time.
-                  If the second element of the tuple is an integer, it is interpreted as the column index. Otherwise,
-                  it is interpreted as the column name.
+                - tuple of Sequence of integers and Sequence of integers: slice rows and extract columns at the same time.
         Notes:
+            Integers are always interpreted as positions, and strings always as column names.
+
             In contrast with Polars, pandas allows non-string column names.
             If you don't know whether the column name you're trying to extract
             is definitely a string (e.g. `df[df.columns[0]]`) then you should
@@ -508,6 +520,12 @@ class DataFrame(BaseFrame[FrameT]):
                 2
             ]
         """
+        if (
+            isinstance(item, tuple)
+            and len(item) == 2
+            and isinstance(item[1], (list, tuple))
+        ):
+            return self._from_compliant_dataframe(self._compliant_frame[item])
         if isinstance(item, str) or (isinstance(item, tuple) and len(item) == 2):
             from narwhals.series import Series
 
diff --git a/narwhals/stable/v1.py b/narwhals/stable/v1.py
index 9f68b36a..b4bc713a 100644
--- a/narwhals/stable/v1.py
+++ b/narwhals/stable/v1.py
@@ -67,7 +67,15 @@ class DataFrame(NwDataFrame[IntoDataFrameT]):
     """
 
     @overload
-    def __getitem__(self, item: tuple[Sequence[int], str | int]) -> Series: ...  # type: ignore[overload-overlap]
+    def __getitem__(self, item: tuple[Sequence[int], Sequence[int]]) -> Self: ...
+
+    @overload
+    def __getitem__(self, item: tuple[Sequence[int], str]) -> Series: ...  # type: ignore[overload-overlap]
+    @overload
+    def __getitem__(self, item: tuple[Sequence[int], Sequence[str]]) -> Self: ...
+
+    @overload
+    def __getitem__(self, item: tuple[Sequence[int], int]) -> Series: ...  # type: ignore[overload-overlap]
 
     @overload
     def __getitem__(self, item: Sequence[int]) -> Self: ...


Test patch:
diff --git a/tests/frame/slice_test.py b/tests/frame/slice_test.py
index 54c1a672..1feaaf5c 100644
--- a/tests/frame/slice_test.py
+++ b/tests/frame/slice_test.py
@@ -74,6 +74,10 @@ def test_gather_pandas_index() -> None:
     expected = {"a": [1, 2], "b": [4, 2]}
     compare_dicts(result, expected)
 
+    result = nw.from_native(df, eager_only=True)[[1, 2], "a"].to_frame()
+    expected = {"a": [1, 2]}
+    compare_dicts(result, expected)
+
 
 def test_gather_rows_cols(constructor: Any) -> None:
     native_df = constructor(data)
@@ -93,3 +97,19 @@ def test_gather_rows_cols(constructor: Any) -> None:
         assert_series_equal(result, expected, check_dtype=False)
         result = df[np.array([0, 3, 1]), "b"].to_pandas()
         assert_series_equal(result, expected, check_dtype=False)
+
+
+def test_slice_both_tuples_of_ints(constructor: Any) -> None:
+    data = {"a": [1, 2, 3], "b": [4, 5, 6], "c": [7, 8, 9]}
+    df = nw.from_native(constructor(data), eager_only=True)
+    result = df[[0, 1], [0, 2]]
+    expected = {"a": [1, 2], "c": [7, 8]}
+    compare_dicts(result, expected)
+
+
+def test_slice_int_rows_str_columns(constructor: Any) -> None:
+    data = {"a": [1, 2, 3], "b": [4, 5, 6], "c": [7, 8, 9]}
+    df = nw.from_native(constructor(data), eager_only=True)
+    result = df[[0, 1], ["a", "c"]]
+    expected = {"a": [1, 2], "c": [7, 8]}
+    compare_dicts(result, expected)