Spark Sql 源碼剖析(二): TreeNode

零、前置知識 Scala Product trait

// 所有 products 的基trait,至少包含 [[scala.Product1]] 至 [[scala.Product22]] 及 [[scala.Tuple1]] 至 [[scala.Tuple22]]
trait Product extends Any with Equals {
  // 第 n 個元素,從0開始
  def productElement(n: Int): Any

  // product size
  def productArity: Int

  // product 遍及所有元素的迭代器
  def productIterator: Iterator[Any] = new scala.collection.AbstractIterator[Any] {
    private var c: Int = 0
    private val cmax = productArity
    def hasNext = c < cmax
    def next() = { val result = productElement(c); c += 1; result }
  }
}

一、CurrentOrigin

使用 object CurrentOrigin 為 TreeNodes 提供一個可以查找上下文的地方,比如當前正在解析哪行 code。

// Origin 表示第幾行第幾列
case class Origin(
  line: Option[Int] = None,
  startPosition: Option[Int] = None)

object CurrentOrigin 主要包含一個 private val value = new ThreadLocal[Origin]() ,目前 CurrentOrigin 僅在 parser 中使用,在 visit 每個節點的時候都會使用,記錄當前 parse 的節點是哪行哪列

另外,從 value 是 ThreadLocal 類型可以看出,在 Spark SQL 中,parse sql 時都是在單獨的 thread 里進行的(不同的 sql 不同的 thread)

二、重要方法

2.1、children: Seq[BaseType](由子類實現)

返回該節點的 seq of children,children 是不可變的。有三種情況:

  • LeafNode:無 children
  • UnaryNode:包含一個 child
  • BinaryNode:包含 left、right 兩個 child

2.2、find(f: BaseType => Boolean): Option[BaseType]

查找第一個符合 f 條件(比如某個類型的)的 TreeNode,先序遍歷。

2.3、foreach(f: BaseType => Unit): Unit

  def foreach(f: BaseType => Unit): Unit = {
    f(this)
    children.foreach(_.foreach(f))
  }

將函數 f 遞歸應用于節點及其子節點

2.4、foreachUp(f: BaseType => Unit): Unit

與 foreach 不同的是,foreach 先應用于 parent,再應用與 child;而 foreachUp 是先應用于 child 再應用與 parent

2.5、map[A](f: BaseType => A): Seq[A]

  def map[A](f: BaseType => A): Seq[A] = {
    val ret = new collection.mutable.ArrayBuffer[A]()
    foreach(ret += f(_))
    ret
  }

調用 foreach,foreach 中應用的函數是 ret += f(_) ,最終返回一個 seq,包含將 f 通過 foreach 方式應用于所有節點并 add 到 ret。其中 f 本身是 BaseType => A 類型

2.6、flatMap[A](f: BaseType => TraversableOnce[A]): Seq[A]

原理與 map 一致,只是 f 變成了 BaseType => TraversableOnce[A]

2.5、collect[B](pf: PartialFunction[BaseType, B]): Seq[B]

  def collect[B](pf: PartialFunction[BaseType, B]): Seq[B] = {
    val ret = new collection.mutable.ArrayBuffer[B]()
    val lifted = pf.lift
    foreach(node => lifted(node).foreach(ret.+=))
    ret
  }

PartialFunction#lift:將 partial func 轉換為一個返回 Option 結果的函數。將 pf 函數應用于符合 pf 定義的節點(即 pf.lift(node)返回的 Option 不是 None )并都 add 到 ret = new collection.mutable.ArrayBuffer[B] 以 Seq 形式返回

2.6、collectLeaves(): Seq[BaseType]

以 Seq 的形式返回 tree 的所有葉子節點

def collectFirst[B](pf: PartialFunction[BaseType, B]): Option[B]:注意,因為可能沒有符合 pf 定義的節點,所有返回的 Option 可能是 None

2.7、mapProductIterator[B: ClassTag](f: Any => B): Array[B]

相當于 productIterator.map(f).toArray ,即對于 productIterator 每個元素執行 f 然后將 ret 組成一個 arr 返回

注意:TreeNode 沒有實現 Product 相關方法,都由其子類自行實現

2.8、withNewChildren

使用 new children 替換并返回該節點的拷貝。該方法會對 productElement 每個元素進行模式匹配,根據節點類型及一定規則進行替換。

2.9、transform(rule: PartialFunction[BaseType, BaseType]): BaseType

調用 transformDown

2.10、transformDown(rule: PartialFunction[BaseType, BaseType]): BaseType

