为何Pandas的groupby.apply在小数据集上运行缓慢？

问题场景 #

需要通过两个分组变量聚合Pandas DataFrame，因涉及多列操作使用了groupby.apply方法，代码可正常运行但速度极慢：4000行数据聚合耗时约2.77秒，单分组变量时耗时减半。

原代码及运行耗时 #

import random
import pandas as pd
import numpy as np
import time

df = pd.DataFrame(np.random.rand(4000,4), columns=list('abcd'))
df['group'] = random.choices([0, 0, 1, 1],k=4000)
df["grupp"]=  random.choices([2, 3, 4, 2],k=4000)

def f(x):
    d = {}
    d['c_d_prodsum'] = (x['c'] * x['d']).sum()
    return pd.Series(d, index=['c_d_prodsum'])

start = time.time()
%timeit b=df.groupby(['group','grupp']).apply(f)
end = time.time()
print(end - start)

本地运行结果：33.2 ms ± 2.03 ms per loop，总耗时约2.77秒

速度缓慢的原因 #

• apply的逐组开销：apply会遍历每个分组执行自定义函数，多分组变量会导致分组总数大幅增加（比如两个变量的组合数远多于单个变量的分组数），每个分组的函数调用、临时对象创建都会累积额外开销。
• 自定义函数的冗余操作：原函数里手动创建字典再转换为Series，多了不必要的中间步骤，进一步增加了计算耗时。
• 多分组的拆分重组成本：多分组变量下，DataFrame的拆分和结果重组本身就比单分组消耗更多资源。

优化方案 #

方案1：使用内置矢量化聚合（推荐） #

先计算c*d的衍生列，再用groupby.sum直接聚合，完全利用Pandas的矢量化操作，避免逐组循环的开销：

import random
import pandas as pd
import numpy as np
import time

df = pd.DataFrame(np.random.rand(4000,4), columns=list('abcd'))
df['group'] = random.choices([0, 0, 1, 1],k=4000)
df["grupp"]=  random.choices([2, 3, 4, 2],k=4000)

start = time.time()
%timeit b = df.assign(c_d=df['c']*df['d']).groupby(['group','grupp'])['c_d'].sum().rename('c_d_prodsum')
end = time.time()
print(end - start)

这种方式的速度会比原代码快几十倍，因为矢量化操作是Pandas底层优化的，效率远高于逐组循环。

方案2：简化自定义函数（仅适用于必须用apply的复杂场景） #

如果因业务逻辑复杂必须使用apply，可以去掉冗余的字典创建步骤，直接返回计算结果：

def f(x):
    return (x['c'] * x['d']).sum()

%timeit b = df.groupby(['group','grupp']).apply(f).rename('c_d_prodsum')

该方法能小幅提升速度，但仍远不如内置聚合方法。

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