如何在Polars中实现连续1的累加计数（遇0重置）？

可以利用Polars的窗口函数和累积求和实现这个需求，全程无需循环，性能适配大数据量场景，代码如下：

import polars as pl

pl.Config(tbl_rows=16)

data = [0, 1, 1, 0, 1, 1, 1, 0, 0, 1, 0, 1, 1, 1, 1, 0]

df = pl.DataFrame({"test": data})

# 生成目标结果
result_df = df.with_columns(
    # 生成分组ID：每次遇到0时分组ID递增，将连续的1划分为同一组
    group_id=pl.col("test").eq(0).cumsum(),
    # 生成辅助列：1保持1，0保持0
    ones=pl.col("test")
).with_columns(
    # 组内累加得到连续1的计数，再乘以原test列将0对应的结果置0
    result=pl.col("ones").cumsum().over("group_id") * pl.col("test")
).drop("group_id", "ones")

print(result_df)

步骤解释 #

1. 创建分组键：pl.col("test").eq(0).cumsum()会把每一个0的位置标记为分组的分界点，连续的1会被归到同一个group_id下，确保后续的累加只在连续1的区间内进行。
2. 辅助列生成：ones列直接复用原test列的值，为后续的累加提供基础数值。
3. 组内累加与修正：cumsum().over("group_id")实现每个分组内的累积求和，这样连续的1会得到1、2、3...的递增计数；最后乘以原test列，将所有原数值为0的行结果强制置为0。
4. 清理中间列：删除用于计算的group_id和ones辅助列，得到最终的结果结构。

执行后输出的结果与需求完全一致：

┌──────┬────────┐
│ test ┆ result │
│ ---  ┆ ---    │
│ i64  ┆ i64    │
╞══════╪════════╡
│ 0    ┆ 0      │
│ 1    ┆ 1      │
│ 1    ┆ 2      │
│ 0    ┆ 0      │
│ 1    ┆ 1      │
│ 1    ┆ 2      │
│ 1    ┆ 3      │
│ 0    ┆ 0      │
│ 0    ┆ 0      │
│ 1    ┆ 1      │
│ 0    ┆ 0      │
│ 1    ┆ 1      │
│ 1    ┆ 2      │
│ 1    ┆ 3      │
│ 1    ┆ 4      │
│ 0    ┆ 0      │
└──────┴────────┘

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