Pandas 掩码 / where 方法与 NumPy np.where

Question

在有条件地更新系列中的值时，我经常使用 Pandas mask 和 where 方法来实现更清晰的逻辑。但是，对于相对性能关键的代码，我注意到相对于numpy.where 的性能显着下降。

虽然我很高兴在特定情况下接受这一点，但我很想知道：

1. Pandas mask / where 方法是否提供任何其他功能，除了 inplace / errors / try-cast 参数？我了解这 3 个参数，但很少使用它们。例如，我不知道level 参数指的是什么。
2. 是否有任何重要的反例，mask / where 优于 numpy.where？如果存在这样的例子，它可能会影响我今后如何选择合适的方法。

作为参考，这里是 Pandas 0.19.2 / Python 3.6.0 的一些基准测试：

np.random.seed(0)

n = 10000000
df = pd.DataFrame(np.random.random(n))

assert (df[0].mask(df[0] > 0.5, 1).values == np.where(df[0] > 0.5, 1, df[0])).all()

%timeit df[0].mask(df[0] > 0.5, 1)       # 145 ms per loop
%timeit np.where(df[0] > 0.5, 1, df[0])  # 113 ms per loop

对于非标量值，性能似乎进一步：

%timeit df[0].mask(df[0] > 0.5, df[0]*2)       # 338 ms per loop
%timeit np.where(df[0] > 0.5, df[0]*2, df[0])  # 153 ms per loop

Answer 1

我使用的是 pandas 0.23.3 和 Python 3.6，因此仅在您的第二个示例中，我可以看到运行时间的真正差异。

但是，让我们研究您的第二个示例的稍微不同的版本（因此我们将2*df[0] 排除在外）。这是我机器上的基线：

twice = df[0]*2
mask = df[0] > 0.5
%timeit np.where(mask, twice, df[0])  
# 61.4 ms ± 1.51 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

%timeit df[0].mask(mask, twice)
# 143 ms ± 5.27 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

Numpy 的版本比 pandas 快大约 2.3 倍。

因此，让我们对这两个函数进行剖析以了解区别 - 当您对代码基础不太熟悉时，剖析是一种了解全局的好方法：它比调试更快，并且比试图弄清楚更不容易出错仅通过阅读代码会发生什么。

我在 Linux 上使用 perf。对于我们得到的 numpy 版本（列表见附录 A）：

>>> perf record python np_where.py
>>> perf report

Overhead  Command  Shared Object                                Symbol                              
  68,50%  python   multiarray.cpython-36m-x86_64-linux-gnu.so   [.] PyArray_Where
   8,96%  python   [unknown]                                    [k] 0xffffffff8140290c
   1,57%  python   mtrand.cpython-36m-x86_64-linux-gnu.so       [.] rk_random

我们可以看到，大部分时间花在PyArray_Where 上——大约 69%。未知符号是一个内核函数（事实上 clear_page） - 我在没有 root 权限的情况下运行，因此无法解析符号。

对于 pandas，我们得到（代码见附录 B）：

>>> perf record python pd_mask.py
>>> perf report

Overhead  Command  Shared Object                                Symbol                                                                                               
  37,12%  python   interpreter.cpython-36m-x86_64-linux-gnu.so  [.] vm_engine_iter_task
  23,36%  python   libc-2.23.so                                 [.] __memmove_ssse3_back
  19,78%  python   [unknown]                                    [k] 0xffffffff8140290c
   3,32%  python   umath.cpython-36m-x86_64-linux-gnu.so        [.] DOUBLE_isnan
   1,48%  python   umath.cpython-36m-x86_64-linux-gnu.so        [.] BOOL_logical_not

完全不同的情况：

• pandas 不会在后台使用PyArray_Where - 最显着的耗时是vm_engine_iter_task，即numexpr-functionality。
• 正在进行一些繁重的内存复制 - __memmove_ssse3_back 大约使用了 25% 的时间！可能内核的一些功能也与内存访问有关。

实际上，pandas-0.19 在后台使用了 PyArray_Where，对于旧版本，性能报告看起来像：

