PySpark计算数组列各子项全局均值并将胜出值写入全表新列

核心思路：全程不调用collect()拉取全量明细数据到driver，利用Spark对小结果集的自动广播优化，匹配原业务逻辑的同时完全符合PySpark开发规范，大表场景下性能较原实现提升数倍至数十倍。

实现步骤 #

1. 分布式计算全局聚合值
   不需要拉取原始明细数据到driver，直接在executor端完成所有均值计算，全程只传输最终的单行聚合结果：

   from pyspark.sql import functions as F
   # 直接取Col2数组指定位置的元素计算全局均值，不需要explode炸行
   global_agg = df.agg(
       F.avg(F.col("Col2")[0]).alias("avg_first"),
       F.avg(F.col("Col2")[1]).alias("avg_second")
   )
   # 取两个均值的最大值，生成目标Col4列的单行结果
   col4_df = global_agg.select(
       F.greatest(F.col("avg_first"), F.col("avg_second")).alias("Col4")
   )
   

   针对示例数据，这步计算得到的两个均值分别为1.33和2.67，取最大值后Col4值为2.67，和预期输出完全一致。
2. 关联原表完成固定值赋值
   因为col4_df是仅1行的极小数据集，Spark会自动触发广播嵌套循环连接，不会产生全量shuffle，性能和手动调用lit()赋值几乎无差异：

   # 单行结果表和原表做cross join，自动广播小表
   final_df = df.crossJoin(col4_df)
   # 也可以手动标记广播，效果完全一致
   # final_df = df.crossJoin(F.broadcast(col4_df))
   

性能对比说明 #

• 原collect实现的问题：需要把全量Col2的数组明细从所有executor拉取到driver，网络传输量大，driver内存压力高，大表场景下极易出现任务卡顿、OOM的问题
• 优化后实现的优势：所有均值计算在executor端分布式完成，仅传输最终的单行聚合结果，网络传输量可忽略不计；关联阶段无shuffle开销，完全适配超大规模数据集的计算场景。

注意事项 #

• 不需要额外用explode拆分数组、也不需要用全表窗口函数计算赋值，这套逻辑的执行计划非常精简，没有多余计算环节
• 该逻辑和原实现的计算结果完全一致，没有业务偏差

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