手记

Spark-Core源码精读(14)、Shuffle--Write部分

前面我们分析了Spark中具体的Task的提交和运行过程,从本文开始我们开始进入Shuffle的世界,Shuffle对于分布式计算来说是至关重要的部分,它直接影响了分布式系统的性能,所以我将尽可能进行详细的分析。

我们首先来看Shuffle中的Write部分:

override def runTask(context: TaskContext): MapStatus = {  // Deserialize the RDD using the broadcast variable.
  val deserializeStartTime = System.currentTimeMillis()  val ser = SparkEnv.get.closureSerializer.newInstance()  val (rdd, dep) = ser.deserialize[(RDD[_], ShuffleDependency[_, _, _])](    ByteBuffer.wrap(taskBinary.value), Thread.currentThread.getContextClassLoader)
  _executorDeserializeTime = System.currentTimeMillis() - deserializeStartTime
  metrics = Some(context.taskMetrics)  var writer: ShuffleWriter[Any, Any] = null
  try {    val manager = SparkEnv.get.shuffleManager
    writer = manager.getWriter[Any, Any](dep.shuffleHandle, partitionId, context)
    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
  }
}

首先根据SparkEnv获得ShuffleManager,ShuffleManager是为Spark shuffle系统而抽象的可插拔的接口,它被创建在Driver和Executor上,具体是在SparkEnv实例化的时候进行配置的,源码如下:

val shortShuffleMgrNames = Map(  "hash" -> "org.apache.spark.shuffle.hash.HashShuffleManager",  "sort" -> "org.apache.spark.shuffle.sort.SortShuffleManager",  "tungsten-sort" -> "org.apache.spark.shuffle.sort.SortShuffleManager")val shuffleMgrName = conf.get("spark.shuffle.manager", "sort")val shuffleMgrClass = shortShuffleMgrNames.getOrElse(shuffleMgrName.toLowerCase, shuffleMgrName)val shuffleManager = instantiateClass[ShuffleManager](shuffleMgrClass)

可以看到是由"spark.shuffle.manager"配置项来决定具体使用哪种实现方式,默认情况下使用的是sort的方式(本文参考的是Spark 1.6.3版本的源码)。Driver通过ShuffleManager来注册shuffles,Executors可以通过它来读写数据。

获得到ShuffleManager后,就根据它来获得ShuffleWriter(根据具体ShuffleManager的getWriter方法获得),顾名思义就是用来写数据,而接下来的工作就是调用具体的ShuffleWriter的write方法来进行写数据的工作。

先来看一下getWriter方法,这里的第一个参数dep.shuffleHandle是ShuffleDependency的一个成员变量:

val shuffleHandle: ShuffleHandle = _rdd.context.env.shuffleManager.registerShuffle(
  shuffleId, _rdd.partitions.size, this)

这里的registerShuffle方法用来向ShuffleManager注册一个shuffle并且获得一个用来传递任务的句柄,会根据不同ShuffleManager有不同的实现,HashShuffleManager返回的是BaseShuffleHandle,而SortShuffleManager又会根据不同的情况返回BypassMergeSortShuffleHandle、SerializedShuffleHandle或者BaseShuffleHandle。

HashShuffleManager:

override def registerShuffle[K, V, C](
    shuffleId: Int,
    numMaps: Int,
    dependency: ShuffleDependency[K, V, C]): ShuffleHandle = {  new BaseShuffleHandle(shuffleId, numMaps, dependency)
}

SortShuffleManager:

override def registerShuffle[K, V, C](
    shuffleId: Int,
    numMaps: Int,
    dependency: ShuffleDependency[K, V, C]): ShuffleHandle = {  if (SortShuffleWriter.shouldBypassMergeSort(SparkEnv.get.conf, dependency)) {    // 这里注释说的很清楚,根据spark.shuffle.sort.bypassMergeThreshold的值(默认是200)判断是否需要进行Map端的聚合操作
    // 如果partitions的个数小于200就不进行该操作
    // If there are fewer than spark.shuffle.sort.bypassMergeThreshold partitions and we don't
    // need map-side aggregation, then write numPartitions files directly and just concatenate
    // them at the end. This avoids doing serialization and deserialization twice to merge
    // together the spilled files, which would happen with the normal code path. The downside is
    // having multiple files open at a time and thus more memory allocated to buffers.
    new BypassMergeSortShuffleHandle[K, V](
      shuffleId, numMaps, dependency.asInstanceOf[ShuffleDependency[K, V, V]])  // 这里是判断是否使用tungsten的方式
  } else if (SortShuffleManager.canUseSerializedShuffle(dependency)) {    // Otherwise, try to buffer map outputs in a serialized form, since this is more efficient:
    new SerializedShuffleHandle[K, V](
      shuffleId, numMaps, dependency.asInstanceOf[ShuffleDependency[K, V, V]])
  } else {    // 如果不是上述两种方式,就使用默认的方式
    // Otherwise, buffer map outputs in a deserialized form:
    new BaseShuffleHandle(shuffleId, numMaps, dependency)
  }
}

使用一张图来总结一下上面的过程:

然后我们来看ShuffleManager的getWriter方法:

HashShuffleManager:

/** Get a writer for a given partition. Called on executors by map tasks. */override def getWriter[K, V](handle: ShuffleHandle, mapId: Int, context: TaskContext)
    : ShuffleWriter[K, V] = {  new HashShuffleWriter(
    shuffleBlockResolver, handle.asInstanceOf[BaseShuffleHandle[K, V, _]], mapId, context)
}

SortShuffleManager:

