How to get current index of element in polars list

Question

When evaluating list elements I would like to know and use the current index. Is there already a way of doing it?

Something like pl.element().idx() ?

import polars as pl

data = {"a": [[1,2,3],[4,5,6]]}
schema = {"a": pl.List(pl.Int8)}

df = pl.DataFrame(data, schema=schema)

df = df.with_columns(
    pl.col("a").list.eval(pl.element() * pl.element().idx())
)
# AttributeError: 'Expr' object has no attribute 'idx'

Expected result:

+-------------+
¦ a           ¦
¦ ---         ¦
¦ list[u8]    ¦
¦-------------¦
¦ [0, 2, 6]   ¦
¦ [0, 5, 12]  ¦
+-------------+

Answer 1

The index of an element within a list can easily be accessed using pl.element().int_range() as follows:

df.with_columns(
    pl.col("a").list.eval(pl.element() * pl.int_range(pl.len()))
)

Output.

shape: (2, 1)
┌────────────┐
│ a          │
│ ---        │
│ list[i64]  │
╞════════════╡
│ [0, 2, 6]  │
│ [0, 5, 12] │
└────────────┘

Answer 2

The best way (that I know of) is to make a row index, explode, use int_range with a window function to create the idx (I'm calling it j), and then put it back together with group_by/agg

(
    df
        .with_row_index('i')
        .explode('a')
        .with_columns(j=pl.int_range(pl.len()).over('i'))
        .with_columns(new=pl.col('a')*pl.col('j'))
        .group_by('i', maintain_order=True)
        .agg(pl.col('new'))
        .drop('i')
)