优化 pandas 多列/多索引查询

Question

我有一个非常大的 table（目前有 5500 万行，可能更多），我需要 select 它的子集并对这些子集执行非常简单的操作，很多很多次。看起来 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')

这 selects（对于我的随机数据集）1848 行：

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

每次查询大约需要 .1-.15 秒：

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

所以我想我一定可以通过对 table 进行排序和索引来做得更好，对吗？而且我可以利用一切都是 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 之类的事情，但我想不出如何在多列上做到这一点。 pandas 是错误的选择吗？

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

优化 pandas 多列/多索引查询

optimize pandas query on multiple columns / multiindex

python

numpy

bigdata

pandas