Spark SortShuffleWriter

这是三种ShuffleWriter中最通用的情况,对应BaseShuffleHandle,此时可以在map端进行数据合并,否则不向排序工具ExternalSorter传入排序相关参数,只会根据key值获取对应的分区id,来划分数据,不会在分区内排序,如果结果需要排序,例如sortByKey,会在reduce端获取shuffle数据后进行

  override def write(records: Iterator[Product2[K, V]]): Unit = {
    //根据是否在map端进行数据合并初始化ExternalSorter
    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 {
      // 不进行聚合,就不进行排序,如果有需要reduce端再进行排序
      new ExternalSorter[K, V, V](
        context, aggregator = None, Some(dep.partitioner), ordering = None, dep.serializer)
    }
    sorter.insertAll(records)

    // shuffle输出文件
    val output = shuffleBlockResolver.getDataFile(dep.shuffleId, mapId)
    val tmp = Utils.tempFileWith(output)
    try {
      val blockId = ShuffleBlockId(dep.shuffleId, mapId, IndexShuffleBlockResolver.NOOP_REDUCE_ID)
      //sorter中的数据写出到该文件中
      val partitionLengths = sorter.writePartitionedFile(blockId, tmp)
      // 写出对应的index文件,纪录每个Partition对应的偏移量
      shuffleBlockResolver.writeIndexFileAndCommit(dep.shuffleId, mapId, partitionLengths, tmp)
      //shuffleWriter的返回结果
      mapStatus = MapStatus(blockManager.shuffleServerId, partitionLengths)
    } finally {
      if (tmp.exists() && !tmp.delete()) {
        logError(s"Error while deleting temp file ${tmp.getAbsolutePath}")
      }
    }
  }

ExternalSorter[K, V, C]

继承关系

MemoryConsumer (org.apache.spark.memory)
  Spillable (org.apache.spark.util.collection)
    ExternalSorter (org.apache.spark.util.collection)

对[K,V]键值对数据排序,支持聚合,可以按照键值,转换为(K, C)类型的数据

首先将键值K根据分区器分配到相应的分区中,如果需要聚合数据,就对每个分区中的键值进行排序,否则,则类型C必须等于V,不会进行分区内部排序

参数:

  • aggregator 聚合器 - 可选,用于合并数据
  • partitioner 分区器 - 优先按分区id排序
  • ordering 排序 - 可选,为每个分区内的键值排序
  • serializer 序列化器 - 写出数据到磁盘时使用

只有在确实需要输出的K为有序的时候才提供Ordering。如果map端的ShuffleMapTask不需要合并,排序参数传递Null,以避免额外排序;另一方面,如果确实想要合并,有Ordering参数会更加高效,否则使用的hash值排序,然后hash值相同再判定key值本身是否相同进行合并

  1. 实例化一个ExternalSorter
  2. 用一组记录调用insertAll()
  3. 调用writePartitionedFile()写出一个文件,作为Shuffle输出,也支持调用iterator方法,返回排序合并过的所有元素

类在内部工作原理如下:

  • 重复填充内存数据缓冲区,如果要按K组合,使用PartitionedAppendOnlyMap,否则使用PartitionedPairBuffer。在这些缓冲区中,按分区ID对元素进行排序,然后也可能通过K值来排序。为避免每个K值多次调用分区器获取分区id,将分区ID与每条记录的K一起存储,作为实际的键值
  • 当每个缓冲区达到内存限制时,会将其写出(spill)到一个文件中,这个文件内存会按分区id顺序写出,如果需要进行聚合,排序时会先比较id,相同时再按K值或K值的哈希码(没有传递K的比较器的时候)排序
  • 当用户请求迭代器或文件输出时,溢出文件将以及剩余的内存数据合写成一个有序的文件。如果需要按K值进行聚合,可以使用排序参数决定顺序,或者先根据键值的哈希码排序,然后将相同hash的相互比较,如果相等就进行合并
  • 最后调用stop()来删除所有中间文件

源码分析:

Aggregator

Aggregator内部包含三个函数,分别用来初始化C,合并V到C,合并C,在map端合并时会用到

case class Aggregator[K, V, C] (
    createCombiner: V => C,
    mergeValue: (C, V) => C,
    mergeCombiners: (C, C) => C)
存储对象的数据结构
@volatile private var map = new PartitionedAppendOnlyMap[K, C]
@volatile private var buffer = new PartitionedPairBuffer[K, C]

