Spark RDD sortBy算子什么情况会触发shuffle

在 Spark 的 RDD 中，sortBy 是一个排序算子，虽然它在某些场景下可能看起来是分区内排序，但实际上在需要全局排序时会触发 Shuffle。这里我们分析其底层逻辑，结合源码和原理来解释为什么会有 Shuffle 的发生。

1. 为什么 `sortBy` 会触发 Shuffle？

关键点 1：全局有序性要求

sortBy 并非单纯的分区内排序。它的目标是按照用户指定的键对整个 RDD 的数据进行排序，这种操作需要保证全局顺序。为实现这一点，必须：

对数据进行 重新分区（Repartition），确保每个分区中的数据按照全局范围内的排序键正确分布；
每个分区内部再完成排序。

这些步骤不可避免地引入了 Shuffle，因为数据需要从一个分区转移到另一个分区以保证全局有序性。

关键点 2：底层调用 `repartitionAndSortWithinPartitions`

sortBy 的底层实现会调用 repartitionAndSortWithinPartitions 方法：

this.keyBy(f).repartitionAndSortWithinPartitions(
  new RangePartitioner(numPartitions, this, ascending))(ordInverse).values

keyBy(f)：
- 将数据转化为 (key, value) 格式，key 是排序的关键字，value 是原始数据。
RangePartitioner：
- 使用 RangePartitioner 将数据根据排序键重新分区（这一步需要 Shuffle）。
repartitionAndSortWithinPartitions：
- 先 Shuffle 数据以保证每个分区内的 key 是按范围划分的；
- 然后对每个分区内的数据进行排序。

Shuffle 的触发

当目标分区数量与当前分区数量不一致时（用户指定分区数或默认分区数），会触发 Shuffle；
即使目标分区数一致，只要需要保证全局有序，也需要重新分布数据来确保各分区内数据按键范围划分。

2. Shuffle 的作用

全局排序：分区间重新分布数据，确保所有分区的排序键范围是连续的。
负载均衡：通过 RangePartitioner 分布数据，避免某些分区过大或过小的问题。
分区内排序：确保每个分区内部数据按键排序。

3. 源码分析

repartitionAndSortWithinPartitions 的核心逻辑如下：

def repartitionAndSortWithinPartitions(
    partitioner: Partitioner)(
    implicit ord: Ordering[K]): RDD[(K, V)] = withScope {
  val shuffled = new ShuffledRDD[K, V, V](this, partitioner)
  shuffled.setKeyOrdering(ord)
  new MapPartitionsRDD(shuffled, (context, pid, iter) => {
    val sorter = new ExternalSorter[K, V, V](context, Some(partitioner), Some(ord))
    sorter.insertAll(iter)
    context.taskMetrics().incMemoryBytesSpilled(sorter.memoryBytesSpilled)
    context.taskMetrics().incDiskBytesSpilled(sorter.diskBytesSpilled)
    context.taskMetrics().incPeakExecutionMemory(sorter.peakMemoryUsedBytes)
    sorter.iterator
  })
}

ShuffledRDD：
- 触发 Shuffle，将数据根据分区器重新分布。
ExternalSorter：
- 对每个分区内的数据进行排序（如果数据超出内存，会使用磁盘作为临时存储）。

4. 举例说明 Shuffle 的发生

sortBy 的行为取决于传递的参数。为了实现分区内排序，你需要明确控制 sortBy 的参数设置。如果不显式指定目标分区数（numPartitions 参数），sortBy 默认不会触发 Shuffle，因此只会在分区内排序。

例子 1：带 Shuffle 的全局排序

val rdd = sc.parallelize(Seq(5, 2, 4, 3, 1), numSlices = 2)
val sortedRdd = rdd.sortBy(x => x, ascending = true, numPartitions = 3)// 指定目标分区数
println(sortedRdd.collect().mkString(", "))

初始数据分区：
分区 1：[5, 2]，分区 2：[4, 3, 1]
重新分区和排序后：
分区 1：[1, 2]，分区 2：[3, 4]，分区 3：[5]
Shuffle 触发原因：
数据必须重新分布，确保分区键范围（[1-2], [3-4], [5]）。
特点：
触发 Shuffle 操作，数据按照 RangePartitioner 进行分区。
每个分区内局部排序后，实现全局排序。

例子 2：分区内排序（无 Shuffle）

val rdd = sc.parallelize(Seq(5, 2, 4, 3, 1), numSlices = 2) // 两个分区
// 如果只需要分区内排序，mapPartitions 提供了无 Shuffle 的选择。
val sorted = rdd.mapPartitions(partition => partition.toList.sorted.iterator)
sorted.collect().foreach(println)