Spark ShuffleMap任务的生成、执行及数据跟踪流程
RDD抽象类的定义
我们知道Spark中的一个可以用户计算的数据集，被抽象成了一个RDD，如下是RDD的类定义（这里只保留了类中关键的成员变量）：
abstract class RDD[T: ClassTag](
@transient private var _sc: SparkContext,
@transient private var deps: Seq[Dependency[_]]
) extends Serializable with Logging {
/** A unique ID for this RDD (within its SparkContext). */
val id: Int = sc.newRddId()
/** A friendly name for this RDD */
@transient var name: String = _
// Our dependencies and partitions will be gotten by calling subclass's methods below, and will
// be overwritten when we're checkpointed
private var dependencies_ : Seq[Dependency[_]] = _
@transient private var partitions_ : Array[Partition] = _
private var storageLevel: StorageLevel = StorageLevel.NONE
private[spark] var checkpointData: Option[RDDCheckpointData[T]] = None
private val checkpointAllMarkedAncestors =
Option(sc.getLocalProperty(RDD.CHECKPOINT_ALL_MARKED_ANCESTORS)).exists(_.toBoolean)
// From performance concern, cache the value to avoid repeatedly compute `isBarrier()` on a long
// RDD chain.
@transient protected lazy val isBarrier_ : Boolean =
dependencies.filter(!_.isInstanceOf[ShuffleDependency[_, _, _]]).exists(_.rdd.isBarrier())
private[spark] final lazy val outputDeterministicLevel: DeterministicLevel.Value = {
if (isReliablyCheckpointed) {
DeterministicLevel.DETERMINATE
} else {
getOutputDeterministicLevel
}
}
}
从上面的代码可以看到RDD在序列化时，仅仅会序列化id、dependencies_、storageLevel、checkpointAllMarkedAncestors、outputDeterministicLevel这些成员变量，当然还包括RD实现类的其它可以被序列化的成员变量。
ShuffledRDD创建
当我们通过RDD API调用带有Shuffle动作的方法时，就会生成这样的实例，partitionBy(…)方法的定义如下：
/**
* Return a copy of the RDD partitioned using the specified partitioner.
*/
def partitionBy(partitioner: Partitioner): RDD[(K, V)] = self.withScope {
if (keyClass.isArray && partitioner.isInstanceOf[HashPartitioner]) {
throw new SparkException("HashPartitioner cannot partition array keys.")
}
if (self.partitioner == Some(partitioner)) {
self
} else {
new ShuffledRDD[K, V, V](self, partitioner)
}
}
ShuffledRDD是整个RDD依赖图（一个有向图）中的一个结点，因此在下面的类定义中，可以看到ShuffledRDD会有一个成员变量prev，它就是前驱RDD的引用，同时注意这个变量是@transient类型的，因此它不能被序列化。
Spark会根据RDD依赖图，生成一系列的任务，对于ShuffledRDD来说，就会对应一个ShuffleMapTask（后面篇章详细说明），通常需要跨结点进行数据交换。
为了在创建任务时，方便解析RDD的依赖关系，RDD抽象类提供了getDenpendencies方法，用于生成RDD的依赖关系类ShuffleDependency，这个关系除了记录当前RDD以及前驱RDD，还记录了当前RDD的数据分区信息part、数据序列化器serializer、数据排序信息keyOrdering、聚合函数aggregator、合并函数mapSideCombine。
class ShuffledRDD[K: ClassTag, V: ClassTag, C: ClassTag](
@transient var prev: RDD[_ <: Product2[K, V]],
part: Partitioner)
extends RDD[(K, C)](prev.context, Nil) {
override def getDependencies: Seq[Dependency[_]] = {
val serializer = userSpecifiedSerializer.getOrElse {
val serializerManager = SparkEnv.get.serializerManager
if (mapSideCombine) {
serializerManager.getSerializer(implicitly[ClassTag[K]], implicitly[ClassTag[C]])
} else {
serializerManager.getSerializer(implicitly[ClassTag[K]], implicitly[ClassTag[V]])
}
}
List(new ShuffleDependency(prev, part, serializer, keyOrdering, aggregator, mapSideCombine))
}
}
ShuffleDependency创建
ShuffleDependency是对一个RDD及其直接父继RDD的关系的描述，下面的代码展示了它的定义。这里主要关心的是ShuffleHandle的生成以及shuffle数据清理逻辑的注册，细节见代码中的注释。
class ShuffleDependency[K: ClassTag, V: ClassTag, C: ClassTag](
@transient private val _rdd: RDD[_ <: Product2[K, V]],
val partitioner: Partitioner,
val serializer: Serializer = SparkEnv.get.serializer,
val keyOrdering: Option[Ordering[K]] = None,
val aggregator: Option[Aggregator[K, V, C]] = None,
val mapSideCombine: Boolean = false)
extends Dependency[Product2[K, V]] {
if (mapSideCombine) {
require(aggregator.isDefined, "Map-side combine without Aggregator specified!")
}
override def rdd: RDD[Product2[K, V]] = _rdd.asInstanceOf[RDD[Product2[K, V]]]
private[spark] val keyClassName: String = reflect.classTag[K].runtimeClass.getName
private[spark] val valueClassName: String = reflect.classTag[V].runtimeClass.getName
// Note: It's possible that the combiner class tag is null, if the combineByKey
// methods in PairRDDFunctions are used instead of combineByKeyWithClassTag.
private[spark] val combinerClassName: Option[String] =
Option(reflect.classTag[C]).map(_.runtimeClass.getName)
// 在创建ShuffleDependency实例时，会生成一个在当前用户环境下的唯一ID
// 以便能够在其它地方通过这个id映射到此实例
val shuffleId: Int = _rdd.context.newShuffleId()
// 同时在实例化时，会向当前执行环境中的ShuffleManager中注册自己
// 实际是返回一个适合的ShuffleHandle，以便在Map端完成数据处理工作时，决定如何
// 缓存这些数据。
// BypassMergeSortShuffleHandle：序列化并直接将数据写出到文件文件
// SerializedShuffleHandle：对应于UnsafeShuffleWriter，数据会被序列化，先缓存，如果走过内存则溢出到文件
// BaseShuffleHandle：对应于SortShuffleWriter写出器，不序列化，先缓存，如果超过内存则溢出到文件
val shuffleHandle: ShuffleHandle = _rdd.context.env.shuffleManager.registerShuffle(
shuffleId, _rdd.partitions.length, this)
// 在这里会注册清理方法，当这个ShuffleDependency对应的数据不再使用时，清理掉
_rdd.sparkContext.cleaner.foreach(_.registerShuffleForCleanup(this))
}
ShuffleMapTask创建
当调用RDD的action行为的方法时，会触发DAGScheduler::handleJobSubmitted(…)方法的调用，然后遍历前面讲到的ShuffledRDD依赖图，创建任务ShuffleMapTask和FinalTask两种类型的任务，最后提交集群运行。提交任务方法的核心代码如下：
private[scheduler] def handleJobSubmitted(jobId: Int,
finalRDD: RDD[_],
func: (TaskContext, Iterator[_]) => _,
partitions: Array[Int],
callSite: CallSite,
listener: JobListener,
properties: Properties) {
var finalStage: ResultStage = null
// New stage creation may throw an exception if, for example, jobs are run on a
// HadoopRDD whose underlying HDFS files have been deleted.
// 从最后一个RDD，递归生成Stage图，本人其它Spark文章已经有介绍，这里仅简单说明：
// 这个方法会根据RDD的依赖图，即ShuffleDependency依赖图，生成一个
// ShuffleMapStage的依赖图，ShuffleDependency依赖图中的每一个结点就是一个
// ShuffleMapStage
finalStage = createResultStage(finalRDD, func, partitions, jobId, callSite)
...
// 创建一个ActiveJob实例，也就是我们在WEBUI页面上看到的JOB，每一个
// 依赖图的都会产生一个新的JOB。
val job = new ActiveJob(jobId, finalStage, callSite, listener, properties)
...
val jobSubmissionTime = clock.getTimeMillis()
// 缓存ActiveJob信息
jobIdToActiveJob(jobId) = job
activeJobs += job
// 设置finalStage关联的Job
finalStage.setActiveJob(job)
val stageIds = jobIdToStageIds(jobId).toArray
val stageInfos = stageIds.flatMap(id => stageIdToStage.get(id).map(_.latestInfo))
// 通知所有的事件监听者，有新的任务生成
listenerBus.post(
SparkListenerJobStart(job.jobId, jobSubmissionTime, stageInfos, properties))
// 提交sage
submitStage(finalStage)
}
ShuffleMapStage的创建
createShuffleMapStage是在DAGScheduler中定义的方法用于从当前的ShuffleDenpendency实例递归地构建每一个ShuffleMapStage，同时会将它们注册到各个缓存数据结构中，以便在任务运行时交换状态及数据，其定义如下：
// 可以看到每一个ShuffleDependency实例就对应一个ShuffleMapStage实例
def createShuffleMapStage(shuffleDep: ShuffleDependency[_, _, _], jobId: Int): ShuffleMapStage = {
val rdd = shuffleDep.rdd
checkBarrierStageWithDynamicAllocation(rdd)
checkBarrierStageWithNumSlots(rdd)
checkBarrierStageWithRDDChainPattern(rdd, rdd.getNumPartitions)
val numTasks = rdd.partitions.length
// 递归创建父依赖（ShuffleMapStage)的实例
val parents = getOrCreateParentStages(rdd, jobId)
val id = nextStageId.getAndIncrement()
val stage = new ShuffleMapStage(
id, rdd, numTasks, parents, jobId, rdd.creationSite, shuffleDep, mapOutputTracker)
// 缓存当前stage的映射信息
stageIdToStage(id) = stage
// 缓存父stage到当前stage的映射信息
shuffleIdToMapStage(shuffleDep.shuffleId) = stage
// 将当前stage添加到jobId对应的集合
updateJobIdStageIdMaps(jobId, stage)
if (!mapOutputTracker.containsShuffle(shuffleDep.shuffleId)) {
logInfo("Registering RDD " + rdd.id + " (" + rdd.getCreationSite + ")")
// 已经处理完当前stage所有父依赖，向MapOutputTracker中注册当前
// shuffle stage，实际上就是缓存当前stage与ShuffleStatus实例的映射，
// 而ShuffleStatus用来记录上游rdd即将输出的每一个partion的状态
mapOutputTracker.registerShuffle(shuffleDep.shuffleId, rdd.partitions.length)
}
stage
}
提交Stage图
在前面的过程中，DAGScheduler根据RDD依赖图，最终生成了Stage依赖，如果想要Stage依次执行，就需要将整个依赖图，转换成一个个有依赖关系的Task。而submitStage(stage: Stage)方法，就是从Stage依赖图的最后一个Stage，即FinalStage开始构建不同类型的任务，具体的方法定义如下：
/** Submits stage, but first recursively submits any missing parents. */
private def submitStage(stage: Stage) {
val jobId = activeJobForStage(stage)
if (jobId.isDefined) {
logDebug("submitStage(" + stage + ")")
// 如果当前stage还没有被执行，同时如果有父Stage且还没准备好，
// 就递归地、尝试执行构建祖先任务，将提交执行。
if (!waitingStages(stage) && !runningStages(stage) && !failedStages(stage)) {
val missing = getMissingParentStages(stage).sortBy(_.id)
logDebug("missing: " + missing)
if (missing.isEmpty) {
logInfo("Submitting " + stage + " (" + stage.rdd + "), which has no missing parents")
// 如果发现所有依赖的stage都已经完成了，那么就可以从当前stage开始跑任务啦
submitMissingTasks(stage, jobId.get)
} else {
for (parent <- missing) {
submitStage(parent)
}
waitingStages += stage
}
}
} else {
abortStage(stage, "No active job for stage " + stage.id, None)
}
}
从当前ShuffleStage构建任务
下面的方法，首先查找当前stage要计算的所有的Partitions信息，然后根据自己的类型分别创建ShuffleMapTask或是ResultTask，并填充任务实例所需要的各种信息，最后在有任务还没跑时，调用TaskScheduler::submitTasks(…)方法提交任务。
/** Called when stage's parents are available and we can now do its task. */
private def submitMissingTasks(stage: Stage, jobId: Int) {
logDebug("submitMissingTasks(" + stage + ")")
// First figure out the indexes of partition ids to compute.
// 找到所有要计算的partition信息
val partitionsToCompute: Seq[Int] = stage.findMissingPartitions()
// Use the scheduling pool, job group, description, etc. from an ActiveJob associated
// with this Stage
val properties = jobIdToActiveJob(jobId).properties
runningStages += stage
// SparkListenerStageSubmitted should be posted before testing whether tasks are
// serializable. If tasks are not serializable, a SparkListenerStageCompleted event
// will be posted, which should always come after a corresponding SparkListenerStageSubmitted
// event.
stage match {
case s: ShuffleMapStage =>
outputCommitCoordinator.stageStart(stage = s.id, maxPartitionId = s.numPartitions - 1)
case s: ResultStage =>
// 如果是
outputCommitCoordinator.stageStart(
stage = s.id, maxPartitionId = s.rdd.partitions.length - 1)
}
val taskIdToLocations: Map[Int, Seq[TaskLocation]] = try {
stage match {
case s: ShuffleMapStage =>
partitionsToCompute.map {
id => (id, getPreferredLocs(stage.rdd, id))}.toMap
case s: ResultStage =>
partitionsToCompute.map {
id =>
val p = s.partitions(id)
(id, getPreferredLocs(stage.rdd, p))
}.toMap
}
} catch {
case NonFatal(e) =>
stage.makeNewStageAttempt(partitionsToCompute.size)
listenerBus.post(SparkListenerStageSubmitted(stage.latestInfo, properties))
abortStage(stage, s"Task creation failed: $e\n${Utils.exceptionString(e)}", Some(e))
runningStages -= stage
return
}
stage.makeNewStageAttempt(partitionsToCompute.size, taskIdToLocations.values.toSeq)
// If there are tasks to execute, record the submission time of the stage. Otherwise,
// post the even without the submission time, which indicates that this stage was
// skipped.
if (partitionsToCompute.nonEmpty) {
stage.latestInfo.submissionTime = Some(clock.getTimeMillis())
}
listenerBus.post(SparkListenerStageSubmitted(stage.latestInfo, properties))
// TODO: Maybe we can keep the taskBinary in Stage to avoid serializing it multiple times.
// Broadcasted binary for the task, used to dispatch tasks to executors. Note that we broadcast
// the serialized copy of the RDD and for each task we will deserialize it, which means each
// task gets a different copy of the RDD. This provides stronger isolation between tasks that
// might modify state of objects referenced in their closures. This is necessary in Hadoop
// where the JobConf/Configuration object is not thread-safe.
var taskBinary: Broadcast[Array[Byte]] = null
var partitions: Array[Partition] = null
try {
// For ShuffleMapTask, serialize and broadcast (rdd, shuffleDep).
// For ResultTask, serialize and broadcast (rdd, func).
var taskBinaryBytes: Array[Byte] = null
// taskBinaryBytes and partitions are both effected by the checkpoint status. We need
// this synchronization in case another concurrent job is checkpointing this RDD, so we get a
// consistent view of both variables.
RDDCheckpointData.synchronized {
taskBinaryBytes = stage match {
case stage: ShuffleMapStage =>
JavaUtils.bufferToArray(
closureSerializer.serialize((stage.rdd, stage.shuffleDep): AnyRef))
case stage: ResultStage =>
JavaUtils.bufferToArray(closureSerializer.serialize((stage.rdd, stage.func): AnyRef))
}
partitions = stage.rdd.partitions
}
// 序列化任务图，包含当前stage.rdd信息以及依赖关系信息stage.shuffleDep
taskBinary = sc.broadcast(taskBinaryBytes)
} catch {
// In the case of a failure during serialization, abort the stage.
case e: NotSerializableException =>
abortStage(stage, "Task not serializable: " + e.toString, Some(e))
runningStages -= stage
// Abort execution
return
case e: Throwable =>
abortStage(stage, s"Task serialization failed: $e\n${Utils.exceptionString(e)}", Some(e))
runningStages -= stage
// Abort execution
return
}
// 创建任务
val tasks: Seq[Task[_]] = try {
val serializedTaskMetrics = closureSerializer.serialize(stage.latestInfo.taskMetrics).array()
stage match {
// 如果是ShuffleMapStage类型的任务，则创建ShuffleMapTask
// 对每一个partition，都创建一个任务，因此可以看到ShuffleMapTask会带有
// stage.id：当前任务属于哪个stage
// stage.lastInfo.attemptNumber：任务的重试次数
// taskBinary：任务执行时，需要执行在parition数据集上需要进行的操作
// part：当前任务要在哪个Partition数据集上的运行
// locs：任务要读取的parition数据的位置，可以在本地，可能在远端
// ...
case stage: ShuffleMapStage =>
stage.pendingPartitions.clear()
partitionsToCompute.map {
id =>
val locs = taskIdToLocations(id)
val part = partitions(id)
stage.pendingPartitions += id
new ShuffleMapTask(stage.id, stage.latestInfo.attemptNumber,
taskBinary, part, locs, properties, serializedTaskMetrics, Option(jobId),
Option(sc.applicationId), sc.applicationAttemptId, stage.rdd.isBarrier())
}
// 创建ResultTask类型的任务
case stage: ResultStage =>
partitionsToCompute.map {
id =>
val p: Int = stage.partitions(id)
val part = partitions(p)
val locs = taskIdToLocations(id)
new ResultTask(stage.id, stage.latestInfo.attemptNumber,
taskBinary, part, locs, id, properties, serializedTaskMetrics,
Option(jobId), Option(sc.applicationId), sc.applicationAttemptId,
stage.rdd.isBarrier())
}
}
} catch {
case NonFatal(e) =>
abortStage(stage, s"Task creation failed: $e\n${Utils.exceptionString(e)}", Some(e))
runningStages -= stage
return
}
if (tasks.size > 0) {
logInfo(s"Submitting ${tasks.size} missing tasks from $stage (${stage.rdd}) (first 15 " +
s"tasks are for partitions ${tasks.take(15).map(_.partitionId)})")
// 当前Stage所依赖的所有Paritions，都会生成对应的ShuffleMapTask或是ResultTask，
// 这些任务又被组织成了TaskSet类，它表示了一组可以并行 计算的任务集合，至于如何分配这些任务
//到哪些Worker上执行，以及如何管理状态等，都会交给TaskSetManager去完成。
taskScheduler.submitTasks(new TaskSet(
tasks.toArray, stage.id, stage.latestInfo.attemptNumber, jobId, properties))
} else {
// Because we posted SparkListenerStageSubmitted earlier, we should mark
// the stage as completed here in case there are no tasks to run
markStageAsFinished(stage, None)
stage match {
case stage: ShuffleMapStage =>
logDebug(s"Stage ${stage} is actually done; " +
s"(available: ${stage.isAvailable}," +
s"available outputs: ${stage.numAvailableOutputs}," +
s"partitions: ${stage.numPartitions})")
markMapStageJobsAsFinished(stage)
case stage : ResultStage =>
logDebug(s"Stage ${stage} is actually done; (partitions: ${stage.numPartitions})")
}
submitWaitingChildStages(stage)
}
}
ShuffleMapTask的执行
前面已经讲到，在Driver端，DAGScheduler会遍历RDD依赖图，最终生成一组组的TaskSet，并将由TaskSetManager进行管理，最终每一个ShuffleMapTask会给分配到空闲的Worker上运行，而任务的运行过程定义在ShuffleMapTask中，代码如下：
class ShuffleMapTask {
override def runTask(context: TaskContext): MapStatus = {
// Deserialize the RDD using the broadcast variable.
val threadMXBean = ManagementFactory.getThreadMXBean
val deserializeStartTime = System.currentTimeMillis()
val deserializeStartCpuTime = if (threadMXBean.isCurrentThreadCpuTimeSupported) {
threadMXBean.getCurrentThreadCpuTime
} else 0L
val ser = SparkEnv.get.closureSerializer.newInstance()
// 从任务携带的taskBinary广播变量中，反序列化出当前任务对应的RDD，以及依赖信息
val (rdd, dep) = ser.deserialize[(RDD[_], ShuffleDependency[_, _, _])](
ByteBuffer.wrap(taskBinary.value), Thread.currentThread.getContextClassLoader)
_executorDeserializeTime = System.currentTimeMillis() - deserializeStartTime
_executorDeserializeCpuTime = if (threadMXBean.isCurrentThreadCpuTimeSupported) {
threadMXBean.getCurrentThreadCpuTime - deserializeStartCpuTime
} else 0L
var writer: ShuffleWriter[Any, Any] = null
try {
val manager = SparkEnv.get.shuffleManager
// 根据当前任务指定的shuffle处理器类型，返回一个可用的、绑定到partitionId的writer，
writer = manager.getWriter[Any, Any](dep.shuffleHandle, partitionId, context)
// 通过rdd的迭代iterator操作，处理绑定的parition数据，此RDD可能包含了多个流水线操作，
// 例如有依赖如下（其中<--表示右边依赖左边）：
// ShuffledRDD <-- MapParitionsRDD <-- MapParitionsRDD <-- current_rdd
// 因此rdd.iterator在迭代时，会以递归的形式，从最左边的ShuffledRDD开始，通过
// ShuffleManager远程读取数据，然后再今次执行后面的两个MapParitionsRDD上绑定的处理函数，
// 最终由当前RDD，即ShuffledRDD类型，根据dep.shuffleHandle类型，将处理后的Parition
// 数据输出到指定位置。
writer.write(rdd.iterator(partition, context).asInstanceOf[Iterator[_ <: Product2[Any, Any]]])
writer.stop(success = true).get
} catch {
case e: Exception =>
try {
if (writer != null) {
writer.stop(success = false)
}
} catch {
case e: Exception =>
log.debug("Could not stop writer", e)
}
throw e
}
}
}
Shuffle数据合并及写出
BypassMergeSortShuffleWriter
ShuffleMapTask数据处理完成后，如果它的数据处理方法是BypassMergeSortShuffleHandle类型的，那么就会使用此Writer写出Shuffle数据，并将写出位置注册到当前环境的BlockManager中。
此Writer写出数据时，会首先为每一个Partition都创建一个单独的文件，并开启一个文件流；然后顺序将每一个Parition数据直接写出到文件；最后将所有单独的parition文件拼接（concatenate）成一个文件，并写出索引文件。
可以看到这个写出器并不会Map端进行数据的合并、排序等，同时不会缓存数据在内存中，因此性能比较底，会在下面的场景中使用：
- 没有指定排序键
- 没有指定聚合函数
- partitions的数量小于
spark.shuffle.sort.bypassMergeThreshold指定的值
UnsafeShuffleWriter
ShuffleMapTask数据处理完成后，如果它的数据处理方法是SerializedShuffleHandle类型的，那么就会使用此Writer先尝试在内存中保存所有的记录，(实际上内部通过ShuffleExternalSorter维护数据的缓存与写出)，如果当前任务完成或是内存要溢出时，会触发写数据到文件的动作。
此Writer会在触发写文件时，先在内存中将所有的记录按ParitionID排序；然后将数据序列化后写出到一个单独的文件中（如果有内存溢出发生时，可能会写出多个溢出文件，即是将内存中的数据序列化后写出到另外一个溢出文件）；最后所有的数据序列化后的溢出文件会被合并为一个文件。
SortShuffleWriter
ShuffleMapTask数据处理完成后，如果它的数据处理方法是BaseShuffleHandle类型的，那么就会使用此Writer，其基本的数据溢出逻辑与UnsafeShuffleWriter相似，不同的是它内部会创建ExternalSorter实例，并绑定Combine方法，在必要的情况下先在内存中对数据进行聚合。
We repeatedly fill up buffers of in-memory data, using either a PartitionedAppendOnlyMap if we want to combine by key, or a PartitionedPairBuffer if we don’t.
Map中的数据是PartitionID作为key。
同样在处理完成所有的数据后，会将所有可能的溢出文件与内存中的数据按ParitionID进行合并，并执行聚合操作，最终写出到一个shuffle文件；当然也会产生一个索引文件。
ShuffleMapTask任务的更新
当一个ShuffleMapTask任务完成后，会生成一个状态实例MapStatus，记录了这个任务的产生的数据信息。
随后 Executor就会map status转发给driver端，driver收到executor发来的任务完成消息时后，再会将这个消息转发给真正的任务调度器TaskSchedulerImpl，同时回收executor的上分配的资源，但不会移除它。
TaskSchedulerImpl收到任务完成的消息后，会更新内部的状态器，同时判断是不是这个任务所属的TaskSet中的任务都执行完成了，如果都完成了，那么就触发调度后置的TaskSet执行。
private[scheduler] def handleTaskCompletion(event: CompletionEvent) {
...
// 获取到完成任务对应的Stage
val stage = stageIdToStage(task.stageId)
...
stage match {
// 如果任务是一个ShuffleMapTask任务
case smt: ShuffleMapTask =>
// 则将对应的stage转换为ShuffleMapStage
val shuffleStage = stage.asInstanceOf[ShuffleMapStage]
if (!shuffleStage.isAvailable) {
// Some tasks had failed; let's resubmit this shuffleStage.
// TODO: Lower-level scheduler should also deal with this
logInfo("Resubmitting " + shuffleStage + " (" + shuffleStage.name +
") because some of its tasks had failed: " +
shuffleStage.findMissingPartitions().mkString(", "))
submitStage(shuffleStage)
} else {
// 标记这个Stage的状态的执行成功
markMapStageJobsAsFinished(shuffleStage)
// 提交后置Stage执行
submitWaitingChildStages(shuffleStage)
}
...
}
}
Shuffle数据的读取
SortShuffleManager
一个ShuffleMapTask任务在依赖图时，即ShuffleDependency时，就会将自己注册到ShuffleManager上，是便DAGScheduler在调度时，能够跟踪一个JOB所有Stage的状态信息，进而选择合适的时机、合适的Stage进行调度执行。
Shuffle数据读取
从ShuffleManager构建上游Stage包含的某个Partition所使用的Shuffle数据写出器，通过Spark中实现的Iterator实例，从当前任务RDD上绑定的函数开始执行，递归调用可能的上游RDD上的处理函数执行，直到ShuffledRDD。
而ShuffledRDD的工作并不是直接完成计算了，而是从本地或是远程的BlockManager上拉取当前任务所需的Partition的数据。
当整个数据都被处理完成后，当前的ShuffleMapTask又会将新产生的数据通过自己的ShuffleWriter写出到合适的位置，最终完成整个任务的执行。
var writer: ShuffleWriter[Any, Any] = null
val manager = SparkEnv.get.shuffleManager
writer = manager.getWriter[Any, Any](dep.shuffleHandle, partitionId, context)
writer.write(rdd.iterator(partition, context).asInstanceOf[Iterator[_ <: Product2[Any, Any]]])
writer.stop(success = true).get
总结
从RDD依赖图到任务执行的流程简化如下（其中<-- 表示右边依赖左边的完成）：
调用RDD或DataFrame API，生成一个RDD依赖图，例如df.read().filter(…).partitionBy(…).write()，
将生成ShuffledRDD <-- MapPartitionsRDD <-- ShuffledRDD <-- ShuffledRDD
通过DAGScheduler，生成依赖图：
ShuffleDependency <-- ShuffleDependency <-- ShuffleDependency
通过DAGScheduler，生成Stage依赖图
ShuffleMapStage <-- ShuffleMapStage <-- ResultStage
通过DAGScheduler，生成任务集TaskSet
ShuffleMapTask_1 <-- ShuffleMapTask_2 <-- ResultTask
通过TaskSetManager管理任务，将任务发送到Worker端执行
在Worker上执行ShuffleMapTask_1
在Worker上执行ShuffleMapTask_2
在Driver上执行ResultTask