
本文介绍如何在 Polars 中为满足特定起止条件(如 “Life” 开始、”Death” 结束)的连续行区间,高效计算目标列(如 “column A”)的最大值,并将结果填充至起始行,其余行置空。核心在于构造语义化分组 ID 并结合 over() 窗口聚合。
本文介绍如何在 polars 中为满足特定起止条件(如 “life” 开始、”death” 结束)的连续行区间,高效计算目标列(如 “column a”)的最大值,并将结果填充至起始行,其余行置空。核心在于构造语义化分组 id 并结合 `over()` 窗口聚合。
在数据处理中,常需对逻辑上连续的“事件区间”进行聚合——例如,从某个标记(如 “Life”)开始,到下一个匹配标记(如 “Death”)结束之间的所有行构成一个有效区间,要求该区间内某数值列的最大值仅出现在起始行。Polars 本身不提供原生的“区间扫描”API,但可通过组合布尔掩码、累积运算与分组机制优雅实现,且全程保持惰性计算与高性能。
以下为完整解决方案(基于 Polars v0.20+):
import polars as pl df = pl.DataFrame({ "Column A": [2, 3, 1, 4, 1, 3, 3, 2, 1, 0], "Column B": ["Life", None, None, None, "Death", None, "Life", None, None, "Death"] }) # 步骤:构建语义分组 ID start = pl.col("Column B") == "Life" end = pl.col("Column B") == "Death" # 初步分组:利用 cum_sum() 对 Life/Death 分别计数,相加后 forward_fill 处理缺失 group_id = (start.cum_sum() + end.cum_sum()).forward_fill() # 关键修正:将每个 "Death" 行的 group_id 替换为其前一行的值(即归属到前一个 Life 区间) group_id = pl.when(end).then(group_id.shift()).otherwise(group_id) # 最终:仅在 "Life" 行注入对应区间的 Column A 最大值;其余行自动为 NULL result = df.with_columns( pl.when(start) .then(pl.col("Column A").max().over(group_id)) .alias("Column C") ) print(result)
输出结果与预期一致:
shape: (10, 3) ┌──────────┬──────────┬──────────┐ │ Column A ┆ Column B ┆ Column C │ │ --- ┆ --- ┆ --- │ │ i64 ┆ str ┆ i64 │ ╞══════════╪══════════╪══════════╡ │ 2 ┆ Life ┆ 4 │ │ 3 ┆ null ┆ null │ │ 1 ┆ null ┆ null │ │ 4 ┆ null ┆ null │ │ 1 ┆ Death ┆ null │ │ 3 ┆ null ┆ null │ │ 3 ┆ Life ┆ 3 │ │ 2 ┆ null ┆ null │ │ 1 ┆ null ┆ null │ │ 0 ┆ Death ┆ null │ └──────────┴──────────┴──────────┘
✅ 关键原理说明
- cum_sum().forward_fill():将 “Life” 和 “Death” 分别转为 0/1 序列并累加,使每个新出现的 Life 或 Death 都推动计数器递增;forward_fill() 确保中间 null 行继承最近的有效组号。
- group_id.shift() 配合 pl.when():因 “Death” 行本身属于当前 Life 区间的终点(而非下一区间起点),需将其 group_id 替换为前一行值,从而让整个 [Life, …, Death] 范围共享同一 ID。
- .max().over(group_id):Polars 的 over 实现高效分组聚合,无需显式 group_by().agg(),避免打乱原始行序。
- pl.when(…).then(…).alias():确保仅 “Life” 行被赋值,其余位置自然保留 null,语义清晰且零冗余赋值。
⚠️ 注意事项
- 若存在 “Life” 后无对应 “Death” 的情况(如末尾未闭合区间),该方案会将后续所有行归入同一组,导致最大值覆盖范围扩大。如需严格配对,建议先用 pl.scan_csv() 或自定义函数预校验 Life/Death 序列完整性。
- null 值在布尔比较中默认返回 False,因此 pl.col(“Column B”) == “Life” 自动跳过 null 行,无需额外 is_not_null() 过滤。
- 所有操作均为表达式链式调用,完全兼容 Polars 的 LazyFrame 模式,适用于大规模数据集。
该方法兼具可读性、性能与可扩展性,是 Polars 中处理“标记驱动区间聚合”的典型范式。