/** Get a writer for a given partition. Called on executors by map tasks. */override def getWriter[K, V](
    handle: ShuffleHandle,
    mapId: Int,
    context: TaskContext): ShuffleWriter[K, V] = {
  numMapsForShuffle.putIfAbsent(
    handle.shuffleId, handle.asInstanceOf[BaseShuffleHandle[_, _, _]].numMaps)  val env = SparkEnv.get
  handle match {    case unsafeShuffleHandle: SerializedShuffleHandle[K @unchecked, V @unchecked] =>      new UnsafeShuffleWriter(
        env.blockManager,
        shuffleBlockResolver.asInstanceOf[IndexShuffleBlockResolver],
        context.taskMemoryManager(),
        unsafeShuffleHandle,
        mapId,
        context,
        env.conf)    case bypassMergeSortHandle: BypassMergeSortShuffleHandle[K @unchecked, V @unchecked] =>      new BypassMergeSortShuffleWriter(
        env.blockManager,
        shuffleBlockResolver.asInstanceOf[IndexShuffleBlockResolver],
        bypassMergeSortHandle,
        mapId,
        context,
        env.conf)    case other: BaseShuffleHandle[K @unchecked, V @unchecked, _] =>      new SortShuffleWriter(shuffleBlockResolver, other, mapId, context)
  }
}

从上述源码可以看出,getWriter方法内部实际上是根据传进来的ShuffleHandle的具体类型来判断使用哪种ShuffleWriter的,然后最终执行ShuffleWriter的write方法,下面我们就分为HashShuffleManager和SortShuffleManager两种类型来进行分析。

1、HashShuffleManager

从上面的源码中可以看到HashShuffleManager最终实例化的是HashShuffleWriter,实例化的时候有一行比较重要的代码:

private val shuffle = shuffleBlockResolver.forMapTask(dep.shuffleId, mapId, numOutputSplits, ser,
  writeMetrics)

我们来看forMapTask这个方法:

def forMapTask(shuffleId: Int, mapId: Int, numReducers: Int, serializer: Serializer,
    writeMetrics: ShuffleWriteMetrics): ShuffleWriterGroup = {  // 为每个ShuffleMapTask实例化了一个ShuffleWriterGroup
  new ShuffleWriterGroup {    // 实例化ShuffleState并保存shuffleId和ShuffleState的对应关系
    shuffleStates.putIfAbsent(shuffleId, new ShuffleState(numReducers))    // 根据shuffleId获得对应的ShuffleState
    private val shuffleState = shuffleStates(shuffleId)    val openStartTime = System.nanoTime    val serializerInstance = serializer.newInstance()    // 获得该ShuffleWriterGroup的writers
    val writers: Array[DiskBlockObjectWriter] = {      Array.tabulate[DiskBlockObjectWriter](numReducers) { bucketId =>        // 生成ShuffleBlockId,是一个case class,我们可以通过name方法看到其具体的组成:
        // override def name: String = "shuffle_" + shuffleId + "_" + mapId + "_" + reduceId
        val blockId = ShuffleBlockId(shuffleId, mapId, bucketId)        // 通过DiskBlockManager的getFile方法获得File
        val blockFile = blockManager.diskBlockManager.getFile(blockId)        // 临时目录
        val tmp = Utils.tempFileWith(blockFile)        // 使用BlockManager的getDiskWriter方法获得DiskBlockObjectWriter
        // 注意这里的bufferSize默认情况下是32kb,可以通过spark.shuffle.file.buffer进行配置
        // private val bufferSize = conf.getSizeAsKb("spark.shuffle.file.buffer", "32k").toInt * 1024
        blockManager.getDiskWriter(blockId, tmp, serializerInstance, bufferSize, writeMetrics)
      }
    }    // Creating the file to write to and creating a disk writer both involve interacting with
    // the disk, so should be included in the shuffle write time.
    writeMetrics.incShuffleWriteTime(System.nanoTime - openStartTime)    override def releaseWriters(success: Boolean) {
      shuffleState.completedMapTasks.add(mapId)
    }
  }
}

为每个ShuffleMapTask实例化一个ShuffleWriterGroup,其中包含了一组writers,每个writer对应一个reducer。

然后我们进入到HashShuffleWriter的write方法:

/** Write a bunch of records to this task's output */override def write(records: Iterator[Product2[K, V]]): Unit = {  val iter = if (dep.aggregator.isDefined) {    // 判断是否进行map端的combine操作
    if (dep.mapSideCombine) {
      dep.aggregator.get.combineValuesByKey(records, context)
    } else {
      records
    }
  } else {
    require(!dep.mapSideCombine, "Map-side combine without Aggregator specified!")
    records
  }  for (elem <- iter) {    val bucketId = dep.partitioner.getPartition(elem._1)
    shuffle.writers(bucketId).write(elem._1, elem._2)
  }
}

如果没有进行Map端的combine操作,根据key获得bucketId,实际上是进行取模运算:

def nonNegativeMod(x: Int, mod: Int): Int = {  val rawMod = x % mod
  rawMod + (if (rawMod < 0) mod else 0)
}

下面就是根据bucketId获得ShuffleWriterGroup中对应的writer,然后执行其write方法,将key和value写入到对应的block file中:

/**
 * Writes a key-value pair.
 */def write(key: Any, value: Any) {  if (!initialized) {
    open()
  }
  objOut.writeKey(key)
  objOut.writeValue(value)
  recordWritten()
}

写入完成后,我们回到ShuffleMapTask的runTask方法中,接下来执行的是:

writer.stop(success = true).get

即HashShuffleWriter的stop方法,该方法返回的是Option[MapStatus],最主要的一句代码为:

Some(commitWritesAndBuildStatus())

进入到commitWritesAndBuildStatus方法:

