如何在Spark SQL逻辑计划中仅替换FROM子句表名（排除字面量）

给定Spark SQL语句：

SELECT id, "delta.`/example/table/path`" FROM delta.`/example/table/path` WHERE str LIKE "%delta.`/example/table/path`"

需要仅替换FROM子句中的表引用delta./example/table/path``为new_table_name，不能修改SELECT和WHERE子句中作为字符串字面量的相同内容。此前通过文本匹配替换会误改字符串字面量，因此需要通过Spark逻辑计划树实现精准替换。

Spark的逻辑计划树已经明确区分了表引用（LogicalRelation节点）和字符串字面量（Literal节点），因此可以通过遍历并修改逻辑计划中的目标LogicalRelation节点，再将修改后的计划转回SQL语句，实现精准替换。

1. 解析原SQL为逻辑计划 #

使用Spark的SQL解析器和分析器生成逻辑计划：

val plan = spark.sessionState.sqlParser.parsePlan(query)
val analyzedPlan = spark.sessionState.analyzer.executeAndCheck(plan, new QueryPlanningTracker)

2. 编写逻辑计划替换规则 #

自定义Rule[LogicalPlan]，遍历计划树找到目标表对应的LogicalRelation，替换为新表的引用：

import org.apache.spark.sql.catalyst.plans.logical.{LogicalPlan, LogicalRelation, SubqueryAlias}
import org.apache.spark.sql.catalyst.rules.Rule
import org.apache.spark.sql.catalyst.TableIdentifier
import org.apache.spark.sql.catalyst.catalog.CatalogTable

class ReplaceTargetTableRule(targetTablePath: String, newTableName: String) extends Rule[LogicalPlan] {
  override def apply(plan: LogicalPlan): LogicalPlan = plan.transform {
    // 处理带别名的表引用
    case alias: SubqueryAlias if isTargetTable(alias.child, targetTablePath) =>
      val newTableId = TableIdentifier(newTableName)
      val newCatalogTable = spark.sessionState.catalog.getTableMetadata(newTableId)
      val newRelation = LogicalRelation(
        spark.sessionState.catalog.getTable(newTableId).table,
        alias.child.output,
        Some(newCatalogTable),
        isStreaming = false
      )
      SubqueryAlias(alias.alias, newRelation)

    // 处理直接的表引用
    case relation: LogicalRelation if isTargetTable(relation, targetTablePath) =>
      val newTableId = TableIdentifier(newTableName)
      val newCatalogTable = spark.sessionState.catalog.getTableMetadata(newTableId)
      LogicalRelation(
        spark.sessionState.catalog.getTable(newTableId).table,
        relation.output,
        Some(newCatalogTable),
        isStreaming = false
      )
  }

  // 判断当前节点是否为目标表（通过存储路径匹配）
  private def isTargetTable(plan: LogicalPlan, targetPath: String): Boolean = plan match {
    case lr: LogicalRelation =>
      lr.catalogTable.exists(ct => ct.storage.locationUri.exists(_.getPath == targetPath))
    case _ => false
  }
}

3. 应用规则并生成最终SQL #

将规则应用到分析后的逻辑计划，再将修改后的计划转回SQL：

// 原SQL
val originalSql = """SELECT id, "delta.`/example/table/path`" FROM delta.`/example/table/path` WHERE str LIKE "%delta.`/example/table/path`""""

// 生成分析后的逻辑计划
val parsedPlan = spark.sessionState.sqlParser.parsePlan(originalSql)
val analyzedPlan = spark.sessionState.analyzer.executeAndCheck(parsedPlan, new QueryPlanningTracker)

// 替换目标表
val rule = new ReplaceTargetTableRule("/example/table/path", "new_table_name")
val modifiedPlan = rule.apply(analyzedPlan)

// 生成最终SQL
val finalSql = spark.sessionState.sqlGenerator.generateSql(modifiedPlan)

// 输出结果：SELECT id, "delta.`/example/table/path`" FROM new_table_name WHERE str LIKE "%delta.`/example/table/path`"
println(finalSql)

• 确保new_table_name已在Spark Catalog中注册，若为路径表，需调整LogicalRelation的构建逻辑（直接使用DataSource而非Catalog表）。
• 若目标表有多层SubqueryAlias嵌套，transform方法会递归遍历所有子节点，无需额外处理。
• 可根据实际需求调整isTargetTable的匹配逻辑，比如通过表的标识符（identifier）而非路径匹配。

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