Pandas：如何使DataFrame的apply方法更快？

Question

考虑这个pandas示例，我正在使用apply和一个lambda函数计算列C，通过将A与B相乘并乘以一个float，只有在满足特定条件时才进行计算：

import pandas as pd
df = pd.DataFrame({'A':[1,2,3,4,5,6,7,8,9],'B':[9,8,7,6,5,4,3,2,1]})

df['C'] = df.apply(lambda x: x.A if x.B > 5 else 0.1*x.A*x.B, axis=1)

预期结果将是：

   A  B    C
0  1  9  1.0
1  2  8  2.0
2  3  7  3.0
3  4  6  4.0
4  5  5  2.5
5  6  4  2.4
6  7  3  2.1
7  8  2  1.6
8  9  1  0.9

问题在于这段代码运行速度缓慢，而我需要对一个大约有 5600 万行的数据框执行此操作。

上述 lambda 操作的 %timeit 结果为：

1000 loops, best of 3: 1.63 ms per loop

根据我在大型数据框上进行计算时的计算时间和内存使用情况，我推测这个操作在执行计算时使用了中间序列。

我尝试以不同的方式来表述它，包括使用临时列，但我想到的每个替代方案都更慢。

是否有一种不同且更快的方法可以得到我需要的结果，例如使用numpy？

- Khris

---

你应该查阅numpy.where。 - IanS

Answer 1

为了提高性能，你最好使用NumPy数组并使用np.where函数 -

a = df.values # Assuming you have two columns A and B
df['C'] = np.where(a[:,1]>5,a[:,0],0.1*a[:,0]*a[:,1])

---

运行时间测试

def numpy_based(df):
    a = df.values # Assuming you have two columns A and B
    df['C'] = np.where(a[:,1]>5,a[:,0],0.1*a[:,0]*a[:,1])

时间 -

In [271]: df = pd.DataFrame(np.random.randint(0,9,(10000,2)),columns=[['A','B']])

In [272]: %timeit numpy_based(df)
1000 loops, best of 3: 380 µs per loop

In [273]: df = pd.DataFrame(np.random.randint(0,9,(10000,2)),columns=[['A','B']])

In [274]: %timeit df['C'] = df.A.where(df.B.gt(5), df[['A', 'B']].prod(1).mul(.1))
100 loops, best of 3: 3.39 ms per loop

In [275]: df = pd.DataFrame(np.random.randint(0,9,(10000,2)),columns=[['A','B']])

In [276]: %timeit df['C'] = np.where(df['B'] > 5, df['A'], 0.1 * df['A'] * df['B'])
1000 loops, best of 3: 1.12 ms per loop

In [277]: df = pd.DataFrame(np.random.randint(0,9,(10000,2)),columns=[['A','B']])

In [278]: %timeit df['C'] = np.where(df.B > 5, df.A, df.A.mul(df.B).mul(.1))
1000 loops, best of 3: 1.19 ms per loop

仔细观察

让我们仔细观察NumPy的数值计算能力，并将其与pandas进行比较 -

# Extract out as array (its a view, so not really expensive
#   .. as compared to the later computations themselves)

In [291]: a = df.values 

In [296]: %timeit df.values
10000 loops, best of 3: 107 µs per loop

案例#1：使用NumPy数组并使用numpy.where：

In [292]: %timeit np.where(a[:,1]>5,a[:,0],0.1*a[:,0]*a[:,1])
10000 loops, best of 3: 86.5 µs per loop

再次，将值赋给新的列df['C']也不会很昂贵 -

In [300]: %timeit df['C'] = np.where(a[:,1]>5,a[:,0],0.1*a[:,0]*a[:,1])
1000 loops, best of 3: 323 µs per loop

案例#2：与 Pandas 数据框架一起使用其 .where 方法（不使用 NumPy）

In [293]: %timeit df.A.where(df.B.gt(5), df[['A', 'B']].prod(1).mul(.1))
100 loops, best of 3: 3.4 ms per loop

案例 #3：使用pandas数据帧（不使用NumPy数组），但使用numpy.where -

In [294]: %timeit np.where(df['B'] > 5, df['A'], 0.1 * df['A'] * df['B'])
1000 loops, best of 3: 764 µs per loop

案例#4：再次使用pandas dataframe（而不是NumPy数组），但使用numpy.where -

In [295]: %timeit np.where(df.B > 5, df.A, df.A.mul(df.B).mul(.1))
1000 loops, best of 3: 830 µs per loop

- Divakar

---

你比我先说了，但是numpy.where与pandas系列很搭配，我相信我的版本更易读 ;) - IanS

@IanS OP要求更快。这将所有内容都带入numpy，从而实现更高效的操作。 - piRSquared

@IanS NumPy 在数字计算方面表现出色，至少我知道它与数据框架很搭配！;) - Divakar

我从这个解决方案中得到了“1000次循环，3次中的最佳：每个循环255微秒”的结果，非常感谢。 - Khris