PartitionedAppendOnlyMap底层是一个AppendOnlyMap类型的简单哈希表存储结构,特点是只会添加keys,不会删除key,数据存储在数组private var data = new Array[AnyRef](2 * capacity)中,每对key和value连续存放,PartitionedAppendOnlyMap对此进行了封装,它的键值是Tuple2类型的(Int, K),值是类型C,其中Int表示的是分区id,如果key值冲突,采用线性探测再散列处理冲突,同时混入了SizeTracker特质,能够评估自身大致占用的内存数量

PartitionedPairBuffer底层直接是一个数组存储数据,数据类型与PartitionedAppendOnlyMap相同,将每条记录顺序插入即可,在不需要进行map端聚合的时候使用,也混入了SizeTracker特质,用于评估自身大致占用的内存数量,其初始容量为64,即128长度的AnyRef类型数组,实际相邻存储((partitionId, K), V)

比较函数

排序使用的是TimSort,源自归并排序和插入排序

没有定义聚合使用WritablePartitionedPairCollection.partitionComparator,只比较分区id
如果定义了排序,使用WritablePartitionedPairCollection.partitionKeyComparator,先比较分区,后比较Key值

比较key值时,如果传入了自定义的Key值比较器,就是用该比较器进行Partition内部的比较,否则使用hash值比较,后续需要进行聚合时判断相同hash的key是否相同

// ExternalSorter 类
private val keyComparator: Comparator[K] = ordering.getOrElse(new Comparator[K] {
    override def compare(a: K, b: K): Int = {
      val h1 = if (a == null) 0 else a.hashCode()
      val h2 = if (b == null) 0 else b.hashCode()
      if (h1 < h2) -1 else if (h1 == h2) 0 else 1
    }
  })

private def comparator: Option[Comparator[K]] = {
  if (ordering.isDefined || aggregator.isDefined) {
    Some(keyComparator)
  } else {
    None
  }
}

// WritablePartitionedPairCollection类,两种存储结构都混入了该伴生特质,如果Comparator[K]非空,先比较键值中的分区id,相等时比较实际的key
def partitionKeyComparator[K](keyComparator: Comparator[K]): Comparator[(Int, K)] = {
  new Comparator[(Int, K)] {
    override def compare(a: (Int, K), b: (Int, K)): Int = {
      val partitionDiff = a._1 - b._1
      if (partitionDiff != 0) {
        partitionDiff
      } else {
        keyComparator.compare(a._2, b._2)
      }
    }
  }
}

//只比较分区
def partitionComparator[K]: Comparator[(Int, K)] = new Comparator[(Int, K)] {
  override def compare(a: (Int, K), b: (Int, K)): Int = {
    a._1 - b._1
  }
}
写入数据

具体调用:

  def insertAll(records: Iterator[Product2[K, V]]): Unit = {
    val shouldCombine = aggregator.isDefined

    if (shouldCombine) {
      // Combine values in-memory first using our AppendOnlyMap
      // 定义update函数,根据不同情况,决定是合并V值到C中,还是创建一个新的C
      val mergeValue = aggregator.get.mergeValue
      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)
      }
      while (records.hasNext) {
        //计数,元素数目加1
        addElementsRead()
        kv = records.next()
        //写入数据
        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)
      }
    }
  }

每次插入数据以后,都会调用maybeSpillCollection判断是否需要将数据写入磁盘,内部会调用maybeSpill,当底层存储的数据的集合内存不足申请内存时,通过spill(collection: C)写出

spill写出

SpillableExternalSorter的父类,当前集合中的记录数目是elementsRead,估计大小为currentMemory,当集合中记录条数为32的整数倍,且估计大小大于内存阈值时,先尝试申请两倍内存大小的空间,如果成功增大阈值不进行落盘,如果内存不足时,则写出集合中的数据到磁盘中,内存阈值初始大小通过spark.shuffle.spill.initialMemoryThreshold设定,初始为5MB

