Polars中按组条件累计求和，是否有更优雅实现方式？

你可以将条件判断与分组累计求和的逻辑直接链式组合，无需创建中间临时列，让代码更简洁优雅：

import polars as pl

# 原始DataFrame创建
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"]
})

# 优化后的写法
df = df.with_columns(
    pl.when(pl.col("days") > 2)
      .then(pl.col("amount"))
      .otherwise(0)
      .cum_sum()
      .over("group")
      .alias("group_cumsum")
)

print(df)

输出结果 #

shape: (13, 4)
┌──────┬────────┬───────┬──────────────┐
│ days ┆ amount ┆ group ┆ group_cumsum │
│ ---  ┆ ---    ┆ ---   ┆ ---          │
│ i64  ┆ i64    ┆ str   ┆ i64          │
╞══════╪════════╪═══════╪══════════════╡
│ 0    ┆ 100    ┆ A     ┆ 0            │
│ 1    ┆ 200    ┆ B     ┆ 0            │
│ 2    ┆ 150    ┆ A     ┆ 0            │
│ 3    ┆ 300    ┆ C     ┆ 300          │
│ 4    ┆ 250    ┆ B     ┆ 250          │
│ …    ┆ …      ┆ …     ┆ …            │
│ 3    ┆ 210    ┆ C     ┆ 510          │
│ 4    ┆ 320    ┆ A     ┆ 320          │
│ 6    ┆ 21     ┆ C     ┆ 531          │
│ 7    ┆ 456    ┆ B     ┆ 706          │
│ 1    ┆ 111    ┆ B     ┆ 706          │
└──────┴────────┴───────┴──────────────┘

逻辑说明 #

Polars支持表达式的链式调用，我们直接在cum_sum().over("group")之前嵌入条件判断逻辑：

• 当days > 2时取对应的amount值，否则取0
• 对这个条件化后的序列按group分组做累计求和

如果需要保留原来的3+days_amount临时列，也可以把两个列的生成合并到同一个with_columns中，避免多次调用：

df = df.with_columns(
    pl.when(pl.col("days") > 2)
      .then(pl.col("amount"))
      .otherwise(0).alias("3+days_amount"),
    pl.when(pl.col("days") > 2)
      .then(pl.col("amount"))
      .otherwise(0)
      .cum_sum()
      .over("group")
      .alias("group_cumsum")
)

还可以用filter+fill_null替代when/then/otherwise，写法更紧凑：

df = df.with_columns(
    pl.col("amount").filter(pl.col("days") > 2).fill_null(0)
      .cum_sum().over("group")
      .alias("group_cumsum")
)

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