PySpark中如何遍历数组列并关联映射对应值？

我来帮你搞定这个数组映射的问题！你之前尝试用UDF但没成功，主要是写法上有误区，而且其实还有更高效的内置函数方案，我给你详细说明：

先准备测试数据 #

首先我们先把你给出的两个DataFrame构建出来，方便后续测试：

from pyspark.sql import SparkSession
from pyspark.sql import functions as F

spark = SparkSession.builder.appName("array_mapping").getOrCreate()

# 构建dataframe_a
data_a = [
    ("John", ["mango", "apple"]),
    ("Tom", ["mango", "orange"]),
    ("Matteo", ["apple", "banana"])
]
df_a = spark.createDataFrame(data_a, ["str1", "array_of_str"])

# 构建dataframe_b
data_b = [
    ("mango", 1),
    ("apple", 2),
    ("orange", 3)
]
df_b = spark.createDataFrame(data_b, ["key", "value"])

方案1：推荐用内置函数（无需UDF，效率更高） #

Spark 3.0及以上版本支持transform和filter内置函数，结合广播变量可以高效完成映射，避免UDF的性能损耗：

1. 先把dataframe_b转换成字典并广播（广播变量能让每个Executor只加载一次映射关系，提升分布式环境下的效率）
2. 用transform把数组每个元素替换成对应的值，不存在的元素会转为null
3. 用filter过滤掉数组中的null值，得到最终结果

# 将dataframe_b转为字典并广播
key_value_map = {row.key: row.value for row in df_b.collect()}
broadcast_map = spark.sparkContext.broadcast(key_value_map)

# 生成joined_result列
df_result = df_a.withColumn(
    "joined_result",
    F.filter(
        F.transform(
            "array_of_str",
            lambda elem: F.lit(broadcast_map.value.get(elem))
        ),
        lambda val: val.isNotNull()
    )
)

# 查看结果
df_result.show(truncate=False)

执行后就能得到你预期的结果：

+------+----------------+-------------+
|str1  |array_of_str    |joined_result|
+------+----------------+-------------+
|John  |[mango, apple]  |[1, 2]       |
|Tom   |[mango, orange] |[1, 3]       |
|Matteo|[apple, banana] |[2]          |
+------+----------------+-------------+

方案2：正确的UDF写法 #

如果你一定要用UDF，那需要修正之前的写法，注意要结合广播变量，并且正确定义UDF的输入输出类型：

from pyspark.sql.types import ArrayType, IntegerType

# 定义映射函数：遍历数组元素，只保留能找到对应值的元素
def map_array_elements(arr):
    return [broadcast_map.value.get(item) for item in arr if broadcast_map.value.get(item) is not None]

# 注册UDF，指定返回类型为整数数组
map_udf = F.udf(map_array_elements, ArrayType(IntegerType()))

# 应用UDF生成新列
df_result_udf = df_a.withColumn("joined_result", map_udf(F.col("array_of_str")))

df_result_udf.show(truncate=False)

你之前代码的问题分析 #

你之前的UDF写法存在几个问题：

• 直接用map(lambda x: ...)作为UDF的逻辑是错误的，UDF需要接收一个完整的函数，而不是map对象
• 没有正确指定UDF的返回类型（你写的ArrayType(StringType)不对，因为value是整数类型）
• 没有使用广播变量，在分布式环境下会重复加载dataframe_b的数据，效率极低
• 没有处理不存在的元素（比如banana），导致结果中会出现null或者错误值

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