如何在Polars中实现子集列的分组累计计数？

示例数据集 #

import polars as pl

df = pl.from_repr("""
┌─────────┬──────────┬───────────────┬─────────────────┐
│ item_id ┆ store_id ┆ sold_quantity ┆ total_sku_count │
│ ---     ┆ ---      ┆ ---           ┆ ---             │
│ str     ┆ i64      ┆ i64           ┆ i64             │
╞═════════╪══════════╪═══════════════╪═════════════════╡
│ A       ┆ 2        ┆ 3             ┆ 1               │
│ B       ┆ 2        ┆ 3             ┆ 1               │
│ B       ┆ 1        ┆ 3             ┆ 1               │
│ B       ┆ 1        ┆ 4             ┆ 0               │
│ B       ┆ 2        ┆ 1             ┆ 0               │
│ B       ┆ 2        ┆ 1             ┆ 0               │
│ A       ┆ 1        ┆ 0             ┆ 1               │
│ A       ┆ 2        ┆ 4             ┆ 1               │
└─────────┴──────────┴───────────────┴─────────────────┘
""")

Pandas中的实现方式 #

在Pandas中，可通过groupby结合cumcount()实现分组内从1开始的累计计数：

subset = ["item_id", "store_id"]

df.to_pandas().groupby(subset).cumcount() + 1

# 输出结果：
# 0    1
# 1    1
# 2    1
# 3    2
# 4    2
# 5    3
# 6    1
# 7    2
# dtype: int64

尝试的Polars代码及问题 #

尝试的代码如下，但结果不符合预期（出现全局连续计数而非分组内计数）：

subset = ["item_id", "store_id"]

df.with_columns((pl.struct(subset).over(subset).cum_count()).alias("cum_counts"))

得到的错误结果：

shape: (8, 5)
┌─────────┬──────────┬───────────────┬─────────────────┬────────────┐
│ item_id ┆ store_id ┆ sold_quantity ┆ total_sku_count ┆ cum_counts │
│ ---     ┆ ---      ┆ ---           ┆ ---             ┆ ---        │
│ str     ┆ i64      ┆ i64           ┆ i64             ┆ u32        │
╞═════════╪══════════╪═══════════════╪═════════════════╪════════════╡
│ A       ┆ 2        ┆ 3             ┆ 1               ┆ 1          │
│ B       ┆ 2        ┆ 3             ┆ 1               ┆ 2          │
│ B       ┆ 1        ┆ 3             ┆ 1               ┆ 3          │
│ B       ┆ 1        ┆ 4             ┆ 0               ┆ 4          │
│ B       ┆ 2        ┆ 1             ┆ 0               ┆ 5          │
│ B       ┆ 2        ┆ 1             ┆ 0               ┆ 6          │
│ A       ┆ 1        ┆ 0             ┆ 1               ┆ 7          │
│ A       ┆ 2        ┆ 4             ┆ 1               ┆ 8          │
└─────────┴──────────┴───────────────┴─────────────────┴────────────┘

正确的Polars实现方法 #

方法1：使用cum_count()结合窗口函数 #

Polars的cum_count()默认从0开始计数，在分组窗口（over(subset)）内调用后加1，即可实现和Pandas一致的效果：

subset = ["item_id", "store_id"]

result = df.with_columns(
    pl.cum_count().over(subset).add(1).alias("cum_counts")
)
print(result)

输出结果：

shape: (8, 5)
┌─────────┬──────────┬───────────────┬─────────────────┬────────────┐
│ item_id ┆ store_id ┆ sold_quantity ┆ total_sku_count ┆ cum_counts │
│ ---     ┆ ---      ┆ ---           ┆ ---             ┆ ---        │
│ str     ┆ i64      ┆ i64           ┆ i64             ┆ u32        │
╞═════════╪══════════╪═══════════════╪═════════════════╪════════════╡
│ A       ┆ 2        ┆ 3             ┆ 1               ┆ 1          │
│ B       ┆ 2        ┆ 3             ┆ 1               ┆ 1          │
│ B       ┆ 1        ┆ 3             ┆ 1               ┆ 1          │
│ B       ┆ 1        ┆ 4             ┆ 0               ┆ 2          │
│ B       ┆ 2        ┆ 1             ┆ 0               ┆ 2          │
│ B       ┆ 2        ┆ 1             ┆ 0               ┆ 3          │
│ A       ┆ 1        ┆ 0             ┆ 1               ┆ 1          │
│ A       ┆ 2        ┆ 4             ┆ 1               ┆ 2          │
└─────────┴──────────┴───────────────┴─────────────────┴────────────┘

方法2：使用int_range()生成分组内序号 #

通过int_range()结合窗口内的count()生成从1到组大小的序列，再映射到对应分组：

subset = ["item_id", "store_id"]

result = df.with_columns(
    pl.int_range(1, pl.count() + 1).over(subset).alias("cum_counts")
)
print(result)

该方法同样能得到与Pandas一致的结果。

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