Spark 技术 | Spark Shuffle 解读

// BypassMergeSortShuffleWriter.java
// Step 1 : 为每个分区创建一个分区文件写
partitionWriters = new DiskBlockObjectWriter[numPartitions];    // 对象写文件工具
partitionWriterSegments = new FileSegment[numPartitions];       // 文件元信息

for (int i = 0; i < numPartitions; i++) {
  // 对于每个分区， 创建一个临时文件用来存放shuffle数据
  final Tuple2<TempShuffleBlockId, File> tempShuffleBlockIdPlusFile = blockManager.diskBlockManager().createTempShuffleBlock();
  final File file = tempShuffleBlockIdPlusFile._2();
  final BlockId blockId = tempShuffleBlockIdPlusFile._1();

  // 通过 BlockManager 封装了一个磁盘写入工具 [[org.apache.spark.storage.DiskBlockObjectWriter]]
  partitionWriters[i] = blockManager.getDiskWriter(blockId, file, serInstance, fileBufferSize, writeMetrics);
}

// BypassMergeSortShuffleWriter.java
// Step 2 : 遍历迭代数据记录，将记录内容写出
while (records.hasNext()) {
  final Product2<K, V> record = records.next();
  final K key = record._1();
  final int partition = partitioner.getPartition(key);    // 定位分区
  // 找到对应的写文件， 数据写出
  partitionWriters[partition].write(key, record._2());
}

// 将每个文件都提交上
for (int i = 0; i < numPartitions; i++) {
  final DiskBlockObjectWriter writer = partitionWriters[i];
  partitionWriterSegments[i] = writer.commitAndGet();
  writer.close();
}

// BypassMergeSortShuffleWriter.java
// Step 3 : 合并临时文件，生成目标文件
File output = shuffleBlockResolver.getDataFile(shuffleId, mapId);   // 创建shuffle 写出文件 : shuffle_{shuffleId}_{mapId}_0.data
File tmp = Utils.tempFileWith(output);                              // {output}.xxxxx -- 临时文件
partitionLengths = writePartitionedFile(tmp);                     // 将每个分区文件合并写到目标文件， 并返回每个分区的文件长度
// 提交 Shuffle 数据
shuffleBlockResolver.writeIndexFileAndCommit(shuffleId, mapId, partitionLengths, tmp);

/**
 * 将所有分区文件连接成一个单独的组合文件。
 * @return 每个分区文件的长度数组，以字节为单位（由map输出跟踪器使用）。
 */
private long[] writePartitionedFile(File outputFile) throws IOException {
  ...
  final FileOutputStream out = new FileOutputStream(outputFile, true);  // 创建增量输出流
  
  for (int i = 0; i < numPartitions; i++) {
    final File file = partitionWriterSegments[i].file();               // 拿到每个分区的写出文件
    lengths[i] = Utils.copyStream(in, out, false, transferToEnabled);  // 将分区文件数据拷贝到目标文件中
  }
  ...
}

// UnsafeShuffleWriter.insertRecordIntoSorter

// STEP 1 : 将数据序列化
serBuffer.reset();
serOutputStream.writeKey(key, OBJECT_CLASS_TAG);
serOutputStream.writeValue(record._2(), OBJECT_CLASS_TAG);
serOutputStream.flush();

final int serializedRecordSize = serBuffer.size();

// STEP 2 : 将序列化后的对象写入 ShuffleExternalSorter
sorter.insertRecord(serBuffer.getBuf(), Platform.BYTE_ARRAY_OFFSET, serializedRecordSize, partitionId);

// ShuffleExternalSorter.insertRecord

private void acquireNewPageIfNecessary(int required) {
// 如果当前没有分配页，或者分配的内存页剩余空间不足，则分配一个新的页
	if (currentPage == null ||
	  pageCursor + required > currentPage.getBaseOffset() + currentPage.size() ) {
	  // TODO: try to find space in previous pages
	  currentPage = allocatePage(required);
	  pageCursor = currentPage.getBaseOffset();   // 拿到页内偏移
	  allocatedPages.add(currentPage);
	}
}
// 需求内存总大小 = 序列化对象长度 + UAO 对象大小
final int required = length + uaoSize;

// 检查内存页是否满足条件，不满足条件则准备一个新的内存页
acquireNewPageIfNecessary(required);

// 拿到底层的存储对象 long[]
final Object base = currentPage.getBaseObject();    

// 计算数据要写入的位置(用户后期数据定位): 页号左移51位 | 页内偏移低51位
final long recordAddress = taskMemoryManager.encodePageNumberAndOffset(currentPage, pageCursor);

