如何递归为PySpark DataFrame嵌套数组添加哈希列？

问题背景 #

需要处理包含多层嵌套结构体数组的PySpark DataFrame，目标是给每个嵌套数组的元素添加对应层级的哈希列，同时给整个记录添加顶层哈希列。由于嵌套层级不确定，必须通过递归实现。

输入Schema示例 #

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

schema = StructType([
    StructField("name", StringType()),
    StructField("experience", ArrayType(StructType([
        StructField("role", StringType()),
        StructField("duration", StringType()),
        StructField("company", ArrayType(StructType([
            StructField("company_name", StringType()),
            StructField("location", StringType())
        ])))
    ])))
])

期望输出Schema #

hashed_schema = StructType([
    StructField("name", StringType()),
    StructField("experience", ArrayType(StructType([
        StructField("role", StringType()),
        StructField("duration", StringType()),
        StructField("experience_hash", StringType()),  # 为experience数组元素添加的哈希列
        StructField("company", ArrayType(StructType([
            StructField("company_name", StringType()),
            StructField("location", StringType()),
            StructField("company_hash", StringType())  # 为company数组元素添加的哈希列
        ])))
    ]))),
    StructField("employee_hash", StringType()),  # 整条记录的顶层哈希列
])

现有尝试及问题 #

我编写了递归代码，但仅对第一层嵌套数组生效，还出现了递归过深的错误，代码如下：

from pyspark.sql.functions import md5, concat_ws, lit, col, transform

def hash_for_level(level_path):
    return md5(concat_ws("_", *[lit(elem) for elem in level_path]))

def add_hash_columns(df, level_path, current_struct, root_hash_col=None):
    # 如果是根层级，创建根哈希列
    if not level_path and root_hash_col is None:
        root_hash_col = 'employee_hash'
        df = df.withColumn(root_hash_col, hash_for_level(['employee']))
    
    # 遍历当前结构添加哈希列
    for field in current_struct.fields:
        new_level_path = level_path + [field.name]
        # 如果字段是结构体数组，给数组每个元素添加哈希列
        if isinstance(field.dataType, ArrayType):
            nested_struct = field.dataType.elementType
            hash_expr = transform(
                col('.'.join(level_path + [field.name])),
                lambda x: x.withField(new_level_path[-1] + '_hash', hash_for_level(new_level_path))
                    .withField(root_hash_col, col(root_hash_col))  # 包含根哈希列
            )
            # 给数组元素添加哈希列
            df = df.withColumn('.'.join(level_path + [field.name]), hash_expr)
            # 递归处理嵌套数组
            df = add_hash_columns(df, new_level_path, nested_struct, root_hash_col)
            
    # 给当前层级添加哈希列
    if level_path:
        hash_col_name = '_'.join(level_path) + '_hash'
        df = df.withColumn(hash_col_name, hash_for_level(level_path))
        if root_hash_col:
            # 确保每个结构体层级都包含根哈希列
            df = df.withColumn(root_hash_col, col(root_hash_col))
            
    return df

# 测试代码
df = spark.createDataFrame([], schema)
df = add_hash_columns(df, [], df.schema)
df

解决方案 #

原代码存在两个核心问题：一是哈希计算基于层级路径而非实际字段值，二是递归处理数组时的逻辑混乱导致重复操作。以下是修正后的实现：

from pyspark.sql import SparkSession
from pyspark.sql.types import StructType, StructField, StringType, ArrayType
from pyspark.sql.functions import md5, concat_ws, col, transform, struct

def get_struct_fields_hash_expr(struct_col):
    """生成结构体所有字段值拼接后的MD5哈希表达式"""
    fields = struct_col.fields
    if not fields:
        return md5(lit(""))
    # 拼接当前结构体所有字段的值，生成哈希
    concat_expr = concat_ws("_", *[struct_col[field.name] for field in fields])
    return md5(concat_expr)

def add_nested_hashes(struct_type, root_hash_col):
    """递归生成添加哈希列后的结构体类型及转换表达式"""
    new_fields = []
    transform_exprs = {}
    
    for field in struct_type.fields:
        field_name = field.name
        if isinstance(field.dataType, ArrayType) and isinstance(field.dataType.elementType, StructType):
            # 处理结构体数组：先递归处理内部嵌套结构
            nested_struct_type, nested_transform = add_nested_hashes(field.dataType.elementType, root_hash_col)
            # 定义数组元素的哈希列名称
            array_hash_col = f"{field_name}_hash"
            # 生成数组转换逻辑：递归处理内部后添加当前层级哈希
            array_transform = transform(
                col(field_name),
                lambda elem: nested_transform(elem).withField(array_hash_col, get_struct_fields_hash_expr(elem))
            )
            transform_exprs[field_name] = array_transform
            # 更新数组元素的结构体类型（添加哈希列）
            new_fields.append(StructField(field_name, ArrayType(nested_struct_type.add(array_hash_col, StringType()))))
        elif isinstance(field.dataType, StructType):
            # 处理嵌套结构体：递归生成转换逻辑
            nested_struct_type, nested_transform = add_nested_hashes(field.dataType, root_hash_col)
            transform_exprs[field_name] = nested_transform(col(field_name))
            new_fields.append(StructField(field_name, nested_struct_type))
        else:
            # 普通字段直接保留
            new_fields.append(field)
    
    # 生成当前结构体的转换函数
    def struct_transform(struct_col):
        transformed = struct_col
        for name, expr in transform_exprs.items():
            transformed = transformed.withField(name, expr)
        return transformed
    
    new_struct_type = StructType(new_fields)
    return new_struct_type, struct_transform

def add_all_hashes(df, root_hash_name="employee_hash"):
    # 递归处理根结构体的所有嵌套层级
    root_struct_type, root_transform = add_nested_hashes(df.schema, root_hash_name)
    # 转换整个DataFrame的结构
    transformed_df = df.select(root_transform(struct("*")).alias("data")).select("data.*")
    # 添加顶层全局哈希列（基于整条记录的所有字段值）
    root_hash_expr = get_struct_fields_hash_expr(struct("*"))
    transformed_df = transformed_df.withColumn(root_hash_name, root_hash_expr)
    return transformed_df

# 测试代码
spark = SparkSession.builder.appName("NestedHash").getOrCreate()

# 构造测试数据
test_data = [
    (
        "Alice",
        [
            {
                "role": "Engineer",
                "duration": "3y",
                "company": [
                    {"company_name": "Tech Corp", "location": "NY"},
                    {"company_name": "Startup Inc", "location": "CA"}
                ]
            },
            {
                "role": "Senior Engineer",
                "duration": "2y",
                "company": [
                    {"company_name": "Big Tech", "location": "WA"}
                ]
            }
        ]
    )
]

df = spark.createDataFrame(test_data, schema)
hashed_df = add_all_hashes(df)

# 查看结果Schema和数据
hashed_df.printSchema()
hashed_df.show(truncate=False)

代码说明 #

1. get_struct_fields_hash_expr：负责基于结构体的实际字段值生成哈希，确保哈希能反映数据内容而非路径标识
2. add_nested_hashes：递归遍历结构体类型，对每个数组字段生成转换逻辑，自动处理任意深度的嵌套数组，给每个数组元素添加对应层级的哈希列
3. add_all_hashes：入口函数，先转换整个DataFrame的嵌套结构，再添加基于整条记录的顶层哈希列

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