Catalyst规则处理CTE时表名大写不生效问题求助

我编写了如下Catalyst规则代码，期望将指定列表中的表名转换为大写：

def apply(plan: LogicalPlan): LogicalPlan = {
    plan transform {
      case unresolvedRelation: UnresolvedRelation =>
        val tblSchemaName: Array[String] = unresolvedRelation.tableName.split("\\.")
        if (tblSchemaName.length == 1) return plan

        val schema = tblSchemaName.apply(0)
        val tblName = tblSchemaName.apply(1)
        for (ref <- this.refs) {
          if (tblName == ref.nqName) {
            return unresolvedRelation.copy(multipartIdentifier = Seq(schema.toUpperCase, tblName.toUpperCase), unresolvedRelation.options, unresolvedRelation.isStreaming)
          }
        }
        unresolvedRelation

      case unresolvedWith: UnresolvedWith =>
        val newCteRelations = unresolvedWith.cteRelations.map {
          case (aliasName, subqueryAlias) =>
            val newSubqueryAlias = apply(subqueryAlias).asInstanceOf[SubqueryAlias]
            (aliasName, newSubqueryAlias)
        }

        val modifiedCTEUnresolvedWith = unresolvedWith.copy(cteRelations = newCteRelations)
        modifiedCTEUnresolvedWith
        
      case otherPlan: LogicalPlan =>
        val newChildren = otherPlan.children.map(child => apply(child))
        val modified_plan = otherPlan.withNewChildren(newChildren)

        modified_plan
    }
  }
}

rules += UpperCaseTableRule(refs)

val optimizer = new RuleExecutor[LogicalPlan] {
    val batches = Seq(
      Batch("Rewrite", Once, rules.toList: _*)
    )
}

val logicalPlanRewrite = optimizer.execute(logicalPlan)

该规则在普通SELECT语句（如SELECT * FROM schema.t1 a JOIN schema.t2 b ON a.id = b.id）中可正常将表名转为SCHEMA.T1和SCHEMA.T2，但处理CTE语句（如WITH x AS (SELECT * FROM schema.t1),y AS (SELECT * FROM schema.t2) SELECT x.id, y.date FROM x JOIN y ON x.id = y.id）时，调试可见UnresolvedWith分支已将CTE内表名转为大写，但optimizer.execute返回的LogicalPlan中表名又变回小写。

问题原因分析 #

核心问题出在UnresolvedRelation分支中的return语句破坏了Catalyst transform的递归遍历流程：

1. Catalyst的transform方法是基于递归的深度优先遍历，每个case分支应返回修改后的节点，让遍历自动继续处理后续节点。但你在UnresolvedRelation分支中使用return，会直接跳出当前的apply函数调用，中断遍历流程——比如处理完CTE中某个UnresolvedRelation后，直接return会导致后续的CTE节点或主查询节点无法被完整处理，最终修改后的表名在后续遍历中被未处理的节点覆盖。

2. 次要问题：在UnresolvedWith分支中使用asInstanceOf[SubqueryAlias]进行强制类型转换，存在潜在的类型不匹配风险，可能导致后续节点处理异常，间接引发表名回退的问题。

修复方案 #

1. 移除UnresolvedRelation分支中的return语句 #

将分支内的return替换为直接返回节点，让transform的递归遍历正常完成：

case unresolvedRelation: UnresolvedRelation =>
  val tblSchemaName: Array[String] = unresolvedRelation.tableName.split("\\.")
  if (tblSchemaName.length == 1) {
    plan
  } else {
    val schema = tblSchemaName(0)
    val tblName = tblSchemaName(1)
    val matchedRef = refs.find(_.nqName == tblName)
    matchedRef match {
      case Some(_) =>
        unresolvedRelation.copy(
          multipartIdentifier = Seq(schema.toUpperCase, tblName.toUpperCase),
          options = unresolvedRelation.options,
          isStreaming = unresolvedRelation.isStreaming
        )
      case None =>
        unresolvedRelation
    }
  }

（注：这里用find替代原有的for循环，逻辑更简洁且避免不必要的遍历）

2. 优化UnresolvedWith分支的类型处理 #

用模式匹配替代强制类型转换，确保类型安全：

case unresolvedWith: UnresolvedWith =>
  val newCteRelations = unresolvedWith.cteRelations.map {
    case (aliasName, subqueryAlias) =>
      apply(subqueryAlias) match {
        case updatedAlias: SubqueryAlias => (aliasName, updatedAlias)
        case other => throw new IllegalArgumentException(s"CTE subquery must be SubqueryAlias, got ${other.getClass.getSimpleName}")
      }
  }
  unresolvedWith.copy(cteRelations = newCteRelations)

3. 验证遍历流程 #

修复后，transform会完整遍历整个LogicalPlan树：

• 处理UnresolvedWith时，先递归处理所有CTE子查询中的表名，再自动处理主查询部分
• 所有匹配的UnresolvedRelation都会被正确修改为大写，且不会因return中断遍历

内容的提问来源于stack exchange，提问作者user2289345