pandas合并DataFrame时仅保留双方共有的重复值的实现方法

问题背景 #

现有两个pandas DataFrame 格式的同源数据，存在以下特征：

• 受旧软件读取错误影响，来源1（左表）的数据仅是来源2（右表）数据的子集
• 来源2的数据包含1个额外列
• 两个数据源都可能包含有效的重复值
• 两张表之间没有共享的唯一键

期望效果 #

• 合并两个表，保留数据交集
• 仅保留在两个DataFrame中都存在的重复值
• 第二张表的额外列需要保留在新的合并表中

现有方案问题 #

测试用例 #

df1 = pd.DataFrame( 
    { # 最后一行是第一行的重复值
        'var1' : [1, 2, 3, 1,], 
        'var2'  : ['a', 'b', 'c', 'a']
    }
)

df2 = pd.DataFrame(  
    { # 最后两行是前两行的重复值
    'cat': [1, 2, 3, 4, 1, 2],
    'var1' : [1, 2, 3, 4, 1, 2],
    'var2'  : ['a', 'b', 'c', 'd', 'a', 'b']
    }
)

直接合并问题 #

如果直接用公共列做内连接，会产生笛卡尔积式的多余重复：

merged = pd.merge(df1, df2, on=["var1", "var2"], how="inner")

得到的结果会多出不需要的重复行：

var1 var2 cat
0   1   a   1 
1   1   a   1 
2   1   a   1 
3   1   a   1 
4   2   b   2 
5   2   b   2 
6   3   c   3 

直接去重问题 #

如果直接对合并结果去重，会丢失原本有效的重复值：

merged.drop_duplicates()

得到的结果会少掉有效重复行：

var1 var2 cat
0   1   a   1
4   2   b   2
6   3   c   3

可行解决方案 #

核心思路：给每个公共键分组下的重复行加组内序号，将序号也作为合并键，即可实现仅保留两边都存在的重复次数的行，避免笛卡尔积多余重复。

实现代码 #

import pandas as pd

# 测试用例定义同上
df1 = pd.DataFrame( 
    { 
        'var1' : [1, 2, 3, 1,], 
        'var2'  : ['a', 'b', 'c', 'a']
    }
)

df2 = pd.DataFrame(  
    { 
    'cat': [1, 2, 3, 4, 1, 2],
    'var1' : [1, 2, 3, 4, 1, 2],
    'var2'  : ['a', 'b', 'c', 'd', 'a', 'b']
    }
)

# 给两个表的公共键分组下的行加组内序号
df1['group_idx'] = df1.groupby(['var1', 'var2']).cumcount()
df2['group_idx'] = df2.groupby(['var1', 'var2']).cumcount()

# 合并时将组内序号也作为关联条件，合并后删除临时序号列
merged = pd.merge(df1, df2, on=['var1', 'var2', 'group_idx'], how='inner').drop('group_idx', axis=1)
print(merged)

输出结果 #

var1 var2  cat
0     1    a    1
1     2    b    2
2     3    c    3
3     1    a    1

逻辑说明 #

• groupby([公共键]).cumcount() 会给每个分组下的重复行依次标记0、1、2...的序号，代表该行是该公共键下的第N次出现
• 合并时将该序号也作为关联条件，就能保证：某公共键在左表出现M次、右表出现N次时，合并后只会保留min(M,N)条数据，刚好匹配两边都存在的重复次数
• 不需要额外做去重操作，也不会丢失有效的重复值

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