如何编写通用函数实现两个PySpark DataFrame的内连接？

可以实现无需手动指定连接列的方案，核心思路是自动识别两个DataFrame的所有公共列作为连接键，利用PySpark的join特性自动合并这些公共列（不会重复保留），从而完成内连接并去除重复公共列的需求。

实现函数 #

def inner_join_remove_duplicate_cols(df1, df2):
    # 获取两个DataFrame的公共列
    common_cols = list(set(df1.columns) & set(df2.columns))
    if not common_cols:
        raise ValueError("两个DataFrame没有公共列，无法执行基于公共列的内连接")
    
    # 检查公共列的数据类型是否一致（避免连接失败）
    for col in common_cols:
        dtype_df1 = df1.schema[col].dataType
        dtype_df2 = df2.schema[col].dataType
        if dtype_df1 != dtype_df2:
            raise TypeError(f"公共列 `{col}` 在两个DataFrame中的数据类型不匹配：{dtype_df1} vs {dtype_df2}")
    
    # 执行内连接，公共列仅保留一份
    joined_df = df1.join(df2, on=common_cols, how="inner")
    
    return joined_df

示例用法 #

from pyspark.sql import SparkSession

# 初始化SparkSession
spark = SparkSession.builder.appName("AutoJoinExample").getOrCreate()

# 创建测试DataFrame
data1 = [("Alice", 25, "New York"), ("Bob", 30, "London")]
df1 = spark.createDataFrame(data1, ["name", "age", "city"])

data2 = [("Alice", 25, "Data Engineer"), ("Bob", 30, "Data Scientist")]
df2 = spark.createDataFrame(data2, ["name", "age", "job"])

# 调用函数执行自动连接
result_df = inner_join_remove_duplicate_cols(df1, df2)
result_df.show()

输出结果 #

+-----+---+--------+-------------+
| name|age|    city|          job|
+-----+---+--------+-------------+
|Alice| 25|New York|Data Engineer|
|  Bob| 30|  London|Data Scientist|
+-----+---+--------+-------------+

关键说明 #

• 自动连接逻辑：通过集合交集获取两个DF的所有公共列，将其作为连接键执行内连接。PySpark在join时，on参数指定的列会自动合并，不会重复保留来自两个DF的同名连接键列。
• 类型检查：加入数据类型校验是因为PySpark要求连接键的数据类型必须一致，否则会抛出异常。
• 异常处理：当两个DF无公共列时，直接抛出异常，因为无法基于公共列完成内连接。

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