private def commitWritesAndBuildStatus(): MapStatus = {  // Commit the writes. Get the size of each bucket block (total block size).
  val sizes: Array[Long] = shuffle.writers.map { writer: DiskBlockObjectWriter =>    // 调用DiskBlockObjectWriter的commitAndClose方法
    writer.commitAndClose()    // 获得每个bucket block的大小
    writer.fileSegment().length
  }  // 重命名所有的shuffle文件,每个executor只有一个ShuffleBlockResolver,所以使用了synchronized关键字
  // rename all shuffle files to final paths
  // Note: there is only one ShuffleBlockResolver in executor
  shuffleBlockResolver.synchronized {
    shuffle.writers.zipWithIndex.foreach { case (writer, i) =>      val output = blockManager.diskBlockManager.getFile(writer.blockId)      if (sizes(i) > 0) {        if (output.exists()) {          // Use length of existing file and delete our own temporary one
          sizes(i) = output.length()
          writer.file.delete()
        } else {          // Commit by renaming our temporary file to something the fetcher expects
          if (!writer.file.renameTo(output)) {            throw new IOException(s"fail to rename ${writer.file} to $output")
          }
        }
      } else {        if (output.exists()) {
          output.delete()
        }
      }
    }
  }  MapStatus(blockManager.shuffleServerId, sizes)
}

这里我们需要注意最终返回的是封装的MapStatus,它记录了产生的磁盘文件的位置,然后Executor中的MapOutputTrackerWorker将MapStatus信息发送给Driver中的MapOutputTrackerMaster,后面Shuffle Read的之后就会从Driver的MapOutputTrackerMaster获取MapStatus的信息,也就是获取对应的上一个ShuffleMapTask的计算结果的输出的文件位置信息。

Map端combine的情况

再来补充一下Map端combine的情况:

if (dep.mapSideCombine) {
  dep.aggregator.get.combineValuesByKey(records, context)
} else {
  records
}

进入到Aggregator的combineValuesByKey方法:

def combineValuesByKey(
    iter: Iterator[_ <: Product2[K, V]],
    context: TaskContext): Iterator[(K, C)] = {  // 首先实例化ExternalAppendOnlyMap
  val combiners = new ExternalAppendOnlyMap[K, V, C](createCombiner, mergeValue, mergeCombiners)  // 执行ExternalAppendOnlyMap的insertAll方法
  combiners.insertAll(iter)
  updateMetrics(context, combiners)
  combiners.iterator
}

首先实例化了ExternalAppendOnlyMap,然后执行ExternalAppendOnlyMap的insertAll方法:

def insertAll(entries: Iterator[Product2[K, V]]): Unit = {  if (currentMap == null) {    throw new IllegalStateException(      "Cannot insert new elements into a map after calling iterator")
  }  // An update function for the map that we reuse across entries to avoid allocating
  // a new closure each time
  var curEntry: Product2[K, V] = null
  val update: (Boolean, C) => C = (hadVal, oldVal) => {    if (hadVal) mergeValue(oldVal, curEntry._2) else createCombiner(curEntry._2)
  }  while (entries.hasNext) {
    curEntry = entries.next()    val estimatedSize = currentMap.estimateSize()    if (estimatedSize > _peakMemoryUsedBytes) {
      _peakMemoryUsedBytes = estimatedSize
    }    if (maybeSpill(currentMap, estimatedSize)) {
      currentMap = new SizeTrackingAppendOnlyMap[K, C]
    }
    currentMap.changeValue(curEntry._1, update)
    addElementsRead()
  }
}

具体的实现方式就不再解释了,简单的说就是将key相同的value进行合并,如果某个key有对应的值就执行merge(也可以理解为更新)操作,如果没有对应的值就新建一个combiner,需要注意的是如果内存不够的话就会将数据spill到磁盘。

HashShuffle方式的Shuffle Write部分至此结束,使用一张图概括一下:

接下来看一下SortShuffle方式的具体流程。

2、SortShuffleManager

为了解决Hash Shuffle产生小文件过多的问题,产生了Sort Shuffle,解下来我们就一起看一下Sort Shuffle的Write部分。

上文中我们已经提到,SortShuffleManager中的getWriter会根据不同的ShuffleHandle产生相应的ShuffleWriter:

  • SerializedShuffleHandle 对应 UnsafeShuffleWriter

  • BypassMergeSortShuffleHandle 对应 BypassMergeSortShuffleWriter

  • BaseShuffleHandle 对应 SortShuffleWriter

下面我们分别进行分析:

2.1、BaseShuffleHandle & SortShuffleWriter

首先来看一下SortShuffleWriter,直接来看它的write方法:

/** Write a bunch of records to this task's output */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 {    // In this case we pass neither an aggregator nor an ordering to the sorter, because we don't
    // care whether the keys get sorted in each partition; that will be done on the reduce side
    // if the operation being run is sortByKey.
    new ExternalSorter[K, V, V](
      context, aggregator = None, Some(dep.partitioner), ordering = None, dep.serializer)
  }
  sorter.insertAll(records)  // Don't bother including the time to open the merged output file in the shuffle write time,
  // because it just opens a single file, so is typically too fast to measure accurately
  // (see SPARK-3570).
  val output = shuffleBlockResolver.getDataFile(dep.shuffleId, mapId)  val tmp = Utils.tempFileWith(output)  try {    val blockId = ShuffleBlockId(dep.shuffleId, mapId, IndexShuffleBlockResolver.NOOP_REDUCE_ID)    val partitionLengths = sorter.writePartitionedFile(blockId, tmp)
    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}")
    }
  }
}

我们看到,内部有一个非常重要的部分,即ExternalSorter,而关于ExternalSorter的使用,源码中的注释说的很清楚了,这里就不做翻译了:

/**
* Users interact with this class in the following way:
*
* 1. Instantiate an ExternalSorter.
*
* 2. Call insertAll() with a set of records.
*
* 3. Request an iterator() back to traverse sorted/aggregated records.
*     - or -
*    Invoke writePartitionedFile() to create a file containing sorted/aggregated outputs
*    that can be used in Spark's sort shuffle.
*/

我们就根据这三个步骤进行说明:

2.1.1 第一步

首先就是实例化ExternalSorter,这里有一个判断,如果要进行map端的combine操作的话就需要指定Aggregator和Ordering,否则这两个参数为None。我们熟悉的reduceByKey就进行了Map端的combine操作,如下图所示:

