零、前置知識 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