Python-Polars中rolling_map线性拟合运行缓慢且内存溢出问题求助

核心问题原因 #

1. rolling_map的低效性：逐窗口调用自定义ols_slope/ols_offset函数，每个窗口都要重复计算rank、均值、求和，产生大量临时对象，内存占用激增且CPU开销极大。
2. 错误的自变量x定义：用y.rank("ordinal")作为线性回归的x，每个窗口都要对y排序，这是时序预测中的不合理操作（时序预测x应为时间步），进一步加剧性能损耗。

优化方案 #

1. 修正自变量x的定义

时序局部线性预测中，x应为窗口内的相对时间步（如1到60，对应60分钟窗口的每个分钟），无需对y排序计算rank。若业务确实需要用rank作为x，也需改为组内批量计算。

2. 用group_by_dynamic替换rolling_map

group_by_dynamic是Polars专为时序数据设计的高效分组工具，支持批量处理滚动窗口，避免逐窗口函数调用的开销。

3. 批量处理多列

针对300列的场景，使用列选择语法批量生成计算逻辑，减少冗余代码。

完整实现代码 #

场景1：x为窗口内相对时间步（推荐）

import polars as pl

# 读取并排序数据（时序数据必须按时间排序）
data = pl.read_ipc("../data_vault_1min.feather").sort("date_utc")

# 预计算窗口大小为60时的x统计量（固定值）
n_full = 60
sum_x_full = n_full * (n_full + 1) // 2
sum_x2_full = n_full * (n_full + 1) * (2 * n_full + 1) // 6

# 定义要处理的列（示例为所有以A开头的列，可根据实际调整）
target_cols = [col for col in data.columns if col.startswith("A")]

# 批量生成聚合表达式
agg_exprs = []
for col in target_cols:
    agg_exprs.extend([
        pl.col(col).sum().alias(f"{col}_sum_y"),
        # 计算x*y的和，x为组内相对位置（1到组内行数）
        pl.col(col).zip_with(pl.int_range(1, pl.count() + 1)).sum().alias(f"{col}_sum_xy")
    ])

# 滚动窗口分组计算
result = data.group_by_dynamic(
    index="date_utc",
    every="1m",  # 每1分钟生成一个窗口
    period="60m",  # 窗口覆盖过去60分钟
    offset="-59m",  # 窗口起始为当前时间前59分钟，确保窗口包含当前时间
    min_periods=3
).agg(
    *agg_exprs,
    count=pl.count()  # 记录每个窗口的实际行数
).with_columns(
    # 根据窗口实际行数计算x的统计量
    sum_x=pl.when(pl.col("count") < n_full)
            .then(pl.col("count") * (pl.col("count") + 1) // 2)
            .otherwise(sum_x_full),
    sum_x2=pl.when(pl.col("count") < n_full)
             .then(pl.col("count") * (pl.col("count") + 1) * (2 * pl.col("count") + 1) // 6)
             .otherwise(sum_x2_full)
)

# 批量计算斜率和截距，并生成预测列
for col in target_cols:
    result = result.with_columns(
        # 计算斜率
        ((pl.col("count") * pl.col(f"{col}_sum_xy") - pl.col("sum_x") * pl.col(f"{col}_sum_y")) / 
         (pl.col("count") * pl.col("sum_x2") - pl.col("sum_x") ** 2)).alias(f"{col}_hourly_lin_slope"),
        # 计算截距
        ((pl.col(f"{col}_sum_y") - pl.col(f"{col}_hourly_lin_slope") * pl.col("sum_x")) / 
         pl.col("count")).alias(f"{col}_hourly_lin_intercept"),
        # 计算1小时超前预测值（x为窗口行数+1）
        (pl.col(f"{col}_hourly_lin_slope") * (pl.col("count") + 1) + 
         pl.col(f"{col}_hourly_lin_intercept")).alias(f"{col}_hourly_lin_pred")
    )

# 将预测结果与原数据合并（按时间对齐）
final_data = data.join(result, on="date_utc", how="left")

场景2：x必须为y的rank（不推荐，仅适配业务特殊需求）

import polars as pl

data = pl.read_ipc("../data_vault_1min.feather").sort("date_utc")
target_cols = [col for col in data.columns if col.startswith("A")]

# 批量生成聚合表达式
agg_exprs = []
for col in target_cols:
    agg_exprs.append(pl.col(col).alias(f"{col}_values"))

# 滚动窗口分组计算
result = data.group_by_dynamic(
    index="date_utc",
    every="1m",
    period="60m",
    offset="-59m",
    min_periods=3
).agg(
    *agg_exprs,
    count=pl.count()
).with_columns(
    # 批量计算每个列的rank、sum_x、sum_y、sum_xy、sum_x2
    **{
        f"{col}_x": pl.col(f"{col}_values").rank("ordinal"),
        f"{col}_sum_y": pl.col(f"{col}_values").sum(),
        f"{col}_sum_x": pl.col(f"{col}_x").sum(),
        f"{col}_sum_xy": (pl.col(f"{col}_x") * pl.col(f"{col}_values")).sum(),
        f"{col}_sum_x2": (pl.col(f"{col}_x") ** 2).sum()
    } for col in target_cols
).with_columns(
    # 批量计算斜率和预测值
    **{
        f"{col}_hourly_lin_slope": (pl.col("count") * pl.col(f"{col}_sum_xy") - pl.col(f"{col}_sum_x") * pl.col(f"{col}_sum_y")) / 
                                   (pl.col("count") * pl.col(f"{col}_sum_x2") - pl.col(f"{col}_sum_x") ** 2),
        f"{col}_hourly_lin_pred": pl.col(f"{col}_hourly_lin_slope") * pl.col(f"{col}_x").mean() - pl.col(f"{col}_sum_y")/pl.col("count")
    } for col in target_cols
).drop([col for cols in [(f"{col}_values", f"{col}_x", f"{col}_sum_y", f"{col}_sum_x", f"{col}_sum_xy", f"{col}_sum_x2") for col in target_cols] for col in cols])

final_data = data.join(result, on="date_utc", how="left")

额外优化建议 #

• 升级Polars版本：0.20.6是较旧版本，新版本对group_by_dynamic和滚动窗口操作有性能优化，可尝试升级到最新稳定版。
• 内存分片处理：若数据量仍过大，可按时间分片（如按天拆分），处理后再合并结果，降低单批次内存占用。

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