2.1.2 第二步(这一步非常重要)

通过判断是否进行Map端combine操作而实例化出不同的ExternalSorter后,就会调用insertAll方法,将输入的记录写入到内存中,如果内存不足就spill到磁盘中,具体的实现我们来看insertAll方法:

def insertAll(records: Iterator[Product2[K, V]]): Unit = {  // TODO: stop combining if we find that the reduction factor isn't high
  // 首先判断是否需要进行Map端的combine操作
  val shouldCombine = aggregator.isDefined  if (shouldCombine) {    // 如果需要进行map端的combine操作,使用PartitionedAppendOnlyMap作为缓存
    // 将record根据key对value按照获得的聚合函数进行聚合操作(combine)
    // Combine values in-memory first using our AppendOnlyMap
    // 获得聚合函数,例如我们使用reduceByKey时编写的函数
    val mergeValue = aggregator.get.mergeValue    // 获取createCombiner函数
    val createCombiner = aggregator.get.createCombiner    var kv: Product2[K, V] = null
    // 定义update函数,主要的逻辑是:如果某个key已经存在记录(record)就使用上面获取
    // 的聚合函数进行聚合操作,如果还不存在记录就使用createCombiner方法进行初始化操作
    val update = (hadValue: Boolean, oldValue: C) => {      if (hadValue) mergeValue(oldValue, kv._2) else createCombiner(kv._2)
    }    // 循环遍历所有的records(记录)
    while (records.hasNext) {      // 记录spill的频率,每当read一条record的时候都会记录一次
      addElementsRead()      // 使用kv储存当前读的record
      kv = records.next()      // 这里的map和下面else中的buffer都是用来缓存的数据结构
      // 如果进行Map端的聚合操作,使用的就是PartitionedAppendOnlyMap[K, C]
      // 如果不进行Map端的聚合操作,使用的是PartitionedPairBuffer[K, C]
      // 调用上面定义的update函数将记录插入到map中
      map.changeValue((getPartition(kv._1), kv._1), update)      // 判断是否要进行spill操作
      maybeSpillCollection(usingMap = true)
    }
  } else {    // 如果不需要进行Map端的聚合操作,就直接将记录放到buffer(PartitionedPairBuffer)中
    // Stick values into our buffer
    while (records.hasNext) {
      addElementsRead()      val kv = records.next()
      buffer.insert(getPartition(kv._1), kv._1, kv._2.asInstanceOf[C])
      maybeSpillCollection(usingMap = false)
    }
  }
}

具体的流程用注释的方式写在了上面的源码中,这里我们先来看一下PartitionedAppendOnlyMap和PartitionedPairBuffer分别是如何工作的:

PartitionedAppendOnlyMap:

首先来看PartitionedAppendOnlyMap的changeValue实现,实际上,PartitionedAppendOnlyMap是继承自SizeTrackingAppendOnlyMap,而SizeTrackingAppendOnlyMap又继承自AppendOnlyMap,这里调用的changeValue方法实际上是SizeTrackingAppendOnlyMap的changeValue方法:

override def changeValue(key: K, updateFunc: (Boolean, V) => V): V = {  // 首先调用父类的changeValue方法
  val newValue = super.changeValue(key, updateFunc)  // 然后调用SizeTracker接口的afterUpdate方法
  super.afterUpdate()  // 返回newValue
  newValue
}

父类(AppendOnlyMap)的changeValue方法:

def changeValue(key: K, updateFunc: (Boolean, V) => V): V = {
  assert(!destroyed, destructionMessage)  val k = key.asInstanceOf[AnyRef]  // key为空时候的处理,增加长度
  if (k.eq(null)) {    if (!haveNullValue) {
      incrementSize()
    }
    nullValue = updateFunc(haveNullValue, nullValue)
    haveNullValue = true
    return nullValue
  }  var pos = rehash(k.hashCode) & mask  var i = 1
  while (true) {    // 这里的data是一个数组,用来同时存储key和value:key0, value0, key1, value1, key2, value2, etc.
    // 即2 * pos上存储的是key的值,2 * pos + 1上存储的是value的值
    val curKey = data(2 * pos)    // 如果key已经存在,就调用updateFunc方法更新value
    if (k.eq(curKey) || k.equals(curKey)) {      val newValue = updateFunc(true, data(2 * pos + 1).asInstanceOf[V])
      data(2 * pos + 1) = newValue.asInstanceOf[AnyRef]      return newValue
    } else if (curKey.eq(null)) {      // 如果key不存在就将该key和对应的value添加到data这个数组中
      val newValue = updateFunc(false, null.asInstanceOf[V])
      data(2 * pos) = k
      data(2 * pos + 1) = newValue.asInstanceOf[AnyRef]
      incrementSize()      return newValue
    } else {      // 否则继续计算位置(pos)
      val delta = i
      pos = (pos + delta) & mask
      i += 1
    }
  }  null.asInstanceOf[V] // Never reached but needed to keep compiler happy}

然后是SizeTracker接口的afterUpdate方法

protected def afterUpdate(): Unit = {
  numUpdates += 1
  if (nextSampleNum == numUpdates) {
    takeSample()
  }
}

更新数据的更新次数,如果更新的次数达到nextSampleNum,就执行采样操作,主要用来评估内存的使用情况。

PartitionedPairBuffer:

再来看PartitionedPairBuffer的insert方法,也就是不进行Map端combine操作的情况:

def insert(partition: Int, key: K, value: V): Unit = {  // 如果当前的大小达到了capacity的值就需要扩大该数组
  if (curSize == capacity) {
    growArray()
  }  // 存储key,这里存储的是(partition Id, key)的格式
  data(2 * curSize) = (partition, key.asInstanceOf[AnyRef])  // 存储value
  data(2 * curSize + 1) = value.asInstanceOf[AnyRef]
  curSize += 1
  // 参考上面PartitionedAppendOnlyMap的部分
  afterUpdate()
}

直接将数据存储到buffer中。

执行完上面的更新数据操作后,就要判断是否要将数据spill到磁盘,即maybeSpillCollection方法:

private def maybeSpillCollection(usingMap: Boolean): Unit = {  var estimatedSize = 0L  // 这里需要判断使用的是map(PartitionedAppendOnlyMap)还是buffer(PartitionedPairBuffer)
  // 如果true就是map,false就是buffer
  if (usingMap) {    // 估计当前map的内存占用大小
    estimatedSize = map.estimateSize()    // 如果超过内存的限制,就将缓存中的数据spill到磁盘
    if (maybeSpill(map, estimatedSize)) {      // spill到磁盘后,重置缓存
      map = new PartitionedAppendOnlyMap[K, C]
    }
  } else {    // 不进行Map端聚合操作的情况
    estimatedSize = buffer.estimateSize()    if (maybeSpill(buffer, estimatedSize)) {
      buffer = new PartitionedPairBuffer[K, C]
    }
  }  if (estimatedSize > _peakMemoryUsedBytes) {
    _peakMemoryUsedBytes = estimatedSize
  }
}

上面代码的主要作用就是估计当前缓存(map或者buffer)使用内存的大小,如果超过了内存使用的限制,就要将缓存中的数据spill到磁盘中,同时重置当前的缓存。

下面就来看一下maybeSpill方法:

// 如果成功spill到磁盘就返回true,否则返回falseprotected def maybeSpill(collection: C, currentMemory: Long): Boolean = {  var shouldSpill = false
  // 在进行真正的spill操作之前向TaskMemoryManager申请再多分配一些内存
  if (elementsRead % 32 == 0 && currentMemory >= myMemoryThreshold) {    // Claim up to double our current memory from the shuffle memory pool
    val amountToRequest = 2 * currentMemory - myMemoryThreshold    val granted =
      taskMemoryManager.acquireExecutionMemory(amountToRequest, MemoryMode.ON_HEAP, null)
    myMemoryThreshold += granted    // If we were granted too little memory to grow further (either tryToAcquire returned 0,
    // or we already had more memory than myMemoryThreshold), spill the current collection
    // 如果内存仍然不够用,就认定为需要spill到磁盘
    shouldSpill = currentMemory >= myMemoryThreshold
  }  // 如果内存中元素的个数超过了强制spill的上限也会认定为需要进行spill操作
  shouldSpill = shouldSpill || _elementsRead > numElementsForceSpillThreshold  // Actually spill
  // 接下来就真正将数据spill到磁盘
  if (shouldSpill) {
    _spillCount += 1
    logSpillage(currentMemory)    // spill操作
    spill(collection)
    _elementsRead = 0
    _memoryBytesSpilled += currentMemory
    releaseMemory()
  }
  shouldSpill
}

在进行真正的spill操作之前会向TaskMemoryManager申请再多分配一些内存,如果还不能够满足,或者不能分配更多的内存,或者内存中元素的个数超过了强制spill的上限,最终就会执行spill操作,接下来进入spill方法:

// 这里的collection就是指的map或者bufferoverride protected[this] def spill(collection: WritablePartitionedPairCollection[K, C]): Unit = {  // Because these files may be read during shuffle, their compression must be controlled by
  // spark.shuffle.compress instead of spark.shuffle.spill.compress, so we need to use
  // createTempShuffleBlock here; see SPARK-3426 for more context.
  // 获取临时的BlockId(TempShuffleBlockId)及对应的File
  val (blockId, file) = diskBlockManager.createTempShuffleBlock()  // These variables are reset after each flush
  var objectsWritten: Long = 0
  var spillMetrics: ShuffleWriteMetrics = null
  var writer: DiskBlockObjectWriter = null
  def openWriter(): Unit = {
    assert (writer == null && spillMetrics == null)
    spillMetrics = new ShuffleWriteMetrics
    writer = blockManager.getDiskWriter(blockId, file, serInstance, fileBufferSize, spillMetrics)
  }  // 获得DiskWriter(DiskBlockObjectWriter)
  openWriter()  // List of batch sizes (bytes) in the order they are written to disk
  // 用来储存每个batch对应的size
  val batchSizes = new ArrayBuffer[Long]  // How many elements we have in each partition
  // 用来储存每个partition有多少元素
  val elementsPerPartition = new Array[Long](numPartitions)  // Flush the disk writer's contents to disk, and update relevant variables.
  // The writer is closed at the end of this process, and cannot be reused.
  def flush(): Unit = {    val w = writer
    writer = null
    w.commitAndClose()
    _diskBytesSpilled += spillMetrics.shuffleBytesWritten
    batchSizes.append(spillMetrics.shuffleBytesWritten)
    spillMetrics = null
    objectsWritten = 0
  }  var success = false
  try {    // 排序部分的操作,返回迭代器
    val it = collection.destructiveSortedWritablePartitionedIterator(comparator)    // 循环的到的迭代器,执行write操作
    while (it.hasNext) {      val partitionId = it.nextPartition()
      require(partitionId >= 0 && partitionId < numPartitions,        s"partition Id: ${partitionId} should be in the range [0, ${numPartitions})")
      it.writeNext(writer)
      elementsPerPartition(partitionId) += 1
      objectsWritten += 1
      // 如果写的对象达到serializerBatchSize的个数时就进行flush操作
      if (objectsWritten == serializerBatchSize) {
        flush()
        openWriter()
      }
    }    if (objectsWritten > 0) {
      flush()
    } else if (writer != null) {      val w = writer
      writer = null
      w.revertPartialWritesAndClose()
    }
    success = true
  } finally {    if (!success) {      // This code path only happens if an exception was thrown above before we set success;
      // close our stuff and let the exception be thrown further
      if (writer != null) {
        writer.revertPartialWritesAndClose()
      }      if (file.exists()) {        if (!file.delete()) {
          logWarning(s"Error deleting ${file}")
        }
      }
    }
  }  // 实例化SpilledFile,并保存在数据结构ArrayBuffer[SpilledFile]中
  spills.append(SpilledFile(file, blockId, batchSizes.toArray, elementsPerPartition))
}

