[SPARK-32796][SQL] Make withField API support nested struct in array

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

@@ -371,6 +371,7 @@ class Analyzer(
       ResolveRandomSeed ::
       ResolveBinaryArithmetic ::
       ResolveUnion ::
+      ResolveWithFields ::
       TypeCoercion.typeCoercionRules(conf) ++
       extendedResolutionRules : _*),
     Batch("Post-Hoc Resolution", Once, postHocResolutionRules: _*),

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

@@ -0,0 +1,38 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql.catalyst.analysis
+
+import org.apache.spark.sql.catalyst.expressions.WithFields
+import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan
+import org.apache.spark.sql.catalyst.rules.Rule
+
+/**
+ * Resolves `UnresolvedWithFields`.
+ */
+object ResolveWithFields extends Rule[LogicalPlan] {
+  override def apply(plan: LogicalPlan): LogicalPlan = plan.resolveOperatorsUp {
+    case e if !e.childrenResolved => e
+
+    case q: LogicalPlan =>
+      q.transformExpressions {
+        case expr if !expr.childrenResolved => expr
+        case e: UnresolvedWithFields => WithFields(e.col, e.fieldName, e.expr)
+        case others => others
+      }
+  }
+}

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

@@ -554,3 +554,14 @@ case class UnresolvedHaving(
   override lazy val resolved: Boolean = false
   override def output: Seq[Attribute] = child.output
 }
+
+case class UnresolvedWithFields(
+    col: Expression,
+    fieldName: String,
+    expr: Expression) extends Unevaluable with NonSQLExpression {
+  override def children: Seq[Expression] = Seq(col, expr)
+  override def dataType: DataType = throw new UnresolvedException(this, "dataType")
+  override def foldable: Boolean = throw new UnresolvedException(this, "foldable")
+  override def nullable: Boolean = throw new UnresolvedException(this, "nullable")
+  override lazy val resolved = false
+}

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

@@ -17,11 +17,13 @@
 
 package org.apache.spark.sql.catalyst.expressions
 
+import org.apache.spark.sql.AnalysisException
 import org.apache.spark.sql.catalyst.InternalRow
-import org.apache.spark.sql.catalyst.analysis.{TypeCheckResult, TypeCoercion}
+import org.apache.spark.sql.catalyst.analysis.{TypeCheckResult, TypeCoercion, UnresolvedExtractValue}
 import org.apache.spark.sql.catalyst.analysis.FunctionRegistry.{FUNC_ALIAS, FunctionBuilder}
 import org.apache.spark.sql.catalyst.expressions.codegen._
 import org.apache.spark.sql.catalyst.expressions.codegen.Block._
+import org.apache.spark.sql.catalyst.parser.CatalystSqlParser
 import org.apache.spark.sql.catalyst.util._
 import org.apache.spark.sql.internal.SQLConf
 import org.apache.spark.sql.types._
@@ -595,3 +597,68 @@ case class WithFields(
     }
   }
 }
+
+object WithFields {
+  /**
+   * Adds/replaces field in `StructType` into `col` expression by name.
+   */
+  def apply(
+      col: Expression,
+      fieldName: String,
+      expr: Expression): Expression = {
+    val nameParts = if (fieldName.isEmpty) {
+      fieldName :: Nil
+    } else {
+      CatalystSqlParser.parseMultipartIdentifier(fieldName)
+    }
+    withFieldHelper(col, nameParts, expr)
+  }
+
+  /**
+   * Recursively builds expressions for adding/replacing (nested) field in `StructType` into
+   * `col` expression by name. Supports nested struct in array and struct.
+   */
+  private def withFieldHelper(
+      col: Expression,
+      namePartsRemaining: Seq[String],
+      value: Expression) : Expression = {
+    val name = namePartsRemaining.head
+    if (namePartsRemaining.length == 1) {
+      col.dataType match {
+        case ArrayType(et, containsNull) =>
+          val lv = NamedLambdaVariable("arg", et, containsNull)
+          val function = withFieldHelper(lv, name :: Nil, value)
+          ArrayTransform(col, LambdaFunction(function, Seq(lv)))
+
+        case _: StructType =>
+          WithFields(col, name :: Nil, value :: Nil)
+
+        case dt =>
+          throw new AnalysisException(s"WithFields's argument does not support ${dt.catalogString}")
+      }
+    } else {
+      val newNamesRemaining = namePartsRemaining.tail
+
+      col.dataType match {
+        case ArrayType(et, containsNull) =>
+          val lv = NamedLambdaVariable("arg", et, containsNull)
+          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 =>
+          throw new AnalysisException(s"WithFields's argument does not support ${dt.catalogString}")
+      }
+
+    }
+  }
+}

sql/core/src/main/scala/org/apache/spark/sql/Column.scala

@@ -909,31 +909,10 @@ class Column(val expr: Expression) extends Logging {
     require(fieldName != null, "fieldName cannot be null")
     require(col != null, "col cannot be null")
 
-    val nameParts = if (fieldName.isEmpty) {
-      fieldName :: Nil
+    if (expr.resolved) {
+      WithFields(expr, fieldName, col.expr)
     } else {
-      CatalystSqlParser.parseMultipartIdentifier(fieldName)
-    }
-    withFieldHelper(expr, nameParts, Nil, col.expr)
-  }
-
-  private def withFieldHelper(
-      struct: Expression,
-      namePartsRemaining: Seq[String],
-      namePartsDone: Seq[String],
-      value: Expression) : WithFields = {
-    val name = namePartsRemaining.head
-    if (namePartsRemaining.length == 1) {
-      WithFields(struct, name :: Nil, value :: Nil)
-    } else {
-      val newNamesRemaining = namePartsRemaining.tail
-      val newNamesDone = namePartsDone :+ name
-      val newValue = withFieldHelper(
-        struct = UnresolvedExtractValue(struct, Literal(name)),
-        namePartsRemaining = newNamesRemaining,
-        namePartsDone = newNamesDone,
-        value = value)
-      WithFields(struct, name :: Nil, newValue :: Nil)
+      UnresolvedWithFields(expr, fieldName, col.expr)
     }
   }