如何在Polars的DataFrame分组窗口内添加行号？

要实现分组窗口内的行号计数（每个组内从1开始递增），可以利用Polars的窗口函数来完成，以下是具体实现方法：

方法1：使用pl.row_number()窗口函数 #

这是最直接的方式，row_number()会在指定的分组窗口内自动生成从1开始的递增序号：

import polars as pl

# 构造原始DataFrame
df = pl.DataFrame([
    {'groupings': 'a', 'target_count_over_windows': 1},
    {'groupings': 'a', 'target_count_over_windows': 2},
    {'groupings': 'a', 'target_count_over_windows': 3},
    {'groupings': 'b', 'target_count_over_windows': 1},
    {'groupings': 'c', 'target_count_over_windows': 1},
    {'groupings': 'c', 'target_count_over_windows': 2},
    {'groupings': 'd', 'target_count_over_windows': 1},
    {'groupings': 'd', 'target_count_over_windows': 2},
    {'groupings': 'd', 'target_count_over_windows': 3}
])

# 添加分组内的行号列
df_with_counter = df.with_columns(
    pl.row_number().over("groupings").alias("counter")
)

print(df_with_counter)

执行后输出：

shape: (9, 3)
┌───────────┬───────────────────────────┬─────────┐
│ groupings ┆ target_count_over_windows ┆ counter │
│ ---       ┆ ---                       ┆ ---     │
│ str       ┆ i64                       ┆ u32     │
╞═══════════╪═══════════════════════════╪═════════╡
│ a         ┆ 1                         ┆ 1       │
│ a         ┆ 2                         ┆ 2       │
│ a         ┆ 3                         ┆ 3       │
│ b         ┆ 1                         ┆ 1       │
│ c         ┆ 1                         ┆ 1       │
│ c         ┆ 2                         ┆ 2       │
│ d         ┆ 1                         ┆ 1       │
│ d         ┆ 2                         ┆ 2       │
│ d         ┆ 3                         ┆ 3       │
└───────────┴───────────────────────────┴─────────┘

方法2：使用pl.int_range()（支持自定义起始值） #

如果需要自定义行号的起始值（比如从0开始），可以用int_range()结合窗口内的计数来实现：

df_with_counter = df.with_columns(
    pl.int_range(1, pl.count().over("groupings") + 1).over("groupings").alias("counter")
)

这里pl.count().over("groupings")会计算每个分组的行数，int_range(1, count+1)就生成了从1到分组行数的连续整数序列。

关键说明 #

• over("groupings")是核心：它指定了窗口的分组维度，确保行号仅在同一个groupings值的范围内递增
• pl.row_number()默认从1开始计数，如果需要从0开始，只需改为pl.row_number() - 1即可

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