Polars 中基于区间分组求列最大值的高效实现方法

2次阅读

Polars 中基于区间分组求列最大值的高效实现方法

本文介绍如何在 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 中处理“标记驱动区间聚合”的典型范式。

text=ZqhQzanResources