Polars中按分组计算日期滚动均值的问题求助

我之前看到过类似的问题，但照着方法试下来在我的场景里不管用，想请教大家怎么解决。

我有一个包含date、groups、prob三列的DataFrame，需求是按groups分组，计算prob列的3天滚动均值。我先按照之前看到的方法写了代码，但返回的结果全是null：

import polars as pl
from datetime import date
import numpy as np

dates = pl.date_range(date(2024, 12, 1), date(2024, 12, 30), "1d", eager=True).alias("date")
    
days = pl.concat([dates,dates])
groups = pl.concat([pl.select(pl.repeat("B", n = 30)).to_series(),
           pl.select(pl.repeat("A", n = 30)).to_series()]).alias('groups')

data = pl.DataFrame([days, groups])

data2 = data.with_columns(pl.lit(np.random.rand(data.height)).alias("prob"))
data2.with_columns(
    pl.col('prob')
    .rolling_mean(window_size = 3)
    .over('date','groups')
    .alias('rolling_mean')
)

运行结果：

shape: (60, 4)
┌────────────┬────────┬──────────┬──────────────┐
│ date       ┆ groups ┆ prob     ┆ rolling_mean │
│ ---        ┆ ---    ┆ ---      ┆ ---          │
│ date       ┆ str    ┆ f64      ┆ f64          │
╞════════════╪════════╪══════════╪══════════════╡
│ 2024-12-01 ┆ B      ┆ 0.938982 ┆ null         │
│ 2024-12-02 ┆ B      ┆ 0.103133 ┆ null         │
│ 2024-12-03 ┆ B      ┆ 0.724672 ┆ null         │
│ 2024-12-04 ┆ B      ┆ 0.495868 ┆ null         │
│ 2024-12-05 ┆ B      ┆ 0.621124 ┆ null         │
│ …          ┆ …      ┆ …        ┆ …            │
│ 2024-12-26 ┆ A      ┆ 0.762529 ┆ null         │
│ 2024-12-27 ┆ A      ┆ 0.766366 ┆ null         │
│ 2024-12-28 ┆ A      ┆ 0.272936 ┆ null         │
│ 2024-12-29 ┆ A      ┆ 0.28709  ┆ null         │
│ 2024-12-30 ┆ A      ┆ 0.403478 ┆ null         │
└────────────┴────────┴──────────┴──────────────┘

后来我在文档里找到了rolling_mean_by方法，试着替换后运行，但结果不对——它没有计算滚动均值，反而直接返回了每行的prob原值：

data2.with_columns(
    pl.col('prob')
    .rolling_mean_by(window_size = '3d', by = 'date')
    .over('groups', 'date')
    .alias('rolling_mean')
)

运行结果：

shape: (60, 4)
┌────────────┬────────┬──────────┬──────────────┐
│ date       ┆ groups ┆ prob     ┆ rolling_mean │
│ ---        ┆ ---    ┆ ---      ┆ ---          │
│ date       ┆ str    ┆ f64      ┆ f64          │
╞════════════╪════════╪══════════╪══════════════╡
│ 2024-12-01 ┆ B      ┆ 0.938982 ┆ 0.938982     │
│ 2024-12-02 ┆ B      ┆ 0.103133 ┆ 0.103133     │
│ 2024-12-03 ┆ B      ┆ 0.724672 ┆ 0.724672     │
│ 2024-12-04 ┆ B      ┆ 0.495868 ┆ 0.495868     │
│ 2024-12-05 ┆ B      ┆ 0.621124 ┆ 0.621124     │
│ …          ┆ …      ┆ …        ┆ …            │
│ 2024-12-26 ┆ A      ┆ 0.762529 ┆ 0.762529     │
│ 2024-12-27 ┆ A      ┆ 0.766366 ┆ 0.766366     │
│ 2024-12-28 ┆ A      ┆ 0.272936 ┆ 0.272936     │
│ 2024-12-29 ┆ A      ┆ 0.28709  ┆ 0.28709      │
│ 2024-12-30 ┆ A      ┆ 0.403478 ┆ 0.403478     │
└────────────┴────────┴──────────┴──────────────┘

想请教各位，我哪里写错了？应该怎么正确实现按分组计算日期滚动均值的需求？

备注：内容来源于stack exchange，提问作者AColoredReptile