From 323bd56e2a63b72d0b6fe0ab29afbab36bb13ee6 Mon Sep 17 00:00:00 2001
From: Vaibhav Garg <gargvboo@amazon.com>
Date: Fri, 17 Jul 2026 23:12:26 +0000
Subject: [PATCH] [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") {