Spark JSON转关系型CSV工具开发：嵌套数组转列需求及代码优化

需求说明 #

开发一款工具接收任意JSON输入并转换为扁平化表格，需处理嵌套对象与数组：

• 嵌套对象正常展开为列
• 嵌套数组需将其中的两个值转换为新列，而非现有代码那样将数组元素展开为多行，最终结构便于导入关系型数据库或分析。

修改后的代码 #

from typing import Dict
from pyspark.sql.session import SparkSession
from pyspark.sql import DataFrame as SDF
from pyspark.sql.functions import *
from pyspark.sql.types import StructType, ArrayType

def rename_dataframe_cols(df: SDF, col_names: Dict[str, str]) -> SDF:
    """重命名DataFrame指定列"""
    return df.select(*[col(col_name).alias(col_names.get(col_name, col_name)) for col_name in df.columns])

def update_column_names(df: SDF, index: int) -> SDF:
    """为所有列添加层级索引后缀"""
    new_cols = {column: f"{column}*{index}" for column in df.columns}
    return df.transform(lambda df_x: rename_dataframe_cols(df_x, new_cols))

def flatten_json(df_arg: SDF, index: int = 1) -> SDF:
    """递归扁平化JSON结构，将数组的两个元素转为列而非多行"""
    # 初始调用时为所有列添加层级1后缀
    df = update_column_names(df_arg, index) if index == 1 else df_arg

    fields = df.schema.fields

    for field in fields:
        data_type = str(field.dataType)
        column_name = field.name
        first_10_chars = data_type[:10]

        # 处理数组类型：提取前两个元素转为新列
        if first_10_chars == 'ArrayType(':
            # 提取数组第0、1位元素（Spark数组索引从1开始），支持空数组或长度不足的情况
            col_0 = element_at(col(column_name), 1).alias(f"{column_name}_0*{index+1}")
            col_1 = element_at(col(column_name), 2).alias(f"{column_name}_1*{index+1}")
            
            # 添加新列并删除原数组列
            df_temp = df.withColumn(col_0._alias, col_0).withColumn(col_1._alias, col_1).drop(column_name)
            return flatten_json(df_temp, index + 1)

        # 处理嵌套结构体：递归展开
        elif first_10_chars == 'StructType':
            current_col = column_name
            append_str = current_col

            # 避免列名冲突，若结构体包含自身同名字段则重命名
            if column_name in str(df.schema[current_col].dataType):
                df_temp = df.withColumnRenamed(column_name, f"{column_name}#1")
                current_col = f"{column_name}#1"
            else:
                df_temp = df

            # 展开结构体字段
            struct_fields = df_temp.select(f"{current_col}.*").columns
            # 获取当前层级，计算下一层级
            level = int(append_str.split('*')[-1])
            next_level = level + 1
            # 为展开后的字段添加层级前缀
            custom_cols = {field: f"{append_str}->{field}*{next_level}" for field in struct_fields}
            
            # 合并原字段与展开后的结构体字段，删除原结构体列
            df_temp2 = df_temp.select("*", f"{current_col}.*").drop(current_col)
            df_temp3 = df_temp2.transform(lambda df_x: rename_dataframe_cols(df_x, custom_cols))
            return flatten_json(df_temp3, index + 1)

    return df

关键改动说明 #

• 数组处理逻辑替换：移除原有的explode_outer多行展开逻辑，改用element_at提取数组前两位元素生成新列，空数组或长度不足时自动填充null，保持与结构体一致的层级命名规则。
• 结构体展开优化：简化层级计算逻辑，通过拆分列名后缀获取当前层级，保留原有的列名冲突处理逻辑，避免嵌套字段歧义。

使用示例 #

输入JSON：

{"id":1,"info":{"name":"test"},"values":[100,200]}

处理后得到的扁平列：

• id*1: 1
• info->name*2: "test"
• values_0*2: 100
• values_1*2: 200

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