提高具有多个条件的 Pandas 子集/过滤器的性能

Question

在我的 Python 脚本中，我有一个大约 520 万行和 26 列的 Pandas DataFrame。

我正在运行的当前分析针对此 DataFrame 测试 150 万个不同的子集/过滤器，使用多个条件（列的组合）来计算出现的总数。我已经尝试了多种方法，并比我的原始代码提高了一点速度，但我相信这可以做得更快。

下面是带有虚拟 DataFrame（520 万行和 26 列 A-Z）的代码和我最初使用的子集方法的一个示例，具有多个条件：

import pandas as pd
import numpy as np
import time

np.random.seed(101)
df = pd.DataFrame(np.random.randint(0,11,size=(5200000, 26)), columns=list('ABCDEFGHIJKLMNOPQRSTUVWXYZ'))

cond1 = (df.A > 5)
cond2 = (df.B.shift(1).between(8,10,inclusive=True))
cond3 = (df.C.shift(1) < 5)
cond4 = (df.D.shift(1) == 1)
cond5 = (df.E.shift(1).between(2,5,inclusive=True))

start = time.time()
subsetLength = len(df.loc[cond1 & cond2 & cond3 & cond4 & cond5])
stop = time.time()

print("Length:",subsetLength)
print("Time:", stop - start)

对于每个子集/过滤器，在我的笔记本电脑（Macbook Pro 2020 Intel）上检索此子集长度的时间约为 0.030 秒：

Length: 9781
Time: 0.029639244079589844

注意：我正在使用 .shift(1) 检查 B - E 列的前几行，而 A 列的情况并非如此。

现在 0.030 秒可能看起来相当快，但由于我正在运行 150 万个这些子集，因此脚本目前大约需要 12-13 小时。因此，我可以在单个子集/过滤器上获得的每一次速度提升都可以节省大量时间。

我也尝试过 .shape[0] 和 .sum() 而不是 .loc[]：

.shape[0]

subsetLength = df[cond1 & cond2 & cond3 & cond4 & cond5].shape[0]

这似乎快了一点，但仍然在 0.027s - 0.028s 的范围内。

.sum()

subsetLength = (cond1 & cond2 & cond3 & cond4 & cond5).sum()

这大约是 0.021 秒，与可以节省大约 4 小时的 .loc[] 相比，这已经是相当不错的改进了。但仍然需要很多小时（大约 8 或 9 个）。

np.where()

最后我尝试使用 np.where()，为此我不得不稍微改变一下条件：

cond1 = (df['A'] > 5)
cond2 = (df['B'].shift(1).between(8,10,inclusive=True))
cond3 = (df['C'].shift(1) < 5)
cond4 = (df['D'].shift(1) == 1)
cond5 = (df['E'].shift(1).between(2,5,inclusive=True))

subsetLength = len(np.where(cond1 & cond2 & cond3 & cond4 & cond5)[0])

这大约花费了 0.0195 秒，是迄今为止我得到的最快的方法。

最好我正在寻找 0.002 秒或更快范围内的性能，因此脚本需要大约 1 小时才能完成，但我不知道这是否可能。关于如何提高这个子集/过滤大 DataFrame 以进行计数的速度的任何建议？我是否可能使用缓慢的方法对具有多个条件的子集进行计数？或者这种分析只是繁重，需要时间。

Answer 1

一种解决方案是就地执行逻辑操作以避免创建许多临时数组：

tmp = np.empty(cond1.shape, dtype=np.bool8)
np.logical_and(cond1, cond2, out=tmp)
np.logical_and(tmp, cond3, out=tmp)
np.logical_and(tmp, cond4, out=tmp)
np.logical_and(tmp, cond5, out=tmp)
subsetLength = tmp.sum()

这个解决方案在我的机器上快 2.2 倍。

上述解决方案效率不高，因为它仍然在内存中读写大数组。但是，我怀疑仅使用 Numpy 是否可以更快地执行某些操作（因为 logical_and 显然不提供一次计算超过 2 个元素的方法）。

另一种解决方案是使用Numba来避免临时数组的读/写，甚至可以并行运行代码。方法如下：

import numba as nb

@nb.njit('int64(bool_[::1],bool_[::1],bool_[::1],bool_[::1],bool_[::1])', parallel=True)
def computeSubsetLength(cond1, cond2, cond3, cond4, cond5):
    n = len(cond1)
    assert len(cond2) == n and len(cond3) == n and len(cond4) == n and len(cond5) == n
    subsetLength = 0
    for i in nb.prange(n):
        subsetLength += cond1[i] & cond2[i] & cond3[i] & cond4[i] & cond5[i]
    return subsetLength

subsetLength = computeSubsetLength(cond1.to_numpy(), cond2.to_numpy(), cond3.to_numpy(), cond4.to_numpy(), cond5.to_numpy())

这个新版本快了 31 倍。它在我的机器上接近最优，因为读取内存中的数组大约需要时间（需要0.000624 s，而仅从内存中读取输入数据的时间是0.000484 s）。