spark UnsafeShuffleWriter

spark中 shuffleWriter有

1 ByPassMergeSortShuffleWriter 跳过在内存中缓冲排序 直接向对应分区小文件中怼数据 最后线性合并的 这种适用与下游分区数量较少(分区太多,小文件过多影响性能) map端不需要聚合的
2 SortShuffleWriter 适用于多种情况 作为基本处理器 当下游分区过多 或需要map端聚合时使用 其内部的缓冲区分为buffer和map(在一个数组中使用线性探测存储map结构 在存储数据时就进行了 map聚合 对相同key的数据进行处理)两种 溢写的小文件有多个分区的数据 需要走归并
3 unsafeShuffleWriter使用堆外内存优化 当支持relocation序列化 不在map端聚合 下游分区数少于一定数量(2的24次方,这是因为框架用long做索引 long64位 前24位记录分区号 后面的记录数据地址) 满足上述条件可以使用

可以看到spark为了提速 相比于hadoop(数据和索引写入环形缓冲区,溢写前排序,归并等)有多种方案实现shuffleWrite进行优化 今天看较特殊的unsafeShuffleWriter:

unsafeShuffleWriter:

  @Override
  public void write(scala.collection.Iterator> records) throws IOException {
    // Keep track of success so we know if we encountered an exception
    // We do this rather than a standard try/catch/re-throw to handle
    // generic throwables.
    boolean success = false;
    try {
    
    //write方法中 对于每条记录执行insertRecordIntoSorter方法
      while (records.hasNext()) {
        insertRecordIntoSorter(records.next());
      }
     
  }
  
  //堆内字节数组做buffer
  private MyByteArrayOutputStream serBuffer;
  //输出流
  private SerializationStream serOutputStream;
  
  @Nullable private ShuffleExternalSorter sorter;
  
  serBuffer = new MyByteArrayOutputStream(DEFAULT_INITIAL_SER_BUFFER_SIZE);
  //将buffer 绑定输出流
  serOutputStream = serializer.serializeStream(serBuffer);
  
  @VisibleForTesting
  void insertRecordIntoSorter(Product2 record) throws IOException {
    assert(sorter != null);
    final K key = record._1();
    final int partitionId = partitioner.getPartition(key);
    serBuffer.reset();
    serOutputStream.writeKey(key, OBJECT_CLASS_TAG);
    
    //这里将对象转为字节数组 存入buffer中  buffer包装了一个字节数组
    serOutputStream.writeValue(record._2(), OBJECT_CLASS_TAG);
    serOutputStream.flush();

    //返回字节数组的长度
    final int serializedRecordSize = serBuffer.size();
    assert (serializedRecordSize > 0);

    //核心 每条记录插入sorter中
    sorter.insertRecord(
      serBuffer.getBuf(), Platform.BYTE_ARRAY_OFFSET, serializedRecordSize, partitionId);
  }  

ShuffleExternalSorter 核心的sorter

  //ShuffleExternalSorter中有ShuffleInMemorySorter 用于存储 排序记录索引
  //可以存储在堆内 可以存储在堆外
  @Nullable private ShuffleInMemorySorter inMemSorter; 

  //注意  
  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();
    // Need 4 bytes to store the record length.
    
    //字节数组的长度+4个字节 4字节用于描述这条数据的大小 总共需要申请length+4的空间 存储数据(字节数组)+数据大小作为一条记录
    final int required = length + 4;
    //每条记录来了判断是否需要申请新的page
    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;
    inMemSorter.insertRecord(recordAddress, partitionId);
  }     
  
  
   private void acquireNewPageIfNecessary(int required) {
   //如果当前page是空
     if (currentPage == null ||
     //如果页的游标+需要申请的数据空间大小 大于 页起始位置+页的大小 线性空间地址存储 数组下标>数组下标
       pageCursor + required > currentPage.getBaseOffset() + currentPage.size() ) {
       // TODO: try to find space in previous pages
       //分配新的page
       currentPage = allocatePage(required);
       pageCursor = currentPage.getBaseOffset();
       allocatedPages.add(currentPage);
     }
   } 

ShuffleExternalSorter extends MemoryConsume:

