python - 如何加速 pandas groupby - 应用函数与 R 的 data.table 相媲美

Question

我有这样的数据

   location  sales  store
0        68    583     17
1        28    857      2
2        55    190     59
3        98    517     64
4        94    892     79
...

对于每个独特的组合(位置、商店)，都有 1 次或多次销售。我想添加一列 pcnt_sales，它显示给定行中的销售额占该对(位置、商店)总销售额的百分比。

   location  sales  store  pcnt_sales
0        68    583     17    0.254363
1        28    857      2    0.346543
2        55    190     59    1.000000
3        98    517     64    0.272105
4        94    892     79    1.000000
...

这有效，但是很慢

import pandas as pd
import numpy as np

df = pd.DataFrame({'location':np.random.randint(0, 100, 10000), 'store':np.random.randint(0, 100, 10000), 'sales': np.random.randint(0, 1000, 10000)})

import timeit
start_time = timeit.default_timer()
df['pcnt_sales'] = df.groupby(['location', 'store'])['sales'].apply(lambda x: x/x.sum())
print(timeit.default_timer() - start_time)  # 1.46 seconds

相比之下，R 的 data.table 做起来 super 快

library(data.table)

dt <- data.table(location=sample(100, size=10000, replace=TRUE), store=sample(100, size=10000, replace=TRUE), sales=sample(1000, size=10000, replace=TRUE))

ptm <- proc.time()
dt[, pcnt_sales:=sales/sum(sales), by=c("location", "store")]
proc.time() - ptm  # 0.007 seconds

我如何在 Pandas 中高效地执行此操作(特别是考虑到我的真实数据集有数百万行)？

Answer 1

为了性能你想避免apply。您可以使用 transform 将 groupby 的结果扩展为原始索引，此时除法将以矢量化速度工作:

>>> %timeit df['pcnt_sales'] = df.groupby(['location', 'store'])['sales'].apply(lambda x: x/x.sum())
1 loop, best of 3: 2.27 s per loop
>>> %timeit df['pcnt_sales2'] = (df["sales"] /
            df.groupby(['location', 'store'])['sales'].transform(sum))
100 loops, best of 3: 6.25 ms per loop
>>> df["pcnt_sales"].equals(df["pcnt_sales2"])
True