在多列/多索引上優化大熊貓查詢

Question

我有一個非常大的表（目前有5500萬行，可能更多），我需要選擇它的子集並對這些子集，很多次執行非常簡單的操作。看起來像大熊貓可能是在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-0.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

所以這裏有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() 

這是使用多指標切片

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