From 53de4f384a2ecaa3981b769e90edb09c2c89aee9 Mon Sep 17 00:00:00 2001 From: devdanzin <74280297+devdanzin@users.noreply.github.com> Date: Sat, 7 Dec 2024 09:21:09 -0300 Subject: [PATCH 1/7] enh: allow passing Series to Series.__getitem__. --- narwhals/_polars/series.py | 9 +++++++-- narwhals/series.py | 2 +- tests/series_only/__getitem___test.py | 17 +++++++++++++++++ 3 files changed, 25 insertions(+), 3 deletions(-) create mode 100644 tests/series_only/__getitem___test.py diff --git a/narwhals/_polars/series.py b/narwhals/_polars/series.py index 0bfea20adf..3497096c94 100644 --- a/narwhals/_polars/series.py +++ b/narwhals/_polars/series.py @@ -21,6 +21,7 @@ from narwhals._polars.dataframe import PolarsDataFrame from narwhals.dtypes import DType + from narwhals.series import Series as NWSeries from narwhals.utils import Version T = TypeVar("T") @@ -122,8 +123,12 @@ def __getitem__(self: Self, item: int) -> Any: ... @overload def __getitem__(self: Self, item: slice | Sequence[int]) -> Self: ... - def __getitem__(self: Self, item: int | slice | Sequence[int]) -> Any | Self: - return self._from_native_object(self._native_series.__getitem__(item)) + def __getitem__( + self: Self, item: int | slice | Sequence[int] | NWSeries + ) -> Any | Self: + if isinstance(item, (int, slice, Sequence)): + return self._from_native_object(self._native_series.__getitem__(item)) + return self._from_native_object(self._native_series.__getitem__(item.to_numpy())) def cast(self: Self, dtype: DType) -> Self: ser = self._native_series diff --git a/narwhals/series.py b/narwhals/series.py index 2846aebea9..f5d05bc67e 100644 --- a/narwhals/series.py +++ b/narwhals/series.py @@ -66,7 +66,7 @@ def __getitem__(self: Self, idx: int) -> Any: ... @overload def __getitem__(self: Self, idx: slice | Sequence[int]) -> Self: ... - def __getitem__(self: Self, idx: int | slice | Sequence[int]) -> Any | Self: + def __getitem__(self: Self, idx: int | slice | Sequence[int] | Series) -> Any | Self: if isinstance(idx, int): return self._compliant_series[idx] return self._from_compliant_series(self._compliant_series[idx]) diff --git a/tests/series_only/__getitem___test.py b/tests/series_only/__getitem___test.py new file mode 100644 index 0000000000..5da0f51b8f --- /dev/null +++ b/tests/series_only/__getitem___test.py @@ -0,0 +1,17 @@ +from __future__ import annotations + +import pytest +import polars as pl + +import narwhals.stable.v1 as nw + +def test_getitem() -> None: + spl = pl.Series([1, 2, 3]) + assert spl[spl[0, 1]].equals(pl.Series([2, 3])) + + snw = nw.from_native(spl, series_only=True) + assert snw[snw[0, 1]].to_native().equals(pl.Series([2, 3])) + + spl = pl.Series([1, 2, 3]) + snw = nw.from_native(spl, series_only=True) + assert pytest.raises(TypeError, lambda: snw[snw[True, False]]) \ No newline at end of file From 9f0139cadd3152425c0e84b50833ef48653d9abf Mon Sep 17 00:00:00 2001 From: devdanzin <74280297+devdanzin@users.noreply.github.com> Date: Sat, 7 Dec 2024 10:00:17 -0300 Subject: [PATCH 2/7] Address pre-commit issues. --- narwhals/_polars/series.py | 4 ++-- narwhals/series.py | 6 ++++-- tests/series_only/__getitem___test.py | 5 +++-- 3 files changed, 9 insertions(+), 6 deletions(-) diff --git a/narwhals/_polars/series.py b/narwhals/_polars/series.py index 3497096c94..13da2c1362 100644 --- a/narwhals/_polars/series.py +++ b/narwhals/_polars/series.py @@ -121,10 +121,10 @@ def dtype(self: Self) -> DType: def __getitem__(self: Self, item: int) -> Any: ... @overload - def __getitem__(self: Self, item: slice | Sequence[int]) -> Self: ... + def __getitem__(self: Self, item: slice | Sequence[int] | NWSeries[Any]) -> Self: ... def __getitem__( - self: Self, item: int | slice | Sequence[int] | NWSeries + self: Self, item: int | slice | Sequence[int] | NWSeries[Any] ) -> Any | Self: if isinstance(item, (int, slice, Sequence)): return self._from_native_object(self._native_series.__getitem__(item)) diff --git a/narwhals/series.py b/narwhals/series.py index f5d05bc67e..96cc2c009b 100644 --- a/narwhals/series.py +++ b/narwhals/series.py @@ -64,9 +64,11 @@ def __array__(self: Self, dtype: Any = None, copy: bool | None = None) -> np.nda def __getitem__(self: Self, idx: int) -> Any: ... @overload - def __getitem__(self: Self, idx: slice | Sequence[int]) -> Self: ... + def __getitem__(self: Self, idx: slice | Sequence[int] | Series[Any]) -> Self: ... - def __getitem__(self: Self, idx: int | slice | Sequence[int] | Series) -> Any | Self: + def __getitem__( + self: Self, idx: int | slice | Sequence[int] | Series[Any] + ) -> Any | Self: if isinstance(idx, int): return self._compliant_series[idx] return self._from_compliant_series(self._compliant_series[idx]) diff --git a/tests/series_only/__getitem___test.py b/tests/series_only/__getitem___test.py index 5da0f51b8f..c6dc125e23 100644 --- a/tests/series_only/__getitem___test.py +++ b/tests/series_only/__getitem___test.py @@ -1,10 +1,11 @@ from __future__ import annotations -import pytest import polars as pl +import pytest import narwhals.stable.v1 as nw + def test_getitem() -> None: spl = pl.Series([1, 2, 3]) assert spl[spl[0, 1]].equals(pl.Series([2, 3])) @@ -14,4 +15,4 @@ def test_getitem() -> None: spl = pl.Series([1, 2, 3]) snw = nw.from_native(spl, series_only=True) - assert pytest.raises(TypeError, lambda: snw[snw[True, False]]) \ No newline at end of file + assert pytest.raises(TypeError, lambda: snw[snw[True, False]]) From a6942adabb3874434cbf39a56e7f15f40b86782d Mon Sep 17 00:00:00 2001 From: FBruzzesi Date: Fri, 14 Feb 2025 19:48:37 +0100 Subject: [PATCH 3/7] docstring --- narwhals/series.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/narwhals/series.py b/narwhals/series.py index f76749b965..d9e836600a 100644 --- a/narwhals/series.py +++ b/narwhals/series.py @@ -130,8 +130,8 @@ def __getitem__( idx: The index, slice, or sequence of indices to retrieve. - If `idx` is an integer, a single element is returned. - - If `idx` is a slice or a sequence of integers, - a subset of the Series is returned. + - If `idx` is a slice, a sequence of integers, or another Series + (with integer values) a subset of the Series is returned. Returns: A single element if `idx` is an integer, else a subset of the Series. From 64c125ce0b4ec85a5ea79d7cd6778a954cb9d9bc Mon Sep 17 00:00:00 2001 From: FBruzzesi Date: Fri, 14 Feb 2025 19:51:49 +0100 Subject: [PATCH 4/7] mypy --- narwhals/series.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/narwhals/series.py b/narwhals/series.py index d9e836600a..f3889c07d2 100644 --- a/narwhals/series.py +++ b/narwhals/series.py @@ -119,10 +119,10 @@ def __array__(self: Self, dtype: Any = None, copy: bool | None = None) -> _1DArr def __getitem__(self: Self, idx: int) -> Any: ... @overload - def __getitem__(self: Self, idx: slice | Sequence[int] | Self[Any]) -> Self: ... + def __getitem__(self: Self, idx: slice | Sequence[int] | Self) -> Self: ... def __getitem__( - self: Self, idx: int | slice | Sequence[int] | Self[Any] + self: Self, idx: int | slice | Sequence[int] | Self ) -> Any | Self: """Retrieve elements from the object using integer indexing or slicing. From 1b9ff14f52d752ada5d0a32a4ca5aebd0797b6a2 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 14 Feb 2025 18:57:03 +0000 Subject: [PATCH 5/7] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- narwhals/series.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/narwhals/series.py b/narwhals/series.py index f3889c07d2..e174b492ca 100644 --- a/narwhals/series.py +++ b/narwhals/series.py @@ -121,9 +121,7 @@ def __getitem__(self: Self, idx: int) -> Any: ... @overload def __getitem__(self: Self, idx: slice | Sequence[int] | Self) -> Self: ... - def __getitem__( - self: Self, idx: int | slice | Sequence[int] | Self - ) -> Any | Self: + def __getitem__(self: Self, idx: int | slice | Sequence[int] | Self) -> Any | Self: """Retrieve elements from the object using integer indexing or slicing. Arguments: From 9914a92fa8c2155c1ac24e2ed707d709a5f1fb5b Mon Sep 17 00:00:00 2001 From: FBruzzesi Date: Fri, 14 Feb 2025 20:25:34 +0100 Subject: [PATCH 6/7] xfail modin pyarrow --- tests/series_only/__getitem___test.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/tests/series_only/__getitem___test.py b/tests/series_only/__getitem___test.py index beb790883b..5d6572a6be 100644 --- a/tests/series_only/__getitem___test.py +++ b/tests/series_only/__getitem___test.py @@ -3,6 +3,7 @@ from typing import TYPE_CHECKING import pyarrow as pa +import pytest import narwhals.stable.v1 as nw from tests.utils import assert_equal_data @@ -51,7 +52,12 @@ def test_index(constructor_eager: ConstructorEager) -> None: assert snw[snw[0]] == 0 -def test_getitem_other_series(constructor_eager: ConstructorEager) -> None: +def test_getitem_other_series( + request: pytest.FixtureRequest, constructor_eager: ConstructorEager +) -> None: + if "modin_pyarrow" in str(constructor_eager): + request.applymarker(pytest.mark.xfail) + series = nw.from_native(constructor_eager({"a": [1, None, 2, 3]}), eager_only=True)[ "a" ] From 69b460f829bd2e4c704e97f206d5648ad463f09c Mon Sep 17 00:00:00 2001 From: FBruzzesi Date: Fri, 14 Feb 2025 20:27:51 +0100 Subject: [PATCH 7/7] enough to ignore warning --- tests/series_only/__getitem___test.py | 10 ++++------ 1 file changed, 4 insertions(+), 6 deletions(-) diff --git a/tests/series_only/__getitem___test.py b/tests/series_only/__getitem___test.py index 5d6572a6be..247444f656 100644 --- a/tests/series_only/__getitem___test.py +++ b/tests/series_only/__getitem___test.py @@ -52,12 +52,10 @@ def test_index(constructor_eager: ConstructorEager) -> None: assert snw[snw[0]] == 0 -def test_getitem_other_series( - request: pytest.FixtureRequest, constructor_eager: ConstructorEager -) -> None: - if "modin_pyarrow" in str(constructor_eager): - request.applymarker(pytest.mark.xfail) - +@pytest.mark.filterwarnings( + "ignore:.*_array__ implementation doesn't accept a copy keyword.*:DeprecationWarning:modin" +) +def test_getitem_other_series(constructor_eager: ConstructorEager) -> None: series = nw.from_native(constructor_eager({"a": [1, None, 2, 3]}), eager_only=True)[ "a" ]