Spark中连接DataFrame后批量重命名重复列的方法

问题场景 #

当用Spark做DataFrame关联（比如示例里的full join）时，若关联的DataFrame存在同名非关联列，结果会出现重复列名。示例代码如下：

vals1 = [(1, "a"), 
        (2, "b"), 
      ]
columns1 = ["id","name"]
df1 = spark.createDataFrame(data=vals1, schema=columns1)

vals2 = [(1, "k"), 
      ]
columns2 = ["id","name"]
df2 = spark.createDataFrame(data=vals2, schema=columns2)

df1 = df1.alias('df1').join(df2.alias('df2'), 'id', 'full')
df1.show()

执行后结果包含1个id列和2个name列，当存在数十个这类重复列时，可通过以下方法批量处理：

解决方案 #

方法1：关联前提前重命名列 #

在join操作前给重复列加上来源表前缀，从根源避免重复，这是最高效的方式：

# 给df2的非关联列统一加前缀
df2_renamed = df2.withColumnsRenamed({col: f"df2_{col}" for col in df2.columns if col != "id"})
# 执行关联
joined_df = df1.join(df2_renamed, on="id", how="full")
joined_df.show()

处理后结果列名为id、name、df2_name，无重复问题。

方法2：关联后批量重命名重复列 #

若已完成join，可结合DataFrame的别名，遍历列名批量添加前缀：

from pyspark.sql.functions import col

new_columns = []
for c in df1.columns:
    if c == "id":
        new_columns.append(c)
    else:
        # 给两个来源表的重复列分别加前缀
        new_columns.append(col(f"df1.{c}").alias(f"df1_{c}"))
        new_columns.append(col(f"df2.{c}").alias(f"df2_{c}"))

# 生成重命名后的DataFrame
renamed_df = df1.select(*new_columns)
renamed_df.show()

此方法自动给所有重复列加上对应表的别名前缀，适合批量处理大量重复列。

方法3：自动检测重复列并重命名 #

若不确定哪些列重复，可先统计列名出现次数，再对重复列批量重命名：

from collections import Counter

# 统计列名出现频次
col_counts = Counter(df1.columns)
# 筛选出重复列名
duplicate_cols = [col for col, count in col_counts.items() if count > 1]

# 生成重命名规则
renamed_cols = []
for col_name in df1.columns:
    if col_name in duplicate_cols:
        # 给两个来源表的重复列加前缀
        for alias in ["df1", "df2"]:
            renamed_cols.append(f"{alias}.{col_name} as {alias}_{col_name}")
    else:
        renamed_cols.append(col_name)

# 通过selectExpr执行重命名
final_df = df1.selectExpr(*renamed_cols)
final_df.show()

这种方法无需提前知晓重复列，适合动态处理不同关联场景。

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