spark shuffle

435 阅读8分钟

shuffle

shuffle即对数据进行重组,并不是spark特有的知识点,之前也写过mapreduce中shuffle的处理流程可以查看:mapreduce之shuffle。不过spark中shuffle的处理更为复杂,但核心也都是数据在一个任务中处理完如何落地,另一个任务如何获取数据进行处理的过程。

spark shuffle

因为mapreduce会将任务分为map task和reduce task,之前关于mapreduce的shuffle中简单地将shuffle定义为map函数之后,reduce函数之前的操作。在spark中任务是根据ShuffleDependence切分成一个个stage执行的,shuffle可以简单理解为上一个stage处理完将数据落地以及下一个stage拉取上一个stage结果数据的过程。

在spark中将shuffle分为shuffle写和shuffle读,两个加起来就是完整的shuffle过程:

ShuffleWriter

ShuffleWriter负责将数据落地磁盘,在spark中将ShuffleWriter分为三种:UnsafeShuffleWriter、BypassMergeSortShuffleWriter和SortShuffleWriter。

spark是依据ShuffleDependency进行stage切分,在action算子中会new出ShuffleDependency,里面包含了ShuffleHandle属性,而ShuffleHandle分为BypassMergeSortShuffleHandle、SerializedShuffleHandle和BaseShuffleHandle,分别与三种ShuffleWriter一一对应。

获取ShuffleWriter整体的流程如下:

主要的代码为org.apache.spark.shuffle.sort.SortShuffleManager.registerShuffle():

  /**
   * Obtains a [[ShuffleHandle]] to pass to tasks.
   */
  override def registerShuffle[K, V, C](
      shuffleId: Int,
      numMaps: Int,
      dependency: ShuffleDependency[K, V, C]): ShuffleHandle = {
    if (SortShuffleWriter.shouldBypassMergeSort(conf, dependency)) {
      // If there are fewer than spark.shuffle.sort.bypassMergeThreshold partitions and we don't
      // need map-side aggregation, then write numPartitions files directly and just concatenate
      // them at the end. This avoids doing serialization and deserialization twice to merge
      // together the spilled files, which would happen with the normal code path. The downside is
      // having multiple files open at a time and thus more memory allocated to buffers.
      new BypassMergeSortShuffleHandle[K, V](
        shuffleId, numMaps, dependency.asInstanceOf[ShuffleDependency[K, V, V]])
    } else if (SortShuffleManager.canUseSerializedShuffle(dependency)) {
      // Otherwise, try to buffer map outputs in a serialized form, since this is more efficient:
      new SerializedShuffleHandle[K, V](
        shuffleId, numMaps, dependency.asInstanceOf[ShuffleDependency[K, V, V]])
    } else {
      // Otherwise, buffer map outputs in a deserialized form:
      new BaseShuffleHandle(shuffleId, numMaps, dependency)
    }
  }

org.apache.spark.shuffle.sort.SortShuffleManager.getWriter()

  /** Get a writer for a given partition. Called on executors by map tasks. */
  override def getWriter[K, V](
      handle: ShuffleHandle,
      mapId: Int,
      context: TaskContext): ShuffleWriter[K, V] = {
    numMapsForShuffle.putIfAbsent(
      handle.shuffleId, handle.asInstanceOf[BaseShuffleHandle[_, _, _]].numMaps)
    val env = SparkEnv.get
    handle match {
      case unsafeShuffleHandle: SerializedShuffleHandle[K @unchecked, V @unchecked] =>
        new UnsafeShuffleWriter(
          env.blockManager,
          shuffleBlockResolver.asInstanceOf[IndexShuffleBlockResolver],
          context.taskMemoryManager(),
          unsafeShuffleHandle,
          mapId,
          context,
          env.conf)
      case bypassMergeSortHandle: BypassMergeSortShuffleHandle[K @unchecked, V @unchecked] =>
        new BypassMergeSortShuffleWriter(
          env.blockManager,
          shuffleBlockResolver.asInstanceOf[IndexShuffleBlockResolver],
          bypassMergeSortHandle,
          mapId,
          context,
          env.conf)
      case other: BaseShuffleHandle[K @unchecked, V @unchecked, _] =>
        new SortShuffleWriter(shuffleBlockResolver, other, mapId, context)
    }
  }

