[SPARK-12503] [SQL] Pushing Limit Through Union All

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

@@ -86,8 +86,8 @@ object SamplePushDown extends Rule[LogicalPlan] {
  * 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
- * safe to pushdown Filters and Projections through it. Once we add UNION DISTINCT,
- * we will not be able to pushdown Projections.
+ * safe to pushdown Filters, Projections and Limits through it. Once we add UNION DISTINCT,
+ * we will not be able to pushdown Projections and Limits.
  *
  * Intersect:
  * It is not safe to pushdown Projections through it because we need to get the
@@ -153,6 +153,17 @@ object SetOperationPushDown extends Rule[LogicalPlan] with PredicateHelper {
         )
       )
 
+    // Adding extra Limit below UNION ALL iff both left and right childs are not Limit and no Limit
+    // was pushed down before. This heuristic is valid assuming there does not exist any Limit
+    // push-down rule that is unable to infer the value of maxRows. Any operator that a Limit can
+    // be pushed passed should override this function.
+    case Limit(exp, Union(left, right))
+      if left.maxRows.isEmpty || right.maxRows.isEmpty =>
+      Limit(exp,
+        Union(
+          Limit(exp, left),
+          Limit(exp, right)))
+
     // Push down deterministic projection through UNION ALL
     case p @ Project(projectList, u @ Union(left, right)) =>
       if (projectList.forall(_.deterministic)) {

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

@@ -90,6 +90,13 @@ abstract class LogicalPlan extends QueryPlan[LogicalPlan] with Logging {
     Statistics(sizeInBytes = children.map(_.statistics.sizeInBytes).product)
   }
 
+  /**
+   * Returns the limited number of rows to be returned.
+   *
+   * Any operator that a Limit can be pushed passed should override this function.
+   */
+  def maxRows: Option[Expression] = None
+
   /**
    * Returns true if this expression and all its children have been resolved to a specific schema
    * and false if it still contains any unresolved placeholders. Implementations of LogicalPlan

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

@@ -28,6 +28,8 @@ import scala.collection.mutable.ArrayBuffer
 case class Project(projectList: Seq[NamedExpression], child: LogicalPlan) extends UnaryNode {
   override def output: Seq[Attribute] = projectList.map(_.toAttribute)
 
+  override def maxRows: Option[Expression] = child.maxRows
+
   override lazy val resolved: Boolean = {
     val hasSpecialExpressions = projectList.exists ( _.collect {
         case agg: AggregateExpression => agg
@@ -451,6 +453,8 @@ case class Pivot(
 case class Limit(limitExpr: Expression, child: LogicalPlan) extends UnaryNode {
   override def output: Seq[Attribute] = child.output
 
+  override def maxRows: Option[Expression] = Option(limitExpr)
+
   override lazy val statistics: Statistics = {
     val limit = limitExpr.eval().asInstanceOf[Int]
     val sizeInBytes = (limit: Long) * output.map(a => a.dataType.defaultSize).sum

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

@@ -31,6 +31,9 @@ class SetOperationPushDownSuite extends PlanTest {
         EliminateSubQueries) ::
       Batch("Union Pushdown", Once,
         SetOperationPushDown,
+        CombineLimits,
+        ConstantFolding,
+        BooleanSimplification,
         SimplifyFilters) :: Nil
   }
 
@@ -61,6 +64,33 @@ class SetOperationPushDownSuite extends PlanTest {
     comparePlans(exceptOptimized, exceptCorrectAnswer)
   }
 
+  test("union: limit to each side") {
+    val unionQuery = testUnion.limit(1)
+    val unionOptimized = Optimize.execute(unionQuery.analyze)
+    val unionCorrectAnswer =
+      Limit(1, Union(testRelation.limit(1), testRelation2.limit(1))).analyze
+    comparePlans(unionOptimized, unionCorrectAnswer)
+  }
+
+  test("union: limit to each side with the new limit number") {
+    val testLimitUnion = Union(testRelation, testRelation2.limit(3))
+    val unionQuery = testLimitUnion.limit(1)
+    val unionOptimized = Optimize.execute(unionQuery.analyze)
+    val unionCorrectAnswer =
+      Limit(1, Union(testRelation.limit(1), testRelation2.limit(1))).analyze
+    comparePlans(unionOptimized, unionCorrectAnswer)
+  }
+
+  // If users already manually added the Limit, we do not add extra Limit
+  test("union: no limit to both sides") {
+    val testLimitUnion = Union(testRelation.limit(2), testRelation2.select('d).limit(3))
+    val unionQuery = testLimitUnion.limit(2)
+    val unionOptimized = Optimize.execute(unionQuery.analyze)
+    val unionCorrectAnswer =
+      Limit(2, Union(testRelation.limit(2), testRelation2.select('d).limit(3))).analyze
+    comparePlans(unionOptimized, unionCorrectAnswer)
+  }
+
   test("union: project to each side") {
     val unionQuery = testUnion.select('a)
     val unionOptimized = Optimize.execute(unionQuery.analyze)