spark sql解析过程中对tree的遍历（源码详解）

SELECT A,B FROM TESTDATA2 WHERE A>2

def transformUp(rule: PartialFunction[BaseType, BaseType]): BaseType = {
  // 先遍历子节点，得到叶子节点
  val afterRuleOnChildren = mapChildren(_.transformUp(rule))
 
 //对节点执行规则
  val newNode = if (this fastEquals afterRuleOnChildren) {
    CurrentOrigin.withOrigin(origin) {
      //这里用到了PartialFunction的applyOrElse方法，用来避免undefined的情况发生。如果当前节点应用rule没有匹配的话，则返回默认的当前节点本身
      rule.applyOrElse(this, identity[BaseType])
    }
  } else {
    CurrentOrigin.withOrigin(origin) {
      rule.applyOrElse(afterRuleOnChildren, identity[BaseType])
    }
  }
  // If the transform function replaces this node with a new one, carry over the tags.
  newNode.copyTagsFrom(this)
  newNode
}

def transformDown(rule: PartialFunction[BaseType, BaseType]): BaseType = {
      // 对当前节点，调用rule函数。
    val afterRule = CurrentOrigin.withOrigin(origin) {
    // 这里rule函数有可能会生成新的节点，新节点的子节点可能不一样
    rule.applyOrElse(this, identity[BaseType])
  }

  // Check if unchanged and then possibly return old copy to avoid gc churn.
 //再遍历子节点 
 if (this fastEquals afterRule) {
// 如果当前节点没有变化，则继续遍历它的子节点
    mapChildren(_.transformDown(rule))
  } else {
  
// 如果当前节点发生改变，需要对改变后的节点进行遍历
    afterRule.copyTagsFrom(this)
    afterRule.mapChildren(_.transformDown(rule))
  }
}

def mapChildren(f: BaseType => BaseType): BaseType = {
//如果不是叶子节点，那么会执行mapChildren(f, forceCopy = false)方法，遍历构造函数的参数。如果参数是子节点，那么递归遍历
if (containsChild.nonEmpty) {
    mapChildren(f, forceCopy = false)
  } else {
//如果是叶子节点，则返回自身节点 
    this
  }
}


