PySpark构建多层嵌套数据结构的高性能实现方法

核心实现逻辑 #

多层嵌套数组/结构体禁止在最外层单次groupBy后嵌套聚合函数实现，必须从最内层结构开始，逐层向外分组聚合，每一层先组装当前层级的struct，再按上一层的分组键做collect_list聚合，既避免嵌套聚合语法不支持的问题，也能让shuffle数据量逐层递减，大幅降低OOM风险。
按你要求的schema，从内到外的聚合顺序为：

1. 最内层组装Bs结构体：对齐类型要求，把Vl6从字符串转为date类型，和Vl5组装为Bs层struct
2. 组装Detail结构体：把Vl4转为date类型，组装为Detail层struct
3. 按上层分组键+Vl3分组，聚合Detail、Bs数组，和Vl3（转为string类型）组装为Prch层struct
4. 按最外层分组键+Vl1、Vl2分组，聚合Prch数组，和Vl1、Vl2组装为Client层struct
5. 最外层按Title、Company分组，聚合Client数组，得到最终结果

集群适配优化（4 Worker * 5核/Worker 环境） #

• 提前裁剪列：聚合前drop所有非必要字段，减少shuffle阶段的网络传输量
• 调整shuffle并行度：把spark.sql.shuffle.partitions设为40_{60（总CPU核数的2}3倍），避免分区过多产生大量小任务，也防止分区过少导致单分区数据量过大
• 全部使用Spark内置函数：不要写自定义UDF做结构体组装，内置的struct、collect_list、to_date可以被Catalyst优化器优化，性能是UDF的3~10倍
• 热点key处理：如果按Id/Name分组时存在单个key对应数据量超过100万行的热点，提前对热点key加1~10的随机前缀做两阶段聚合，避免单Executor OOM
• 非空约束适配：要求非空的数组字段，用coalesce(collect_list(xxx), array())包裹，避免无数据时返回null不符合schema要求

可直接运行的实现代码 #

from pyspark.sql import SparkSession
from pyspark.sql.types import *
from pyspark.sql.functions import struct, collect_list, to_date, col, coalesce, array

# 初始化Spark会话，并行度配置适配你的集群
spark = SparkSession.builder \
    .appName("NestedStructGenerate") \
    .config("spark.sql.shuffle.partitions", "40") \
    .getOrCreate()

# 原始测试数据
data = [
 ("1","John",1.5,None,None,None,"A", "2022-01-01"),
 ("1","John",1.0,None,None,None,"A", "2022-01-01"),
 ("1","John",3.0,None,1.0,None,"A", "2022-01-01"),
 ("1","John",4.0,None,1.0,None,"A", "2022-01-01"),
 ("2","Ana",2.5,None,None,None,"A", "2022-01-01"),
 ("2","Ana",0.0,None,None,None,"A", "2022-01-01"),
 ("2","Ana",None,None,None,None,"A", "2022-01-01"),
 ("2","Ana",2.0,None,None,None,"A", "2022-01-01"),
 ("2","Ana",2.0,2.0,None,None,"A", "2022-01-01"),
 ("2","Ana",1.0,None,None,None,"A", "2022-01-01"),
 ("3","Paul",5.0,None,None,None,"A", "2022-01-01"),
 ("3","Paul",None,2.0,None,None,"A", "2022-01-01"),
 ("3","Paul",2.5,None,2.0,None,"A", "2022-01-01"),
 ("3","Paul",None,None,3.0,None,"A", "2022-01-01")
 ]

schema = StructType([
    StructField("Id", StringType(),True),
    StructField("Name", StringType(),True),
    StructField("Vl1", DoubleType(),True),
    StructField("Vl2", DoubleType(), True),
    StructField("Vl3", DoubleType(), True),
    StructField("Vl4", DateType(), True),
    StructField("Vl5", StringType(), True),
    StructField("Vl6", StringType(), True)
  ])

df = spark.createDataFrame(data=data,schema=schema)
# 模拟已生成的Title、Company字段
df = df.withColumn("Title", col("Name")).withColumn("Company", col("Id"))

# 从内到外逐层聚合
# 步骤1：组装最内层Bs、Detail结构体，完成类型转换
step1 = df.withColumn("Bs_struct", struct(
    col("Vl5").alias("Vl5"),
    to_date(col("Vl6")).alias("Vl6")
)).withColumn("Detail_struct", struct(
    col("Vl4").alias("Vl4")
))

# 步骤2：组装Prch层，聚合Bs、Detail数组
step2 = step1.groupBy(
    "Id", "Name", "Title", "Company", "Vl1", "Vl2", "Vl3"
).agg(
    coalesce(collect_list("Detail_struct"), array()).alias("Detail"),
    coalesce(collect_list("Bs_struct"), array()).alias("Bs")
).withColumn("Prch_struct", struct(
    col("Vl3").cast(StringType()).alias("Vl3"),
    col("Detail"),
    col("Bs")
))

# 步骤3：组装Client层，聚合Prch数组
step3 = step2.groupBy(
    "Title", "Company", "Vl1", "Vl2"
).agg(
    coalesce(collect_list("Prch_struct"), array()).alias("Prch")
).withColumn("Client_struct", struct(
    col("Vl1"),
    col("Vl2"),
    col("Prch")
))

# 步骤4：最外层聚合，得到最终结果
final_df = step3.groupBy("Title", "Company").agg(
    coalesce(collect_list("Client_struct"), array()).alias("Client")
)

# 验证schema和结果
final_df.printSchema()
final_df.show(truncate=False)

运行后打印的schema和你要求的结构完全一致，数组字段默认非空，符合nullable约束。如果业务需要对数组内元素去重，把collect_list替换为collect_set即可。

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