PySpark实现嵌套XML部分扁平化：_Name值转列名的方法

当前使用PySpark处理嵌套XML文件时，已实现的flatten_explode_df工具函数仅支持全量扁平化所有嵌套结构，或完全不做扁平化处理，无法满足定制化部分扁平化的需求。

实现嵌套XML的部分扁平化，特殊处理规则为：针对包含*_Name*（列名标识）和*_VALUE*（对应取值）的结构体数组，需将*_Name的字段值作为最终输出的列名，对应的_VALUE*作为该列的取值，而非按原有逻辑展开嵌套结构。

基础调用代码 #

XML读取与全量扁平化原有调用逻辑：

df = spark.read.format("xml").options(rowTag="MAT").load("file.xml")
df_flat = flatten_explode_df(nested_df=df)

原有扁平化工具函数 #

import pyspark.sql.functions as F

def type_cols(df_dtypes, filter_type):
    cols = []
    for col_name, col_type in df_dtypes:
        if col_type.startswith(filter_type):
            cols.append(col_name)
    return cols

def flatten_df(nested_df, sep="_"):
    nested_cols = type_cols(nested_df.dtypes, "struct")
    flatten_cols = [fc for fc, _ in nested_df.dtypes if fc not in nested_cols]
    for nc in nested_cols:
        for cc in nested_df.select(f"{nc}.*").columns:
            if sep is None:
                flatten_cols.append(F.col(f"{nc}.{cc}").alias(f"{cc}"))
            else:
                flatten_cols.append(F.col(f"{nc}.{cc}").alias(f"{nc}{sep}{cc}"))
    return nested_df.select(flatten_cols)

def explode_df(nested_df):
    nested_cols = type_cols(nested_df.dtypes, "array")
    exploded_df = nested_df
    for nc in nested_cols:
        exploded_df = exploded_df.withColumn(nc, F.explode(F.col(nc)))
    return exploded_df

def flatten_explode_df(nested_df):
    df = nested_df
    struct_cols = type_cols(nested_df.dtypes, "struct")
    array_cols = type_cols(nested_df.dtypes, "array")
    if struct_cols:
        df = flatten_df(df)
        return flatten_explode_df(df)
    if array_cols:
        df = explode_df(df)
        return flatten_explode_df(df)
    return df

在原有递归扁平化逻辑前，新增特殊键值对数组的识别与转列处理，核心逻辑如下：

• 遍历所有数组类型字段，识别数组元素为结构体、且同时包含_Name和_VALUE子字段的特殊数组
• 对匹配到的特殊数组，先转换为Map结构，再将Map中的所有键展开为独立列，对应值取_VALUE字段内容
• 特殊数组处理完成后，剩余普通嵌套结构沿用原有递归展开、扁平化逻辑处理

完整实现代码：

from pyspark.sql.types import StructType, ArrayType

def is_kv_array(df, col_name):
    """校验字段是否为符合_Name+_VALUE结构的键值对结构体数组"""
    col_dtype = dict(df.dtypes)[col_name]
    if not col_dtype.startswith("array<struct<"):
        return False
    inner_fields = [f.name for f in df.schema[col_name].dataType.elementType.fields]
    return "_Name" in inner_fields and "_VALUE" in inner_fields

def process_kv_arrays(df):
    """将所有符合规则的键值对数组转换为独立列"""
    array_cols = type_cols(df.dtypes, "array")
    kv_cols = []
    # 先将所有KV数组转为Map结构
    for col_name in array_cols:
        if is_kv_array(df, col_name):
            df = df.withColumn(col_name, F.map_from_entries(F.col(col_name)))
            kv_cols.append(col_name)
    # 展开Map为独立列
    for kv_col in kv_cols:
        # 动态获取所有存在的键，无需提前硬编码字段
        all_keys = df.select(F.explode(F.col(kv_col))).select("key").distinct().rdd.flatMap(lambda x: x).collect()
        for key in all_keys:
            df = df.withColumn(key, F.col(f"{kv_col}.{key}"))
        df = df.drop(kv_col)
    return df

def custom_flatten_explode_df(nested_df):
    df = process_kv_arrays(nested_df)
    struct_cols = type_cols(df.dtypes, "struct")
    array_cols = type_cols(df.dtypes, "array")
    if struct_cols:
        df = flatten_df(df)
        return custom_flatten_explode_df(df)
    if array_cols:
        df = explode_df(df)
        return custom_flatten_explode_df(df)
    return df

替换原有调用方法即可使用：

df = spark.read.format("xml").options(rowTag="MAT").load("file.xml")
df_flat = custom_flatten_explode_df(df)

提示：如果需要固定输出列顺序，可在最终返回结果前通过select方法指定列的展示顺序；如果KV数组中存在空值，可根据业务需要增加fillna逻辑补全默认值。

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