Polars中按连续Group分组、Date维度的累积求和优化问询

给定如下Polars DataFrame：

import polars as pl

df = pl.from_repr("""
┌────────────┬───────┬───────┐
│ Date       ┆ Group ┆ Value │
│ ---        ┆ ---   ┆ ---   │
│ date       ┆ i64   ┆ i64   │
╞════════════╪═══════╪═══════╡
│ 2020-01-01 ┆ 0     ┆ 5     │
│ 2020-01-02 ┆ 0     ┆ 8     │
│ 2020-01-03 ┆ 0     ┆ 9     │
│ 2020-01-01 ┆ 1     ┆ 5     │
│ 2020-01-02 ┆ 1     ┆ -1    │
│ 2020-01-03 ┆ 1     ┆ 2     │
│ 2020-01-01 ┆ 2     ┆ -2    │
│ 2020-01-02 ┆ 2     ┆ -1    │
│ 2020-01-03 ┆ 2     ┆ 7     │
└────────────┴───────┴───────┘
""")

需要按Date分组，以Group的顺序（0→1→2）进行连续累积求和，规则如下：

• Group 0的累积和为自身Value
• Group 1的累积和为同日期下Group 0与Group 1的Value之和
• Group 2的累积和为同日期下Group 0、Group 1与Group 2的Value之和

预期结果：

| Date       | Group | Value            |
|------------|-------|------------------|
| 2020-01-01 | 0     | 5                |
| 2020-01-02 | 0     | 8                |
| 2020-01-03 | 0     | 9                |
| 2020-01-01 | 1     | 10 (= 5 + 5)     |
| 2020-01-02 | 1     | 7  (= 8 - 1)     |
| 2020-01-03 | 1     | 11 (= 9 + 2)     |
| 2020-01-01 | 2     | 8  (= 5 + 5 - 2) |
| 2020-01-02 | 2     | 6  (= 8 - 1 - 1) |
| 2020-01-03 | 2     | 18 (= 9 + 2 + 7) |

当前实现方式（带循环）：

ddf = df.pivot(on='Group', index='Date', values='Value')

new_vals = []
for i in range(df['Group'].max() + 1):
    new_vals.extend(
        ddf.select([pl.col(f'{j}') for j in range(i+1)])
           .sum_horizontal()
           .to_list()
    )

df.with_columns(pl.Series(new_vals).alias('CumSumValue'))

请问是否存在无需循环的更优雅实现方式？

方案1：透视+横向累积求和+重塑 #

利用Polars的横向累积求和能力，再将数据重塑回原格式：

result = (
    df.pivot(on='Group', index='Date', values='Value')
      .select(pl.all().cum_sum(axis=1))
      .unpivot(index='Date', variable_name='Group', value_name='CumSumValue')
      .with_columns(pl.col('Group').cast(pl.Int64))
      .join(df, on=['Date', 'Group'], how='left')
      .select(df.columns + ['CumSumValue'])
)
print(result)

方案2：窗口函数+分组累积求和 #

通过排序确保Group顺序后，在Date分组内直接做累积求和：

result = (
    df.sort(['Date', 'Group'])
      .with_columns(
          CumSumValue=pl.col('Value').cumsum().over('Date')
      )
)

说明：因为每个Date内的Group按0、1、2顺序排列，所以Date分组内的cumsum正好匹配连续累积求和的规则。

方案3：分组转换实现 #

直接通过groupby结合窗口函数完成计算：

result = (
    df.groupby('Date')
      .with_columns(
          CumSumValue=pl.col('Value').cumsum().over(pl.col('Group').sort())
      )
)

以上方案均无需手动循环，完全利用Polars的向量化操作，代码更简洁且性能更优。

内容的提问来源于stack exchange，提问作者Syafiq Kamarul Azman