优化多列/多索引的熊猫查询

Question

我有一个非常大的表（目前有 5500 万行，可能更多），我需要选择它的子集并对这些子集执行非常简单的操作，很多次。看起来 pandas 可能是在 python 中执行此操作的最佳方法，但我遇到了优化问题。

我试图创建一个与我的真实数据集非常匹配的假数据集（尽管它要小约 5-10 倍）。这仍然很大，占用大量内存等。我要查询四列，还有两列用于计算。

import pandas
import numpy as np
import timeit

n=10000000
mdt = pandas.DataFrame()
mdt['A'] = np.random.choice(range(10000,45000,1000), n)
mdt['B'] = np.random.choice(range(10,400), n)
mdt['C'] = np.random.choice(range(1,150), n)
mdt['D'] = np.random.choice(range(10000,45000), n)
mdt['x'] = np.random.choice(range(400), n)
mdt['y'] = np.random.choice(range(25), n)


test_A = 25000
test_B = 25
test_C = 40
test_D = 35000

eps_A = 5000
eps_B = 5
eps_C = 5
eps_D = 5000


f1 = lambda : mdt.query('@test_A-@eps_A <= A <= @test_A+@eps_A  &  ' +
                        '@test_B-@eps_B <= B <= @test_B+@eps_B  &  ' +
                        '@test_C-@eps_C <= C <= @test_C+@eps_C  &  ' +
                        '@test_D-@eps_D <= D <= @test_D+@eps_D')

这会选择（对于我的随机数据集）1848 行：

len(f1())
Out[289]: 1848

每个查询大约需要 0.1-.15 秒：

timeit.timeit(f1,number=10)/10
Out[290]: 0.10734589099884033

所以我认为我必须能够通过对表格进行排序和索引来做得更好，对吧？而且我可以利用一切都是 int 的事实，所以我可以做切片..

mdt2 = mdt.set_index(['A', 'B', 'C', 'D']).sortlevel()

f2 = lambda : mdt2.loc[(slice(test_A-eps_A, test_A+eps_A),
                        slice(test_B-eps_B, test_B+eps_B),
                        slice(test_C-eps_C, test_C+eps_C),
                        slice(test_D-eps_D, test_D+eps_D)), :]

len(f2())
Out[299]: 1848

而且它需要很多更长的时间：

timeit.timeit(f2,number=10)/10
Out[295]: 7.335134506225586

我在这里遗漏了什么吗？似乎我可以做类似 numpy.searchsorted 的事情，但我想不出如何在多个列上做到这一点。熊猫是错误的选择吗？

Answer 1

所以这里有两个问题。

这是一个让语法更好看的技巧

In [111]: idx = pd.IndexSlice

1) 您的.query 没有正确的优先级。 & 运算符的优先级高于 <= 等比较运算符，并且需要在其左右操作数两边加上括号。

In [102]: result3 = mdt.query("(@test_A-@eps_A <= A <= @test_A+@eps_A) & (@test_B-@eps_B <= B <= @test_B+@eps_B) & (@test_C-@eps_C <= C <= @test_C+@eps_C) & (@test_D-@eps_D <= D <= @test_D+@eps_D)").set_index(['A','B','C','D']).sortlevel()

这是您使用 MultiIndex 切片器的原始查询

In [103]: result1 = mdt2.loc[idx[test_A-eps_A:test_A+eps_A,test_B-eps_B:test_B+eps_B,test_C-eps_C:test_C+eps_C,test_D-eps_D:test_D+eps_D],:]

这是此查询的链接版本。 IOW 它是对结果集的重复选择。

In [104]: result2 = mdt2.loc[idx[test_A-eps_A:test_A+eps_A],:].loc[idx[:,test_B-eps_B:test_B+eps_B],:].loc[idx[:,:,test_C-eps_C:test_C+eps_C],:].loc[idx[:,:,:,test_D-eps_D:test_D+eps_D],:]

在处理性能之前始终确认正确性

In [109]: (result1==result2).all().all()
Out[109]: True

In [110]: (result1==result3).all().all()
Out[110]: True

性能

.query 恕我直言实际上可以很好地扩展并使用多核。对于大型选择集，这将是要走的路

In [107]: %timeit mdt.query("(@test_A-@eps_A <= A <= @test_A+@eps_A) & (@test_B-@eps_B <= B <= @test_B+@eps_B) & (@test_C-@eps_C <= C <= @test_C+@eps_C) & (@test_D-@eps_D <= D <= @test_D+@eps_D)").set_index(['A','B','C','D']).sortlevel()
10 loops, best of 3: 107 ms per loop

2) 原来的多索引切片。这里有一个问题，见下文。我不确定这究竟是什么原因，并将对此进行调查here

In [106]: %timeit  mdt2.loc[idx[test_A-eps_A:test_A+eps_A,test_B-eps_B:test_B+eps_B,test_C-eps_C:test_C+eps_C,test_D-eps_D:test_D+eps_D],:]
1 loops, best of 3: 4.34 s per loop

重复的选择使这非常有效。请注意，我通常不建议您这样做，因为您无法分配给它，但出于此目的，它是可以的。

In [105]: %timeit mdt2.loc[idx[test_A-eps_A:test_A+eps_A],:].loc[idx[:,test_B-eps_B:test_B+eps_B],:].loc[idx[:,:,test_C-eps_C:test_C+eps_C],:].loc[idx[:,:,:,test_D-eps_D:test_D+eps_D],:]
10 loops, best of 3: 140 ms per loop