Pandas升级后可变窗口滚动计算的替代实现方案咨询

问题背景 #

在Pandas 1.4.1中，通过自定义BaseIndexer并传入window_size=1的变通方法，可实现每个数据点对应不同大小的滚动窗口计算，但升级到1.5.2后，新增的窗口大小校验机制会因返回窗口与指定window_size不匹配抛出错误，且纯遍历数据的方式速度过慢。

方案一：调整自定义索引器适配Pandas校验 #

修改索引器初始化逻辑，将window_size设为所有窗口的最大尺寸，让Pandas的校验通过，同时保留原可变窗口计算逻辑：

import numpy as np
from pandas.api.indexers import BaseIndexer

class CustomVariableWindowIndexer(BaseIndexer):
    def __init__(self, true_window_size, **kwargs):
        # 将window_size设为窗口最大值，通过Pandas大小校验
        self.true_window_size = np.array(true_window_size, dtype=np.int64)
        super().__init__(window_size=self.true_window_size.max(), **kwargs)
    
    def get_window_bounds(self, num_values, min_periods, center, closed):
        # 生成每个窗口的起始索引
        start = np.arange(num_values, dtype=np.int64)
        # 生成结束索引，确保不超出数据总长度
        end = start + self.true_window_size
        end = np.minimum(end, num_values)
        return start, end

使用方式：

# list_of_window_sizes为每个数据点对应的窗口大小列表
indexer = CustomVariableWindowIndexer(true_window_size=list_of_window_sizes)
# 执行滚动均值计算示例
result = df['target_col'].rolling(window=indexer).mean()

方案二：Numba加速的向量化计算 #

用Numba编译循环逻辑，比纯Python遍历快一个数量级以上，同时避开Pandas的窗口校验限制：

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

def variable_window_calculate(series, window_sizes, agg_func):
    arr = series.to_numpy()
    n = len(arr)
    starts = np.arange(n)
    ends = starts + window_sizes
    ends = np.minimum(ends, n)
    
    # Numba编译计算逻辑
    @jit(nopython=True)
    def compute_window_results(arr, starts, ends, func):
        results = np.empty(n, dtype=np.float64)
        for i in range(n):
            window_data = arr[starts[i]:ends[i]]
            results[i] = func(window_data)
        return results
    
    # 转回Pandas Series
    return pd.Series(compute_window_results(arr, starts, ends, agg_func), index=series.index)

使用方式：

# 计算每个可变窗口的均值
mean_results = variable_window_calculate(df['target_col'], list_of_window_sizes, np.mean)
# 计算每个可变窗口的总和
sum_results = variable_window_calculate(df['target_col'], list_of_window_sizes, np.sum)

方案三：中小规模数据适配的切片+apply方法 #

如果数据量不大，可预先生成每个窗口的切片，再用apply完成计算：

def get_window_slice(i, window_size, arr):
    end_idx = min(i + window_size, len(arr))
    return arr[i:end_idx]

arr = df['target_col'].to_numpy()
# 生成每个数据点对应的窗口数据
window_data = pd.Series(range(len(arr))).apply(lambda x: get_window_slice(x, list_of_window_sizes[x], arr))
# 应用聚合函数
result = window_data.apply(np.mean)

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