BypassMergeSortShuffleWriter

ShuffleWariter是由ShuffleHandle模式匹配出来的,BypassMergeSortShuffleWriter对应的是BypassMergeSortShuffleHandle,而BypassMergeSortShuffleHandle需要满足两个条件:

  1. 没有mapSideCombine,如groupByKey;
  2. 下游stage任务数小于等于200(stage任务数由最后RDD的分区数决定),该数值由spark.shuffle.sort.bypassMergeThreshold控制,不配置默认200。

BypassMergeSortShuffleWriter的逻辑相对简单,它会为每一个分区创建一个临时文件写入该分区的数据,最终将所有文件合并成一个单独的文件,并记录一个索引文件,在索引文件中记录分区的起始偏移量和分区数据字节长度。

UnsafeShuffleWriter

UnsafeShuffleWriter对应ShuffleHandle是SerializedShuffleHandle,SerializedShuffleHandle需要满足三个条件:

  1. ShuffleDependency的序列号器(serializer)支持可寻址序列号,即supportsRelocationOfSerializedObjects;
  2. 没有定义聚合器(aggregator),reduceByKey也不行(没有map端聚合但是有reduce端聚合)
  3. 下游stage分区数小于1 << 2416777215

UnsafeShuffleWriter是通过jdk自带的sun.misc.Unsafe直接对堆内或者对外内存进行操控写入数据,至于是堆内还是堆外内存源码中由一个叫做钨丝计划的模式控制(tungstenMemoryMode),要想使用堆外内存需要满足三个条件:

  1. 开启堆外内存,由spark.memory.offHeap.enabled控制,默认为false;
  2. 堆外内存大于0,spark.memory.offHeap.size
  3. 运行spark的环境架构支持,linux的ppc64leppc64之类的,以及^(i[3-6]86|x86(_64)?|x64|amd64|aarch64)$

所以默认情况下UnsafeShuffleWriter用的是堆内内存进行数据存储。UnsafeShuffleWriter是以内存页的形式存储数据,内存页的大小计算比较复杂:

