如何在Polars中实现最优雅的连续计数（Streak Counter）

一、连续相同值的计数（对应你的示例需求） #

你的示例需求是按person_id分组、按order_id排序，统计每个用户连续购买相同item的次数，当item变化时重置计数。Polars可以用窗口函数+运行长度编码（RLE）分组的方式实现，代码比PySpark更紧凑：

import polars as pl

# 构造示例数据
df = pl.DataFrame({
    "person_id": ["alice", "bob", "alice", "bob", "alice", "bob", "alice", "bob"],
    "order_id": [1, 2, 3, 4, 5, 6, 7, 8],
    "item": ["A", "B", "B", "B", "A", "B", "A", "A"]
})

# 核心实现
result = (
    df
    .sort("person_id", "order_id")  # 先按用户和订单号排序，保证顺序正确
    .with_columns(
        # 生成streak分组ID：当当前item与前一个不同时，计数+1
        streak_id=pl.col("item").ne(pl.col("item").shift()).over("person_id").cumsum()
    )
    .with_columns(
        # 在每个用户+streak分组内，按顺序计数
        streak_counter=pl.col("order_id").cum_count().over(["person_id", "streak_id"])
    )
    .sort("order_id")  # 恢复原订单顺序
    .drop("streak_id")  # 移除中间分组ID
)

print(result)

输出结果完全匹配你的示例：

shape: (8, 4)
┌───────────┬──────────┬──────┬────────────────┐
│ person_id ┆ order_id ┆ item ┆ streak_counter │
│ ---       ┆ ---      ┆ ---  ┆ ---            │
│ str       ┆ i64      ┆ str  ┆ u32            │
╞═══════════╪══════════╪══════╪════════════════╡
│ alice     ┆ 1        ┆ A    ┆ 1              │
│ bob       ┆ 2        ┆ B    ┆ 1              │
│ alice     ┆ 3        ┆ B    ┆ 1              │
│ bob       ┆ 4        ┆ B    ┆ 2              │
│ alice     ┆ 5        ┆ A    ┆ 1              │
│ bob       ┆ 6        ┆ B    ┆ 3              │
│ alice     ┆ 7        ┆ A    ┆ 2              │
│ bob       ┆ 8        ┆ A    ┆ 1              │
└───────────┴──────────┴──────┴────────────────┘

二、满足指定条件的连续行数统计 #

如果需要统计满足自定义条件的连续行数（比如连续购买item='A'的次数，不满足时重置），只需调整条件判断逻辑即可：

# 示例：按person_id分组，统计连续购买item='A'的次数，非A时计数重置为0
result_cond = (
    df
    .sort("person_id", "order_id")
    .with_columns(
        # 标记是否满足条件
        is_target=pl.col("item") == "A",
        # 生成streak分组ID：当当前状态与前一个不同时，分组ID+1
        streak_id=pl.col("is_target").ne(pl.col("is_target").shift()).over("person_id").cumsum()
    )
    .with_columns(
        # 仅在满足条件的分组内计数，否则设为0
        streak_counter=pl.when(pl.col("is_target"))
                        .then(pl.col("order_id").cum_count().over(["person_id", "streak_id"]))
                        .otherwise(0)
    )
    .sort("order_id")
    .drop("is_target", "streak_id")
)

print(result_cond)

输出：

shape: (8, 4)
┌───────────┬──────────┬──────┬────────────────┐
│ person_id ┆ order_id ┆ item ┆ streak_counter │
│ ---       ┆ ---      ┆ ---  ┆ ---            │
│ str       ┆ i64      ┆ str  ┆ u32            │
╞═══════════╪══════════╪══════╪════════════════╡
│ alice     ┆ 1        ┆ A    ┆ 1              │
│ bob       ┆ 2        ┆ B    ┆ 0              │
│ alice     ┆ 3        ┆ B    ┆ 0              │
│ bob       ┆ 4        ┆ B    ┆ 0              │
│ alice     ┆ 5        ┆ A    ┆ 1              │
│ bob       ┆ 6        ┆ B    ┆ 0              │
│ alice     ┆ 7        ┆ A    ┆ 2              │
│ bob       ┆ 8        ┆ A    ┆ 1              │
└───────────┴──────────┴──────┴────────────────┘

与PySpark实现的对比 #

Polars的实现无需显式定义多个WindowSpec，通过over()直接指定分区范围，结合cumsum()生成分组ID、cum_count()完成连续计数，代码更简洁直观，同时保持了逻辑的清晰性——核心思路和PySpark一致（标记重置点→生成分组→分组内计数），但语法更紧凑。

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