Pandas apply和np.vectorize在从现有列创建新列方面的性能比较

Question

我正在使用Pandas数据框架，并希望创建一个新列作为现有列的函数。我没有看到关于df.apply()和np.vectorize()速度差异的好讨论，所以我想在这里问一下。

Pandas的apply()函数较慢。根据我测量得出的结果（在下面的一些实验中显示），使用np.vectorize()比使用DataFrame函数apply()快25倍或更多，至少在我的2016 MacBook Pro上是这样。这是预期结果吗？为什么？

例如，假设我有以下具有N行的数据表：

N = 10
A_list = np.random.randint(1, 100, N)
B_list = np.random.randint(1, 100, N)
df = pd.DataFrame({'A': A_list, 'B': B_list})
df.head()
#     A   B
# 0  78  50
# 1  23  91
# 2  55  62
# 3  82  64
# 4  99  80

假设我想根据两列“ A”和“ B”创建一个新列。 在下面的示例中，我将使用一个简单的函数“ divide（）”。 要应用该函数，我可以使用“ df.apply（）”或“ np.vectorize（）”：

Suppose further that I want to create a new column as a function of the two columns A and B. In the example below, I'll use a simple function divide(). To apply the function, I can use either df.apply() or np.vectorize():

def divide(a, b):
    if b == 0:
        return 0.0
    return float(a)/b

df['result'] = df.apply(lambda row: divide(row['A'], row['B']), axis=1)

df['result2'] = np.vectorize(divide)(df['A'], df['B'])

df.head()
#     A   B    result   result2
# 0  78  50  1.560000  1.560000
# 1  23  91  0.252747  0.252747
# 2  55  62  0.887097  0.887097
# 3  82  64  1.281250  1.281250
# 4  99  80  1.237500  1.237500

如果我将N增加到一百万或更多这样的真实世界规模，那么我发现np.vectorize()比df.apply()快25倍甚至更多。

以下是完整的基准测试代码：

import pandas as pd
import numpy as np
import time

def divide(a, b):
    if b == 0:
        return 0.0
    return float(a)/b

for N in [1000, 10000, 100000, 1000000, 10000000]:    

    print ''
    A_list = np.random.randint(1, 100, N)
    B_list = np.random.randint(1, 100, N)
    df = pd.DataFrame({'A': A_list, 'B': B_list})

    start_epoch_sec = int(time.time())
    df['result'] = df.apply(lambda row: divide(row['A'], row['B']), axis=1)
    end_epoch_sec = int(time.time())
    result_apply = end_epoch_sec - start_epoch_sec

    start_epoch_sec = int(time.time())
    df['result2'] = np.vectorize(divide)(df['A'], df['B'])
    end_epoch_sec = int(time.time())
    result_vectorize = end_epoch_sec - start_epoch_sec


    print 'N=%d, df.apply: %d sec, np.vectorize: %d sec' % \
            (N, result_apply, result_vectorize)

    # Make sure results from df.apply and np.vectorize match.
    assert(df['result'].equals(df['result2']))

结果如下：

N=1000, df.apply: 0 sec, np.vectorize: 0 sec

N=10000, df.apply: 1 sec, np.vectorize: 0 sec

N=100000, df.apply: 2 sec, np.vectorize: 0 sec

N=1000000, df.apply: 24 sec, np.vectorize: 1 sec

N=10000000, df.apply: 262 sec, np.vectorize: 4 sec

如果np.vectorize()通常比df.apply()更快，为什么就没有更多提到np.vectorize()呢？我只看到与df.apply()相关的StackOverflow帖子，例如：

基于其他列的值创建新列

如何在Pandas中使用“apply”函数处理多列？

如何将函数应用于Pandas数据框的两列

- stackoverflowuser2010

---

@PMende 但是 np.vectorize 并不是矢量化的。这是一个众所周知的用词错误。 - roganjosh

@PMende 我不知道你在那里说明什么。iterrows不是Python方法。你应该更关心的是库在底层做了什么。 - roganjosh

@jpp 如果您想要的话，可以在实际的numpy数组上做同样的事情。例如，A_arr, B_arr = df['A'].values, df['B'].values 然后调用 %timeit for a, b in zip(A_arr, B_arr): foo(a, b) 将计时缩短到：1.88 s ± 57 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)。这里发生的事情不仅仅是“python循环”。 - PMende

