PySpark中如何实现带条件判断的多表Join操作

实现逻辑说明 #

PySpark 原生支持在 join 关联条件中写入分支判断逻辑，和现有SQL的逻辑完全对齐，不需要拆分多步关联再合并，直接用when/otherwise语法对应SQL中的if判断即可。
需要注意两个表的国家字段名大小写不一致（表X为country、表Y为Country），关联时要明确指定字段所属表，避免字段歧义；另外输出结果中非city匹配场景下city字段需要置空，要在select阶段做对应处理。

完整可运行代码 #

from pyspark.sql import SparkSession
from pyspark.sql.functions import col, when, lit

# 初始化Spark会话
spark = SparkSession.builder.appName("conditional_join_demo").getOrCreate()

# 构造表X测试数据集
x_data = [
    ("USA", "Boston", "David"),
    ("USA", "Miami", "John"),
    ("France", "Paris", "Peter")
]
df_x = spark.createDataFrame(x_data, schema=["country", "city", "user"])

# 构造表Y测试数据集
y_data = [
    ("USA", "city", "Boston", 1),
    ("USA", None, None, 2),
    ("France", None, None, 3)
]
df_y = spark.createDataFrame(y_data, schema=["Country", "detail", "value", "id"])

# 执行带条件的内关联
result_df = df_x.alias("x").join(
    df_y.alias("y"),
    on=(
        # 第一层关联条件：国家字段相等
        (col("x.country") == col("y.Country"))
        # 第二层条件：detail为city时匹配city和value，否则恒为真
        & when(
            col("y.detail") == lit("city"),
            col("x.city") == col("y.value")
        ).otherwise(lit(True))
    ),
    how="inner"
).select(
    col("y.Country").alias("Country"),
    col("y.id"),
    # 非city匹配场景下city字段置空，匹配预期输出
    when(col("y.detail") == lit("city"), col("x.city")).otherwise(lit(None)).alias("city"),
    col("x.user")
)

# 查看输出结果
result_df.show()

运行结果 #

执行上述代码后，输出的结果完全符合预期，共4条数据：

+-------+---+------+-----+
|Country| id|  city| user|
+-------+---+------+-----+
|    USA|  1|Boston|David|
|    USA|  2|  null|David|
|    USA|  2|  null| John|
| France|  3|  null|Peter|
+-------+---+------+-----+

注意事项 #

• 关联条件中多个判断逻辑拼接时，必须使用&表示逻辑与、|表示逻辑或，每个独立的判断表达式需要用括号包裹，否则会因为Python运算符优先级规则报错
• 不要把这类条件关联拆成「先匹配city规则再关联非city规则最后union」的写法，直接在on条件里写分支判断的执行性能更优
• 如果select阶段不对city字段做空值处理，非city匹配的行会直接带出表X的city值，和期望的输出格式不符

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