[SPARK-34474][SQL] Remove unnecessary Union under Distinct/Deduplicate

### What changes were proposed in this pull request?

This patch proposes to let optimizer to remove unnecessary `Union` under `Distinct`/`Deduplicate`.

### Why are the changes needed?

For an `Union` under `Distinct`/`Deduplicate`, if its children are all the same, we can just keep one among them and remove the `Union`.

### Does this PR introduce _any_ user-facing change?

No

### How was this patch tested?

Unit tests.

Closes #31595 from viirya/remove-union.

Authored-by: Liang-Chi Hsieh <viirya@gmail.com>
Signed-off-by: Liang-Chi Hsieh <viirya@gmail.com>

Author: viirya

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

@@ -157,6 +157,7 @@ abstract class Optimizer(catalogManager: CatalogManager)
     //   since the other rules might make two separate Unions operators adjacent.
     Batch("Union", Once,
       RemoveNoopOperators,
+      RemoveNoopUnion,
       CombineUnions) ::
     Batch("OptimizeLimitZero", Once,
       OptimizeLimitZero) ::
@@ -501,6 +502,48 @@ object RemoveNoopOperators extends Rule[LogicalPlan] {
   }
 }
 
+/**
+ * Remove no-op `Union` from the query plan that do not make any modifications.
+ */
+object RemoveNoopUnion extends Rule[LogicalPlan] {
+  /**
+   * This only removes the `Project` that has only attributes or aliased attributes
+   * from its child.
+   */
+  private def removeAliasOnlyProject(plan: LogicalPlan): LogicalPlan = plan match {
+    case p @ Project(projectList, child) =>
+      val aliasOnly = projectList.length == child.output.length &&
+        projectList.zip(child.output).forall {
+          case (Alias(left: Attribute, _), right) => left.semanticEquals(right)
+          case (left: Attribute, right) => left.semanticEquals(right)
+          case _ => false
+        }
+      if (aliasOnly) {
+        child
+      } else {
+        p
+      }
+    case _ => plan
+  }
+
+  private def removeUnion(u: Union): Option[LogicalPlan] = {
+    val unionChildren = u.children.map(removeAliasOnlyProject)
+    if (unionChildren.tail.forall(unionChildren.head.sameResult(_))) {
+      Some(u.children.head)
+    } else {
+      None
+    }
+  }
+
+  def apply(plan: LogicalPlan): LogicalPlan = plan transformUp {
+    case d @ Distinct(u: Union) =>
+      removeUnion(u).map(c => d.withNewChildren(Seq(c))).getOrElse(d)
+
+    case d @ Deduplicate(_, u: Union) =>
+      removeUnion(u).map(c => d.withNewChildren(Seq(c))).getOrElse(d)
+  }
+}
+
 /**
  * Pushes down [[LocalLimit]] beneath UNION ALL and joins.
  */

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