private def mapChildren(
    f: BaseType => BaseType,
    forceCopy: Boolean): BaseType = {
  var changed = false

  def mapChild(child: Any): Any = child match {
    case arg: TreeNode[_] if containsChild(arg) =>
      val newChild = f(arg.asInstanceOf[BaseType])
      if (forceCopy || !(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 (forceCopy || !(newChild1 fastEquals arg1) || !(newChild2 fastEquals arg2)) {
        changed = true
        (newChild1, newChild2)
      } else {
        tuple
      }
    case other => other
  }

// 调用了mapProductIterator方法，遍历构造函数的参数，返回新的构造参数
  val newArgs = mapProductIterator {
// 如果参数是TreeNode子类，并且是当前节点的子节点
    case arg: TreeNode[_] if containsChild(arg) =>
// 递归调用函数遍历 这里的f可能是 transformUp or transformDown
      val newChild = f(arg.asInstanceOf[BaseType])
// 如果子节点发生变化了，则更改changed标识
      if (forceCopy || !(newChild fastEquals arg)) {
        changed = true
        newChild
      } else {
        arg
      }
    case Some(arg: TreeNode[_]) if containsChild(arg) =>
      val newChild = f(arg.asInstanceOf[BaseType])
      if (forceCopy || !(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 (forceCopy || !(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: Stream[_] => args.map(mapChild).force // Force materialization on stream
    case args: Iterable[_] => args.map(mapChild)
    case nonChild: AnyRef => nonChild
    case null => null
  }
 // 如果子节点发生变化，则利用新的构造参数，实例化新的节点
  if (forceCopy || changed) makeCopy(newArgs, forceCopy) else this
}

//ClassTag用法
 def mapProductIterator[B: ClassTag](f: Any => B): Array[B] = {
// protected def mapProductIterator[B: ClassTag](f: Any => B): Array[B] = {
  val arr = Array.ofDim[B](productArity)
  var i = 0
  while (i < arr.length) {
    arr(i) = f(productElement(i))
    i += 1
  }
  arr
}

SELECT A,B FROM TESTDATA2 WHERE A>2

== Parsed Logical Plan ==
'Project ['A, 'B]
+- 'Filter ('A > 2)
   +- 'UnresolvedRelation [TESTDATA2]

// 入参为rule，偏函数
def resolveOperatorsUp(rule: PartialFunction[LogicalPlan, LogicalPlan]): LogicalPlan = {
  if (!analyzed) {
    AnalysisHelper.allowInvokingTransformsInAnalyzer {
     // 1、先遍历子节点，得到叶子节点
      val afterRuleOnChildren = mapChildren(_.resolveOperatorsUp(rule))
      
    //2、为节点执行规则
    if (self fastEquals afterRuleOnChildren) {
        CurrentOrigin.withOrigin(origin) {
    // 如果遍历后当前节点没有发生变化，对当前的plan执行rule规则
          rule.applyOrElse(self, identity[LogicalPlan])
        }
      } else {
        CurrentOrigin.withOrigin(origin) {
    // 如果遍历后 当前 节点发 生了变化，则新负值的afterRuleOnChildren执行rule规则
          rule.applyOrElse(afterRuleOnChildren, identity[LogicalPlan])
        }
      }
    }
  } else {
    self
  }
}

// 该规则主要是针对Hint节点起作用 ，目前是UnresolvedRelation节点
def apply(plan: LogicalPlan): LogicalPlan = plan resolveOperatorsUp {
  case h: UnresolvedHint if STRATEGY_HINT_NAMES.contains(h.name.toUpperCase(Locale.ROOT)) =>
    if (h.parameters.isEmpty) {
      // If there is no table alias specified, apply the hint on the entire subtree.
      ResolvedHint(h.child, createHintInfo(h.name))
    } else {
      // Otherwise, find within the subtree query plans to apply the hint.
      val relationNamesInHint = h.parameters.map {
        case tableName: String => UnresolvedAttribute.parseAttributeName(tableName)
        case tableId: UnresolvedAttribute => tableId.nameParts
        case unsupported => throw new AnalysisException("Join strategy hint parameter " +
          s"should be an identifier or string but was $unsupported (${unsupported.getClass}")
      }.toSet
      val relationsInHintWithMatch = new mutable.HashSet[Seq[String]]
      val applied = applyJoinStrategyHint(
        h.child, relationNamesInHint, relationsInHintWithMatch, h.name)

      // Filters unmatched relation identifiers in the hint
      val unmatchedIdents = relationNamesInHint -- relationsInHintWithMatch
      hintErrorHandler.hintRelationsNotFound(h.name, h.parameters, unmatchedIdents)
      applied
    }
}

  rule.applyOrElse(self, identity[LogicalPlan])

/** Similar to [[resolveOperatorsUp]], but does it top-down. */
def resolveOperatorsDown(rule: PartialFunction[LogicalPlan, LogicalPlan]): LogicalPlan = {
  if (!analyzed) {
    AnalysisHelper.allowInvokingTransformsInAnalyzer {   
      val afterRule = CurrentOrigin.withOrigin(origin) {
       // 1、为当前节点执行规则
        rule.applyOrElse(self, identity[LogicalPlan])
      }

      //  2、对执行完规则后的新节点遍历迭代
      if (self fastEquals afterRule) {
        //如果执行完规则后的节点没有变化（即规则没有起到作用），则对节点遍历迭代
        mapChildren(_.resolveOperatorsDown(rule))
      } else {
       //如果执行完规则后的节点发生变化，则对新节点遍历迭代
        afterRule.mapChildren(_.resolveOperatorsDown(rule))
      }
    }
  } else {
    self
  }
}