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..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,16 +710,22 @@ 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) } - 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..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 @@ -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,22 @@ 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))) + // `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") { def findIMRInnerChild(p: SparkPlan): SparkPlan = { val tableCache = find(p) {