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

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

@@ -153,6 +153,15 @@ object SetOperationPushDown extends Rule[LogicalPlan] with PredicateHelper {
         )
       )
 
+    // Push down limit into union
+    case Limit(exp, Union(left, right)) =>
+      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/test/scala/org/apache/spark/sql/catalyst/optimizer/SetOperationPushDownSuite.scala

@@ -61,6 +61,14 @@ 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: project to each side") {
     val unionQuery = testUnion.select('a)
     val unionOptimized = Optimize.execute(unionQuery.analyze)