Spark 3.2.1嵌套结构体转结构体数组并提取结构体键

针对你需要将结构体a下的多个同结构子结构体（如b、c）批量转换为包含原结构体名称的结构体数组需求，可以通过Spark内置函数实现批量处理，无需硬编码子字段名。

方法1：动态生成结构体表达式 #

这种方法通过读取Schema动态获取所有子字段，批量生成数组元素表达式，灵活性强，适合子结构体数量不确定的场景。

步骤与代码示例 #

1. 导入Spark函数库

from pyspark.sql import functions as F

2. 读取JSON数据并获取目标结构体的子字段信息

# 读取JSON数据（替换为你的数据路径）
df = spark.read.json("path/to/your/data.json")

# 获取结构体a下的所有子字段（StructField对象列表）
a_sub_fields = df.schema["a"].dataType.fields

3. 构造数组元素表达式并转换字段

# 遍历每个子字段，生成包含struct_key和原结构体字段的结构体表达式
array_items = []
for sub_field in a_sub_fields:
    # 提取当前子结构体的所有字段，避免硬编码
    sub_struct_fields = [
        F.col(f"a.{sub_field.name}.{field.name}").alias(field.name) 
        for field in sub_field.dataType.fields
    ]
    # 组合成新的结构体：struct_key存储原字段名，其余为原结构体字段
    item_expr = F.struct(
        F.lit(sub_field.name).alias("struct_key"),
        *sub_struct_fields
    )
    array_items.append(item_expr)

# 将所有结构体表达式组合成数组，替换原a字段或新增字段
df_transformed = df.withColumn("a_array", F.array(*array_items)).drop("a")

方法2：利用Map转换简化代码 #

通过struct_to_map将结构体转为键值对Map，再通过map_entries和transform批量处理，代码更简洁。

代码示例 #

from pyspark.sql import functions as F

df = spark.read.json("path/to/your/data.json")
a_sub_fields = df.schema["a"].dataType.fields

# 假设所有子结构体结构一致，取第一个子结构体的字段列表
sub_struct_field_names = [field.name for field in a_sub_fields[0].dataType.fields]

df_transformed = df.withColumn(
    "a_array",
    F.transform(
        # 将结构体a转为Map，再转为键值对数组
        F.map_entries(F.struct_to_map(F.col("a"))),
        # 遍历每个键值对，生成目标结构体
        lambda entry: F.struct(
            entry["key"].alias("struct_key"),
            *[entry["value"][field].alias(field) for field in sub_struct_field_names]
        )
    )
).drop("a")

效果说明 #

假设原数据结构为：

{"a": {"b": {"x": 1, "y": "foo"}, "c": {"x": 2, "y": "bar"}, "d": {"x": 3, "y": "baz"}}}

转换后a_array字段的输出为：

[
  {"struct_key": "b", "x": 1, "y": "foo"},
  {"struct_key": "c", "x": 2, "y": "bar"},
  {"struct_key": "d", "x": 3, "y": "baz"}
]

两种方法均支持批量处理任意数量的同结构子结构体，无需修改代码适配子字段数量。

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