[SPARK-12616] [SQL] Making Logical Operator Union Support Arbitrary Number of Children

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

@@ -189,6 +189,16 @@ trait CheckAnalysis {
                s"but the left table has ${left.output.length} columns and the right has " +
                s"${right.output.length}")
 
+          case s: Unions if s.children.exists(_.output.length != s.children.head.output.length) =>
+            s.children.filter(_.output.length != s.children.head.output.length).foreach { child =>
+              failAnalysis(
+                s"""
+                  |Unions can only be performed on tables with the same number of columns,
+                  | but the table '${child.simpleString}' has '${child.output.length}' columns
+                  | and the first table '${s.children.head.simpleString}' has
+                  | '${s.children.head.output.length}' columns""".stripMargin)
+            }
+
           case _ => // Fallbacks to the following checks
         }
 

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

@@ -196,7 +196,7 @@ object HiveTypeCoercion {
    * - LongType to DoubleType
    * - DecimalType to Double
    *
-   * This rule is only applied to Union/Except/Intersect
+   * This rule is only applied to Unions/Except/Intersect
    */
   object WidenSetOperationTypes extends Rule[LogicalPlan] {
 
@@ -212,29 +212,59 @@ object HiveTypeCoercion {
         case other => None
       }
 
-      def castOutput(plan: LogicalPlan): LogicalPlan = {
-        val casted = plan.output.zip(castedTypes).map {
-          case (e, Some(dt)) if e.dataType != dt =>
-            Alias(Cast(e, dt), e.name)()
-          case (e, _) => e
-        }
-        Project(casted, plan)
+      if (castedTypes.exists(_.isDefined)) {
+        (castOutput(left, castedTypes), castOutput(right, castedTypes))
+      } else {
+        (left, right)
       }
+    }
+
+    private[this] def widenOutputTypes(
+        planName: String,
+        children: Seq[LogicalPlan]): Seq[LogicalPlan] = {
+      require(children.forall(_.output.length == children.head.output.length))
+
+      val castedTypes: Seq[Option[DataType]] =
+        children.tail.foldLeft(children.head.output.map(a => Option(a.dataType))) {
+          case (currentOutputDataTypes, child) => {
+            currentOutputDataTypes.zip(child.output).map {
+              case (Some(dt), a2) if dt != a2.dataType =>
+                findWiderTypeForTwo(dt, a2.dataType)
+              case other => None
+            }
+          }
+        }
 
       if (castedTypes.exists(_.isDefined)) {
-        (castOutput(left), castOutput(right))
+        children.map(castOutput(_, castedTypes))
       } else {
-        (left, right)
+        children
       }
     }
 
+    private[this] def castOutput(
+        plan: LogicalPlan,
+        castedTypes: Seq[Option[DataType]]): LogicalPlan = {
+      val casted = plan.output.zip(castedTypes).map {
+        case (e, Some(dt)) if e.dataType != dt =>
+          Alias(Cast(e, dt), e.name)()
+        case (e, _) => e
+      }
+      Project(casted, plan)
+    }
+
     def apply(plan: LogicalPlan): LogicalPlan = plan resolveOperators {
       case p if p.analyzed => p
 
       case s @ SetOperation(left, right) if s.childrenResolved
           && left.output.length == right.output.length && !s.resolved =>
         val (newLeft, newRight) = widenOutputTypes(s.nodeName, left, right)
         s.makeCopy(Array(newLeft, newRight))
+
+      case s: Unions if s.childrenResolved &&
+        s.children.forall(_.output.length == s.children.head.output.length) && !s.resolved =>
+        val newChildren: Seq[LogicalPlan] = widenOutputTypes(s.nodeName, s.children)
+        s.makeCopy(Array(newChildren))
     }
   }
 

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

@@ -40,6 +40,8 @@ abstract class Optimizer extends RuleExecutor[LogicalPlan] {
     Batch("Aggregate", FixedPoint(100),
       ReplaceDistinctWithAggregate,
       RemoveLiteralFromGroupExpressions) ::
+    Batch("Unions", FixedPoint(100),
+      CombineUnions) ::
     Batch("Operator Optimizations", FixedPoint(100),
       // Operator push down
       SetOperationPushDown,
@@ -54,6 +56,7 @@ abstract class Optimizer extends RuleExecutor[LogicalPlan] {
       ProjectCollapsing,
       CombineFilters,
       CombineLimits,
+      CombineUnions,
       // Constant folding
       NullPropagation,
       OptimizeIn,
@@ -95,8 +98,8 @@ object SamplePushDown extends Rule[LogicalPlan] {
 /**
  * Pushes certain operations to both sides of a Union, Intersect or Except operator.
  * Operations that are safe to pushdown are listed as follows.
- * Union:
- * Right now, Union means UNION ALL, which does not de-duplicate rows. So, it is
+ * Unions:
+ * Right now, Unions means UNION ALL, which does not de-duplicate rows. So, it is
  * safe to pushdown Filters and Projections through it. Once we add UNION DISTINCT,
  * we will not be able to pushdown Projections.
  *
@@ -116,15 +119,15 @@ object SetOperationPushDown extends Rule[LogicalPlan] with PredicateHelper {
    * Maps Attributes from the left side to the corresponding Attribute on the right side.
    */
   private def buildRewrites(bn: BinaryNode): AttributeMap[Attribute] = {
-    assert(bn.isInstanceOf[Union] || bn.isInstanceOf[Intersect] || bn.isInstanceOf[Except])
+    assert(bn.isInstanceOf[Intersect] || bn.isInstanceOf[Except])
     assert(bn.left.output.size == bn.right.output.size)
 
     AttributeMap(bn.left.output.zip(bn.right.output))
   }
 
   /**
    * Rewrites an expression so that it can be pushed to the right side of a
-   * Union, Intersect or Except operator. This method relies on the fact that the output attributes
+   * Unions, Intersect or Except operator. This method relies on the fact that the output attributes
    * of a union/intersect/except are always equal to the left child's output.
    */
   private def pushToRight[A <: Expression](e: A, rewrites: AttributeMap[Attribute]) = {
@@ -153,27 +156,6 @@ object SetOperationPushDown extends Rule[LogicalPlan] with PredicateHelper {
   }
 
   def apply(plan: LogicalPlan): LogicalPlan = plan transform {
-    // Push down filter into union
-    case Filter(condition, u @ Union(left, right)) =>
-      val (deterministic, nondeterministic) = partitionByDeterministic(condition)
-      val rewrites = buildRewrites(u)
-      Filter(nondeterministic,
-        Union(
-          Filter(deterministic, left),
-          Filter(pushToRight(deterministic, rewrites), right)
-        )
-      )
-
-    // Push down deterministic projection through UNION ALL
-    case p @ Project(projectList, u @ Union(left, right)) =>
-      if (projectList.forall(_.deterministic)) {
-        val rewrites = buildRewrites(u)
-        Union(
-          Project(projectList, left),
-          Project(projectList.map(pushToRight(_, rewrites)), right))
-      } else {
-        p
-      }
 
     // Push down filter through INTERSECT
     case Filter(condition, i @ Intersect(left, right)) =>
@@ -594,6 +576,22 @@ object BooleanSimplification extends Rule[LogicalPlan] with PredicateHelper {
   }
 }
 
+/**
+ * Combines all adjacent [[Unions]] operators into a single [[Unions]].
+ */
+object CombineUnions extends Rule[LogicalPlan] {
+  private def buildUnionChildren(children: Seq[LogicalPlan]): Seq[LogicalPlan] =
+    children.foldLeft(Seq.empty[LogicalPlan]) { (newChildren, child) => child match {
+      case Unions(grandchildren) => newChildren ++ grandchildren
+      case other => newChildren ++ Seq(other)
+    }
+  }
+
+  def apply(plan: LogicalPlan): LogicalPlan = plan transformUp {
+    case u @ Unions(children) => Unions(buildUnionChildren(children))
+  }
+}
+
 /**
  * Combines two adjacent [[Filter]] operators into one, merging the
  * conditions into one conjunctive predicate.

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/planning/patterns.scala

@@ -169,18 +169,3 @@ object ExtractFiltersAndInnerJoins extends PredicateHelper {
     case _ => None
   }
 }
-
-/**
- * A pattern that collects all adjacent unions and returns their children as a Seq.
- */
-object Unions {
-  def unapply(plan: LogicalPlan): Option[Seq[LogicalPlan]] = plan match {
-    case u: Union => Some(collectUnionChildren(u))
-    case _ => None
-  }
-
-  private def collectUnionChildren(plan: LogicalPlan): Seq[LogicalPlan] = plan match {
-    case Union(l, r) => collectUnionChildren(l) ++ collectUnionChildren(r)
-    case other => other :: Nil
-  }
-}

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/basicOperators.scala

@@ -107,21 +107,42 @@ private[sql] object SetOperation {
   def unapply(p: SetOperation): Option[(LogicalPlan, LogicalPlan)] = Some((p.left, p.right))
 }
 
-case class Union(left: LogicalPlan, right: LogicalPlan) extends SetOperation(left, right) {
-
-  override def statistics: Statistics = {
-    val sizeInBytes = left.statistics.sizeInBytes + right.statistics.sizeInBytes
-    Statistics(sizeInBytes = sizeInBytes)
-  }
-}
-
 case class Intersect(left: LogicalPlan, right: LogicalPlan) extends SetOperation(left, right)
 
 case class Except(left: LogicalPlan, right: LogicalPlan) extends SetOperation(left, right) {
   /** We don't use right.output because those rows get excluded from the set. */
   override def output: Seq[Attribute] = left.output
 }
 
+/** Factory for constructing new `Unions` nodes. */
+object Union {
+  def apply(left: LogicalPlan, right: LogicalPlan): Unions = {
+    Unions (left :: right :: Nil)
+  }
+}
+
+case class Unions(children: Seq[LogicalPlan]) extends LogicalPlan {
+
+  override def output: Seq[Attribute] = {
+    children.tail.foldLeft(children.head.output) { case (currentOutput, child) =>
+      currentOutput.zip(child.output).map { case (a1, a2) =>
+        a1.withNullability(a1.nullable || a2.nullable)
+      }
+    }
+  }
+
+  override lazy val resolved: Boolean =
+    childrenResolved &&
+      children.forall(_.output.length == children.head.output.length) &&
+      children.forall(_.output.zip(children.head.output).forall {
+        case (l, r) => l.dataType == r.dataType })
+
+  override def statistics: Statistics = {
+    val sizeInBytes = children.map(_.statistics.sizeInBytes).sum
+    Statistics(sizeInBytes = sizeInBytes)
+  }
+}
+
 case class Join(
   left: LogicalPlan,
   right: LogicalPlan,

sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/analysis/DecimalPrecisionSuite.scala

@@ -22,9 +22,9 @@ import org.scalatest.BeforeAndAfter
 import org.apache.spark.SparkFunSuite
 import org.apache.spark.sql.catalyst.expressions._
 import org.apache.spark.sql.catalyst.expressions.aggregate._
-import org.apache.spark.sql.catalyst.plans.logical.{Union, Project, LocalRelation}
+import org.apache.spark.sql.catalyst.plans.logical.{LocalRelation, Project, Union, Unions}
+import org.apache.spark.sql.catalyst.{SimpleCatalystConf, TableIdentifier}
 import org.apache.spark.sql.types._
-import org.apache.spark.sql.catalyst.{TableIdentifier, SimpleCatalystConf}
 
 class DecimalPrecisionSuite extends SparkFunSuite with BeforeAndAfter {
   val conf = new SimpleCatalystConf(true)
@@ -70,7 +70,7 @@ class DecimalPrecisionSuite extends SparkFunSuite with BeforeAndAfter {
       Union(Project(Seq(Alias(left, "l")()), relation),
         Project(Seq(Alias(right, "r")()), relation))
     val (l, r) = analyzer.execute(plan).collect {
-      case Union(left, right) => (left.output.head, right.output.head)
+      case Unions(Seq(child1, child2)) => (child1.output.head, child2.output.head)
     }.head
     assert(l.dataType === expectedType)
     assert(r.dataType === expectedType)

sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/analysis/HiveTypeCoercionSuite.scala

@@ -384,11 +384,11 @@ class HiveTypeCoercionSuite extends PlanTest {
     val wt = HiveTypeCoercion.WidenSetOperationTypes
     val expectedTypes = Seq(StringType, DecimalType.SYSTEM_DEFAULT, FloatType, DoubleType)
 
-    val r1 = wt(Union(left, right)).asInstanceOf[Union]
+    val r1 = wt(Union(left, right)).asInstanceOf[Unions]
     val r2 = wt(Except(left, right)).asInstanceOf[Except]
     val r3 = wt(Intersect(left, right)).asInstanceOf[Intersect]
-    checkOutput(r1.left, expectedTypes)
-    checkOutput(r1.right, expectedTypes)
+    checkOutput(r1.children.head, expectedTypes)
+    checkOutput(r1.children.last, expectedTypes)
     checkOutput(r2.left, expectedTypes)
     checkOutput(r2.right, expectedTypes)
     checkOutput(r3.left, expectedTypes)
@@ -410,12 +410,12 @@ class HiveTypeCoercionSuite extends PlanTest {
       AttributeReference("r", DecimalType(5, 5))())
     val expectedType1 = Seq(DecimalType(10, 8))
 
-    val r1 = dp(Union(left1, right1)).asInstanceOf[Union]
+    val r1 = dp(Union(left1, right1)).asInstanceOf[Unions]
     val r2 = dp(Except(left1, right1)).asInstanceOf[Except]
     val r3 = dp(Intersect(left1, right1)).asInstanceOf[Intersect]
 
-    checkOutput(r1.left, expectedType1)
-    checkOutput(r1.right, expectedType1)
+    checkOutput(r1.children.head, expectedType1)
+    checkOutput(r1.children.last, expectedType1)
     checkOutput(r2.left, expectedType1)
     checkOutput(r2.right, expectedType1)
     checkOutput(r3.left, expectedType1)
@@ -427,23 +427,23 @@ class HiveTypeCoercionSuite extends PlanTest {
     val expectedTypes = Seq(DecimalType(10, 5), DecimalType(10, 5), DecimalType(15, 5),
       DecimalType(25, 5), DoubleType, DoubleType)
 
-    rightTypes.zip(expectedTypes).map { case (rType, expectedType) =>
+    rightTypes.zip(expectedTypes).foreach { case (rType, expectedType) =>
       val plan2 = LocalRelation(
         AttributeReference("r", rType)())
 
-      val r1 = dp(Union(plan1, plan2)).asInstanceOf[Union]
+      val r1 = dp(Union(plan1, plan2)).asInstanceOf[Unions]
       val r2 = dp(Except(plan1, plan2)).asInstanceOf[Except]
       val r3 = dp(Intersect(plan1, plan2)).asInstanceOf[Intersect]
 
-      checkOutput(r1.right, Seq(expectedType))
+      checkOutput(r1.children.last, Seq(expectedType))
       checkOutput(r2.right, Seq(expectedType))
       checkOutput(r3.right, Seq(expectedType))
 
-      val r4 = dp(Union(plan2, plan1)).asInstanceOf[Union]
+      val r4 = dp(Union(plan2, plan1)).asInstanceOf[Unions]
       val r5 = dp(Except(plan2, plan1)).asInstanceOf[Except]
       val r6 = dp(Intersect(plan2, plan1)).asInstanceOf[Intersect]
 
-      checkOutput(r4.left, Seq(expectedType))
+      checkOutput(r4.children.last, Seq(expectedType))
       checkOutput(r5.left, Seq(expectedType))
       checkOutput(r6.left, Seq(expectedType))
     }