Overhead  Command        Shared Object                     Symbol                                                                                                     
  32,42%  python         multiarray.so                     [.] PyArray_Where
  30,25%  python         libc-2.23.so                      [.] __memmove_ssse3_back
  21,31%  python         [kernel.kallsyms]                 [k] clear_page
   1,72%  python         [kernel.kallsyms]                 [k] __schedule

所以基本上它会在后台使用np.where + 一些开销（所有上述数据复制，请参阅__memmove_ssse3_back）。

在 pandas 的 0.19 版中，我没有看到 pandas 比 numpy 更快的情况——它只是增加了 numpy 功能的开销。 Pandas 的 0.23.3 版是完全不同的故事——这里使用了 numexpr-module，很有可能在某些情况下 pandas 的版本（至少稍微）更快。

我不确定这种内存复制是否真的需要/必要 - 也许有人甚至可以称之为性能错误，但我只是不知道足以确定。

我们可以通过剥离一些间接方式（通过np.array 而不是pd.Series）来帮助熊猫不要复制。例如：

%timeit df[0].mask(mask.values > 0.5, twice.values)
# 75.7 ms ± 1.5 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

现在，pandas 只慢了 25%。性能说：

Overhead  Command  Shared Object                                Symbol                                                                                                
  50,81%  python   interpreter.cpython-36m-x86_64-linux-gnu.so  [.] vm_engine_iter_task
  14,12%  python   [unknown]                                    [k] 0xffffffff8140290c
   9,93%  python   libc-2.23.so                                 [.] __memmove_ssse3_back
   4,61%  python   umath.cpython-36m-x86_64-linux-gnu.so        [.] DOUBLE_isnan
   2,01%  python   umath.cpython-36m-x86_64-linux-gnu.so        [.] BOOL_logical_not

数据复制要少得多，但仍然比主要负责开销的 numpy 版本多。

我的主要收获：

• pandas 有可能至少比 numpy 快一点（因为它可能更快）。但是，pandas 对数据复制的处理有些不透明，因此很难预测这种潜力何时会被（不必要的）数据复制所掩盖。

• 当where/mask 的性能成为瓶颈时，我会使用 numba/cython 来提高性能 - 请参阅下面我相当天真的尝试使用 numba 和 cython。

---

想法是采取

np.where(df[0] > 0.5, df[0]*2, df[0])

版本并消除创建临时的需要 - 即df[0]*2。

按照@max9111 的建议，使用 numba：

import numba as nb
@nb.njit
def nb_where(df):
    n = len(df)
    output = np.empty(n, dtype=np.float64)
    for i in range(n):
        if df[i]>0.5:
            output[i] = 2.0*df[i]
        else:
            output[i] = df[i]
    return output

assert(np.where(df[0] > 0.5, twice, df[0])==nb_where(df[0].values)).all()
%timeit np.where(df[0] > 0.5, df[0]*2, df[0])
# 85.1 ms ± 1.61 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

%timeit nb_where(df[0].values)
# 17.4 ms ± 673 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

这比 numpy 的版本快大约 5 倍！

这是我在 Cython 的帮助下提高性能的成功尝试：

%%cython -a
cimport numpy as np
import numpy as np
cimport cython

@cython.boundscheck(False)
@cython.wraparound(False)
def cy_where(double[::1] df):
    cdef int i
    cdef int n = len(df)
    cdef np.ndarray[np.float64_t] output = np.empty(n, dtype=np.float64)
    for i in range(n):
        if df[i]>0.5:
            output[i] = 2.0*df[i]
        else:
            output[i] = df[i]
    return output

assert (df[0].mask(df[0] > 0.5, 2*df[0]).values == cy_where(df[0].values)).all()

%timeit cy_where(df[0].values)
# 66.7± 753 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

提供 25% 的加速。不确定，为什么 cython 比 numba 慢得多。

---

列表：

答： np_where.py：

import pandas as pd
import numpy as np

np.random.seed(0)

n = 10000000
df = pd.DataFrame(np.random.random(n))

twice = df[0]*2
for _ in range(50):
      np.where(df[0] > 0.5, twice, df[0])  

B: pd_mask.py:

import pandas as pd
import numpy as np

np.random.seed(0)

n = 10000000
df = pd.DataFrame(np.random.random(n))

twice = df[0]*2
mask = df[0] > 0.5
for _ in range(50):
      df[0].mask(mask, twice)