Polars中能否按特定条件重置cumsum()？无需逐行循环

当然可以不用逐行循环实现这个需求，Polars的矢量化分组累加就能高效完成。核心思路是先生成一个分组标识，将每次a=True的位置作为新分组的起点，再在每个分组内对b执行累加求和。

实现代码 #

import polars as pl

df = pl.from_repr("""
┌───────┬─────┐
│ a     ┆ b   │
│ ---   ┆ --- │
│ bool  ┆ i64 │
╞═══════╪═════╡
│ false ┆ 1   │
│ false ┆ 2   │
│ true  ┆ 3   │
│ false ┆ 4   │
└───────┴─────┘
""")

# 生成分组标识 + 分组内累加
result = df.with_columns(
    # 分组键：每次a为True时，分组编号递增（True会被Polars视为1，False视为0）
    group_id=pl.col("a").cumsum(),
    # 在每个分组内对b执行累加求和
    sum=pl.col("b").cumsum().over("group_id")
).drop("group_id")  # 移除临时分组键

print(result)

输出结果 #

┌───────┬─────┬─────┐
│ a     ┆ b   ┆ sum │
│ ---   ┆ --- ┆ --- │
│ bool  ┆ i64 ┆ i64 │
╞═══════╪═════╪═════╡
│ false ┆ 1   ┆ 1   │
│ false ┆ 2   ┆ 3   │
│ true  ┆ 3   ┆ 3   │
│ false ┆ 4   ┆ 7   │
└───────┴─────┴─────┘

原理说明 #

1. 分组键生成：pl.col("a").cumsum()会生成一个序列，示例中为[0, 0, 1, 1]。每次遇到a=True时，累加值加1，自然将数据划分为[行0-行1]、[行2-行3]两个独立分组。
2. 分组内累加：使用.over("group_id")指定在每个分组内执行cumsum，实现“重置后继续累加”的效果。

这种方法完全基于Polars的矢量化操作，性能远高于逐行循环，尤其适合处理大规模数据集。

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