PySpark如何从数组列生成元素带固定公共键的字典数组列

问题根因 #

你原有代码的错误主要来自3个方面：

• 逐行过滤赋值逻辑错误：每次df.where(df.increasing_id == j)后生成的新df仅保留当前id的行，多次赋值后只会保留最后一次循环的行，其余行全部丢失
• 索引越界原因：你固定循环3次，若数据中不存在对应id的行，collect()[0][0]会直接拿空报错
• 类型和逻辑不符：你定义的UDF返回单个Map类型，和你需要的Array[Map]类型不匹配，另外UDF里写的key是name，和需求要求的mapname不符

最优实现方案（无需UDF，性能最优） #

Spark 2.4及以上版本可以直接用内置的transform高阶函数处理数组列，不需要自定义UDF也不需要循环：

from pyspark.sql import functions as F

df = df.withColumn(
    "mappings",
    F.transform(
        F.col("array_list"),
        lambda item: F.create_map(F.lit("mapname"), item)
    )
)

这个方法会自动遍历array_list的每一个元素，为每个元素生成{"mapname": 元素值}的map，最终拼接成数组，完全符合你要求的输出结构和schema。

低版本Spark兼容方案（UDF实现） #

如果你的Spark版本低于2.4不支持transform函数，可以用自定义UDF实现，不需要循环逐行处理：

from pyspark.sql import functions as F
from pyspark.sql import types as T

# 定义UDF，返回类型为数组嵌套Map
@udf(T.ArrayType(T.MapType(T.StringType(), T.StringType())))
def generate_mappings(arr):
    if not arr:
        return []
    return [{"mapname": item} for item in arr]

df = df.withColumn("mappings", generate_mappings(F.col("array_list")))

验证结果 #

执行完后查看schema和数据，和你预期的完全一致：

root
 |-- array_list: array (nullable = true)
 |    |-- element: string (containsNull = true)
 |-- len_of_array: integer (nullable = false)
 |-- mappings: array (nullable = true)
 |    |-- element: map (containsNull = true)
 |    |    |-- key: string
 |    |    |-- value: string (valueContainsNull = true)

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