Эффективно выбирайте строки, соответствующие одному из нескольких значений в Pandas DataFrame

Проблема

Данные в Pandas DataFrame, как показано ниже:

Name     Amount
---------------
Alice       100
Bob          50
Charlie     200
Alice        30
Charlie      10

Я хочу выбрать все строки, где Name является одним из нескольких значений в коллекции {Alice, Bob}

Name     Amount
---------------
Alice       100
Bob          50
Alice        30

Вопрос

Каков эффективный способ сделать это в Pandas?

Параметры, которые я вижу в них

Перебирать строки, обрабатывая логику с помощью Python

Выберите и слейте много операторов, например,

merge(df[df.name = specific_name] for specific_name in names) # something like this

Выполните какое-то соединение

Каковы компромиссы производительности здесь? Когда одно решение лучше других? Какие решения мне не хватает?

В то время как в приведенном выше примере используются строки, мое фактическое задание использует совпадения на 10-100 целых чисел по миллионам строк и поэтому очень важны операции NumPy.

Ответ 1

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

In [11]: df['Name'].isin(['Alice', 'Bob'])
Out[11]: 
0     True
1     True
2    False
3     True
4    False
Name: Name, dtype: bool

In [12]: df[df.Name.isin(['Alice', 'Bob'])]
Out[12]: 
    Name  Amount
0  Alice     100
1    Bob      50
3  Alice      30

Ответ 2

Так как в вашем фактическом прецеденте значения в df['Name'] равны ints, вы можете быстрее генерировать булевскую маску, используя индексирование NumPy вместо Series.isin.

idx = np.zeros(N, dtype='bool')
idx[names] = True
df[idx[df['Name'].values]]

Например, с учетом этой настройки:

import pandas as pd
import numpy as np

N = 100000
df = pd.DataFrame(np.random.randint(N, size=(10**6, 2)), columns=['Name', 'Amount'])
names = np.random.choice(np.arange(N), size=100, replace=False)

In [81]: %timeit idx = np.zeros(N, dtype='bool'); idx[names] = True; df[idx[df['Name'].values]]
100 loops, best of 3: 9.88 ms per loop

In [82]: %timeit df[df.Name.isin(names)]
10 loops, best of 3: 107 ms per loop

In [83]: 107/9.88
Out[83]: 10.82995951417004

N (по существу) максимальное значение, которое может достичь df['Names']. Если N меньше, то преимущество в скорости не так велико. С N = 200,

In [93]: %timeit idx = np.zeros(N, dtype='bool'); idx[names] = True; df[idx[df['Name'].values]]
10 loops, best of 3: 62.6 ms per loop

In [94]: %timeit df[df.Name.isin(names)]
10 loops, best of 3: 178 ms per loop

In [95]: 178/62.6
Out[95]: 2.8434504792332267

Предостережение: Как показано выше, кажется, есть преимущество в скорости, особенно когда N становится большим. Однако, если N слишком велико, то формирование idx = np.zeros(N, dtype='bool') может оказаться невозможным.

Проверка работоспособности:

expected = df[df.Name.isin(names)]
idx = np.zeros(N, dtype='bool')
idx[names] = True
result = df[idx[df['Name'].values]]
assert expected.equals(result)