python - 提高pandas groupby的性能

Question

我有一个用 Python 编写的机器学习应用程序，其中包括一个数据处理步骤。当我写它的时候，我最初是在 Pandas DataFrames 上进行数据处理，但是当这导致糟糕的性能时，我最终使用普通 Python 重写了它，用 for 循环代替矢量化操作，用列表和字典代替 DataFrames 和 Series。令我惊讶的是，用 vanilla Python 编写的代码的性能最终远 高于使用 Pandas 编写的代码。

由于我的手写数据处理代码比原来的 Pandas 代码更大更困惑，我还没有完全放弃使用 Pandas，目前我正在尝试优化 Pandas 代码，但没有取得太大成功。

数据处理步骤的核心包括以下内容:我首先将行分成几组，因为数据由几千个时间序列组成(每个“个人”)，然后进行相同的数据处理在每个组上:大量总结，将不同的列组合成新的列等。

我使用 Jupyter Notebook 的 lprun 分析了我的代码，大部分时间花在了以下和其他类似的行上:

grouped_data = data.groupby('pk')
data[[v + 'Diff' for v in val_cols]] = grouped_data[val_cols].transform(lambda x: x - x.shift(1)).fillna(0)
data[[v + 'Mean' for v in val_cols]] = grouped_data[val_cols].rolling(4).mean().shift(1).reset_index()[val_cols]
(...)

...向量化和非向量化处理的混合。我知道非矢量化操作不会比我手写的 for 循环快，因为这基本上就是它们的内幕，但它们怎么会慢得多呢？我们正在谈论我的手写代码和 Pandas 代码之间的性能下降 10-20 倍。

我是不是做错了非常非常错误的事情？

Answer 1

不，我认为您不应该放弃 pandas。肯定有更好的方法来做你想做的事。诀窍是尽可能避免任何形式的apply/transform。像躲避瘟疫一样躲避它们。它们基本上是作为 for 循环实现的，因此您不妨直接使用 python for 循环，它以 C 速度运行并为您提供更好的性能。

真正的速度增益是您摆脱循环并使用 pandas 的隐式向量化操作的函数。例如，您的第一行代码可以大大简化，我很快就会向您展示。

在这篇文章中，我概述了设置过程，然后，针对您问题中的每一行，提供改进，以及时间和正确性的并排比较。

设置

data = {'pk' : np.random.choice(10, 1000)} 
data.update({'Val{}'.format(i) : np.random.randn(1000) for i in range(100)})

df = pd.DataFrame(data)

g = df.groupby('pk')
c = ['Val{}'.format(i) for i in range(100)]

---

transform + sub + shift → diff

您的第一行代码可以用简单的diff 语句替换:

v1 = df.groupby('pk')[c].diff().fillna(0)

完整性检查

v2 = df.groupby('pk')[c].transform(lambda x: x - x.shift(1)).fillna(0)

np.allclose(v1, v2)
True

性能

%timeit df.groupby('pk')[c].transform(lambda x: x - x.shift(1)).fillna(0)
10 loops, best of 3: 44.3 ms per loop

%timeit df.groupby('pk')[c].diff(-1).fillna(0)
100 loops, best of 3: 9.63 ms per loop

---

去除多余的索引操作

就您的第二行代码而言，我认为没有太多改进空间，尽管您可以摆脱 reset_index() + [val_cols] 如果您的 groupby 语句未将 pk 视为索引，则调用:

g = df.groupby('pk', as_index=False)

您的第二行代码将缩减为:

v3 = g[c].rolling(4).mean().shift(1)

完整性检查

g2 = df.groupby('pk')
v4 = g2[c].rolling(4).mean().shift(1).reset_index()[c]

np.allclose(v3.fillna(0), v4.fillna(0))
True

性能

%timeit df.groupby('pk')[c].rolling(4).mean().shift(1).reset_index()[c]
10 loops, best of 3: 46.5 ms per loop

%timeit df.groupby('pk', as_index=False)[c].rolling(4).mean().shift(1)
10 loops, best of 3: 41.7 ms per loop

请注意，时间因机器而异，因此请务必彻底测试您的代码，以确保您的数据确实有所改善。

虽然这次的差别不大，但您可以欣赏到您可以做出改进的事实!这可能会对更大的数据产生更大的影响。

---

后记

总而言之，大多数操作之所以缓慢，是因为它们可以加速。关键是摆脱任何不使用矢量化的方法。

为此，有时走出 pandas 空间进入 numpy 空间是有益的。对 numpy 数组的操作或使用 numpy 往往比 pandas 等价物快得多(例如，np.sum 比 pd.DataFrame.sum 快，而 np .where 比 pd.DataFrame.where 等更快。

有时候，循环是不可避免的。在这种情况下，您可以创建一个基本的循环函数，然后可以使用 numba 或 cython 对其进行矢量化。这方面的例子在这里 Enhancing Performance , 直接来自马口。

在其他情况下，您的数据太大而无法合理地放入 numpy 数组中。在这种情况下，是时候放弃并切换到 dask 了。或 spark ，两者都提供了用于处理大数据的高性能分布式计算框架。