优化多列/多索引的pandas查询

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到1.5秒:

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）独创的多索引切片。这里有一个问题，见下文。我不确定为什么这是无效的，并将在这里调查这一点

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