rule: PartialFunction[BaseType, BaseType]

  def transformDown(rule: PartialFunction[BaseType, BaseType]): BaseType = {
    val afterRule = CurrentOrigin.withOrigin(origin) {
        // 如果 this 是 BaseType 或其子類,則對 this 應用 rule 再返回應用 rule 后的結果,否則返回 this 
      rule.applyOrElse(this, identity[BaseType])
    }

    // Check if unchanged and then possibly return old copy to avoid gc churn.
    if (this fastEquals afterRule) {
        // 如果應用了 rule 后節點無變化,則遞歸將 rule 應用于 children
      mapChildren(_.transformDown(rule))
    } else {
        // 如果應用了 rule 后節點有變化,則本節點換成變化后的節點(children 不變),再將 rule 遞歸應用于子節點。也就是從根節點往下來應用 rule 替換節點
      afterRule.mapChildren(_.transformDown(rule))
    }
  }

2.11、mapChildren(f: BaseType => BaseType): BaseType

返回 f 應用于所有子節點(非遞歸,一般將遞歸操作放在調用該函數的地方)后該節點的 copy。其內部的原理是調用 mapProductIterator,對每一個 productElement(i) 進行各種模式匹配,若能匹配上某個再根據一定規則進行轉換,核心匹配轉換如下:

case arg: TreeNode[_] if containsChild(arg) =>
  val newChild = f(arg.asInstanceOf[BaseType])
  if (!(newChild fastEquals arg)) {
    changed = true
    newChild
  } else {
    arg
  }
case Some(arg: TreeNode[_]) if containsChild(arg) =>
  val newChild = f(arg.asInstanceOf[BaseType])
  if (!(newChild fastEquals arg)) {
    changed = true
    Some(newChild)
  } else {
    Some(arg)
  }
case m: Map[_, _] => m.mapValues {
  case arg: TreeNode[_] if containsChild(arg) =>
    val newChild = f(arg.asInstanceOf[BaseType])
    if (!(newChild fastEquals arg)) {
      changed = true
      newChild
    } else {
      arg
    }
  case other => other
}.view.force // `mapValues` is lazy and we need to force it to materialize
case d: DataType => d // Avoid unpacking Structs
case args: Traversable[_] => args.map {
  case arg: TreeNode[_] if containsChild(arg) =>
    val newChild = f(arg.asInstanceOf[BaseType])
    if (!(newChild fastEquals arg)) {
      changed = true
      newChild
    } else {
      arg
    }
  case tuple@(arg1: TreeNode[_], arg2: TreeNode[_]) =>
    val newChild1 = if (containsChild(arg1)) {
      f(arg1.asInstanceOf[BaseType])
    } else {
      arg1.asInstanceOf[BaseType]
    }

    val newChild2 = if (containsChild(arg2)) {
      f(arg2.asInstanceOf[BaseType])
    } else {
      arg2.asInstanceOf[BaseType]
    }

    if (!(newChild1 fastEquals arg1) || !(newChild2 fastEquals arg2)) {
      changed = true
      (newChild1, newChild2)
    } else {
      tuple
    }
  case other => other
}
case nonChild: AnyRef => nonChild
case null => null

以上都是適用于有 children 的 node,如果是 children 為 null 的 node 直接返回

2.12、makeCopy(newArgs: Array[AnyRef]): BaseType

反射生成節點副本

2.13、nodeName: String

返回該類型 TreeNode 的 name,默認為 class name;注意,會移除物理操作的 Exec$ 前綴

2.14、innerChildren: Seq[TreeNode[_]]

所有應該以該節點內嵌套樹表示的 nodes,比如,可以被用來表示 sub-queries

2.15、 allChildren: Set[TreeNode[_]]

(children ++ innerChildren).toSet[TreeNode[_]]

2.16、node string 相關

  • 用一行表示該節點
  • 一行更細致的
  • 帶 suffix 的
  • tree 形狀的
  • tree 形狀帶 num 的
  • to json
  • pretty json 等 json 相關的

2.17、apply(number: Int): TreeNode[_]

主要用于交互式 debug,返回該 tree 指定下標的節點,num 可以在 numberedTreeString 找到。最終調用的

private def getNodeNumbered(number: MutableInt): Option[TreeNode[_]] = {
    if (number.i < 0) {
      None
    } else if (number.i == 0) {
      // 返回根節點
      Some(this)
    } else {
      number.i -= 1
      // 注意,此遍歷順序必須與numberedTreeString相同
      innerChildren.map(_.getNodeNumbered(number)).find(_ != None).getOrElse {
        children.map(_.getNodeNumbered(number)).find(_ != None).flatten
      }
    }
  }

我的博客即將搬運同步至騰訊云+社區,邀請大家一同入駐:https://cloud.tencent.com/developer/support-plan?invite_code=x2lzoxh4s5hi

最后編輯于
?著作權歸作者所有,轉載或內容合作請聯系作者
平臺聲明:文章內容(如有圖片或視頻亦包括在內)由作者上傳并發布,文章內容僅代表作者本人觀點,簡書系信息發布平臺,僅提供信息存儲服務。

推薦閱讀更多精彩內容