一文读懂 超简单的spark structured stream 源码解读

为了让大家理解structured stream的运行流程，我将根据一个代码例子，讲述structured stream的基本运行流程和原理。

下面是一段简单的代码:

val spark = SparkSession
      .builder
      .appName("StructuredNetworkWordCount")
      .master("local[4]")

      .getOrCreate()
    spark.conf.set("spark.sql.shuffle.partitions", 4)

    import spark.implicits._
    val words = spark.readStream
      .format("socket")
      .option("host", "localhost")
      .option("port", 9999)
      .load()

    val df1 = words.as[String]
      .flatMap(_.split(" "))
      .toDF("word")
      .groupBy("word")
      .count()

    df1.writeStream
      .outputMode("complete")
      .format("console")
      .trigger(ProcessingTime(10))
      .start()

    spark.streams.awaitAnyTermination()

　　这段代码就是单词计数。先从一个socket数据源读入数据，然后以" " 为分隔符把一行文本转换成单词的DataSet，然后转换成有标签("word")的DataFrame，接着按word列进行分组，聚合计算每个word的个数。最后输出到控制台，以10秒为批处理执行周期。

　　

现在来分析它的原理。spark的逻辑里面有一个惰性计算的概念，以上面的例子来说，在第22行代码以前，程序都不会对数据进行真正的计算，而是将计算的公式（或者函数）保存在DataFrame里面，在22行开始的writeStream.start调用后才开始真正的计算。为什么？

因为：

这可以让spark内核做一些优化。

例如：

数据库中存放着人的名字和年龄，我想要在控制台打印出前十个年龄大于20岁的人的名字，那么我的spark代码会这么写：

df.fileter{row=>
row._2>20}
.show(10)

假如说我每执行一行代码就进行一次计算，那么在第二行的时候，我就会把df里面所有的数据进行过滤，筛选出其中年龄大于20的，然后在第3行执行的时候，从第2行里面的结果中选前面10个进行打印。

看出问题了么？这里的输出仅仅只需要10个年龄大于20的人，但是我却把所有人都筛选了一遍，其实我只需要筛选出10个，后面的就不必要筛选了。这就是spark的惰性计算进行优化的地方。

在spark的计算中，在真正的输出函数之前，都不会进行真正的计算，而会在输出函数之前进行优化后再进行计算。我们来看源代码。

这里我贴的是structured stream每次批处理周期到达时会运行的代码:

 private def runBatch(sparkSessionToRunBatch: SparkSession): Unit = {
    // Request unprocessed data from all sources.
    newData = reportTimeTaken("getBatch") {
      availableOffsets.flatMap {
        case (source, available)
          if committedOffsets.get(source).map(_ != available).getOrElse(true) =>
          val current = committedOffsets.get(source)
          val batch = source.getBatch(current, available)
          logDebug(s"Retrieving data from $source: $current -> $available")
          Some(source -> batch)
        case _ => None
      }
    }

    // A list of attributes that will need to be updated.
    var replacements = new ArrayBuffer[(Attribute, Attribute)]
    // Replace sources in the logical plan with data that has arrived since the last batch.
    val withNewSources = logicalPlan transform {
      case StreamingExecutionRelation(source, output) =>
        newData.get(source).map { data =>
          val newPlan = data.logicalPlan
          assert(output.size == newPlan.output.size,
            s"Invalid batch: ${Utils.truncatedString(output, ",")} != " +
            s"${Utils.truncatedString(newPlan.output, ",")}")
          replacements ++= output.zip(newPlan.output)
          newPlan
        }.getOrElse {
          LocalRelation(output)
        }
    }

    // Rewire the plan to use the new attributes that were returned by the source.
    val replacementMap = AttributeMap(replacements)
    val triggerLogicalPlan = withNewSources transformAllExpressions {
      case a: Attribute if replacementMap.contains(a) => replacementMap(a)
      case ct: CurrentTimestamp =>
        CurrentBatchTimestamp(offsetSeqMetadata.batchTimestampMs,
          ct.dataType)
      case cd: CurrentDate =>
        CurrentBatchTimestamp(offsetSeqMetadata.batchTimestampMs,
          cd.dataType, cd.timeZoneId)
    }

    reportTimeTaken("queryPlanning") {
      lastExecution = new IncrementalExecution(
        sparkSessionToRunBatch,
        triggerLogicalPlan,
        outputMode,
        checkpointFile("state"),
        currentBatchId,
        offsetSeqMetadata)
      lastExecution.executedPlan // Force the lazy generation of execution plan
    }

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

    reportTimeTaken("addBatch") {
      sink.addBatch(currentBatchId, nextBatch)
    }

    awaitBatchLock.lock()
    try {
      // Wake up any threads that are waiting for the stream to progress.
      awaitBatchLockCondition.signalAll()
    } finally {
      awaitBatchLock.unlock()
    }
  }

