如何加速df.query？为何其性能远低于原生numexpr

测试验证过程 #

1. 构建测试用大型DataFrame #

from random import shuffle
import pandas as pd
l1 = list(range(1000000))
l2 = list(range(1000000))
shuffle(l1)
shuffle(l2)
df = pd.DataFrame([l1, l2]).T
df = df.sample(frac=1)
df = df.rename(columns={0: 'A', 1:'B'})

2. Pandas原生列对比查询性能 #

执行复合条件的列对比查询：

%timeit (df.A == df.B) | (df.A / df.B < 1) |  (df.A * df.B > 3000)

结果：10.8 ms ± 309 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
耗时约10.8毫秒。

3. 原生numexpr执行相同查询的性能 #

直接调用numexpr执行相同逻辑：

import numexpr
a = df.A.__array__()
b = df.B.__array__()
%timeit numexpr.evaluate('((a == b) | (a / b < 1) | (a * b > 3000))')

结果：1.95 ms ± 25.5 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
numexpr的速度是Pandas原生查询的6倍。

4. df.loc与df.query的性能对比 #

对比三种筛选方式的性能：

%timeit df.loc[numexpr.evaluate('((a == b) | (a / b < 1) | (a * b > 3000))')]
# 结果：20.5 ms ± 155 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

%timeit df.loc[(df.A == df.B) | (df.A / df.B < 1) |  (df.A * df.B > 3000)]
# 结果：27 ms ± 296 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

%timeit df.query('((A == B) | (A / B < 1) | (A * B > 3000))')
# 结果：32.5 ms ± 80.4 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

可以看到numexpr配合df.loc的性能仍显著优于纯Pandas查询，但df.query的速度却是三者中最慢的。

核心疑问与优化需求 #

df.query底层明明基于numexpr实现，为何性能反而远不如直接调用numexpr？有没有无需大量调整代码，就能在Pandas中高效利用numexpr的优化方法？

内容的提问来源于stack exchange，提问作者Hans