james-willis opened a new pull request, #58388: URL: https://github.com/apache/spark/pull/58388
Backport of #58293 to branch-4.1. ### 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](https://issues.apache.org/jira/browse/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: ```python 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](https://issues.apache.org/jira/browse/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) -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
