如何在Pandas中无循环为Rolling Window应用自定义函数并解决列限制？

核心思路 #

利用pandas的rolling.apply结合method="table"和result_type="expand"参数，直接处理窗口DataFrame并生成任意列数的结果，无需循环或raw=True参数。

解决问题1：按列名访问窗口数据 #

不设置raw=True时，rolling(window=..., method="table").apply()会将每个窗口的DataFrame传入自定义函数，你可以直接通过列名（如window_df["A"]）引用数据进行计算。

解决问题2：返回任意列数的结果 #

添加result_type="expand"参数，函数返回的多元素数组/Series会被自动扩展为多列，列数可与输入列数不同。

示例代码1：基础实现（无numba加速） #

import numpy as np
import pandas as pd

def process_window(window_df):
    # 直接用列名访问窗口数据，自定义计算逻辑
    col_d = window_df["A"].mean()
    col_f = window_df["B"].sum()
    col_g = window_df["C"].max()
    # 返回3个值，对应3列新数据
    return pd.Series([col_d, col_f, col_g])

df = pd.DataFrame({"A": range(10), "B": range(10, 20), "C": range(20, 30)})

# 生成新列，result_type="expand"将结果扩展为多列
new_columns = df.rolling(window=2, method="table").apply(process_window, result_type="expand")
new_columns.columns = ["D", "F", "G"]

# 合并到原DataFrame
df = pd.concat([df, new_columns], axis=1)
print(df)

示例代码2：兼容numba加速的折中方案 #

如果需要numba加速（处理大数据量），可以在函数内先通过列名提取所需数据，再转为numpy数组计算（兼顾列名可读性和numba性能）：

import numpy as np
import pandas as pd
from numba import jit

# 用numba装饰纯numpy运算的函数
@jit(nopython=True)
def numba_calc(a_arr, b_arr):
    return np.array([a_arr.mean(), b_arr.sum()])

def process_window(window_df):
    # 按列名提取数据转numpy数组
    a = window_df["A"].to_numpy()
    b = window_df["B"].to_numpy()
    return numba_calc(a, b)

df = pd.DataFrame({"A": range(10), "B": range(10, 20), "C": range(20, 30)})

# 生成2列新数据，列数与输入列数（3列）不同
new_columns = df.rolling(window=2, method="table").apply(process_window, result_type="expand", engine="numba")
new_columns.columns = ["D", "F"]

df = pd.concat([df, new_columns], axis=1)
print(df)

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