如何为指定结构的Spark DataFrame创建键值对数组列？

没问题！要在你的Spark DataFrame里创建包含键值对的数组列，用PySpark的内置函数就能轻松实现，我给你分几种常见场景来演示：

1. 手动指定所有列生成键值对数组 #

如果你的列不多，可以直接手动构造每个键值对的结构体，再组合成数组：

首先，先确保导入需要的函数：

from pyspark.sql.functions import array, struct, lit

然后对原DataFrame进行转换：

# 给原DataFrame添加键值对数组列，命名为key_value_array
df_with_kv = df.withColumn(
    "key_value_array",
    array(
        struct(lit("Id").alias("key"), df["Id"].alias("value")),
        struct(lit("Size").alias("key"), df["Size"].alias("value")),
        struct(lit("Amt").alias("key"), df["Amt"].alias("value"))
    )
)

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

执行后得到的结果如下：

+---+----+----+----------------------------------------------------------+
|Id |Size|Amt |key_value_array                                           |
+---+----+----+----------------------------------------------------------+
|a1 |1   |55.0|[{Id, a1}, {Size, 1}, {Amt, 55.0}]                        |
|a2 |2   |48.0|[{Id, a2}, {Size, 2}, {Amt, 48.0}]                        |
|a3 |3   |28.0|[{Id, a3}, {Size, 3}, {Amt, 28.0}]                        |
+---+----+----+----------------------------------------------------------+

2. 动态生成所有列的键值对数组 #

如果DataFrame的列很多，手动写每个结构体太麻烦，可以用列表推导式动态生成：

# 遍历所有列，自动生成每个列对应的键值对结构体
kv_structs = [struct(lit(col).alias("key"), df[col].alias("value")) for col in df.columns]

# 添加数组列
df_with_kv = df.withColumn("key_value_array", array(*kv_structs))

这种方式不管后续列怎么变化，都能自动适配，非常灵活。

3. 只生成指定列的键值对数组 #

如果不需要所有列，只想选取部分列生成键值对数组，比如只保留Id和Amt：

# 指定需要转换的列
target_cols = ["Id", "Amt"]
kv_structs = [struct(lit(col).alias("key"), df[col].alias("value")) for col in target_cols]

df_with_selected_kv = df.withColumn("selected_key_value_array", array(*kv_structs))
df_with_selected_kv.show(truncate=False)

输出结果：

+---+----+----+--------------------------+
|Id |Size|Amt |selected_key_value_array  |
+---+----+----+--------------------------+
|a1 |1   |55.0|[{Id, a1}, {Amt, 55.0}]   |
|a2 |2   |48.0|[{Id, a2}, {Amt, 48.0}]   |
|a3 |3   |28.0|[{Id, a3}, {Amt, 28.0}]   |
+---+----+----+--------------------------+

输出JSON的效果 #

如果把生成后的DataFrame写入JSON，每条数据的格式会是这样（以第一行为例）：

{"Id":"a1","Size":1,"Amt":55.0,"key_value_array":[{"key":"Id","value":"a1"},{"key":"Size","value":1},{"key":"Amt","value":55.0}]}

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