不,我认为你不应该放弃使用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
+ 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
减少冗余的索引操作
就你的第二行代码而言,我没有看到太多改进的空间,但是如果你的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.sum
比pd.DataFrame.sum
快,np.where
比pd.DataFrame.where
快,等等)。
有时,无法避免使用循环。在这种情况下,可以创建一个基本的循环函数,然后使用numba或cython对其进行矢量化。这些示例在Enhancing Performance中都有介绍。
在其他一些情况下,您的数据太大而无法合理地装入numpy数组中。此时,是时候放弃并切换到dask
或spark
了,它们都提供了用于处理大数据的高性能分布式计算框架。