在多列/多索引上优化pandas查询

Question

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

我试图创建一个与真实数据集非常匹配的假数据集（尽管它比实际数据集小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行：

^{pr2}$

每个查询大约需要0.1到0.15秒：

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

所以我想我必须通过对表进行排序和索引来做得更好，对吗？我可以利用一切都是整数的事实，所以我可以做切片。。在

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没有正确的优先级。&运算符比比较运算符（如<=）具有更高的优先级，并且需要在其左右操作数周围加上括号。在

^{pr2}$

这是使用多索引切片器的原始查询

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],:]

下面是这个查询的链式版本。这是对结果集的重复选择。在

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

性能

实际上，.queryIMHO将非常好地扩展并使用多核。对于一个大的选择集，这将是一个好办法

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