polars 滚动选项不允许？

Question

我有一个数据框，内容大概是这样的：

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]})


day type    value
i64 str i64
1   "a" 1
2   "a" 0
4   "a" 3
5   "a" 4
2   "b" 2
3   "b" 2
5   "b" 0
6   "b" 1

我想创建一个滚动求和的变量，也就是在每种不同的“类型”中，计算一个两天的窗口内的值的总和。理想情况下，最终的数据集应该是这样的：

天数	类型	数值	滚动总和
1	a	1	1
2	a	0	1
4	a	3	3
5	a	4	7
2	b	2	2
3	b	2	4
5	b	0	0
6	b	1	1

我尝试使用以下代码：

df = df.with_columns(pl.col("value")
                     .rolling(index_column="day", by="type", period="2i")
                     .sum().alias("rolling_sum"))

但是我遇到了一个错误：“TypeError: rolling() got an unexpected keyword argument 'by'”。

你能帮我解决这个问题吗？

Answer 1

这是因为在你的代码中，你试图使用 Expr.rolling()，但这个方法没有 by 参数（奇怪的是，它在文档中提到了 check_sorted 参数——难道这个功能还没实现吗？），而是应该使用 DataFrame.rolling()。

如果你把代码改成用后者，那么就能正常工作了：

(
    df.rolling(
        index_column="day", by="type", period="2i"
    )
    .agg(
        pl.col('value').sum().alias("rolling_sum")
    )
)

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

如果你想在结果中包含 value 列，可以使用 Expr.rolling_sum()，并结合 Expr.over()（假设你的 DataFrame 已经按 day 排序）：

df.with_columns(
    pl.col("value")
    .rolling_sum(window_size=2,min_periods=0)
    .over("type")
    .alias('rolling_sum')
)

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

理想情况下，我可能会期待 Expr.rolling 和 Expr.over 一起工作：

# something like this
df.with_columns(
    pl.col("value")
    .rolling(index_column="day", period="2i")
    .sum()
    .over("type")
    .alias('rolling_sum')
)

# or this
df.set_sorted(['type','day']).with_columns(
    pl.col("value")
    .sum()
    .over('type')
    .rolling(index_column="day", period="2i")
    .alias('rolling_sum')
)

但不幸的是，它并没有这样做：

InvalidOperationError: rolling expression not allowed in aggregation

更新

如果你打算根据天/周等来设置窗口，使用 rolling_sum() 可能不是你想要的。在这种情况下，你仍然可以使用 DataFrame.rolling()，并结合 Expr.last()，放在 GroupBy.agg() 中，以获取窗口中的最后一个值：

(
    df.rolling(
        index_column="day", by="type", period="2i"
    )
    .agg(
        pl.col('value').last(),
        pl.col('value').sum().alias("rolling_sum")
    )
)

┌──────┬─────┬───────┬─────────────┐
│ 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           │
└──────┴─────┴───────┴─────────────┘