PySpark DataFrame数组列关联时，如何保持映射列数组值的顺序？

问题场景 #

原代码尝试通过array_contains关联两个DataFrame，再用groupBy+collect_list生成映射后的数组，但无法保持原数组的元素顺序：

import pyspark.sql.functions as F

df1 = spark.createDataFrame([(2, [3, 4]), (3, [4]),(4, [3,5]),(5, [4,5]),(6, [5,4])], ["a", "b"])
df2 = spark.createDataFrame([(3, "Three"), (4, "Four"),(5, "Five")], ["b", "c"])

df3 = df1.alias("df1").join(
    df2.alias("df2"),
    F.expr("array_contains(df1.b, df2.b)"),
    "left"
).groupBy("df1.a").agg(
    F.first("df1.b").alias("b"),
    F.collect_list("df2.c").alias("c")
)

df3.show()

执行结果中，a=6的行原数组b=[5,4]，但映射后的c数组为[Four, Five]，不符合预期的[Five, Four]。

问题原因 #

collect_list在groupBy聚合时不保证顺序，因为Spark分布式计算的shuffle和分区过程不会保留数据的原始顺序，join后的数据顺序是随机的，导致最终生成的数组顺序与原数组不一致。

解决方案 #

方法一：使用transform+map_from_arrays直接映射（推荐，性能更优） #

将df2转换为键值对Map，再通过transform遍历原数组，按顺序获取对应映射值，无需join和groupBy：

import pyspark.sql.functions as F

df1 = spark.createDataFrame([(2, [3, 4]), (3, [4]),(4, [3,5]),(5, [4,5]),(6, [5,4])], ["a", "b"])
df2 = spark.createDataFrame([(3, "Three"), (4, "Four"),(5, "Five")], ["b", "c"])

# 将df2转换为Map类型的常量
b_c_map = df2.agg(F.map_from_arrays(F.collect_list("b"), F.collect_list("c")).alias("b_c_map")).first()["b_c_map"]

# 遍历原数组b，按顺序映射得到c数组
df3 = df1.withColumn("c", F.transform("b", lambda x: F.lit(b_c_map).getItem(x)))

df3.show()

执行结果：

+---+------+-------------+
|  a|     b|            c|
+---+------+-------------+
|  2|[3, 4]|[Three, Four]|
|  3|   [4]|       [Four]|
|  4|[3, 5]|[Three, Five]|
|  5|[4, 5]| [Four, Five]|
|  6|[5, 4]| [Five, Four]|
+---+------+-------------+

方法二：保留数组元素索引后聚合（适合动态映射场景） #

如果df2数据需动态更新，可通过explode拆分数组时记录元素的原始索引，关联后按索引排序再聚合：

import pyspark.sql.functions as F

df1 = spark.createDataFrame([(2, [3, 4]), (3, [4]),(4, [3,5]),(5, [4,5]),(6, [5,4])], ["a", "b"])
df2 = spark.createDataFrame([(3, "Three"), (4, "Four"),(5, "Five")], ["b", "c"])

# 拆分数组并记录每个元素在原数组中的索引
df1_exploded = df1.withColumn("b_element", F.explode("b")) \
                  .withColumn("temp_id", F.monotonically_increasing_id()) \
                  .withColumn("array_index", F.row_number().over(F.partitionBy("a").orderBy("temp_id")) - 1) \
                  .drop("temp_id")

# 关联后按索引排序，再聚合回数组
df3 = df1_exploded.join(df2, df1_exploded.b_element == df2.b, "left") \
                  .orderBy("a", "array_index") \
                  .groupBy("a", "b") \
                  .agg(F.collect_list("c").alias("c"))

df3.show()

该方法通过array_index固定原数组元素的位置，确保聚合时的顺序与原数组一致。

内容的提问来源于stack exchange，提问作者Mohammad Mahfooz Alam