源码--spark2.X的shuffle-CFANZ编程社区

一、源码流程

1、首先sparkContext被实例化的时候,会创建sparkEnv,在sparkEnv的create方法中对shuffleManager进行了初始化

val shortShuffleMgrNames = Map(
      "sort" -> classOf[org.apache.spark.shuffle.sort.SortShuffleManager].getName,
      "tungsten-sort" -> classOf[org.apache.spark.shuffle.sort.SortShuffleManager].getName)
val shuffleMgrName = conf.get("spark.shuffle.manager", "sort")
val shuffleMgrClass =
      shortShuffleMgrNames.getOrElse(shuffleMgrName.toLowerCase(Locale.ROOT), shuffleMgrName)
//这句代码其实创建了ShuffleManager的子类SortShuffleManager
val shuffleManager = instantiateClass[ShuffleManager](shuffleMgrClass)

2、创建SortShuffleManager

这个类比较重要的方法

2.1、registerShuffle,作用是注册不同的handle,然后运行不同的shuffle机制;

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.
      //如果partitions的数目比spark.shuffle.sort.bypassMergeThreshold(默认200)参数设置的阈值少,
      //并且我们也不需要做map端的聚合操作(跟写代码的时候调用的算子有关系),
      //然后就直接写numPartitions个文件,最后把这些文件合并成一个
      //(归并排序-->归并排序,并不是原地排序算法,最后一次合并会占用大约2倍的空间,这点要注意,在分析Shuffle使用内存的时候要考虑到),
    //可以避免在合并溢写文件的时候进行两次的序列化和反序列化(?),正常的代码路径会反正这个情况(BaseShuffleHandle)
    //确定是一次会打开多个文件,因此需要分配给buffer更多的内存
      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)
    }
  }

private[spark] object SortShuffleWriter {
  def shouldBypassMergeSort(conf: SparkConf, dep: ShuffleDependency[_, _, _]): Boolean = {
    // We cannot bypass sorting if we need to do map-side aggregation.
    if (dep.mapSideCombine) {
      false
    } else {
      val bypassMergeThreshold: Int = conf.getInt("spark.shuffle.sort.bypassMergeThreshold", 200)
      dep.partitioner.numPartitions <= bypassMergeThreshold
    }
  }
}

def canUseSerializedShuffle(dependency: ShuffleDependency[_, _, _]): Boolean = {
    val shufId = dependency.shuffleId
    val numPartitions = dependency.partitioner.numPartitions
    if (!dependency.serializer.supportsRelocationOfSerializedObjects) {
      log.debug(s"Can't use serialized shuffle for shuffle $shufId because the serializer, " +
        s"${dependency.serializer.getClass.getName}, does not support object relocation")
      false
    } else if (dependency.mapSideCombine) {
      log.debug(s"Can't use serialized shuffle for shuffle $shufId because we need to do " +
        s"map-side aggregation")
      false
    } else if (numPartitions > MAX_SHUFFLE_OUTPUT_PARTITIONS_FOR_SERIALIZED_MODE) {
      log.debug(s"Can't use serialized shuffle for shuffle $shufId because it has more than " +
        s"$MAX_SHUFFLE_OUTPUT_PARTITIONS_FOR_SERIALIZED_MODE partitions")
      false
    } else {
      log.debug(s"Can use serialized shuffle for shuffle $shufId")
      true
    }
  }

2.2、getReader

当我们写的shuffle类的算子,在reduce阶段,读shuffle数据的时候调用这个方法来读

override def getReader[K, C](handle: ShuffleHandle, startPartition: Int, endPartition: Int, context: TaskContext): ShuffleReader[K, C] = {
        new BlockStoreShuffleReader(handle.asInstanceOf[BaseShuffleHandle[K, _, C]], startPartition, endPartition, context)
    }

2.3、getWriter

当我们写的shuffle类的算子,在map阶段,写shuffle数据的时候调用这个方法来写,其实是调用的Writer方法

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 {
            //在不需要map端聚合、partition数量小于16777216，Serializer支持relocation的情况下
            case unsafeShuffleHandle: SerializedShuffleHandle[K@unchecked, V@unchecked] => new UnsafeShuffleWriter(env.blockManager,
                shuffleBlockResolver.asInstanceOf[IndexShuffleBlockResolver], context.taskMemoryManager(), unsafeShuffleHandle, mapId, context,
                env.conf)
            
            // 在不需要map端聚合、partition数量小于200的情况返回BypassMergeSortShuffleHandle对象
            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)
        }
    }