Polars中按组添加行号：无需先创建全1列的实现方法

合并两步with_columns为一步 #

你可以直接在单个with_columns调用中完成逻辑，不需要预先创建全1列，用pl.lit(1)生成临时的全1序列即可：

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=pl.lit(1).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   │
└─────┴─────┴─────┘

更简洁的原生方案 #

Polars提供了原生的row_number()函数，专门用于生成行号。它默认从1开始计数，若要匹配原代码从0开始的结果，只需对结果减1：

df = df.with_columns(
    rn=pl.row_number().over('g') - 1
)

这个方法代码更简洁、可读性更强，是Polars官方推荐的实现方式，性能也更优。

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