Spark SQL加载kafka数据并将查询结果写出到kafka的执行过程

说明：本文意在对源码进行分析，说明Spark SQL加载kafka数据并将查询结果写出到kafka的过程，如果错误，欢迎指出，大家共同进步 ^ _ ^。

一、加载kafka数据的代码样例

// 这段代码加载kafka数据源，指定了数据源格式，topic、bootstrap server、offset信息等
Dataset dataset = this.sparkSession.readStream().format("kafka")
		.options(this.getSparkKafkaCommonOptions(sparkSession))
		.option("kafka.bootstrap.servers", ConfigConstMdt.MDT_STREAMING_KAFKA_SERVERS)
		.option("subscribe", ConfigConstMdt.MDT_STREAMING_KAFKA_TOPICS)
		.option("startingOffsets", ConfigConstMdt.MDT_STREAMING_KAFKA_OFFSETS)
		.load();

我们看看load源码

def load(): DataFrame = {
    // 这里的source就是.format("kafka")中的"kafka"
    if (source.toLowerCase(Locale.ROOT) == DDLUtils.HIVE_PROVIDER) {
      throw new AnalysisException("Hive data source can only be used with tables, you can not " +
        "read files of Hive data source directly.")
    }

    // 1 、这里得到的ds是一个KafkaSourceProvider实例
    val ds = DataSource.lookupDataSource(source, sparkSession.sqlContext.conf).newInstance()
    // 2、v1DataSource 确切地讲是对kafka数据源基本信息（配置、schema）的一种封装
    val v1DataSource = DataSource(
      sparkSession,
      userSpecifiedSchema = userSpecifiedSchema,
      className = source,
      options = extraOptions.toMap)
    val v1Relation = ds match {
      // 3、因为KafkaSourceProvider继承StreamSourceProvider，所以匹配到这里，StreamingRelation继承了LeafNode，作用是
      // 将source和逻辑计划连接，逻辑计划用于创建流DataFrame，在传递到StreamExecution进行流查询的时候需要转换成
      // StreamingExecutionRelation。
      case _: StreamSourceProvider => Some(StreamingRelation(v1DataSource))
      case _ => None
    }
    ds match {
      // 以微批次为例
      case s: MicroBatchReadSupport =>
        ......
        // 4、主要看这里做了什么？ofRows告诉我们StreamingRelationV2就是用于组装一个逻辑计划，因为Dataset是懒执行的，
        // 其代表用于生产所需结果的计算过程的描述信息。因此，到这里数据并没有被加载，而是准备好了加载数据的逻辑计划。
        Dataset.ofRows(
          sparkSession,
          StreamingRelationV2(
            s, source, options,
            schema.toAttributes, v1Relation)(sparkSession))
    }
  }

我们应该可以发现，到这里程序并没有真的加载kafka数据。然后我们再看看DataStreamWriter的start，也就是用于启动流处理的函数调用，我们直接看kafka数据源匹配的地方。

二、结果写出到kafka的代码样例

resultDataSet.writeStream()
      .format("kafka")
      .options(this.context.getKafkaOptions())
      .option("topic", this.context.getTargetTopic())
      .option("checkpointLocation", this.context.getCheckPointPath())
      .outputMode(OutputMode.Update())
      .trigger(Trigger.Continuous(1, TimeUnit.MINUTES))
      .start();

在load时形成的Dataset，在写出时通过如下语法转换成了DataFrame

private val df = ds.toDF()

往下看

