Добавить количество строк в группе в полярах ⇐ Python

[Anonymous]

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

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

import numpy
import polars

df = (polars
.DataFrame(dict(
j=numpy.random.randint(10, 99, 20),
))
.with_row_index()
.select(
g=polars.col('index') // 3,
j='j'
)
.with_columns(rn=1)
.with_columns(
rn=polars.col('rn').shift().fill_null(0).cum_sum().over('g')
)
)
print(df)

g (u32)  j (i64)  rn (i32)
0        47       0
0        22       1
0        82       2
1        19       0
1        85       1
1        15       2
2        89       0
2        74       1
2        26       2
3        11       0
3        86       1
3        81       2
4        16       0
4        35       1
4        60       2
5        30       0
5        28       1
5        94       2
6        21       0
6        38       1
shape: (20, 3)

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

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

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

чтобы:

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

  .with_columns(rn=1)

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

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