@PMende，好的，我并没有暗示其他意思。你不应该从时间上推断出对实现的看法。是的，这些信息很有见地。但它们可能会让你假设一些不正确的事情。 - jpp

请尝试使用pandas的.str属性。在许多情况下，它们比列表解析慢。我们经常过于自以为是。 - roganjosh

Answer 1

我将从强调Pandas和NumPy数组的威力开始，这是由于对数字数组进行高性能向量化计算所带来的。整个向量化计算的重点在于通过将计算移动到高度优化的C代码并利用连续的内存块来避免Python级别的循环。

Python级别的循环

现在我们可以看一些时间。下面是所有产生包含相同值的pd.Series、np.ndarray或list对象的Python级别循环。为了分配到数据框中的系列，结果是可比较的。

# Python 3.6.5, NumPy 1.14.3, Pandas 0.23.0

np.random.seed(0)
N = 10**5

%timeit list(map(divide, df['A'], df['B']))                                   # 43.9 ms
%timeit np.vectorize(divide)(df['A'], df['B'])                                # 48.1 ms
%timeit [divide(a, b) for a, b in zip(df['A'], df['B'])]                      # 49.4 ms
%timeit [divide(a, b) for a, b in df[['A', 'B']].itertuples(index=False)]     # 112 ms
%timeit df.apply(lambda row: divide(*row), axis=1, raw=True)                  # 760 ms
%timeit df.apply(lambda row: divide(row['A'], row['B']), axis=1)              # 4.83 s
%timeit [divide(row['A'], row['B']) for _, row in df[['A', 'B']].iterrows()]  # 11.6 s

一些要点：

1. 基于元组的方法（前4个）比基于pd.Series的方法（后3个）更高效。
2. np.vectorize、列表推导式+zip和map方法，即前三者，性能大致相同。这是因为它们使用了tuple并绕过了一些从pd.DataFrame.itertuples中产生的Pandas开销。
3. 使用raw=True与不使用相比，pd.DataFrame.apply可以显著提高速度。该选项将NumPy数组传递给自定义函数，而不是pd.Series对象。

pd.DataFrame.apply: 只是另一个循环

要确切地看到Pandas传递的对象，您只需轻松修改函数即可：

def foo(row):
    print(type(row))
    assert False  # because you only need to see this once
df.apply(lambda row: foo(row), axis=1)

输出：

输出：<class 'pandas.core.series.Series'>。与NumPy数组相比，创建、传递和查询Pandas系列对象会带来显着的开销。这不应该令人惊讶：Pandas系列包括相当数量的脚手架来保存索引、值、属性等。

使用raw=True再次进行相同的练习，你将看到<class 'numpy.ndarray'>。所有这些都在文档中有描述，但看到它更加令人信服。

np.vectorize:虚假向量化

np.vectorize的文档有以下注释：

向量化函数像Python map函数一样评估pyfunc在输入数组的连续元组上，除了它使用NumPy的广播规则。

“广播规则”在这里不相关，因为输入数组具有相同的维度。与map并行处理类似，因为上面的map版本几乎具有相同的性能。源代码显示了正在发生的情况：np.vectorize通过np.frompyfunc将您的输入函数转换为通用函数（"ufunc"）。有一些优化，例如缓存，可以带来一些性能提升。

简而言之，np.vectorize执行Python级循环应该做的事情，但是pd.DataFrame.apply增加了笨重的开销。没有像numba那样的JIT编译。它只是一个方便。

真正的向量化：你应该使用什么

为什么以上差异没有被提及？因为真正向量化计算的性能使它们不相关：

%timeit np.where(df['B'] == 0, 0, df['A'] / df['B'])       # 1.17 ms
%timeit (df['A'] / df['B']).replace([np.inf, -np.inf], 0)  # 1.96 ms

是的，这比上述循环解决方案中最快的解决方案快约40倍。这两种方法都可以接受。在我看来，第一种方法简洁、易读且高效。仅在性能关键且这是瓶颈的一部分时，才查看其他方法，例如下面的numba。

numba.njit: 更高的效率

当循环被认为可行时，通常会通过使用底层NumPy数组的numba进行优化，以尽可能多地将代码移植到C中。

的确，numba将性能提高到了微秒级别。除非进行一些繁琐的工作，否则将很难获得比此更高效的结果。

from numba import njit

