(spark) branch branch-4.x updated: [SPARK-57956][SQL] Fix AQE incorrectly eliminating a global LIMIT over a non-materialized query stage

[we...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

cloud-fan pushed a commit to branch branch-4.x
in repository https://gitbox.apache.org/repos/asf/spark.git


The following commit(s) were added to refs/heads/branch-4.x by this push:
     new bc51e9a16de6 [SPARK-57956][SQL] Fix AQE incorrectly eliminating a global LIMIT over a non-materialized query stage
bc51e9a16de6 is described below

commit bc51e9a16de6f58110d3c4b4d09ca754ddb5e510
Author: Vaibhav Garg <89...@users.noreply.github.com>
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 <89...@users.noreply.github.com>
    Co-authored-by: Vaibhav Garg <ga...@amazon.com>
    Co-authored-by: Wenchen Fan <we...@databricks.com>
    Signed-off-by: Wenchen Fan <we...@databricks.com>
    (cherry picked from commit 5623395102680617162ec8e22e67485f44aa095e)
    Signed-off-by: Wenchen Fan <we...@databricks.com>
---
 .../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: commits-unsubscribe@spark.apache.org
For additional commands, e-mail: commits-help@spark.apache.org