加速 Pandas 中 Apply 函数百倍的技巧

引言

尽管出现了像 dask、cudf 等包，加速了数据处理过程，但并非每个人都有强大的 GPU 资源。因此，许多人仍然依赖于 Pandas 工具包。然而，Pandas 中的 apply 函数通常效率较低。本文将介绍如何通过一些技巧，将 apply 函数的执行速度提升 600 倍。

实验对比

1. Apply (基准)

原始的 apply 函数处理一个简单问题需要 18.4 秒。

<code>import pandas as pd import numpy as np df = pd.DataFrame(np.random.randint(0, 11, size=(1000000, 5)), columns=('a','b','c','d','e')) def func(a, b, c, d, e): if e == 10: return c * d elif 5 <= e < 10: return c + d elif e < 5: return a + b %%time df['new'] = df.apply(lambda x: func(x['a'], x['b'], x['c'], x['d'], x['e']), axis=1)</code>

执行时间：18.4 秒

2. 使用 Swift 加速

使用 Swift 可以并行处理数据，相同的操作在作者的机器上只需要 7.67 秒。

<code>%%time import swifter df['new'] = df.swifter.apply(lambda x: func(x['a'], x['b'], x['c'], x['d'], x['e']), axis=1)</code>

执行时间：7.67 秒

3. 向量化

将函数向量化是 Pandas 和 Numpy 中最快的方法之一。将操作向量化后，执行时间缩短为 421 毫秒。

<code>%%time df['new'] = df['c'] * df['d'] # 默认情况下 e == 10 mask = df['e'] < 10 df.loc[mask, 'new'] = df['c'] + df['d'] mask = df['e'] < 5 df.loc[mask, 'new'] = df['a'] + df['b']</code>

执行时间：421 毫秒

4. 类别转化 + 向量化

首先将数据转化为 int16 类型，然后进行相同的向量化操作，时间缩短为 116 毫秒。

<code>for col in ('a', 'b', 'c', 'd'): df[col] = df[col].astype(np.int16) %%time df['new'] = df['c'] * df['d'] # 默认情况下 e == 10 mask = df['e'] < 10 df.loc[mask, 'new'] = df['c'] + df['d'] mask = df['e'] < 5 df.loc[mask, 'new'] = df['a'] + df['b']</code>

执行时间：116 毫秒

5. 转化为 values 处理

尽可能将数据转化为 .values，再进行操作。这样可以更快地执行向量化操作。转化为 values 后，操作时间缩短为 74.9 毫秒。

<code>%%time df['new'] = df['c'].values * df['d'].values # 默认情况下 e == 10 mask = df['e'].values < 10 df.loc[mask, 'new'] = df['c'] + df['d'] mask = df['e'].values < 5 df.loc[mask, 'new'] = df['a'] + df['b']</code>

执行时间：74.9 毫秒

实验汇总

通过上述技巧，简单的 apply 函数的执行效率提升了数百倍：

• Apply: 18.4 秒
• Apply + Swifter: 7.67 秒
• Pandas 向量化: 421 毫秒
• Pandas 向量化 + 数据类型优化: 116 毫秒
• Pandas 向量化 + values + 数据类型优化: 74.9 毫秒

以上是针对简单问题的优化，但在实际应用中，这些技巧也许会带来更大的效率提升。