[spark] Shuffle Write解析 (Sort Based Shuffle)

本文基于 Spark 2.1 进行解析

前言

从 Spark 2.0 开始移除了Hash Based Shuffle，想要了解可参考Shuffle 过程，本文将讲解 Sort Based Shuffle。

ShuffleMapTask的结果（ShuffleMapStage中FinalRDD的数据）都将写入磁盘，以供后续Stage拉取，即整个Shuffle包括前Stage的Shuffle Write和后Stage的Shuffle Read，由于内容较多，本文先解析Shuffle Write。

概述：

• 写records到内存缓冲区（一个数组维护的map），每次insert&update都需要检查是否达到溢写条件。
• 若需要溢写，将集合中的数据根据partitionId和key（若需要）排序后顺序溢写到一个临时的磁盘文件，并释放内存新建一个map放数据，每次溢写都是写一个新的临时文件。
• 一个task最终对应一个文件，将还在内存中的数据和已经spill的文件根据reduce端的partitionId进行合并，合并后需要再次聚合排序（若需要），再根据partition的顺序写入最终文件，并返回每个partition在文件中的偏移量，最后以MapStatus对象返回给driver并注册到MapOutputTrackerMaster中，后续reduce好通过它来访问。

入口

执行一个ShuffleMapTask最终的执行逻辑是调用了ShuffleMapTask类的runTask()方法：

override def runTask(context: TaskContext): MapStatus = {
    // Deserialize the RDD using the broadcast variable.
    val deserializeStartTime = System.currentTimeMillis()
    val ser = SparkEnv.get.closureSerializer.newInstance()
    // 从广播变量中反序列化出finalRDD和dependency
    val (rdd, dep) = ser.deserialize[(RDD[_], ShuffleDependency[_, _, _])](
      ByteBuffer.wrap(taskBinary.value), Thread.currentThread.getContextClassLoader)
    _executorDeserializeTime = System.currentTimeMillis() - deserializeStartTime

    var writer: ShuffleWriter[Any, Any] = null
    try {
      // 获取shuffleManager
      val manager = SparkEnv.get.shuffleManager
      // 通过shuffleManager的getWriter()方法，获得shuffle的writer
      writer = manager.getWriter[Any, Any](dep.shuffleHandle, partitionId, context)
       // 通过rdd指定分区的迭代器iterator方法来遍历每一条数据，再之上再调用writer的write方法以写数据
      writer.write(rdd.iterator(partition, context).asInstanceOf[Iterator[_ <: Product2[Any, Any]]])
      writer.stop(success = true).get
    } catch {
      case e: Exception =>
        try {
          if (writer != null) {
            writer.stop(success = false)
          }
        } catch {
          case e: Exception =>
            log.debug("Could not stop writer", e)
        }
        throw e
    }
  }

其中的finalRDD和dependency是在Driver端DAGScheluer中提交Stage的时候加入广播变量的。

接着通过SparkEnv获取shuffleManager，默认使用的是sort（对应的是org.apache.spark.shuffle.sort.SortShuffleManager），可通过spark.shuffle.manager设置。

然后调用了manager.getWriter方法，该方法中检测到满足Unsafe Shuffle条件会自动采用Unsafe Shuffle，否则采用Sort Shuffle。使用Unsafe Shuffle有几个限制，shuffle阶段不能有aggregate操作，分区数不能超过一定大小($2^{24}$−1，这是可编码的最大parition id)，所以像reduceByKey这类有aggregate操作的算子是不能使用Unsafe Shuffle。

这里暂时讨论Sort Shuffle的情况，即getWriter返回的是SortShuffleWriter，我们直接看writer.write发生了什么：

override def write(records: Iterator[Product2[K, V]]): Unit = {
    sorter = if (dep.mapSideCombine) {
      require(dep.aggregator.isDefined, "Map-side combine without Aggregator specified!")
      new ExternalSorter[K, V, C](
        context, dep.aggregator, Some(dep.partitioner), dep.keyOrdering, dep.serializer)
    } else {
      new ExternalSorter[K, V, V](
        context, aggregator = None, Some(dep.partitioner), ordering = None, dep.serializer)
    }
    // 写内存缓冲区，超过阈值则溢写到磁盘文件
    sorter.insertAll(records)
    // 获取该task的最终输出文件
    val output = shuffleBlockResolver.getDataFile(dep.shuffleId, mapId)
    val tmp = Utils.tempFileWith(output)
    try {
      val blockId = ShuffleBlockId(dep.shuffleId, mapId, IndexShuffleBlockResolver.NOOP_REDUCE_ID)
      // merge后写到data文件
      val partitionLengths = sorter.writePartitionedFile(blockId, tmp)
      // 写index文件
      shuffleBlockResolver.writeIndexFileAndCommit(dep.shuffleId, mapId, partitionLengths, tmp)
      mapStatus = MapStatus(blockManager.shuffleServerId, partitionLengths)
    } finally {
      if (tmp.exists() && !tmp.delete()) {
        logError(s"Error while deleting temp file ${tmp.getAbsolutePath}")
      }
    }
  }

• 通过判断是否有map端的combine来创建不同的ExternalSorter，若有则将对应的aggregator和keyOrdering作为参数传入。
• 调用sorter.insertAll(records)，将records写入内存缓冲区，超过阈值则溢写到磁盘文件。
• Merge内存记录和所有被spill到磁盘的文件，并写到最终的数据文件.data中。
• 将每个partition的偏移量写到index文件中。

先细看sorter.inster是怎么写到内存，并spil