This is an automated email from the ASF dual-hosted git repository.
cloud-fan pushed a commit to branch branch-4.3
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/branch-4.3 by this push:
new 3c6ecb2f07cc [SPARK-57956][SQL] Fix AQE incorrectly eliminating a
global LIMIT over a non-materialized query stage
3c6ecb2f07cc is described below
commit 3c6ecb2f07cc28a8e4a7a33c155026648e7bb786
Author: Vaibhav Garg <[email protected]>
AuthorDate: Tue Aug 18 01:29:36 2026 +0800
[SPARK-57956][SQL] Fix AQE incorrectly eliminating a global LIMIT over a
non-materialized query stage
### What changes were proposed in this pull request?
This PR fixes `LogicalQueryStage.maxRows` to only use runtime row counts
when the underlying query stage is materialized. Before materialization,
`computeStats()` returns `None` and the code falls back to `logicalPlan.stats`,
which is a cost estimate, not a hard upper bound. That estimate was being
surfaced as `maxRows`, causing `EliminateLimits` to incorrectly remove a global
LIMIT.
The fix adds a materialization check: if the stage is materialized, use the
runtime `rowCount` (an exact upper bound). Otherwise, fall back to
`logicalPlan.maxRows`, which is always a sound upper bound derived from plan
structure rather than cardinality estimates.
### Why are the changes needed?
With AQE enabled, a root global LIMIT can be dropped during adaptive
re-optimization, returning too many rows. For example:
```sql
WITH top AS (SELECT id % 1 AS k FROM range(0, 2, 1, 1) OFFSET 1),
b AS (SELECT * FROM VALUES (0), (0) AS b(k))
SELECT 1 AS c FROM top LEFT JOIN b USING (k) LIMIT 1
```
This query returns 2 rows instead of 1.
The root cause: `EliminateLimits` removes a `GlobalLimit` when
`child.maxRows <= limit`. The child here is a `Project` over a `LeftOuter Join`
whose left input is a `LogicalQueryStage` wrapping `Offset 1 +- Range`. Before
that stage materializes, `computeStats()` falls back to `logicalPlan.stats` and
yields `rowCount = 0`. `Join.maxRows` for `LeftOuter` is `max(left * right,
left)`, which correctly propagates the 0 upward. `EliminateLimits` then sees
`maxRows = 0 <= 1` and drops the LIMIT.
The failure is timing-dependent because it depends on whether the stage has
materialized when re-optimization runs. This was introduced by SPARK-36424 (AQE
EliminateLimits support).
### Does this PR introduce _any_ user-facing change?
No API or behavior change. This fixes a correctness bug where queries with
a LIMIT over certain AQE plans could return too many rows.
### How was this patch tested?
Added a regression test in `AdaptiveQueryExecSuite`: "SPARK-57956: AQE must
not eliminate a global LIMIT over a row-increasing outer join". Because the
trigger is timing-dependent, the test runs the query 20 times and asserts each
result has exactly 1 row.
Verified:
- Without the fix, the test fails (returns 2 rows).
- With the fix, the test passes.
- The existing test "SPARK-36424: Support eliminate limits in AQE
Optimizer" still passes, confirming that legitimate limit elimination over
materialized aggregate stages is preserved.
Ran `AdaptiveQueryExecSuite` locally. Full test matrix relies on GitHub
Actions.
### Was this patch authored or co-authored using generative AI tooling?
Generated using Kiro (Claude Opus 4.8)
Closes #57347 from vboo123/SPARK-57956.
Lead-authored-by: Vaibhav Garg <[email protected]>
Co-authored-by: Vaibhav Garg <[email protected]>
Co-authored-by: Wenchen Fan <[email protected]>
Signed-off-by: Wenchen Fan <[email protected]>
(cherry picked from commit 5623395102680617162ec8e22e67485f44aa095e)
Signed-off-by: Wenchen Fan <[email protected]>
---
.../sql/execution/adaptive/LogicalQueryStage.scala | 16 ++++++++-
.../adaptive/AdaptiveQueryExecSuite.scala | 42 ++++++++++++++++++++--
2 files changed, 54 insertions(+), 4 deletions(-)
diff --git
a/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/LogicalQueryStage.scala
b/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/LogicalQueryStage.scala
index 62e00d1ea6ed..09fc97bb46fa 100644
---
a/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/LogicalQueryStage.scala
+++
b/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/LogicalQueryStage.scala
@@ -71,7 +71,21 @@ case class LogicalQueryStage(
physicalStats.getOrElse(logicalPlan.stats)
}
- override def maxRows: Option[Long] =
stats.rowCount.map(_.min(Long.MaxValue).toLong)
+ override def maxRows: Option[Long] = {
+ // A query stage's `rowCount` is an exact, valid upper bound only when it
comes from the
+ // runtime statistics of a materialized stage. Checking `isMaterialized`
alone is not enough:
+ // `computeStats()` can still fall back to `logicalPlan.stats` (a cost
estimate, with
+ // `isRuntime = false`) when the physical-stage lookup yields no
statistics, and that estimate
+ // can under-count (e.g., returning 0). Treating such an estimate as a
hard upper bound lets
+ // rules such as EliminateLimits wrongly drop a LIMIT and change the
result cardinality
+ // (SPARK-57956). Only trust `rowCount` when the stats are materialized
runtime stats;
+ // otherwise fall back to `logicalPlan.maxRows`, which is always a sound
upper bound.
+ if (isMaterialized && stats.isRuntime) {
+ stats.rowCount.map(_.min(Long.MaxValue).toLong)
+ } else {
+ logicalPlan.maxRows
+ }
+ }
override def isMaterialized: Boolean = physicalPlan.exists {
case s: QueryStageExec => s.isMaterialized
diff --git
a/sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/AdaptiveQueryExecSuite.scala
b/sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/AdaptiveQueryExecSuite.scala
index f84f3ad5cb6c..f17b56fda894 100644
---
a/sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/AdaptiveQueryExecSuite.scala
+++
b/sql/core/src/test/scala/org/apache/spark/sql/execution/adaptive/AdaptiveQueryExecSuite.scala
@@ -31,10 +31,10 @@ import org.apache.spark.scheduler.{SparkListener,
SparkListenerEvent, SparkListe
import org.apache.spark.shuffle.sort.SortShuffleManager
import org.apache.spark.sql.{DataFrame, Dataset, Row, SparkSession}
import org.apache.spark.sql.catalyst.InternalRow
-import org.apache.spark.sql.catalyst.expressions.{Ascending, Attribute,
AttributeReference, EqualTo, IsNull, Or, SortOrder}
-import org.apache.spark.sql.catalyst.optimizer.{BuildLeft, BuildRight}
+import org.apache.spark.sql.catalyst.expressions.{Ascending, Attribute,
AttributeReference, EqualTo, IsNull, Literal, Or, SortOrder}
+import org.apache.spark.sql.catalyst.optimizer.{BuildLeft, BuildRight,
EliminateLimits}
import org.apache.spark.sql.catalyst.plans.{Inner, LeftAnti}
-import org.apache.spark.sql.catalyst.plans.logical.{Aggregate, Join, JoinHint,
LocalRelation, LogicalPlan}
+import org.apache.spark.sql.catalyst.plans.logical.{Aggregate, GlobalLimit,
Join, JoinHint, LeafNode, LocalRelation, LogicalPlan, Statistics}
import
org.apache.spark.sql.catalyst.plans.physical.{CoalescedNullAwareHashPartitioning,
SinglePartition}
import org.apache.spark.sql.classic.Strategy
import org.apache.spark.sql.execution._
@@ -63,6 +63,16 @@ import org.apache.spark.tags.SlowSQLTest
import org.apache.spark.util.ArrayImplicits._
import org.apache.spark.util.Utils
+/**
+ * A leaf whose cost estimate under-counts its structural row bound. Used by
SPARK-57956 to verify
+ * that an unmaterialized query stage exposes the structural `maxRows` instead
of an estimate.
+ */
+private case class UnderCountLeaf(output: Seq[Attribute]) extends LeafNode {
+ override def maxRows: Option[Long] = Some(2L)
+ override def computeStats(): Statistics =
+ Statistics(sizeInBytes = BigInt(1), rowCount = Some(BigInt(0)))
+}
+
@SlowSQLTest
class AdaptiveQueryExecSuite
extends SharedSparkSession
@@ -3398,6 +3408,32 @@ class AdaptiveQueryExecSuite
}
}
+ test("SPARK-57956: unmaterialized query stage exposes structural maxRows,
not the estimate") {
+ // Build a LogicalQueryStage whose underlying stage is never materialized.
Its computeStats()
+ // falls back to the logical plan's cost estimate - here an under-count of
0 rows - which must
+ // NOT be promoted to a hard maxRows bound. Otherwise EliminateLimits
would drop a LIMIT that
+ // still needs to be applied once the stage runs (SPARK-57956).
+ val output = Seq(AttributeReference("a", IntegerType)())
+ val logical = UnderCountLeaf(output)
+ val scan = LocalTableScanExec(output, Nil, None)
+ val exchange = BroadcastExchangeExec(
+ HashedRelationBroadcastMode(output, isNullAware = false), scan)
+ val queryStage = LogicalQueryStage(logical, BroadcastQueryStageExec(0,
exchange, exchange))
+
+ assert(!queryStage.isMaterialized)
+ // The under-counted estimate is what computeStats() surfaces...
+ assert(queryStage.stats.rowCount.contains(BigInt(0)))
+ // ...but maxRows must remain the structural bound (2), not the estimate
(0).
+ assert(queryStage.maxRows.contains(2L),
+ "an unmaterialized stage must expose its structural maxRows, not the
row-count estimate")
+
+ // Since the structural bound (2) exceeds the limit (1), the LIMIT must be
retained. With the
+ // estimate wrongly promoted (maxRows = 0), EliminateLimits would instead
drop it.
+ val limited = GlobalLimit(Literal(1), queryStage)
+ assert(EliminateLimits(limited).isInstanceOf[GlobalLimit],
+ "LIMIT must be retained when the child's structural row bound exceeds
the limit")
+ }
+
test("SPARK-48037: Fix SortShuffleWriter lacks shuffle write related metrics
" +
"resulting in potentially inaccurate data") {
withTable("t3") {
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
