Добавить номер строки/индекс для каждой группы в полярах ⇐ Python

[Anonymous]

Есть ли способ переписать это:

Код: Выделить всё

import polars as pl

df = pl.DataFrame({
'g': [0, 0, 0, 1, 1, 1, 2, 2, 2, 3],
'j': [37, 87, 56, 28, 49, 65, 37, 91, 45, 95]
})

df = (
df.with_columns(rn=1)
.with_columns(
rn=pl.col('rn').shift().fill_null(0).cum_sum().over('g')
)
)

Код: Выделить всё

┌─────┬─────┬─────┐
│ g   ┆ j   ┆ rn  │
│ --- ┆ --- ┆ --- │
│ u32 ┆ i64 ┆ i32 │
╞═════╪═════╪═════╡
│ 0   ┆ 37  ┆ 0   │
│ 0   ┆ 87  ┆ 1   │
│ 0   ┆ 56  ┆ 2   │
│ 1   ┆ 28  ┆ 0   │
│ 1   ┆ 49  ┆ 1   │
│ 1   ┆ 65  ┆ 2   │
│ 2   ┆ 37  ┆ 0   │
│ 2   ┆ 91  ┆ 1   │
│ 2   ┆ 45  ┆ 2   │
│ 3   ┆ 95  ┆ 0   │
└─────┴─────┴─────┘

чтобы он добавлял столбец rn, не требуя сначала добавить столбец, полный единиц? Т.е. как-нибудь перепишите эту часть:

Код: Выделить всё

  .with_columns(rn=1)
.with_columns(
rn=pl.col('rn').shift().fill_null(0).cum_sum().over('g')
)

чтобы:

Код: Выделить всё

  .with_columns(rn=1)

не требуется? По сути, сократите два выражения до одного.
Или любой другой/лучший способ добавить номер строки для каждой группы?

Подробнее здесь: https://stackoverflow.com/questions/768 ... -in-polars