Polars中分组滚动（Rolling Per Group）的高效实现方案咨询

问题背景 #

给定以下Polars DataFrame：

import numpy as np
import polars as pl

df = pl.DataFrame({
    "date": pl.date_range(pl.date(2022, 1, 1), pl.date(2022, 1, 30), interval="1d", eager=True),
    "group": ["A", "B", "C"]*10,
    "total": np.random.randint(100, size=30)
})

需求是按group字段分组，计算每个组内date列5天滚动窗口中total列的求和（需支持多列求和）。原实现采用循环遍历每个分组的方式，代码繁琐且在百万行数据、超10000个分组的场景下内存占用过高：

原低效实现 #

df_list = []

for table_group in df["group"].unique():
    temp_df = (
       df
        .filter(pl.col("group")==table_group)
        .sort("date")
        .rolling("date", period="5d")
        .agg(pl.col("total").sum().name.prefix("roll_"))
    )
    
    temp_df2 = (
        df
        .filter(pl.col("group")==table_group)
        .sort("date")
        .select(
            pl.all(),
            temp_df.get_column("roll_total")
        )
    )
    df_list.append(temp_df2)

pl.concat(df_list)

高效优雅的解决方案 #

Polars原生支持分组滚动窗口操作，无需手动循环分组，以下两种方式均能大幅提升性能并简化代码：

方式一：group_by().rolling().agg（适合仅保留核心结果） #

通过链式调用直接完成分组、滚动、聚合全流程，内存占用和性能均经过深度优化：

# 多列求和只需在agg中添加对应列的聚合逻辑
result = (
    df
    .sort(["group", "date"])  # 确保每个组内日期有序（滚动窗口依赖顺序）
    .group_by("group", maintain_order=True)  # maintain_order保留分组原顺序
    .rolling(
        index_column="date",
        period="5d",
        closed="both"  # 可根据需求调整窗口闭合方式，默认值为'both'
    )
    .agg(
        pl.col("total").sum().alias("roll_total"),
        # 示例：添加第二列的滚动求和
        # pl.col("other_col").sum().alias("roll_other_col")
    )
)

方式二：窗口函数over结合滚动函数（适合保留原所有列） #

直接在原DataFrame上添加滚动求和列，代码最简洁，自动处理分组内的窗口计算：

result = (
    df
    .sort(["group", "date"])
    .with_columns(
        pl.col("total").rolling_sum(period="5d").over("group").alias("roll_total"),
        # 示例：多列求和直接添加
        # pl.col("other_col").rolling_sum(period="5d").over("group").alias("roll_other_col")
    )
)

方案优势 #

• 性能优异：Polars内部对分组滚动操作做了矢量化优化，内存占用远低于循环拼接方式，尤其适配大规模数据和大量分组的场景
• 代码简洁：无需手动遍历分组，链式调用逻辑清晰，多列求和只需添加对应聚合项，扩展性强
• 灵活性高：两种方式分别适配仅保留核心结果或保留原全量列的需求

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