Fix and test.

Author: viirya

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/complexTypeCreator.scala

@@ -625,9 +625,9 @@ object WithFields {
     val name = namePartsRemaining.head
     if (namePartsRemaining.length == 1) {
       col.dataType match {
-        case ArrayType(et: StructType, containsNull) =>
+        case ArrayType(et, containsNull) =>
           val lv = NamedLambdaVariable("arg", et, containsNull)
-          val function = WithFields(lv, name :: Nil, value :: Nil)
+          val function = withFieldHelper(lv, name :: Nil, value)
           ArrayTransform(col, LambdaFunction(function, Seq(lv)))
 
         case _: StructType =>
@@ -638,21 +638,21 @@ object WithFields {
       }
     } else {
       val newNamesRemaining = namePartsRemaining.tail
-      val resolver = SQLConf.get.resolver
-
-      val newCol = ExtractValue(col, Literal(name), resolver)
-      val newValue = withFieldHelper(
-        col = newCol,
-        namePartsRemaining = newNamesRemaining,
-        value = value)
 
       col.dataType match {
         case ArrayType(et, containsNull) =>
           val lv = NamedLambdaVariable("arg", et, containsNull)
-          val function = withFieldHelper(lv, namePartsRemaining, newValue)
+          val function = withFieldHelper(lv, namePartsRemaining, value)
           ArrayTransform(col, LambdaFunction(function, Seq(lv)))
 
         case _: StructType =>
+          val resolver = SQLConf.get.resolver
+          val newCol = ExtractValue(col, Literal(name), resolver)
+          val newValue = withFieldHelper(
+            col = newCol,
+            namePartsRemaining = newNamesRemaining,
+            value = value)
+
           WithFields(col, name :: Nil, newValue :: Nil)
 
         case dt =>

sql/core/src/test/scala/org/apache/spark/sql/ColumnExpressionSuite.scala

@@ -1802,4 +1802,77 @@ class ColumnExpressionSuite extends QueryTest with SharedSparkSession {
           StructField("b", IntegerType, nullable = false))),
           containsNull = false), nullable = false))))
   }
+
+  private lazy val arrayStructArrayLevel1: DataFrame = spark.createDataFrame(
+    sparkContext.parallelize(Row(Array(Row(Array(Row(1, null, 3)), null, 3))) :: Nil),
+    StructType(
+      Seq(StructField("a", ArrayType(
+        StructType(Seq(
+          StructField("a", arrayType, nullable = false),
+          StructField("b", IntegerType, nullable = true),
+          StructField("c", IntegerType, nullable = false))),
+        containsNull = false)))))
+
+  test("withField should add and replace field to struct of array of array") {
+    checkAnswerAndSchema(
+      arrayStructArrayLevel1.withColumn("a", $"a".withField("a.d", lit(2))),
+      Row(Seq(Row(Seq(Row(1, null, 3, 2)), null, 3))) :: Nil,
+      StructType(
+        Seq(StructField("a", ArrayType(
+          StructType(Seq(
+            StructField("a", ArrayType(
+              StructType(Seq(
+                StructField("a", IntegerType, nullable = false),
+                StructField("b", IntegerType, nullable = true),
+                StructField("c", IntegerType, nullable = false),
+                StructField("d", IntegerType, nullable = false))),
+              containsNull = true), nullable = false),
+            StructField("b", IntegerType, nullable = true),
+            StructField("c", IntegerType, nullable = false))),
+          containsNull = false)))))
+
+    checkAnswerAndSchema(
+      arrayStructArrayLevel1.withColumn("a", $"a.a".withField("d", lit(2))),
+      Row(Seq(Seq(Row(1, null, 3, 2)))) :: Nil,
+      StructType(
+        Seq(StructField("a", ArrayType(
+          ArrayType(
+            StructType(Seq(
+              StructField("a", IntegerType, nullable = false),
+              StructField("b", IntegerType, nullable = true),
+              StructField("c", IntegerType, nullable = false),
+              StructField("d", IntegerType, nullable = false))),
+            containsNull = true),
+          containsNull = false)))))
+
+    checkAnswerAndSchema(
+      arrayStructArrayLevel1.withColumn("a", $"a".withField("a.b", lit(2))),
+      Row(Seq(Row(Seq(Row(1, 2, 3)), null, 3))) :: Nil,
+      StructType(
+        Seq(StructField("a", ArrayType(
+          StructType(Seq(
+            StructField("a", ArrayType(
+              StructType(Seq(
+                StructField("a", IntegerType, nullable = false),
+                StructField("b", IntegerType, nullable = false),
+                StructField("c", IntegerType, nullable = false))),
+              containsNull = true), nullable = false),
+            StructField("b", IntegerType, nullable = true),
+            StructField("c", IntegerType, nullable = false))),
+          containsNull = false)))))
+
+    checkAnswerAndSchema(
+      arrayStructArrayLevel1.withColumn("a", $"a.a".withField("b", lit(2))),
+      Row(Seq(Seq(Row(1, 2, 3)))) :: Nil,
+      StructType(
+        Seq(StructField("a", ArrayType(
+          ArrayType(
+            StructType(Seq(
+              StructField("a", IntegerType, nullable = false),
+              StructField("b", IntegerType, nullable = false),
+              StructField("c", IntegerType, nullable = false))),
+            containsNull = true),
+          containsNull = false)))))
+
+  }
 }