Spark Shuffle源码分析系列之ShuffleReader 一

Shuffle涉及到三方面问题：Shuffle write写过程，中间数据记录过程以及Shuffle read读过程，上面几节我们分析了write和中间记录过程，本文将聚焦在Shuffle read部分。ShffuleRead什么时候进行数据读取？ShuffleMap产生的数据如何拉取过来？拉取过来的数据如何存储和处理？本节我将通过源码分析来解答这些问题。

问题

1. ShuffleRDD是何时产生的?
2. Shuffle-ReduceTask在什么时候开始执行，是等parent stage中的一个ShuffleMapTask 执行完还是等全部ShuffleMapTasks执行完？
3. Shuffle-ReduceTask怎么获得当前Reduce任务所需要数据的存放位置？
4. Shuffle-ReduceTask一边获取数据一边处理还是一次性获取完再处理？
5. Shuffle-ReduceTask获取到的数据存放到哪里，内存？磁盘？如何协调的？
6. Shuffle-ReduceTask如何处理具有聚合，排序要求的算子?

ShuffleRDD

ShuffleRDD是如何产生的

首先我们来看下ShuffleRDD是如何产生的，我们知道，RDD的计算链根据Shuffle被切分为不同的stage，一个stage的开始阶段一般是读取上一阶段的数据<数据有可能是从hdfs读取的HadoopRDD或者是上一个Shuffle产出的RDD，如果是上一个shuffle的结果，则该stage读取数据的过程其实就是reduce过程>，然后经过该stage的计算链后得到结果数据，再然后就会把这些数据写入到磁盘供下一个stage读取。我们来看一个使用简单的Shuffle算子的例子，我们创建一个RDD，然后使用groupByKey进行分组，然后对每个key对应的数据拿取数据迭代器，最后使用collect这个action算子来触发真实计算，如下所示:

scala> var rdd1 = sc.makeRDD(Array(("A",0),("A",2),("B",1),("B",2),("C",1)))
rdd1: org.apache.spark.rdd.RDD[(String, Int)] = ParallelCollectionRDD[0] at makeRDD at <console>:24

scala> val rdd2 = rdd1.groupByKey()
rdd2: org.apache.spark.rdd.RDD[(String, Iterable[Int])] = ShuffledRDD[1] at groupByKey at <console>:25

scala> val rdd3 = rdd2.map(row => row._2)
rdd3: org.apache.spark.rdd.RDD[Iterable[Int]] = MapPartitionsRDD[2] at map at <console>:25

scala> rdd3.collect
res1: Array[Iterable[Int]] = Array(CompactBuffer(0, 2), CompactBuffer(2, 1), CompactBuffer(1))

我们在Spark UI上面看下该job是如何划分stage的，可以看出来划分了两个stage，groupByKey算子由于涉及到Shuffle操作，将这个job划分为两个stage，Stage 0对应的ShuffleMapTask，Stage 1对应的是ShuffleReadTask也有可能是下一个Shuffle的MapTask，如下所示:

Stage 0主要是读取RDD，然后按照key进行合并数据，写入到磁盘中:

Stage 1是读取ShuffleMapTask的数据，然后进行map算子计算，获取数据，可以看出来这个时候ReduceTask的初始数据RDD是ShuffleRDD。

我们最后来看下执行toDebugString打印相应RDD的Lineage，和上面DAG图是一样的:

scala> rdd3.toDebugString
res0: String =
(5) MapPartitionsRDD[2] at map at <console>:25 []
 |  ShuffledRDD[1] at groupByKey at <console>:25 []
 +-(5) ParallelCollectionRDD[0] at makeRDD at <console>:24 []

所以我们可以看出来，在Shuffle产生时候，上游进行ShuffleMapTask，写入数据到文件中，下游stage读取这些数据形成的RDD即为ShuffleRDD。

ShuffleRDD何时去读取数据

在MapReduce中，可以通过设置mapred.reduce.slowstart.completed.maps来控制map端任务执行到多少比例后可以启动reduce端任务的计算，Spark中是不是也是一定比例的ShuffleMapTask结束后就可以进行Reduce任务执行，从而拉取数据，答案是否定的，Spark中Shuffle将job划分为多个stage，一个stage要想开始执行，必须满足它所依赖的stages都执行完成，也就是ShuffleMapTask计算结束后。

ShuffleRDD是如何得到数据

ShuffleRDD继承自RDD，实现了compute逻辑，实现了如何获取数据的，源码如下，可以看出来先通过ShuffleManager获取相应的Reader<默认是BlockStoreShuffleReader>，然后通过read方法，获取数据迭代器，给下游算子计算使用。

override def compute(split: Partition, context: TaskContext): Iterator[(K, C)] = {
   
  // 获取第一个依赖
  val dep = dependencies.head.asInstanceOf[ShuffleDependency[K, V, C]]
  // 首先调用SortShuffleManager的getReader()方法获取BlockStoreShuffleReader
  SparkEnv.get.shuffleManager.getReader(dep.shuffleHandle, split.index, split.index + 1, context)
  // 调用BlockStoreShuffleReader的read()方法获取map任务输出的Block并在reduce端进行聚合或排序
  .read()
  .asInstanceOf[Iterator[(K, C)]]
}

我们再来看下Reader的初始化操作:

// 用于获取对map任务输出的分区数据文件中从startPartition到endPartition-1范围内的数据
// 进行读取的读取器（即BlockStoreShuffleReader），供reduce任务使用。
override def getReader[K, C](
  handle: ShuffleHandle,
  startPartition: Int,
  endPartition: Int,
  context: TaskContext): ShuffleReader[K, C] = {
   
  new BlockStoreShuffleReader(
    handle.asInstanceOf[BaseShuffleHandle[K, _, C]], startPartition, endPartition, context)
}

另外我们看下ShuffleRDD中的一些成员变量，记录了是否包含自定义的序列化器，是否需要聚合数据，是否需要排序，是否进行了mapSide聚合操作，这些都是上游算子定义的:

private var userSpecifiedSerializer: Option[Serializer] = None
private var keyOrdering: Option[Ordering[K]] = None
private var aggregator: Option[Aggregator[K, V, C]] = None
private var mapSideCombine: Boolean = false

def setSerializer(serializer: Serializer): ShuffledRDD[K, V, C] = {
   
  this.userSpecifiedSerializer = Option(serializer)
  this
}

def setKeyOrdering(keyOrdering: Ordering[K]): ShuffledRDD[K, V, C] = {
   
  this.keyOrdering = Option(keyOrdering)
  this
}

def setAggregator(aggregator: Aggregator[K, V, C]