Spark Shuffle源码分析系列之UnsafeShuffleWriter

最新推荐文章于 2024-07-10 17:25:40 发布

小白数据猿

最新推荐文章于 2024-07-10 17:25:40 发布

阅读量634

点赞数 1

分类专栏： Spark 文章标签：大数据 spark

本文链接：https://blog.csdn.net/lidongmeng0213/article/details/110633370

版权

前面我们介绍了BypassMergeSortShuffleWriter和SortShuffleWriter，知道了它们的应用场景和实现方式，本节我们来看下UnsafeShuffleWriter，它使用了Tungsten优化，排序的是二进制的数据，不会对数据进行反序列化操作，在某些情况下会加速Shuffle过程。

概述

UnsafeShuffleWriter提供了以序列化方式写入文件的方式，写入的内存既可以是on-heap也可以是off-heap，当数据到来时候，先序列化，然后使用序列化的数据排序，要想使用该Writer需要满足以下条件:

ShuffleDependency的序列化器需要支持对流中输出的序列化后的对象的字节进行排序；
ShuffleDependency不能指定指定聚合器，map阶段不进行聚合操作，用到了Tungsten优化，排序的是二进制的数据，不会对数据进行反序列化操作(反序列化就是转成Java对象)，所以不支持aggregation；
Shuffle过程产生的分区数大于16777216 (1 << 24)，内部存储的PackedRecordPointer 中8字节的前24bit表示分区数目，最大是16777216。

UnsafeShuffleWriter基于序列化的二进制数据进行排序，而不是基于java object，可以节省内存和GC开销，加速Shuffle过程，接下来我们来分析下具体的实现以及依赖的组件。

依赖

MemoryLocation

在tungsten-sort机制中，也存在数据缓存和溢写，这与sort shuffle是类似的。但是，这里不再借助像PartitionedPairBuffer之类的高级数据结构，而是由程序自己完成，并且是直接操作内存空间，而且可以使用堆内内存或者堆外内存，当使用堆外内存时候，由于不受JVM GC的影响，可以直接使用内存地址来进行寻址；但是使用堆内内存时候，内存地址可以由一个base对象和一个offset对象组合起来表示，但是对于一些数据结构比如在hashmap或者是sorting buffer中的记录的指针，尽管我们决定使用128位来寻址，我们不能只存base对象的地址，因为由于gc的存在，这个地址不能保证是稳定不变的。为了解决这个问题，Tungsten使用页表来管理内存，使用64位的高13位来保存内存页数，低51位来保存这个页中的offset，使用page表来保存base对象，其在page表中的索引就是该内存的内存页数。另外为了统一堆内堆外内存，统一抽象页为MemoryLocation，包含了obj对象和offset:

obj: 如果是堆内存模式时，数据作为对象存储在JVM的堆上，此时的obj不为空；处于堆外内存模式时，数据存储在JVM的堆外内存中，因而不会在JVM中存在对象；
offset属性主要用来定位数据，堆内存模式时，首先从堆内找到对象，然后使用offset定位数据的具体位置；堆外内存模式时，则直接使用offset从堆外内存中定位。

MemoryBlock继承自MemoryLocation，代表从obj和offset定位的起始位置开始，固定长度的连续内存块，是申请来的具体的内存空间，长度由length来决定，可以作为TaskMemoryManager页表中的一个页的抽象。

TaskMemoryManager

前面在讲解SortShuffleWriter时候，执行内存的申请与释放是通过TaskMemoryManager的acquireExecutionMemory和releaseExecutionMemory进行的，UnsafeShuffleWriter为了统一堆内和堆外内存，使用页表来管理内存的申请和释放，TaskMemoryManager提供的是allocatePage和freePage，并提供了一系列寻址的方法。

页表

TaskMemoryManager用页表来管理内存，维护了一个MemoryBlock数组用于存放该TMM分配得到的pages[pageTable]，逻辑地址用一个Long类型(64-bit)来表示，高13位来保存内存页数，低51位来保存这个页中的offset，使用page表来保存base对象，其在page表中的索引就是该内存的内存页数。页数最多有8192页，理论上允许索引 8192 * (2^31 -1)* 8 bytes，相当于140TB的数据。其中 2^31 -1 是整数的最大值，因为page表中记录索引的是一个long型数组，这个数组的最大长度是2^31 -1。实际上没有那么大。因为64位中除了用来设计页数和页内偏移量外还用于存放数据的分区信息。

