[SPARK-30814][SQL] ALTER TABLE ... ADD COLUMN position should be able to reference columns being added

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

@@ -3023,9 +3023,27 @@ class Analyzer(
   object ResolveAlterTableChanges extends Rule[LogicalPlan] {
     def apply(plan: LogicalPlan): LogicalPlan = plan.resolveOperatorsUp {
       case a @ AlterTable(_, _, t: NamedRelation, changes) if t.resolved =>
+        // Keeps track of new columns being added. Keys are normalized parent names and
+        // values are field names that belong to their parent.
+        val colsToAdd = mutable.Map.empty[Seq[String], Seq[String]]
         val schema = t.schema
         val normalizedChanges = changes.flatMap {
           case add: AddColumn =>
+            def addColumn(
+                parentSchema: StructType,
+                parentName: String,
+                normalizedParentName: Seq[String]): TableChange = {
+              val fieldsAdded = colsToAdd.get(normalizedParentName).getOrElse(Nil)
+              val pos = findColumnPosition(add.position(), parentName, parentSchema, fieldsAdded)
+              val field = add.fieldNames().last
+              colsToAdd(normalizedParentName) = fieldsAdded :+ field
+              TableChange.addColumn(
+                (normalizedParentName :+ field).toArray,
+                add.dataType(),
+                add.isNullable,
+                add.comment,
+                pos)
+            }
             val parent = add.fieldNames().init
             if (parent.nonEmpty) {
               // Adding a nested field, need to normalize the parent column and position
@@ -3037,27 +3055,14 @@ class Analyzer(
                 val (normalizedName, sf) = target.get
                 sf.dataType match {
                   case struct: StructType =>
-                    val pos = findColumnPosition(add.position(), parent.quoted, struct)
-                    Some(TableChange.addColumn(
-                      (normalizedName ++ Seq(sf.name, add.fieldNames().last)).toArray,
-                      add.dataType(),
-                      add.isNullable,
-                      add.comment,
-                      pos))
-
+                    Some(addColumn(struct, parent.quoted, normalizedName :+ sf.name))
                   case other =>
                     Some(add)
                 }
               }
             } else {
               // Adding to the root. Just need to normalize position
-              val pos = findColumnPosition(add.position(), "root", schema)
-              Some(TableChange.addColumn(
-                add.fieldNames(),
-                add.dataType(),
-                add.isNullable,
-                add.comment,
-                pos))
+              Some(addColumn(schema, "root", Nil))
             }
 
           case typeChange: UpdateColumnType =>
@@ -3156,17 +3161,18 @@ class Analyzer(
 
     private def findColumnPosition(
         position: ColumnPosition,
-        field: String,
-        struct: StructType): ColumnPosition = {
+        parentName: String,
+        struct: StructType,
+        fieldsAdded: Seq[String]): ColumnPosition = {
       position match {
         case null => null
         case after: After =>
-          struct.fieldNames.find(n => conf.resolver(n, after.column())) match {
+          (struct.fieldNames ++ fieldsAdded).find(n => conf.resolver(n, after.column())) match {
             case Some(colName) =>
               ColumnPosition.after(colName)
             case None =>
               throw new AnalysisException("Couldn't find the reference column for " +
-                s"$after at $field")
+                s"$after at $parentName")
           }
         case other => other
       }

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

@@ -474,6 +474,7 @@ trait CheckAnalysis extends PredicateHelper {
             }
 
             val colsToDelete = mutable.Set.empty[Seq[String]]
+            val colsToAdd = mutable.Set.empty[Seq[String]]
 
             alter.changes.foreach {
               case add: AddColumn =>
@@ -483,8 +484,16 @@ trait CheckAnalysis extends PredicateHelper {
                   checkColumnNotExists("add", add.fieldNames(), table.schema)
                 }
                 val parent = findParentStruct("add", add.fieldNames())
-                positionArgumentExists(add.position(), parent)
+                add.position match {
+                  case after: After =>
+                    // Handle the case where column position is referencing new columns being added.
+                    if (!colsToAdd.contains(add.fieldNames().init :+ after.column)) {
+                      positionArgumentExists(add.position(), parent)
+                    }
+                  case _ =>
+                }
                 TypeUtils.failWithIntervalType(add.dataType())
+                colsToAdd += add.fieldNames()
               case update: UpdateColumnType =>
                 val field = findField("update", update.fieldNames)
                 val fieldName = update.fieldNames.quoted

sql/core/src/test/scala/org/apache/spark/sql/connector/AlterTableTests.scala

@@ -173,6 +173,42 @@ trait AlterTableTests extends SharedSparkSession {
     }
   }
 
+  test("SPARK-30814: add column with position referencing new columns being added") {
+    val t = s"${catalogAndNamespace}table_name"
+    withTable(t) {
+      sql(s"CREATE TABLE $t (a string, b int, point struct<x: double, y: double>) USING $v2Format")
+      sql(s"ALTER TABLE $t ADD COLUMNS (x int AFTER a, y int AFTER x, z int AFTER y)")
+
+      assert(getTableMetadata(t).schema === new StructType()
+        .add("a", StringType)
+        .add("x", IntegerType)
+        .add("y", IntegerType)
+        .add("z", IntegerType)
+        .add("b", IntegerType)
+        .add("point", new StructType()
+          .add("x", DoubleType)
+          .add("y", DoubleType)))
+
+      sql(s"ALTER TABLE $t ADD COLUMNS (point.z double AFTER x, point.zz double AFTER z)")
+      assert(getTableMetadata(t).schema === new StructType()
+        .add("a", StringType)
+        .add("x", IntegerType)
+        .add("y", IntegerType)
+        .add("z", IntegerType)
+        .add("b", IntegerType)
+        .add("point", new StructType()
+          .add("x", DoubleType)
+          .add("z", DoubleType)
+          .add("zz", DoubleType)
+          .add("y", DoubleType)))
+
+      // The new column being referenced should come before being referenced.
+      val e = intercept[AnalysisException](
+        sql(s"ALTER TABLE $t ADD COLUMNS (yy int AFTER xx, xx int)"))
+      assert(e.getMessage().contains("Couldn't find the reference column for AFTER xx at root"))
+    }
+  }
+
   test("AlterTable: add multiple columns") {
     val t = s"${catalogAndNamespace}table_name"
     withTable(t) {

sql/core/src/test/scala/org/apache/spark/sql/connector/V2CommandsCaseSensitivitySuite.scala

@@ -151,6 +151,17 @@ class V2CommandsCaseSensitivitySuite extends SharedSparkSession with AnalysisTes
     }
   }
 
+  test("AlterTable: add column resolution - column position referencing new column") {
+    alterTableTest(
+      Seq(
+        TableChange.addColumn(
+          Array("x"), LongType, true, null, ColumnPosition.after("id")),
+        TableChange.addColumn(
+          Array("y"), LongType, true, null, ColumnPosition.after("X"))),
+      Seq("Couldn't find the reference column for AFTER X at root")
+    )
+  }
+
   test("AlterTable: add column resolution - nested positional") {
     Seq("X", "Y").foreach { ref =>
       alterTableTest(
@@ -161,6 +172,17 @@ class V2CommandsCaseSensitivitySuite extends SharedSparkSession with AnalysisTes
     }
   }
 
+  test("AlterTable: add column resolution - column position referencing new nested column") {
+    alterTableTest(
+      Seq(
+        TableChange.addColumn(
+          Array("point", "z"), LongType, true, null),
+        TableChange.addColumn(
+          Array("point", "zz"), LongType, true, null, ColumnPosition.after("Z"))),
+      Seq("Couldn't find the reference column for AFTER Z at point")
+    )
+  }
+
   test("AlterTable: drop column resolution") {
     Seq(Array("ID"), Array("point", "X"), Array("POINT", "X"), Array("POINT", "x")).foreach { ref =>
       alterTableTest(
@@ -207,13 +229,17 @@ class V2CommandsCaseSensitivitySuite extends SharedSparkSession with AnalysisTes
   }
 
   private def alterTableTest(change: TableChange, error: Seq[String]): Unit = {
+    alterTableTest(Seq(change), error)
+  }
+
+  private def alterTableTest(changes: Seq[TableChange], error: Seq[String]): Unit = {
     Seq(true, false).foreach { caseSensitive =>
       withSQLConf(SQLConf.CASE_SENSITIVE.key -> caseSensitive.toString) {
         val plan = AlterTable(
           catalog,
           Identifier.of(Array(), "table_name"),
           TestRelation2,
-          Seq(change)
+          changes
         )
 
         if (caseSensitive) {