def start(): StreamingQuery = {
   if(......){
   ......
    } else {
      // 本段与前文分析相似，不再重复分析
      val ds = DataSource.lookupDataSource(source, df.sparkSession.sessionState.conf)
      val disabledSources = df.sparkSession.sqlContext.conf.disabledV2StreamingWriters.split(",")
      var options = extraOptions.toMap
      val sink = ds.newInstance() match {
        case w: StreamWriteSupport if !disabledSources.contains(w.getClass.getCanonicalName) =>
          val sessionOptions = DataSourceV2Utils.extractSessionConfigs(
            w, df.sparkSession.sessionState.conf)
          options = sessionOptions ++ extraOptions
          w
        case _ =>
          val ds = DataSource(
            df.sparkSession,
            className = source,
            options = options,
            partitionColumns = normalizedParCols.getOrElse(Nil))
          ds.createSink(outputMode)
      }
      // 这里开始执行查询
      df.sparkSession.sessionState.streamingQueryManager.startQuery(
        options.get("queryName"),
        options.get("checkpointLocation"),
        df,  // 这里的df就是通过上边toDF()得来
        options,
        sink,
        outputMode,
        useTempCheckpointLocation = source == "console",
        recoverFromCheckpointLocation = true,
        trigger = trigger)
    }
  }

继续往下，我们主要看看startQuery中的两个操作

• createQuery

createQuery(......){
    ......
    // 可以看到，在这里逻辑执行计划转换成为解析的逻辑执行计划
    val analyzedPlan = df.queryExecution.analyzed
    df.queryExecution.assertAnalyzed()
    ......
    // 以及如下操作，这一操作中就完成了前文中所说的StreamingRelation到StreamingExecutionRelation的转换。
    new MicroBatchExecution(...analyzedPlan...)
}

• start查询

query.streamingQuery.start()

接下进入到start里边看start过程

runStream() --->runActivatedStream(sparkSessionForStream) ---> runBatch(sparkSessionForStream) 

重点关注runBatch，首先是从source中获取批数据，这一步得到的是真实的数据，请求数据过程参照KafkaSource.scala中的getBatch即可。也就是说在这一步之前，不管是load还是其他操作，都没有真正的去kafka中拉取数据

val batch = source.getBatch(current, available)

然后用获取的batch数据替代逻辑计划logicalPlan中的sources，这一步操作使得DF中逻辑计划相对完整了，有了真实数据基础，查询计划在此数据集上完成

val newBatchesPlan = logicalPlan transform {......}

如下代码，强制生成执行计划

reportTimeTaken("queryPlanning") {
lastExecution = new IncrementalExecution(
   sparkSessionToRunBatch,
   triggerLogicalPlan,
   outputMode,
   checkpointFile("state"),
   runId,
   currentBatchId,
   offsetSeqMetadata)
 lastExecution.executedPlan // 强制生成执行计划
}

接着，同样生成一个带执行计划的懒执行的Dataset，其中包含了得到最终结果所需的计算过程的描述信息。

val nextBatch =
      new Dataset(sparkSessionToRunBatch, lastExecution, RowEncoder(lastExecution.analyzed.schema))

最后看看写出操作

val nextBatch =
    new Dataset(sparkSessionToRunBatch, lastExecution, RowEncoder(lastExecution.analyzed.schema))

    reportTimeTaken("addBatch") {
      SQLExecution.withNewExecutionId(sparkSessionToRunBatch, lastExecution) {
        sink match {
          // 这里的addBatch操作便是完成数据写入到对应sink的功能，由不同的Sink来实现
          case s: Sink => s.addBatch(currentBatchId, nextBatch)
          case _: StreamWriteSupport =>
            // This doesn't accumulate any data - it just forces execution of the microbatch writer.
            nextBatch.collect()
        }
      }
    }

KafkaSink.scala的addBatch实现

override def addBatch(batchId: Long, data: DataFrame): Unit = {
    if (batchId <= latestBatchId) {
      logInfo(s"Skipping already committed batch $batchId")
    } else {
      // 完成对应数据到kafka的写出  data.queryExecution触发Spark执行查询计划，执行完成后得到最终需要写出的结果
      KafkaWriter.write(sqlContext.sparkSession, data.queryExecution, executorKafkaParams, topic)
      latestBatchId = batchId
    }
  }

以上过程也体现了Spark的懒执行特性，也就是在KafkaWriter.write被调用时，data.queryExecution才被执行，于是Spark真正开始执行各stage中的计算组成的查询计划。