Spark DataFrame转指定嵌套JSON失败，求正确实现方案

问题背景 #

需要将给定的Spark DataFrame转换为特定嵌套结构的JSON并写入文件，当前使用的代码无法生成预期结果。

原始DataFrame #

|       Trans_Type   |Payment_Due_Dt|BILL_CYCLE_ID|BALANCE_AMT|
+--------------------+--------------+-------------+-----------+
|  New Transaction   |    26-01-2015|   31-12-2014|      61.20|
|Current Transactions|    26-01-2015|   31-12-2014|     289.02|
|  Age Transaction   |              |             |       0.00|
|Clear Transaction   |              |             |       0.00|
|      Remittances   |              |             |     -61.20|
+--------------------+--------------+-------------+-----------+

DataFrame Schema #

spark_df.printSchema()
root
 |-- Trans_Type: string (nullable = false)
 |-- Payment_Due_Dt: string (nullable = true)
 |-- BILL_CYCLE_ID: string (nullable = true)
 |-- BALANCE_AMT: decimal(25,2) (nullable = false)

计算规则与预期JSON #

Account Balance为sum(balance_amt) = 647.00，预期输出JSON结构如下：

{
    "Summary": [{
            "Age Transactions": {
                "Payment_Due_Dt": " ",
                "BILL_CYCLE_ID": " ",
                "BALANCE_AMT": 0.00
            }
        },
        {
            "Clear Transactions": {
                "Payment_Due_Dt": " ",
                "BILL_CYCLE_ID": " ",
                "BALANCE_AMT": 0.00
            }
        },
        {
            "Current Transactions": {
                "Payment_Due_Dt": "26-01-2015",
                "BILL_CYCLE_ID": "31-12-2014",
                "BALANCE_AMT": 289.02
            }
        },
        {
            "New Transactions": {
                "Payment_Due_Dt": "26-01-2015",
                "BILL_CYCLE_ID": "31-12-2014",
                "BALANCE_AMT": 61.20
            }   
        },
        {
            "Remittances" : {
                "BALANCE_AMT": 61.20
            }
            
        },
        {
            "Account Balance": {
            "BAL_AM": "647.00"
            }
        }
    ]
}

用户当前代码（未得到预期结果） #

list1 = [Txn_Sum_base2_DF["Payment_Due_Dt"], Txn_Sum_base2_DF["BILL_CYCLE_ID"], Txn_Sum_base2_DF["BALANCE_AMT"] ]
df = Txn_Sum_base2_DF.groupBy("Trans_Type") \
            .agg(collect_list(struct(list1)).alias("summary")).toJSON().collect()

正确实现步骤 #

要生成预期的嵌套JSON，需分阶段处理数据并构造结构：

完整代码实现 #

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

# 1. 预处理原始DataFrame
processed_df = spark_df \
    # 统一交易类型名称
    .withColumn("Trans_Type", 
                F.when(F.col("Trans_Type") == "New Transaction", "New Transactions")
                .when(F.col("Trans_Type") == "Age Transaction", "Age Transactions")
                .when(F.col("Trans_Type") == "Clear Transaction", "Clear Transactions")
                .otherwise(F.col("Trans_Type"))) \
    # 空值替换为空格
    .withColumn("Payment_Due_Dt", F.coalesce(F.col("Payment_Due_Dt"), F.lit(" "))) \
    .withColumn("BILL_CYCLE_ID", F.coalesce(F.col("BILL_CYCLE_ID"), F.lit(" "))) \
    # Remittances金额取绝对值
    .withColumn("BALANCE_AMT", 
                F.when(F.col("Trans_Type") == "Remittances", F.abs(F.col("BALANCE_AMT")))
                .otherwise(F.col("BALANCE_AMT")))

# 2. 计算并生成Account Balance行
total_balance = processed_df.agg(F.sum("BALANCE_AMT").alias("BAL_AM")).first()["BAL_AM"]
balance_schema = StructType([
    StructField("Trans_Type", StringType(), False),
    StructField("BAL_AM", StringType(), False)
])
balance_row = spark.createDataFrame([("Account Balance", str(total_balance))], schema=balance_schema)

# 3. 构造每个交易类型的嵌套结构
# 处理普通交易行
transaction_items = processed_df \
    .withColumn("details", 
                F.struct(
                    F.when(F.col("Trans_Type") != "Remittances", F.col("Payment_Due_Dt")).alias("Payment_Due_Dt"),
                    F.when(F.col("Trans_Type") != "Remittances", F.col("BILL_CYCLE_ID")).alias("BILL_CYCLE_ID"),
                    F.col("BALANCE_AMT")
                )) \
    # 过滤Remittances不需要的空字段
    .withColumn("details", F.expr("transform(details, (k, v) -> if(v is not null, named_struct(k, v), null))")) \
    .withColumn("details", F.expr("aggregate(details, cast(map() as map<string, string>), (acc, x) -> if(x is not null, map_concat(acc, x), acc))")) \
    .withColumn("summary_item", F.create_map(F.col("Trans_Type"), F.col("details")))

# 处理Account Balance行
balance_item = balance_row \
    .withColumn("details", F.struct(F.col("BAL_AM"))) \
    .withColumn("summary_item", F.create_map(F.col("Trans_Type"), F.col("details")))

# 4. 聚合所有条目并生成最终结构
final_df = transaction_items.select("summary_item").union(balance_item.select("summary_item")) \
    .agg(F.collect_list("summary_item").alias("Summary"))

# 写入JSON文件
final_df.write.mode("overwrite").json("/path/to/your/output.json")

# 查看结果（可选）
result_json = final_df.toJSON().collect()[0]
print(result_json)

代码说明 #

• 预处理阶段：统一交易类型命名、替换空值、调整Remittances金额，确保数据符合预期格式。
• Account Balance处理：单独计算总余额并生成对应行，保证最终JSON包含该汇总条目。
• 嵌套结构构造：使用create_map将交易类型作为键，详情作为值；通过transform和aggregate过滤Remittances不需要的字段，避免生成冗余键值对。
• 最终聚合：将所有条目收集到Summary数组中，生成符合要求的嵌套JSON结构。

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