如何在Polars中按列分组计算带最小周期的高效滚动均值？

需求说明 #

现有一份每日每小时的价格数据集，需按Group和Hour分组，计算Price字段过去30天的滚动均值（即同组同时段的过去30条数据求平均）。当前使用rolling方法无法指定min_periods参数，而直接用rolling_mean效率偏低，需最优实现方案。

示例数据集 #

import polars as pl

df = pl.from_repr("""
┌─────────┬─────────────────────┬─────────┬──────┐
│ Group   ┆ Time                ┆ Price   ┆ Hour │
│ ---     ┆ ---                 ┆ ---     ┆ ---  │
│ str     ┆ datetime[ns]        ┆ f64     ┆ i8   │
╞═════════╪═════════════════════╪═════════╪══════╡
│ Group1  ┆ 2021-01-03 00:00:00 ┆ 15.6328 ┆ 0    │
│ Group1  ┆ 2021-01-03 05:00:00 ┆ 17.1562 ┆ 5    │
│ Group1  ┆ 2021-01-03 08:00:00 ┆ 13.9062 ┆ 8    │
│ Group2  ┆ 2021-01-03 10:00:00 ┆ 18.5625 ┆ 10   │
│ Group2  ┆ 2021-01-03 00:00:00 ┆ 28.375  ┆ 0    │
│ Group2  ┆ 2021-01-03 13:00:00 ┆ 15.4219 ┆ 13   │
└─────────┴─────────────────────┴─────────┴──────┘
""")

局限的现有实现 #

当前代码无法指定最小周期限制：

df.rolling('Time', group_by=['Group', 'Hour'], period='30d').agg(
    pl.col('Price').mean().alias('PriceAvg')
)

最优实现方案 #

方案一：滚动聚合结合条件过滤（灵活可控） #

先通过滚动聚合统计数据量与均值，再通过条件判断过滤掉数据量不足的记录，实现min_periods的效果：

min_periods = 5  # 根据需求设置最小需要的样本数量

# 必须先按时间排序，滚动计算依赖有序序列
result = df.sort('Time').rolling(
    'Time',
    group_by=['Group', 'Hour'],
    period='30d'
).agg(
    pl.col('Price').count().alias('sample_count'),
    pl.col('Price').mean().alias('PriceAvg')
).with_columns(
    # 仅当样本数达标时保留均值，否则设为None
    PriceAvg=pl.when(pl.col('sample_count') >= min_periods)
               .then(pl.col('PriceAvg'))
               .otherwise(None)
).drop('sample_count')

print(result)

方案二：滚动求和+计数手动计算均值（大数据量更高效） #

对于超大规模数据集，手动计算滚动求和与计数再推导均值，性能优于直接调用rolling_mean：

min_periods = 5

result = df.sort('Time').with_columns(
    # 计算滚动窗口内的价格总和与样本数
    rolling_sum=pl.col('Price').rolling_sum(
        window_size='30d',
        by='Time',
        group_by=['Group', 'Hour']
    ),
    rolling_count=pl.col('Price').rolling_count(
        window_size='30d',
        by='Time',
        group_by=['Group', 'Hour']
    )
).with_columns(
    # 样本数达标时计算均值，否则设为None
    PriceAvg=pl.when(pl.col('rolling_count') >= min_periods)
               .then(pl.col('rolling_sum') / pl.col('rolling_count'))
               .otherwise(None)
).drop(['rolling_sum', 'rolling_count'])

print(result)

关键注意事项 #

• 所有滚动计算必须先对Time列排序，否则结果会出错。
• min_periods可根据业务需求调整，比如设为30时，只有当窗口内有至少30条同组同时段数据时才计算均值。
• 若使用Polars 0.19及以上版本，方案一的性能已足够优秀；大数据量场景优先选择方案二。

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