pandas groupby apply 真的很慢

Posted 2023-03-11

【中文标题】pandas groupby apply 真的很慢

【英文标题】：pandas groupby apply is really slow

【发布时间】：2018-04-17 19:14:53

【问题描述】：

当我打电话给df.groupby([...]).apply(lambda x: ...) 时，表现很糟糕。有没有更快/更直接的方法来完成这个简单的查询？

为了证明我的观点，这里有一些设置 DataFrame 的代码：

import pandas as pd

df = pd.DataFrame(data=
    'ticker': ['AAPL','AAPL','AAPL','IBM','IBM','IBM'],
       'side': ['B','B','S','S','S','B'],
       'size': [100, 200, 300, 400, 100, 200],
      'price': [10.12, 10.13, 10.14, 20.3, 20.2, 20.1])


    price   side     size   ticker
0   10.12   B        100    AAPL
1   10.13   B        200    AAPL
2   10.14   S        300    AAPL
3   20.30   S        400    IBM
4   20.20   S        100    IBM
5   20.10   B        200    IBM

现在这是我需要加快速度的非常慢的部分：

%timeit avgpx = df.groupby(['ticker','side']) \
.apply(lambda group: (group['size'] * group['price']).sum() / group['size'].sum())

3.23 ms ± 148 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

这会产生正确的结果，但正如您在上面看到的，需要很长时间（3.23ms 看起来并不多，但这只有 6 行......当我在真实数据集上使用它时，它需要很长时间）。

ticker  side
AAPL    B       10.126667
        S       10.140000
IBM     B       20.100000
        S       20.280000
dtype: float64

【问题讨论】：

在开始代码优化之前了解您的目标是非常有帮助。

我正在尝试做的只是在那里......计算每个股票/每边的平均价格。当我使用 DataFrame 中的 40K 行执行此操作时……性能非常慢（几秒钟）。我想我可能只是使用 lambda 做错了？

【参考方案1】：

您可以通过预先计算产品并摆脱apply 来节省一些时间。

df['scaled_size'] = df['size'] * df['price']
g = df.groupby(['ticker', 'side'])

g['scaled_size'].sum() / g['size'].sum()

ticker  side
AAPL    B       10.126667
        S       10.140000
IBM     B       20.100000
        S       20.280000
dtype: float64

100 loops, best of 3: 2.58 ms per loop

健全性检查

df.groupby(['ticker','side']).apply(
    lambda group: (group['size'] * group['price']).sum() / group['size'].sum())

ticker  side
AAPL    B       10.126667
        S       10.140000
IBM     B       20.100000
        S       20.280000
dtype: float64

100 loops, best of 3: 5.02 ms per loop

摆脱apply 似乎会导致我的机器上的速度提高了 2 倍。