Spark ShuffleWriter之UnsafeShuffleWriter详解

最新推荐文章于 2024-09-13 15:19:32 发布

原创

最新推荐文章于 2024-09-13 15:19:32 发布 · 828 阅读

2 ·

CC 4.0 BY-SA版权

文章标签：

#Spark ShuffleWriter #UnsafeShuffleWriter #Spark shuffle

本文详细介绍了Spark Shuffle过程中UnsafeShuffleWriter的工作流程，包括将records缓存、排序、溢写到磁盘文件，以及如何合并文件和生成索引。还探讨了分区数限制的原因，即内存寻址的限制。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

前文提到，当序列化器支持relocation序列化对象，map端无需Combine，分区数不大于16777215 ((1 << 24) - 1)。这时候便匹配上UnsafeShuffleWriter将shuffle数据写出到磁盘。

过程

1、先将records插入到ShuffleExternalSorter缓存页中。

2、当插入内存的records达到shuffle内存限制或者记录已经全部插入到内存时，使用ShuffleInMemorySorter并依据记录的分区ID给records排序。

3、记录被溢出到多个磁盘临时文件（每次溢出都产生一个临时文件）。

4、合并多个磁盘文件为一个最终文件。

5、生成索引文件。

细节

1、插入过程

先将record序列化，放入到ByteArrayOutputStream的子类MyByteArrayOutputStream的buffer中。

void insertRecordIntoSorter(Product2<K, V> record) throws IOException {
   
   
  assert(sorter != null);
  final K key = record._1();
  final int partitionId = partitioner.getPartition(key);
  serBuffer.reset();
  // 序列化key及value到输出流MyByteArrayOutputStream
  serOutputStream.writeKey(key, OBJECT_CLASS_TAG);
  serOutputStream.writeValue(record._2(), OBJECT_CLASS_TAG);
  serOutputStream.flush();

  final int serializedRecordSize = serBuffer.size();
  assert (serializedRecordSize > 0);

  sorter.insertRecord(
    serBuffer.getBuf(), Platform.BYTE_ARRAY_OFFSET, serializedRecordSize, partitionId);
}

插入到shuffle sorter

public void insertRecord(Object recordBase, long recordOffset, int length, int partitionId)
  throws IOException {
   
   

  // for tests
  assert(inMemSorter != null);
  // 溢出条件中的大于设置门限值溢出，溢出到磁盘
  if (inMemSorter.numRecords() >= numElementsForSpillThreshold) {
   
   
    logger.info("Spilling data because number of spilledRecords crossed the threshold " +
      numElementsForSpillThreshold);
    spill();
  }

  // 这里是指针数组，当没有更多空间来存储指针时，内存数据也溢出到磁盘
  growPointerArrayIfNecessary();
  final int uaoSize = UnsafeAlignedOffset.getUaoSize();
  // Need 4 or 8 bytes to store the record length.
  final int required = length + uaoSize;
  // 如果没有足够空间容纳record，则尝试申请新的页内存
  acquireNewPageIfNecessary(required);

  assert(currentPage != null);
  final Object base = currentPage.getBaseObject();
  final long recordAddress = taskMemoryManager.encodePageNumberAndOffset(currentPage, pageCursor);
  UnsafeAlignedOffset.putSize(base, pageCursor, length);
  pageCursor += uaoSize;
  // 将记录拷贝到内存页中的对应位置（由页地址base + 起始偏移pageCursor + record长度length共同确定）
  Platform.copyMemory(recordBase, recordOffset, base, pageCursor, length);
  pageCursor += length;
  // 将record的地址和其对应的分区ID加入到指针数组里，供排序使用
  inMemSorter.insertRecord(recordAddress, partitionId);
}

溢写

public long spill(long size, MemoryConsumer trigger) throws IOException {
   
   
  if (trigger != this || inMemSorter == null || inMemSorter.numRecords() == 0) {
   
   
    return 0L;
  }
  // 排序并写出到磁盘
  writeSortedFile(false);
  // 溢出到磁盘之后，释放内存，重置指针数组
  final long spillSize = freeMemory();
  inMemSorter.reset();
  taskContext.taskMetrics().incMemoryBytesSpilled(spillSize);
  return spillSize;
}

排序并写出

排序并写出到文件，出于写出效率考虑，并不是一个一个记录写出，而是一次写出diskWriteBufferSize大小的记录。

private void writeSortedFile(boolean isLastFile) {
   
   

  final ShuffleWriteMetrics writeMetricsToUse;

  if (isLastFile) {
   
   
    writeMetricsToUse = writeMetrics;
  } else {
   
   
    writeMetricsToUse = new ShuffleWriteMetrics();
  }

  // getSortedIterator执行具体的排序逻辑，可以看看是怎么根据record指针和分区ID进行排序的
  final ShuffleInMemorySorter.ShuffleSorterIterator sortedRecords =
    inMemSorter.getSortedIterator();

  // 由于每次往磁盘写出很少量的数据，非常影响效率，因此用一个缓存buffer，一次写一批数据
  final byte[] writeBuffer = new byte[diskWriteBufferSize];
  
  // 写出磁盘的临时文件名称和块信息blockId的生成
  final Tuple2<TempShuffleBlockId, File> spilledFileInfo =
    blockManager.diskBlockManager().createTempShuffleBlock();
  final File file = spilledFileInfo._2();<