[SPARK-54172][SQL] Merge Into Schema Evolution should only add referenced columns

Author: szehon-ho

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

@@ -1669,7 +1669,7 @@ class Analyzer(override val catalogManager: CatalogManager) extends RuleExecutor
 
       case u: UpdateTable => resolveReferencesInUpdate(u)
 
-      case m @ MergeIntoTable(targetTable, sourceTable, _, _, _, _, _)
+      case m @ MergeIntoTable(targetTable, sourceTable, _, _, _, _, _, _)
         if !m.resolved && targetTable.resolved && sourceTable.resolved && !m.needSchemaEvolution =>
 
         EliminateSubqueryAliases(targetTable) match {

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

@@ -34,24 +34,26 @@ import org.apache.spark.sql.execution.datasources.v2.DataSourceV2Relation
 object ResolveMergeIntoSchemaEvolution extends Rule[LogicalPlan] {
 
   override def apply(plan: LogicalPlan): LogicalPlan = plan resolveOperators {
-    case m @ MergeIntoTable(_, _, _, _, _, _, _)
+    case m @ MergeIntoTable(_, _, _, _, _, _, _, _)
       if m.needSchemaEvolution =>
         val newTarget = m.targetTable.transform {
-          case r : DataSourceV2Relation => performSchemaEvolution(r, m.sourceTable)
+          case r : DataSourceV2Relation => performSchemaEvolution(r, m)
         }
         m.copy(targetTable = newTarget)
   }
 
-  private def performSchemaEvolution(relation: DataSourceV2Relation, source: LogicalPlan)
+  private def performSchemaEvolution(relation: DataSourceV2Relation, m: MergeIntoTable)
     : DataSourceV2Relation = {
     (relation.catalog, relation.identifier) match {
       case (Some(c: TableCatalog), Some(i)) =>
-        val changes = MergeIntoTable.schemaChanges(relation.schema, source.schema)
+        val referencedSourceSchema = MergeIntoTable.referencedSourceSchema(m)
+
+        val changes = MergeIntoTable.schemaChanges(relation.schema, referencedSourceSchema)
         c.alterTable(i, changes: _*)
         val newTable = c.loadTable(i)
         val newSchema = CatalogV2Util.v2ColumnsToStructType(newTable.columns())
         // Check if there are any remaining changes not applied.
-        val remainingChanges = MergeIntoTable.schemaChanges(newSchema, source.schema)
+        val remainingChanges = MergeIntoTable.schemaChanges(newSchema, referencedSourceSchema)
         if (remainingChanges.nonEmpty) {
           throw QueryCompilationErrors.unsupportedTableChangesInAutoSchemaEvolutionError(
             remainingChanges, i.toQualifiedNameParts(c))

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

@@ -60,11 +60,18 @@ object ResolveRowLevelCommandAssignments extends Rule[LogicalPlan] {
         notMatchedActions = alignActions(m.targetTable.output, m.notMatchedActions,
           coerceNestedTypes),
         notMatchedBySourceActions = alignActions(m.targetTable.output, m.notMatchedBySourceActions,
-          coerceNestedTypes))
+          coerceNestedTypes),
+        preservedSourceActions = Some(m.matchedActions ++ m.notMatchedActions)
+      )
 
     case m: MergeIntoTable if !m.skipSchemaResolution && m.resolved && !m.aligned
       && !m.needSchemaEvolution =>
-      resolveAssignments(m)
+      m.copy(
+        matchedActions = m.notMatchedActions.map(resolveMergeAction),
+        notMatchedActions = m.notMatchedActions.map(resolveMergeAction),
+        notMatchedBySourceActions = m.matchedActions.map(resolveMergeAction),
+        preservedSourceActions = Some(m.matchedActions ++ m.notMatchedActions)
+      )
   }
 
   private def validateStoreAssignmentPolicy(): Unit = {
@@ -83,33 +90,51 @@ object ResolveRowLevelCommandAssignments extends Rule[LogicalPlan] {
 
   private def resolveAssignments(p: LogicalPlan): LogicalPlan = {
     p.transformExpressions {
-      case assignment: Assignment =>
-        val nullHandled = if (!assignment.key.nullable && assignment.value.nullable) {
-          AssertNotNull(assignment.value)
-        } else {
-          assignment.value
-        }
-        val casted = if (assignment.key.dataType != nullHandled.dataType) {
-          val cast = Cast(nullHandled, assignment.key.dataType, ansiEnabled = true)
-          cast.setTagValue(Cast.BY_TABLE_INSERTION, ())
-          cast
-        } else {
-          nullHandled
-        }
-        val rawKeyType = assignment.key.transform {
-          case a: AttributeReference =>
-            CharVarcharUtils.getRawType(a.metadata).map(a.withDataType).getOrElse(a)
-        }.dataType
-        val finalValue = if (CharVarcharUtils.hasCharVarchar(rawKeyType)) {
-          CharVarcharUtils.stringLengthCheck(casted, rawKeyType)
-        } else {
-          casted
-        }
-        val cleanedKey = assignment.key.transform {
-          case a: AttributeReference => CharVarcharUtils.cleanAttrMetadata(a)
-        }
-        Assignment(cleanedKey, finalValue)
+      case assignment: Assignment => resolveAssignment(assignment)
+    }
+  }
+
+  private def resolveMergeAction(mergeAction: MergeAction) = {
+    mergeAction match {
+      case u @ UpdateAction(_, assignments) =>
+        u.copy(assignments = assignments.map(resolveAssignment))
+      case i @ InsertAction(_, assignments) =>
+        i.copy(assignments = assignments.map(resolveAssignment))
+      case d: DeleteAction =>
+        d
+      case other =>
+        throw new AnalysisException(
+          errorClass = "_LEGACY_ERROR_TEMP_3053",
+          messageParameters = Map("other" -> other.toString))
+    }
+  }
+
+  private def resolveAssignment(assignment: Assignment) = {
+    val nullHandled = if (!assignment.key.nullable && assignment.value.nullable) {
+      AssertNotNull(assignment.value)
+    } else {
+      assignment.value
+    }
+    val casted = if (assignment.key.dataType != nullHandled.dataType) {
+      val cast = Cast(nullHandled, assignment.key.dataType, ansiEnabled = true)
+      cast.setTagValue(Cast.BY_TABLE_INSERTION, ())
+      cast
+    } else {
+      nullHandled
+    }
+    val rawKeyType = assignment.key.transform {
+      case a: AttributeReference =>
+        CharVarcharUtils.getRawType(a.metadata).map(a.withDataType).getOrElse(a)
+    }.dataType
+    val finalValue = if (CharVarcharUtils.hasCharVarchar(rawKeyType)) {
+      CharVarcharUtils.stringLengthCheck(casted, rawKeyType)
+    } else {
+      casted
+    }
+    val cleanedKey = assignment.key.transform {
+      case a: AttributeReference => CharVarcharUtils.cleanAttrMetadata(a)
     }
+    Assignment(cleanedKey, finalValue)
   }
 
   private def alignActions(

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

@@ -45,7 +45,7 @@ object RewriteMergeIntoTable extends RewriteRowLevelCommand with PredicateHelper
 
   override def apply(plan: LogicalPlan): LogicalPlan = plan resolveOperators {
     case m @ MergeIntoTable(aliasedTable, source, cond, matchedActions, notMatchedActions,
-      notMatchedBySourceActions, _) if m.resolved && m.rewritable && m.aligned &&
+      notMatchedBySourceActions, _, _) if m.resolved && m.rewritable && m.aligned &&
         !m.needSchemaEvolution && matchedActions.isEmpty && notMatchedActions.size == 1 &&
         notMatchedBySourceActions.isEmpty =>
 
@@ -79,7 +79,7 @@ object RewriteMergeIntoTable extends RewriteRowLevelCommand with PredicateHelper
       }
 
     case m @ MergeIntoTable(aliasedTable, source, cond, matchedActions, notMatchedActions,
-        notMatchedBySourceActions, _)
+        notMatchedBySourceActions, _, _)
       if m.resolved && m.rewritable && m.aligned && !m.needSchemaEvolution &&
         matchedActions.isEmpty && notMatchedBySourceActions.isEmpty =>
 
@@ -121,7 +121,7 @@ object RewriteMergeIntoTable extends RewriteRowLevelCommand with PredicateHelper
       }
 
     case m @ MergeIntoTable(aliasedTable, source, cond, matchedActions, notMatchedActions,
-        notMatchedBySourceActions, _)
+        notMatchedBySourceActions, _, _)
       if m.resolved && m.rewritable && m.aligned && !m.needSchemaEvolution =>
 
       EliminateSubqueryAliases(aliasedTable) match {

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/v2Commands.scala

@@ -19,7 +19,7 @@ package org.apache.spark.sql.catalyst.plans.logical
 
 import org.apache.spark.{SparkIllegalArgumentException, SparkUnsupportedOperationException}
 import org.apache.spark.sql.AnalysisException
-import org.apache.spark.sql.catalyst.analysis.{AnalysisContext, AssignmentUtils, EliminateSubqueryAliases, FieldName, NamedRelation, PartitionSpec, ResolvedIdentifier, ResolvedProcedure, TypeCheckResult, UnresolvedException, UnresolvedProcedure, ViewSchemaMode}
+import org.apache.spark.sql.catalyst.analysis.{AnalysisContext, AssignmentUtils, EliminateSubqueryAliases, FieldName, NamedRelation, PartitionSpec, ResolvedIdentifier, ResolvedProcedure, TypeCheckResult, UnresolvedAttribute, UnresolvedException, UnresolvedProcedure, ViewSchemaMode}
 import org.apache.spark.sql.catalyst.analysis.TypeCheckResult.{DataTypeMismatch, TypeCheckSuccess}
 import org.apache.spark.sql.catalyst.catalog.{FunctionResource, RoutineLanguage}
 import org.apache.spark.sql.catalyst.catalog.CatalogTypes.TablePartitionSpec
@@ -860,7 +860,10 @@ case class MergeIntoTable(
     matchedActions: Seq[MergeAction],
     notMatchedActions: Seq[MergeAction],
     notMatchedBySourceActions: Seq[MergeAction],
-    withSchemaEvolution: Boolean) extends BinaryCommand with SupportsSubquery {
+    withSchemaEvolution: Boolean,
+    // Preserves original pre-aligned actions for source matches
+    preservedSourceActions: Option[Seq[MergeAction]] = None)
+  extends BinaryCommand with SupportsSubquery {
 
   lazy val aligned: Boolean = {
     val actions = matchedActions ++ notMatchedActions ++ notMatchedBySourceActions
@@ -892,9 +895,13 @@ case class MergeIntoTable(
     case _ => false
   }
 
-  lazy val needSchemaEvolution: Boolean =
+  private lazy val migrationSchema: StructType =
+    MergeIntoTable.referencedSourceSchema(this)
+
+  lazy val needSchemaEvolution: Boolean = {
     schemaEvolutionEnabled &&
-      MergeIntoTable.schemaChanges(targetTable.schema, sourceTable.schema).nonEmpty
+      MergeIntoTable.schemaChanges(targetTable.schema, migrationSchema).nonEmpty
+  }
 
   private def schemaEvolutionEnabled: Boolean = withSchemaEvolution && {
     EliminateSubqueryAliases(targetTable) match {
@@ -948,11 +955,12 @@ object MergeIntoTable {
             case currentField: StructField if newFieldMap.contains(currentField.name) =>
               schemaChanges(currentField.dataType, newFieldMap(currentField.name).dataType,
                 originalTarget, originalSource, fieldPath ++ Seq(currentField.name))
-          }}.flatten
+          }
+        }.flatten
 
         // Identify the newly added fields and append to the end
         val currentFieldMap = toFieldMap(currentFields)
-        val adds = newFields.filterNot (f => currentFieldMap.contains (f.name))
+        val adds = newFields.filterNot(f => currentFieldMap.contains(f.name))
           .map(f => TableChange.addColumn(fieldPath ++ Set(f.name), f.dataType))
 
         updates ++ adds
@@ -990,8 +998,81 @@ object MergeIntoTable {
       CaseInsensitiveMap(fieldMap)
     }
   }
+
+  // Filter the source schema to retain only fields that are referenced
+  // by at least one merge action
+  def referencedSourceSchema(merge: MergeIntoTable): StructType = {
+
+    val actions = merge.preservedSourceActions match {
+      case Some(preserved) => preserved
+      case None => merge.matchedActions ++ merge.notMatchedActions
+    }
+
+    val assignments = actions.collect {
+      case a: UpdateAction => a.assignments.map(_.key)
+      case a: InsertAction => a.assignments.map(_.key)
+    }.flatten
+
+    val containsStarAction = actions.exists {
+      case _: UpdateStarAction => true
+      case _: InsertStarAction => true
+      case _ => false
+    }
+
+    def filterSchema(sourceSchema: StructType, basePath: Seq[String]): StructType =
+      StructType(sourceSchema.flatMap { field =>
+        val fieldPath = basePath :+ field.name
+
+        field.dataType match {
+          // Specifically assigned to in one clause:
+          // always keep, including all nested attributes
+          case _ if assignments.exists(isEqual(_, fieldPath)) => Some(field)
+          // If this is a struct and one of the children is being assigned to in a merge clause,
+          // keep it and continue filtering children.
+          case struct: StructType if assignments.exists(assign =>
+            isPrefix(fieldPath, extractFieldPath(assign))) =>
+            Some(field.copy(dataType = filterSchema(struct, fieldPath)))
+          // The field isn't assigned to directly or indirectly (i.e. its children) in any non-*
+          // clause. Check if it should be kept with any * action.
+          case struct: StructType if containsStarAction =>
+            Some(field.copy(dataType = filterSchema(struct, fieldPath)))
+          case _ if containsStarAction => Some(field)
+          // The field and its children are not assigned to in any * or non-* action, drop it.
+          case _ => None
+        }
+      })
+
+    val sourceSchema = merge.sourceTable.schema
+    val targetSchema = merge.targetTable.schema
+    val res = filterSchema(merge.sourceTable.schema, Seq.empty)
+    res
+  }
+
+  // Helper method to extract field path from an Expression.
+  private def extractFieldPath(expr: Expression): Seq[String] = expr match {
+    case UnresolvedAttribute(nameParts) => nameParts
+    case a: AttributeReference => Seq(a.name)
+    case GetStructField(child, ordinal, nameOpt) =>
+      extractFieldPath(child) :+ nameOpt.getOrElse(s"col$ordinal")
+    case _ => Seq.empty
+  }
+
+  // Helper method to check if a given field path is a prefix of another path. Delegates
+  // equality to conf.resolver to correctly handle case sensitivity.
+  private def isPrefix(prefix: Seq[String], path: Seq[String]): Boolean =
+    prefix.length <= path.length && prefix.zip(path).forall {
+      case (prefixNamePart, pathNamePart) =>
+        SQLConf.get.resolver(prefixNamePart, pathNamePart)
+    }
+
+  // Helper method to check if an assignment Expression's field path is equal to a path.
+  def isEqual(assignmentExpr: Expression, path: Seq[String]): Boolean = {
+    val exprPath = extractFieldPath(assignmentExpr)
+    exprPath.length == path.length && isPrefix(exprPath, path)
+  }
 }
 
+
 sealed abstract class MergeAction extends Expression with Unevaluable {
   def condition: Option[Expression]
   override def nullable: Boolean = false

sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/PullupCorrelatedPredicatesSuite.scala

@@ -167,7 +167,7 @@ class PullupCorrelatedPredicatesSuite extends PlanTest {
     assert(optimized.resolved)
 
     optimized match {
-      case MergeIntoTable(_, _, s: InSubquery, _, _, _, _) =>
+      case MergeIntoTable(_, _, s: InSubquery, _, _, _, _, _) =>
         val outerRefs = SubExprUtils.getOuterReferences(s.query.plan)
         assert(outerRefs.isEmpty, "should be no outer refs")
       case other =>