@@ -0,0 +1,70 @@
+/*
+ * 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.plans.PlanTest
+import org.apache.spark.sql.catalyst.plans.logical._
+import org.apache.spark.sql.catalyst.rules.RuleExecutor
+
+class RemoveNoopUnionSuite extends PlanTest {
+
+  object Optimize extends RuleExecutor[LogicalPlan] {
+    val batches =
+      Batch("CollapseProject", Once,
+        CollapseProject) ::
+        Batch("RemoveNoopUnion", Once,
+          RemoveNoopUnion) :: Nil
+  }
+
+  val testRelation = LocalRelation('a.int, 'b.int)
+
+  test("SPARK-34474: Remove redundant Union under Distinct") {
+    val union = Union(testRelation :: testRelation :: Nil)
+    val distinct = Distinct(union)
+    val optimized = Optimize.execute(distinct)
+    comparePlans(optimized, Distinct(testRelation))
+  }
+
+  test("SPARK-34474: Remove redundant Union under Deduplicate") {
+    val union = Union(testRelation :: testRelation :: Nil)
+    val deduplicate = Deduplicate(testRelation.output, union)
+    val optimized = Optimize.execute(deduplicate)
+    comparePlans(optimized, Deduplicate(testRelation.output, testRelation))
+  }
+
+  test("SPARK-34474: Do not remove necessary Project 1") {
+    val child1 = Project(Seq(testRelation.output(0), testRelation.output(1),
+      (testRelation.output(0) + 1).as("expr")), testRelation)
+    val child2 = Project(Seq(testRelation.output(0), testRelation.output(1),
+      (testRelation.output(0) + 2).as("expr")), testRelation)
+    val union = Union(child1 :: child2 :: Nil)
+    val distinct = Distinct(union)
+    val optimized = Optimize.execute(distinct)
+    comparePlans(optimized, distinct)
+  }
+
+  test("SPARK-34474: Do not remove necessary Project 2") {
+    val child1 = Project(Seq(testRelation.output(0), testRelation.output(1)), testRelation)
+    val child2 = Project(Seq(testRelation.output(1), testRelation.output(0)), testRelation)
+    val union = Union(child1 :: child2 :: Nil)
+    val distinct = Distinct(union)
+    val optimized = Optimize.execute(distinct)
+    comparePlans(optimized, distinct)
+  }
+}

sql/core/src/test/resources/sql-tests/inputs/explain.sql

@@ -34,7 +34,7 @@ EXPLAIN FORMATTED
 EXPLAIN FORMATTED
   SELECT key, val FROM explain_temp1 WHERE key > 0
   UNION 
-  SELECT key, val FROM explain_temp1 WHERE key > 0;
+  SELECT key, val FROM explain_temp1 WHERE key > 1;
 
 -- Join
 EXPLAIN FORMATTED

sql/core/src/test/resources/sql-tests/results/explain-aqe.sql.out

@@ -205,7 +205,7 @@ Arguments: isFinalPlan=false
 EXPLAIN FORMATTED
   SELECT key, val FROM explain_temp1 WHERE key > 0
   UNION 
-  SELECT key, val FROM explain_temp1 WHERE key > 0
+  SELECT key, val FROM explain_temp1 WHERE key > 1
 -- !query schema
 struct<plan:string>
 -- !query output
@@ -236,12 +236,12 @@ Condition : (isnotnull(key#x) AND (key#x > 0))
 Output [2]: [key#x, val#x]
 Batched: true
 Location [not included in comparison]/{warehouse_dir}/explain_temp1]
-PushedFilters: [IsNotNull(key), GreaterThan(key,0)]
+PushedFilters: [IsNotNull(key), GreaterThan(key,1)]
 ReadSchema: struct<key:int,val:int>
 
 (4) Filter
 Input [2]: [key#x, val#x]
-Condition : (isnotnull(key#x) AND (key#x > 0))
+Condition : (isnotnull(key#x) AND (key#x > 1))
 
 (5) Union
 

sql/core/src/test/resources/sql-tests/results/explain.sql.out

@@ -203,7 +203,7 @@ Input [3]: [key#x, max(val)#x, max(val#x)#x]
 EXPLAIN FORMATTED
   SELECT key, val FROM explain_temp1 WHERE key > 0
   UNION 
-  SELECT key, val FROM explain_temp1 WHERE key > 0
+  SELECT key, val FROM explain_temp1 WHERE key > 1
 -- !query schema
 struct<plan:string>
 -- !query output
@@ -238,15 +238,15 @@ Condition : (isnotnull(key#x) AND (key#x > 0))
 Output [2]: [key#x, val#x]
 Batched: true
 Location [not included in comparison]/{warehouse_dir}/explain_temp1]
-PushedFilters: [IsNotNull(key), GreaterThan(key,0)]
+PushedFilters: [IsNotNull(key), GreaterThan(key,1)]
 ReadSchema: struct<key:int,val:int>
 
 (5) ColumnarToRow [codegen id : 2]
 Input [2]: [key#x, val#x]
 
 (6) Filter [codegen id : 2]
 Input [2]: [key#x, val#x]
-Condition : (isnotnull(key#x) AND (key#x > 0))
+Condition : (isnotnull(key#x) AND (key#x > 1))
 
 (7) Union