如何高效执行pandas分组运算及向量化实现相关问题咨询

现有结构如下的DataFrame：

import pandas as pd
import numpy as np

df = pd.DataFrame([[1,22], [1,23], [1,44], [2, 33], [2, 55]], columns=['id', 'delay'])

数据预览：

id  delay
0   1     22
1   1     23
2   1     44
3   2     33
4   2     55

当前实现逻辑为按id字段分组后，对delay列执行固定权重的滚动运算，代码如下：

k = [0.1, 0.5, 1]

def f(d):
    d['new_delay'] = pd.Series([0,0]).append(d['delay']).rolling(window=3).apply(lambda x: np.sum(x*k)).iloc[2:]
    return d

df.groupby(['id']).apply(f)

运行结果符合预期：

id  delay  new_delay
0   1     22       22.0
1   1     23       34.0
2   1     44       57.7
3   2     33       33.0
4   2     55       71.5

1. 分组后的DataFrame调用.apply是否属于向量化操作？ #

不属于。当前实现存在两层Python级别的循环开销：

• 第一层groupby.apply会遍历每个id分组，逐个把分组数据传入自定义函数执行，每次调用都有Python层的函数调用开销
• 第二层rolling.apply会遍历每个滑动窗口，逐个把窗口数据传入lambda函数执行，同样走Python循环逻辑
  本质和手动写for循环遍历分组、遍历窗口的开销差异不大，不属于真正的向量化操作。

2. 大数据量下更高效的向量化实现方案 #

该需求本质是分组后对delay列做固定权重的滑动加权求和，前导补2个0，可以用两种性能远高于原实现的方案：

方案1：移位运算直接向量化计算（无额外依赖，性能最高） #

直接用分组移位代替滚动窗口，完全避免Python层循环：

k = np.array([0.1, 0.5, 1])
df['new_delay'] = df.groupby('id')['delay'].transform(
    lambda x: x.shift(2, fill_value=0)*k[0] + x.shift(1, fill_value=0)*k[1] + x*k[2]
)

该实现全走pandas底层C语言实现的运算逻辑，比原实现快10~100倍，数据量越大性能优势越明显，同时还规避了pd.Series.append在高版本pandas中的废弃警告。

方案2：Numba加速滚动运算（适合窗口/权重动态变化的场景） #

如果窗口长度、权重是动态可变的，可以开启pandas rolling的Numba引擎，把滚动运算编译成机器码执行：

k = np.array([0.1, 0.5, 1])
def rolling_weighted_sum(x):
    return np.sum(x * k)

df['new_delay'] = df.groupby('id')['delay'].apply(
    lambda x: pd.concat([pd.Series([0,0]), x]).rolling(3).apply(
        rolling_weighted_sum, engine='numba', raw=True
    ).iloc[2:]
)

首次运行需要编译函数，后续执行速度接近原生C语言水平。

3. Python单线程环境下pandas、numpy的向量化实现逻辑 #

虽然Python本身受GIL（全局解释器锁）限制默认是单线程执行，但pandas、numpy的核心运算逻辑都不在Python层实现：

• 两者的数组都是连续内存块存储，运算时直接把整块数据交给底层的C/Fortran实现执行，完全避开Python层的循环开销和GIL限制
• 底层运算依赖高度优化的线性代数库（如MKL、OpenBLAS），这些库原生支持SIMD（单指令多数据流）指令集，一次CPU指令就可以完成多个数据的并行计算
• 大部分底层运算库自带多线程优化，可以自动利用CPU多核资源并行计算，不需要Python层面启动多线程，所以即使Python本身是单线程的，也能实现高效的并行向量化计算。

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