Spark SQL CLI 實現分析

背景

本文主要介紹了Spark SQL裏目前的CLI實現，代碼之後肯定會有不少變動，所以我關注的是比較核心的邏輯。主要是對比了Hive CLI的實現方式，比較Spark SQL在哪塊地方做了修改，哪些地方與Hive CLI是保持一致的。可以先看下總結一節裏的內容。

Spark SQL的hive-thriftserver項目裏是其CLI實現代碼，下面先說明Hive CLI的主要實現類和關係，再說明Spark SQL CLI的做法。

Hive CLI

核心啓動類是org.apache.hive.service.server.HiveServer2，啓動方式：

    try {
      ServerOptionsProcessor oproc = new ServerOptionsProcessor("hiveserver2");
      if (!oproc.process(args)) {
        LOG.fatal("Error starting HiveServer2 with given arguments");
        System.exit(-1);
      }
      HiveConf hiveConf = new HiveConf();
      HiveServer2 server = new HiveServer2();
      server.init(hiveConf);
      server.start();
    } catch (Throwable t) {
      LOG.fatal("Error starting HiveServer2", t);
      System.exit(-1);
    }

HiveServer2繼承CompositeService類，CompositeService類內部維護一個serviceList，能夠加入、刪除、啓動、停止不同的服務。HiveServer2在init(hiveConf)的時候，會加入CLIService和ThriftCLIService兩個Service。根據傳輸模式，如果是http或https的話，就使用ThriftHttpCLIService，否則使用ThriftBinaryCLIService。無論是哪個ThriftCLIService，都傳入了CLIService的引用，thrift只是一個封裝。

加入了這些服務後，把服務都啓動起來。

CLIService也繼承自CompositeService，CLIService 在init的時候會加入SessionManager服務，並且根據hiveConf，從 hadoop shims裏得到UGI裏的serverUsername。

SessionManager管理hive連接的開啓、關閉等管理功能，已有的連接會維護在一個HashMap裏，value爲HiveSession類，裏面大致是用戶名、密碼、hive配置等info。

所以CLIService裏幾乎所有的事情都是委託給SessionManager做的。

 

SessionManager內主要是OperationManager這個服務，是最重要的和執行邏輯有關的類，下面會具體說。

 

另外，關於ThriftCLIService，有兩個實現子類，子類只複寫了run()方法，設置thrift server相關的網絡連接，其他對CLIService的調用邏輯都在父類ThriftCLIService本身裏面。

實際上，ThriftCLIService裏很多事情也是委託給CLIService做的。

 

那麼上面大致是Hive CLI、Thrift server啓動的流程，以及幾個主要類的相互關係。

Spark SQL CLI

根據上面Hive CLI的邏輯，看看Spark SQL的CLI是怎麼做的。

Spark裏的HiveThriftServer2(這個類名看起來有點奇怪)繼承了Hive的HiveServer2，並且複寫了init方法，其初始化的時候加入的是SparkSQLCLIService和ThriftBinaryCLIService兩個服務。前者繼承了Hive的CLIService，有一些不同的邏輯；後者直接使用的是Hive的類，但傳入的是SparkSQLCLIService的引用。

SparkSQLCLIService內部，類似Hive的CLIService，有一個SparkSQLSessionManager，繼承自Hive的SessionManager。也有得到serverUsername的邏輯，代碼和CLIService是一樣的。

 

SparkSQLSessionManager複寫了init這個方法，裏面有Spark自己的SparkSQLOperationManager服務，繼承自Hive的OperationManager類。

 

可能上面這幾個類有點看暈了，本質上都是一些封裝而已，沒什麼大的區別。真正重要的是SparkSQLOperationManager這個類裏面，定義瞭如何使用Spark SQL來處理query操作。

SparkSQLOperationManager關鍵邏輯

Hive的CLI Operation父類有如下的子類繼承體系，代表hive cli會處理的不同操作類型：

上半部分ExecuteStatementOperation子類體系是實際和查詢相關的操作，下半部分是一些元數據讀取操作。SparkSQLOperationManager實際改寫的就是ExecuteStatementOperation子類的執行邏輯，而元數據相關的操作還是沿用hive本來的處理邏輯。

 

