提高pandas groupby的性能

39

我有一个用Python编写的机器学习应用程序,其中包括数据处理步骤。当我最初使用Pandas DataFrames进行数据处理时,性能非常糟糕,因此我最终改用普通的Python代码进行重写,使用for循环代替向量化操作,使用列表和字典代替DataFrames和Series。令我惊讶的是,使用普通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循环更快,因为它们本质上就是在底层执行for循环,但是它们为什么可以慢那么多?我们谈论的是我的手写代码和Pandas代码之间性能下降10-20倍。

我是不是做错了什么非常严重的事情?

1个回答

50

不,我认为你不应该放弃使用pandas。肯定有更好的方法可以达成你想要的目标。关键是尽可能避免使用任何形式的apply/transform。像瘟疫一样避免它们。它们基本上是实现为for循环,因此您最好直接使用在C速度下运行并提供更好性能的Python for循环。

真正的加速是消除循环并使用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 + shiftdiff

你的第一行代码可以替换为一个简单的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

减少冗余的索引操作

就你的第二行代码而言,我没有看到太多改进的空间,但是如果你的groupby语句没有将pk作为索引考虑,你可以摆脱reset_index() + [val_cols]调用:

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.sumpd.DataFrame.sum快,np.wherepd.DataFrame.where快,等等)。

有时,无法避免使用循环。在这种情况下,可以创建一个基本的循环函数,然后使用numba或cython对其进行矢量化。这些示例在Enhancing Performance中都有介绍。

在其他一些情况下,您的数据太大而无法合理地装入numpy数组中。此时,是时候放弃并切换到daskspark了,它们都提供了用于处理大数据的高性能分布式计算框架。


1
感谢您的出色回答,Coldspeed。我已经用矢量化操作替换了大部分剩余的apply和transforms,虽然现在比以前快了,但仍然比纯Python代码慢。话虽如此,我认为真正的原因是无法从我的问题中推断出来的 -- 即组非常丰富而且非常小(大约有400,000个组,每个组平均有约10行)。我可能会稍后发布一个新问题,但现在我将授予您奖励。 - haroba
1
@haroba 很抱歉听到这个消息!请确保在下一个问题中提供足够的细节。祝你好运。 - cs95

网页内容由stack overflow 提供, 点击上面的
可以查看英文原文,
原文链接