如果内存足够,即MemoryManager始终能够提供足够大小的内存空间,一般就不需要落盘,除非达到了底层数组存储数目的上限

  protected def maybeSpill(collection: C, currentMemory: Long): Boolean = {
    var shouldSpill = false
    // 集合中元素的数目为32的整数倍,且大于内存阈值时,判断是否需要写出
    if (elementsRead % 32 == 0 && currentMemory >= myMemoryThreshold) {
      // 内存扩容为当前两倍时需要额外的空间
      val amountToRequest = 2 * currentMemory - myMemoryThreshold
      // 确保有足够的空间,可以增加阈值,否则进行写出
      val granted = acquireMemory(amountToRequest)
      myMemoryThreshold += granted
      shouldSpill = currentMemory >= myMemoryThreshold
    }
    // 元素数目大于阈值,强制写出,通过"spark.shuffle.spill.numElementsForceSpillThreshold"设定,默认是Long.MaxValue
    shouldSpill = shouldSpill || _elementsRead > numElementsForceSpillThreshold
    // 实际进行写出,调用ExternalSorter的spill方法
    if (shouldSpill) {
      _spillCount += 1
      logSpillage(currentMemory)
      spill(collection)
      _elementsRead = 0
      _memoryBytesSpilled += currentMemory
      releaseMemory()
    }
    shouldSpill
  }
private val spills = new ArrayBuffer[SpilledFile]

  override protected[this] def spill(collection: WritablePartitionedPairCollection[K, C]): Unit = {
    //传入比较器,对内存数据进行排序,返回一个排序后结果的迭代器
    val inMemoryIterator = collection.destructiveSortedWritablePartitionedIterator(comparator)
    //写出到文件
    val spillFile = spillMemoryIteratorToDisk(inMemoryIterator)
    spills += spillFile
  }

写出函数spillMemoryIteratorToDisk通过DiskBlockManager获得一个文件,以及对应的DiskBlockObjectWriter,缓存是32K,这里分批(batch,通过spark.shuffle.spill.batchSize设定,默认10000纪录算一批)写出串行化后的数据,同时记录相关信息,封装返回,返回对象为SpilledFile(File, blockId, 每个batch写出的数据量大小 : Array[Long], 每个分区元素数目 : Array[Long])

最终写出数据

写出添加到ExternalSorter的所有数据:

  def writePartitionedFile(
      blockId: BlockId,
      outputFile: File): Array[Long] = {

    // 纪录每个分区的数据长度
    val lengths = new Array[Long](numPartitions)
    val writer = blockManager.getDiskWriter(blockId, outputFile, serInstance, fileBufferSize,
      context.taskMetrics().shuffleWriteMetrics)

    if (spills.isEmpty) {
      // 只有内存中的数据,没有溢出文件,那么顺序写出,记录每个分区大小即可
      val collection = if (aggregator.isDefined) map else buffer
      val it = collection.destructiveSortedWritablePartitionedIterator(comparator)
      while (it.hasNext) {
        val partitionId = it.nextPartition()
        while (it.hasNext && it.nextPartition() == partitionId) {
          it.writeNext(writer)
        }
        val segment = writer.commitAndGet()
        lengths(partitionId) = segment.length
      }
    } else {
      // 归并排序,把所有文件和当前内存中的数据合并,然后写出
      // 读取时是按照partition顺序逐个处理,一个分区一个分区的合并
      for ((id, elements) <- this.partitionedIterator) {
        if (elements.hasNext) {
          for (elem <- elements) {
            writer.write(elem._1, elem._2)
          }
          val segment = writer.commitAndGet()
          lengths(id) = segment.length
        }
      }
    }

    writer.close()
    ...
    lengths
  }

这里的核心是partitionedIterator,将已排序的文件序列和内存中的数据合并,返回Iterator[(Int, Iterator[Product2[K, C]])]迭代器,按分区分组,对每个分区,都有一个遍历其内容的迭代器,按顺序访问数据

  private def merge(spills: Seq[SpilledFile], inMemory: Iterator[((Int, K), C)])
      : Iterator[(Int, Iterator[Product2[K, C]])] = {
    // 文件读取器
    val readers = spills.map(new SpillReader(_))
    val inMemBuffered = inMemory.buffered
    //逐个分区处理
    (0 until numPartitions).iterator.map { p =>
      // 内存中的数据已经按照partition排序好了,所有可以通过迭代器顺序输出
      val inMemIterator = new IteratorForPartition(p, inMemBuffered)
      // 合并该分区ID下,文件数据迭代器和内存数据迭代器,类型为 Seq[Iterator[Product2[K, C]]]
      val iterators = readers.map(_.readNextPartition()) ++ Seq(inMemIterator)
      if (aggregator.isDefined) {
        // key值相同进行聚合
        (p, mergeWithAggregation(
          iterators, aggregator.get.mergeCombiners, keyComparator, ordering.isDefined))
      } else if (ordering.isDefined) {
        // 没有聚合,但指定顺序的情况
        (p, mergeSort(iterators, ordering.get))
      } else {
        (p, iterators.iterator.flatten)
      }
    }
  }
  • mergeSort:指定ordering时,算法采用的是经典的将多个有序数据,合并为一个有序序列的算法,底层采用优先级队列,将多个迭代器作为元素,根据迭代器内的第一个元素作为大小比较的依据,读取时将元素最小的迭代器出队列,取出元素,再加入队列
  • 什么都没有指定时,将Seq[Iterator[Product2[K, C]]]直接转换为Iterator[Product2[K, C]]