原本hive的ExecuteStatementOperation處理邏輯是這樣的：

  public static ExecuteStatementOperation newExecuteStatementOperation(
      HiveSession parentSession, String statement, Map<String, String> confOverlay, boolean runAsync) {
    String[] tokens = statement.trim().split("\\s+");
    String command = tokens[0].toLowerCase();

    if ("set".equals(command)) {
      return new SetOperation(parentSession, statement, confOverlay);
    } else if ("dfs".equals(command)) {
      return new DfsOperation(parentSession, statement, confOverlay);
    } else if ("add".equals(command)) {
      return new AddResourceOperation(parentSession, statement, confOverlay);
    } else if ("delete".equals(command)) {
      return new DeleteResourceOperation(parentSession, statement, confOverlay);
    } else {
      return new SQLOperation(parentSession, statement, confOverlay, runAsync);
    }
  }

ExecuteStatementOperation也分兩部分，HiveCommandOperation和SQLOperation。

不同的ExecuteStatementOperation子類最終由對應的CommandProcessor子類來完成操作請求。

那Spark是如何改寫ExecuteStatementOperation的執行邏輯的呢？

最核心的邏輯如下：

      def run(): Unit = {
        logInfo(s"Running query '$statement'")
        setState(OperationState.RUNNING)
        try {
          result = hiveContext.sql(statement)
          logDebug(result.queryExecution.toString())
          val groupId = round(random * 1000000).toString
          hiveContext.sparkContext.setJobGroup(groupId, statement)
          iter = result.queryExecution.toRdd.toLocalIterator
          dataTypes = result.queryExecution.analyzed.output.map(_.dataType).toArray
          setHasResultSet(true)
        } catch {
          // Actually do need to catch Throwable as some failures don't inherit from Exception and
          // HiveServer will silently swallow them.
          case e: Throwable =>
            logError("Error executing query:",e)
            throw new HiveSQLException(e.toString)
        }
        setState(OperationState.FINISHED)
      }

statement是一個String，即query本身，調用HiveContext的sql()方法，返回的是一個SchemaRDD。HiveContext的這段邏輯如下：

  override def sql(sqlText: String): SchemaRDD = {
    // TODO: Create a framework for registering parsers instead of just hardcoding if statements.
    if (dialect == "sql") {
      super.sql(sqlText)
    } else if (dialect == "hiveql") {
      new SchemaRDD(this, HiveQl.parseSql(sqlText))
    }  else {
      sys.error(s"Unsupported SQL dialect: $dialect.  Try 'sql' or 'hiveql'")
    }
  }

調完sql()後返回的是一個帶被解析過了的基礎邏輯計劃的SchemaRDD。後續，

logDebug(result.queryExecution.toString())

這一步觸發了邏輯執行計劃的進一步分析、優化和變成物理執行計劃的幾個過程。之後，

result.queryExecution.toRdd

toRdd這步是觸發計算並返回結果。這幾個邏輯在之前Spark SQL源碼分析的文章裏都提到過。

除了上面這部分，還有一些schema轉化、數據類型轉化的邏輯，是因爲Catalyst這邊，有自己的數據行表示方法，也有自己的dataType，而且schema這塊呢，在生成SchemaRDD的時候也轉化過一次。所以在返回執行結果的時候，需要有轉換回Hive的TableSchema、FieldSchema的邏輯。

 

以上說明了Spark SQL是如何把query的執行轉換到Spark SQL裏的。

總結

基本上Spark SQL在CLI這塊的實現很靠近Hive Service項目裏的CLI模塊，主要類繼承體系、執行邏輯差不多都一樣。Spark SQL修改的關鍵邏輯在CLIService內的SessionManager內的OperationManager裏，將非元數據查詢操作的query丟給了Spark SQL的Hive工程裏的HiveContext.sql()來完成，通過返回的SchemaRDD，來進一步得到結果數據、得到中間執行計劃的Schema信息。

全文完 :)