[SPARK-58093][SQL] Push down LIMIT through UNION ALL to Data Source V2 scan

sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/V2ScanRelationPushDown.scala

@@ -28,7 +28,7 @@ import org.apache.spark.sql.catalyst.expressions.{aggregate, Alias, And, Attribu
 import org.apache.spark.sql.catalyst.expressions.aggregate.AggregateExpression
 import org.apache.spark.sql.catalyst.optimizer.{CollapseGroupedSumOfCount, CollapseProject}
 import org.apache.spark.sql.catalyst.planning.{PhysicalOperation, ScanOperation}
-import org.apache.spark.sql.catalyst.plans.logical.{Aggregate, Filter, Join, LeafNode, Limit, LimitAndOffset, LocalLimit, LogicalPlan, Offset, OffsetAndLimit, Project, Sample, SampleMethod, Sort}
+import org.apache.spark.sql.catalyst.plans.logical.{Aggregate, Filter, Join, LeafNode, Limit, LimitAndOffset, LocalLimit, LogicalPlan, Offset, OffsetAndLimit, Project, Sample, SampleMethod, Sort, Union}
 import org.apache.spark.sql.catalyst.rules.Rule
 import org.apache.spark.sql.catalyst.types.DataTypeUtils.toAttributes
 import org.apache.spark.sql.connector.expressions.{SortOrder => V2SortOrder}
@@ -1047,6 +1047,17 @@ object V2ScanRelationPushDown extends Rule[LogicalPlan] with PredicateHelper {
     case p: Project =>
       val (newChild, isPartiallyPushed) = pushDownLimit(p.child, limit)
       (p.withNewChildren(Seq(newChild)), isPartiallyPushed)
+    case u: Union =>
+      // The union of branches that each keep at most `limit` rows still contains the first `limit`
+      // rows of the whole union, as UNION ALL does not de-duplicate. The union may return more than
+      // `limit` rows, so this is only a partial push.
+      val newChildren = u.children.map(child => pushDownLimit(child, limit)._1)
+      (u.withNewChildren(newChildren), false)
+    case l @ LocalLimit(IntegerLiteral(_), _) =>
+      // `LimitPushDown` inserts this per-branch limit with the same value as the limit above the
+      // union, so the incoming limit can be pushed through it.
+      val (newChild, _) = pushDownLimit(l.child, limit)
+      (l.withNewChildren(Seq(newChild)), false)
     case other => (other, false)
   }
 

sql/core/src/test/scala/org/apache/spark/sql/jdbc/JDBCV2Suite.scala

@@ -402,6 +402,93 @@ class JDBCV2Suite extends SharedSparkSession with ExplainSuiteHelper {
     }
   }
 
+  private def checkPushedLimits(df: DataFrame, limits: Option[Int]*): Unit = {
+    val pushedLimits = df.queryExecution.optimizedPlan.collect {
+      case relation: DataSourceV2ScanRelation => relation.scan match {
+        case v1: V1ScanWrapper => v1.pushedDownOperators.limit
+        case other => fail(s"Expected a V1ScanWrapper but got $other")
+      }
+    }
+    assert(pushedLimits === limits)
+  }
+
+  test("SPARK-58093: push down LIMIT through UNION ALL to each data source scan") {
+    // dept = 6 matches a single employee, so the union of the two branches is deterministic
+    val df = sql(
+      """
+        |SELECT name FROM h2.test.employee WHERE dept = 6
+        |UNION ALL
+        |SELECT name FROM h2.test.employee WHERE dept = 6
+        |LIMIT 1
+        |""".stripMargin)
+    checkPushedLimits(df, Some(1), Some(1))
+    checkPushedInfo(df, "PushedFilters: [DEPT IS NOT NULL, DEPT = 6]", "PushedLimit: LIMIT 1")
+    checkLimitRemoved(df, removed = false)
+    checkAnswer(df, Seq(Row("jen")))
+  }
+
+  test("SPARK-58093: push down LIMIT through UNION ALL with more than two branches") {
+    val df = sql(
+      """
+        |SELECT name FROM h2.test.employee WHERE dept = 6
+        |UNION ALL
+        |SELECT name FROM h2.test.employee WHERE dept = 6
+        |UNION ALL
+        |SELECT name FROM h2.test.employee WHERE dept = 6
+        |LIMIT 2
+        |""".stripMargin)
+    checkPushedLimits(df, Some(2), Some(2), Some(2))
+    checkAnswer(df, Seq(Row("jen"), Row("jen")))
+  }
+
+  test("SPARK-58093: push down LIMIT through UNION ALL over a sorted branch") {
+    // the sort survives above the scan of the first branch, so the limit reaches it as a top N
+    val df = spark.read.table("h2.test.employee").select($"name").orderBy($"salary")
+      .union(spark.read.table("h2.test.employee").select($"name"))
+      .limit(1)
+    checkPushedLimits(df, Some(1), Some(1))
+    checkSortRemoved(df)
+    checkLimitRemoved(df, removed = false)
+    assert(df.collect().length == 1)
+  }
+
+  test("SPARK-58093: push down LIMIT through UNION ALL together with OFFSET") {
+    // `LIMIT n OFFSET m` skips before it takes, so each branch has to keep `n + m` rows
+    val df1 = sql(
+      """
+        |SELECT name FROM h2.test.employee
+        |UNION ALL
+        |SELECT name FROM h2.test.employee
+        |LIMIT 4 OFFSET 2
+        |""".stripMargin)
+    checkPushedLimits(df1, Some(6), Some(6))
+    checkOffsetRemoved(df1, removed = false)
+    assert(df1.collect().length == 4)
+
+    // `limit(n).offset(m)` takes before it skips, so `n` rows per branch are enough
+    val df2 = spark.read.table("h2.test.employee").select($"name")
+      .union(spark.read.table("h2.test.employee").select($"name"))
+      .limit(4)
+      .offset(2)
+    checkPushedLimits(df2, Some(4), Some(4))
+    checkOffsetRemoved(df2, removed = false)
+    assert(df2.collect().length == 2)
+  }
+
+  test("SPARK-58093: LIMIT with ORDER BY on top of UNION ALL is not pushed down") {
+    // the limit must not be pushed below the sort: a branch capped before the global sort could
+    // drop rows that belong to the overall top N
+    val df = sql(
+      """
+        |SELECT name FROM h2.test.employee
+        |UNION ALL
+        |SELECT name FROM h2.test.employee
+        |ORDER BY name LIMIT 2
+        |""".stripMargin)
+    checkPushedLimits(df, None, None)
+    checkAnswer(df, Seq(Row("alex"), Row("alex")))
+  }
+
   private def checkOffsetRemoved(df: DataFrame, removed: Boolean = true): Unit = {
     val offsets = df.queryExecution.optimizedPlan.collectFirst {
       case offset: Offset => offset