// TODO : 将数据写入 currentPage 对应的块中:  [ len(recordBase) ] [recordBase]
// 写入序列化数据的长度
UnsafeAlignedOffset.putSize(base, pageCursor, length);
pageCursor += uaoSize; // 更新偏移位置
Platform.copyMemory(recordBase, recordOffset, base, pageCursor, length);
pageCursor += length;

// inMemSorter 记录数据 : 数据位置， 数据所在分区
inMemSorter.insertRecord(recordAddress, partitionId);

// ShuffleInMemorySorter.java
public static long packPointer(long recordPointer, int partitionId) {
  final long pageNumber = (recordPointer & MASK_LONG_UPPER_13_BITS) >>> 24;
  final long compressedAddress = pageNumber | (recordPointer & MASK_LONG_LOWER_27_BITS);
  // 高40位位分区号 | 第 40 位是压缩地址
  return (((long) partitionId) << 40) | compressedAddress;
}
public void insertRecord(long recordPointer, int partitionId) {  
  // 对象指针跟分区号的封装结构，
  array.set(pos, PackedRecordPointer.packPointer(recordPointer, partitionId));
  pos++;
}

// ShuffleExternalSorter.writeSortedFile

// 将 inMemSorter 排序并返回一个迭代器
final ShuffleInMemorySorter.ShuffleSorterIterator sortedRecords = inMemSorter.getSortedIterator();

// 对象拷贝到writeBuffer用于写磁盘
final byte[] writeBuffer = new byte[diskWriteBufferSize];

// 因为这个输出将在shuffle期间被读取，所以它的压缩编解码器必须由spark.shuffle.compress控制
// 而不是由spark.shuffle.spill.compress控制，因此我们需要在这里使用createTempShuffleBlock
// 更多细节请参见SPARK-3426。
final Tuple2<TempShuffleBlockId, File> spilledFileInfo = blockManager.diskBlockManager().createTempShuffleBlock();
final File file = spilledFileInfo._2();
final TempShuffleBlockId blockId = spilledFileInfo._1();

final SpillInfo spillInfo = new SpillInfo(numPartitions, file, blockId);
final DiskBlockObjectWriter writer = blockManager.getDiskWriter(blockId, file, ser, fileBufferSizeBytes, writeMetricsToUse);

// 遍历有序数据集
while (sortedRecords.hasNext()) {
  sortedRecords.loadNext();   // 获取迭代数据
  final int partition = sortedRecords.packedRecordPointer.getPartitionId();    // 获取数据对应分区

  if (partition != currentPartition) {                                       // 如果切换新的分区，则记录上一个分区的数据长度
    if (currentPartition != -1) {
      final FileSegment fileSegment = writer.commitAndGet();
      spillInfo.partitionLengths[currentPartition] = fileSegment.length();  // 记录每个分区的长度
    }
    currentPartition = partition;
  }
    
  // 取出序列化对象数据，并写到文件中
  final long recordPointer = sortedRecords.packedRecordPointer.getRecordPointer();    // 拿到数据索引
  final Object recordPage = taskMemoryManager.getPage(recordPointer);
  final long recordOffsetInPage = taskMemoryManager.getOffsetInPage(recordPointer);
  int dataRemaining = UnsafeAlignedOffset.getSize(recordPage, recordOffsetInPage);
  long recordReadPosition = recordOffsetInPage + uaoSize; // skip over record length
    
  while (dataRemaining > 0) {
    final int toTransfer = Math.min(diskWriteBufferSize, dataRemaining);
    Platform.copyMemory(recordPage, recordReadPosition, writeBuffer, Platform.BYTE_ARRAY_OFFSET, toTransfer);

    writer.write(writeBuffer, 0, toTransfer);
    recordReadPosition += toTransfer;
    dataRemaining -= toTransfer;
  }
   // 记录数据写
  writer.recordWritten();
}
final FileSegment committedSegment = writer.commitAndGet();
writer.close();
if (currentPartition != -1) {
  spillInfo.partitionLengths[currentPartition] = committedSegment.length();
  spills.add(spillInfo);
}

// 使用计数输出流，以避免在每个分区写入后关闭底层文件并向文件系统请求其大小。
final CountingOutputStream mergedFileOutputStream = new CountingOutputStream(bos);

