对比两个DataFrame列名筛选共同列，填充字典结果不符求助

需求与问题场景 #

目标是将两个DataFrame（dfUpdates和df_metadata）共有的列，按主键（PRIMARY KEY）和非主键分类存入dictionary_of_columns，期望输出为{'BK_COLUMNS': ['CODE'], 'COLUMNS': ['DESCRIPTION']}，但当前代码执行后字典包含过多列，不符合预期。

原代码 #

list_of_columns = []
for column in dfUpdates.schema:
    list_of_columns.append(column.jsonValue()["name"].upper())

dictionary_of_columns = {}
dictionary_of_columns['BK_COLUMNS'] = []
dictionary_of_columns['COLUMNS'] = []
for row in df_metadata.dropDuplicates(['COLUMN_NAME', 'KeyType']).collect():
    if row.KeyType == 'PRIMARY KEY' and row.COLUMN_NAME.upper() in list_of_columns:
        dictionary_of_columns['BK_COLUMNS'].append(row.COLUMN_NAME.upper())
    elif row.KeyType  != 'PRIMARY KEY' and row.COLUMN_NAME.upper() in list_of_columns:
        dictionary_of_columns['COLUMNS'].append(row.COLUMN_NAME.upper())

问题根源分析 #

1. 重复条目未彻底清理：df_metadata.dropDuplicates(['COLUMN_NAME', 'KeyType'])会保留同一列名对应不同KeyType的记录，比如某列同时被标记为PRIMARY KEY和其他类型时，会被多次处理，导致同一列重复出现在结果中。
2. 列名查询效率低且易出错：用列表存储dfUpdates列名，查询匹配时效率低；同时遍历逻辑中重复判断列名是否存在，冗余且可能引入逻辑漏洞。
3. 未确保列的唯一性：同一列名在df_metadata中可能对应多条记录，导致重复添加。

修正后的代码 #

# 提取dfUpdates的列名并转为大写集合，大幅提高查询效率
df_updates_columns = {col.jsonValue()["name"].upper() for col in dfUpdates.schema}

dictionary_of_columns = {'BK_COLUMNS': [], 'COLUMNS': []}

# 先按KeyType排序（优先保留PRIMARY KEY记录），再按COLUMN_NAME去重，确保同一列只处理一次
df_metadata_clean = df_metadata.orderBy(df_metadata.KeyType == 'PRIMARY KEY', ascending=False)\
                               .dropDuplicates(['COLUMN_NAME'])

for row in df_metadata_clean.collect():
    col_name_upper = row.COLUMN_NAME.upper()
    if col_name_upper in df_updates_columns:
        if row.KeyType == 'PRIMARY KEY':
            dictionary_of_columns['BK_COLUMNS'].append(col_name_upper)
        else:
            dictionary_of_columns['COLUMNS'].append(col_name_upper)

print(dictionary_of_columns)

额外排查步骤 #

• 打印df_updates_columns，确认是否与dfUpdates实际列名一致，排除列名提取错误。
• 打印df_metadata_clean内容，验证每个列名仅对应一个KeyType，无重复条目。
• 检查df_metadata中是否存在同一列名对应多个KeyType的情况，这是导致多余列的常见原因。

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