From 9060a5a84c4bb6d2e57fc21ad4a8287d7528aa02 Mon Sep 17 00:00:00 2001 From: James Willis Date: Tue, 25 Aug 2026 14:31:47 -0700 Subject: [PATCH 1/2] [SPARK-59009][SQL] Re-map outputOrdering in InMemoryRelation.newInstance() `InMemoryRelation.newInstance()` gave `output` fresh exprIds but passed `outputOrdering` through unchanged, leaving the ordering referencing the old attributes. Since SPARK-53738 routed `doCanonicalize` through `withOutput`, which re-maps the ordering with a strict `AttributeMap` lookup, canonicalizing such a relation throws NoSuchElementException. Route `newInstance()` through `withOutput` so the ordering is re-mapped onto the fresh attributes. --- .../execution/columnar/InMemoryRelation.scala | 12 ++++----- .../apache/spark/sql/CachedTableSuite.scala | 27 +++++++++++++++++++ .../columnar/InMemoryRelationSuite.scala | 15 +++++++++++ 3 files changed, 47 insertions(+), 7 deletions(-) diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/columnar/InMemoryRelation.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/columnar/InMemoryRelation.scala index e6bf65ec89e6b..cc3c89a948f68 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/columnar/InMemoryRelation.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/columnar/InMemoryRelation.scala @@ -713,13 +713,11 @@ case class InMemoryRelation( InMemoryRelation(newOutput, cacheBuilder, newOutputOrdering, statsOfPlanToCache) } - override def newInstance(): this.type = { - InMemoryRelation( - output.map(_.newInstance()), - cacheBuilder, - outputOrdering, - statsOfPlanToCache).asInstanceOf[this.type] - } + // Goes through `withOutput` so that `outputOrdering` is re-mapped onto the fresh exprIds. + // Returning a relation whose `outputOrdering` still references the old attributes would break + // canonicalization, which re-maps the ordering through the relation's own `output`. + override def newInstance(): this.type = + withOutput(output.map(_.newInstance())).asInstanceOf[this.type] // override `clone` since the default implementation won't carry over mutable states. override def clone(): LogicalPlan = { diff --git a/sql/core/src/test/scala/org/apache/spark/sql/CachedTableSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/CachedTableSuite.scala index 90f991479801d..66732a9e52b00 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/CachedTableSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/CachedTableSuite.scala @@ -2727,6 +2727,33 @@ class CachedTableSuite extends SharedSparkSession } } + test("SPARK-59009: cached relation with outputOrdering can be referenced more than once") { + withTempView("t", "ordered_t") { + spark.range(0, 20).selectExpr("id", "id % 3 AS k").createOrReplaceTempView("t") + val ordered = sql("SELECT id, k FROM t ORDER BY k, id") + ordered.persist() + try { + ordered.createOrReplaceTempView("ordered_t") + // The CTE below is referenced twice, so InlineCTE deduplicates the second copy and calls + // newInstance() on the cached relation. The join then canonicalizes it, which used to + // throw because newInstance() left `outputOrdering` on the old attributes. + checkAnswer( + sql( + """ + |WITH r AS ( + | SELECT *, row_number() OVER (PARTITION BY k ORDER BY id DESC) AS rn + | FROM ordered_t + |) + |SELECT x.id AS x_id, y.id AS y_id + |FROM r x JOIN r y ON x.k = y.k AND x.rn = 1 AND y.rn = 2 + """.stripMargin), + Row(18L, 15L) :: Row(19L, 16L) :: Row(17L, 14L) :: Nil) + } finally { + ordered.unpersist() + } + } + } + private def cacheManager = spark.sharedState.cacheManager private def pinTable( diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/InMemoryRelationSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/InMemoryRelationSuite.scala index 2c73622739a51..1fc6568ccbb74 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/InMemoryRelationSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/InMemoryRelationSuite.scala @@ -18,6 +18,7 @@ package org.apache.spark.sql.execution.columnar import org.apache.spark.SparkFunSuite +import org.apache.spark.sql.catalyst.expressions.AttributeSet import org.apache.spark.sql.execution.SparkPlan import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper import org.apache.spark.sql.functions.expr @@ -34,6 +35,20 @@ class InMemoryRelationSuite extends SparkFunSuite assert(r1.sameResult(r2)) } + test("SPARK-59009: newInstance() re-maps outputOrdering onto the new attributes") { + val d = spark.range(10).selectExpr("id", "id % 3 AS k").orderBy("k", "id") + val r1 = InMemoryRelation(StorageLevel.MEMORY_ONLY, d.queryExecution, None) + assert(r1.outputOrdering.nonEmpty) + + val r2 = r1.newInstance() + assert(r2.output.map(_.exprId) != r1.output.map(_.exprId)) + // The ordering must be re-mapped onto the new attributes, not left referencing the old ones. + assert(r2.outputOrdering.nonEmpty) + assert(AttributeSet(r2.outputOrdering.flatMap(_.references)).subsetOf(AttributeSet(r2.output))) + // Canonicalization re-maps `outputOrdering` through `output`, so a stale ordering would throw. + r2.canonicalized + } + test("SPARK-47177: Cached SQL plan do not display final AQE plan in explain string") { def findIMRInnerChild(p: SparkPlan): SparkPlan = { val tableCache = find(p) { From 9cba6ebf440909d6846338a5d94b2d0a9773537d Mon Sep 17 00:00:00 2001 From: James Willis Date: Wed, 26 Aug 2026 01:26:27 -0700 Subject: [PATCH 2/2] Address review: re-map statsOfPlanToCache, assert sameResult `withOutput` re-mapped `outputOrdering` but passed `statsOfPlanToCache` through unchanged, leaving `Statistics.attributeStats` keyed by the old attributes. Every column stat lookup then missed for the new relation and CBO estimates fell back to the un-filtered defaults. Re-key the stats with `LogicalRDD.rewriteStatistics`, matching `LogicalRDD.newInstance()`. Also assert `sameResult` in the unit test rather than only that canonicalization does not throw, mirroring the SPARK-46779 test. --- .../sql/execution/columnar/InMemoryRelation.scala | 10 +++++++++- .../sql/execution/columnar/InMemoryRelationSuite.scala | 6 ++++-- 2 files changed, 13 insertions(+), 3 deletions(-) diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/columnar/InMemoryRelation.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/columnar/InMemoryRelation.scala index cc3c89a948f68..584eaa9ecbc82 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/columnar/InMemoryRelation.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/columnar/InMemoryRelation.scala @@ -710,7 +710,15 @@ case class InMemoryRelation( val newOutputOrdering = outputOrdering .map(_.transform { case a: Attribute => map(a) }) .asInstanceOf[Seq[SortOrder]] - InMemoryRelation(newOutput, cacheBuilder, newOutputOrdering, statsOfPlanToCache) + // `attributeStats` is keyed by attribute, so it has to be re-keyed onto `newOutput` as well, + // otherwise every column stat lookup misses for the new relation and the estimates silently + // fall back to the un-filtered defaults. `statsOfPlanToCache` is a `var` that starts as null. + val newStatsOfPlanToCache = if (statsOfPlanToCache == null) { + null + } else { + LogicalRDD.rewriteStatistics(statsOfPlanToCache, map) + } + InMemoryRelation(newOutput, cacheBuilder, newOutputOrdering, newStatsOfPlanToCache) } // Goes through `withOutput` so that `outputOrdering` is re-mapped onto the fresh exprIds. diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/InMemoryRelationSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/InMemoryRelationSuite.scala index 1fc6568ccbb74..6a9c8964e4dbf 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/InMemoryRelationSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/InMemoryRelationSuite.scala @@ -45,8 +45,10 @@ class InMemoryRelationSuite extends SparkFunSuite // The ordering must be re-mapped onto the new attributes, not left referencing the old ones. assert(r2.outputOrdering.nonEmpty) assert(AttributeSet(r2.outputOrdering.flatMap(_.references)).subsetOf(AttributeSet(r2.output))) - // Canonicalization re-maps `outputOrdering` through `output`, so a stale ordering would throw. - r2.canonicalized + // `sameResult` is what CacheManager lookups and exchange reuse rely on. It goes through + // `doCanonicalize`, which re-maps `outputOrdering` through `output`, so a stale ordering + // throws here rather than merely producing an unequal plan. + assert(r1.sameResult(r2)) } test("SPARK-47177: Cached SQL plan do not display final AQE plan in explain string") {