Spark中mapPartition引发Worker节点数据分布倾斜，如何按分区大小分配数据至Worker？

这确实是Spark默认调度机制里很容易踩的坑——它默认是按分区数量来把任务分配到Worker节点，完全不考虑每个分区的实际数据量大小，就像你遇到的这样：有的Worker扛着两个超大分区，有的却只处理几个小分区，负载严重失衡。

要实现按数据量均衡分配，你可以试试下面这几种方案：

1. 自定义基于数据量的分区器 #

这是最直接的解决方案，核心思路是先统计每个tag的总数据量，然后根据Worker节点数（你这里是3个）计算每个节点应该承载的目标数据量，再自定义分区逻辑，把数据累计到接近目标值时就划分到下一个分区。

举个Scala的示例代码：

// 第一步：先统计每个tag的总数据量
val tagCountMap = df.groupBy("tag").count().collectAsMap()
val totalData = tagCountMap.values.sum
val targetPerWorker = totalData / 3 // 每个Worker的目标数据量

// 自定义分区器
class SizeBasedPartitioner(targetSize: Long, tagCounts: Map[Int, Long]) extends Partitioner {
  override def numPartitions: Int = 3

  override def getPartition(key: Any): Int = {
    val tag = key.asInstanceOf[Int]
    var currentCumulative = 0L
    var partitionId = 0

    // 按tag的数据量累加，达到阈值就切换分区
    for ((currentTag, count) <- tagCounts if partitionId < numPartitions) {
      currentCumulative += count
      if (currentCumulative > targetSize) {
        partitionId += 1
        currentCumulative = count
      }
      if (currentTag == tag) return partitionId
    }
    partitionId
  }
}

// 应用自定义分区
val balancedDF = df.repartition(new SizeBasedPartitioner(targetPerWorker, tagCountMap), $"tag")

⚠️ 注意：这种方法依赖于预先统计的tag数据量，如果你的数据是动态更新的，需要重新执行统计步骤。另外，如果某个tag的数据量本身就超过了单Worker的目标值，那这个大分区只能单独分配给一个Worker，你可以考虑再按其他字段把这个大tag的数据拆分成更小的子分区。

2. 合并小分区后再重新分配 #

如果你的小分区数量多但数据量小，可以先把这些小分区合并成几个大分区，再分配到Worker节点。比如先将数据转换为RDD，然后使用coalesce配合自定义逻辑来合并：

// 将DataFrame转为RDD，以tag为key
val tagRDD = df.rdd.map(row => (row.getAs[Int]("tag"), row))

// 先按tag聚合（根据你的实际处理逻辑调整）
val aggregatedRDD = tagRDD.reduceByKey((a, b) => a)

// 使用自定义分区器重新分区
val balancedRDD = aggregatedRDD.partitionBy(new SizeBasedPartitioner(targetPerWorker, tagCountMap))

这种方式适合小分区较多的场景，能减少Task数量同时均衡负载。

3. 基于分桶表的持久化数据均衡 #

如果你的数据是需要持久化存储的，可以使用Spark SQL的分桶表来实现长期的负载均衡：

1. 先手动给每个tag分配桶ID：比如把最大的tag22分配到桶0，中等的tag23分配到桶1，剩下的所有小tag分配到桶2。
2. 创建分桶表：

CREATE TABLE balanced_tag_table (
  tag INT,
  count BIGINT,
  bucket_id INT
)
CLUSTERED BY (bucket_id) INTO 3 BUCKETS
STORED AS PARQUET;

3. 插入数据时带上预先分配的bucket_id，这样每个桶的数据量就会和你期望的一致，后续查询时Spark会自动按桶分配任务到Worker节点。

关键注意点 #

• Spark的Task是和分区一一对应的，所以要实现负载均衡，核心是让每个Task处理的数据量尽量接近，而不是单纯调整分区数量。
• 如果存在单个超大分区（比如你的tag22），唯一能拆分它的方式是基于其他字段再细分这个tag的数据，比如如果有时间字段，可以按时间拆分成多个子分区，让多个Worker并行处理这个大tag的数据。

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