How to query MultiIndex index columns values in pandas

Solution 1:

To query the df by the MultiIndex values, for example where (A > 1.7) and (B < 666):

In [536]: result_df = df.loc[(df.index.get_level_values('A') > 1.7) & (df.index.get_level_values('B') < 666)]

In [537]: result_df
Out[537]: 
          C
A   B      
3.3 222  43
    333  59
5.5 333  56

Hence, to get for example the 'A' index values, if still required:

In [538]: result_df.index.get_level_values('A')
Out[538]: Index([3.3, 3.3, 5.5], dtype=object)

The problem is, that in large data frames the performance of by index selection worse by 10% than the sorted regular rows selection. And in repetitive work, looping, the delay accumulated. See example:

In [558]: df = store.select(STORE_EXTENT_BURSTS_DF_KEY)

In [559]: len(df)
Out[559]: 12857

In [560]: df.sort(inplace=True)

In [561]: df_without_index = df.reset_index()

In [562]: %timeit df.loc[(df.index.get_level_values('END_TIME') > 358200) & (df.index.get_level_values('START_TIME') < 361680)]
1000 loops, best of 3: 562 µs per loop

In [563]: %timeit df_without_index[(df_without_index.END_TIME > 358200) & (df_without_index.START_TIME < 361680)]
1000 loops, best of 3: 507 µs per loop

Solution 2:

For better readability, we can simply use the query() Method, to avoid the lengthy df.index.get_level_values() and reset_index/set_index to and fro.

Here is the target DataFrame:

In [12]: df                                                                    
Out[12]:                                                                       
          C                                                                    
A   B                                                                          
1.1 111  68                                                                    
    222  40                                                                    
3.3 222  20                                                                    
    333  11                                                                    
5.5 333  80                                                                    
6.6 777  51 

Answer for Q1 (A in range [3.3, 6.6]):

In [13]: df.query('3.3 <= A <= 6.6') # for closed interval                       
Out[13]:                                                                       
          C                                                                    
A   B                                                                          
3.3 222  20                                                                    
    333  11                                                                    
5.5 333  80                                                                    
6.6 777  51                                                                    

In [14]: df.query('3.3 < A < 6.6') # for open interval                         
Out[14]:                                                                       
          C                                                                    
A   B                                                                          
5.5 333  80

and of course one can play around with <, <=, >, >= for any kind of inclusion.


Similarly, answer for Q2 (A in range [2.0, 4.0]):

In [15]: df.query('2.0 <= A <= 4.0')                                        
Out[15]:                                                                    
          C                                                                 
A   B                                                                       
3.3 222  20                                                                 
    333  11 

Answer for Q3 (B in range [111, 500]):

In [16]: df.query('111 <= B <= 500')                                        
Out[16]:                                                                    
          C                                                                 
A   B                                                                       
1.1 111  68                                                                 
    222  40                                                                 
3.3 222  20                                                                 
    333  11                                                                 
5.5 333  80

And moreover, you can COMBINE the query for col A and B very naturally!

In [17]: df.query('0 < A < 4 and 150 < B < 400')                            
Out[17]:                                                                    
          C                                                                 
A   B                                                                       
1.1 222  40                                                                 
3.3 222  20                                                                 
    333  11

Solution 3:

With a 'float' like index you always want to use it as a column rather than a direct indexing action. These will all work whether the endpoints exist or not.

In [11]: df
Out[11]: 
          C
A   B      
1.1 111  81
    222  45
3.3 222  98
    333  13
5.5 333  89
6.6 777  98

In [12]: x = df.reset_index()

Q1

In [13]: x.loc[(x.A>=3.3)&(x.A<=6.6)]
Out[13]: 
     A    B   C
2  3.3  222  98
3  3.3  333  13
4  5.5  333  89
5  6.6  777  98

Q2

In [14]: x.loc[(x.A>=2.0)&(x.A<=4.0)]
Out[14]: 
     A    B   C
2  3.3  222  98
3  3.3  333  13

Q3

In [15]: x.loc[(x.B>=111.0)&(x.B<=500.0)]
Out[15]: 
     A    B   C
0  1.1  111  81
1  1.1  222  45
2  3.3  222  98
3  3.3  333  13
4  5.5  333  89

If you want the indices back, just set them. This is a cheap operation.

In [16]: x.loc[(x.B>=111.0)&(x.B<=500.0)].set_index(['A','B'])
Out[16]: 
          C
A   B      
1.1 111  81
    222  45
3.3 222  98
    333  13
5.5 333  89

If you REALLY want the actual index values

In [5]: x.loc[(x.B>=111.0)&(x.B<=500.0)].set_index(['A','B']).index
Out[5]: 
MultiIndex
[(1.1, 111), (1.1, 222), (3.3, 222), (3.3, 333), (5.5, 333)]