Pandas multiindex subset selection

Question

У меня есть Python 3.6 и Pandas 19.0. Я играю с несколькими индексами для фреймов данных.

iterables = [['bar', 'baz', 'foo', 'qux'], ['one', 'two', 'three','four']]
pd.MultiIndex.from_product(iterables, names=['first', 'second'])
arrays = [np.array(['bar', 'bar', 'baz', 'baz', 'foo', 'foo', 'qux', 'qux', 'bar', 'foo']),
          np.array(['one', 'two', 'one', 'two', 'one', 'two', 'one', 'two', 'three', 'four'])]
s = pd.DataFrame(np.random.randn(10), index=arrays)

Я знаю, как получить подмножество на основе одного значения для одного из индексов, например,

s.loc[s.index.get_level_values(0)=='bar']
Out[16]: 
                  0
bar one    1.409395
    two    0.837486
    three  1.290018

Как я могу получить подмножество на основе набора значений, соответствующих одному индексу? Очевидно, что синтаксис ниже не работает:

my_subset = set(['three', 'one'])
s.loc[s.index.get_level_values(1) in my_subset]

EDIT:

Что было бы самым быстрым решением для большого фрейма данных?

Answer 1

Вы можете query ваш фрейм данных. Предполагая, как в вашем примере, второй уровень вашего индекса называется second:

my_subset = ['three', 'one']

res = s.query('second in @my_subset')

Answer 2

Используйте Index.isin и выберите второй уровень с помощью 1:

my_subset = set(['three', 'one'])
a = s.loc[s.index.get_level_values(1).isin(my_subset)]
print (a)

                  0
bar one   -0.372206
baz one    0.886271
foo one   -2.231380
qux one    0.960636
bar three  1.272873

Perrfomance : зависит от количества совпадающих значений и количества строк:

N = 10000
a = ['bar', 'baz', 'foo', 'qux']
b = ['one', 'two', 'three','four']
arrays = pd.MultiIndex.from_arrays([np.random.choice(a, size=N),
                                     np.random.choice(b, size=N)], names=['first', 'second'])

s = pd.DataFrame(np.random.randn(N), index=arrays).sort_index()
print (s)

my_subset1 = set(['three', 'one'])
my_subset2 = ['three', 'one']

In [209]: %timeit s.loc[s.index.get_level_values(1).isin(my_subset1)]
866 µs ± 59.8 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

In [210]: %timeit s.query('second in @my_subset2')
2.19 ms ± 47.6 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)