用一个DataFrame查询另一个DataFrame

Question

我希望能够通过第二个DataFrame中多列和多行的值来过滤主要的DataFrame。我们可以将第二个DataFrame称为搜索条件。

原始数据如下：

   a  b  c
0  A  0  b
1  A  1  b
2  A  2  b
3  B  3  a
4  B  4  x

我想筛选/搜索所有满足条件 a=='B' 和 c=='a' | c=='x' 的行。这可以表示为一个 DataFrame。

search_for = df.loc[:, ['a', 'c']].iloc[3:5]

   a  c
3  B  a
4  B  x

关键点是我不知道具体的搜索值，因为它们是另一种“查询”的结果并存储在DataFrame中。

可以这样做，并且结果是正确的。

search_for = df.loc[:, ['a', 'c']].iloc[3:5]
query_mask = (df.a.isin(search_for.a) & df.c.isin(search_for.c))
result_a = df[query_mask]

但是实际数据中我有更多的列，因此我更偏向于以下伪代码：

df[search_for]

这里有一种方法，但它并不能完全奏效。

search_for = df.loc[:, ['a', 'c']].iloc[3:5]
query_mask = df.loc[:, search_for.columns].isin(search_for)
result_b = df[query_mask]

结果为

     a   b    c
0  NaN NaN  NaN
1  NaN NaN  NaN
2  NaN NaN  NaN
3    B NaN    a
4    B NaN    x

这是完整的 MWE

#!/usr/bin/env python3
import pandas as pd

print(pd.__version__)

# initial data
df = pd.DataFrame({
    'a': list('AAABB'),
    'b': range(5),
    'c': list('bbbax')
})
print(f'\nPrimary data:\n{df}')

# the values I need to search for
search_for = df.loc[:, ['a', 'c']].iloc[3:5]
print(f'\nSearch for:\n{search_for}')

# IMHO bad solution
result_a = df[df.a.isin(search_for.a) & df.c.isin(search_for.c)]
print(f'Result-A\n{result_a}')  # correct result

# An approach not fully working
search_for = df.loc[:, ['a', 'c']].iloc[3:5]
query_mask = df.loc[:, search_for.columns].isin(search_for)
result_b = df[query_mask]
print(f'Result-B\n{result_b}')

完整的 MWE 输出

1.2.5

Primary data:
   a  b  c
0  A  0  b
1  A  1  b
2  A  2  b
3  B  3  a
4  B  4  x

Search for:
   a  c
3  B  a
4  B  x
Result-A
   a  b  c
3  B  3  a
4  B  4  x
Result-B
     a   b    c
0  NaN NaN  NaN
1  NaN NaN  NaN
2  NaN NaN  NaN
3    B NaN    a
4    B NaN    x

- buhtz

---

像 df.merge(search_for)[search_for.columns] 这样合并的怎么样？ - Ben.T

这是正确的答案。你应该只是合并。 - alparslan mimaroğlu

这就是解决方案。但我不明白为什么你不把它发布为答案呢？我可以替你完成，但那样的话我会获得积分和声望值。;) - buhtz

Answer 1

使用 DataFrame.merge() 是 pandas 的方法。

df.merge(search_for)[search_for.columns]

- buhtz