Skip to content

enh: allow passing Series to Series.__getitem__ #1525

Closed
@MarcoGorelli

Description

@MarcoGorelli

With Polars we can do:

In [8]: spl = pl.Series([1,2,3])

In [9]: spl[spl[0, 1]]
Out[9]:
shape: (2,)
Series: '' [i64]
[
        2
        3
]

But Narwhals gives an incomprehensible error message:

In [13]: snw = nw.from_native(spl, series_only=True)

In [14]: snw[snw[0, 1]]
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
Cell In[14], line 1
----> 1 snw[snw[0, 1]]

File ~/scratch/.venv/lib/python3.12/site-packages/narwhals/series.py:72, in Series.__getitem__(self, idx)
     70 if isinstance(idx, int):
     71     return self._compliant_series[idx]
---> 72 return self._from_compliant_series(self._compliant_series[idx])

File ~/scratch/.venv/lib/python3.12/site-packages/narwhals/_polars/series.py:126, in PolarsSeries.__getitem__(self, item)
    125 def __getitem__(self: Self, item: int | slice | Sequence[int]) -> Any | Self:
--> 126     return self._from_native_object(self._native_series.__getitem__(item))

File ~/scratch/.venv/lib/python3.12/site-packages/polars/series/series.py:1282, in Series.__getitem__(self, key)
   1254 def __getitem__(
   1255     self, key: SingleIndexSelector | MultiIndexSelector
   1256 ) -> Any | Series:
   1257     """
   1258     Get part of the Series as a new Series or scalar.
   1259
   (...)
   1280     ]
   1281     """
-> 1282     return get_series_item_by_key(self, key)

File ~/scratch/.venv/lib/python3.12/site-packages/polars/_utils/getitem.py:90, in get_series_item_by_key(s, key)
     87     return _select_elements_by_index(s, indices)
     89 msg = f"cannot select elements using key of type {type(key).__name__!r}: {key!r}"
---> 90 raise TypeError(msg)

TypeError: cannot select elements using key of type 'Series': ┌─────────────────────────────────────────┐
| Narwhals Series                         |
| Use `.to_native()` to see native output |
└─────────────────────────────────────────┘

I think we should:

  • allow this if the Series is of integral dtypes
  • raise if the Series is boolean, advising users to use Series.filter instead

Metadata

Metadata

Assignees

No one assigned

    Labels

    enhancementNew feature or requesthelp wantedExtra attention is neededhigh priorityYour PR will be reviewed very quickly if you address this

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions