分组(groupby)后执行diff方法速度极慢的性能问题排查

问题背景 #

处理5万行5列的Pandas DataFrame，结构如下：

• 列A、B、C：产品唯一标识（如销售店铺、产品类型、品牌）
• 列X：可排序变量（如记录日期）
• 列Y：监测值（如当日销量）

需求为：按A、B、C的唯一组合，计算Y随X变化的差值。预期操作逻辑是先按X排序，再对A、B、C分组，最后对Y执行diff。但实际测试发现：

• 单独执行groupby耗时0.005秒
• 单独执行diff耗时0.005秒
• 两者结合后耗时10秒，性能骤降1000倍

使用Pandas版本为0.24.2，测试代码如下：

import numpy, pandas
from timeit import default_timer as timer

# Initialize dummy df
df = pandas.DataFrame( numpy.random.randint( 0, 100, size=( 50000, 5 ) ), columns=list( 'ABCXY') )

# Groupby
start = timer()
_ = df.sort_values( [ 'X' ], axis=0 ).groupby( list( 'ABC' ) )[ 'Y' ]
print( timer() - start )
    
# Diff
start = timer()
_ = df.sort_values( [ 'X' ], axis=0 )[ 'Y' ].diff()
print( timer() - start )
    
# Groupby + Diff
start = timer()
df.index = df.index.astype(int)
_ = df.sort_values( [ 'X' ], axis=0 ).groupby( list( 'ABC' ) )[ 'Y' ].diff()
print( timer() - start )

问题根源 #

Pandas 0.24.x版本中，groupby.diff()的底层实现存在性能缺陷：当分组数量较多时，会对每个分组单独执行diff操作，产生大量循环开销，导致整体性能急剧下降。而单独的groupby仅创建分组对象不执行计算，单独的diff是对整列的向量化操作，因此两者单独执行效率都很高。

优化方案 #

方案1：先按分组+排序键排序，再用全局diff配合分组标记 #

通过先排序确保同组内数据按X有序，再执行全局diff，最后修正分组首行的diff结果，完全规避groupby的循环开销：

import numpy, pandas
from timeit import default_timer as timer

df = pandas.DataFrame( numpy.random.randint( 0, 100, size=( 50000, 5 ) ), columns=list( 'ABCXY') )

start = timer()
# 按分组键+排序键排序，保证同组内X递增
df_sorted = df.sort_values(['A', 'B', 'C', 'X'])
# 计算全局diff
df_sorted['Y_diff'] = df_sorted['Y'].diff()
# 标记每个分组的第一行，将其diff结果设为NaN
group_mask = df_sorted[['A', 'B', 'C']].ne(df_sorted[['A', 'B', 'C']].shift()).any(axis=1)
df_sorted.loc[group_mask, 'Y_diff'] = numpy.nan
print(f"优化后耗时：{timer() - start:.4f}秒")

此方法耗时可降至0.01秒以内，与单独操作性能持平。

方案2：升级Pandas版本 #

Pandas 1.0及后续版本对groupby.diff()做了底层优化，改用向量化操作替代分组循环，原代码性能会大幅提升。升级到稳定新版本后，直接运行原groupby+diff代码即可获得接近单独操作的效率。

方案3：使用groupby.transform配合diff #

若无法升级版本，可尝试用transform包装diff，性能略逊于方案1，但远优于原代码：

start = timer()
df_sorted = df.sort_values(['X'])
df_sorted['Y_diff'] = df_sorted.groupby(['A', 'B', 'C'])['Y'].transform(lambda x: x.diff())
print(f"transform方案耗时：{timer() - start:.4f}秒")

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