如何提取与另一DataFrame基因ID前缀匹配的DataFrame行？

问题背景 #

现有两个Pandas DataFrame：

import pandas as pd

# df数据
df = pd.DataFrame({
    'gene_ids': ['LOC_Os03g58080', 'LOC_Os06g17030', 'LOC_Os11g01380'],
    'baseMean': [15.8597614967997, 15.8597614967997, 779.386943289599],
    'log2FoldChange': [-20.7189803704191, -20.7189803704191, -12.3657001542033]
})

# df2数据
df2 = pd.DataFrame({
    'gene_id': ['LOC_Os01g01010.1', 'LOC_Os06g17030.1', 'LOC_Os06g17030.2', 'LOC_Os01g01010.1'],
    'GO': ['GO:0030234', 'GO:0007165', 'GO:0006139', 'GO:0009056']
})

需求：当df['gene_ids']的字符串与df2['gene_id']中.之前的前缀匹配时，提取df2中的对应行，预期输出：

gene_id          GO
1  LOC_Os06g17030.1  GO:0007165
2  LOC_Os06g17030.2  GO:0006139

解决方案 #

分步实现 #

1. 提取df2['gene_id']的前缀部分（.之前的内容）
2. 生成df中基因ID的集合，提升匹配查询效率
3. 通过布尔索引筛选出匹配的行

完整代码 #

# 提取df2基因ID的前缀
df2['gene_prefix'] = df2['gene_id'].str.split('.').str[0]

# 获取df中的目标基因ID集合
target_genes = set(df['gene_ids'])

# 筛选匹配行并移除临时前缀列
result_df = df2[df2['gene_prefix'].isin(target_genes)].drop('gene_prefix', axis=1)

# 打印结果
print(result_df)

简化写法（无需临时列） #

result_df = df2[df2['gene_id'].str.split('.').str[0].isin(set(df['gene_ids']))]

运行后即可得到符合预期的结果DataFrame。

内容的提问来源于stack exchange，提问作者sumitra sivaprakasam