groupByKey&mapGroups后repartition无法解决Dataset数据倾斜问题

1. 先定位倾斜根源：从groupByKey阶段入手 #

你的数据倾斜大概率在groupByKey阶段就已经产生了，mapGroups只是延续了这个倾斜状态。先排查核心key的分布：

• 执行以下代码查看id_field的数据分布，确认是否存在热点key（少数key对应远超平均的数据量）：

  ds.groupBy("id_field")
    .count()
    .orderBy(desc("count"))
    .show(20)
  

如果发现某个/某些id_field对应的数据量是其他key的几十甚至几百倍，那就是倾斜的核心原因。

2. 热点key拆分（加盐法） #

针对热点key，在groupByKey前添加随机后缀拆分，让热点数据分散到多个分区处理：

// 假设已通过排查得到热点key列表hotKeys
val saltedDs = ds.withColumn(
  "salted_id",
  when(col("id_field").isin(hotKeys: _*), concat(col("id_field"), lit("_"), floor(rand() * 10)))
    .otherwise(col("id_field"))
)

// 基于加盐后的key执行groupByKey和mapGroups
val processedDs = saltedDs.groupByKey(_.salted_id)
  .mapGroups((key, iter) => {
    // 这里保留原有的mapGroups逻辑，处理后记得把加盐key还原
    val originalId = key.replaceAll("_\\d+$", "")
    (originalId, /* 你的处理结果 */)
  })
  .toDF("id_field", "result")

这种方法能把单个热点key的计算压力分散到10个（可调整随机数范围）分区，从根源解决倾斜。

3. 调整repartition策略 #

如果必须在mapGroups后做分区，避免用本身就倾斜的id_field作为分区键：

• 改用随机分区打散数据：

  val evenlyPartitionedDs = mapGroupsResult.repartition(2001, rand())
  

• 检查shuffle配置：确保spark.sql.shuffle.partitions设置为2001（与你的分区数一致），同时可以调优spark.shuffle.file.buffer、spark.reducer.maxSizeInFlight等参数，保障shuffle过程中数据能均匀分配。

4. 检查mapGroups逻辑是否放大数据量 #

如果你的mapGroups逻辑存在数据膨胀（比如展开超大数组），会导致单个分区数据量暴增：

• 提前展开数组：用explode把数组字段拆分成单条数据后再做groupByKey，而不是在mapGroups内部处理大数组：

  val explodedDs = ds.withColumn("array_element", explode(col("your_array_field")))
  val groupedDs = explodedDs.groupByKey(_.id_field).mapGroups(...)
  

• 如果必须在mapGroups内处理大数组，可在内部做分批处理，避免单条数据过度膨胀。

5. 强制触发shuffle #

如果Spark优化器合并了操作导致repartition未生效，可通过缓存+释放的方式强制触发shuffle：

val tempDs = mapGroupsResult.cache()
tempDs.count() // 触发缓存写入
tempDs.unpersist() // 释放缓存，断开与上游操作的依赖
val finalDs = tempDs.repartition(2001, rand())

也可以关闭自适应执行：spark.sql.adaptive.enabled=false，避免Spark自动调整分区覆盖你的设置。

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