Pandas如何基于当前行与下一行id条件实现两个DataFrame关联查询

实现方法 #

你可以通过新增下一行id辅助列+左连接的方式快速实现需求，代码如下：

import pandas as pd

# 构造示例数据
df1 = pd.DataFrame([[0, 17], [1, 5], [2, 11], [3, 15], [4, 10]], columns = ['seq', 'id'])
df2 = pd.DataFrame([[17, 7, 'a'], [17, 5, 'b'], [17, 8, 'c'], [5, 4, 'd'], [5, 11, 'e'], [11, 9, 'f'], [11, 15, 'g'], [15, 21, 'h'], [15, 10, 'i']], columns = ['id_1', 'id_2', 'x1'])

# 步骤1：生成下一行id的辅助列
df1['next_id'] = df1['id'].shift(-1)

# 步骤2：按照匹配规则左连接df2获取x1
df_result = df1.merge(
    df2,
    left_on=['id', 'next_id'],
    right_on=['id_1', 'id_2'],
    how='left'
)[['seq', 'id', 'x1']]

print(df_result)

运行后输出的df_result就是你需要的结果。

大数据量优化方案 #

如果你的数据量较大，可以提前把df2的匹配关系转为字典映射，避免merge的性能开销：

# 生成id组合到x1的映射字典
id2x1 = df2.set_index(['id_1', 'id_2'])['x1'].to_dict()
next_id_list = df1['id'].shift(-1)

# 直接匹配赋值
df1['x1'] = df1.apply(lambda row: id2x1.get((row['id'], next_id_list[row.name])), axis=1)
df_result = df1[['seq', 'id', 'x1']]

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