Pandas中sum(level=...)与groupby(level=...).sum()的性能差异及优化

问题描述 #

我在使用Pandas时发现两段聚合代码性能差异明显：

df.sum(level=['My', 'Index', 'Levels']) # 底层用向量化的numpy.sum

和

df.groupby(level=['My', 'Index', 'Levels']).sum() # 速度很慢...

根据官方文档，前者使用了向量化的numpy.sum，但遗憾的是sum(level=...)已被API弃用，执行时会触发以下警告：

FutureWarning: Using the level keyword in DataFrame and Series
aggregations is deprecated and will be removed in a future version.
Use groupby instead. df.sum(level=1) should use
df.groupby(level=1).sum()

我不想因为改用非向量化的groupby版本导致性能下降，请问如何在groupby中结合numpy.sum来保持原有的性能？

补充测试验证 #

针对这个性能差异，我做了两组测试（基于Pandas 1.4.4）：

基础测试（1万行数据，10级索引） #

测试代码：

import pandas as pd
import numpy as np
print('pandas:', pd.__version__)

nb_lines = int(1e4)
nb_levels = 10
# 生成10k行、每行10个[0,8]范围内的随机整数
ix = np.random.randint(0, nb_levels-1, size=(nb_lines, nb_levels))

cols = [chr(65+i) for i in range(nb_levels)] # 列名A、B、C...
df = pd.DataFrame(ix, columns=cols)
df = df.set_index(cols)
df['VALUE'] = np.random.rand(nb_lines) # 生成待聚合的随机值

print('with groupby:')
%timeit -n 300 df.groupby(level=cols).sum()
print('without groupby:')
%timeit -n 300 df.sum(level=cols)

测试结果：

pandas: 1.4.4
with groupby:
5.51 ms ± 1.06 µs per loop (mean ± std. dev. of 7 runs, 500 loops each)
without groupby:
<magic-timeit>:1: FutureWarning: Using the level keyword in DataFrame and Series aggregations is deprecated and will be removed in a future version. Use groupby instead. df.sum(level=1) should use df.groupby(level=1).sum().
4.93 ms ± 40.1 µs per loop (mean ± std. dev. of 7 runs, 500 loops each)

可以看到，不带groupby的版本始终更快，调整层级、分组列等参数后结果一致，大数据集下两者的性能差异会更显著。

进阶测试（10万行数据，10级索引） #

针对@mozway提到的“性能相近”的测试结果，我调整了测试规模（10级索引+10万行数据），代码如下：

import numpy as np
from string import ascii_uppercase as UP

np.random.seed(0)
N = 100_000
nb_levels = 10
cols = [chr(65+i) for i in range(nb_levels)]

d = {cols[i]: np.random.choice(list(UP), size=N) for i in range(nb_levels)}
d.update({'num': np.random.random(size=N)})
df = pd.DataFrame(d).set_index(cols)

print(pd.__version__)

print('with groupby:')
%timeit -n 300 df.groupby(level=cols).sum()
print('without groupby:')
%timeit -n 300 df.sum(level=cols)

测试结果：

1.4.4
with groupby:
50.8 ms ± 536 µs per loop (mean ± std. dev. of 7 runs, 300 loops each)
without groupby:
<magic-timeit>:1: FutureWarning: Using the level keyword in DataFrame and Series aggregations is deprecated and will be removed in a future version. Use groupby instead. df.sum(level=1) should use df.groupby(level=1).sum().
42 ms ± 506 µs per loop (mean ± std. dev. of 7 runs, 300 loops each)

可以看到，当数据规模和索引层级增加后，groupby版本的性能下降更明显。

解决方案 #

要在groupby中实现类似sum(level=...)的向量化性能，可以通过以下几种方式：

• 使用groupby.agg直接调用numpy.sum
  通过agg方法指定使用np.sum，可以强制使用向量化的numpy求和逻辑，绕过Pandas默认的聚合路径：

  df.groupby(level=['My', 'Index', 'Levels']).agg(np.sum)
  

• 利用get_level_values+numpy底层分组求和（极致性能）
  如果需要极致性能，可以手动提取索引层级的分组键，结合numpy的unique和bincount实现分组求和，完全基于向量化操作，性能最接近原sum(level=...)：

  # 获取分组键的唯一标识
  keys = df.index.get_level_values(['My', 'Index', 'Levels']).to_numpy()
  unique_keys, indices = np.unique(keys, return_inverse=True, axis=0)
  
  # 对数值列进行分组求和
  sum_values = np.bincount(indices, weights=df['VALUE'].values)
  
  # 重新构建结果DataFrame
  result = pd.DataFrame(sum_values, index=pd.MultiIndex.from_tuples(unique_keys, names=['My', 'Index', 'Levels']), columns=['VALUE'])
  

• 升级到较新版本的Pandas
  Pandas在1.5+版本中对groupby的性能做了针对性优化，尤其是MultiIndex的分组聚合，升级后groupby.sum()的性能会更接近原sum(level=...)的水平：

  pip install --upgrade pandas
  

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