Repository: narwhals-dev/narwhals
Pull request #923 (2024-09-07 21:17:45)

fix: Fixes arrow support for df[:, list[int|str]]
<!--
# 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
- [x] 🐛 Bug Fix
- [ ] 🔧 Optimization
- [ ] 📝 Documentation
- [x] ✅ Test
- [ ] 🐳 Other

## Related issues 

- Related issue https://github.com/narwhals-dev/narwhals/issues/914 

## Checklist

- [x] Code follows style guide (ruff)
- [x] Tests added 
- [x] Documented the changes

## If you have comments or can explain your changes, please do so below.
This PR adds new tests to check `df[:, list[int, str]` and adjusts `pa.Table` to pass the new test.



Patch:
diff --git a/docs/why.md b/docs/why.md
index adf8f39b..4ec605d1 100644
--- a/docs/why.md
+++ b/docs/why.md
@@ -27,7 +27,7 @@ pl_df_right = pl.DataFrame({"a": [1, 2, 3], "c": [4, 5, 6]})
 pl_left_merge = pl_df_left.join(pl_df_right, left_on="b", right_on="c", how="left")
 
 print(pd_left_merge.columns)
-print(pl_df_right.columns)
+print(pl_left_merge.columns)
 ```
 
 There are several such subtle difference between the libraries. Writing dataframe-agnostic code is hard!
diff --git a/narwhals/_arrow/dataframe.py b/narwhals/_arrow/dataframe.py
index 960d833a..064903d7 100644
--- a/narwhals/_arrow/dataframe.py
+++ b/narwhals/_arrow/dataframe.py
@@ -136,9 +136,12 @@ class ArrowDataFrame:
             and len(item) == 2
             and isinstance(item[1], (list, tuple))
         ):
-            return self._from_native_frame(
-                self._native_frame.take(item[0]).select(item[1])
-            )
+            if item[0] == slice(None):
+                selected_rows = self._native_frame
+            else:
+                selected_rows = self._native_frame.take(item[0])
+
+            return self._from_native_frame(selected_rows.select(item[1]))
 
         elif isinstance(item, tuple) and len(item) == 2:
             if isinstance(item[1], slice):
diff --git a/narwhals/dataframe.py b/narwhals/dataframe.py
index 165b6598..f6fccb27 100644
--- a/narwhals/dataframe.py
+++ b/narwhals/dataframe.py
@@ -585,10 +585,14 @@ class DataFrame(BaseFrame[FrameT]):
     @overload
     def __getitem__(self, item: tuple[Sequence[int], Sequence[int]]) -> Self: ...
     @overload
+    def __getitem__(self, item: tuple[slice, 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[slice, Sequence[str]]) -> Self: ...
+    @overload
     def __getitem__(self, item: tuple[Sequence[int], int]) -> Series: ...  # type: ignore[overload-overlap]
 
     @overload
@@ -606,7 +610,7 @@ class DataFrame(BaseFrame[FrameT]):
         | slice
         | Sequence[int]
         | tuple[Sequence[int], str | int]
-        | tuple[Sequence[int], Sequence[int] | Sequence[str] | slice],
+        | tuple[slice | Sequence[int], Sequence[int] | Sequence[str] | slice],
     ) -> Series | Self:
         """
         Extract column or slice of DataFrame.
@@ -623,6 +627,10 @@ class DataFrame(BaseFrame[FrameT]):
                     a `Series`.
                 - `df[[0, 1], [0, 1, 2]]` extracts the first two rows and the first three columns
                     and returns a `DataFrame`
+                - `df[:, [0, 1, 2]]` extracts all rows from the first three columns and returns a
+                  `DataFrame`.
+                - `df[:, ['a', 'c']]` extracts all rows and columns `'a'` and `'c'` and returns a
+                  `DataFrame`.
                 - `df[0: 2, ['a', 'c']]` extracts the first two rows and columns `'a'` and `'c'` and
                     returns a `DataFrame`
                 - `df[:, 0: 2]` extracts all rows from the first two columns and returns a `DataFrame`
diff --git a/narwhals/stable/v1.py b/narwhals/stable/v1.py
index 98ee12a7..1af7a26f 100644
--- a/narwhals/stable/v1.py
+++ b/narwhals/stable/v1.py
@@ -76,11 +76,15 @@ class DataFrame(NwDataFrame[IntoDataFrameT]):
     def __getitem__(self, item: tuple[Sequence[int], slice]) -> Self: ...
     @overload
     def __getitem__(self, item: tuple[Sequence[int], Sequence[int]]) -> Self: ...
+    @overload
+    def __getitem__(self, item: tuple[slice, 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[slice, Sequence[str]]) -> Self: ...
 
     @overload
     def __getitem__(self, item: tuple[Sequence[int], int]) -> Series: ...  # type: ignore[overload-overlap]


Test patch:
diff --git a/tests/frame/slice_test.py b/tests/frame/slice_test.py
index eea94d44..18b05bf3 100644
--- a/tests/frame/slice_test.py
+++ b/tests/frame/slice_test.py
@@ -141,6 +141,12 @@ def test_slice_slice_columns(constructor_eager: Any) -> None:
     result = df[[0, 1], 1:]
     expected = {"b": [4, 5], "c": [7, 8], "d": [1, 4]}
     compare_dicts(result, expected)
+    result = df[:, ["b", "d"]]
+    expected = {"b": [4, 5, 6], "d": [1, 4, 2]}
+    compare_dicts(result, expected)
+    result = df[:, [0, 2]]
+    expected = {"a": [1, 2, 3], "c": [7, 8, 9]}
+    compare_dicts(result, expected)
 
 
 def test_slice_invalid(constructor_eager: Any) -> None: