tối ưu hóa truy vấn gấu trúc trên nhiều cột/multiindex

Question

Tôi có một bảng rất lớn (hiện tại là 55 triệu hàng, có thể nhiều hơn) và tôi cần chọn tập con của nó và thực hiện các thao tác rất đơn giản trên các tập con đó, rất nhiều lần . Dường như gấu trúc có thể là cách tốt nhất để làm điều này trong python, nhưng tôi đang chạy vào các vấn đề tối ưu hóa.

Tôi đã cố tạo một tập dữ liệu giả phù hợp chặt chẽ với tập dữ liệu thực của tôi (mặc dù nó nhỏ hơn ~ 5-10 lần). Điều này vẫn còn lớn, mất rất nhiều bộ nhớ, vv Có bốn cột mà tôi đang truy vấn và hai cột mà tôi đang sử dụng để tính toán.

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') 

Lựa chọn này (đối với dữ liệu ngẫu nhiên của tôi) 1848 hàng:

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

Và nó mất khoảng 0,1-0,15 giây cho mỗi truy vấn:

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

Vì vậy, tôi nghĩ rằng tôi phải có thể làm tốt hơn bằng cách sắp xếp và lập chỉ mục bảng, đúng không? Và tôi có thể tận dụng thực tế là tất cả mọi thứ là một int, vì vậy tôi có thể làm lát ..

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 

và phải mất rất nhiều còn:

timeit.timeit(f2,number=10)/10 
Out[295]: 7.335134506225586 

Am tôi thiếu cái gì ở đây? Nó có vẻ như tôi có thể làm một cái gì đó như numpy.searchsorted, nhưng tôi không thể nghĩ làm thế nào để làm điều đó trên nhiều cột. Gấu trúc là sự lựa chọn sai?

Answer 1

Vì vậy, có 2 vấn đề ở đây.

Đây là một kỹ xảo mà làm cho cú pháp một chút đẹp hơn

In [111]: idx = pd.IndexSlice 

1) của bạn .query không có ưu tiên chính xác. Toán tử & có mức độ ưu tiên cao hơn so với toán tử so sánh như <= và cần dấu ngoặc đơn xung quanh toán hạng trái và phải của nó.

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() 

Đây là truy vấn ban đầu của bạn sử dụng máy thái 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],:] 

Đây là một phiên bản xích của truy vấn này. IOW của nó một lựa chọn lặp đi lặp lại trên tập kết quả.

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

Luôn luôn khẳng định tính đúng đắn trước khi làm việc trên hiệu suất

In [109]: (result1==result2).all().all() 
Out[109]: True 

In [110]: (result1==result3).all().all() 
Out[110]: True 

Performance

Các .query IMHO sẽ thực sự quy mô rất tốt và sử dụng đa lõi. Đối với một bộ lựa chọn lớn, đây sẽ là cách để đi

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) Cắt nhiều chỉ mục gốc. Có một vấn đề ở đây, xem bên dưới.Tôi không chắc chắn chính xác tại sao điều này là phi performant, và sẽ điều tra này 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 

lựa chọn lặp đi lặp lại làm cho điều này khá performant. Lưu ý rằng tôi sẽ không bình thường khuyên bạn nên làm điều này vì bạn không thể gán cho nó, nhưng cho mục đích này nó là ok.

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