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()方法：其中的finalRDD和dependency是在Driver端DAGScheluer中提交Stage的时候加入广播变量的。
  接着通过SparkEnv获取shuffleManager，默认使用的是sort（对应的是org.apache.spark.shuffle.sort.SortShuffleManager），可通过spark.shuffle.manager设置。然后调用了manager.getWriter方法，getWriter返回的是SortShuffleWriter，直接看writer.write发生了什么：

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

先细看sorter.inster是怎么写到内存，并spill到磁盘文件的：

• 需要聚合的情况，遍历records拿到record的KV，通过map的changeValue方法并根据update函数来对相同K的V进行聚合，这里的map是PartitionedAppendOnlyMap类型，只能添加数据不能删除数据，底层实现是一个数组，数组中存KV键值对的方式是[K1,V1,K2,V2…]，每一次操作后都会判断是否要spill到磁盘
• 不需要聚合的情况，直接将record放入buffer，然后判断是否要溢写到磁盘

先看map.changeValue方法到底是怎么通过map实现对数据combine的：

override def changeValue(key: K, updateFunc: (Boolean, V) => V): V = {
    // 通过聚合算法得到newValue
    val newValue = super.changeValue(key, updateFunc)
    // 跟新对map的大小采样
    super.afterUpdate()
    newValue
  }

super.changeValue的实现：
根据K的hashCode再哈希与上掩码 得到 pos，2 * pos 为 k 应该所在的位置，2 * pos + 1 为 k 对应的 v 所在的位置，获取k应该所在位置的原来的key：

• 若原来的key和当前的 k 相等，则通过update函数将两个v进行聚合并更新该位置的value
• 若原来的key存在但不和当前的k 相等，则说明hash冲突了，更新pos继续遍历
• 若原来的key不存在，则将当前k作为该位置的key，并通过update函数初始化该k对应的聚合结果，接着会通过incrementSize()方法进行扩容。更新curSize，若当前大小超过了阈值growThreshold（growThreshold是当前容量capacity的0.7倍），则通过growTable()来扩容：这里重新创建了一个两倍capacity 的数组来存放数据，将原来数组中的数据通过重新计算位置放到新数组里，将data替换为新的数组，并跟新一些变量。

此时聚合已经完成，回到changeValue方面里面，接下来会执行super.afterUpdate()方法来对map的大小进行采样：

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

若每遍历跟新一条record，都来对map进行采样估计大小，假设采样一次需要1ms，100w次采样就会花上16.7分钟，性能大大降低。所以这里只有当update次数达到nextSampleNum 的时候才通过takeSample()采样一次：这里估计每次跟新的变化量的逻辑是：(当前map大小-上次采样的时候的大小) / (当前update的次数 - 上次采样的时候的update次数)。

  接着计算下次需要采样的update次数，该次数是指数级增长的，基数是1.1，第一次采样后，要1.1次进行第二次采样，第1.1*1.1次后进行第三次采样，以此类推，开始增长慢，后面增长跨度会非常大。这里采样完成后回到insetAll方法，接着通过maybeSpillCollection方法判断是否需要spill：通过集合的estimateSize方法估计map的大小，若需要spill则将集合中的数据spill到磁盘文件，并且为集合创建一个新的对象放数据。
  先看看估计大小的方法estimateSize：以上次采样完更新的bytePerUpdate 作为最近平均每次跟新的大小，估计当前占用内存：（当前update次数-上次采样时的update次数） * 每次跟新大小 + 上次采样记录的大小。

获取到当前集合的大小后调用maybeSpill判断是否需要spill，这里有两种情况都可导致spill：

• 当前集合包含的records数超过了 numElementsForceSpillThreshold（默认为Long.MaxValue，可通过spark.shuffle.spill.numElementsForceSpillThreshold设置）
• 当前集合包含的records数为32的整数倍，并且当前集合的大小超过了申请的内存myMemoryThreshold（第一次申请默认为5 * 1024 * 1024，可通过spark.shuffle.spill.initialMemoryThreshold设置），此时并不会立即spill，会尝试申请更多的内存避免spill，这里尝试申请的内存为2倍集合大小减去当前已经申请的内存大小（实际申请到的内存为granted），若加上原来的内存还是比当前集合的大小要小则需要spill。

