[SPARK-57956][SQL] Fix AQE incorrectly eliminating a global LIMIT over a non-materialized query stage

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. 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
+    }
+  }
 
   override def isMaterialized: Boolean = physicalPlan.exists {
     case s: QueryStageExec => s.isMaterialized

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,6 +3390,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") {