如何将Spark DataFrame的JSON列嵌套对象转为对象数组

需求说明 #

现有仅含TEST列的Spark DataFrame，列中JSON结构如下：

{
"TAXES": {
"TAX":
{
"TAXFREEFLAG": 0,
"TAXINCLUDEDFLAG": 1,
"TAXPERCENT": 5.5,
"TAXTYPE": 49,
"TICKETTAXABLEAMOUNT": 6.14,
"TICKETTAXAMOUNT": 0.34,
"TICKETTAXID": 1
}
}
}

需要将TAX字段从单个对象转换为包含该对象的数组，最终JSON结构如下：

{
"TAXES": {
"TAX": [
{
"TAXFREEFLAG": 0,
"TAXINCLUDEDFLAG": 1,
"TAXPERCENT": 5.5,
"TAXTYPE": 49,
"TICKETTAXABLEAMOUNT": 6.14,
"TICKETTAXAMOUNT": 0.34,
"TICKETTAXID": 1
}
]
}
}

方案一：Schema解析法（推荐） #

通过定义JSON结构Schema，解析后重构数组字段再转回JSON格式，稳定性高，适配字段变化场景。

PySpark 实现代码

from pyspark.sql import functions as F
from pyspark.sql.types import StructType, StructField, IntegerType, DoubleType

# 定义TAX字段的Schema
tax_schema = StructType([
    StructField("TAXFREEFLAG", IntegerType()),
    StructField("TAXINCLUDEDFLAG", IntegerType()),
    StructField("TAXPERCENT", DoubleType()),
    StructField("TAXTYPE", IntegerType()),
    StructField("TICKETTAXABLEAMOUNT", DoubleType()),
    StructField("TICKETTAXAMOUNT", DoubleType()),
    StructField("TICKETTAXID", IntegerType())
])

# 定义TAXES字段的Schema
taxes_schema = StructType([StructField("TAX", tax_schema)])

# 定义根级Schema
root_schema = StructType([StructField("TAXES", taxes_schema)])

# 解析TEST列的JSON结构
df_parsed = df.withColumn("parsed_json", F.from_json(F.col("TEST"), root_schema))

# 将单个TAX对象包装为数组
df_with_array = df_parsed.withColumn(
    "tax_array",
    F.array(F.col("parsed_json.TAXES.TAX"))
)

# 重构JSON结构并转回字符串格式
df_result = df_with_array.withColumn(
    "TEST",
    F.to_json(
        F.struct(
            F.struct(F.col("tax_array").alias("TAX")).alias("TAXES")
        )
    )
).drop("parsed_json", "tax_array")

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

Scala 实现代码

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

// 定义TAX字段的Schema
val taxSchema = StructType(Seq(
  StructField("TAXFREEFLAG", IntegerType),
  StructField("TAXINCLUDEDFLAG", IntegerType),
  StructField("TAXPERCENT", DoubleType),
  StructField("TAXTYPE", IntegerType),
  StructField("TICKETTAXABLEAMOUNT", DoubleType),
  StructField("TICKETTAXAMOUNT", DoubleType),
  StructField("TICKETTAXID", IntegerType)
))

// 定义TAXES字段的Schema
val taxesSchema = StructType(Seq(StructField("TAX", taxSchema)))

// 定义根级Schema
val rootSchema = StructType(Seq(StructField("TAXES", taxesSchema)))

// 解析TEST列的JSON结构
val dfParsed = df.withColumn("parsedJson", from_json(col("TEST"), rootSchema))

// 将单个TAX对象包装为数组
val dfWithArray = dfParsed.withColumn(
  "taxArray",
  array(col("parsedJson.TAXES.TAX"))
)

// 重构JSON结构并转回字符串格式
val dfResult = dfWithArray.withColumn(
  "TEST",
  to_json(
    struct(
      struct(col("taxArray").alias("TAX")).alias("TAXES")
    )
  )
).drop("parsedJson", "taxArray")

// 查看结果
dfResult.show(false)

方案二：字符串拼接法（快速实现） #

如果确认原始JSON结构固定、无特殊字符干扰，可直接通过字符串拼接快速实现，缺点是容错性低。

PySpark 实现代码

from pyspark.sql import functions as F

df_result = df.withColumn(
    "TEST",
    F.concat(
        F.lit('{"TAXES": {"TAX": ['),
        F.get_json_object(F.col("TEST"), "$.TAXES.TAX"),
        F.lit(']}}')
    )
)

df_result.show(truncate=False)

Scala 实现代码

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

val dfResult = df.withColumn(
  "TEST",
  concat(
    lit("{\"TAXES\": {\"TAX\": ["),
    get_json_object(col("TEST"), "$.TAXES.TAX"),
    lit("]}}")
  )
)

dfResult.show(false)

内容的提问来源于stack exchange，提问作者Asri Med Amin