[SPARK-32941][SQL] Optimize UpdateFields expression chain and put the rule early in Analysis phase

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/ComplexTypes.scala

@@ -48,7 +48,11 @@ object SimplifyExtractValueOps extends Rule[LogicalPlan] {
           // `$"struct_col".withField("b", lit(1)).withField("b", lit(2)).getField("b")`
           // we want to return `lit(2)` (and not `lit(1)`).
           val expr = matches.last._2
-          If(IsNull(struct), Literal(null, expr.dataType), expr)
+          if (struct.nullable) {
+            If(IsNull(struct), Literal(null, expr.dataType), expr)
+          } else {
+            expr
+          }
         } else {
           GetStructField(struct, ordinal, maybeName)
         }

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/Optimizer.scala

@@ -109,7 +109,7 @@ abstract class Optimizer(catalogManager: CatalogManager)
         RemoveRedundantAliases,
         UnwrapCastInBinaryComparison,
         RemoveNoopOperators,
-        CombineWithFields,
+        OptimizeWithFields,
         SimplifyExtractValueOps,
         CombineConcats) ++
         extendedOperatorOptimizationRules

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/WithFields.scala

@@ -17,16 +17,49 @@
 
 package org.apache.spark.sql.catalyst.optimizer
 
-import org.apache.spark.sql.catalyst.expressions.WithFields
+import java.util.Locale
+
+import scala.collection.mutable
+
+import org.apache.spark.sql.catalyst.expressions.{Expression, WithFields}
 import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan
 import org.apache.spark.sql.catalyst.rules.Rule
+import org.apache.spark.sql.internal.SQLConf
 
 
 /**
- * Combines all adjacent [[WithFields]] expression into a single [[WithFields]] expression.
+ * Optimizes [[WithFields]] expression chains.
  */
