如何在Polars DataFrame中按item_id分组并添加计数器列？

可以用Polars的窗口函数，通过row_number()结合分组指定，为每个item_id组内的记录生成递增计数。

代码实现 #

import polars as pl

df = pl.from_repr("""
┌─────────┬─────────────────────────────┐
│ item_id ┆ num_days_after_first_review │
│ ---     ┆ ---                         │
│ i64     ┆ i64                         │
╞═════════╪═════════════════════════════╡
│ 1       ┆ 3                           │
│ 1       ┆ 3                           │
│ 1       ┆ 10                          │
│ 2       ┆ 2                           │
│ 2       ┆ 2                           │
│ 3       ┆ 1                           │
│ 3       ┆ 5                           │
└─────────┴─────────────────────────────┘
""")

# 添加计数列
result_df = df.with_columns(
    pl.row_number().over("item_id").alias("num")
)

print(result_df)

输出结果 #

shape: (7, 3)
┌─────────┬─────────────────────────────┬─────┐
│ item_id ┆ num_days_after_first_review ┆ num │
│ ---     ┆ ---                         ┆ --- │
│ i64     ┆ i64                         ┆ i64 │
╞═════════╪═════════════════════════════╪═════╡
│ 1       ┆ 3                           ┆ 1   │
│ 1       ┆ 3                           ┆ 2   │
│ 1       ┆ 10                          ┆ 3   │
│ 2       ┆ 2                           ┆ 1   │
│ 2       ┆ 2                           ┆ 2   │
│ 3       ┆ 1                           ┆ 1   │
│ 3       ┆ 5                           ┆ 2   │
└─────────┴─────────────────────────────┴─────┘

关键说明 #

• row_number()：生成当前行在分组内的序号，从1开始递增
• over("item_id")：指定按item_id分组，序号在每个组内独立计算
• alias("num")：将生成的计数列命名为需求中的num

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