若需要spill，则跟新spill次数，调用spill(collection)方法进行溢写磁盘，并释放内存。跟进spill方法看看其具体实现：

override protected[this] def spill(collection: WritablePartitionedPairCollection[K, C]): Unit = {
    // 传入comparator将集合中的数据先根据partition排序再通过key排序后返回一个迭代器
    val inMemoryIterator = collection.destructiveSortedWritablePartitionedIterator(comparator)
    // 写到磁盘文件，并返回一个对该文件的描述对象SpilledFile
    val spillFile = spillMemoryIteratorToDisk(inMemoryIterator)
    // 添加到spill文件数组
    spills.append(spillFile)
  }

继续跟进看看spillMemoryIteratorToDisk的实现：
  通过diskBlockManager创建临时文件和blockID，临时文件名格式为是 “temp_shuffle_” + id，遍历内存数据迭代器，并调用Writer（DiskBlockObjectWriter）的write方法，当写的次数达到序列化大小则flush到磁盘文件，并重新打开writer，及跟新batchSizes等信息。
  最后返回一个SpilledFile对象，该对象包含了溢写的临时文件File，blockId，每次flush的到磁盘的大小，每个partition对应的数据条数。

spill完成，并且insertAll方法也执行完成，回到开始的SortShuffleWriter的write方法：获取最后的输出文件名及blockId，文件格式：

 "shuffle_" + shuffleId + "_" + mapId + "_" + reduceId + ".data"

接着通过sorter.writePartitionedFile方法来写文件，其中包括内存及所有spill文件的merge操作，看看起具体实现：

• 数据只存在内存中而没有spill文件，根据传入的比较函数comparator来对集合里的数据先根据partition排序再对里面的key排序并返回一个迭代器，遍历该迭代器得到所有recored，每一个partition对应一个writer，一个partition的数据写完后再flush到磁盘文件，并记录该partition的数据长度
• 数据有spill文件，通过方法partitionedIterator对内存和spill文件的数据进行merge-sort后返回一个（partitionId，对应分区的数据的迭代器）的迭代器，也是一个partition对应一个Writer，写完一个partition再flush到磁盘，并记录该partition数据的长度

接下来看看通过this.partitionedIterator方法是怎么将内存及spill文件的数据进行merge-sort的：
  这里在有spill文件的情况下会执行下面的merge方法，传入的是spill文件数组和内存中的数据进过partitionId和key排序后的数据迭代器，看看merge方法
  merge方法将属于同一个reduce端的partition的内存数据和spill文件数据合并起来，再进行聚合排序（有需要的话），最后返回（reduce对应的partitionId，该分区数据迭代器）

  将数据merge-sort后写入最终的文件后，需要将每个partition的偏移量持久化到文件以供后续每个reduce根据偏移量获取自己的数据，写偏移量的逻辑很简单，就是根据前面得到的partition长度的数组将偏移量写到index文件中。根据shuffleId和mapId获取index文件并创建一个写文件的文件流，按照reduce端partition对应的offset依次写到index文件中，最后创建一个MapStatus实例返回，包含了reduce端每个partition对应的偏移量。
  该对象将返回到Driver端的DAGScheluer处理，被添加到对应stage的OutputLoc里，当该stage的所有task完成的时候会将这些结果注册到MapOutputTrackerMaster，以便下一个stage的task就可以通过它来获取shuffle的结果的元数据信息。

Shuffle Read

  shuffle的下游Stage的第一个RDD是ShuffleRDD，通过其compute方法来获取上游Stage Shuffle Write溢写到磁盘文件数据的一个迭代器：从SparkEnv中获取shuffleManager（这里是SortShuffleManager），通过manager获取Reader并调用其read方法来得到一个迭代器。
  getReader方法实例化了一个BlockStoreShuffleReader，参数有需要获取分区对应的partitionId。看看read方法，首先实例化了ShuffleBlockFetcherIterator对象，其中一个参数：

