Repository: narwhals-dev/narwhals
Pull request #839 (2024-08-21 13:39:37)

feat: support `__getitem__` with slices for columns
I kinda wish we didn't have to go overboard with `__getitem__` overloads, but scikit-learn are using this, so I think we should

My general thinking on this is: make sure that Narwhals is easily usable by scikit-learn. Then, it's their decision whether or not they adopt - but we should at least be ready

<!--
# 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: good points, thanks! I think this is easier to explain with examples

Patch:
diff --git a/narwhals/_arrow/dataframe.py b/narwhals/_arrow/dataframe.py
index b5a14f5a..65a155e0 100644
--- a/narwhals/_arrow/dataframe.py
+++ b/narwhals/_arrow/dataframe.py
@@ -137,6 +137,32 @@ class ArrowDataFrame:
             )
 
         elif isinstance(item, tuple) and len(item) == 2:
+            if isinstance(item[1], slice):
+                columns = self.columns
+                if isinstance(item[1].start, str) or isinstance(item[1].stop, str):
+                    start = (
+                        columns.index(item[1].start)
+                        if item[1].start is not None
+                        else None
+                    )
+                    stop = (
+                        columns.index(item[1].stop) + 1
+                        if item[1].stop is not None
+                        else None
+                    )
+                    step = item[1].step
+                    return self._from_native_frame(
+                        self._native_frame.take(item[0]).select(columns[start:stop:step])
+                    )
+                if isinstance(item[1].start, int) or isinstance(item[1].stop, int):
+                    return self._from_native_frame(
+                        self._native_frame.take(item[0]).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
+
             from narwhals._arrow.series import ArrowSeries
 
             # PyArrow columns are always strings
diff --git a/narwhals/_pandas_like/dataframe.py b/narwhals/_pandas_like/dataframe.py
index 817d28cd..c6a51be2 100644
--- a/narwhals/_pandas_like/dataframe.py
+++ b/narwhals/_pandas_like/dataframe.py
@@ -146,6 +146,30 @@ class PandasLikeDataFrame:
             )
             raise TypeError(msg)  # pragma: no cover
 
