将Key列转换为JSON结构：Spark DataFrame添加指定格式JSON列

现有DataFrame #

业务数据表df1 #

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

data2 = [("James","","Smith","36636","M",3000),
    ("Robert","","Williams","42114","M",4000),
    ("Maria","Anne","Jones","39192","F",4000),
  ]

schema = StructType([ \
    StructField("firstname",StringType(),True), \
    StructField("middlename",StringType(),True), \
    StructField("lastname",StringType(),True), \
    StructField("id", StringType(), True), \
    StructField("gender", StringType(), True), \
    StructField("salary", IntegerType(), True) \
  ])
  
df1 = spark.createDataFrame(data=data2,schema=schema)
df1.show(truncate=False)

输出结果：

+---------+----------+--------+-----+------+------+
|firstname|middlename|lastname|id   |gender|salary|
+---------+----------+--------+-----+------+------+
|James    |          |Smith   |36636|M     |3000  |
|Robert   |          |Williams|42114|M     |4000  |
|Maria    |Anne      |Jones   |39192|F     |4000  |
+---------+----------+--------+-----+------+------+

元数据表df2 #

df2 = spark.createDataFrame(
    [
    (["firstname", "lastname"],["001", "002"])    ],
    ["col_lst", "id_lst"]
)

df2.show(truncate=False)

输出结果：

+---------------------+----------+
|col_lst              |id_lst    |
+---------------------+----------+
|[firstname, lastname]|[001, 002]|
+---------------------+----------+

需求 #

为df1新增一列（命名为JSON），该列需以指定JSON结构存储df2中对应(id, value)的组合，具体结构示例：

{"business_key":[{"id":"001","value":"对应字段值"},{"id":"002","value":"对应字段值"}]}

解决方案 #

# 提取df2中的字段列表和ID列表
col_list = df2.select(F.col("col_lst")).first()[0]
id_list = df2.select(F.col("id_lst")).first()[0]

# 构造每个(id, value)的结构体
struct_list = [F.struct(F.lit(id).alias("id"), F.col(col).alias("value")) for col, id in zip(col_list, id_list)]

# 生成目标JSON列
result_df = df1.withColumn(
    "JSON",
    F.to_json(F.struct(F.array(*struct_list).alias("business_key")))
)

result_df.show(truncate=False)

最终输出结果 #

+---------+----------+--------+-----+------+------+--------------------------------------------------------------------------------+
|firstname|middlename|lastname|id   |gender|salary|JSON                                                                            |
+---------+----------+--------+-----+------+------+--------------------------------------------------------------------------------+
|James    |          |Smith   |36636|M     |3000  |{"business_key":[{"id":"001","value":"James"},{"id":"002","value":"Smith"}]}    |
|Robert   |          |Williams|42114|M     |4000  |{"business_key":[{"id":"001","value":"Robert"},{"id":"002","value":"Williams"}]}|
|Maria    |Anne      |Jones   |39192|F     |4000  |{"business_key":[{"id":"001","value":"Maria"},{"id":"002","value":"Jones"}]}    |
+---------+----------+--------+-----+------+------+--------------------------------------------------------------------------------+

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