/** 13位用来表示能存储的最大页数:8092 */
private static final int PAGE_NUMBER_BITS = 13;

/** 最大页数 8092*/
private static final int PAGE_TABLE_SIZE = 1 << PAGE_NUMBER_BITS;

/** 用于保存编码后的偏移量的位数。静态常量OFFSET_BITS的值为51。 */
static final int OFFSET_BITS = 64 - PAGE_NUMBER_BITS; 

/** 具体的内存存储是由一个数组存储，数组最大是2^32-1，数组元素是long类型,所以是8字节;
    所以得到最大的Page大小。静态常量MAXIMUM_PAGE_SIZE_BYTES的值为17179869176，即（2^32-1）× 8==17G。*/
public static final long MAXIMUM_PAGE_SIZE_BYTES = ((1L << 31) - 1) * 8L;

/** 长整型的低51位的位掩码。静态常量MASK_LONG_LOWER_51_BITS的值为2251799813685247 */
private static final long MASK_LONG_LOWER_51_BITS = 0x7FFFFFFFFFFFFL;

/* 维护一个Page表。pageTable实际为Page（即MemoryBlock）的数组，数组长度为PAGE_TABLE_SIZE。*/
private final MemoryBlock[] pageTable = new MemoryBlock[PAGE_TABLE_SIZE];

寻址

编码

将page和在page中的偏移量来编码为内存地址:

如果是堆外内存，由于是操作系统中的地址，所以直接是跟起始地址的偏移量；
如果是堆内内存，需要根据页号和偏移量进行组合得到。

public long encodePageNumberAndOffset(MemoryBlock page, long offsetInPage) {
   
  if (tungstenMemoryMode == MemoryMode.OFF_HEAP) {
    // Tungsten的内存模式是堆外内存
    // 此时的参数offsetInPage是操作系统内存的绝对地址，offsetInPage与MemoryBlock的起始地址之差就是相对于起始地址的偏移量
    offsetInPage -= page.getBaseOffset();
  }
  // 通过位运算将页号存储到64位长整型的高13位中，并将偏移量存储到64位长整型的低51位中，返回生成的64位的长整型。
  return encodePageNumberAndOffset(page.pageNumber, offsetInPage);
}

// 获取页号相对于内存块起始地址的偏移量
public static long encodePageNumberAndOffset(int pageNumber, long offsetInPage) {
   
  return (((long) pageNumber) << OFFSET_BITS) | (offsetInPage & MASK_LONG_LOWER_51_BITS);
}

解码

根据逻辑地址获取页号和偏移量，页号是8字节的13位，而偏移量是低51位。

// 用于解码页号，将64位的长整型右移51位（只剩下页号），然后转换为整型以获得Page的页号。
public static int decodePageNumber(long pagePlusOffsetAddress) {
   
  // 右移51位
  return (int) (pagePlusOffsetAddress >>> OFFSET_BITS);
}

// 解码偏移量，用于将64位的长整型与51位的掩码按位进行与运算，以获得在Page中的偏移量。
private static long decodeOffset(long pagePlusOffsetAddress) {
   
  // 与上MASK_LONG_LOWER_51_BITS掩码，即取pagePlusOffsetAddress的低51位
  return (pagePlusOffsetAddress & MASK_LONG_LOWER_51_BITS);
}

LongArray

LongArray本质是对MemoryBlock的封装，持有一块内存，它是通过ShuffleExternalSorter申请来的一个内存块MemoryBlock构建的，也就是申请一个页表，然后将内存空间按照8字节来进行切割，形成数组，所以能装下size / 8个元素，另外提供了添加数据以及获取数据的方法，数组里面的元素是8字节的PackedRecordPointer。

// 设置数据
public void set(int index, long value) {
   
  // 插入到baseObj的指定位置，baseObj就是MemoryBlock的obj
  Platform.putLong(baseObj, baseOffset + index * WIDTH, value);
}

