Pandas中应用自定义滚动函数遇AttributeError的解决方法

问题场景 #

现有mtcars DataFrame，需要用自定义函数fun处理多列数据。全量数据通过apply可正常计算，但尝试用滚动窗口实现相同逻辑时触发错误。

原始代码：

import statsmodels.api as sm
import numpy as np
import pandas as pd
mtcars=pd.DataFrame(sm.datasets.get_rdataset("mtcars", "datasets", cache=True).data)

def fun(col1, col2, w1=10, w2=2):
    return(np.mean(w1 * col1 + w2 * col2))

# 全量数据计算正常
mtcars.apply(lambda x: fun(x.cyl, x.mpg), axis=1)

# 滚动窗口写法报错
mtcars.rolling(3).apply(lambda x: fun(x.cyl, x.mpg))

报错信息：

AttributeError: 'Series' object has no attribute 'cyl'

核心原因 #

rolling().apply()默认针对单列执行滚动计算，每次传入自定义函数的是当前窗口的单列Series，而非包含多列的DataFrame切片，因此无法通过x.cyl这种方式访问列名。

解决方案 #

针对Pandas 1.5.2版本，可通过将多列打包为二维数组，让滚动窗口处理整个数组，再在自定义函数中拆分列进行计算：

import statsmodels.api as sm
import numpy as np
import pandas as pd

mtcars=pd.DataFrame(sm.datasets.get_rdataset("mtcars", "datasets", cache=True).data)

# 定义适配滚动窗口的函数，接收窗口内的二维数组切片
def rolling_fun(window_array, w1=10, w2=2):
    # 拆分出cyl和mpg列
    col1 = window_array[:, 0]
    col2 = window_array[:, 1]
    return np.mean(w1 * col1 + w2 * col2)

# 提取需要的两列转为二维数组
target_cols = mtcars[['cyl', 'mpg']].to_numpy()

# 应用滚动窗口，center=True对应R中的.before=1,.after=1（中心窗口）
# raw=True确保直接传递数组而非Series
roll_result = pd.Series(target_cols.tolist()).rolling(3, center=True).apply(rolling_fun, raw=True)

# 将结果合并到原DataFrame
mtcars['roll'] = roll_result

执行后得到的roll列结果与R中slider::slide2的输出一致，例如前3行结果为：

Mazda RX4           102.00000
Mazda RX4 Wag        96.53333
Datsun 710           96.80000

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