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

[james-willis]

What changes were proposed in this pull request?

InMemoryRelation.newInstance() now goes through withOutput, so that outputOrdering is re-mapped onto the freshly-instantiated attributes instead of being carried over unchanged.

withOutput also re-keys statsOfPlanToCache onto the new attributes via LogicalRDD.rewriteStatistics. It previously passed the stats through unchanged, so Statistics.attributeStats stayed keyed by the old attributes and every column-stat lookup missed on the new relation, silently dropping CBO estimates back to the un-filtered defaults. This matches what LogicalRDD.newInstance() already does. Thanks to @peter-toth for catching it.

Why are the changes needed?

InMemoryRelation has an implicit invariant that outputOrdering may only reference attributes present in output. newInstance() violates it: it gives output fresh exprIds but passes outputOrdering through unchanged, so the returned relation's ordering still points at the old attributes.

That was harmless until SPARK-53738, which routed doCanonicalize through withOutput and made withOutput re-map the ordering with a strict AttributeMap lookup. Since then, any InMemoryRelation that has been through newInstance() fails as soon as anything canonicalizes it:

java.util.NoSuchElementException: key not found: k#1L
  at scala.collection.MapOps.default(Map.scala:289)
  at org.apache.spark.sql.catalyst.expressions.AttributeMap.apply(AttributeMap.scala:41)
  at org.apache.spark.sql.execution.columnar.InMemoryRelation.$anonfun$withOutput$1(InMemoryRelation.scala:711)
  at org.apache.spark.sql.execution.columnar.InMemoryRelation.withOutput(InMemoryRelation.scala:711)
  at org.apache.spark.sql.execution.columnar.InMemoryRelation.doCanonicalize(InMemoryRelation.scala:672)
  ...
  at org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanExec.createNonResultQueryStages(AdaptiveSparkPlanExec.scala:589)

This is reachable from ordinary SQL. Cache substitution (CacheManager.useCachedData) runs on the analyzed plan, before the optimizer. InlineCTE then inlines a CTE that is referenced more than once and, to get a fresh-exprId copy, runs DeduplicateRelations over a synthetic self-join. By that point the plan already contains the InMemoryRelation, which is a MultiInstanceRelation, so newInstance() is called on it.

A user-facing repro — a persisted DataFrame with a global ORDER BY, window-ranked and then self-joined:

spark.range(0, 20).selectExpr("id", "id % 3 AS k").createOrReplaceTempView("t")

b = spark.sql("SELECT id, k FROM t ORDER BY k, id")
b.persist()
b.count()
b.createOrReplaceTempView("b")

spark.sql("""
  WITH r AS (SELECT *, row_number() OVER (PARTITION BY k ORDER BY id DESC) AS rn FROM b)
  SELECT x.id AS p, y.id AS q
  FROM r x JOIN r y ON x.k = y.k AND x.rn = 1 AND y.rn = 2
""").show()

This fails on 4.0.2 and later. It succeeds on 4.0.1, which predates SPARK-53738. The user sees only an internal NoSuchElementException at the first action, with nothing actionable in it — the query itself is well formed. I verified the failure on 4.0.4, 4.1.3 and 4.2.0.

See SPARK-59009 for the full analysis.

Does this PR introduce any user-facing change?

No, other than the bug fix itself: queries that reference a cached relation with a non-empty outputOrdering more than once now succeed instead of failing with an internal error.

newInstance() also preserves the ordering now rather than returning a relation with a stale one, so the ordering remains usable as an optimization hint for the new instance.

How was this patch tested?

Two new tests, both of which fail on unmodified master and pass with the change:

• InMemoryRelationSuite, a unit test asserting that after newInstance() the ordering references the new attributes and that the result canonicalizes without throwing.
• CachedTableSuite, an end-to-end test running the CTE self-join over a cached, ordered relation and checking the answer. Without the change it fails with NoSuchElementException: key not found: k#....

InMemoryRelationSuite, CachedTableSuite and DatasetCacheSuite are green with the change.

build/sbt "sql/testOnly org.apache.spark.sql.execution.columnar.InMemoryRelationSuite"
build/sbt "sql/testOnly org.apache.spark.sql.CachedTableSuite"

Was this patch authored or co-authored using generative AI tooling?

Generated-by: Claude Code (model claude-opus-5)

[james-willis]

@pan3793 or @peter-toth may be good reviewers for this change as they have recently worked with this code.

[pan3793]

LGTM, cc @ulysses-you

[peter-toth]

Thanks for the PR, @james-willis!

I traced this independently and land on the same diagnosis. SPARK-53738 made withOutput re-map outputOrdering through a strict AttributeMap and routed doCanonicalize through it, and newInstance() was the one place left that refreshed output without refreshing the ordering. Routing it through withOutput is the right minimal fix, and it makes newInstance() do exactly what the existing SPARK-46779 test above already does by hand. I reverted just the one-line main change in a worktree and confirmed both new tests fail on base, the CachedTableSuite one with the exact NoSuchElementException: key not found: k#77L. With the fix, CachedTableSuite + DatasetCacheSuite + InMemoryColumnarQuerySuite are green (159 tests). Nothing blocking from me.

Non-blocking

• 1. statsOfPlanToCache is still keyed by the old attributes: withOutput re-maps the ordering but passes the stats through unchanged, so after newInstance() every attributeStats lookup misses. Measured a 10x row-count overestimate on the new copy. The sibling LogicalRDD.newInstance() re-maps its stats. [inline: sql/core/src/main/scala/org/apache/spark/sql/execution/columnar/InMemoryRelation.scala:720]

Minor

• 2. Assert the canonical form, not just that it doesn't throw: assert(r1.sameResult(r2)) is the property that actually matters and it still fails on base. [inline: sql/core/src/test/scala/org/apache/spark/sql/execution/columnar/InMemoryRelationSuite.scala:49]

[uros-b]

+1, thank you @james-willis!

[james-willis]

Thanks all for your approvals. Is it possible for someone to merge this PR? I cannot as I am not a committer.

[peter-toth]

Merge Summary:

• merged into master 38fc867
• merged into branch-4.x 0b67dc2
• merged into branch-4.3 bf536dc
• merged into branch-4.2 77a245a

Posted by merge_spark_pr.py

[peter-toth]

Thank you @james-willis for the fix and everyone for the review.

@james-willis, can you please open backport PRs to 4.1 and 4.0 as those were not clean cherry-picks...

[james-willis]

Edit: these are the wrong backports, for a different bug. Please ignore

@peter-toth
4.1: #58358
4.0: #58359

[peter-toth]

@james-willis, those PR seems to be very different to this one.

[james-willis]

@peter-toth sorry got my wires crossed! Will raise the right PR

[james-willis]

@peter-toth Ok here are the correct backports. CI is still running.

#58388
#58389