mapOutputTracker.getMapSizesByExecutorId(handle.shuffleId, startPartition, endPartition)

该方法获取reduce端数据的来源的元数据，返回的是 Seq[(BlockManagerId, Seq[(BlockId, Long)])]，即数据是来自于哪个节点的哪些block的，并且block的数据大小是多少，看看getMapSizesByExecutorId是怎么实现的：

• 传入shuffleId获取对应shuffle的所有元数据信息
• 转换格式并获取指定partition的元数据

跟进getStatuses：
  若能从mapStatuses获取到则直接返回，若不能则向mapOutputTrackerMaster通信发送GetMapOutputStatuses消息来获取元数据。
  因为一个Executor对应一个CoarseGrainedExecutorBackend，构建CoarseGrainedExecutorBackend的时候会创建一个SparkEnv，创建SparkEnv的时候会创建一个mapOutputTracker，即mapOutputTracker和Executor一一对应，也就是每一个Executor都有一个mapOutputTracker来维护元数据信息。
  这里的mapStatuses就是mapOutputTracker保存元数据信息的，mapOutputTracker和Executor一一对应，在该Executor上完成的Shuffle Write的元数据信息都会保存在其mapStatus里面，另外通过远程获取的其他Executor上完成的Shuffle Write的元数据信息也会在当前的mapStatuses中保存。
  Executor对应的是mapOutputTrackerWorker，而Driver对应的是mapOutputTrackerMaster，两者都是在实例化SparkEnv的时候创建的，每个在Executor上完成的Shuffle Task的结果都会注册到driver端的mapOutputTrackerMaster中，即driver端的mapOutputTrackerMaster的mapStatuses保存这所有元数据信息，所以当一个Executor上的任务需要获取一个shuffle的输出时，会先在自己的mapStatuses中查找，找不到再和mapOutputTrackerMaster通信获取元数据。

mapOutputTrackerMaster收到消息后的处理逻辑：
  调用了tracker的post方法：将该Message加入了mapOutputRequests中，mapOutputRequests是一个链式阻塞队列，在mapOutputTrackerMaster初始化的时候专门启动了一个线程池来执行这些请求。看看线程处理类MessageLoop的run方法是怎么定义的：通过shuffleId获取对应序列化后的元数据信息并返回。
  具体看看getSerializedMapOutputStatuses的实现：大体思路是先从缓存中获取元数据（MapStatuses），获取到直接返回，若没有则从mapStatuses获取，获取到后将其序列化后返回，随后返回给mapOutputTrackerWorker（刚才与之通信的节点），mapOutputTracker收到回复后又将元数据序列化并加入当前Executor的mapStatuses中。
  再回到getMapSizesByExecutorId方法中，getStatuses得到shuffleID对应的所有的元数据信息后，通过convertMapStatuses方法将获得的元数据信息转化成形如Seq[(BlockManagerId, Seq[(BlockId, Long)])]格式的位置信息，用来读取指定的分区的数据：这里的参数statuses:Array[MapStatus]是前面获取的上游stage所有的shuffle Write 文件的元数据，并且是按map端的partitionId排序的，通过zipWithIndex将元素和这个元素在数组中的ID（索引号）组合成键/值对，这里的索引号即是map端的partitionId，再根据shuffleId、mapPartitionId、reducePartitionId来构建ShuffleBlockId（在map端的ShuffleBlockId构建中的reducePartitionId始终是0，因为一个ShuffleMapTask就一个Block，而这里加入的真正的reducePartitionId在后面通过index文件获取对应reduce端partition偏移量的时候需要用到），并估算得到对应数据的大小，因为后面获取远程数据的时候需要限制大小，最后返回位置信息。
  至此mapOutputTracker.getMapSizesByExecutorId(handle.shuffleId, startPartition, endPartition)方法完成，返回了指定分区对应的元数据MapStatus信息。

