[SPARK-32241][SQL] Remove empty children of union

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

@@ -161,7 +161,8 @@ abstract class Optimizer(catalogManager: CatalogManager)
     // LocalRelation and does not trigger many rules.
     Batch("LocalRelation early", fixedPoint,
       ConvertToLocalRelation,
-      PropagateEmptyRelation) ::
+      PropagateEmptyRelation,
+      UpdateAttributeNullability) ::
     Batch("Pullup Correlated Expressions", Once,
       PullupCorrelatedPredicates) ::
     // Subquery batch applies the optimizer rules recursively. Therefore, it makes no sense
@@ -198,7 +199,8 @@ abstract class Optimizer(catalogManager: CatalogManager)
       ReassignLambdaVariableID) :+
     Batch("LocalRelation", fixedPoint,
       ConvertToLocalRelation,
-      PropagateEmptyRelation) :+
+      PropagateEmptyRelation,
+      UpdateAttributeNullability) :+
     // The following batch should be executed after batch "Join Reorder" and "LocalRelation".
     Batch("Check Cartesian Products", Once,
       CheckCartesianProducts) :+

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

@@ -50,8 +50,24 @@ object PropagateEmptyRelation extends Rule[LogicalPlan] with PredicateHelper wit
   override def conf: SQLConf = SQLConf.get
 
   def apply(plan: LogicalPlan): LogicalPlan = plan transformUp {
-    case p: Union if p.children.forall(isEmptyLocalRelation) =>
-      empty(p)
+    case p @ Union(children) if children.exists(isEmptyLocalRelation) =>
+      val newChildren = children.filterNot(isEmptyLocalRelation)
+      if (newChildren.isEmpty) {
+        empty(p)
+      } else {
+        val newPlan = if (newChildren.size > 1) Union(newChildren) else newChildren.head
+        val outputs = newPlan.output.zip(p.output)
+        // the original Union may produce different output attributes than the new one so we alias
+        // them if needed
+        if (outputs.forall { case (newAttr, oldAttr) => newAttr.exprId == oldAttr.exprId }) {
+          newPlan
+        } else {
+          val outputAliases = outputs.map { case (newAttr, oldAttr) =>
+            Alias(newAttr, oldAttr.name)(oldAttr.exprId)
+          }
+          Project(outputAliases, newPlan)
+        }
+      }
 
     // Joins on empty LocalRelations generated from streaming sources are not eliminated
     // as stateful streaming joins need to perform other state management operations other than

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

@@ -55,6 +55,7 @@ class PropagateEmptyRelationSuite extends PlanTest {
 
   val testRelation1 = LocalRelation.fromExternalRows(Seq('a.int), data = Seq(Row(1)))
   val testRelation2 = LocalRelation.fromExternalRows(Seq('b.int), data = Seq(Row(1)))
+  val testRelation3 = LocalRelation.fromExternalRows(Seq('c.int.notNull), data = Seq(Row(1)))
 
   test("propagate empty relation through Union") {
     val query = testRelation1
@@ -67,6 +68,34 @@ class PropagateEmptyRelationSuite extends PlanTest {
     comparePlans(optimized, correctAnswer)
   }
 
+  test("SPARK-32241: remove empty relation children from Union") {
+    val query = testRelation1.union(testRelation2.where(false))
+    val optimized = Optimize.execute(query.analyze)
+    val correctAnswer = testRelation1
+    comparePlans(optimized, correctAnswer)
+
+    val query2 = testRelation1.where(false).union(testRelation2)
+    val optimized2 = Optimize.execute(query2.analyze)
+    val correctAnswer2 = testRelation2.select('b.as('a)).analyze
+    comparePlans(optimized2, correctAnswer2)
+
+    val query3 = testRelation1.union(testRelation2.where(false)).union(testRelation3)
+    val optimized3 = Optimize.execute(query3.analyze)
+    val correctAnswer3 = testRelation1.union(testRelation3)
+    comparePlans(optimized3, correctAnswer3)
+
+    val query4 = testRelation1.where(false).union(testRelation2).union(testRelation3)
+    val optimized4 = Optimize.execute(query4.analyze)
+    val correctAnswer4 = testRelation2.union(testRelation3).select('b.as('a)).analyze
+    comparePlans(optimized4, correctAnswer4)
+
+    // Nullability can change from nullable to non-nullable
+    val query5 = testRelation1.where(false).union(testRelation3).select('a.as('x)).analyze
+    val optimized5 = Optimize.execute(query5.analyze)
+    assert(query5.output.head.nullable, "Original output should be nullable")
+    assert(!optimized5.output.head.nullable, "The new output should be non-nullable")
+  }
+
   test("propagate empty relation through Join") {
     // Testcases are tuples of (left predicate, right predicate, joinType, correct answer)
     // Note that `None` is used to compare with OptimizeWithoutPropagateEmptyRelation.