Polars中rolling方法不支持by参数？求分组滚动求和解决方案

错误原因 #

你遇到的TypeError是因为Polars的rolling()方法不支持by参数，无法直接通过该参数指定分组列。要实现不同type分组内的滚动求和，需要先通过group_by()完成分组，再在分组内执行滚动计算。

正确代码 #

import polars as pl

df = pl.LazyFrame({"day": [1,2,4,5,2,3,5,6], 'type': ['a', 'a', 'a', 'a', 'b', 'b', 'b', 'b'], "value": [1, 0, 3, 4, 2, 2, 0, 1]})

# 分组后执行滚动求和
result = df.group_by("type", maintain_order=True).agg(
    pl.col("day"),
    pl.col("value"),
    pl.col("value")
      .rolling(index_column="day", period="2i")
      .sum()
      .alias("rolling_sum")
).explode(["day", "value", "rolling_sum"])

# 查看结果
print(result.collect())

代码说明 #

1. 分组：使用group_by("type", maintain_order=True)按type分组，maintain_order=True保证结果顺序和原数据一致。
2. 滚动计算：在分组内对value列调用rolling()，指定index_column="day"以整数天作为窗口索引，period="2i"表示窗口覆盖当前行及前2个单位（天）内的行。
3. 展开结果：agg()会生成嵌套结构，用explode()将嵌套列展开为原数据的行结构。

输出结果 #

执行后会得到你期望的结果：

shape: (8, 4)
┌──────┬─────┬───────┬─────────────┐
│ type ┆ day ┆ value ┆ rolling_sum │
│ ---  ┆ --- ┆ ---   ┆ ---         │
│ str  ┆ i64 ┆ i64   ┆ i64         │
╞══════╪═════╪═══════╪═════════════╡
│ a    ┆ 1   ┆ 1     ┆ 1           │
│ a    ┆ 2   ┆ 0     ┆ 1           │
│ a    ┆ 4   ┆ 3     ┆ 3           │
│ a    ┆ 5   ┆ 4     ┆ 7           │
│ b    ┆ 2   ┆ 2     ┆ 2           │
│ b    ┆ 3   ┆ 2     ┆ 4           │
│ b    ┆ 5   ┆ 0     ┆ 0           │
│ b    ┆ 6   ┆ 1     ┆ 1           │
└──────┴─────┴───────┴─────────────┘

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