使用Numba加速pandas rolling.apply时遇NumbaPerformanceWarning求助

警告原因 #

你遇到的NumbaPerformanceWarning核心原因是：Pandas的rolling.apply配合engine="numba"时，parallel=True参数无法触发并行执行。

Pandas在调用你的自定义ewm函数时，会自动生成一个包装函数（即警告里提到的impl），这个函数的逻辑是逐个遍历滚动窗口，每次调用ewm计算单个窗口结果。由于该包装函数没有可被Numba识别的并行循环结构，即使指定parallel=True，Numba也无法进行并行转换，因此抛出警告。

正确的并行加速实现方式 #

要利用Numba实现滚动窗口计算的并行加速，需要直接在Numba层面处理滚动窗口的循环，而非依赖Pandas的rolling.apply机制。具体实现如下：

1. 编写Numba并行加速的滚动EWM函数 #

用numba.jit(parallel=True)装饰器，结合numba.prange标记可并行的循环（每个窗口的计算相互独立，适合并行）：

import numba
import numpy as np

@numba.jit(parallel=True)
def numba_rolling_ewm(arr, window, half_life=63):
    n = len(arr)
    result = np.empty(n, dtype=np.float64)
    min_periods = window // 2
    
    # 用prange让Numba并行处理每个元素对应的窗口
    for i in numba.prange(n):
        # 确定当前窗口的范围
        start_idx = max(0, i - window + 1)
        current_window = arr[start_idx:i+1]
        
        # 处理min_periods逻辑
        if len(current_window) < min_periods:
            result[i] = np.nan
            continue
        
        # 计算EWM权重与结果
        weights = 0.5 ** (np.arange(current_window.size-1, -1, -1) / half_life)
        weighted_sum = np.sum(current_window * weights)
        sum_weights = np.sum(weights)
        result[i] = weighted_sum / sum_weights
    
    return result

2. 结合Pandas Groupby调用并行函数 #

修改ts_ewm函数，通过groupby.apply对每个分组的时间序列应用上述并行函数：

def ts_ewm(s, window):
    # 按people分组，对每组的数值数组应用并行滚动EWM计算
    grouped_result = s.groupby('people').apply(lambda x: numba_rolling_ewm(x.values, window))
    
    # 整理索引，匹配原Series的结构
    grouped_result = grouped_result.explode()
    grouped_result.index = grouped_result.index.reorder_levels(['date', 'people'])
    return grouped_result.reindex_like(s)

替代方案：仅加速单个窗口计算（无并行） #

如果坚持使用Pandas的rolling.apply，可以去掉parallel=True，仅用Numba加速单个窗口的计算逻辑，虽无法并行处理窗口，但能提升单个窗口的计算速度：

@numba.jit
def ewm(s):
    half_life = 63
    weight = 0.5 ** (np.arange(s.size-1, -1, -1)/half_life)
    return np.sum(s * weight) / np.sum(weight)

def ts_ewm(s, window):
    res = s.groupby('people').rolling(window, min_periods=window//2)
    # 移除parallel=True，仅启用Numba加速单个窗口计算
    res = res.apply(ewm, raw=True, engine="numba")
    res = res.droplevel(0)
    return res.reindex_like(s)

关键总结 #

• Pandas的rolling.apply的parallel=True参数无法实现窗口间的并行，因为Pandas的窗口遍历逻辑无法被Numba识别为可并行循环。
• 要实现真正的并行加速，必须将滚动窗口的循环逻辑放入Numba函数中，用prange标记并行循环，再通过groupby.apply调用。

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