MemoryConsume:allocatePage调用父类的方法

   protected MemoryBlock allocatePage(long required) {
       MemoryBlock page = taskMemoryManager.allocatePage(Math.max(pageSize, required), this);
       if (page == null || page.size() < required) {
         throwOom(page, required);
       }
       used += page.size();
       return page;
   }

每个task中都有taskMemoryManager 对executor上的memoryManager封装调用 用于申请内存
也就是每个task线程共享了executor上的内存
task线程之间需要使用taskMemoryManager 互相协调分配内存

taskMemoryManager:

  //java的位图对象 直接操作位数组  
  private final BitSet allocatedPages = new BitSet(PAGE_TABLE_SIZE);  

  public MemoryBlock allocatePage(long size, MemoryConsumer consumer) {
    ...省略
        final int pageNumber;
        synchronized (this) {
              //nextClearBit 从传入的index开始寻找下一个为0的下标位置
              //用位图存储可用页的信息 返回为0的下标 可用页作为页的number
              pageNumber = allocatedPages.nextClearBit(0);
              //标记该页码被使用
              allocatedPages.set(pageNumber);
              
              //使用memoryManager中的钨丝计划 分配页
              MemoryBlock page = memoryManager.tungstenMemoryAllocator().allocate(acquired);
              page.pageNumber = pageNumber;
  }

总结

unsafeShuffleWriter的write方法:

首先unsafeShuffleWriter对象的属性 有一个堆上的字节数组作为buffer 和绑定该buffer的输出流

write方法对iterator迭代 对于每条记录 调用insertRecordIntoSorter

insertRecordIntoSorter内部 首先对传入的记录对象转为字节数组存入buffer中 之后就是核心的ShuffleExternalSorter中的insertRecord方法 由buffer传输字节数组数据

ShuffleExternalSorter对象中有ShuffleInMemorySorter 用于存储排序数据索引信息 可以存储在堆内 可以存储在堆外

insertRecord方法接受每条记录转为的字节数组对象作为参数 每条数据来的时候

growPointerArrayIfNecessary //判断是否需要 并扩容索引空间
final int required = length + 4; //所有的记录放入一个字节数组中 需要额外的4个字节记录该条记录的长度
//每条记录来了判断是否需要申请新的page  如果当前currentPage是null 或者当前记录的大小超过了 当前页的空间 新申请一个页
acquireNewPageIfNecessary(required);  //调用了父类 MemoryConsume:allocatePage  内部调用了taskMemoryManager.allocatePage

taskMemoryManager.allocatePage中

1 使用 BitSet java的位图存储可用页码信息   
2 使用 page = memoryManager.tungstenMemoryAllocator().allocate(acquired); 申请页空间
3 page.pageNumber = pageNumber; 设置页码 返回

分析

memoryManager.tungstenMemoryAllocator().allocate(acquired)

有两类实现: 堆上分配 堆外分配

返回的表示页的具体对象是MemoryBlock

public class MemoryBlock extends MemoryLocation {
    ...属性
    //数组  偏移量 长度
     public MemoryBlock(@Nullable Object obj, long offset, long length) {
        super(obj, offset);
        this.length = length;
      }
}

public class MemoryLocation {

  @Nullable
  Object obj;

  long offset;

  public MemoryLocation(@Nullable Object obj, long offset) {
    this.obj = obj;
    this.offset = offset;
  }
} 

Platform是由spark包装
底层使用unsafe类操作堆外内存(自己如果要使用堆外内存 可以仿写platForm类 模板写法)

HeapMemoryAllocator 堆上分配

 @Override
 public MemoryBlock allocate(long size) throws OutOfMemoryError {
        ...
    //堆上开辟一个long类型的数组 数组长度最少为 1  此时是一个空数组
    //size是字节数组的长度 一个long类型占用了8个字节 计算总共需要多少个long数组的空间
    long[] array = new long[(int) ((size + 7) / 8)];
    //new了的一个MemoryBlock 是表示页的具体对象  内部包括了空数组 第一个元素的偏移量 数据长度(原数据长度+4[记录数据长度的空间])
    MemoryBlock memory = new MemoryBlock(array, Platform.LONG_ARRAY_OFFSET, size);
    if (MemoryAllocator.MEMORY_DEBUG_FILL_ENABLED) {
      memory.fill(MemoryAllocator.MEMORY_DEBUG_FILL_CLEAN_VALUE);
    }
    return memory;
}

Platform:

 //使用unsafe的方法 获取到了
 //long[]实际是一个对象 该对象占用了一定宽度空间   
 //long[] 第一个元素在该对象中有一个偏移量(前面存储对象头等信息)
 //这个偏移量是16个字节
 LONG_ARRAY_OFFSET = _UNSAFE.arrayBaseOffset(long[].class); 

向memoryBlock传入一个字节数组 第一个元素的偏移量位置 数据长度 以及页码

ShuffleExternalSorter:中

 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 默认值为-1 游标存储页的BaseOffset对象  第一个对象的偏移量
      pageCursor = currentPage.getBaseOffset();
      allocatedPages.add(currentPage);
    }
  }

