如何在PySpark中处理DataFrame并输出指定结构的JSON结果？

这事儿在PySpark里其实挺容易实现的，核心思路就是分组聚合收集数组，再构造目标结构并转JSON，我给你一步步来演示：

1. 准备测试数据（复现你的场景） #

首先我们先创建和你示例一致的DataFrame：

from pyspark.sql import SparkSession
from pyspark.sql.functions import collect_list, struct, to_json, col

# 初始化SparkSession
spark = SparkSession.builder.appName("EmpJSONTransform").getOrCreate()

# 测试数据
sample_data = [(1, "A", 1, 1), (1, "A", 1, 1)]
emp_df = spark.createDataFrame(sample_data, schema=["empid", "empname", "in", "out"])

2. 分组聚合，收集in/out为数组 #

我们需要把同一个员工的in和out值分别收集成数组，这里用groupBy+collect_list就能搞定：

# 按empid和empname分组（假设empid和empname一一对应）
aggregated_df = emp_df.groupBy("empid", "empname").agg(
    collect_list("in").alias("in"),  # 收集所有in值为数组
    collect_list("out").alias("out") # 收集所有out值为数组
)

这一步之后，DataFrame的结构就变成了：empid | empname | in | out，对应数据就是1 | A | [1,1] | [1,1]。

3. 构造目标结构并转换为JSON #

接下来我们要把字段映射成你需要的id/name，再转成JSON格式，用struct定义结构，to_json完成转换：

# 构造目标结构体并转为JSON字符串
result_df = aggregated_df.withColumn(
    "target_json",
    to_json(
        struct(
            col("empid").alias("id"),    # 把empid重命名为id
            col("empname").alias("name"),# 把empname重命名为name
            col("in"),
            col("out")
        )
    )
)

# 查看最终结果
result_df.select("target_json").show(truncate=False)

执行后你会看到输出的JSON就是你要的格式：{"id":1,"name":"A","in":[1,1],"out":[1,1]}。

4. 输出到JSON文件（可选） #

如果需要把结果保存成JSON文件，直接用write.json即可：

# 覆盖模式写入指定路径
result_df.select("target_json").write.mode("overwrite").json("/path/to/your/output/dir")

小提醒 #

如果你的数据中存在empid对应多个empname的情况，建议先做数据清洗（比如按empid取唯一的empname），避免聚合后出现异常。

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