+        elif isinstance(item, tuple) and len(item) == 2 and isinstance(item[1], slice):
+            columns = self._native_frame.columns
+            if isinstance(item[1].start, str) or isinstance(item[1].stop, str):
+                start = (
+                    columns.get_loc(item[1].start) if item[1].start is not None else None
+                )
+                stop = (
+                    columns.get_loc(item[1].stop) + 1
+                    if item[1].stop is not None
+                    else None
+                )
+                step = item[1].step
+                return self._from_native_frame(
+                    self._native_frame.iloc[item[0], slice(start, stop, step)]
+                )
+            if isinstance(item[1].start, int) or isinstance(item[1].stop, int):
+                return self._from_native_frame(
+                    self._native_frame.iloc[
+                        item[0], slice(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
+
         elif isinstance(item, tuple) and len(item) == 2:
             from narwhals._pandas_like.series import PandasLikeSeries
 
diff --git a/narwhals/_polars/dataframe.py b/narwhals/_polars/dataframe.py
index c2b771f9..c5164147 100644
--- a/narwhals/_polars/dataframe.py
+++ b/narwhals/_polars/dataframe.py
@@ -84,6 +84,31 @@ class PolarsDataFrame:
         return self._native_frame.shape  # type: ignore[no-any-return]
 
     def __getitem__(self, item: Any) -> Any:
+        if isinstance(item, tuple) and len(item) == 2 and isinstance(item[1], slice):
+            # TODO(marco): we can delete this branch after Polars==0.20.30 becomes the minimum
+            # Polars version we support
+            columns = self.columns
+            if isinstance(item[1].start, str) or isinstance(item[1].stop, str):
+                start = (
+                    columns.index(item[1].start) if item[1].start is not None else None
+                )
+                stop = (
+                    columns.index(item[1].stop) + 1 if item[1].stop is not None else None
+                )
+                step = item[1].step
+                return self._from_native_frame(
+                    self._native_frame.select(columns[start:stop:step]).__getitem__(
+                        item[0]
+                    )
+                )
+            if isinstance(item[1].start, int) or isinstance(item[1].stop, int):
+                return self._from_native_frame(
+                    self._native_frame.select(
+                        columns[item[1].start : item[1].stop : item[1].step]
+                    ).__getitem__(item[0])
+                )
+            msg = f"Expected slice of integers or strings, got: {type(item[1])}"  # pragma: no cover
+            raise TypeError(msg)  # pragma: no cover
         pl = get_polars()
         result = self._native_frame.__getitem__(item)
         if isinstance(result, pl.Series):
diff --git a/narwhals/dataframe.py b/narwhals/dataframe.py
index 4fa7e913..16120374 100644
--- a/narwhals/dataframe.py
+++ b/narwhals/dataframe.py
@@ -512,6 +512,8 @@ class DataFrame(BaseFrame[FrameT]):
             level=self._level,
         )
 
+    @overload
+    def __getitem__(self, item: tuple[Sequence[int], slice]) -> Self: ...
     @overload
     def __getitem__(self, item: tuple[Sequence[int], Sequence[int]]) -> Self: ...
     @overload
@@ -536,20 +538,33 @@ class DataFrame(BaseFrame[FrameT]):
         | slice
         | Sequence[int]
         | tuple[Sequence[int], str | int]
-        | tuple[Sequence[int], Sequence[int] | Sequence[str]],
+        | tuple[Sequence[int], Sequence[int] | Sequence[str] | slice],
     ) -> Series | Self:
         """
         Extract column or slice of DataFrame.
 
         Arguments:
-            item: how to slice dataframe:
+            item: how to slice dataframe. What happens depends on what is passed. It's easiest
+                to explain by example. Suppose we have a Dataframe `df`:
+
+                    - `df['a']` extracts column `'a'` and returns a `Series`.
+                    - `df[0:2]` extracts the first two rows and returns a `DataFrame`.
+                    - `df[0:2, 'a']` extracts the first two rows from column `'a'` and returns
+                        a `Series`.
+                    - `df[0:2, 0]` extracts the first two rows from the first column and returns
+                        a `Series`.
+                    - `df[[0, 1], [0, 1, 2]]` extracts the first two rows and the first three columns
+                        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`
+                    - `df[:, 'a': 'c']` extracts all rows and all columns positioned between `'a'` and `'c'`
+                        _inclusive_ and returns a `DataFrame`. For example, if the columns are
+                        `'a', 'b', 'c', 'd'`, then that would extract columns `'a'`, `'b'`, and `'c'`.
 
-                - 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.
-                - 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.
+            - Integers are always interpreted as positions
+            - Strings are always interpreted 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
@@ -598,7 +613,7 @@ class DataFrame(BaseFrame[FrameT]):
         if (
             isinstance(item, tuple)
             and len(item) == 2
-            and isinstance(item[1], (list, tuple))
+            and isinstance(item[1], (list, tuple, slice))
         ):
             return self._from_compliant_dataframe(self._compliant_frame[item])
         if isinstance(item, str) or (isinstance(item, tuple) and len(item) == 2):
diff --git a/narwhals/stable/v1.py b/narwhals/stable/v1.py
index 65d42080..03aa13b8 100644
--- a/narwhals/stable/v1.py
+++ b/narwhals/stable/v1.py
@@ -68,6 +68,8 @@ class DataFrame(NwDataFrame[IntoDataFrameT]):
     `narwhals.from_native`.
     """
 
+    @overload
+    def __getitem__(self, item: tuple[Sequence[int], slice]) -> Self: ...
     @overload
     def __getitem__(self, item: tuple[Sequence[int], Sequence[int]]) -> Self: ...
 


Test patch:
diff --git a/tests/expr_and_series/dt/to_string_test.py b/tests/expr_and_series/dt/to_string_test.py
index 8cd7ae4c..7cbbf72f 100644
--- a/tests/expr_and_series/dt/to_string_test.py
+++ b/tests/expr_and_series/dt/to_string_test.py
@@ -105,16 +105,16 @@ def test_dt_to_string_iso_local_datetime_series(
     result = (
         nw.from_native(df, eager_only=True)["a"]
         .dt.to_string("%Y-%m-%dT%H:%M:%S.%f")
-        .to_list()[0]
+        .item(0)
     )
-    assert _clean_string(result) == _clean_string(expected)
+    assert _clean_string(str(result)) == _clean_string(expected)
 
     result = (
         nw.from_native(df, eager_only=True)["a"]
         .dt.to_string("%Y-%m-%dT%H:%M:%S%.f")
-        .to_list()[0]
+        .item(0)
     )
-    assert _clean_string(result) == _clean_string(expected)
+    assert _clean_string(str(result)) == _clean_string(expected)
 
 
 @pytest.mark.parametrize(
@@ -156,10 +156,8 @@ def test_dt_to_string_iso_local_date_series(
     constructor_eager: Any, data: datetime, expected: str
 ) -> None:
     df = constructor_eager({"a": [data]})
-    result = (
-        nw.from_native(df, eager_only=True)["a"].dt.to_string("%Y-%m-%d").to_list()[0]
-    )
-    assert result == expected
+    result = nw.from_native(df, eager_only=True)["a"].dt.to_string("%Y-%m-%d").item(0)
+    assert str(result) == expected
 
 
 @pytest.mark.parametrize(
diff --git a/tests/frame/slice_test.py b/tests/frame/slice_test.py
index 4fa5d86b..222717d1 100644
--- a/tests/frame/slice_test.py
+++ b/tests/frame/slice_test.py
@@ -5,10 +5,8 @@ import pandas as pd
 import polars as pl
 import pyarrow as pa
 import pytest
-from pandas.testing import assert_series_equal
 
 import narwhals.stable.v1 as nw
-from narwhals.utils import parse_version
 from tests.utils import compare_dicts
 
 data = {
@@ -83,21 +81,14 @@ def test_gather_pandas_index() -> None:
 def test_gather_rows_cols(constructor_eager: Any) -> None:
     native_df = constructor_eager(data)
     df = nw.from_native(native_df, eager_only=True)
-    is_pandas_wo_pyarrow = parse_version(pd.__version__) < parse_version("1.0.0")
-    if isinstance(native_df, pa.Table) or is_pandas_wo_pyarrow:
-        # PyArrowSeries do not have `to_pandas`
-        result = df[[0, 3, 1], 1].to_numpy()
-        expected = np.array([11, 14, 12])
-        assert np.array_equal(result, expected)
-        result = df[np.array([0, 3, 1]), "b"].to_numpy()
-        assert np.array_equal(result, expected)
-    else:
-        result = df[[0, 3, 1], 1].to_pandas()
-        expected_index = range(3) if isinstance(native_df, pl.DataFrame) else [0, 3, 1]
-        expected = pd.Series([11, 14, 12], name="b", index=expected_index)
-        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)
+
+    expected = {"b": [11, 14, 12]}
+
+    result = {"b": df[[0, 3, 1], 1]}
+    compare_dicts(result, expected)
+
+    result = {"b": df[np.array([0, 3, 1]), "b"]}
+    compare_dicts(result, expected)
 
 
 def test_slice_both_tuples_of_ints(constructor_eager: Any) -> None:
@@ -116,6 +107,35 @@ def test_slice_int_rows_str_columns(constructor_eager: Any) -> None:
     compare_dicts(result, expected)
 
 
+def test_slice_slice_columns(constructor_eager: Any) -> None:
+    data = {"a": [1, 2, 3], "b": [4, 5, 6], "c": [7, 8, 9], "d": [1, 4, 2]}
+    df = nw.from_native(constructor_eager(data), eager_only=True)
+    result = df[[0, 1], "b":"c"]  # type: ignore[misc]
+    expected = {"b": [4, 5], "c": [7, 8]}
+    compare_dicts(result, expected)
+    result = df[[0, 1], :"c"]  # type: ignore[misc]
+    expected = {"a": [1, 2], "b": [4, 5], "c": [7, 8]}
+    compare_dicts(result, expected)
+    result = df[[0, 1], "a":"d":2]  # type: ignore[misc]
+    expected = {"a": [1, 2], "c": [7, 8]}
+    compare_dicts(result, expected)
+    result = df[[0, 1], "b":]  # type: ignore[misc]
+    expected = {"b": [4, 5], "c": [7, 8], "d": [1, 4]}
+    compare_dicts(result, expected)
+    result = df[[0, 1], 1:3]
+    expected = {"b": [4, 5], "c": [7, 8]}
+    compare_dicts(result, expected)
+    result = df[[0, 1], :3]
+    expected = {"a": [1, 2], "b": [4, 5], "c": [7, 8]}
+    compare_dicts(result, expected)
+    result = df[[0, 1], 0:4:2]
+    expected = {"a": [1, 2], "c": [7, 8]}
+    compare_dicts(result, expected)
+    result = df[[0, 1], 1:]
+    expected = {"b": [4, 5], "c": [7, 8], "d": [1, 4]}
+    compare_dicts(result, expected)
+
+
 def test_slice_invalid(constructor_eager: Any) -> None:
     data = {"a": [1, 2], "b": [4, 5]}
     df = nw.from_native(constructor_eager(data), eager_only=True)
diff --git a/tests/series_only/to_list_test.py b/tests/series_only/to_list_test.py
index adb0ceb1..10e41591 100644
--- a/tests/series_only/to_list_test.py
+++ b/tests/series_only/to_list_test.py
@@ -12,4 +12,4 @@ def test_to_list(constructor_eager: Any, request: Any) -> None:
     if "cudf" in str(constructor_eager):  # pragma: no cover
         request.applymarker(pytest.mark.xfail)
     s = nw.from_native(constructor_eager({"a": data}), eager_only=True)["a"]
-    compare_dicts({"a": s}, {"a": [1, 2, 3]})
+    compare_dicts({"a": s.to_list()}, {"a": [1, 2, 3]})