如何在PySpark中合并两个DataFrame：数组列与字符串列关联

这个问题其实很常见，处理数组类型的映射替换，核心思路是先拆再合——把数组拆成单个元素完成关联，再重新聚合回数组。下面给你一步步的解决方案，附带可运行的代码：

1. 准备测试数据 #

首先我们先创建和你示例一致的两个DataFrame，方便后续测试：

from pyspark.sql import SparkSession
from pyspark.sql.functions import explode, collect_list, col, posexplode

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

# 创建df1
data1 = [(1, ["A", "F"]), (2, ["G"]), (3, ["A"])]
df1 = spark.createDataFrame(data1, ["id", "code"])

# 创建df2
data2 = [("Apple", "A"), ("Google", "G"), ("Facebook", "F")]
df2 = spark.createDataFrame(data2, ["col1", "col2"])

2. 拆分数组并关联映射表 #

因为数组无法直接和另一张表的单个字段关联，所以我们先用explode把code数组拆分成多行，每个数组元素单独对应一行数据，然后和df2做关联：

# 拆分code数组，生成每个id对应的单个code元素
df_exploded = df1.select("id", explode(col("code")).alias("code_element"))

# 和df2关联，匹配code_element与col2，拿到对应的col1值
df_joined = df_exploded.join(df2, df_exploded.code_element == df2.col2, "left")

3. 聚合回数组并合并原数据 #

关联完成后，我们按id分组，用collect_list把匹配到的col1值重新聚合成数组，最后再和原df1关联，把code列加回来：

# 按id分组聚合，生成changed列
df_aggregated = df_joined.groupBy("id").agg(collect_list(col("col1")).alias("changed"))

# 关联原df1，恢复code列，得到目标DataFrame
df3 = df_aggregated.join(df1, on="id", how="inner").select("id", "code", "changed")

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

运行后就能得到你想要的结果：

+---+------+-----------------+
|id |code  |changed          |
+---+------+-----------------+
|1  |[A, F]|[Apple, Facebook]|
|2  |[G]   |[Google]         |
|3  |[A]   |[Apple]          |
+---+------+-----------------+

可选：保证数组顺序一致 #

如果你的业务需要严格保证changed数组的顺序和原code数组完全一致（因为collect_list的顺序可能受数据分区影响），可以改用posexplode来保留数组元素的位置索引，关联后按索引排序再聚合：

# 使用posexplode保留数组元素的位置
df_exploded = df1.select("id", posexplode(col("code")).alias("position", "code_element"))

# 关联后按id和位置排序，确保聚合顺序正确
df_joined = df_exploded.join(df2, df_exploded.code_element == df2.col2, "left")
df_aggregated = df_joined.orderBy("id", "position").groupBy("id").agg(collect_list(col("col1")).alias("changed"))

# 关联原df1得到结果
df3 = df_aggregated.join(df1, on="id", how="inner").select("id", "code", "changed")

这种方法能确保changed数组的元素顺序和原code数组完全对应，适合对顺序有严格要求的场景。

为什么不用UDF？ #

可能你会想到用自定义UDF来遍历数组元素做映射，但PySpark的内置函数（explode/collect_list）在性能上远优于UDF，尤其是处理大数据量时，内置函数是基于Spark的优化引擎执行，而UDF会触发Python-JVM的序列化开销，效率低很多。所以优先用内置函数方案。

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