Как добавить столбец диапазона к группам в полярах

Нашел ответ на проблему GitHub с полярами: используйте cum_count

df.with_columns(pl.col('c').cum_count().over("key").alias("r"))

Вы можете использовать int_range() + pl.len() в сочетании с .over()

df = pl.DataFrame({
    "key":["a","a","a","a","b","b","b","b"],
    "a":[2,4,6,8,1,None,None,4]
})

df.with_columns(r = pl.int_range(pl.len()).over("key") + 1)

shape: (8, 3)
┌─────┬──────┬─────┐
│ key ┆ a    ┆ r   │
│ --- ┆ ---  ┆ --- │
│ str ┆ i64  ┆ i64 │
╞═════╪══════╪═════╡
│ a   ┆ 2    ┆ 1   │
│ a   ┆ 4    ┆ 2   │
│ a   ┆ 6    ┆ 3   │
│ a   ┆ 8    ┆ 4   │
│ b   ┆ 1    ┆ 1   │
│ b   ┆ null ┆ 2   │
│ b   ┆ null ┆ 3   │
│ b   ┆ 4    ┆ 4   │
└─────┴──────┴─────┘

pl.int_range(pl.len()) эквивалентно pl.int_range(start=0, end=pl.len())

• end: Конец диапазона (эксклюзивный). Если установлено значение None (по умолчанию), используется значение start, а start устанавливается на 0.

pl.len() дает нам «количество строк в контексте».

>>> df.select(pl.len())
shape: (1, 1)
┌─────┐
│ len │
│ --- │
│ u32 │
╞═════╡
│ 8   │
└─────┘

В групповом контексте (т. е. .agg() или .over()) мы получаем количество строк в каждой группе.

>>> df.group_by("key").agg(pl.len())
shape: (2, 2)
┌─────┬─────┐
│ key ┆ len │
│ --- ┆ --- │
│ str ┆ u32 │
╞═════╪═════╡
│ b   ┆ 4   │
│ a   ┆ 4   │
└─────┴─────┘

.cum_count() недавно изменился способ обработки нулевых значений.

Он больше не гарантирует «уникальный счетчик», если вы используете его для данных, содержащих значения NULL.

>>> df.with_columns(r = pl.col("a").cum_count().over("key"))
shape: (8, 3)
┌─────┬──────┬─────┐
│ key ┆ a    ┆ r   │
│ --- ┆ ---  ┆ --- │
│ str ┆ i64  ┆ u32 │
╞═════╪══════╪═════╡
│ a   ┆ 2    ┆ 1   │
│ a   ┆ 4    ┆ 2   │
│ a   ┆ 6    ┆ 3   │
│ a   ┆ 8    ┆ 4   │
│ b   ┆ 1    ┆ 1   │
│ b   ┆ null ┆ 1   │ # <-
│ b   ┆ null ┆ 1   │ # <-
│ b   ┆ 4    ┆ 2   │
└─────┴──────┴─────┘