// 遍历所有的分区
for (int partition = 0; partition < numPartitions; partition++) {

    // 基于计数器的写流
    OutputStream partitionOutput = new CloseAndFlushShieldOutputStream(new TimeTrackingOutputStream(writeMetrics, mergedFileOutputStream));
	partitionOutput = blockManager.serializerManager().wrapForEncryption(partitionOutput);

    if (compressionCodec != null) {   // 判断是否要压缩写出
      partitionOutput = compressionCodec.compressedOutputStream(partitionOutput);
    }

    for (int i = 0; i < spills.length; i++) {     // 遍历所有的Spill文件
      final long partitionLengthInSpill = spills[i].partitionLengths[partition];   // 定位当前当前spill中当前分区的数据位置
      if (partitionLengthInSpill > 0) {
        InputStream partitionInputStream = new LimitedInputStream(spillInputStreams[i], partitionLengthInSpill, false);
        try {
          partitionInputStream = blockManager.serializerManager().wrapForEncryption(partitionInputStream);
          if (compressionCodec != null) {
            partitionInputStream = compressionCodec.compressedInputStream(partitionInputStream);
          }
          // 将SpillFile的当前分区写出到目标文件
          ByteStreams.copy(partitionInputStream, partitionOutput);
        } finally {
          partitionInputStream.close();
        }
      }
    }
	partitionOutput.flush();
	partitionOutput.close();

    // 记录每个分区的数据长度
    partitionLengths[partition] = (mergedFileOutputStream.getByteCount() - initialFileLength);
}

// SortShuffleWriter.write()

sorter = if (dep.mapSideCombine) {
  // 如果需要combine 操作，则开启写出聚合操作
  new ExternalSorter[K, V, C](context, dep.aggregator, Some(dep.partitioner), dep.keyOrdering, dep.serializer)
} else {
  // 如果不需要combine操作
  // 在这种情况下，我们既不向排序器传递聚合器，也不向排序器传递排序，因为我们不关心每个分区中的键是否已排序；
  // 这将在reduce方面完成
  new ExternalSorter[K, V, V](context, aggregator = None, Some(dep.partitioner), ordering = None, dep.serializer)
}

sorter.insertAll(records)

// ExternalSorter.insertAll()
while (records.hasNext) {
  val kv = records.next()   
  buffer.insert(getPartition(kv._1), kv._1, kv._2.asInstanceOf[C])  // 数据组织成 // (partitionId, recordKey) -> record 写入 PartitionedPairBuffer
  maybeSpillCollection(usingMap = false)                            // 判断内存空间是否足够，是否需要将数据刷到磁盘  
}

// PartitionedPairBuffer.java

def partitionComparator[K]: Comparator[(Int, K)] = new Comparator[(Int, K)] {
  override def compare(a: (Int, K), b: (Int, K)): Int = {
    a._1 - b._1  // partitionId
  }
}
// 如果没有开启combine操作， 则partitionKeyComparator一直是None
// 参照 SortShuffleWriter.write() 中ExternalSort 构造过程
val comparator = keyComparator.map(partitionKeyComparator).getOrElse(partitionComparator)  // 返回 partitionComparator 比较器

// ExternalSorter.scala
// 数据spill
override protected[this] def spill(collection: WritablePartitionedPairCollection[K, C]): Unit = {
  val inMemoryIterator = collection.destructiveSortedWritablePartitionedIterator(comparator)   // 对数据按照分区进行排序
  val spillFile = spillMemoryIteratorToDisk(inMemoryIterator)                                  // 将内存数据写到磁盘，形成 SpillFile 文件
  spills += spillFile                                                                          // 记录所有的 SpillFile 文件集合
}

// ExternalSorter.java
def insertAll(records: Iterator[Product2[K, V]]): Unit = {
  // 如果需要进行与合并， 则直接在内存中进行合并处理
  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) {
    addElementsRead()    // 记录数据读取
    kv = records.next()  // 提取出 kv 数据记录, key -> (partitinId, recordKey)
    map.changeValue((getPartition(kv._1), kv._1), update)    // 插入的同时完成合并工作
    maybeSpillCollection(usingMap = true)
  }
}

// AppendOnlyMap.java
var pos = rehash(key.hashCode) & mask
var i = 1
while (true) {
  val curKey = data(2 * pos)
  if (curKey.eq(null)) {
    data(2 * pos) = k
    data(2 * pos + 1) = value.asInstanceOf[AnyRef]
    incrementSize()  // Since we added a new key
    return
  } else if (k.eq(curKey) || k.equals(curKey)) {
    data(2 * pos + 1) = value.asInstanceOf[AnyRef]
    return
  } else {
    val delta = i
    pos = (pos + delta) & mask
    i += 1
  }
}