@njit
def divide(a, b):
    res = np.empty(a.shape)
    for i in range(len(a)):
        if b[i] != 0:
            res[i] = a[i] / b[i]
        else:
            res[i] = 0
    return res

%timeit divide(df['A'].values, df['B'].values)  # 717 µs

使用@njit(parallel=True)可能会为更大的数组提供进一步的提升。

---

¹ 数字类型包括: int, float, datetime, bool, category。它们不包括 object dtype，并且可以保存在连续的内存块中。

² NumPy操作比Python高效的至少有两个原因：

• Python中的所有内容都是对象。这包括了数字，这与C不同。因此，Python类型具有C类型没有的开销。
• NumPy方法通常基于C。此外，在可能的情况下使用优化算法。

- jpp

---

@PMende，你在这里错过了重点。使用apply创建的系列是逐行的，即每个系列都有来自A和B的一个元素。而使用列表推导式，df ['A']/df['B']是*唯一的2个系列并且它们并没有以任何方式被“创建”，它们已经存在。zip可以与生成一个tuple进行比较，后者要便宜得多。 - jpp

是的，我尝试了 10**8 和 10**9，发现它更快。 - Sheldore

你对 b[i] != 0 进行了双重检查。正常的 Python 和 Numba 行为是检查 0 并抛出错误。这可能会破坏任何 SIMD 向量化，并且通常对执行速度有很大影响。但是，你可以在 Numba 中通过 @njit(error_model='numpy') 来更改这种情况，以避免除以 0 的双重检查。同时建议使用 np.empty 分配内存，并在 else 语句中将结果设置为 0。 - max9111

错误模型numpy在除以0时使用处理器提供的内容-> NaN。至少在Numba 0.41dev中，两个版本都使用SIMD向量化。您可以按照此处描述的方式进行检查https://numba.pydata.org/numba-doc/dev/user/faq.html（1.16.2.3.为什么我的循环没有向量化？）我会简单地向您的函数添加一个else语句（res [i] = 0.），并使用np.empty分配内存。这应该与error_model ='numpy'一起将性能提高约20％。在旧版Numba上，对性能的影响更大... - max9111

@stackoverflowuser2010，没有针对“任意函数”的通用答案。您必须为正确的工作选择正确的工具，这是理解编程/算法的一部分。 - jpp

Answer 2

你的函数越复杂（即，numpy 无法将其移动到自己的内部），你会发现性能差别不大。例如：

name_series = pd.Series(np.random.choice(['adam', 'chang', 'eliza', 'odom'], replace=True, size=100000))

def parse_name(name):
    if name.lower().startswith('a'):
        return 'A'
    elif name.lower().startswith('e'):
        return 'E'
    elif name.lower().startswith('i'):
        return 'I'
    elif name.lower().startswith('o'):
        return 'O'
    elif name.lower().startswith('u'):
        return 'U'
    return name

parse_name_vec = np.vectorize(parse_name)

进行一些时间测试:

使用 apply 方法

%timeit name_series.apply(parse_name)

结果：

76.2 ms ± 626 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

使用 np.vectorize

%timeit parse_name_vec(name_series)

结果:

77.3 ms ± 216 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

Numpy尝试将Python函数转换为numpy ufunc对象，当您调用np.vectorize时。它是如何做到这一点的，我实际上并不知道 - 您必须深入挖掘numpy的内部，而我目前没有这个意愿。话虽如此，它似乎在简单的数字函数上比这里基于字符串的函数表现得更好。

将大小增加到1,000,000：

name_series = pd.Series(np.random.choice(['adam', 'chang', 'eliza', 'odom'], replace=True, size=1000000))

apply

%timeit name_series.apply(parse_name)

结果：

769 ms ± 5.88 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

np.vectorize

%timeit parse_name_vec(name_series)

结果：

794 ms ± 4.85 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

使用 np.select 的更好（向量化）的方法：

cases = [
    name_series.str.lower().str.startswith('a'), name_series.str.lower().str.startswith('e'),
    name_series.str.lower().str.startswith('i'), name_series.str.lower().str.startswith('o'),
    name_series.str.lower().str.startswith('u')
]
replacements = 'A E I O U'.split()

时间：

%timeit np.select(cases, replacements, default=name_series)

结果：

67.2 ms ± 683 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

- PMende

---

我相信你在这里的断言是错误的。目前我无法用代码来支持这个说法，希望其他人可以。 - roganjosh