在初始化对象ShuffleBlockFetcherIterator的时候调用了其初始化方法initialize()：

• 区分local blocks和remote blocks，并返回远程请求FetchRequest加入到fetchRequests队列中
• 从fetchRequests取出远程请求，并使用sendRequest方法发送请求，获取远程数据
• 获取本地blocks

先看是怎么区分local blocks和remote blocks的：

• 为了增加在远程节点获取数据的并行度，将一个请求的大小限制除以5作为最终的大小限制，即每次最多启动5个线程去最多5个节点上读取数据
• 判断是否是本地blocks的条件是block所在的executor和当前executor是否是同一个
• 遍历远程数据节点（Executor节点）的blocks，在一个节点上的请求数据超过大小限制则构建一个FetchRequest并加入到remoteRequests中，最后返回远程请求remoteRequests，这里的FetchRequest是对一个请求数据的包装，包括地址和blockId及大小

区分完local remote blocks后加入到了队列fetchRequests中，并调用fetchUpToMaxBytes()来获取远程数据：
  从fetchRequests中取出FetchRequest，并调用了sendRequest方法。通过shuffleClient的fetchBlocks方法来获取对应远程节点上的数据，默认是通过NettyBlockTransferService的fetchBlocks方法实现的，不管是成功还是失败都将构建SuccessFetchResult & FailureFetchResult 结果放入results中。获取完远程的数据接着通过fetchLocalBlocks()方法来获取本地的blocks信息：迭代需要获取的block，直接从blockManager中获取数据，并通过结果数据构建SuccessFetchResult或者FailureFetchResult放入results中，看看在blockManager.getBlockData(blockId)的实现：

override def getBlockData(blockId: BlockId): ManagedBuffer = {
    if (blockId.isShuffle) {
      shuffleManager.shuffleBlockResolver.getBlockData(blockId.asInstanceOf[ShuffleBlockId])
    } else {
      getLocalBytes(blockId) match {
        case Some(buffer) => new BlockManagerManagedBuffer(blockInfoManager, blockId, buffer)
        case None =>
          reportBlockStatus(blockId, BlockStatus.empty)
          throw new BlockNotFoundException(blockId.toString)
      }
    }
  }

再看看getBlockData方法：根据shuffleId和mapId获取index文件，并创建一个读文件的文件流，根据block的reduceId（上面获取对应partition元数据的时候提到过）跳过对应的Block的数据区，先后获取开始和结束的offset，然后在数据文件中读取数据。

  得到所有数据结果result后，再回到read()方法中：这里的ShuffleBlockFetcherIterator继承了Iterator，results可以被迭代，在其next()方法中将FetchResult以（blockId，inputStream）的形式返回。在read()方法的后半部分会进行聚合和排序，和Shuffle Write部分很类似，这里大致描述一下。在需要聚合的前提下，有map端聚合的时候执行combineCombinersByKey，没有则执行combineValuesByKey，但最终都调用了ExternalAppendOnlyMap的insertAll(iter)方法：在里面的迭代最终都会调用上面提到的ShuffleBlockFetcherIterator的next方法来获取数据。
  每次update&insert也会估算currentMap的大小，并判断是否需要溢写到磁盘文件，若需要则将map中的数据根据定义的keyComparator对key进行排序后返回一个迭代器，然后写到一个临时的磁盘文件，然后新建一个map来放新的数据。
  执行完combiners[ExternalAppendOnlyMap]的insertAll后，调用其iterator来返回一个代表一个完整partition数据（内存及spillFile）的迭代器。跟进ExternalIterator类的实例化：将currentMap中的数据经过排序后和spillFile数据的iterator组合在一起得到inputStreams ，迭代这个inputStreams ，将所有数据都保存在mergeHeadp中，在ExternalIterator方法的next()方法中将被访问到。
  最后若需要对数据进行全局的排序，则通过只有排序参数的ExternalSorter的insertAll方法来进行排序。最终返回一个指定partition所有数据的一个迭代器。