先来简单的看一下排序部分的逻辑:

def destructiveSortedWritablePartitionedIterator(keyComparator: Option[Comparator[K]])
  : WritablePartitionedIterator = {  // 这里的partitionedDestructiveSortedIterator会根据是map或者buffer有不同的实现
  val it = partitionedDestructiveSortedIterator(keyComparator)  // 最后返回的是WritablePartitionedIterator,上面进行写操作的时候就是调用该迭代器中的writeNext方法
  new WritablePartitionedIterator {    private[this] var cur = if (it.hasNext) it.next() else null
    def writeNext(writer: DiskBlockObjectWriter): Unit = {
      writer.write(cur._1._2, cur._2)
      cur = if (it.hasNext) it.next() else null
    }    def hasNext(): Boolean = cur != null
    def nextPartition(): Int = cur._1._1
  }
}

如果collection是map,具体的实现为(PartitionedAppendOnlyMap):

def partitionedDestructiveSortedIterator(keyComparator: Option[Comparator[K]])
  : Iterator[((Int, K), V)] = {  val comparator = keyComparator.map(partitionKeyComparator).getOrElse(partitionComparator)
  destructiveSortedIterator(comparator)
}

如果collection是buffer,具体的实现为(PartitionedPairBuffer):

override def partitionedDestructiveSortedIterator(keyComparator: Option[Comparator[K]])
  : Iterator[((Int, K), V)] = {  val comparator = keyComparator.map(partitionKeyComparator).getOrElse(partitionComparator)  new Sorter(new KVArraySortDataFormat[(Int, K), AnyRef]).sort(data, 0, curSize, comparator)
  iterator
}

这里需要注意,首先都是获取比较器,比较的数据格式为((partition Id, key), value),而两者的区别在于前者才是真正的Destructive级别的,具体的实现在destructiveSortedIterator方法中,而不管采用哪种方式,其底层都是通过timSort算法实现的,具体的排序逻辑就不详细说明了,有兴趣的朋友可以深入研究下去。

接下来就进入到第三步。

2.1.3 第三步

先贴出该步骤的代码:

...// 这里注释说的很清楚了,只打开了一个文件// Don't bother including the time to open the merged output file in the shuffle write time,// because it just opens a single file, so is typically too fast to measure accurately// (see SPARK-3570).val output = shuffleBlockResolver.getDataFile(dep.shuffleId, mapId)val tmp = Utils.tempFileWith(output)try {  // 构造blockId
  val blockId = ShuffleBlockId(dep.shuffleId, mapId, IndexShuffleBlockResolver.NOOP_REDUCE_ID)  // 写数据
  val partitionLengths = sorter.writePartitionedFile(blockId, tmp)  // 写index文件
  shuffleBlockResolver.writeIndexFileAndCommit(dep.shuffleId, mapId, partitionLengths, tmp)  // 进行Shuffle Read的时候需要参考该信息
  mapStatus = MapStatus(blockManager.shuffleServerId, partitionLengths)
} finally {  if (tmp.exists() && !tmp.delete()) {
    logError(s"Error while deleting temp file ${tmp.getAbsolutePath}")
  }
}

writePartitionedFile:

首先来看writePartitionedFile方法:

def writePartitionedFile(
    blockId: BlockId,
    outputFile: File): Array[Long] = {  // Track location of each range in the output file
  val lengths = new Array[Long](numPartitions)  // 首先判断spills中是否有数据,即判断是否有数据被spill到了磁盘中
  if (spills.isEmpty) {    // 数据只在内存中的情况
    // Case where we only have in-memory data
    val collection = if (aggregator.isDefined) map else buffer    // 获得迭代器
    val it = collection.destructiveSortedWritablePartitionedIterator(comparator)    // 进行迭代并将数据写到磁盘
    while (it.hasNext) {      val writer = blockManager.getDiskWriter(blockId, outputFile, serInstance, fileBufferSize,
        context.taskMetrics.shuffleWriteMetrics.get)      val partitionId = it.nextPartition()      while (it.hasNext && it.nextPartition() == partitionId) {
        it.writeNext(writer)
      }
      writer.commitAndClose()      val segment = writer.fileSegment()      // 最后返回的是每个partition写入的数据的长度
      lengths(partitionId) = segment.length
    }
  } else {    // 如果有数据被spill到了磁盘中,我们就需要进行merge-sort操作
    // We must perform merge-sort; get an iterator by partition and write everything directly.
    for ((id, elements) <- this.partitionedIterator) {      if (elements.hasNext) {        val writer = blockManager.getDiskWriter(blockId, outputFile, serInstance, fileBufferSize,
          context.taskMetrics.shuffleWriteMetrics.get)        for (elem <- elements) {
          writer.write(elem._1, elem._2)
        }
        writer.commitAndClose()        val segment = writer.fileSegment()
        lengths(id) = segment.length
      }
    }
  }
  context.taskMetrics().incMemoryBytesSpilled(memoryBytesSpilled)
  context.taskMetrics().incDiskBytesSpilled(diskBytesSpilled)
  context.internalMetricsToAccumulators(    InternalAccumulator.PEAK_EXECUTION_MEMORY).add(peakMemoryUsedBytes)
  lengths
}

下面我们就看一下this.partitionedIterator即内存和磁盘中的数据是如何合到一起的:

