如何动态转换含同名Struct类型的PySpark DataFrame Schema？

需求 #

将PySpark DataFrame动态转换为符合目标Schema的结构，要求：

• 支持同名Struct类型字段的合并（如输入的manager Struct与输出中包含更多子字段的manager Struct合并）
• 适配任意输入输出Schema（只要输出字段存在于输入DataFrame中）

示例代码与预期效果 #

from pyspark.sql.types import *
from pyspark.sql import functions as f
from pyspark.sql import Row

# 输入Schema
schema_in = StructType([
    StructField("name", StringType(), True),
    StructField("age", LongType(), True),
    StructField("manager", StructType([
                          StructField("manager_name", StringType(), True),
                          StructField("manager_age", LongType(), True)]), True),
    StructField("manager_city", StringType(), True),
    StructField("manager_dept", StringType(), True)])

# 目标Schema：manager新增detail子Struct，包含manager_city和manager_dept
schema_out = StructType([
    StructField("name", StringType(), True),
    StructField("age", LongType(), True),
    StructField("manager", StructType([
                          StructField("manager_name", StringType(), True),
                          StructField("manager_age", LongType(), True),
                          StructField("detail", StructType([
                                               StructField("manager_city", StringType(), True),
                                               StructField("manager_dept", StringType(), True)]), True)]), True)])

# 测试数据
data = [
    Row("Alex", 20, Row("Bob",30),"Madrid","Supply Chain"),
    Row("Cathy", 40, Row("Doge",40),"Barcelona","Human Resources")
]

df = spark.createDataFrame(data, schema=schema_in)
print("输入DataFrame结构：")
df.printSchema()
print("输入DataFrame数据：")
df.show(20, False)

# 手动转换的正确效果
df2 = df.withColumn('manager', f.col('manager')\
    .withField('detail', f.struct(*[f.col('manager_city').alias('manager_city'), f.col('manager_dept').alias('manager_dept')])))\
    .drop("manager_city")\
    .drop("manager_dept")

print("\n手动转换后的结构：")
df2.printSchema()
print("手动转换后的数据：")
df2.show(20, False)

运行后输入DataFrame的结构和数据符合预期，手动转换后的结果正确将manager_city、manager_dept合并到manager.detail中，同时保留原manager的字段。

原递归方法的问题 #

原递归实现错误地将输入中已有的manager Struct嵌套到了detail子Struct中，导致层级混乱。错误代码如下：

def arrange_fields(df_in, schema_out):
    def add_new_fields(schema, prefix=""):
        columnas = []
        for field in schema.fields:
            field_name = f"{prefix}.{field.name}" if prefix else field.name
            
            if isinstance(field.dataType, StructType):
                nested_columns = add_new_fields(field.dataType, field_name)
                if prefix:
                    existing_struct = f.col(prefix)
                    columnas.append(f.struct(*nested_columns, existing_struct).alias(field.name))
                else:
                    columnas.append(f.struct(*nested_columns).alias(field.name))
            else:
                if field.name in df_in.columns:
                    columnas.append(f.col(field.name).alias(field.name))
        return columnas

    columnas = add_new_fields(schema_out)
    return df_in.select(*columnas)

错误输出中，manager_name和manager_age被嵌套在detail内部，完全不符合目标Schema的层级要求。

修正后的递归实现 #

核心思路：对于每个Struct字段，先递归生成其所有子字段的表达式，再检查输入中是否已存在该Struct字段——如果存在，则将原有Struct的所有字段与新生成的子字段合并；如果不存在，则直接使用新生成的子字段构建Struct。

from pyspark.sql import functions as f
from pyspark.sql.types import StructType

def transform_to_target_schema(df_in, schema_out):
    def build_field_expr(field, parent_path=""):
        # 构建当前字段的完整路径
        current_path = f"{parent_path}.{field.name}" if parent_path else field.name
        
        if isinstance(field.dataType, StructType):
            # 递归构建当前Struct的所有子字段表达式
            nested_exprs = []
            for nested_field in field.dataType.fields:
                nested_exprs.append(build_field_expr(nested_field, current_path))
            
            # 检查输入中是否存在当前Struct字段
            if field.name in df_in.columns:
                # 获取原有Struct的所有字段
                original_struct = f.col(field.name)
                original_fields = [original_struct[child.name].alias(child.name) for child in df_in.schema[field.name].dataType.fields]
                # 合并原有字段和新生成的子字段（去重，优先保留目标Schema的字段）
                all_fields = {}
                for expr in original_fields + nested_exprs:
                    all_fields[expr._jc.alias()] = expr
                merged_exprs = list(all_fields.values())
                # 构建合并后的Struct
                return f.struct(*merged_exprs).alias(field.name)
            else:
                # 输入中无该Struct，直接用递归生成的子字段构建
                return f.struct(*nested_exprs).alias(field.name)
        else:
            # 普通字段：支持平级字段映射到嵌套路径（如manager_city -> manager.detail.manager_city）
            flat_field_name = current_path.replace(".", "_")
            if flat_field_name in df_in.columns:
                return f.col(flat_field_name).alias(field.name)
            elif current_path in df_in.columns:
                return f.col(current_path).alias(field.name)
            else:
                # 字段不存在时返回对应类型的null
                return f.lit(None).cast(field.dataType).alias(field.name)
    
    # 构建所有顶层字段的表达式
    top_level_exprs = [build_field_expr(field) for field in schema_out.fields]
    return df_in.select(*top_level_exprs)

测试修正后的方法 #

df_transformed = transform_to_target_schema(df, schema_out)
print("\n动态转换后的结构：")
df_transformed.printSchema()
print("动态转换后的数据：")
df_transformed.show(20, False)

关键改进点 #

1. Struct字段合并逻辑：自动合并输入中已存在的Struct字段与目标Schema的子字段，避免嵌套错误
2. 平级字段映射：将输入中的平级字段（如manager_city）自动映射到目标Schema的嵌套路径
3. 字段去重：合并时自动去重，保证字段唯一性
4. 空值兼容：对输入中不存在的字段返回对应类型的null，确保Schema完全匹配

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