如何在PySpark中将字典列表扁平化为单个字典

问题场景 #

你有一个PySpark DataFrame，其中features列存储的是字典列表，需要将其扁平化为单个字典，且无需显式指定数百个可能存在的键。

示例数据构造 #

首先构造符合场景的示例DataFrame：

from pyspark.sql import SparkSession
from pyspark.sql.types import StructType, StructField, IntegerType, StringType, ArrayType, MapType

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

# 定义数据Schema
schema = StructType([
    StructField("id", IntegerType(), nullable=False),
    StructField("label", IntegerType(), nullable=False),
    StructField("features", ArrayType(MapType(StringType(), StringType())), nullable=False)
])

# 示例数据
sample_data = [
    (1, 1, [{"key1": "1"}, {"key2": "dog"}, {"key4": "jane"}]),
    (3, 0, [{"key3": "3"}, {"key1": "3"}]),
    (4, 1, [{"key1": "4"}, {"key2": "bird"}]),
    (2, 1, [{"key2": "2"}, {"key3": "2"}]),
    (5, 0, [{"key2": "cat"}, {"key3": "5"}]),
    (6, 1, [{"key3": "6"}, {"key1": "6"}])
]

df = spark.createDataFrame(sample_data, schema=schema)
df.show(truncate=False)

解决方案 #

提供两种无需显式指定键的实现方式：

方法1：Python UDF实现

通过字典推导式合并列表中的所有字典，自动处理所有键：

from pyspark.sql.functions import udf

# 定义UDF：合并字典列表为单个字典
flatten_udf = udf(
    lambda dict_list: {k: v for d in dict_list for k, v in d.items()},
    MapType(StringType(), StringType())
)

# 应用UDF更新features列
result_df = df.withColumn("features", flatten_udf("features"))
result_df.show(truncate=False)

方法2：PySpark内置高阶函数（推荐）

使用aggregate和map_concat实现纯Spark原生处理，性能更优，适合大数据场景：

from pyspark.sql.functions import aggregate, map_concat, lit, col

# 聚合合并字典列表
result_df = df.withColumn(
    "features",
    aggregate(
        col("features"),
        lit({}).cast(MapType(StringType(), StringType())),  # 初始化空字典作为累加器
        lambda acc, dict_item: map_concat(acc, dict_item)   # 逐个合并字典
    )
)

result_df.show(truncate=False)

输出结果 #

两种方法都会得到符合目标格式的DataFrame：

+---+-----+------------------------------------------------+
|id |label|features                                        |
+---+-----+------------------------------------------------+
|1  |1    |{key1 -> 1, key2 -> dog, key4 -> jane}          |
|3  |0    |{key3 -> 3, key1 -> 3}                          |
|4  |1    |{key1 -> 4, key2 -> bird}                       |
|2  |1    |{key2 -> 2, key3 -> 2}                          |
|5  |0    |{key2 -> cat, key3 -> 5}                        |
|6  |1    |{key3 -> 6, key1 -> 6}                          |
+---+-----+------------------------------------------------+

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