Pandas: 如何更快地在数据帧上应用?

Pandas: How to make apply on dataframe faster?

考虑这个 pandas 示例,我通过将 AB 相乘来计算列 C,如果满足特定条件则计算 float使用 applylambda 函数:

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?

为了提高性能,您最好使用 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 数据框(没有 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

使用numpy.where:

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

使用:

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

纯粹pandas
使用 pd.Series.where

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

Pandas 是一个很好的数据操作工具,但默认情况下 运行 在单个 CPU 核心上。此外,Pandas 构建为 运行 向量化 API 一次扫描整个列或数据集的函数,但 apply 运行 自定义用户代码。其他答案避免将 apply 与自定义代码一起使用,但这通常可能不是 possible/practical。如果使用 apply 处理大型数据集对您来说是一个痛点,您应该考虑加速和缩放解决方案,例如 Bodo。 Bodo 直接编译您的 apply 代码,以 Pandas 无法做到的方式对其进行优化。除了向量化代码之外,Bodo 还提供自动并行化。使用 Bodo 社区版(免费使用),您可以 运行 您的代码最多 4 个内核。这是 Bodo 安装说明的 link: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

与博多:

%%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 担任数据科学家倡导者。