Spark-RDD迭代器管道计算

一、上下文

《Spark-Task启动流程》中讲到我们提交Stage是传入的是这个Stage最后一个RDD，当Task中触发ShuffleWriter、返回Driver数据或者写入Hadoop文件系统时才触发这个RDD调用它的iterator()，下面我们就来看下RDD.iterator()背后的故事。

二、RDD中的iterator

我们先来看下rdd.iterator()以及后面一些列的调用

  final def iterator(split: Partition, context: TaskContext): Iterator[T] = {
    //如果该rdd是被持久化的那么直接得到它或者根据血缘关系计算它，并返回，到这就终止了，可以直接从这往后计算了
    if (storageLevel != StorageLevel.NONE) {
      getOrCompute(split, context)
    } else {
      //调用compute()或者从Checkpoint读取数据
      computeOrReadCheckpoint(split, context)
    }
  }

  private[spark] def computeOrReadCheckpoint(split: Partition, context: TaskContext): Iterator[T] =
  {
    //如果该RDD是Checkpointed那么直接返回对应分区的迭代器，直接从这里开始计算
    if (isCheckpointedAndMaterialized) {
      firstParent[T].iterator(split, context)
    } else {
      compute(split, context)
    }
  }

  //由子类实现，用于计算给定的分区，下面我们看下重要的几个子类
  def compute(split: Partition, context: TaskContext): Iterator[T]

 三、RDD子类中的compute

MapPartitionsRDD

  //一个RDD调用 map() flatMap() filter() glom() mapPartitions() 都会返回一个 MapPartitionsRDD 且把对应的处理函数传进去
  //这个时候就会返回该函数，且里面会调用父RDD的 iterator 去拿数据来处理
  override def compute(split: Partition, context: TaskContext): Iterator[U] =
    f(context, split.index, firstParent[T].iterator(split, context))

 ShuffledRDD

  //一个RDD 调用 repartition() coalesce() sortByKey() groupByKey() reduceByKey() 时都有可能返回一个ShuffledRDD
  override def compute(split: Partition, context: TaskContext): Iterator[(K, C)] = {
    val dep = dependencies.head.asInstanceOf[ShuffleDependency[K, V, C]]
    val metrics = context.taskMetrics().createTempShuffleReadMetrics()
    //通过SparkEnv 获取 ShuffleManager Reader 来读取上一个Stage Shuffle Writer 的结果，并封装成一个迭代器返回
    SparkEnv.get.shuffleManager.getReader(
      dep.shuffleHandle, split.index, split.index + 1, context, metrics)
      .read()
      .asInstanceOf[Iterator[(K, C)]]
  }

 HadoopRDD

  override def compute(theSplit: Partition, context: TaskContext): InterruptibleIterator[(K, V)] = {
    //制作一个 Iterator 并返回 
    val iter = new NextIterator[(K, V)] {

      //将Spark 的 Partition 转换成 Hadoop 的 Partition
      private val split = theSplit.asInstanceOf[HadoopPartition]
      logInfo("Input split: " + split.inputSplit)
      private val jobConf = getJobConf()

      private val inputMetrics = context.taskMetrics().inputMetrics
      private val existingBytesRead = inputMetrics.bytesRead

      // 为文件的block 信息 设置 InputFileBlockHolder
      split.inputSplit.value match {
        case fs: FileSplit =>
          InputFileBlockHolder.set(fs.getPath.toString, fs.getStart, fs.getLength)
        case _ =>
          InputFileBlockHolder.unset()
      }

      // 找到一个函数，该函数将返回此线程读取的FileSystem字节。在创建RecordReader之前执行此操作，因为RecordReader的构造函数可能会读取一些字节
      private val getBytesReadCallback: Option[() => Long] = split.inputSplit.value match {
        case _: FileSplit | _: CombineFileSplit =>
          Some(SparkHadoopUtil.get.getFSBytesReadOnThreadCallback())
        case _ => None
      }

      // 我们从线程本地Hadoop FileSystem统计数据中获取输入字节。
      //然而，如果我们进行合并，我们可能会在同一任务和同一线程中计算多个分区，
      //在这种情况下，我们需要避免覆盖先前分区写入的值（SPARK-13071）。
      private def updateBytesRead(): Unit = {
        getBytesReadCallback.foreach { getBytesRead =>
          inputMetrics.setBytesRead(existingBytesRead + getBytesRead())
        }
      }

      private var reader: RecordReader[K, V] = null
      private val inputFormat = getInputFormat(jobConf)
      HadoopRDD.addLocalConfiguration(
        new SimpleDateFormat("yyyyMMddHHmmss", Locale.US).format(createTime),
        context.stageId, theSplit.index, context.attemptNumber, jobConf)

      reader =
        try {
          inputFormat.getRecordReader(split.inputSplit.value, jobConf, Reporter.NULL)
        } catch {
          ......
        }
      // 注册任务完成回调以关闭输入流。
      context.addTaskCompletionListener[Unit] { context =>
        // 更新关闭前读取的字节数是为了确保正确添加此线程中的延迟字节读取统计信息
        updateBytesRead()
        closeIfNeeded()
      }

      private val key: K = if (reader == null) null.asInstanceOf[K] else reader.createKey()
      private val value: V = if (reader == null) null.asInstanceOf[V] else reader.createValue()

      override def getNext(): (K, V) = {
        try {
          finished = !reader.next(key, value)
        } catch {
          ...
        }
        if (!finished) {
          inputMetrics.incRecordsRead(1)
        }
        if (inputMetrics.recordsRead % SparkHadoopUtil.UPDATE_INPUT_METRICS_INTERVAL_RECORDS == 0) {
          updateBytesRead()
        }
        (key, value)
      }

      override def close(): Unit = {
        if (reader != null) {
          InputFileBlockHolder.unset()
          try {
            reader.close()
          } catch {
            ......
          } finally {
            reader = null
          }
          if (getBytesReadCallback.isDefined) {
            updateBytesRead()
          } else if (split.inputSplit.value.isInstanceOf[FileSplit] ||
                     split.inputSplit.value.isInstanceOf[CombineFileSplit]) {
            // 如果我们无法从FS统计数据中读取字节，请回退到分割大小，这可能是不准确的。
            try {
              inputMetrics.incBytesRead(split.inputSplit.value.getLength)
            } catch {
             ......
            }
          }
        }
      }
    }
    new InterruptibleIterator[(K, V)](context, iter)
  }

四、总结

1、 一个Stage的最后一个RDD会调用自己的iterator()

2、判断该RDD是否是被持久化的，如果是，直接从这个RDD的持久化开始迭代的计算

3、判断该RDD是否是Checkpoint，如果是，直接从这个RDD开始迭代的计算

4、如果2和3不满足，调用父RDD的iterator()

5、判断父RDD是否满足2和3

6、如果父RDD2和3不满足，调用父的父RDD的iterator()

7、直到该RDD是HadoopRDD或者ShuffledRDD ，也就是该Stage最初的那个RDD

8、从Hadoop文件系统（HDFS或者本地文件系统）或使用 Shuffle Reader 读上一个Stage的Shuffle Writer 结果 再依次调用每个算子的函数，返回结果集

9、使用Shuffle Writer 将结果写入磁盘，这个Stage的Task完成