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..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 @@ -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 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..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,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") {