PySpark大表多字段聚合失败优化问询：适配百亿级数据场景

1. 修复代码中的逻辑错误 #

你的代码存在字段别名重复的严重问题：

sort_array(collected_set("colE")).alias("colE"),
sum("colE").alias("colE"),

第二个sum("colE")的别名会直接覆盖第一个聚合结果，导致前序的collected_set操作完全无效，还额外浪费了计算资源。必须修改别名区分结果，示例如下：

sort_array(collected_set("colE")).alias("colE_unique_list"),
sum("colE").alias("colE_total"),

2. 替换first为更高效的聚合函数 #

由于colL、colM、colN、colO同组内值完全一致，不需要使用first，可以替换为max()或min()：

• max()/min()在Spark中的聚合实现经过高度优化，无需依赖数据顺序，性能远优于first
• 因为同组值一致，max()和min()会返回相同的结果，完全满足业务需求

修改后代码片段：

max("colL").alias("colL"),
max("colM").alias("colM"),
max("colN").alias("colN"),
max("colO").alias("colO"),

3. 优化countDistinct性能 #

countDistinct是开销较高的聚合操作，可根据业务需求选择优化方式：

• 若允许近似值，替换为approx_count_distinct()，性能提升显著（误差约5%）
• 若需要精确值，改用size(collected_set("colK")).alias("colK_distinct")，两者结果一致，但collected_set的聚合逻辑在部分场景下更高效

示例代码：

# 近似去重计数
approx_count_distinct("colK").alias("colK_distinct"),
# 精确去重计数
size(collected_set("colK")).alias("colK_distinct"),

4. 解决数据倾斜（核心问题） #

FetchFailedException几乎都是由数据倾斜导致的：某个分组key对应的行数远超其他key，单个task处理超时或内存不足。

排查倾斜key #

先统计分组key的分布，找出异常key：

from pyspark.sql.functions import desc

df.groupBy("colA","colB","colC","colD")
  .count()
  .orderBy(desc("count"))
  .show(20, truncate=False)

解决倾斜方案 #

• 加盐（Salt）处理：对倾斜的key添加随机后缀，拆分小分组聚合后再合并结果。例如，若发现某个(colA,colB,colC,colD)组合有1000万行，可添加0-9的随机数：

from pyspark.sql.functions import rand, floor

# 对倾斜key加盐
salted_df = df.withColumn("salt", floor(rand() * 10))
# 按加盐后的key聚合
salted_agg = salted_df.groupBy("colA","colB","colC","colD", "salt").agg(
    count(*).alias("numRecords"),
    sort_array(collected_set("colE")).alias("colE_unique_list"),
    sum("colE").alias("colE_total"),
    sum("colF").alias("colF"),
    max("colL").alias("colL"),
    size(collected_set("colK")).alias("colK_distinct")
    # 其他聚合操作...
)
# 合并加盐后的结果
final_agg = salted_agg.groupBy("colA","colB","colC","colD").agg(
    sum("numRecords").alias("numRecords"),
    sort_array(collect_set("colE_unique_list")).alias("colE_unique_list"),
    sum("colE_total").alias("colE_total"),
    sum("colF").alias("colF"),
    max("colL").alias("colL"),
    size(collect_set("colK_distinct")).alias("colK_distinct")
    # 其他聚合操作...
)

• 调整Spark Shuffle参数：针对百亿级数据，需调整以下参数优化shuffle性能：

  # 增加shuffle分区数（默认200，建议设为1000-4000，根据集群规模调整）
  spark.conf.set("spark.sql.shuffle.partitions", "2000")
  # 增大executor内存（根据机器配置调整，例如16G）
  spark.conf.set("spark.executor.memory", "16g")
  # 增大shuffle文件缓冲区
  spark.conf.set("spark.shuffle.file.buffer", "64k")
  # 增大reducer端接收数据的缓冲区
  spark.conf.set("spark.reducer.maxSizeInFlight", "96m")
  # 开启堆外内存优化
  spark.conf.set("spark.memory.offHeap.enabled", "true")
  spark.conf.set("spark.memory.offHeap.size", "8g")
  

5. 提前过滤无效数据（若可行） #

如果业务允许，可在聚合前过滤掉明显不符合后续筛选条件的数据（比如某些字段的极端值、null值），减少聚合的数据量：

# 示例：过滤colE为null的行（若不影响聚合结果）
filtered_df = df.filter(df.colE.isNotNull())

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