如何在PySpark Schema中遍历无字段名的嵌套数组并修改类型？

问题背景 #

当前PySpark DataFrame的Schema如下：

root
 |-- C_0_0: double (nullable = true)
 |-- C_0_1: array (nullable = true)
 |    |-- element: array (containsNull = true)
 |    |    |-- element: double (containsNull = true)
 |-- C_0_2: array (nullable = true)
 |    |-- element: array (containsNull = true)
 |    |    |-- element: double (containsNull = true)

需要将C_0_1字段下第二层数组的元素类型从double改为decimal(8,6)，最终Schema如下：

root
 |-- C_0_0: double (nullable = true)
 |-- C_0_1: array (nullable = true)
 |    |-- element: array (containsNull = true)
 |    |    |-- element: decimal(8,6) (containsNull = true)
 |-- C_0_2: array (nullable = true)
 |    |-- element: array (containsNull = true)
 |    |    |-- element: double (containsNull = true)

解决方案 #

针对嵌套数组无字段名的情况，有两种常用处理方式：

方式一：直接构造目标类型并转换（适合层级明确的场景）

如果已知嵌套数组的层级结构，可以直接构造目标ArrayType，再通过cast方法修改字段类型：

from pyspark.sql import SparkSession
from pyspark.sql.types import ArrayType, DecimalType
from pyspark.sql.functions import col

# 初始化SparkSession并加载得到DataFrame df
spark = SparkSession.builder.appName("NestedArrayTypeUpdate").getOrCreate()

# 构造C_0_1的目标类型：数组嵌套数组，内层元素为decimal(8,6)
new_c01_type = ArrayType(
    ArrayType(DecimalType(8, 6), containsNull=True),
    containsNull=True
)

# 修改C_0_1字段类型，保留其他字段不变
df_updated = df.withColumn("C_0_1", col("C_0_1").cast(new_c01_type))

# 验证修改后的Schema
df_updated.printSchema()

方式二：递归遍历修改Schema（适合复杂嵌套或多路径修改场景）

如果嵌套层级复杂，或者需要批量修改多个路径的元素类型，可以写一个递归函数遍历Schema结构，替换指定路径的类型：

from pyspark.sql.types import StructType, StructField, ArrayType, DecimalType

def update_nested_type(schema, target_path, new_type):
    """
    递归修改嵌套Schema中的指定路径类型
    :param schema: 原始Schema对象
    :param target_path: 目标路径列表，如["C_0_1", 0, 0]表示C_0_1下的两层数组元素
    :param new_type: 替换后的目标类型
    :return: 修改后的Schema
    """
    if not target_path:
        return new_type
    
    current_node = target_path[0]
    remaining_path = target_path[1:]
    
    if isinstance(schema, StructType):
        updated_fields = []
        for field in schema.fields:
            if field.name == current_node:
                updated_field = StructField(
                    field.name,
                    update_nested_type(field.dataType, remaining_path, new_type),
                    field.nullable
                )
                updated_fields.append(updated_field)
            else:
                updated_fields.append(field)
        return StructType(updated_fields)
    elif isinstance(schema, ArrayType):
        # 数组元素无字段名，直接递归处理elementType
        return ArrayType(
            update_nested_type(schema.elementType, remaining_path, new_type),
            schema.containsNull
        )
    else:
        # 到达目标层级，替换类型
        return new_type

# 使用递归函数修改Schema
original_schema = df.schema
# 目标路径：C_0_1 -> 第一层数组 -> 第二层数组元素
new_schema = update_nested_type(original_schema, ["C_0_1", 0, 0], DecimalType(8,6))

# 应用新Schema到DataFrame
df_updated = spark.createDataFrame(df.rdd, new_schema)

# 验证结果
df_updated.printSchema()

说明 #

• 方式一操作简单，适合已知固定嵌套层级的场景，性能更优；
• 方式二灵活性强，可应对复杂嵌套结构或需要批量修改多个路径的需求。

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