有效地选择与 Pandas DataFrame 中多个值之一匹配的行

Question

问题

给定 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 中执行此操作的有效方法是什么？

我看到的选项

1. 遍历行，使用 Python 处理逻辑

2. 选择并合并许多语句，如下所示

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

3. 执行某种连接

这里的性能权衡是什么？什么时候一种解决方案比其他解决方案更好？我缺少什么解决方案？

虽然上面的示例使用字符串，但我的实际工作使用匹配数百万行的 10-100 个整数，因此快速 NumPy 操作可能是相关的。

Answer 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

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