如何在Polars中根据条件重置cum_sum（累积和）值

我有如下LazyFrame，想要构建一个新的累积和列c_sum_n：

import polars as pl
pl.Config(tbl_rows=20)
df= pl.LazyFrame({
    "a" : [0,1,2,3,4,5,6,7,8,9],
    "b" : [33,35,33,32,33,37,33,29,34,36],
    "c" : [33,35,35,35,35,37,37,37,37,37],
})
df = df.with_columns(pl.when(pl.col('b') == pl.col('c')).then(0).otherwise(1).alias('temp'))
df = df.with_columns(pl.when(pl.col('temp') == 0).then(0).otherwise(pl.col('temp').cum_sum()).alias('c_sum_n'))

当前执行后得到如下输出：

>>> df.collect()
shape: (10, 5)
┌─────┬─────┬─────┬──────┬─────────┐
│ a   ┆ b   ┆ c   ┆ temp ┆ c_sum_n │
│ --- ┆ --- ┆ --- ┆ ---  ┆ ---     │
│ i64 ┆ i64 ┆ i64 ┆ i32  ┆ i32     │
╞═════╪═════╪═════╪══════╪═════════╡
│ 0   ┆ 33  ┆ 33  ┆ 0    ┆ 0       │
│ 1   ┆ 35  ┆ 35  ┆ 0    ┆ 0       │
│ 2   ┆ 33  ┆ 35  ┆ 1    ┆ 1       │
│ 3   ┆ 32  ┆ 35  ┆ 1    ┆ 2       │
│ 4   ┆ 33  ┆ 35  ┆ 1    ┆ 3       │
│ 5   ┆ 37  ┆ 37  ┆ 0    ┆ 0       │
│ 6   ┆ 33  ┆ 37  ┆ 1    ┆ 4       │
│ 7   ┆ 29  ┆ 37  ┆ 1    ┆ 5       │
│ 8   ┆ 34  ┆ 37  ┆ 1    ┆ 6       │
│ 9   ┆ 36  ┆ 37  ┆ 1    ┆ 7       │
└─────┴─────┴─────┴──────┴─────────┘

但我期望的输出是当列b等于列c时，累积和值重置为0并重新开始累加，如下：

>>> df.collect()
shape: (10, 5)
┌─────┬─────┬─────┬──────┬─────────┐
│ a   ┆ b   ┆ c   ┆ temp ┆ c_sum_n │
│ --- ┆ --- ┆ --- ┆ ---  ┆ ---     │
│ i64 ┆ i64 ┆ i64 ┆ i32  ┆ i32     │
╞═════╪═════╪═════╪══════╪═════════╡
│ 0   ┆ 33  ┆ 33  ┆ 0    ┆ 0       │
│ 1   ┆ 35  ┆ 35  ┆ 0    ┆ 0       │
│ 2   ┆ 33  ┆ 35  ┆ 1    ┆ 1       │
│ 3   ┆ 32  ┆ 35  ┆ 1    ┆ 2       │
│ 4   ┆ 33  ┆ 35  ┆ 1    ┆ 3       │
│ 5   ┆ 37  ┆ 37  ┆ 0    ┆ 0       │
│ 6   ┆ 33  ┆ 37  ┆ 1    ┆ 1       │
│ 7   ┆ 29  ┆ 37  ┆ 1    ┆ 2       │
│ 8   ┆ 34  ┆ 37  ┆ 1    ┆ 3       │
│ 9   ┆ 36  ┆ 37  ┆ 1    ┆ 4       │
└─────┴─────┴─────┴──────┴─────────┘

寻求Python Polars或Rust Polars的解决方案。

核心思路是通过分组实现累积和的重置：每当b == c时生成新的分组，在每个分组内计算累积和，最后将temp为0的行对应值置为0。

代码实现：

import polars as pl
pl.Config(tbl_rows=20)

df= pl.LazyFrame({
    "a" : [0,1,2,3,4,5,6,7,8,9],
    "b" : [33,35,33,32,33,37,33,29,34,36],
    "c" : [33,35,35,35,35,37,37,37,37,37],
})

df = df.with_columns(
    # 生成temp列
    pl.when(pl.col('b') == pl.col('c')).then(0).otherwise(1).alias('temp'),
    # 生成分组ID：每遇到b==c的行，分组ID递增
    pl.col('b').eq(pl.col('c')).cum_sum().alias('group_id')
).with_columns(
    # 按分组计算累积和，同时将temp为0的行重置为0
    pl.when(pl.col('temp') == 0)
    .then(0)
    .otherwise(pl.col('temp').cum_sum().over('group_id'))
    .alias('c_sum_n')
).drop('group_id')  # 可选：删除中间分组列

print(df.collect())

结果验证 #

执行代码后输出与期望一致：

shape: (10, 5)
┌─────┬─────┬─────┬──────┬─────────┐
│ a   ┆ b   ┆ c   ┆ temp ┆ c_sum_n │
│ --- ┆ --- ┆ --- ┆ ---  ┆ ---     │
│ i64 ┆ i64 ┆ i64 ┆ i32  ┆ i32     │
╞═════╪═════╪═════╪══════╪═════════╡
│ 0   ┆ 33  ┆ 33  ┆ 0    ┆ 0       │
│ 1   ┆ 35  ┆ 35  ┆ 0    ┆ 0       │
│ 2   ┆ 33  ┆ 35  ┆ 1    ┆ 1       │
│ 3   ┆ 32  ┆ 35  ┆ 1    ┆ 2       │
│ 4   ┆ 33  ┆ 35  ┆ 1    ┆ 3       │
│ 5   ┆ 37  ┆ 37  ┆ 0    ┆ 0       │
│ 6   ┆ 33  ┆ 37  ┆ 1    ┆ 1       │
│ 7   ┆ 29  ┆ 37  ┆ 1    ┆ 2       │
│ 8   ┆ 34  ┆ 37  ┆ 1    ┆ 3       │
│ 9   ┆ 36  ┆ 37  ┆ 1    ┆ 4       │
└─────┴─────┴─────┴──────┴─────────┘

逻辑与Python版本一致，通过生成分组ID后计算组内累积和实现重置：

use polars::prelude::*;

fn main() -> Result<()> {
    let df = LazyFrame::new(vec![
        Series::new("a", &[0,1,2,3,4,5,6,7,8,9]),
        Series::new("b", &[33,35,33,32,33,37,33,29,34,36]),
        Series::new("c", &[33,35,35,35,35,37,37,37,37,37]),
    ])?;

    let df = df
        .with_columns([
            // 生成temp列
            col("b")
                .eq(col("c"))
                .then(lit(0))
                .otherwise(lit(1))
                .alias("temp"),
            // 生成分组ID
            col("b")
                .eq(col("c"))
                .cum_sum(false)
                .alias("group_id"),
        ])
        .with_columns([
            // 组内计算累积和并重置
            col("temp")
                .eq(lit(0))
                .then(lit(0))
                .otherwise(col("temp").cum_sum(false).over([col("group_id")]))
                .alias("c_sum_n"),
        ])
        .drop_columns(["group_id"])?;

    println!("{}", df.collect()?);
    Ok(())
}

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