Polars变长分组聚合：求和比值与EWM均值比值实现

问题1：当前id总和与同category前一个id总和的比值 #

核心思路是先聚合每个(category, id)组的value总和，再在同category内按id的原始出现顺序偏移获取前一组的总和，最后计算比值。该方案完全兼容分组长度可变的场景。

示例代码 #

import polars as pl

# 示例输入数据
df = pl.DataFrame({
    "category": ["X", "X", "X", "X", "X", "X", "X", "Y", "Y"],
    "id": ["A", "A", "A", "A", "C", "C", "C", "B", "B"],
    "value": [1, 2, 3, 4, 8, 9, 10, 5, 6]
})

# 问题1实现
result1 = (
    df
    # 按category+id聚合，保留原始出现顺序
    .group_by(["category", "id"], maintain_order=True)
    .agg(sum_value=pl.col("value").sum())
    # 在同category内获取前一个id的总和
    .with_columns(
        prev_sum=pl.col("sum_value").shift(1).over("category"),
        feature=pl.col("sum_value") / pl.col("prev_sum")
    )
    # 首组设为null
    .with_columns(feature=pl.when(pl.col("prev_sum").is_not_null()).then(pl.col("feature")))
    # 合并回原始数据（若需要每行都带feature）
    .join(df, on=["category", "id"], how="left")
)

print(result1)

关键说明 #

• maintain_order=True确保id的处理顺序与原始数据中首次出现的顺序一致
• shift(1).over("category")在每个category分组内单独偏移，获取前一个id的总和
• 最终join回原始数据可让每行都对应到所属id组的feature值

问题2：当前id总和与同category所有之前分组总和的EWM均值的比值 #

核心是利用Polars的EWM窗口函数，在每个category分组内计算当前行之前所有分组总和的指数加权移动均值，再做比值计算，同样兼容分组长度可变场景。

示例代码（以α=0.5为例，可自定义α值） #

# 问题2实现
result2 = (
    df
    .group_by(["category", "id"], maintain_order=True)
    .agg(sum_value=pl.col("value").sum())
    .with_columns(
        # 计算当前行之前所有sum_value的EWM均值（不含当前行）
        prev_ewm_mean=pl.col("sum_value")
        .ewm(alpha=0.5, adjust=True)
        .mean()
        .shift(1)
        .over("category"),
        feature=pl.col("sum_value") / pl.col("prev_ewm_mean")
    )
    # 首组设为null
    .with_columns(feature=pl.when(pl.col("prev_ewm_mean").is_not_null()).then(pl.col("feature")))
    # 合并回原始数据
    .join(df, on=["category", "id"], how="left")
)

print(result2)

关键说明 #

• ewm(alpha=0.5)中的α是平滑系数，取值范围0~1，α越大越重视近期数据
• ewm(...).mean()计算的是包含当前行的累积EWM均值，shift(1)后得到仅包含之前所有分组的均值
• 同样通过over("category")确保每个category分组内独立计算

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