如何为含嵌套结构体与数组的PySpark DataFrame整数字段应用函数

要实现对PySpark DataFrame中所有Integer类型字段（包括嵌套结构体、数组内的整数/结构体）执行自定义操作（如加1），可以通过递归遍历Schema+PySpark内置函数实现，完全不需要UDF，性能更优。

核心思路 #

递归遍历DataFrame的Schema结构，针对不同数据类型做对应处理：

• 对于IntegerType字段：直接执行加1操作
• 对于ArrayType：使用transform函数遍历数组元素，根据元素类型递归处理
• 对于StructType：遍历结构体的每个子字段，递归处理后重新组合成结构体
• 其他数据类型：保持原字段不变

完整实现代码 #

from pyspark.sql import SparkSession
from pyspark.sql.types import (
    StructType, StructField, ArrayType, IntegerType
)
from pyspark.sql.functions import col, struct, transform

# 初始化SparkSession
spark = SparkSession.builder.appName("NestedIntegerTransform").getOrCreate()

# 构造示例输入DataFrame
data = {
    'top_level_field':{
        'array_field': [
            {
                'inner_array_field': [1, 2, 3],
                'inner_array_of_structs': [
                    {
                        'field_1': 2,
                        'field_2': 3
                    }
                ]
            },
            {
                'inner_array_field': [1, 2, 3],
                'inner_array_of_structs': [
                    {
                        'field_1': 2,
                        'field_2': 3,
                        'field_3': [1, 2, 3]
                    }
                ]
            }
        ]
    }
}
df = spark.createDataFrame([data])

def transform_nested_fields(field_ref, field_type):
    """递归转换嵌套字段中的Integer类型"""
    # 处理Integer类型字段
    if isinstance(field_type, IntegerType):
        return field_ref + 1
    # 处理Array类型
    elif isinstance(field_type, ArrayType):
        element_type = field_type.elementType
        # 数组元素是嵌套结构（结构体或数组）
        if isinstance(element_type, (StructType, ArrayType)):
            return transform(field_ref, lambda x: transform_nested_fields(x, element_type))
        # 数组元素是Integer类型
        elif isinstance(element_type, IntegerType):
            return transform(field_ref, lambda x: x + 1)
        # 其他数组元素类型，返回原字段
        else:
            return field_ref
    # 处理Struct类型
    elif isinstance(field_type, StructType):
        struct_fields = []
        for sub_field in field_type.fields:
            struct_fields.append(
                transform_nested_fields(field_ref[sub_field.name], sub_field.dataType).alias(sub_field.name)
            )
        return struct(*struct_fields)
    # 其他数据类型，返回原字段
    else:
        return field_ref

# 生成转换后的DataFrame
transformed_df = df.select(transform_nested_fields(col(""), df.schema).alias("data"))

# 查看结果（可转为字典验证结构）
transformed_df.show(truncate=False)

代码说明 #

1. 递归函数transform_nested_fields：
   • 接收字段引用（或结构体/数组元素的引用）和对应的字段类型
   • 针对每种数据类型分支处理，确保所有层级的Integer字段都被遍历到
2. 数组处理：使用PySpark内置的transform函数遍历数组，避免了UDF的性能开销
3. 结构体处理：遍历结构体的每个子字段，递归转换后用struct函数重新组合成结构体

执行上述代码后，输出的DataFrame将完全符合示例要求：所有Integer类型字段（包括数组内的整数、结构体嵌套的整数、数组内结构体的整数）都被加1，其他结构保持不变。

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