Polars中group_by_dynamic前瞻索引问题及向后聚合实现咨询

设计意图说明 #

Polars的group_by_dynamic采用前瞻式窗口设计是有意为之的，它和rolling_mean的回顾式逻辑差异源于定位场景不同：

• group_by_dynamic最初面向时间序列的前瞻分析场景（比如基于当前及未来N个时间点的数据做预测准备），窗口范围默认是[当前索引, 当前索引+period)；
• rolling_mean则聚焦于回顾式统计（比如计算历史N个周期的均值），窗口范围默认是[当前索引-period+1, 当前索引]。

实现回顾式group_by_dynamic的方法 #

要实现基于低索引（回顾式）的group_by_dynamic，可以通过反转索引方向的方式转换窗口逻辑，具体代码如下：

import polars as pl

df = pl.DataFrame(
    {
        "index": [0, 0, 1, 1],
        "group": ["banana", "pear", "banana", "pear"],
        "weight": [2, 3, 5, 7],
    }
)

# 反转索引将前瞻窗口转为回顾窗口，聚合后恢复原索引
agg = (
    df.with_columns(reversed_index=-pl.col("index"))
    .group_by_dynamic(
        "reversed_index",
        group_by="group",
        every="1i",
        period="2i"
    )
    .agg(pl.col("weight"))
    .with_columns(index=-pl.col("reversed_index"))
    .drop("reversed_index")
    .sort(["group", "index"])
)

# 验证预期结果
assert((
    agg
    .filter(pl.col("index") == 0, pl.col("group") == "banana")
    .select("weight")
    .to_series()
    .to_list()
) == [[2]])

# 查看完整聚合结果
print(agg)

运行后输出的回顾式窗口结果：

shape: (4, 3)
┌────────┬────────┬─────────┐
│ group  ┆ index  ┆ weight  │
│ ---    ┆ ---    ┆ ---     │
│ str    ┆ i64    ┆ list[i64]│
╞════════╪════════╪═════════╡
│ banana ┆ 0      ┆ [2]     │
│ banana ┆ 1      ┆ [2, 5]  │
│ pear   ┆ 0      ┆ [3]     │
│ pear   ┆ 1      ┆ [3, 7]  │
└────────┴────────┴─────────┘

这种方法通过反转索引将group_by_dynamic的前瞻窗口逻辑间接转换为回顾式，既符合API设计逻辑，又能精准实现需求。

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