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

[james-willis]

Backport of #58293 to branch-4.0.

One adaptation for 4.0: InMemoryRelation.scala on branch-4.0 imports from org.apache.spark.sql.execution explicitly rather than with a wildcard, so LogicalRDD is added to that import list for the rewriteStatistics call.

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)

[peter-toth]

Merge Summary:

• merged into branch-4.0 15f4285

Posted by merge_spark_pr.py

[peter-toth]

Thank you @james-willis and @uros-b.