我简直不敢相信，对我的 5600 万行数据框进行的操作只花了大约 1 秒钟。 - Khris

Answer 2

使用 pd.Series.where 实现纯粹的 pandas

df['C'] = df.A.where(df.B.gt(5), df[['A', 'B']].prod(1).mul(.1))

   A  B    C
0  1  9  1.0
1  2  8  2.0
2  3  7  3.0
3  4  6  4.0
4  5  5  2.5
5  6  4  2.4
6  7  3  2.1
7  8  2  1.6
8  9  1  0.9

- piRSquared

---

乘法比加法更快吗？ - jezrael

也许...勉强 - piRSquared

@IanS 因为它是纯粹的pandas...速度较慢。这就是为什么我们都转向numpy的原因。但Divakar比我们更早一步。我提供了这个答案，因为它与众不同。几乎没有人从pandas的角度使用where。它很有趣，因为它假设条件为True时采用现有的值，否则采用替代值。 - piRSquared

Answer 3

Pandas是一个很好的数据操作工具，但默认情况下只在单个CPU核心上运行。此外，Pandas被构建为在整个列或数据集上一次性运行向量化API函数，但apply运行自定义用户代码。其他答案避免使用带有自定义代码的apply，但这在一般情况下可能不可行/实用。如果使用apply处理大型数据集是您的痛点，则应考虑加速和扩展解决方案，例如Bodo。Bodo直接编译您的apply代码以优化它，而Pandas无法做到这一点。除了向量化您的代码外，Bodo还提供自动并行化。您可以使用Bodo社区版（免费使用）在4个核心上运行您的代码。这是Bodo安装说明的链接：https://docs.bodo.ai/latest/source/installation_and_setup/install.html 我生成了一个类似于您的数据集，但包含2000万行，并在一个核心上使用常规Pandas运行代码，在4个核心上使用Bodo运行代码。使用常规Pandas，运行您的代码需要约6.5分钟，而使用Bodo的社区版则只需约半秒钟。

#data generation
import numpy as np
import pandas as pd

df = pd.DataFrame(np.random.randint(1,10,size=(20000000, 2)), columns=list('AB'))
df.to_parquet("data.pq")

常规的Pandas：

import pandas as pd
import time

start = time.time()

df = pd.read_parquet("data.pq")
df['C'] = df.apply(lambda x: x.A if x.B > 5 else 0.1*x.A*x.B, axis=1)

end = time.time()
print("computation time: ", end - start)

print(df.head())

output:
computation time:  378.3832001686096
   A  B    C
0  3  5  1.5
1  8  6  8.0
2  1  7  1.0
3  8  1  0.8
4  4  8  4.0

使用Bodo：

%%px

import pandas as pd
import time
import bodo

@bodo.jit(distributed = ['df'])
def apply():
    start = time.time()
    df = pd.read_parquet("data.pq")
    df['C'] = df.apply(lambda x: x.A if x.B > 5 else 0.1*x.A*x.B, axis=1)
    end = time.time()
    print("computation time: ", end - start)
    print(df.head())
    return df
df = apply()

output:
[stdout:0] 
computation time:  0.3610380489999443
   A  B    C
0  3  5  1.5
1  8  6  8.0
2  1  7  1.0
3  8  1  0.8
4  4  8  4.0

免责声明：我在Bodo.ai担任数据科学家倡导者。

- Niyousha Mohammadshafie

---

我在我的问题上尝试了bodo，但出现了错误：“DataFrame.loc[] getitem (基于位置的索引)使用Tuple(unicode_type, Literalint)尚不支持。” - West

未能成功工作。BodoError：无法从JIT函数调用非JIT函数“loads”（转换为JIT或使用objmode）。 - Alekhya Reddy

Answer 4

使用numpy.where函数：

df['C'] = numpy.where(df['B'] > 5, df['A'], 0.1 * df['A'] * df['B'])

- IanS

Answer 5

使用：

df['C'] = np.where(df.B > 5, df.A, df.A.mul(df.B).mul(.1))
print (df)
   A  B    C
0  1  9  1.0
1  2  8  2.0
2  3  7  3.0
3  4  6  4.0
4  5  5  2.5
5  6  4  2.4
6  7  3  2.1
7  8  2  1.6
8  9  1  0.9

- jezrael

---

“mul”和“*”其实没什么区别，是吧？ ;) - IanS

我做了一些研究，发现如果使用 df.A*df.B 和 df.A.mul(df.B)，那么 mul 更快。但如果乘以常数，则是相同的。 - jezrael

@IanS 同时，这也方便链式调用。 - piRSquared

嗯，这里使用 * 更快，可能的原因是 np.where 与 numpy 数组一起使用。 - jezrael

这也是提到另一次jezrael比我先回答的参考，我们之间唯一的区别是我使用了“/”，而他使用了“div”，他指出这几乎没有什么区别。 - IanS

@jezrael 我也很惊讶，我也预计mul会更快（即使不是太多）。 - IanS