其实很简单，在第58以前都是在解析用户代码，生成logicPlan，优化logicPlan，生成批处理类。第47行的triggerLogicalPlan就是最终优化后的用户逻辑，它被封装在了一个IncrementalExecution类中，这个类连同sparkSessionToRunBatch（运行环境）和RowEncoder（序列化类）一起构成一个新的DataSet，这个DataSet就是最终要发送到worker节点进行执行的代码。第59行代码就是在将它加入到准备发送代码的队列中。我们继续看一段代码，由于我们使用console作为数据下游（sink）所以看看console的addBatch代码：

override def addBatch(batchId: Long, data: DataFrame): Unit = synchronized {
    val batchIdStr = if (batchId <= lastBatchId) {
      s"Rerun batch: $batchId"
    } else {
      lastBatchId = batchId
      s"Batch: $batchId"
    }

    // scalastyle:off println
    println("-------------------------------------------")
    println(batchIdStr)
    println("-------------------------------------------")
    // scalastyle:off println
    data.sparkSession.createDataFrame(
      data.sparkSession.sparkContext.parallelize(data.collect()), data.schema)
      .show(numRowsToShow, isTruncated)
  }

关键代码在16行.show函数，show函数是一个真正的action，在这之前都是一些算子的封装，我们看show的代码:

private[sql] def showString(_numRows: Int, truncate: Int = 20): String = {
    val numRows = _numRows.max(0)
    val takeResult = toDF().take(numRows + 1)
    val hasMoreData = takeResult.length > numRows
    val data = takeResult.take(numRows)

第3行进入take:

  def take(n: Int): Array[T] = head(n)

def head(n: Int): Array[T] = withAction("head", limit(n).queryExecution)(collectFromPlan)

private def withAction[U](name: String, qe: QueryExecution)(action: SparkPlan => U) = {
    try {
      qe.executedPlan.foreach { plan =>
        plan.resetMetrics()
      }
      val start = System.nanoTime()
      val result = SQLExecution.withNewExecutionId(sparkSession, qe) {
        action(qe.executedPlan)
      }
      val end = System.nanoTime()
      sparkSession.listenerManager.onSuccess(name, qe, end - start)
      result
    } catch {
      case e: Exception =>
        sparkSession.listenerManager.onFailure(name, qe, e)
        throw e
    }
  }

这个函数名就告诉我们，这是真正计算要开始了，第7行代码一看就是准备发送代码序列了:

def withNewExecutionId[T](
      sparkSession: SparkSession,
      queryExecution: QueryExecution)(body: => T): T = {
    val sc = sparkSession.sparkContext
    val oldExecutionId = sc.getLocalProperty(EXECUTION_ID_KEY)
    if (oldExecutionId == null) {
      val executionId = SQLExecution.nextExecutionId
      sc.setLocalProperty(EXECUTION_ID_KEY, executionId.toString)
      executionIdToQueryExecution.put(executionId, queryExecution)
      val r = try {
        // sparkContext.getCallSite() would first try to pick up any call site that was previously
        // set, then fall back to Utils.getCallSite(); call Utils.getCallSite() directly on
        // streaming queries would give us call site like "run at <unknown>:0"
        val callSite = sparkSession.sparkContext.getCallSite()

        sparkSession.sparkContext.listenerBus.post(SparkListenerSQLExecutionStart(
          executionId, callSite.shortForm, callSite.longForm, queryExecution.toString,
          SparkPlanInfo.fromSparkPlan(queryExecution.executedPlan), System.currentTimeMillis()))
        try {
          body
        } finally {
          sparkSession.sparkContext.listenerBus.post(SparkListenerSQLExecutionEnd(
            executionId, System.currentTimeMillis()))
        }
      } finally {
        executionIdToQueryExecution.remove(executionId)
        sc.setLocalProperty(EXECUTION_ID_KEY, null)
      }
      r
    } else {
      // Don't support nested `withNewExecutionId`. This is an example of the nested
      // `withNewExecutionId`:
      //
      // class DataFrame {
      //   def foo: T = withNewExecutionId { something.createNewDataFrame().collect() }
      // }
      //
      // Note: `collect` will call withNewExecutionId
      // In this case, only the "executedPlan" for "collect" will be executed. The "executedPlan"
      // for the outer DataFrame won't be executed. So it's meaningless to create a new Execution
      // for the outer DataFrame. Even if we track it, since its "executedPlan" doesn't run,
      // all accumulator metrics will be 0. It will confuse people if we show them in Web UI.
      //
      // A real case is the `DataFrame.count` method.
      throw new IllegalArgumentException(s"$EXECUTION_ID_KEY is already set")
    }
  }

你看第16行，就是在发送数据，包括用户优化后的逻辑，批处理的id，时间戳等等。worker接收到这个事件后根据logicalPlan里面的逻辑就开始干活了。这就是一个很基本很简单的流程，对于spark入门还是挺有帮助的吧。