如何用Apache Spark Scala将DataFrame转换为嵌套多级JSON？

原始DataFrame结构 #

+-------------------+-------+------+--------------------+
|         col_1     |  col_2| col_3|               col_4|
+-------------------+-------+------+--------------------+
|a                  |b      |c     |                  14|
|a                  |e      |f     |                  32|
|k                  |z      |x     |                  15|
+-------------------+-------+------+--------------------+

目标JSON格式 #

{"a": {"b": {"c": 14}, "e": {"f": 32}}, "k": {"z": {"x": 15}}}

解决方案：使用Spark内置函数即可，无需自定义UDF/UDAF #

可以通过Spark提供的create_map、map_concat（或map_merge+aggregate）、collect_list等内置函数组合实现，无需编写自定义UDF/UDAF。

1. Spark 3.0+ 版本实现（推荐） #

以PySpark为例：

from pyspark.sql import functions as F

# 初始化示例DataFrame
df = spark.createDataFrame(
    [("a", "b", "c", 14), ("a", "e", "f", 32), ("k", "z", "x", 15)],
    ["col_1", "col_2", "col_3", "col_4"]
)

# 构建嵌套映射并转为JSON
result_df = df.groupBy("col_1").agg(
    # 同col_1分组内，合并每个col_2对应的嵌套map
    F.map_concat(F.collect_list(F.create_map(F.col("col_2"), F.create_map(F.col("col_3"), F.col("col_4"))))).alias("nested_map")
).agg(
    # 合并所有col_1对应的map，并转为JSON字符串
    F.to_json(F.map_concat(F.collect_list(F.create_map(F.col("col_1"), F.col("nested_map"))))).alias("result_json")
)

# 查看输出
result_df.show(truncate=False)

Scala版本实现：

import org.apache.spark.sql.functions._

val df = spark.createDataFrame(Seq(
  ("a", "b", "c", 14), ("a", "e", "f", 32), ("k", "z", "x", 15)
)).toDF("col_1", "col_2", "col_3", "col_4")

val resultDf = df.groupBy("col_1")
  .agg(map_concat(collect_list(create_map(col("col_2"), create_map(col("col_3"), col("col_4"))))).alias("nested_map"))
  .agg(to_json(map_concat(collect_list(create_map(col("col_1"), col("nested_map"))))).alias("result_json"))

resultDf.show(false)

2. Spark 2.x 兼容版本 #

如果使用Spark 2.x（无map_concat函数），可以用aggregate+map_merge替代：

from pyspark.sql import functions as F

df = spark.createDataFrame(
    [("a", "b", "c", 14), ("a", "e", "f", 32), ("k", "z", "x", 15)],
    ["col_1", "col_2", "col_3", "col_4"]
)

result_df = df.groupBy("col_1").agg(
    F.collect_list(F.create_map(F.col("col_2"), F.create_map(F.col("col_3"), F.col("col_4")))).alias("map_list")
).agg(
    F.to_json(
        # 聚合所有col_1对应的map，合并为一个大map
        F.aggregate(
            F.collect_list(F.create_map(F.col("col_1"), 
                # 聚合组内的多个map，合并为单个map
                F.aggregate(F.col("map_list"), F.lit({}), lambda acc, m: F.map_merge(acc, m))
            )),
            F.lit({}),
            lambda acc, m: F.map_merge(acc, m)
        )
    ).alias("result_json")
)

result_df.show(truncate=False)

核心逻辑说明 #

• create_map：用于构建键值对映射，从最内层的col_3 -> col_4开始，逐层向外嵌套col_2 -> 内层map、col_1 -> 中层map。
• map_concat/map_merge：用于合并同组内的多个map，避免键值冲突（本例中同col_1下的col_2无重复，合并后自然形成嵌套结构）。
• collect_list：用于收集分组内的所有映射结果，为合并做准备。
• to_json：将最终的顶层map转换为符合要求的JSON字符串。

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