如何在Python Polars中高效计算分组滚动中位数绝对偏差(MAD)

针对1300万行、1500个分组的Polars DataFrame，要替代低效的rolling_map实现滚动MAD，核心思路是利用Polars内置的滚动中位数优化操作，规避Python层面的逐窗口计算开销，具体实现如下：

实现逻辑 #

滚动MAD的计算逻辑是「窗口中位数的绝对偏差的中位数」，可以拆分为三次矢量化操作，全程依赖Polars底层优化：

1. 计算每个分组的滚动中位数（窗口大小2016）
2. 计算每个值与对应窗口中位数的绝对偏差
3. 对绝对偏差列再次计算同窗口大小的滚动中位数，得到最终的滚动MAD

代码示例 #

import polars as pl

# 确保DataFrame按时间排序（滚动窗口依赖有序时间序列）
df_sorted = df.sort("datetime")

# 链式计算滚动MAD
result = df_sorted.with_columns(
    # 第一步：分组滚动中位数
    roll_median=pl.col("value").rolling(
        window_size=2016,
        by="group",
        index_column="datetime",
        closed="left"  # 与原rolling_std的窗口闭合规则保持一致
    ).median()
).with_columns(
    # 第二步：计算绝对偏差
    abs_dev=pl.col("value").sub(pl.col("roll_median")).abs()
).with_columns(
    # 第三步：绝对偏差的滚动中位数即MAD
    rolling_mad=pl.col("abs_dev").rolling(
        window_size=2016,
        by="group",
        index_column="datetime",
        closed="left"
    ).median()
)

性能优势说明 #

• 两次rolling.median()都是Polars底层优化的矢量化操作，采用高效的滑动窗口中位数算法，完全避免了rolling_map调用Python lambda的额外开销，性能可接近rolling_std的水平（秒级完成）。
• 保持了与原rolling_std一致的分组、窗口逻辑，结果逻辑完全对齐。

注意事项 #

• 必须先按datetime排序，Polars的滚动操作要求索引列有序。
• 根据原rolling_std的配置调整closed参数（可选left/right/both），确保窗口范围一致。
• 若需处理窗口观测数不足2016的场景，可添加min_periods参数（如min_periods=1），与原逻辑匹配。

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