上述计算后再计算Math.min(PackedRecordPointer.MAXIMUM_PAGE_SIZE_BYTES, memoryManager.pageSizeBytes()MAXIMUM_PAGE_SIZE_BYTES为128M,正常情况下pageSize大小为64M,也可通过spark.buffer.pageSize进行控制。当内存超过一定限制的时候会进行溢写:

UnsafeShuffleWriter整体的数据写入流程如下:

上面的流程省略了很多细节得步骤,比如内存页申请最终是到TaskMemoryManager中申请,会涉及到spark内存管理的一些知识点,这里直接跳过。大概的流程步骤为:

  1. 在UnsafeShuffleWriter拿到一条数据,根据key计算哈希值;

  2. 将key-value数据以字节的形式写入到serBuffer中保存为recordBase,该buffer默认大小为1M;

  3. ShuffleExternalSorter将数据写入内存页,这个时候根据钨丝计划数据在堆内和堆外时写入有些许不同:

    1. 在堆内时,数据直接保存在字节数组中,因为UnSafe是直接操作内存,写入的时候要先跳过16字节的对象头;
    2. 写入数据的时候先写4字节的数据长度size,然后再写入recordBase数据,每次写入都会移动pageCursor到最新的位置;
    3. 堆外和堆内写入逻辑一致,只是因为是堆外直接申请的内存,没有对象头,数据就直接写入了。
  4. 数据写入内存页(MemoryBlock)后,会记录索引信息,索引信息包含分区号和记录地址:

    1. 分区号是long类型左移40位置,所以下游分区最大为2^24;

    2. 数据记录地址包含数据所在内存页编号和偏移量;

    3. 数据溢写的时候会对索引信息进行排序,保证分区间有序。

SortShuffleWriter

SortShuffleWriter对应的ShuffleHandle是BaseShuffleHandle,在前两种情况都不满足的条件下,走的都是SortShuffleWriter。在SortShuffleWriter根据是否有聚合器分为两种情况处理,整体的处理流程如下:

shouldCombine(定义了聚合器):

  1. AppendOnlyMap中定义了个引用类型的数组data,数组默认大小为64 * 2,能存64对key-value键值对,并把这个数组当成map来使用;
  2. 数据写入的时候会根据key生成分区号,在将分区号和key组成的tuple2对象组成新的key:k = (partition, key);
  3. 通过rehash(k.hashCode) & mask的方式获取偏移量(pos),mask为数组键值对容量-1key = array[pos],value = array[pos+1]
  4. 如果pos上数组为null说明第一次处理该数据,调用聚合器的第一条记录处理函数:newValue = updateFunc(false, null.asInstanceOf[V]),并将value写入到pos+1的位置上;
  5. 如果pos上数据不为空且和当前key相同,调用聚合器的后续记录处理函数:newValue = updateFunc(true, data(2 * pos + 1).asInstanceOf[V]),并将新的value更新到pos+1的位置上;
  6. 如果pos上数据不为空但和当前key不相同(哈希碰撞),采用线性探测的方式解决哈希碰撞;
  7. 当数组容量(key-value对)超过当前容量的70%的时候进行扩容,扩容系数为2,扩容后对key进行rehash。

!shouldCombine(没有定义聚合器):

  1. 在PartitionedPairBuffer中定义了个引用类型的数组data,数组默认大小为64 * 2,能存64对key-value键值对;
  2. 数据写入的时候会根据key生成分区号,在将分区号和key组成的tuple2对象组成新的key:k = (partition, key);
  3. 然后往data中追加记录key-value对即可,每记录一条pos自动+1;
  4. 当data容量满了的时候对data进行扩容,扩容系数为2,通过System.arraycopy方式扩容。

当申请不到内存或者data中数据大于内存门槛myMemoryThreshold时,会进行数据溢写,myMemoryThreshold初始值由spark.shuffle.spill.initialMemoryThreshold控制,默认为5M,计算中会不断变化:

在最后会将溢写的临时文件排序合并成一个文件,保证分区间有序。

ShuffleReader

spark只有一个ShuffleReader就是BlockStoreShuffleReader,负责处理所有情况的数据拉取。在BlockStoreShuffleReader会将需要拉取的数据分为远端的和本机的,将远端的数据通过NettyBlockTransferServer拉回到本地后封装成迭代器,再依据各种情况转换成不同的迭代器。整体流程如下:

  1. 在BlockStoreShuffleReader首先会去生成迭代器ShuffleBlockFetcherIterator;

  2. ShuffleBlockFetcherIterator会根据数据中的executorId对比当前executorId是否相同区分出远端的数据和本地的数据,然后将远端数据按照一定大小封装成一个个的remoteRequests,默认每个请求拉去48M/5大小的数据;

  3. 然后通过NettyBlockTransferServer服务将数据从远端拉到本地,将每个block的数据拆分成N个请求的好处是可以从不同的远端拉去部分数据,充分利用本机带宽;

  4. 远端数据拉去后和本地合成一个LinkedBlockingQueue[FetchResult]的result数据;

  5. 生成ShuffleBlockFetcherIterator后会根据sparkEnv中的各种信息,比如是否定义聚合器,是否定义key排序器等将迭代器转成最终的迭代器给到任务调用。

  6. spark中数据是以迭代器嵌套的方式进行处理的,最前的RDD往回迭代的时候就会调用compute方法获取到ShuffleReader中迭代器拉取数据处理,如org.apache.spark.rddShuffledRDD中:

      override def compute(split: Partition, context: TaskContext): Iterator[(K, C)] = {
        val dep = dependencies.head.asInstanceOf[ShuffleDependency[K, V, C]]
        // 获取getReader然后调用read
        SparkEnv.get.shuffleManager.getReader(dep.shuffleHandle, split.index, split.index + 1, context)
          .read()
          .asInstanceOf[Iterator[(K, C)]]
      }