如何在Polars中为分组添加组内递增整数列

你可以通过两种简洁的方式实现需求：

方法一：使用窗口函数row_number() #

利用Polars的窗口函数，按key分组后直接生成组内从1开始的递增序号：

import polars as pl

df = pl.DataFrame({'key':['a','a','a','b','b','b'],'a':[2,4,6,1,2,3]})
df = df.with_columns(
    pl.row_number().over('key').alias('r')
)
print(df)

方法二：使用cumcount()方法 #

通过分组后调用cumcount()，配合add(1)将默认从0开始的计数调整为从1开始，同时保留原始数据顺序：

import polars as pl

df = pl.DataFrame({'key':['a','a','a','b','b','b'],'a':[2,4,6,1,2,3]})
df = df.group_by('key', maintain_order=True).agg(
    pl.col('a'),
    pl.cumcount().add(1).alias('r')
).explode(['a', 'r'])
print(df)

两种方法都会得到目标结果：

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

内容的提问来源于stack exchange，提问作者figs_and_nuts