如何用PySpark编写通用函数扁平化含嵌套数组的复杂JSON Schema

你的问题很典型——原代码只能处理嵌套的Struct类型，完全忽略了Array类型的处理，所以数组内部的Struct元素没有被扁平化。下面我会给出一个通用的递归扁平化函数，可以自动处理任意深度的Struct和Array嵌套结构，完全适配你需要的输出格式。

问题分析 #

原代码的核心局限是：只检查字段类型是否为struct，对array类型直接跳过，导致数组列保留原样，内部的Struct元素无法展开。要解决这个问题，我们需要递归遍历整个Schema，同时处理Struct和Array两种嵌套类型：

• 对于Struct：递归拼接父字段名和子字段名（用下划线分隔）
• 对于Array：先使用explode将数组拆分为多行，再递归处理数组内部的Struct元素（如果有的话）

改进后的通用扁平化函数 #

from pyspark.sql.functions import explode, col
from pyspark.sql.types import StructType, ArrayType

def get_flat_columns(schema, parent_path=""):
    """递归生成所有扁平化的列表达式"""
    columns = []
    for field in schema.fields:
        # 拼接当前字段的完整路径（父路径+当前字段名）
        current_path = f"{parent_path}_{field.name}" if parent_path else field.name
        
        if isinstance(field.dataType, StructType):
            # 递归处理Struct类型的子字段
            columns.extend(get_flat_columns(field.dataType, current_path))
        elif isinstance(field.dataType, ArrayType):
            # 处理Array类型：先explode拆分数组为多行
            exploded_col = explode(col(current_path)).alias(f"{current_path}_elem")
            
            elem_type = field.dataType.elementType
            if isinstance(elem_type, StructType):
                # 递归处理数组内的Struct元素，用数组的完整路径作为前缀
                struct_cols = get_flat_columns(elem_type, current_path)
                # 将递归生成的列引用替换为explode后的元素字段
                struct_cols = [
                    col(f"{current_path}_elem.{col_name.split('_')[-1]}").alias(col_name)
                    for col_name in struct_cols
                ]
                columns.append(exploded_col)
                columns.extend(struct_cols)
            else:
                # 数组元素为普通类型，直接用数组路径作为别名
                columns.append(exploded_col.alias(current_path))
        else:
            # 普通数据类型，直接返回字段引用和完整路径别名
            columns.append(col(current_path).alias(current_path))
    return columns

def flatten_df(df):
    """将传入的DataFrame扁平化，支持任意嵌套的Struct和Array"""
    flat_columns = get_flat_columns(df.schema)
    return df.select(flat_columns)

使用方法 #

直接调用flatten_df即可，不需要手动指定层数（原代码需要传入layers参数）：

my_flattened_df = flatten_df(jsonDF)
my_flattened_df.printSchema()

输出效果 #

运行后你会得到完全符合预期的Schema：

root
 |-- Data_Record_FName: string (nullable = true)
 |-- Data_Record_LName: long (nullable = true)
 |-- Data_Record_Address_Applicant_Id: long (nullable = true)
 |-- Data_Record_Address_Applicant_Type: string (nullable = true)
 |-- Data_Record_Address_Applicant_Option: long (nullable = true)
 |-- Data_Record_Address_Location: string (nullable = true)
 |-- Data_Record_Address_Town: long (nullable = true)
 |-- Data_Record_IsActive: boolean (nullable = true)
 |-- Id: string (nullable = true)

注意事项 #

1. 行数变化：使用explode会将数组的每个元素拆分为单独一行，如果原始数据中有多个数组字段，多次explode会产生笛卡尔积。如果需要保留原始行的关联，可以先给DataFrame添加唯一标识列（比如monotonically_increasing_id()），再执行扁平化。
2. 通用性：这个函数可以处理任意深度的嵌套Struct和Array，不需要针对特定Schema修改代码，完全适配你的通用需求。
3. 空值处理：explode会过滤掉数组为空或null的行，如果需要保留这些行，可以使用explode_outer替代explode（只需将函数中的explode替换为explode_outer即可）。

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