From 323bd56e2a63b72d0b6fe0ab29afbab36bb13ee6 Mon Sep 17 00:00:00 2001 From: Vaibhav Garg Date: Fri, 17 Jul 2026 23:12:26 +0000 Subject: [PATCH 1/2] [SPARK-57956][SQL] Fix AQE incorrectly eliminating a global LIMIT over a non-materialized query stage LogicalQueryStage.maxRows was derived from stats.rowCount, which before a query stage is materialized falls back to the logical plan's cost estimate. That estimate can under-count (e.g. return 0), and EliminateLimits then treats it as a hard upper bound and wrongly removes a global LIMIT, changing the result cardinality. maxRows now trusts the runtime row count only when the stage is materialized, otherwise falling back to logicalPlan.maxRows, which is always a sound upper bound. Generated using Kiro (Claude Opus 4.8) --- .../adaptive/LogicalQueryStage.scala | 14 ++++++++++++- .../adaptive/AdaptiveQueryExecSuite.scala | 20 +++++++++++++++++++ 2 files changed, 33 insertions(+), 1 deletion(-) 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 62e00d1ea6eda..6e43ef7703472 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,19 @@ 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 runtime `rowCount` is an exact, valid upper bound only once the stage + // is materialized. Before materialization, `stats` falls back to the logical plan's cost + // estimate, which can under-count (e.g. return 0) and must not be treated as a hard maximum + // - otherwise rules such as EliminateLimits can wrongly drop a LIMIT and change the result + // cardinality (SPARK-57956). When the stage is not materialized, fall back to the logical + // plan's `maxRows`, which is always a sound upper bound. + if (isMaterialized) { + 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 7446168c84fc8..6d3e5e1c493b5 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 @@ -3380,6 +3380,26 @@ class AdaptiveQueryExecSuite } } + test("SPARK-57956: AQE must not eliminate a global LIMIT over a row-increasing outer join") { + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true", + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", + SQLConf.ADAPTIVE_AUTO_BROADCASTJOIN_THRESHOLD.key -> "1048576") { + // The re-optimization that can wrongly drop the root LIMIT depends on whether the + // OFFSET-side query stage is materialized when EliminateLimits runs, which is timing + // dependent. Repeat the query so the regression is caught deterministically. + (1 to 20).foreach { _ => + val rows = 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 + """.stripMargin).collect() + assert(rows.length == 1, s"LIMIT 1 must cap the result at one row, but got ${rows.length}") + } + } + } + test("SPARK-48037: Fix SortShuffleWriter lacks shuffle write related metrics " + "resulting in potentially inaccurate data") { withTable("t3") { From 43b85591878f29783d194aa3c413c3483e710a74 Mon Sep 17 00:00:00 2001 From: Vaibhav Garg Date: Thu, 30 Jul 2026 01:28:11 +0000 Subject: [PATCH 2/2] [SPARK-57956][SQL] Address review: guard maxRows on stats.isRuntime and add a deterministic test Per review feedback from cloud-fan: - maxRows now trusts the runtime rowCount only when the stats are actually runtime stats (isMaterialized && stats.isRuntime). computeStats() can fall back to logicalPlan.stats (an estimate, isRuntime = false) even for a materialized stage when the physical-stage lookup yields no statistics, so isMaterialized alone was insufficient to prevent promoting an estimate to a hard upper bound. - Replace the timing-dependent 20x repetition with a deterministic unit test that constructs an unmaterialized LogicalQueryStage whose estimate under-counts its structural row bound, then asserts maxRows stays structural and EliminateLimits retains the LIMIT. Verified the test fails on the unguarded implementation. Generated using Kiro (Claude Opus 4.8) --- .../adaptive/LogicalQueryStage.scala | 16 ++--- .../adaptive/AdaptiveQueryExecSuite.scala | 58 ++++++++++++------- 2 files changed, 46 insertions(+), 28 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 6e43ef7703472..8aeba3c252eef 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 @@ -72,13 +72,15 @@ case class LogicalQueryStage( } override def maxRows: Option[Long] = { - // A query stage's runtime `rowCount` is an exact, valid upper bound only once the stage - // is materialized. Before materialization, `stats` falls back to the logical plan's cost - // estimate, which can under-count (e.g. return 0) and must not be treated as a hard maximum - // - otherwise rules such as EliminateLimits can wrongly drop a LIMIT and change the result - // cardinality (SPARK-57956). When the stage is not materialized, fall back to the logical - // plan's `maxRows`, which is always a sound upper bound. - if (isMaterialized) { + // 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. return 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 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 6d3e5e1c493b5..a166ea6b2615c 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 @@ -29,10 +29,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._ @@ -61,6 +61,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 @@ -3380,24 +3390,30 @@ class AdaptiveQueryExecSuite } } - test("SPARK-57956: AQE must not eliminate a global LIMIT over a row-increasing outer join") { - withSQLConf( - SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true", - SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", - SQLConf.ADAPTIVE_AUTO_BROADCASTJOIN_THRESHOLD.key -> "1048576") { - // The re-optimization that can wrongly drop the root LIMIT depends on whether the - // OFFSET-side query stage is materialized when EliminateLimits runs, which is timing - // dependent. Repeat the query so the regression is caught deterministically. - (1 to 20).foreach { _ => - val rows = 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 - """.stripMargin).collect() - assert(rows.length == 1, s"LIMIT 1 must cap the result at one row, but got ${rows.length}") - } - } + 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 " +