ShuffleExternalSorter:

//游标指向16
pageCursor = currentPage.getBaseOffset();


//insertRecord方法中:

//返回了 字节数组
final Object base = currentPage.getBaseObject();
//当前内存页对象 和 游标  0,16
final long recordAddress = taskMemoryManager.encodePageNumberAndOffset(currentPage, pageCursor);

//使用unsafe对象  传入字节数组 页游标16 和 该条数据长度
//前16位存储 数组对象的对象头
//这个就是直接向base 游标后的 的4个空间 存入本条数据的长度
Platform.putInt(base, pageCursor, length);
//游标向后加4
pageCursor += 4;


//使用unsafe对象 拷贝数据  recordBase数据数组 recordOffset去掉对象头
//base目标数组  pageCursor目标数组游标起始复制位置 数据length长度
Platform.copyMemory(recordBase, recordOffset, base, pageCursor, length);
pageCursor += length;
//索引信息放入inMemorySorter
inMemSorter.insertRecord(recordAddress, partitionId);

taskMemoryManager.encodePageNumberAndOffset:

public long encodePageNumberAndOffset(MemoryBlock page, long offsetInPage) {
    ...   
    //返回页码 和 页游标 
 return encodePageNumberAndOffset(page.pageNumber, offsetInPage); 
}  

堆外分配

page = memoryManager.tungstenMemoryAllocator().allocate(acquired):

UnsafeMemoryAllocator:堆外分配器

  @Override
  public MemoryBlock allocate(long size) throws OutOfMemoryError {
    //使用unsafe直接分配堆外内存串 返回一个地址
    //堆外内存 jvm不能管理(不能用jvm的c代码管理)
    //java进程可以访问 相当于开了个后门 绕过kvm虚拟映射 直接操作 os可以使用的地址 
    //速度更快 没有对象包装 更省空间
    long address = Platform.allocateMemory(size);
    //空堆外内存 传入null 堆外地址 和 堆外申请的空间长度
    //堆外字节数组 是纯字节数组没有对象头的
    MemoryBlock memory = new MemoryBlock(null, address, size);
    if (MemoryAllocator.MEMORY_DEBUG_FILL_ENABLED) {
      memory.fill(MemoryAllocator.MEMORY_DEBUG_FILL_CLEAN_VALUE);
    }
    return memory;
  }
  
  ..设置页码等

ShuffleExternalSorter:

//游标指向堆外地址
pageCursor = currentPage.getBaseOffset();

 //insertRecord方法中:
    
//返回了null
final Object base = currentPage.getBaseObject();
//当前内存页对象 和 物理地址  0,address
final long recordAddress = taskMemoryManager.encodePageNumberAndOffset(currentPage, pageCursor);

//base为null 直接向物理地址填充4个字节存储 数据长度
Platform.putInt(base, pageCursor, length);
//游标向后加4
pageCursor += 4;

//base为null 根据游标 直接存储到物理地址上
Platform.copyMemory(recordBase, recordOffset, base, pageCursor, length);
pageCursor += length;
//索引信息放入inMemorySorter
inMemSorter.insertRecord(recordAddress, partitionId);

堆内堆外数据存储的方式是一样的.. 但是效果有很大区别 具体可以去了解下堆外内存的好处..

1 spark这里使用堆外内存的优化在很多地方都可以见到

2 unsafe类支持了许多技术的操作 cas,堆外内存使用等

你可能感兴趣的:(spark UnsafeShuffleWriter)