// 访问数据
public long get(int index) {
   
  // 获取指定数据位位上的Long型数据
  return Platform.getLong(baseObj, baseOffset + index * WIDTH);
}

PackedRecordPointer

上面我们讲到LongArray的每个元素的大小为8字节，这个是因为数组内部存储的是记录指针PackedRecordPointer，它使用8字节来存储记录信息，主要分为三部分: [24 bit 分区号][13 bit 内存页号][27 bit 内存偏移量]，前面我们讲到TaskMemoryManager使用前13位表示内存页号，后51位表示偏移量，这里内存页号是足量的，但内存偏移量只有27位，如果offset in page大于2^27，就会出现高位丟失的问题。所以可寻址的最大偏移量的大小为2^27 bit，即128 MB，所以一个Task的可用内存 = 总页数 * 页的最大大小 = 2^13 * 2^27 = 2^40 = 1 TB字节.

初始化

recordPointer是taskMemoryManager返回的页号和相对于内存块起始地址的偏移量，partitionId是分区号，将他们两个按照上面的格式:[24 bit 分区号][13 bit 内存页号][27 bit 内存偏移量]组装起来。

首先取出来页号，recordPointer的高13位表示页号，由于页号前面24位是分区号的位置，所以需要右移24位；
提取偏移量，recordPointer的低27位是偏移量，将页号和偏移量(offset)组装起来；
最后前面24位是分区号，将分区号和页号偏移量组合起来，形成8字节的long。

// 低51位掩码
private static final long MASK_LONG_LOWER_51_BITS = (1L << 51) - 1;
// 高13位掩码
private static final long MASK_LONG_UPPER_13_BITS = ~MASK_LONG_LOWER_51_BITS;
// 低27位掩码，可以取到offset
private static final long MASK_LONG_LOWER_27_BITS = (1L << 27) - 1;

public static long packPointer(long recordPointer, int partitionId) {
   
  // 取高13位，即页号
  final long pageNumber = (recordPointer & MASK_LONG_UPPER_13_BITS) >>> 24;
  // 取低27位，获取offset,将页号和偏移量进行相或，组成低40位
  final long compressedAddress = pageNumber | (recordPointer & MASK_LONG_LOWER_27_BITS);
  // 将分区号右移40位，然后与页号和偏移量的40位进行相与，取得64位long型值
  return (((long) partitionId) << 40) | compressedAddress;
}

获取分区号

获取分区号，高24位表示分区号。

private static final long MASK_LONG_LOWER_40_BITS = (1L << 40) - 1;
private static final long MASK_LONG_UPPER_24_BITS = ~MASK_LONG_LOWER_40_BITS;

// 获取分区号
public int getPartitionId() {
   
  // 即取packedRecordPointer的高24位，然后右移去掉低40位即可
  return (int) ((packedRecordPointer & MASK_LONG_UPPER_24_BITS) >>> 40);
}

获取页号和偏移量

组装的格式:[24 bit 分区号][13 bit 内存页号][27 bit 内存偏移量]，提取页号和内存偏移量然后组成64位的long。

// 获取TaskMemoryManager需要的页号和偏移量
public long getRecordPointer() {
   
  // 左移24位，去掉分区号，然后取高13位即是页号
  final long pageNumber = (packedRecordPointer << 24) & MASK_LONG_UPPER_13_BITS;
  // 直接取低27位即是偏移量
  final long offsetInPage = packedRecordPointer & MASK_LONG_LOWER_27_BITS;
  // 将高13位和低27位组合成64位的long型整数
  return pageNumber | offsetInPage;
}

ShuffleInMemorySorter

ShuffleInMemorySorter是UnsafeShuffleWriter内存缓冲区使用的缓冲区和排序器，内部是由LongArray来存储数据，该类主要进行插入数据，存储数据，扩充Array以及返回排序好的数据迭代器等，实际存储的是数据的地址和分区组合结果，我们来分析下它是如何实现的。

私有变量

