如何对依赖两个DataFrame的pandas自定义计算操作实现向量化优化

报错原因 #

你设置raw=True后，apply传入lambda的参数x是numpy数组而非pandas Series，不支持字符串索引（比如x['Cat']），只能用位置下标取值，这是触发IndexError的直接原因。如果一定要用raw=True，可以把代码里的x['Cat']改成x[0]（对应df2第一列是Cat的情况），但这不会带来本质的性能提升，因为逐行遍历+每次布尔索引过滤df1的逻辑没有变。

优化方案 #

你不需要基于df2逐行遍历，完全可以通过优化计算函数+原生pandas分组操作实现十倍甚至百倍的性能提升：

1. 重构计算逻辑，移除无必要的pandas操作 #

原some_calc里的设置索引、DataFrame封装都是不必要的，你需要的外积矩阵完全可以用numpy实现，避免pandas的额外开销：

import pandas as pd
import numpy as np

def calc_per_cat_group(group):
    # 直接从分组片段中取数值转为numpy数组，运算速度远快于pandas接口
    data1 = group['data_point1'].values
    data2_sqrt = np.sqrt(group['data_point2'].values)
    # 转为列向量后做矩阵乘法直接得到M*M外积，不需要依赖索引保持维度
    vec = (data1 * data2_sqrt).reshape(-1, 1)
    grid = np.random.rand(vec.shape[0], vec.shape[0])
    # 直接用numpy逐元素乘+求和，没有额外开销
    return (vec @ vec.T * grid).sum()

2. 直接对df1做分组计算，避免逐行过滤 #

你本身就是按Cat分组计算，直接用pandas原生的groupby即可，内部是优化过的分组逻辑，远快于自己写布尔索引逐行过滤：

# 直接对df1按Cat分组计算，结果的索引为Cat分类值
group_calc_result = df1.groupby('Cat').apply(calc_per_cat_group).rename('Apply')
# 把计算结果合并到df2中，自动按Cat匹配
df2 = df2.merge(group_calc_result, on='Cat')

3. 可选：进一步极限优化 #

如果Cat分类数量很少（比如你示例里只有4类），还可以预先把所有组的片段存到字典里，完全避免pandas groupby的微小开销，不过一般场景下上面的方案已经足够满足性能需求。

原写法兼容方案（不推荐） #

如果你不想修改原有some_calc的实现，也可以把raw=True版本的代码改为位置索引即可运行：

# 注意x[0]对应df2的第一列Cat，如果你的df2列顺序变化需要调整下标
df2['ApplyRaw'] = df2.apply(lambda x: some_calc(x=df1[df1['Cat']==x[0]][['Cat','data_point1']], 
                                                y=df1[df1['Cat']==x[0]][['Cat','data_point2']]),axis=1, raw=True)

这个写法只是解决了报错问题，性能没有任何提升，sample_size增大后仍然会很慢。

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