如何用Polars优化多字段多周期分组的Z-score计算？

你的问题与原有代码 #

请问使用Polars计算多字段、多周期且带分组的Z-score的最优方法是什么？我编写了如下代码，想确认是否有更优实现方式：

window = "30d" # 希望支持多周期列表
df = (
    df.sort(["date", "matu", "strike"])
    .rolling(index_column="date", period=window, group_by=["matu", "strike"])
    .agg(
        [
            pl.col(col).mean().alias(f"mean {col} {window}")
            for col in ["value1", "value2", "value3"]
        ]
        + [
            pl.col(col).std().alias(f"std {col} {window}")
            for col in ["value1", "value2", "value3"]
        ]
        + [pl.col(col).first() for col in ["value1", "value2", "value3"]]
    )
    .with_columns(
        (
            (pl.col(f"{col}") - pl.col(f"mean {col} {window}"))
            / pl.col(f"std {col} {window}")
        ).alias(f"z-score {col} {window}")
        for col, window in itertools.product(
            ["value1", "value2", "value3"], [window]
        )
    )
)

优化方案及代码 #

针对多周期批量处理需求和代码简洁性，可从减少冗余、统一逻辑两方面优化，以下是更高效的实现：

import polars as pl
import itertools

# 统一配置参数，方便后续扩展
target_cols = ["value1", "value2", "value3"]
rolling_windows = ["30d", "60d", "90d"]  # 支持多周期列表
group_cols = ["matu", "strike"]
date_col = "date"

# 先按日期和分组键排序
df = df.sort([date_col] + group_cols)

# 构建所有滚动计算表达式
rolling_calcs = []
for win in rolling_windows:
    # 复用滚动窗口上下文，避免重复定义参数
    win_ctx = pl.col(target_cols).rolling(index_column=date_col, period=win, group_by=group_cols)
    
    # 添加均值、标准差计算
    rolling_calcs.extend([
        win_ctx.mean().alias(f"mean_{col}_{win}") for col in target_cols
    ])
    rolling_calcs.extend([
        win_ctx.std().alias(f"std_{col}_{win}") for col in target_cols
    ])
    
    # 直接在滚动上下文里计算Z-score，省去后续单独列运算
    rolling_calcs.extend([
        ((pl.col(col) - win_ctx.mean()) / win_ctx.std()).alias(f"zscore_{col}_{win}")
        for col in target_cols
    ])

# 执行计算，保留原列的同时添加所有统计列
df = df.select(pl.all(), *rolling_calcs)

优化点说明 #

1. 参数集中管理：把列、周期、分组键统一配置，新增字段或周期只需修改参数列表，无需改动核心逻辑
2. 减少重复定义：复用win_ctx滚动窗口上下文，避免重复写rolling参数，代码更简洁
3. 一步计算Z-score：直接在滚动表达式内完成Z-score计算，省去原代码中agg后再with_columns的额外步骤，逻辑更连贯
4. 列名规范：用下划线替代空格作为分隔符，避免后续引用列名时的引号麻烦，也更符合常规命名习惯

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