PySpark循环Join优化咨询：产品编码转名称的最优实现

问题背景与需求 #

我有两个Spark DataFrame：products_df（存储各区域的产品编码）和categories_df（存储产品编码与名称、类别、区域的映射关系），具体结构如下：

products_df 结构 #

+------+------+------+-----------+
|region|laptop|mobile|conditioner|
+------+------+------+-----------+
| North|  L123|  M456|       C789|
|  West|  NULL|  M789|       C123|
|  NULL|  L456|  M123|       C456|
+------+------+------+-----------+

categories_df 结构 #

+------+--------------------+------------+-----------+
|region|        product_name|product_code|      class|
+------+--------------------+------------+-----------+
| North|      Laptop Model X|        L123|electronics|
| North|      Mobile Model Y|        M456|electronics|
|  NULL|      Mobile Model Z|        M789|electronics|
|  NULL|      Laptop Model Z|        L456|electronics|
|  NULL|      Mobile Model A|        M123|electronics|
|  West|  Conditioner Deluxe|        C123| appliances|
| North|     Conditioner Pro|        C789| appliances|
|  NULL|Conditioner Standard|        C456| appliances|
+------+--------------------+------------+-----------+

核心需求是：将products_df中laptop、mobile、conditioner列的产品编码替换为对应的产品名称。匹配规则为：

• 基于产品所属class（electronics或appliances）和region进行匹配
• 优先匹配同region的记录，若无匹配则选用region为NULL的兜底记录

同时，这类产品DataFrame可能包含不同的产品列，需要支持动态映射，因此我通过字典传入产品列与对应class的映射关系。

当前实现方案 #

我编写了如下Spark函数，通过循环对每个产品列单独执行左连接来完成映射：

'''
示例映射字典：
product_class_mapping = {
    "laptop": "electronics",
    "mobile": "electronics",
    "conditioner": "appliances",
}
'''
def set_product_name(products_df, categories_df, product_class_mapping):
    for product_col, product_class in product_class_mapping.items():
        products_df = (
            products_df.alias("p")
            .join(
                categories_df.alias("c"),
                (F.col("c.class") == product_class)
                & (F.col(f"p.{product_col}") == F.col("c.product_code"))
                & (F.col("c.region").isNull() | (F.col("p.region") == F.col("c.region"))),
                how="left",
            )
            .select(F.col("p.*"), F.col("c.product_name"))
            .withColumn(
                product_col,
                F.when(F.col(product_col).isNotNull(), F.col("product_name")).otherwise(
                    F.col(product_col)
                ),
            ).drop("product_name")
        )
        
    print("Final mapped products_df:")
    products_df.show()
    return products_df

该方案能生成预期输出：

Final mapped products_df:
+------+--------------+--------------+--------------------+
|region|        laptop|        mobile|         conditioner|
+------+--------------+--------------+--------------------+
| North|Laptop Model X|Mobile Model Y|     Conditioner Pro|
|  West|          NULL|Mobile Model Z|  Conditioner Deluxe|
|  NULL|Laptop Model Z|Mobile Model A|Conditioner Standard|
+------+--------------+--------------+--------------------+

疑问与求助 #

但当前方案通过循环多次执行Join操作，可能存在性能瓶颈。请问是否有更高效的替代实现方案？

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