pandas多列分组计算：多次groupby还是单次groupby内实现效率更高？

问题背景 #

在使用pandas开展数据处理时，假设需要为现有DataFrame新增A、B、C三列，每列都需要基于原有数据完成对应计算，且所有计算都需按照code列划分的行组执行，对比以下两种实现方式的运行速度：

• 方式一：分三次调用groupby+apply，分别生成对应列：

df.groupby('code').apply(generate_A)
df.groupby('code').apply(generate_B)
df.groupby('code').apply(generate_C)

• 方式二：仅调用一次groupby+apply，在传入的generate_A_B_C函数内对每个分组依次调用generate_A、generate_B、generate_C，一次性生成三列结果：

df.groupby('code').apply(generate_A_B_C)

以下为实际业务场景中的对应代码示例：

df['rsi_1_increases_above_lower_threshold_Close'] = df.groupby('SecuritiesCode').apply(detect_if_RSI_crosses_above_lower_threshold,
                                                                                                          rsi_lower_threshold = rsi_lower_threshold,
                                                                                                          rsi_column='rsi_1_Close')

df['rsi_1_decreases_below_upper_threshold_Close'] = df.groupby('SecuritiesCode').apply(detect_if_RSI_crosses_below_upper_threshold,
                                                                                                          rsi_upper_threshold = rsi_upper_threshold,
                                                                                                          rsi_column='rsi_1_Close')

结论 #

方式二的运行速度明显更快。

原因说明 #

1. 避免了重复的分组构建开销
   每次调用groupby时，pandas都需要遍历全量数据，对分组键做哈希计算、拆分数据块、维护分组和行索引的映射关系，这部分操作在数据量越大、分组数量越多时开销越高。方式一需要重复执行3次完整的分组流程，方式二仅需执行1次。就算提前把分组对象缓存为变量（比如先执行g = df.groupby('code')，再三次调用g.apply()），省去了重复拆分分组的开销，依然存在其他额外成本。
2. 减少了结果合并与索引对齐的重复开销
   每次apply执行完成后，pandas都需要把所有分组返回的计算结果做拼接，再和原DataFrame的索引做对齐，最终生成可以直接赋值给列的序列对象。方式一需要做3次结果拼接、索引对齐操作，方式二仅需做1次。
3. 减少了分组数据的重复遍历
   方式一中每个分组的数据会被三次apply分别读取、计算三次，方式二中每个分组的数据仅需读取一次，就可以在同一个函数内完成所有字段的计算，数据读取层面的开销也更低。

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