Média móvel agrupada em polares

Uma pergunta semelhante é feita aqui

No entanto, não pareceu funcionar no meu caso.

Tenho um dataframe com 3 colunas, data, grupos, prob. O que quero é criar uma média móvel de 3 dias dos valores da coluna prob agrupados por grupos e data. No entanto, seguindo a resposta vinculada acima, obtive todos os nulos retornados.

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")
len(dates)
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(
    rolling_mean = 
    pl.col('prob')
    .rolling_mean(window_size = 3)
    .over('date','groups')
)

"""
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         │
└────────────┴────────┴──────────┴──────────────┘
""""

Na documentação, encontrei .rolling_mean_by e tentei usá-lo, mas em vez de fazer uma média móvel, ele parece apenas retornar o valor prob para cada linha.

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

"""
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     │
└────────────┴────────┴──────────┴──────────────┘
""""