Spark2.x精通:源码剖析UnsafeShuffleWriter具体实现

一、概述

    接上篇文章：Spark2.x精通:源码剖析BypassMergeSortShuffleWriter具体实现，这里将Spark Shuffle的第二种实现UnsafeShuffleWriter，这里回顾下触发条件：

    1).shuffle依赖不带有聚合(aggregation)操作

    2).支持序列化值的重新定位，即使用KryoSerializer或者SparkSQL自定义的一些序列化方式

    3).分区数量不能大于16777216个(2^24）

    由于UnsafeShuffleWriter 对应SortShuffle的tungsten-sort方式，排序的是二进制的数据，不会进数据进行反序列，所以不能进行聚合操作，另一方面PartitionId是占用24位的数,所以要小于16777216，这里不用去深究，知道就行。

二、结合源码剖析UnsafeShuffleWriter

  UnsafeShuffleWriter具体的实现也是先从对应类的write()函数中，我们直接看源码进行剖析

 

1.先看构造函数初始化

 public UnsafeShuffleWriter(      BlockManager blockManager,      IndexShuffleBlockResolver shuffleBlockResolver,      TaskMemoryManager memoryManager,      SerializedShuffleHandle<K, V> handle,      int mapId,      TaskContext taskContext,      SparkConf sparkConf) throws IOException {      //获取下游stage个数，也就是分区个数    final int numPartitions = handle.dependency().partitioner().numPartitions();    //这里判断Partition个数不能大于16777216个(2^24）    if (numPartitions > SortShuffleManager.MAX_SHUFFLE_OUTPUT_PARTITIONS_FOR_SERIALIZED_MODE()) {      throw new IllegalArgumentException(        "UnsafeShuffleWriter can only be used for shuffles with at most " +        SortShuffleManager.MAX_SHUFFLE_OUTPUT_PARTITIONS_FOR_SERIALIZED_MODE() +        " reduce partitions");    }    this.blockManager = blockManager;    this.shuffleBlockResolver = shuffleBlockResolver;    this.memoryManager = memoryManager;    this.mapId = mapId;    final ShuffleDependency<K, V, V> dep = handle.dependency();    this.shuffleId = dep.shuffleId();    this.serializer = dep.serializer().newInstance();    this.partitioner = dep.partitioner();    this.writeMetrics = taskContext.taskMetrics().shuffleWriteMetrics();    this.taskContext = taskContext;    this.sparkConf = sparkConf;    //是否采用NIO的从文件到文件流的复制方式，默认值是true 一般不用修改    this.transferToEnabled = sparkConf.getBoolean("spark.file.transferTo", true);    //Shuffle初始化的排序缓冲区 默认是4KB    this.initialSortBufferSize = sparkConf.getInt("spark.shuffle.sort.initialBufferSize",                                                  DEFAULT_INITIAL_SORT_BUFFER_SIZE);     //将数据写到磁盘文件之前,会先写入buffer缓冲中，待缓冲写满之后，才会溢写到磁盘    //由参数spark.shuffle.file.buffer 配置    this.inputBufferSizeInBytes =      (int) (long) sparkConf.get(package$.MODULE$.SHUFFLE_FILE_BUFFER_SIZE()) * 1024;    //默认大小为32k，大小由参数 spark.shuffle.unsafe.file.output.buffer 配置    this.outputBufferSizeInBytes =      (int) (long) sparkConf.get(package$.MODULE$.SHUFFLE_UNSAFE_FILE_OUTPUT_BUFFER_SIZE()) * 1024;    open();  }

    

2.再看write()函数,源码如下：

 @Override  public void write(scala.collection.Iterator<Product2<K, V>> records) throws IOException {    boolean success = false;    try {      //这个方法一共步：这里第一步，遍历每条数据，写入到sort中      //通过它来对数据进行分Partition之后，各个Partition之间的数据就是有序的了      while (records.hasNext()) {        insertRecordIntoSorter(records.next());      }     //将record进行排序，并在排序完成后写入磁盘文件作为spill file，     // 再将多个spill file合并成一个输出文件。      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);          }        }      }    }  }

3.直接上面插入sort的insertRecordIntoSorter()函数，代码如下：

@VisibleForTesting  void insertRecordIntoSorter(Product2<K, V> record) throws IOException {    assert(sorter != null);    final K key = record._1();    //根据Record对应key进行分区，如果使用的是HashPartitions，这会将key的HashCode    //与partition数目进行取模运行。来确定这条记录属于哪个分区    final int partitionId = partitioner.getPartition(key);    // 清空serBuffer的内存空间，这个内存空间用于存放字节输出流    serBuffer.reset();    //将键值序列化到serOutputStream，保存到BtyeArrayOutputStream底层的buf字段    serOutputStream.writeKey(key, OBJECT_CLASS_TAG);    serOutputStream.writeValue(record._2(), OBJECT_CLASS_TAG);    serOutputStream.flush();
    final int serializedRecordSize = serBuffer.size();    //如果serBuffer不为空，就说明里面是有shufflemaptask输出的序列化后的流数据，那么就会调用insertRecord函数，把序列化后的数据插入到sorter里面    assert (serializedRecordSize > 0);    //被写入到 sorter（ShuffleExternalSorter）中。在sorter中序列化数据被写入到内存中（内存不足会溢出到磁盘中），    // 其地址信息被写入到 ShuffleInMemorySorter 中    //serBuffer.getBuf() 就是获取上面写到buf中的数据    sorter.insertRecord(      serBuffer.getBuf(), Platform.BYTE_ARRAY_OFFSET, serializedRecordSize, partitionId);  }

   

4.我们看下sorter.insertRecord()函数，代码如下：

public void insertRecord(Object recordBase, long recordOffset, int length, int partitionId)    throws IOException {
    // for tests    assert(inMemSorter != null);     // 判断是否大于存储数据记录的上限值，由spark.shuffle.spill.numElementsForceSpillThreshold参数控制，          // 默认是 Long 的最大值，如果内存中的数据超过这个值则对当前内存数据进行排序并写入磁盘临时文件     // 假设可以源源不断的申请到内存，那么 Write 阶段的所有数据将一直保存在内存中    if (inMemSorter.numRecords() >= numElementsForSpillThreshold) {      logger.info("Spilling data because number of spilledRecords crossed the threshold " +        numElementsForSpillThreshold);      spill();    }    //如果没有达到存储的上限值，就判断一下数组指针是否有空间存下一个数据，    // 如果没有，则申请空间，申请不到就会发生溢写。    growPointerArrayIfNecessary();    // Need 4 bytes to store the record length.    final int required = length + 4;    //判断当前的内存页是否有足够的空间，存放下一个数据，如果不够则申请一个新的内存页    acquireNewPageIfNecessary(required);
    assert(currentPage != null);    final Object base = currentPage.getBaseObject();    final long recordAddress = taskMemoryManager.encodePageNumberAndOffset(currentPage, pageCursor);    Platform.putInt(base, pageCursor, length);    pageCursor += 4;     //把数据记录的长度和序列化后的字节序列写入    Platform.copyMemory(recordBase, recordOffset, base, pageCursor, length);    pageCursor += length;     //把这条记录转化为指针，传递到ShuffleinMemSorter中去。    inMemSorter.insertRecord(recordAddress, partitionId);  }

5.回到步骤2中继续向下看closeAndWriteOutput()函数，代码如下：

@VisibleForTesting  void closeAndWriteOutput() throws IOException {    assert(sorter != null);    updatePeakMemoryUsed();    serBuffer = null;    serOutputStream = null;    //将步骤3中写入ShuffleInMemorySorter的内存数据进行排序然后写入spill file    //然后返回spill fill的元数据信息 SpillInfo    final SpillInfo[] spills = sorter.closeAndGetSpills();    sorter = null;    final long[] partitionLengths;       //先根据shuffleID和mapID去构建一个结果文件对象output,    //最终的输出文件名称为："shuffle_" + shuffleId + "_" + mapId + "_" + reduceId    final File output = shuffleBlockResolver.getDataFile(shuffleId, mapId);    //把分区文件内容合并到tmp文件对象中，合并结束后再把tmp文件改成output的文件名    final File tmp = Utils.tempFileWith(output);    try {      try {       //合并所有的溢出文件到最终的shuffle文件，       // mergeSpills函数基于spill文件的数量和IO压缩编解码器选择最合适的合并策略。        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());          }        }      }      //更新shuffle索引文件      shuffleBlockResolver.writeIndexFileAndCommit(shuffleId, mapId, partitionLengths, tmp);    } finally {      if (tmp.exists() && !tmp.delete()) {        logger.error("Error while deleting temp file {}", tmp.getAbsolutePath());      }    }    //将Shuffle结果信息，封装到MapStatus返回    mapStatus = MapStatus$.MODULE$.apply(blockManager.shuffleServerId(), partitionLengths);  }

    总结一下：UnsafeShuffleWriter的代码流程，大体分为三步：

    第一步是ShuffleMapTask的数据输入到一个外部的排序器里面ShuffleExternalSorter；在每次排序比较的时候，只需要线性的查找指针区域的数据，不用根据指针去找真实的记录数据做比较，同时序列化器支持对二进制的数据进行排序比较，不会对数据进行反序列化操作，这样避免了反序列化和随机读取带来的开销，因为不会序列化成对象，可以减少内存的消耗和GC的开销；

        第二步是把内存中放不下的文件溢写到磁盘，在spill时，会根据指针的顺序溢写，这样就保证了每次溢写的文件都是根据Partition来进行排序的。一个文件里不同的partiton的数据用fileSegment来表示，对应的信息存在 SpillInfo 数据结构中；

        第三步是根据分区号去每个溢写文件中去拉取对应分区的数据，然后写入一个输出文件，最终合并成一个依据分区号全局有序的大文件。此外还会将每个partition的offset写入index文件方便reduce端拉取数据。

---

补充知识点：

 

    上面步骤3中涉及到两个概念：ShuffleInMemorySorter和ShuffleExternalSorter, UnsafeShuffleWriter内部使用了和BytesToBytesMap基本相同的数据结构处理map端的输出，不过将其细化为ShuffleExternalSorter和ShuffleInMemorySorter两部分，功能如下：

    ShuffleExternalSorter 使用MemoryBlock存储数据，每条记录包括长度信息和K-V Pair

ShuffleInMemorySorter 使用long数组存储每条记录对应的位置信息(page number + offset)，以及其对应的PartitionId。

    上面这个两个概念后续会写一篇文章单独讲解，这里知道有这么回事就可以了，不用细究。