如何用Spark DataFrame实现复杂嵌套JSON的扁平化处理

我现有一段可正常扁平化JSON数据的Python代码，但因数据结构复杂（输入为字典）运行速度较慢，希望获取一段可在Spark DataFrame上运行的同类代码，能够处理任意复杂结构并将其扁平化至叶子节点。

原Python字典扁平化代码 #

# 扁平化代码
from functools import (partial,
                       singledispatch)
from itertools import chain
from typing import (Dict,
                    List,
                    TypeVar)

Serializable = TypeVar('Serializable', None, int, bool, float, str,
                       dict, list, tuple)
Array = List[Serializable]
Object = Dict[str, Serializable]


def flatten(object_: Object,
            *,
            path_separator: str = '.') -> Array[Object]:
    """
    将给定JSON对象扁平化为不含嵌套值的对象列表。

    >>> flatten({'a': 1})
    [{'a': 1}]
    >>> flatten({'a': [1, 2]})
    [{'a': 1}, {'a': 2}]
    >>> flatten({'a': {'b': None}})
    [{'a.b': None}]
    >>> flatten({'a': [1, 2], 'b': []})
    [{'a': 1}, {'a': 2}]
    """
    keys = set(object_)
    result = [dict(object_)]
    while keys:
        key = keys.pop()
        new_result = []
        for index, record in enumerate(result):
            try:
                value = record[key]
            except KeyError:
                new_result.append(record)
            else:
                if isinstance(value, dict):
                    del record[key]
                    new_value = flatten_nested_objects(
                            value,
                            prefix=key + path_separator,
                            path_separator=path_separator
                    )
                    keys.update(new_value.keys())
                    new_result.append({**new_value, **record})
                elif isinstance(value, list):
                    del record[key]
                    new_records = [
                        flatten_nested_objects(sub_value,
                                               prefix=key + path_separator,
                                               path_separator=path_separator)
                        for sub_value in value
                    ]
                    keys.update(chain.from_iterable(map(dict.keys,
                                                        new_records)))
                    if new_records:
                        new_result.extend({**new_record, **record}
                                          for new_record in new_records)
                    else:
                        new_result.append(record)
                else:
                    new_result.append(record)
        result = new_result
    return result


@singledispatch
def flatten_nested_objects(object_: Serializable,
                           *,
                           prefix: str = '',
                           path_separator: str) -> Object:
    return {prefix[:-len(path_separator)]: object_}


@flatten_nested_objects.register(dict)
def _(object_: Object,
      *,
      prefix: str = '',
      path_separator: str) -> Object:
    result = dict(object_)
    for key in list(result):
        result.update(flatten_nested_objects(result.pop(key),
                                             prefix=(prefix + key
                                                     + path_separator),
                                             path_separator=path_separator))
    return result


@flatten_nested_objects.register(list)
def _(object_: Array,
      *,
      prefix: str = '',
      path_separator: str) -> Object:
    return {prefix[:-len(path_separator)]: list(map(partial(
            flatten_nested_objects,
            path_separator=path_separator),
            object_))}

Spark DataFrame 扁平化实现方案 #

以下是基于Spark API实现的通用扁平化函数，支持处理任意嵌套的结构体、数组，最终将所有叶子节点转为顶级列，列名用指定分隔符拼接路径：

from pyspark.sql import DataFrame
from pyspark.sql.types import StructType, ArrayType, StructField
from pyspark.sql.functions import col, explode, alias

def flatten_df(df: DataFrame, path_separator: str = '.') -> DataFrame:
    """
    扁平化Spark DataFrame中的所有嵌套结构（结构体、数组），将叶子节点转为顶级列，列名用路径分隔符拼接。
    
    参数:
        df: 需要扁平化的Spark DataFrame
        path_separator: 列名路径分隔符，默认是'.'
    
    返回:
        扁平化后的DataFrame
    """
    def get_flattened_columns(schema, prefix: str = ""):
        flattened_cols = []
        for field in schema.fields:
            field_name = f"{prefix}{path_separator}{field.name}" if prefix else field.name
            # 处理结构体类型
            if isinstance(field.dataType, StructType):
                flattened_cols.extend(get_flattened_columns(field.dataType, field_name))
            # 处理数组类型：先explode，再递归处理数组内的结构
            elif isinstance(field.dataType, ArrayType):
                # 先对数组列进行explode，重命名为原字段名
                exploded_col = explode(col(field_name)).alias(field_name)
                # 递归处理数组元素的结构
                array_schema = field.dataType.elementType
                if isinstance(array_schema, StructType):
                    nested_cols = get_flattened_columns(array_schema, field_name)
                    flattened_cols.append(exploded_col)
                    flattened_cols.extend(nested_cols)
                else:
                    flattened_cols.append(exploded_col)
            # 叶子节点直接添加
            else:
                flattened_cols.append(col(field_name).alias(field_name))
        return flattened_cols
    
    # 获取所有扁平化后的列表达式
    flattened_columns = get_flattened_columns(df.schema)
    # 选择这些列生成新的DataFrame
    return df.select(flattened_columns)

使用示例 #

假设你有一个包含嵌套结构的DataFrame：

# 创建示例嵌套DataFrame
data = [
    {"id": 1, "info": {"name": "Alice", "age": 30}, "scores": [{"math": 90, "english": 85}, {"math": 95, "english": 88}]},
    {"id": 2, "info": {"name": "Bob", "age": 25}, "scores": [{"math": 80, "english": 75}]}
]
df = spark.createDataFrame(data)

# 扁平化处理
flattened_df = flatten_df(df)
flattened_df.show(truncate=False)

执行后会得到扁平化的结果，所有叶子节点都成为顶级列，比如info.name、scores.math等，数组会被展开为多行。

代码说明 #

1. 递归遍历Schema：通过get_flattened_columns函数递归遍历DataFrame的Schema，识别结构体、数组和叶子节点。
2. 处理结构体：对于结构体类型，拼接父路径和字段名，继续递归处理内部字段。
3. 处理数组：使用explode函数将数组展开为多行，然后递归处理数组元素的结构（如果元素是结构体）。
4. 列名拼接：用指定的分隔符拼接嵌套路径，确保列名清晰反映原始结构。

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