Pandas groupby apply运行过慢，如何优化代码提升执行速度？

原代码用groupby.apply()逐组处理导致效率低下，核心问题是apply()会逐组调用Python函数，无法利用Pandas的矢量化运算优势。以下是针对性优化方案，完全用矢量化操作替代apply()，能大幅提升运行速度：

优化思路拆解 #

• 分组求和：直接用groupby.sum()计算各组的Call_OI、Call_Volume等字段总和，这是Pandas原生优化的操作，速度远快于逐组遍历。
• 提取首尾行数据：因为数据集已按Strike_Price排序，对above组取每组第一行，below组取每组最后一行。通过groupby.nth()分别获取首尾行，再根据group字段合并对应数据。
• 合并结果：将求和结果与首尾行数据按分组键合并，得到最终的汇总表。

优化后代码 #

# 确保数据已按Strike_Price排序（原代码已完成此步骤，这里保留）
data_v1_sorted = data_v1.sort_values(by="Strike_Price")

# 1. 计算分组总和
sum_df = data_v1_sorted.groupby(['dateTime','close','Index','Expiry','group'])[
    ['Call_OI', 'Call_Volume', 'Put_OI', 'Put_Volume']
].sum().rename(columns={
    'Call_OI': 'call_oi',
    'Call_Volume': 'call_vol',
    'Put_OI': 'put_oi',
    'Put_Volume': 'put_vol'
}).reset_index()

# 2. 获取每组第一行（对应above组）
head_df = data_v1_sorted.groupby(['dateTime','close','Index','Expiry','group'])[
    ['Call_OI', 'Call_Volume', 'Put_OI', 'Put_Volume']
].nth(0).rename(columns={
    'Call_OI': 'call_oi_1',
    'Call_Volume': 'call_vol_1',
    'Put_OI': 'put_oi_1',
    'Put_Volume': 'put_vol_1'
}).reset_index()

# 3. 获取每组最后一行（对应below组）
tail_df = data_v1_sorted.groupby(['dateTime','close','Index','Expiry','group'])[
    ['Call_OI', 'Call_Volume', 'Put_OI', 'Put_Volume']
].nth(-1).rename(columns={
    'Call_OI': 'call_oi_1',
    'Call_Volume': 'call_vol_1',
    'Put_OI': 'put_oi_1',
    'Put_Volume': 'put_vol_1'
}).reset_index()

# 4. 根据group字段筛选对应数据并合并
head_df = head_df[head_df['group'] == 'above']
tail_df = tail_df[tail_df['group'] == 'below']
first_last_df = pd.concat([head_df, tail_df]).reset_index(drop=True)

# 5. 合并总和与首尾行数据
result = pd.merge(sum_df, first_last_df, 
                  on=['dateTime','close','Index','Expiry','group'], 
                  how='left')

关键优化点说明 #

• 抛弃apply()：完全用Pandas原生的分组聚合和索引操作替代自定义函数，利用底层C实现加速运算，效率能提升10~100倍（取决于分组数量）。
• 矢量化批量处理：所有步骤都是对整个DataFrame的批量操作，没有逐组的Python循环，减少了大量函数调用开销。
• 逻辑完全等价：严格保留原代码逻辑：above组取排序后第一行，below组取最后一行，同时计算各组总和。

额外提速建议 #

• 优化分组键类型：比如dateTime用datetime类型而非字符串，close保持数值类型，能加快分组速度。
• 提前处理缺失值：如果Call_OI等字段的NaN不影响求和（sum()会自动忽略NaN），可提前用data_v1.fillna(0)统一处理，避免分组内重复操作。

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