SparkCore — SortShuffle源码分析上

SortShuffle源码分析上

  在上一篇博客中，我们从源码的角度分析了HashShuffle的两种机制的区别，对应之前的理论分析，现在我们分析SortShuffle的两种机制的区别。
  同理我们首先看SortShuffleWriter的write()方法：

override def write(records: Iterator[Product2[K, V]]): Unit = {
    // 创建一个非常重要的组件ExternalSorter，它会对数据进行聚合、分区和排序等操作。
    // 如果需要进行本地聚合，就将aggregator传入，假如不需要聚合，那么就传入none
    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)
    }
    // 这个函数特别重要，对partition数据进行排序、聚合或一些到磁盘等操作
    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).
    // 获取partition数据最终要聚合的那个文件目录
    val output = shuffleBlockResolver.getDataFile(dep.shuffleId, mapId)
    // 生成临时文件名
    val tmp = Utils.tempFileWith(output)
    // 生成一个ShuffleBlockId
    val blockId = ShuffleBlockId(dep.shuffleId, mapId, IndexShuffleBlockResolver.NOOP_REDUCE_ID)
    // 调用ExternalSorter的writePartitionedFile方法
    // 将已经处理的每个partition的数据全部写入磁盘文件中，返回offset数组，数组下标就是partitionID，值是offset
    val partitionLengths = sorter.writePartitionedFile(blockId, tmp)
    // 将写入磁盘文件中的partition的偏移量写入索引文件中，供reduce端进行数据拉取。
    shuffleBlockResolver.writeIndexFileAndCommit(dep.shuffleId, mapId, partitionLengths, tmp)
    // 返回MapStatus，包含了写入每个partition的在输出文件中的offset的信息
    mapStatus = MapStatus(blockManager.shuffleServerId, partitionLengths)
  }

  上面创建了一个非常重要的组件，ExternalSorter，它会对数据进行聚合、分区和排序等操作，根据是否需要本地聚合及排序来创建两种不同的sorter（这是不是bypass和非bypass的区别？），然后使用insertAll()将records也即partition数据进行聚合、排序和溢写到磁盘操作，并生成多个临时文件。接着创建Task最终输出文件名，使用writePartitionedFile将之前溢写到磁盘的临时文件中的数据进行merge，并写入Task输出文件中，也即一个Task就只会创建一个输出文件，在写入的过程中记录每个partition的写入文件的offset，并封装在MapStatus返回（其实这里注意一点，这些spill创建的临时文件，在SortShuffleWriter调用stop的时候会被删除，这点在之前源码中有体现，在runTask()这个方法中）。

insertAll方法

  这里面一个很重要的方法 insertAll()，我们下面分析一下这个方法：

def insertAll(records: Iterator[Product2[K, V]]): Unit = {
   	// 是否需要聚合
    val shouldCombine = aggregator.isDefined
    // 如果需要本地聚合
    if (shouldCombine) {
      // 合并输入类型，比如 (C, V) => C，将C类型和V进行合并得到C。
      val mergeValue = aggregator.get.mergeValue
      // 聚合器，它会在每个分区上都会执行，只要遇到没有处理过的key，就会执行该方法
      val createCombiner = aggregator.get.createCombiner
      var kv: Product2[K, V] = null
      // 给每个数据进行标识，是聚合还是创建
      val update = (hadValue: Boolean, oldValue: C) => {
        if (hadValue) mergeValue(oldValue, kv._2) else createCombiner(kv._2)
      }
      // 遍历partition数据
      while (records.hasNext) {
        // 用于记录写入数据
        addElementsRead()
        // 取出一个数据
        kv = records.next()
        // 将聚合后的值存入本地聚合缓存 PartitionedAppendOnlyMap中
        // 这里的map是很重要的一个组件，后面会反复使用，一会儿单独分析
        map.changeValue((getPartition(kv._1), kv._1), update)
        // 判断是否需要将数据溢写到磁盘
        maybeSpillCollection(usingMap = true)
      }
    } else {
      // 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)
      }
    }
  }

  insertAll()将partition数据取出放入map或buffer中，其中map是一个很重要的组件，它自定义了比较器Comparator，用于对key进行排序。buffer相比于map没有比较器，它只是存储partition的数据。下面我们先看一下map的类型。

// map的定义，注意是可变类型
private var map = new PartitionedAppendOnlyMap[K, C]

private[spark] class PartitionedAppendOnlyMap[K, V]
  extends SizeTrackingAppendOnlyMap[(Int, K), V] with WritablePartitionedPairCollection[K, V] {
  // 定义比较器并返回
  def partitionedDestructiveSortedIterator(keyComparator: Option[Comparator[K]])
    : Iterator[((Int, K), V)] = {
    // partitionKeyComparator: 这个比较器，获取partitionId的差值，或者 key的差值
    // partitionComparator : 比较两个(partitionId, key)的大小
    // 其实这里就有点类似自定义的一个二次排序，先按照partitionID进行排序，相等则按照key进行排序。
    // 说白了就是，不同partition之间按照partitionID排序，相同partition内部数据按照key排序
    val comparator = keyComparator.map(partitionKeyComparator).getOrElse(partitionComparator)
    // 使用自定义的comparator，对数据进行排序，内部使用到的是归并排序。
    destructiveSortedIterator(comparator)
  }

  def insert(partition: Int, key: K, value: V): Unit = {
    update((partition, key), value)
  }
}

