如何在Spark中高效将数组列关联lookup表生成对应ID数组

我有一张包含多列的表base_df，其中一列是数组列FruitNames，数据结构如下：

还有一张表lookup_df，每行对应数组中的一个元素：

需求 #

希望给base_df添加一列FruitIds，该列是lookup_df中对应名称的ID数组，期望结果如下：

现有代码 #

base_df = spark.createDataFrame(
    [
        (1, ["apple", "banana", "orange"], "...", "...", "...", "...", "..."),
        (2, ["banana"], "...", "...", "...", "...", "..."),
        (3, ["grape", "apple"], "...", "...", "...", "...", "..."),
        (4, [], "...", "...", "...", "...", "..."),
        (5, None, "...", "...", "...", "...", "..."),
        (6, ["orange", "grape"], "...", "...", "...", "...", "..."),
    ],
    ["Id", "FruitNames", "Col1", "Col2", "Col3", "...", "Col99"]
)

lookup_df = spark.createDataFrame(
    [
        (1, "apple"),
        (2, "banana"),
        (3, "orange"),
        (4, "grape"),
    ],
    ["Id", "Name"]
)

df_with_fruitids = ???

限制条件 #

• 不能用explode，因为它对空数组和null不安全；
• 虽然可以用explode_outer结合case语句再用collect_list生成列，但需要对大量列做group by，而base_df有1亿行、lookup_df有10万行，这种方法效率太低；
• 考虑过用base_df.withColumn("FruitIds", transform(col("FruitNames"), my_func))，但不知道纯Spark下怎么实现my_func，希望找到无需Python UDF的纯Spark解决方案。

可以通过将lookup_df转换成Map类型常量，结合transform函数实现，全程使用Spark内置函数，无UDF，性能高效：

步骤1：将lookup表转为Name到Id的映射Map #

from pyspark.sql import functions as F

# 生成Name与Id的映射Map，广播到各节点
fruit_name_to_id = lookup_df.agg(
    F.create_map(F.collect_list("Name"), F.collect_list("Id")).alias("fruit_map")
).first()["fruit_map"]

步骤2：用transform函数生成FruitIds列 #

遍历FruitNames数组，通过Map查找对应Id，同时处理空数组和null场景：

df_with_fruitids = base_df.withColumn(
    "FruitIds",
    F.when(
        F.col("FruitNames").isNull(),
        F.lit(None)
    ).otherwise(
        F.transform(
            F.col("FruitNames"),
            lambda fruit_name: F.lit(fruit_name_to_id).getItem(fruit_name)
        )
    )
)

方案说明 #

• 无需展开数组，全程在单条记录层面处理，避免group by带来的性能损耗；
• 生成的Map会被Spark自动广播到各个节点，10万行的lookup表完全适配；
• 自动处理空数组（返回空数组）和null（返回null）的情况，符合需求；
• 所有操作均为Spark内置函数，性能远高于Python UDF。

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