consumer内存消费者，这个实际是ShuffleExternalSorter，可以向TaskMemoryManager申请和释放内存；
array是数据内存存储的地方，内存可以是堆内或者堆外内存；
useRadixSort是否使用基数排序，如果使用基数排序，数组需要留一定的空间，否则使用TimSort进行排序，由参数spark.shuffle.sort.useRadixSort控制，默认是使用的；
pos表示当前array可以写入数据的下标；
usableCapacity是可用内存容量，当使用基数排序时为申请内存的1/2，否则为申请内存的1/1.5；
initialSize是最初的缓冲区大小，有参数spark.shuffle.sort.initialBufferSize控制，默认是4M；

private final MemoryConsumer consumer;
// 申请的内存的主要表现方式；排序器将操作该对象，代替直接操作记录
private LongArray array;

// 是否使用基数排序；基数排序比较快，但需要额外的内存。
private final boolean useRadixSort;

private int pos = 0;
private int usableCapacity = 0;

// 初始的排序缓冲大小
private int initialSize;

内存申请&释放

内存的申请与释放是通过MemoryConsumer来进行的，可以用来向TaskMemoryManager申请和释放内存，我们来看下申请长度为size的内存的步骤:

通过TaskMemoryManger来申请一个页；
如果分配不到页或者分配页的大小比较小，则需要释放申请到的页，抛出OOM异常；
如果分配到的页满足申请要求，就组装为LongArray返回给ShuffleInMemorySorter使用。

// org.apache.spark.memory.MemoryConsumer
public LongArray allocateArray(long size) {
   
  // 计算所需的Page大小。由于长整型占用8个字节，所以需要乘以8。
  long required = size * 8L;
  // 分配指定大小的MemoryBlock
  MemoryBlock page = taskMemoryManager.allocatePage(required, this);
  // 分配得到的MemoryBlock的大小小于所需的大小
  if (page == null || page.size() < required) {
   
    long got = 0;
    if (page != null) {
   
      got = page.size();
      // 释放MemoryBlock
      taskMemoryManager.freePage(page, this);
    }
    // 打印内存使用信息并抛出OutOf MemoryError。
    taskMemoryManager.showMemoryUsage();
    throw new OutOfMemoryError("Unable to acquire " + required + " bytes of memory, got " + got);
  }
  // 将required累加到used，即更新已经使用的内存大小
  used += required;
  // 创建并返回LongArray
  return new LongArray(page);
}

内存释放比较简单，直接释放相应的页即可。

// org.apache.spark.memory.MemoryConsumer
public void freeArray(LongArray array) {
   
  freePage(array.memoryBlock());
}

protected void freePage(MemoryBlock page) {
   
  // 首先更新used
  used -= page.size();
  // 释放MemoryBlock
  taskMemoryManager.freePage(page, this);
}

数据管理

insertRecord插入数据到内存缓冲区中，需要先检查是否还有足够的空间可以插入，然后将分区号和地址进行组装为PackedRecordPointer，放入到LongArray中。

// org.apache.spark.shuffle.sort.ShuffleInMemorySorter
// 判断是否还有剩余空间
public boolean hasSpaceForAnotherRecord() {
   
  return pos < usableCapacity;
}

public void insertRecord(long recordPointer, int partitionId) {
   
  // 检查是否还有空闲空间
  if (!hasSpaceForAnotherRecord()) {
   
    throw new IllegalStateException("There is no space for new record");
  }
  // 将转换后得到的PackedRecordPointer设置到array的pos位置
  array.set(pos, PackedRecordPointer.packPointer(recordPointer, partitionId));
  pos++;
}

// 获取记录数量
public int numRecords() {
   
  return pos;
}

当内存使用完后，需要进行扩容，ShuffleExternalSorter从TaskMemoryManager中申请内存页，并封装为LongArray，将现有的数据复制到新的数组中，释放老的数组内存即可。

// org.apache.spark.shuffle.sort.ShuffleInMemorySorter
// 对内存进行扩容  
public void expandPointerArray(LongArray newArray) {
   
  // 检查，扩容后的内存大小应该大于现有内存大小
  assert(newArray.size() > array.size());
  // 将当前array中的数据拷贝到newArray
  Platform.copyMemory(
    array.getBaseObject(),
    array.getBaseOffset(),
    newArray.getBaseObject(),
    newArray.getBaseOffset(),
    pos * 8L
  )

最低0.47元/天解锁文章