From c9e8a97a60fa78f0fc4876e4b43b11b40aacccf2 Mon Sep 17 00:00:00 2001 From: Jiwon Park Date: Sun, 12 Jul 2026 21:28:13 +0900 Subject: [PATCH] [SPARK-58093][SQL] Push down LIMIT through UNION ALL to Data Source V2 scan `V2ScanRelationPushDown.pushDownLimit` only handled a scan directly under the limit (optionally through `Project`/`Sort`), so a `LIMIT` on top of a `UNION ALL` fell through to `case other` and was never delivered to the underlying Data Source V2 scans via `SupportsPushDownLimit`. `LimitPushDown` already pushes a `LocalLimit` into each union branch, so the operator-level limit reaches the branches while the source-level limit hint does not. Add `Union` and `LocalLimit` cases to `pushDownLimit` that recurse into each branch and push the limit as a partial limit, always keeping the outer limit operators. UNION ALL never de-duplicates rows, so the union of branches each capped at `limit` still contains the first `limit` rows of the whole union. Signed-off-by: Jiwon Park --- .../v2/V2ScanRelationPushDown.scala | 13 ++- .../apache/spark/sql/jdbc/JDBCV2Suite.scala | 87 +++++++++++++++++++ 2 files changed, 99 insertions(+), 1 deletion(-) diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/V2ScanRelationPushDown.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/V2ScanRelationPushDown.scala index 14d9f754f95c5..06052ce15132b 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/V2ScanRelationPushDown.scala +++ b/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) } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/jdbc/JDBCV2Suite.scala b/sql/core/src/test/scala/org/apache/spark/sql/jdbc/JDBCV2Suite.scala index dd3d90d2df052..5f16bdc4f261c 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/jdbc/JDBCV2Suite.scala +++ b/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