Polars高效计算滚动均值：特定窗口内Top5值均值求解

假设你的初始DataFrame中day列为字符串类型，第一步需要先将其转换为日期类型，确保时间计算准确：

import polars as pl

# 模拟多日期场景的示例数据
df = pl.DataFrame({
    "day": ["2023-05-01", "2023-04-15", "2023-04-20", "2023-03-10", "2023-05-01", "2023-04-25"],
    "value": [1, 5, 3, 8, 2, 6]
})

# 转换日期列
df = df.with_columns(pl.col("day").str.to_date())

需求1：计算每个日期前一个月内Top5最高值的均值 #

可以通过窗口过滤+TopK聚合实现，完全避免map_elements，利用Polars的向量化操作保证效率：

result_last_month = df.sort("day").with_columns(
    # 过滤当前日期前30天内的所有value（不包含当前日期，如需包含可修改closed参数为"both"）
    pl.col("value")
    .filter(pl.col("day").is_between(pl.col("day") - pl.duration(days=30), pl.col("day"), closed="left"))
    .top_k(5)  # 取窗口内前5高的值
    .mean()  # 计算均值
    .over(pl.col("day"))  # 按每个日期分组计算
    .alias("top5_mean_last_month")
).unique("day")  # 保留每个唯一日期的结果

print(result_last_month)

关键说明： #

• sort("day")确保窗口计算的时间顺序正确
• is_between精准控制时间范围，pl.duration(days=30)可替换为pl.duration(months=1)适配自然月需求
• over(pl.col("day"))为每个日期单独计算对应窗口的统计值
• unique("day")去重，得到每个唯一日期的结果

需求2：计算每个日期之前所有时段的Top5最高值的均值 #

只需将时间范围扩展为当前日期之前的所有数据，利用无边界窗口实现：

result_all_time = df.sort("day").with_columns(
    pl.col("value")
    .top_k(5)  # 取当前日期及之前所有数据中的前5高值
    .mean()
    .over(pl.col("day"), window=(None, 0))  # 窗口范围：从数据集开头到当前日期（不含当前行可调整为(None, -1)）
    .alias("top5_mean_all_time")
).unique("day")

print(result_all_time)

关键说明： #

• window=(None, 0)表示窗口包含从第一行到当前行的所有数据
• 若需要排除当前日期的数据，可将窗口参数改为(None, -1)
• 同样依赖sort("day")保证时间顺序的正确性

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