如果指定了聚合函数,实现细节:

  • 如果指定了ordering时,采用mergeSort,然后顺序读取时,key相同进行合并
  • 如果没有指定顺序,依然采用mergeSort,不过是使用hash值进行比较的,读取时,将key的hash值相同的记录作为一批,同时读取处理,将key值实际相同的进行合并,最后顺序输出即可
  private def mergeWithAggregation(
      iterators: Seq[Iterator[Product2[K, C]]],
      mergeCombiners: (C, C) => C,
      comparator: Comparator[K],
      totalOrder: Boolean)
      : Iterator[Product2[K, C]] =
  {
    if (!totalOrder) {
      // 没有提供自定义的排序方法,只是部分有序,使用hash值排序,
      // 存在hash值相同,key值不同的情况
      new Iterator[Iterator[Product2[K, C]]] {
        //对此partition内的数据,根据K的hash值进行排序,底层是优先队列实现
        val sorted = mergeSort(iterators, comparator).buffered

        // Buffers reused across elements to decrease memory allocation
        val keys = new ArrayBuffer[K]
        val combiners = new ArrayBuffer[C]

        override def hasNext: Boolean = sorted.hasNext

        override def next(): Iterator[Product2[K, C]] = {
          if (!hasNext) {
            throw new NoSuchElementException
          }
          // 清空缓存
          keys.clear()
          combiners.clear()
          // 获取第一个数据,不一定是sorter中的第一个数据
          val firstPair = sorted.next()
          // 添加到缓存
          keys += firstPair._1
          combiners += firstPair._2
          val key = firstPair._1
          // 获取下一个hash值相同的数据
          while (sorted.hasNext && comparator.compare(sorted.head._1, key) == 0) {
            // 当前数据
            val pair = sorted.next()
            var i = 0
            var foundKey = false
            // 遍历缓存,如果key值与当前数据相同,进行合并
            while (i < keys.size && !foundKey) {
              if (keys(i) == pair._1) {
                combiners(i) = mergeCombiners(combiners(i), pair._2)
                foundKey = true
              }
              i += 1
            }
            // 没有找到当前数据key值相同的记录,将当前记录添加到缓存
            if (!foundKey) {
              keys += pair._1
              combiners += pair._2
            }
          }

          //返回迭代器Iterator[Product2[K, C],他们key的hash值都相等
          keys.iterator.zip(combiners.iterator)
        }
      }.flatMap(i => i)//把二次迭代展开
    } else {
      // 定义了ordering,先进行归并排序,然后key值相同的简单合并起来
      new Iterator[Product2[K, C]] {
        val sorted = mergeSort(iterators, comparator).buffered

        override def hasNext: Boolean = sorted.hasNext

        override def next(): Product2[K, C] = {
          if (!hasNext) {
            throw new NoSuchElementException
          }
          val elem = sorted.next()
          val k = elem._1
          var c = elem._2
          while (sorted.hasNext && sorted.head._1 == k) {
            val pair = sorted.next()
            c = mergeCombiners(c, pair._2)
          }
          (k, c)
        }
      }
    }

推荐阅读更多精彩内容

  • 关于Mongodb的全面总结 MongoDB的内部构造《MongoDB The Definitive Guide》...
    中v中阅读 22,663评论 2 81
  • Spring Cloud为开发人员提供了快速构建分布式系统中一些常见模式的工具(例如配置管理,服务发现,断路器,智...
    卡卡罗2017阅读 89,528评论 13 124
  • 1.ios高性能编程 (1).内层 最小的内层平均值和峰值(2).耗电量 高效的算法和数据结构(3).初始化时...
    欧辰_OSR阅读 17,203评论 8 197
  • Swift1> Swift和OC的区别1.1> Swift没有地址/指针的概念1.2> 泛型1.3> 类型严谨 对...
    cosWriter阅读 6,532评论 1 21
  • 微微一笑很倾城 之前看过几遍小说,很喜欢。今天看到电视剧到最后一集了。在这样的情况下,终于要开始写作业了,从第一集...
    路小fei阅读 22评论 0 0