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