在 apply 中向量化一个非常简单的 pandas lambda 函数

Vectorizing a very simple pandas lambda function in apply

pandas apply/map 是我的克星,即使在小数据集上也可能慢得令人痛苦。下面是一个非常简单的例子,其中速度有近 3 个数量级的差异。下面我创建了一个包含 100 万个值的 Series,并且只想将大于 .5 的值映射到 'Yes',将小于 .5 的值映射到 'No'。我如何对其进行矢量化或显着加快速度?

ser = pd.Series(np.random.rand(1000000))

# vectorized and fast
%%timeit
ser > .5

1000 个循环,3 个循环中的最佳循环:每个循环 477 微秒

%%timeit
ser.map(lambda x: 'Yes' if x > .5 else 'No')

1 个循环,3 个循环中的最佳:每个循环 255 毫秒

np.where(cond, A, B)A if cond else B:

的向量化等价物
import numpy as np
import pandas as pd
ser = pd.Series(np.random.rand(1000000))
mask = ser > 0.5
result = pd.Series(np.where(mask, 'Yes', 'No'))
expected = ser.map(lambda x: 'Yes' if x > .5 else 'No')
assert result.equals(expected)

In [77]: %timeit mask = ser > 0.5
1000 loops, best of 3: 1.44 ms per loop

In [76]: %timeit np.where(mask, 'Yes', 'No')
100 loops, best of 3: 14.8 ms per loop

In [73]: %timeit pd.Series(np.where(mask, 'Yes', 'No'))
10 loops, best of 3: 86.5 ms per loop

In [74]: %timeit ser.map(lambda x: 'Yes' if x > .5 else 'No')
1 loop, best of 3: 223 ms per loop

由于这个系列只有两个值,您可以考虑使用 Categorical 代替:

In [94]: cat = pd.Categorical.from_codes(codes=mask.astype(int), categories=['Yes', 'No']); cat
Out[94]: 
[No, Yes, No, Yes, Yes, ..., Yes, No, Yes, Yes, No]
Length: 1000000
Categories (2, object): [Yes, No]

In [95]: %timeit pd.Categorical.from_codes(codes=mask.astype(int), categories=['Yes', 'No']); cat
100 loops, best of 3: 6.26 ms per loop

这不仅速度更快,而且内存效率更高,因为它避免了创建字符串数组。类别代码是映射到类别的整数数组:

In [96]: cat.codes
Out[96]: array([1, 0, 1, ..., 0, 0, 1], dtype=int8)

In [97]: cat.categories
Out[99]: Index(['Yes', 'No'], dtype='object')