Can I use indexing [start:end] in expressions instead of offset and length in polars - python-polars

Question

In Exp.slice, the only supported syntax is exp.slice(offset,length), but in some cases, something like exp[start:end] would be more convenient and consistent. So how to write exp[1:6] for exp.slice(1,5) just like it in pandas?

Answer 1

You can add polars.internals.expr.expr.Expr.__getitem__():
import polars as pl
from polars.internals.expr.expr import Expr
def expr_slice(self, subscript):
if isinstance(subscript, slice):
if slice.start is not None:
offset = subscript.start
length = subscript.stop - offset + 1 if subscript.stop is not None else None
else:
offset = None
length = subscript.stop
return Expr.slice(self, offset, length)
else:
return Expr.slice(self, subscript, 1)
Expr.__getitem__ = expr_slice
df = pl.DataFrame({'a': range(10)})
print(df.select(pl.col('a')[2:5]))
print(df.select(pl.col('a')[3]))
shape: (4, 1)
┌─────┐
│ a │
│ --- │
│ i64 │
╞═════╡
│ 2 │
├╌╌╌╌╌┤
│ 3 │
├╌╌╌╌╌┤
│ 4 │
├╌╌╌╌╌┤
│ 5 │
└─────┘
shape: (1, 1)
┌─────┐
│ a │
│ --- │
│ i64 │
╞═════╡
│ 3 │
└─────┘
Note that my example doesn't take into account the logic for negative indexing; you can implement that yourself.
Also note that python uses 0-based indexing, so in your example exp[1:6] would raise an error.