Polars中rolling_sum结合over分组结果不符合预期的问题问询

我在用Polars的rolling_sum结合over做计算，原本预期rolling_sum会像sum一样，只是多了时间窗口的限制，但实际结果不符合预期：

• sum搭配over(["X"])和over(["X", "Y"])时，返回的结果符合预期
• rolling_sum搭配over(["X"])和over(["X", "Y"])时，返回结果完全相同，这和我预期的不一样

复现代码 #

首先创建测试用DataFrame：

from itertools import product
import numpy as np
import polars as pl

xs = [f"X{i}" for i in range(2)]
ys = [f"Y{i}" for i in range(2)]
dates = [f"2023-09-{i+1:02d}" for i in range(10)]
all_combinations = list(product(xs, ys, dates))

df = pl.DataFrame(
    {
        "X": [x[0] for x in all_combinations],
        "Y": [x[1] for x in all_combinations],
        "date": [x[2] for x in all_combinations],
        "qty": np.random.randint(0, 11, size=len(all_combinations)).tolist(),
    }
)

df = df.with_columns(pl.col("date").str.strptime(pl.Datetime, "%Y-%m-%d")).sort(
    ["X", "Y", "date"]
)

接着添加聚合列，sum结果符合预期，但rolling_sum未区分分组：

df.with_columns(
    [
        pl.col("qty")
        .sum()
        .over(["X"])
        .alias("sum_over_X"),
        pl.col("qty")
        .sum()
        .over(["X","Y"])
        .alias("sum_over_X_Y"),
        pl.col("qty")
        .rolling_sum("3d", by="date", closed="left")
        .over(["X"])
        .alias("3d_sum_over_X"),
        pl.col("qty")
        .rolling_sum("3d", by="date", closed="left")
        .over(["X","Y"])
        .alias("3d_sum_over_X_Y"),
    ]
)

问题原因与解决方法 #

你的预期完全正确，问题出在Polars当前版本中rolling系列函数与over结合的执行逻辑：rolling_sum的by="date"参数会先对整个DataFrame按日期排序，之后再应用over的分组逻辑，导致滚动窗口计算没有在每个分组内独立进行——over(["X"])和over(["X", "Y"])的滚动计算都基于全局排序后的日期序列，自然返回相同结果。

要实现分组内独立计算滚动和，推荐用over的partition_by参数显式指定分组，并确保滚动窗口在分组内排序：

result = df.with_columns(
    pl.col("qty").sum().over(["X"]).alias("sum_over_X"),
    pl.col("qty").sum().over(["X","Y"]).alias("sum_over_X_Y"),
    # 按X分组，内部按date排序后计算滚动和
    pl.col("qty").rolling_sum("3d", by="date", closed="left").over(
        partition_by="X",
        order_by="date"
    ).alias("3d_sum_over_X"),
    # 按X,Y分组，内部按date排序后计算滚动和
    pl.col("qty").rolling_sum("3d", by="date", closed="left").over(
        partition_by=["X", "Y"],
        order_by="date"
    ).alias("3d_sum_over_X_Y"),
)

也可以通过group_by在每个分组内单独执行滚动操作，确保滚动窗口严格限定在分组内部：

# 生成所有需要的列，合并结果
result = df.with_columns(
    pl.col("qty").sum().over(["X"]).alias("sum_over_X"),
    pl.col("qty").sum().over(["X","Y"]).alias("sum_over_X_Y")
).join(
    # 按X分组计算滚动和
    df.group_by("X", maintain_order=True).agg(
        pl.col("date"),
        pl.col("qty").rolling_sum("3d", by="date", closed="left").alias("3d_sum_over_X")
    ).explode(["date", "qty", "3d_sum_over_X"]),
    on=["X", "date", "qty"]
).join(
    # 按X,Y分组计算滚动和
    df.group_by(["X", "Y"], maintain_order=True).agg(
        pl.col("date"),
        pl.col("qty").rolling_sum("3d", by="date", closed="left").alias("3d_sum_over_X_Y")
    ).explode(["date", "qty", "3d_sum_over_X_Y"]),
    on=["X", "Y", "date", "qty"]
)

两种方式都能得到符合预期的结果：3d_sum_over_X是每个X分组内所有数据的3天滚动和，3d_sum_over_X_Y是每个X-Y子分组内的3天滚动和。

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