def partitionedIterator: Iterator[(Int, Iterator[Product2[K, C]])] = {  val usingMap = aggregator.isDefined  val collection: WritablePartitionedPairCollection[K, C] = if (usingMap) map else buffer  if (spills.isEmpty) {    // Special case: if we have only in-memory data, we don't need to merge streams, and perhaps
    // we don't even need to sort by anything other than partition ID
    if (!ordering.isDefined) {      // The user hasn't requested sorted keys, so only sort by partition ID, not key
      groupByPartition(collection.partitionedDestructiveSortedIterator(None))
    } else {      // We do need to sort by both partition ID and key
      groupByPartition(collection.partitionedDestructiveSortedIterator(Some(keyComparator)))
    }
  } else {    // Merge spilled and in-memory data
    merge(spills, collection.partitionedDestructiveSortedIterator(comparator))
  }
}

我们只考虑spills不为空的情况,即执行merge方法:

private def merge(spills: Seq[SpilledFile], inMemory: Iterator[((Int, K), C)])
    : Iterator[(Int, Iterator[Product2[K, C]])] = {  // 根据每个SpilledFile实例化一个SpillReader,这些SpillReader组成一个Seq
  val readers = spills.map(new SpillReader(_))  // 获得内存BufferedIterator
  val inMemBuffered = inMemory.buffered  // 根据partition的个数进行迭代
  (0 until numPartitions).iterator.map { p =>    // 实例化IteratorForPartition,即当前partition下的Iterator
    val inMemIterator = new IteratorForPartition(p, inMemBuffered)    // 这里就是合并操作
    val iterators = readers.map(_.readNextPartition()) ++ Seq(inMemIterator)    if (aggregator.isDefined) {      // Perform partial aggregation across partitions
      // 如果需要map端的combine操作则需要根据key进行聚合操作
      (p, mergeWithAggregation(
        iterators, aggregator.get.mergeCombiners, keyComparator, ordering.isDefined))
    } else if (ordering.isDefined) {      // No aggregator given, but we have an ordering (e.g. used by reduce tasks in sortByKey);
      // sort the elements without trying to merge them
      // 排序合并,例如sortByKey
      (p, mergeSort(iterators, ordering.get))
    } else {
      (p, iterators.iterator.flatten)
    }
  }
}

具体的mergeWithAggregation和mergeSort就不一一说明了,下面再来看一下writeIndexFileAndCommit

writeIndexFileAndCommit:

再来看writeIndexFileAndCommit方法:

def writeIndexFileAndCommit(
    shuffleId: Int,
    mapId: Int,
    lengths: Array[Long],
    dataTmp: File): Unit = {  val indexFile = getIndexFile(shuffleId, mapId)  val indexTmp = Utils.tempFileWith(indexFile)  try {    val out = new DataOutputStream(new BufferedOutputStream(new FileOutputStream(indexTmp)))    Utils.tryWithSafeFinally {      // We take in lengths of each block, need to convert it to offsets.
      var offset = 0L
      out.writeLong(offset)      for (length <- lengths) {
        offset += length
        out.writeLong(offset)
      }
    } {
      out.close()
    }    val dataFile = getDataFile(shuffleId, mapId)    // There is only one IndexShuffleBlockResolver per executor, this synchronization make sure
    // the following check and rename are atomic.
    synchronized {      val existingLengths = checkIndexAndDataFile(indexFile, dataFile, lengths.length)      if (existingLengths != null) {        // Another attempt for the same task has already written our map outputs successfully,
        // so just use the existing partition lengths and delete our temporary map outputs.
        System.arraycopy(existingLengths, 0, lengths, 0, lengths.length)        if (dataTmp != null && dataTmp.exists()) {
          dataTmp.delete()
        }
        indexTmp.delete()
      } else {        // This is the first successful attempt in writing the map outputs for this task,
        // so override any existing index and data files with the ones we wrote.
        if (indexFile.exists()) {
          indexFile.delete()
        }        if (dataFile.exists()) {
          dataFile.delete()
        }        if (!indexTmp.renameTo(indexFile)) {          throw new IOException("fail to rename file " + indexTmp + " to " + indexFile)
        }        if (dataTmp != null && dataTmp.exists() && !dataTmp.renameTo(dataFile)) {          throw new IOException("fail to rename file " + dataTmp + " to " + dataFile)
        }
      }
    }
  } finally {    if (indexTmp.exists() && !indexTmp.delete()) {
      logError(s"Failed to delete temporary index file at ${indexTmp.getAbsolutePath}")
    }
  }
}

具体的实现就不详细说明了,主要就是根据上一步的到的partition长度的数组将偏移量写入到index文件中。

最后就是实例化MapStatus,shuffle read的时候根据MapStatus获取数据。至此BaseShuffleHandle & SortShuffleWriter的部分就结束了。

BypassMergeSortShuffleHandle & BypassMergeSortShuffleWriter

再来看一下BypassMergeSortShuffleWriter的write方法:

