Review comments

Author: szehon-ho

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/Analyzer.scala

@@ -1672,12 +1672,10 @@ class Analyzer(override val catalogManager: CatalogManager) extends RuleExecutor
       case m @ MergeIntoTable(targetTable, sourceTable, _, _, _, _, _)
         if !m.resolved && targetTable.resolved && sourceTable.resolved =>
 
-        // Do not throw exception for schema evolution case if it has not had a chance to run.
+        // Do not throw exception for schema evolution case.
         // This allows unresolved assignment keys a chance to be resolved by a second pass
         // by newly column/nested fields added by schema evolution.
-        // If schema evolution has already had a chance to run, this will be the final pass
-        val throws = !m.schemaEvolutionEnabled ||
-          (m.canEvaluateSchemaEvolution && !m.schemaChangesNonEmpty)
+        val throws = !m.schemaEvolutionEnabled
 
         EliminateSubqueryAliases(targetTable) match {
           case r: NamedRelation if r.skipSchemaResolution =>

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/ResolveMergeIntoSchemaEvolution.scala

@@ -17,7 +17,6 @@
 
 package org.apache.spark.sql.catalyst.analysis
 
-import org.apache.spark.sql.catalyst.expressions.{Attribute, AttributeReference, AttributeSet, Expression, GetStructField}
 import org.apache.spark.sql.catalyst.plans.logical._
 import org.apache.spark.sql.catalyst.rules.Rule
 import org.apache.spark.sql.catalyst.types.DataTypeUtils
@@ -37,88 +36,17 @@ object ResolveMergeIntoSchemaEvolution extends Rule[LogicalPlan] {
   override def apply(plan: LogicalPlan): LogicalPlan = plan resolveOperators {
     // This rule should run only if all assignments are resolved, except those
     // that will be satisfied by schema evolution
-    case m @ MergeIntoTable(_, _, _, _, _, _, _) if m.needSchemaEvolution =>
-        val newTarget = m.targetTable.transform {
-          case r : DataSourceV2Relation => performSchemaEvolution(r, m)
-        }
-
-        // Unresolve all references based on old target output
-        val targetOutput = m.targetTable.output
-        val unresolvedMergeCondition = unresolveCondition(m.mergeCondition, targetOutput)
-        val unresolvedMatchedActions = unresolveActions(m.matchedActions, targetOutput)
-        val unresolvedNotMatchedActions = unresolveActions(m.notMatchedActions, targetOutput)
-        val unresolvedNotMatchedBySourceActions =
-          unresolveActions(m.notMatchedBySourceActions, targetOutput)
-
-        m.copy(
-          targetTable = newTarget,
-          mergeCondition = unresolvedMergeCondition,
-          matchedActions = unresolvedMatchedActions,
-          notMatchedActions = unresolvedNotMatchedActions,
-          notMatchedBySourceActions = unresolvedNotMatchedBySourceActions)
-  }
-
-  private def unresolveActions(actions: Seq[MergeAction], output: Seq[Attribute]):
-  Seq[MergeAction] = {
-    actions.map {
-      case UpdateAction(condition, assignments) =>
-        UpdateAction(condition.map(unresolveCondition(_, output)),
-          unresolveAssignmentKeys(assignments))
-      case InsertAction(condition, assignments) =>
-        InsertAction(condition.map(unresolveCondition(_, output)),
-          unresolveAssignmentKeys(assignments))
-      case DeleteAction(condition) =>
-        DeleteAction(condition.map(unresolveCondition(_, output)))
-      case other => other
-    }
-  }
-
-  private def unresolveCondition(expr: Expression, output: Seq[Attribute]): Expression = {
-    val outputSet = AttributeSet(output)
-    expr.transform {
-      case attr: AttributeReference if outputSet.contains(attr) =>
-        val nameParts = if (attr.qualifier.nonEmpty) {
-          attr.qualifier ++ Seq(attr.name)
-        } else {
-          Seq(attr.name)
-        }
-        UnresolvedAttribute(nameParts)
-    }
-  }
-
-  private def unresolveAssignmentKeys(assignments: Seq[Assignment]): Seq[Assignment] = {
-    assignments.map { assignment =>
-      val unresolvedKey = assignment.key match {
-        case _: UnresolvedAttribute => assignment.key
-        case gsf: GetStructField =>
-          // Recursively collect all nested GetStructField names and the base AttributeReference
-          val nameParts = collectStructFieldNames(gsf)
-          nameParts match {
-            case Some(names) => UnresolvedAttribute(names)
-            case None => assignment.key
-          }
-        case attr: AttributeReference =>
-          UnresolvedAttribute(Seq(attr.name))
-        case attr: Attribute =>
-          UnresolvedAttribute(Seq(attr.name))
-        case other => other
+    case m@MergeIntoTable(_, _, _, _, _, _, _) if m.needSchemaEvolution =>
+      m transformUpWithNewOutput {
+        case r: DataSourceV2Relation =>
+          val newTarget = performSchemaEvolution(r, m)
+          val oldTargetOutput = m.targetTable.output
+          val newTargetOutput = newTarget.output
+          val attributeMapping = oldTargetOutput.map(
+            oldAttr => (oldAttr, newTargetOutput.find(_.name == oldAttr.name).getOrElse(oldAttr))
+          )
+          newTarget -> attributeMapping
       }
-      Assignment(unresolvedKey, assignment.value)
-    }
-  }
-
-  private def collectStructFieldNames(expr: Expression): Option[Seq[String]] = {
-    expr match {
-      case GetStructField(child, _, Some(fieldName)) =>
-        collectStructFieldNames(child) match {
-          case Some(childNames) => Some(childNames :+ fieldName)
-          case None => None
-        }
-      case attr: AttributeReference =>
-        Some(Seq(attr.name))
-      case _ =>
-        None
-    }
   }
 
   private def performSchemaEvolution(relation: DataSourceV2Relation, m: MergeIntoTable)