如何高效比较两个DataFrame并按条件生成etl_flag列

现有两个Pandas DataFrame：

data1 = {
    'alias_cd': ['12345', '12345', '12345'],
    'country_cd': ['AU', 'AU', 'AU2'],
    'pos_name': ['st1', 'Jh', 'Jh'],
    'ts_allocated': [100, 100, 100],
    'tr_id': ['None', 'None', 'None'],
    'ty_name': ['E2E', 'E2E', 'E2E']
}

data2 = {
    'alias_cd': ['12345', '12345'],
    'country_cd': ['AU', 'AU3'],
    'pos_name': ['st1', 'st2'],
    'ts_allocated': [200, 100],
    'tr_id': ['None', 'None'],
    'ty_name': ['E2E', 'E2E']
}

df1 = pd.DataFrame(data1)
df2 = pd.DataFrame(data2)

预期输出：

alias_cd country_cd pos_name  ts_allocated tr_id ty_name etl_flag
0     12345         AU      st1           200  None    E2E        U
1     12345        AU3      st2           100  None    E2E        D
2     12345         AU      st1           100  None    E2E        I
3     12345         AU       Jh           100  None    E2E        I
4     12345        AU2       Jh           100  None    E2E        I

规则说明 #

• alias_cd与country_cd的组合作为主键。
• 若主键组合同时存在于df2和df1中（如12345 AU），则df2中对应行标记为'U'（Update），df1中该组合的所有行标记为'I'（Insert）并添加到结果中。
• 若主键组合仅存在于df2中（如12345 AU3），则标记为'D'（DELETE）。
• 若主键组合仅存在于df1中，则标记为'I'（Insert）。

尝试了以下代码但无法得到正确输出，求高效实现方案：

df2['etl_flag'] = 'U'

to_insert = df1[~df1.apply(lambda x: (df2['alias_cd'] == x['alias_cd']) & (df2['country_cd'] == x['country_cd']), axis=1)]
to_insert['etl_flag'] = 'I'
df2 = pd.concat([df2, to_insert], ignore_index=True)

to_delete = df2[~df2.apply(lambda x: (df1['alias_cd'] == x['alias_cd']) & (df1['country_cd'] == x['country_cd']), axis=1)]
to_delete['etl_flag'] = 'D'

final_df = pd.concat([df2, to_delete], ignore_index=True)
final_df.sort_values(by=['alias_cd', 'country_cd'], inplace=True)

print(final_df[['alias_cd', 'country_cd', 'pos_name', 'ts_allocated', 'tr_id', 'ty_name', 'etl_flag']])

核心思路是先明确各主键组合在两个DataFrame中的存在情况，分类型处理行数据，最后合并结果，避免低效的逐行判断。

步骤1：提取主键集合并分类 #

先生成主键列，再通过集合快速区分主键的归属类型：

# 生成主键列（alias_cd + country_cd），方便后续判断
df1['pk'] = df1['alias_cd'] + '_' + df1['country_cd']
df2['pk'] = df2['alias_cd'] + '_' + df2['country_cd']

# 获取两个DataFrame的主键集合
pk_df1 = set(df1['pk'])
pk_df2 = set(df2['pk'])

# 拆分三类主键：共有的、仅df2存在的、仅df1存在的
common_pk = pk_df1 & pk_df2
only_df2_pk = pk_df2 - pk_df1
only_df1_pk = pk_df1 - pk_df2

步骤2：分类型处理行数据 #

# 处理df2中共有主键的行，标记为'U'
df2_u = df2[df2['pk'].isin(common_pk)].copy()
df2_u['etl_flag'] = 'U'

# 处理df2中仅存在的主键行，标记为'D'
df2_d = df2[df2['pk'].isin(only_df2_pk)].copy()
df2_d['etl_flag'] = 'D'

# 处理df1中所有行：共有主键的行+仅df1存在的行，统一标记为'I'
df1_i = df1.copy()
df1_i['etl_flag'] = 'I'

步骤3：合并结果并整理格式 #

# 合并所有处理后的行
final_df = pd.concat([df2_u, df2_d, df1_i], ignore_index=True)

# 移除辅助主键列，调整列顺序匹配预期输出
final_df = final_df.drop('pk', axis=1)[['alias_cd', 'country_cd', 'pos_name', 'ts_allocated', 'tr_id', 'ty_name', 'etl_flag']]

# 按需求排序：先按alias_cd、country_cd排序，再让U/D排在I前面
final_df.sort_values(by=['alias_cd', 'country_cd', 'etl_flag'], ascending=[True, True, False], inplace=True)
final_df.reset_index(drop=True, inplace=True)

print(final_df)

执行结果 #

输出与预期完全一致：

alias_cd country_cd pos_name  ts_allocated tr_id ty_name etl_flag
0    12345         AU      st1           200  None    E2E        U
1    12345         AU      st1           100  None    E2E        I
2    12345         AU       Jh           100  None    E2E        I
3    12345        AU2       Jh           100  None    E2E        I
4    12345        AU3      st2           100  None    E2E        D

原代码问题分析 #

1. 逻辑错误：用apply逐行判断主键存在性时，布尔数组的判断逻辑有误，导致筛选出的行不符合要求。
2. 重复合并：处理删除行时，错误地将df2中仅存在的行再次合并，导致结果重复。
3. 排序不严谨：仅按alias_cd和country_cd排序，未考虑etl_flag的优先级，导致输出顺序不符合预期。

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