public void write(Iterator<Product2<K, V>> records) throws IOException {
  assert (partitionWriters == null);  if (!records.hasNext()) {
    partitionLengths = new long[numPartitions];
    shuffleBlockResolver.writeIndexFileAndCommit(shuffleId, mapId, partitionLengths, null);
    mapStatus = MapStatus$.MODULE$.apply(blockManager.shuffleServerId(), partitionLengths);    return;
  }  final SerializerInstance serInstance = serializer.newInstance();  final long openStartTime = System.nanoTime();
  partitionWriters = new DiskBlockObjectWriter[numPartitions];  // 针对每一个reducer建立一个临时文件
  for (int i = 0; i < numPartitions; i++) {    final Tuple2<TempShuffleBlockId, File> tempShuffleBlockIdPlusFile =
      blockManager.diskBlockManager().createTempShuffleBlock();    final File file = tempShuffleBlockIdPlusFile._2();    final BlockId blockId = tempShuffleBlockIdPlusFile._1();
    partitionWriters[i] =
      blockManager.getDiskWriter(blockId, file, serInstance, fileBufferSize, writeMetrics).open();
  }  // Creating the file to write to and creating a disk writer both involve interacting with
  // the disk, and can take a long time in aggregate when we open many files, so should be
  // included in the shuffle write time.
  writeMetrics.incShuffleWriteTime(System.nanoTime() - openStartTime);  // 根据partition将记录写入到不同的临时文件中
  while (records.hasNext()) {    final Product2<K, V> record = records.next();    final K key = record._1();
    partitionWriters[partitioner.getPartition(key)].write(key, record._2());
  }  for (DiskBlockObjectWriter writer : partitionWriters) {
    writer.commitAndClose();
  }  // 将所有的临时文件内容按照partition Id合并到一个文件
  File output = shuffleBlockResolver.getDataFile(shuffleId, mapId);  File tmp = Utils.tempFileWith(output);  try {    // 将记录和partition长度信息分别写入到data文件和index文件中
    partitionLengths = writePartitionedFile(tmp);
    shuffleBlockResolver.writeIndexFileAndCommit(shuffleId, mapId, partitionLengths, tmp);
  } finally {    if (tmp.exists() && !tmp.delete()) {
      logger.error("Error while deleting temp file {}", tmp.getAbsolutePath());
    }
  }
  mapStatus = MapStatus$.MODULE$.apply(blockManager.shuffleServerId(), partitionLengths);
}

该writer在进行写记录之前会根据reducer的个数(例如R个)生成R个临时文件,然后将记录写入对应的临时文件中,最后将这些文件进行合并操作并写入到一个文件中,由于直接将记录写入了临时文件,并没有缓存在内存中,所以如果reducer的个数过多的话,就会为每个reducer打开一个临时文件,如果reducer的数量过多的话就会影响性能,所以使用该种方式需要满足一下条件(下面是源码中的注释):


  1. no Ordering is specified.


  1. no Aggregator is specified.


  1. the number of partitions is less than spark.shuffle.sort.bypassMergeThreshold.

其中spark.shuffle.sort.bypassMergeThreshold的个数默认为200个。

SerializedShuffleHandle & UnsafeShuffleWriter

最后再来看一下UnsafeShuffleWriter,也就是通常所说的Tungsten。

UnsafeShuffleWriter的write方法:

public void write(scala.collection.Iterator<Product2<K, V>> records) throws IOException {  // Keep track of success so we know if we encountered an exception
  // We do this rather than a standard try/catch/re-throw to handle
  // generic throwables.
  boolean success = false;  try {    while (records.hasNext()) {      // 循环便利所有记录,对其作用insertRecordIntoSorter方法
      insertRecordIntoSorter(records.next());
    }    // 将数据输出到磁盘上
    closeAndWriteOutput();
    success = true;
  } finally {    if (sorter != null) {      try {
        sorter.cleanupResources();
      } catch (Exception e) {        // Only throw this error if we won't be masking another
        // error.
        if (success) {          throw e;
        } else {
          logger.error("In addition to a failure during writing, we failed during " +                       "cleanup.", e);
        }
      }
    }
  }
}

首先来看insertRecordIntoSorter:

void insertRecordIntoSorter(Product2<K, V> record) throws IOException {  assert(sorter != null);  final K key = record._1();  final int partitionId = partitioner.getPartition(key);
  serBuffer.reset();
  serOutputStream.writeKey(key, OBJECT_CLASS_TAG);
  serOutputStream.writeValue(record._2(), OBJECT_CLASS_TAG);
  serOutputStream.flush();  final int serializedRecordSize = serBuffer.size();  assert (serializedRecordSize > 0);
  sorter.insertRecord(
    serBuffer.getBuf(), Platform.BYTE_ARRAY_OFFSET, serializedRecordSize, partitionId);
}

可以看出实际上调用的是ShuffleExternalSorter的insertRecord方法,限于篇幅,具体的底层实现暂不说明,以后有时间会单独分析一下Tungsten的部分,接下来看一下closeAndWriteOutput方法:

void closeAndWriteOutput() throws IOException {  assert(sorter != null);
  updatePeakMemoryUsed();
  serBuffer = null;
  serOutputStream = null;  // 获得spilled的文件
  final SpillInfo[] spills = sorter.closeAndGetSpills();
  sorter = null;  final long[] partitionLengths;  // 最终的输出文件
  final File output = shuffleBlockResolver.getDataFile(shuffleId, mapId);  // 临时文件
  final File tmp = Utils.tempFileWith(output);  try {    try {      // 将spilled的文件合并并写入到临时文件
      partitionLengths = mergeSpills(spills, tmp);
    } finally {      for (SpillInfo spill : spills) {        if (spill.file.exists() && ! spill.file.delete()) {
          logger.error("Error while deleting spill file {}", spill.file.getPath());
        }
      }
    }    // 将partition的长度信息写入index文件
    shuffleBlockResolver.writeIndexFileAndCommit(shuffleId, mapId, partitionLengths, tmp);
  } finally {    if (tmp.exists() && !tmp.delete()) {
      logger.error("Error while deleting temp file {}", tmp.getAbsolutePath());
    }
  }  // mapStatus
  mapStatus = MapStatus$.MODULE$.apply(blockManager.shuffleServerId(), partitionLengths);
}

用图来总结一下上面描述的三种方式如下所示:

BaseShuffleHandle & SortShuffleWriter

SerializedShuffleHandle & UnsafeShuffleWriter

BypassMergeSortShuffleHandle & BypassMergeSortShuffleWriter

最后需要说明的是如果采用的是SortShuffleManager,最后每个task产生的文件的个数为2 * M(M代表Mapper端ShuffleMapTask的个数),相对于Hash的方式来说文件的个数明显减少。

至此Shuffle Write的部分就分析完了,下一遍文章会继续分析Shuffle Read的部分。

本文为原创,欢迎转载,转载请注明出处、作者,谢谢!



作者:sun4lower
链接:https://www.jianshu.com/p/766290f06329


0人推荐
随时随地看视频
慕课网APP