[SPARK-59009][SQL] Re-map outputOrdering in InMemoryRelation.newInstance()

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 = {

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(

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) {