本文介绍如何在 polars 中高效实现「按分组、仅当某列满足条件时才参与累积求和」的操作,避免冗余中间列,利用布尔转整数特性与链式表达式达成简洁、高性能的一行解法。
本文介绍如何在 polars 中高效实现「按分组、仅当某列满足条件时才参与累积求和」的操作,避免冗余中间列,利用布尔转整数特性与链式表达式达成简洁、高性能的一行解法。
在 Polars 数据处理中,常需对分组数据执行带条件的累积计算——例如:仅当 days > 2 时,才将对应 amount 纳入该 group 的累计和。初学者易采用“先标记、再累加”的两步法(即创建临时列如 3+days_amount),虽逻辑清晰,但引入不必要的中间列、降低表达力,也违背 Polars 表达式链式编程的设计哲学。
更优雅、更符合 Polars 风格的解法是:直接将布尔条件转化为数值掩码,并与目标列相乘,再对结果执行 cum_sum().over("group")。这是因为 Polars 中布尔值(True/False)在算术上下文中会自动转换为 1/0,从而天然实现“条件筛选 + 值保留/清零”的效果。
以下为完整示例:
import polars as pl
df = pl.DataFrame({
"days": [0, 1, 2, 3, 4, 0, 1, 2, 3, 4, 6, 7, 1],
"amount": [100, 200, 150, 300, 250, 180, 220, 280, 210, 320, 21, 456, 111],
"group": ["A", "B", "A", "C", "B", "A", "C", "B", "C", "A", "C", "B", "B"]
})
# ✅ 推荐写法:单次 with_columns,无中间列
result = df.with_columns(
(pl.col("amount") * (pl.col("days") > 2))
.cum_sum()
.over("group")
.alias("group_cond_cumsum")
)
print(result)输出中 group_cond_cumsum 列将严格满足:
- 同一组内,仅 days > 2 的行贡献其 amount 值;
- 累积和按 group 分区独立计算;
- 不满足条件的行对应位置为前序有效值的延续(即“保持上一个累计值”,而非重置)。
⚠️ 注意事项:
- (pl.col("days") > 2) 返回的是布尔 Series,在与数值列相乘时自动广播为 0/1,无需显式 .cast(pl.Int64)(Polars 会隐式处理);
- 若需「跳过不满足条件的行、使累积索引连续」(即类似 filter → cum_sum → join 的语义),则当前方法不适用——此时应改用 map_groups 或结构化分组聚合,但性能显著下降,绝大多数场景无需此复杂度;
- 该模式可轻松泛化:替换条件(如 pl.col("status") == "active")、更换聚合函数(如 .cum_prod().over("group"))、叠加多条件(如 (cond1 & cond2) * pl.col("value"))。
总结而言,善用布尔值的数值语义与 Polars 表达式的惰性求值特性,不仅能写出更紧凑、可读性更强的代码,还能让查询优化器有更大空间进行向量化加速。告别“胶带式”临时列,拥抱原生、声明式的条件累积表达。










