Pandas分组应用自定义函数的性能优化方案咨询

场景与数据结构 #

我有一个约10万行的DataFrame，包含Category、val1、val2、val3、val4列，示例数据如下：

需求目标 #

需要按Category列分组，对每个分组执行自定义方法（返回float类型的cal值），最终得到如下格式的字典：

{ 
    'A': { 'cal': a },
    'B': { 'cal': b },
    ...
}

当前实现与性能问题 #

最初使用groupby.apply实现，代码如下：

def my_cal(df):
    ret = ...
    return {'cal': ret}

df.groupby('Category').apply(lambda grp: my_cal(grp)).to_dict()

在Jupyter Notebook用timeit测试，耗时超1秒，不符合性能预期。

之后更新代码，将my_cal入参改为数组，减少分组内DataFrame操作：

def my_cal(val1: float, val2: float, val3: float, val4: float):
    ret = inner_cal(val1, val2, val3, val4) # inner_cal是外部库函数
    return {'cal': ret}

df.groupby('Category').apply(lambda grp: my_cal(grp['val1'].to_numpy(),
                                                grp['val2'].to_numpy(),
                                                grp['val3'].to_numpy(),
                                                grp['val4'].to_numpy())).to_dict()

但性能仍未达标，寻求优化方案。

优化方案 #

1. 手动分组+NumPy数组直接操作 #

避开Pandasgroupby.apply的Python层循环开销，直接用NumPy数组切片处理分组：

# 提前提取所有数组，减少重复转换
cat_arr = df['Category'].to_numpy()
val1_arr = df['val1'].to_numpy()
val2_arr = df['val2'].to_numpy()
val3_arr = df['val3'].to_numpy()
val4_arr = df['val4'].to_numpy()

result = {}
for cat in df['Category'].unique():
    # 生成当前分组的掩码
    mask = cat_arr == cat
    # 直接传入切片后的数组给inner_cal
    cal_val = inner_cal(val1_arr[mask], val2_arr[mask], val3_arr[mask], val4_arr[mask])
    result[cat] = {'cal': cal_val}

这种方式避免了每个分组创建子DataFrame的开销，直接操作底层数组，能显著降低耗时。

2. 利用Numba JIT编译（若inner_cal可重写） #

如果inner_cal的逻辑可以自行实现，用Numba对自定义计算函数进行JIT编译，结合groupby.apply的engine='numba'参数加速：

from numba import jit

# 重写inner_cal的逻辑为Numba可编译的版本
@jit(nopython=True)
def numba_inner_cal(val1, val2, val3, val4):
    # 这里实现原inner_cal的计算逻辑
    ...

def my_cal(val1, val2, val3, val4):
    ret = numba_inner_cal(val1, val2, val3, val4)
    return {'cal': ret}

# 使用numba引擎加速apply
df.groupby('Category').apply(my_cal, engine='numba').to_dict()

注意：如果inner_cal是无法修改的外部库函数，此方法不适用。

3. 减少字典转换的额外开销 #

如果最终字典的格式可以调整（或后续可转换），可以先让apply返回Series，再转换为字典，比返回字典后再转的开销更小：

def my_cal(val1, val2, val3, val4):
    return inner_cal(val1, val2, val3, val4)

# 先得到{Category: cal_val}的字典，再包装成目标格式
temp_dict = df.groupby('Category').apply(lambda grp: my_cal(grp['val1'].to_numpy(),
                                                           grp['val2'].to_numpy(),
                                                           grp['val3'].to_numpy(),
                                                           grp['val4'].to_numpy())).to_dict()
result = {k: {'cal': v} for k, v in temp_dict.items()}

这种方式减少了每个分组返回字典的序列化开销，先得到扁平字典再二次包装，速度更快。

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