-object CombineWithFields extends Rule[LogicalPlan] {
+object OptimizeWithFields extends Rule[LogicalPlan] {
   def apply(plan: LogicalPlan): LogicalPlan = plan transformAllExpressions {
+    case WithFields(structExpr, names, values)
+        if names.map(_.toLowerCase(Locale.ROOT)).distinct.length != names.length =>
+      val caseSensitive = SQLConf.get.caseSensitiveAnalysis
+
+      val newNames = mutable.ArrayBuffer.empty[String]
+      val newValues = mutable.ArrayBuffer.empty[Expression]
+
+      if (caseSensitive) {
+        names.zip(values).reverse.foreach { case (name, value) =>
+          if (!newNames.contains(name)) {
+            newNames += name
+            newValues += value
+          }
+        }
+      } else {
+        val nameSet = mutable.HashSet.empty[String]
+        names.zip(values).reverse.foreach { case (name, value) =>
+          val lowercaseName = name.toLowerCase(Locale.ROOT)
+          if (!nameSet.contains(lowercaseName)) {
+            newNames += name
+            newValues += value
+            nameSet += lowercaseName
+          }
+        }
+      }
+
+      WithFields(structExpr, names = newNames.reverse.toSeq, valExprs = newValues.reverse.toSeq)
+
     case WithFields(WithFields(struct, names1, valExprs1), names2, valExprs2) =>
       WithFields(struct, names1 ++ names2, valExprs1 ++ valExprs2)
   }

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

@@ -0,0 +1,154 @@
+/*
+ * 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.optimizer
+
+import org.apache.spark.sql.catalyst.dsl.expressions._
+import org.apache.spark.sql.catalyst.dsl.plans._
+import org.apache.spark.sql.catalyst.expressions.{Alias, GetStructField, Literal, WithFields}
+import org.apache.spark.sql.catalyst.plans.PlanTest
+import org.apache.spark.sql.catalyst.plans.logical._
+import org.apache.spark.sql.catalyst.rules._
+import org.apache.spark.sql.internal.SQLConf
+
+class OptimizeWithFieldsSuite extends PlanTest {
+
+  object Optimize extends RuleExecutor[LogicalPlan] {
+    val batches = Batch("OptimizeWithFields", FixedPoint(10),
+      OptimizeWithFields, SimplifyExtractValueOps) :: Nil
+  }
+
+  private val testRelation = LocalRelation('a.struct('a1.int))
+  private val testRelation2 = LocalRelation('a.struct('a1.int).notNull)
+
+  test("combines two WithFields") {
+    val originalQuery = testRelation
+      .select(Alias(
+        WithFields(
+          WithFields(
+            'a,
+            Seq("b1"),
+            Seq(Literal(4))),
+          Seq("c1"),
+          Seq(Literal(5))), "out")())
+
+    val optimized = Optimize.execute(originalQuery.analyze)
+    val correctAnswer = testRelation
+      .select(Alias(WithFields('a, Seq("b1", "c1"), Seq(Literal(4), Literal(5))), "out")())
+      .analyze
+
+    comparePlans(optimized, correctAnswer)
+  }
+
+  test("combines three WithFields") {
+    val originalQuery = testRelation
+      .select(Alias(
+        WithFields(
+          WithFields(
+            WithFields(
+              'a,
+              Seq("b1"),
+              Seq(Literal(4))),
+            Seq("c1"),
+            Seq(Literal(5))),
+          Seq("d1"),
+          Seq(Literal(6))), "out")())
+
+    val optimized = Optimize.execute(originalQuery.analyze)
+    val correctAnswer = testRelation
+      .select(Alias(WithFields('a, Seq("b1", "c1", "d1"), Seq(4, 5, 6).map(Literal(_))), "out")())
+      .analyze
+
+    comparePlans(optimized, correctAnswer)
+  }
+
+  test("SPARK-32941: optimize WithFields followed by GetStructField") {
+    val originalQuery = testRelation2
+      .select(Alias(
+        GetStructField(WithFields(
+          'a,
+          Seq("b1"),
+          Seq(Literal(4))), 1), "out")())
+
+    val optimized = Optimize.execute(originalQuery.analyze)
+    val correctAnswer = testRelation2
+      .select(Alias(Literal(4), "out")())
+      .analyze
+
+    comparePlans(optimized, correctAnswer)
+  }
+
+  test("SPARK-32941: optimize WithFields chain - case insensitive") {
+    val originalQuery = testRelation
+      .select(Alias(
+        WithFields(
+          WithFields(
+            'a,
+            Seq("b1"),
+            Seq(Literal(4))),
+          Seq("b1"),
+          Seq(Literal(5))), "out1")(),
+        Alias(
+          WithFields(
+            WithFields(
+              'a,
+              Seq("b1"),
+              Seq(Literal(4))),
+            Seq("B1"),
+            Seq(Literal(5))), "out2")())
+
+    val optimized = Optimize.execute(originalQuery.analyze)
+    val correctAnswer = testRelation
+      .select(
+        Alias(WithFields('a, Seq("b1"), Seq(Literal(5))), "out1")(),
+        Alias(WithFields('a, Seq("B1"), Seq(Literal(5))), "out2")())
+      .analyze
+
+    comparePlans(optimized, correctAnswer)
+  }
+
+  test("SPARK-32941: optimize WithFields chain - case sensitive") {
+    withSQLConf(SQLConf.CASE_SENSITIVE.key -> "true") {
+      val originalQuery = testRelation
+        .select(Alias(
+          WithFields(
+            WithFields(
+              'a,
+              Seq("b1"),
+              Seq(Literal(4))),
+            Seq("b1"),
+            Seq(Literal(5))), "out1")(),
+          Alias(
+            WithFields(
+              WithFields(
+                'a,
+                Seq("b1"),
+                Seq(Literal(4))),
+              Seq("B1"),
+              Seq(Literal(5))), "out2")())
+
+      val optimized = Optimize.execute(originalQuery.analyze)
+      val correctAnswer = testRelation
+        .select(
+          Alias(WithFields('a, Seq("b1"), Seq(Literal(5))), "out1")(),
+          Alias(WithFields('a, Seq("b1", "B1"), Seq(Literal(4), Literal(5))), "out2")())
+        .analyze
+
+      comparePlans(optimized, correctAnswer)
+    }
+  }
+}