如何在Polars中实现cum_max_sum累积最大求和计算？

计算逻辑 #

每一行的cum_max_sum遵循以下规则：

• 取prev_cum_max_sum（上一行的cum_max_sum）与当前行minimum的较大值
• 将上述结果加上当前行的to_add得到最终值
• 若prev_cum_max_sum为null，则直接取minimum参与计算

对应示例中的计算：

• max(null, 1) + 5 = 6
• max(6, 4) + 6 = 12
• max(12, 13) + 7 = 20

Polars实现代码 #

首先构造示例DataFrame：

import polars as pl

df = pl.DataFrame({
    "prev_cum_max_sum": [None, 6, 12],
    "minimum": [1, 4, 13],
    "to_add": [5, 6, 7]
})

然后通过with_columns方法计算并填充cum_max_sum列：

df = df.with_columns(
    cum_max_sum=pl.max_horizontal(
        pl.col("prev_cum_max_sum").fill_null(pl.col("minimum")),
        pl.col("minimum")
    ) + pl.col("to_add")
)

执行后得到的结果与预期一致：

┌──────────────────┬─────────┬────────┬─────────────┐
│ prev_cum_max_sum ┆ minimum ┆ to_add ┆ cum_max_sum │
│ ---              ┆ ---     ┆ ---    ┆ ---         │
│ i64              ┆ i64     ┆ i64    ┆ i64         │
╞══════════════════╪═════════╪════════╪═════════════╡
│ null             ┆ 1       ┆ 5      ┆ 6           │
│ 6                ┆ 4       ┆ 6      ┆ 12          │
│ 12               ┆ 13      ┆ 7      ┆ 20          │
└──────────────────┴─────────┴────────┴─────────────┘

代码说明 #

• pl.col("prev_cum_max_sum").fill_null(pl.col("minimum"))：处理第一行prev_cum_max_sum为null的情况，将其替换为当前行的minimum
• pl.max_horizontal(...)：取处理后的prev_cum_max_sum与当前行minimum的最大值
• 最后加上to_add列的值，得到最终的cum_max_sum

内容的提问来源于stack exchange，提问作者T.H Rice