Pandas - какой самый эффективный способ фильтрации DataFrame с использованием pandas .Series.all

Question

Рассмотрим приведенный ниже код -

import pandas as pd
data = []
val = 0
for ind_1 in range(1000):
    for ind_2 in range(1000):
        data.append({'ind_1': ind_1, 'ind_2': ind_2,
                     'val': val})
        val += 1
df_mi = pd.DataFrame(data).set_index(['ind_1', 'ind_2'])

, который создает фрейм данных df_mi с MultiIndex-

In [90]: df_mi                                                                                       
Out[90]: 
                val
ind_1 ind_2        
0     0           0
      1           1
      2           2
      3           3
      4           4
...             ...
999   995    999995
      996    999996
      997    999997
      998    999998
      999    999999

[1000000 rows x 1 columns]

Теперь я хочу отфильтровать строки, применив некоторые условия ко всем значениям для каждого ind_1 -

In [116]: bool_filter_ind_1 = (df_mi['val'] < 999997).all(level='ind_1')                             

In [117]: bool_filter_ind_1                                                                          
Out[117]: 
ind_1
0       True
1       True
2       True
3       True
4       True
       ...  
995     True
996     True
997     True
998     True
999    False
Name: val, Length: 1000, dtype: bool

In [118]: ind_1_filtered = bool_filter_ind_1.index[bool_filter_ind_1]                                

In [119]: ind_1_filtered                                                                             
Out[119]: 
Int64Index([  0,   1,   2,   3,   4,   5,   6,   7,   8,   9,
            ...
            989, 990, 991, 992, 993, 994, 995, 996, 997, 998],
           dtype='int64', name='ind_1', length=999)

Результат верный, но df_mi.loc[ind_1_filtered] относительно медленный -

In [120]: timeit df_mi_filtered = df_mi.loc[ind_1_filtered]                                          
4.73 s ± 10.1 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

In [121]: df_mi_filtered                                                                             
Out[121]: 
                val
ind_1 ind_2        
0     0           0
      1           1
      2           2
      3           3
      4           4
...             ...
998   995    998995
      996    998996
      997    998997
      998    998998
      999    998999

[999000 rows x 1 columns]

Есть ли более быстрый способ выполнить такую ​​же фильтрацию?

Answer 1

Вы можете использовать:

Первая идея - invert Маска для df_mi['val'] >= 999997) и получить все ind_1 индексы для меньшего порога и фильтровать исходные индексы первого уровня по Index.isin для маски и фильтрации по boolean indexing:

def new(df_mi):
    lvl0 = df_mi.index.get_level_values(0)
    return df_mi[~lvl0.isin(lvl0[(df_mi['val'] >= 999997)].unique())]

In [240]: %timeit (new(df_mi))
51.5 ms ± 555 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

Другая идея заключается в использовании GroupBy.transform и GroupBy.all для маска и снова фильтрация по boolean indexing:

In [241]: %timeit df_mi[(df_mi['val'] < 999997).groupby(level='ind_1').transform('all')]
97.3 ms ± 1.04 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

Исходное решение:

def orig(df_mi):
    bool_filter_ind_1 = (df_mi['val'] < 999997).all(level='ind_1')  
    ind_1_filtered = bool_filter_ind_1.index[bool_filter_ind_1]
    return df_mi.loc[ind_1_filtered]

In [242]: %timeit orig(df_mi)
11.2 s ± 405 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)