def destructiveSortedIterator(keyComparator: Comparator[K]): Iterator[(K, V)] = {
    destroyed = true
    // Pack KV pairs into the front of the underlying array
    var keyIndex, newIndex = 0

    // 对data中的数据进行整理，都放在前面
    while (keyIndex < capacity) {
      if (data(2 * keyIndex) != null) {
        data(2 * newIndex) = data(2 * keyIndex)
        data(2 * newIndex + 1) = data(2 * keyIndex + 1)
        newIndex += 1
      }
      keyIndex += 1
    }
    assert(curSize == newIndex + (if (haveNullValue) 1 else 0))

    // 使用前面生成的keyComparator进行排序，内部使用的是TimSort，一种归并排序
    new Sorter(new KVArraySortDataFormat[K, AnyRef]).sort(data, 0, newIndex, keyComparator)

    // 返回数据迭代器，里面的数据是经过归并排序的数据，为什么是归并排序，
    // 可能是因为多个executor并行运行的关系，多个task的数据进行归并排序
    new Iterator[(K, V)] {
      var i = 0
      var nullValueReady = haveNullValue
      def hasNext: Boolean = (i < newIndex || nullValueReady)
      def next(): (K, V) = {
        if (nullValueReady) {
          nullValueReady = false
          (null.asInstanceOf[K], nullValue)
        } else {
          val item = (data(2 * i).asInstanceOf[K], data(2 * i + 1).asInstanceOf[V])
          i += 1
          item
        }
      }
    }
  }

  map的类型是PartitionedAppendOnlyMap，它里面自定义了一个比较器，这个比较器的实现原理也很简单，不同分区之间按照partition ID的大小进行排序，每个partition内部的数据则按照key的大小排序。接着我们分析里面一个很重要的方法maybeSpillCollection()，它是将数据溢写到磁盘的关键。

maybeSpillCollection是否溢写到磁盘

  这个方法就是判断数据是否溢写到磁盘

private def maybeSpillCollection(usingMap: Boolean): Unit = {
    var estimatedSize = 0L
    // 假如在本地聚合
    if (usingMap) {
      // 当前数据大小
      estimatedSize = map.estimateSize()
      // maybeSpill，判断数据是否需要溢写到磁盘
      if (maybeSpill(map, estimatedSize)) {
        // 重新申请缓存
        map = new PartitionedAppendOnlyMap[K, C]
      }
    } else {
      estimatedSize = buffer.estimateSize()
      if (maybeSpill(buffer, estimatedSize)) {
        buffer = new PartitionedPairBuffer[K, C]
      }
    }
	// 更新当前使用的内存大小
    if (estimatedSize > _peakMemoryUsedBytes) {
      _peakMemoryUsedBytes = estimatedSize
    }
  }

  这个函数里面依然封装了一个函数maybeSpill()，用于判断当前数据量是否足够大，是否需要溢写到磁盘，下面分析一下这个方法：

maybeSpill()判断是否需要溢写到磁盘

protected def maybeSpill(collection: C, currentMemory: Long): Boolean = {
    var shouldSpill = false
    // 当前字节是否是32整数倍，并且必须大于5M
    if (elementsRead % 32 == 0 && currentMemory >= myMemoryThreshold) {
      // 首先尝试对当前缓存进行扩容
      val amountToRequest = 2 * currentMemory - myMemoryThreshold
      // 去TaskMemoryManager上去尝试获取内存
      val granted =
        taskMemoryManager.acquireExecutionMemory(amountToRequest, MemoryMode.ON_HEAP, null)
      // 将获取到的内存更新到阈值中
      myMemoryThreshold += granted
      // 假如当前数据大小依然大于扩容后的阈值，那么shouldSpill为true，否则为false
      shouldSpill = currentMemory >= myMemoryThreshold
    }
    // 如果计数器的值大于Long.MaxValue，那么shouldSpill也为true，否则为false
    shouldSpill = shouldSpill || _elementsRead > numElementsForceSpillThreshold
    // 根据上述两个阈值判断结果，决定如果需要溢写到磁盘
    if (shouldSpill) {
      // 溢写计数器，记录产生多少文件
      _spillCount += 1
      // 打印溢写日志信息
      logSpillage(currentMemory)
      // 将数据集合(key-value)溢写到磁盘
      spill(collection)
      // 计数器清零
      _elementsRead = 0
      // 释放内存
      _memoryBytesSpilled += currentMemory
      releaseMemory()
    }
    // 返回是否溢写到磁盘，继而重新申请map
    shouldSpill
  }

  maybeSpill的主要功能就是判断什么情况下需要溢写到磁盘，假设大于扩容后的门限，或者达到Long.MAX_VALUE，那么就会发生溢写操作，否则就不进行溢写，这就对应之前shuffle操作中的达到一定的阈值就写入磁盘。这里就是那个阈值，它通过两种方法判断。
  由于SortShuffle的中方法嵌套调用比较多，这里拆分上下节进行分析；总结一下，insertAll()对整个Task的数据进行处理，包括对数据进行排序、聚合和溢写磁盘等操作，并合并溢写到磁盘的临时文件，以及对应的索引文件等。其中判断什么时候进行溢写操作，主要是通过两个阈值，一个是数据大小，比如超过5M（默认大小，但是当Task内存足够，那么会相应扩容），另一个就是根据数量来进行溢写，默认大小是Long.MaxValue。下篇我们会分析如何进行溢写操作spill(collection)这个方法。