Can I use indexing [start:end] in expressions instead of offset and length in 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.