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

时间:2015-06-05 01:53:59

标签: python numpy pandas bigdata

我有一个非常大的表(目前有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

每个查询需要大约.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这样的事情,但我想不出怎么在多列上做到这一点。熊猫是错误的选择吗?

1 个答案:

答案 0 :(得分:2)

所以这里有两个问题。

这是一种使语法更好的技巧

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