[AURON #2491] Carry Spark's logical link onto converted native plans

[weiqingy]

Which issue does this PR close?

Closes #2491

Rationale for this change

Converting a Spark plan node to a native one builds a new node, and Spark's logicalLink does not travel with it. On Spark 4 that is enough to fail a query.

AdaptiveSparkPlanExec.setLogicalLinkForNewQueryStage asserts that a new query stage's subtree carries a logical link. The assertion is the same in 3.5 and 4.0; what differs is when it sees Auron's output. Spark 3.5 short-circuits a repeated collect, since getFinalPhysicalPlan() returns early once isFinalPlan is set, so no new stage is created. Spark 4 re-enters createQueryStages(..., firstRun = true) on every collect, and the comment on that branch names the case directly ("e.g, when we do df.collect multiple times"). On that pass the plan handed to the assertion is the post-preColumnarTransitions one, and the link is gone.

So executing the same DataFrame twice fails on Spark 4 whenever the final plan is fully native and has no exchange in it:

val df = spark.sql("select c1 from t1 where c2 > (select max(c3) from t2)")
df.collect()
df.collect()   // AssertionError inside setLogicalLinkForNewQueryStage

Any shuffle in the tree masks it, because a ShuffleQueryStageExec satisfies the assertion on its own, and so does any partial fallback that leaves a Spark node carrying a link. Auron sets ADAPTIVE_EXECUTION_FORCE_APPLY, so AQE is always in the path.

The first collect is unaffected: it runs against the pre-columnar plan. QueryStageExec is a LeafExecNode, so link-setting does not descend into a stage's inner plan either.

What changes are included in this PR?

AuronConverters.convertSparkPlanRecursively now carries the original node's logical link onto the node that replaces it, alongside the four tags it already copies there. Shims.setLogicalLink already existed and is unchanged; this adds the call site.

The link is set only when conversion actually returned a different instance. When it returns the same node, that node already carries its link through copyTagsFrom, and setting it again would promote its inherited tag to a primary one and recurse into the original children.

The call sits in the recursive walk rather than in tryConvert for two reasons. tryConvert is not the only path that builds a replacement — Shims.get.convertMoreSparkPlan does so outside it. And convertSparkPlan runs twice per plan: AuronConvertStrategy converts the whole tree once purely to compute tags and discards the result, so a call in tryConvert would also fire on that discarded pass.

Some history, since the call is not new. It existed until 441a1a24 (2022-12-12), which moved setTagValue(convertibleTag, true) from the shim into tryConvert and dropped the setLogicalLink effect along with it. The shim body of that era also recursed over children, though under a guard that made the recursion inert — Spark's own setLogicalLink fans inherited tags downward first, so the guard never passed below the top node. The current shim has no such recursion.

Are there any user-facing changes?

Executing the same DataFrame more than once now succeeds on Spark 4, as it does on Spark 3 and as it does without Auron.

Converted nodes now carry a logical link where they previously carried none. Spark's setLogicalLink writes a primary tag on the node it is called on and fans inherited tags downward, halting at any descendant that already owns a primary tag, so a link is only ever added below a node that already had it — never moved upward. Both consumers that could be affected take the topmost match, so the answers they compute are unchanged. Every other reader of logicalLink in sql/core runs before the columnar rules, and the only rule running after Auron is CollapseCodegenStages, which does not read it.

One consumer was not proven unaffected and is worth watching rather than claiming: RemoveRedundantProjects.canRemove keys off logicalLink.isEmpty, and a pre-columnar ProjectExec that gains an inherited tag stops qualifying. The rule lists were traced and no path was found where it re-runs on nodes Auron has touched, but that is an absence of evidence rather than a proof. TPC-DS plan stability is the check that would catch it.

How was this patch tested?

Three cases in AuronQuerySuite, on the same fully native, exchange-free fixture.

The repeated execution case runs on every profile from spark-3.0 to spark-4.2. The two assertions on the executed plan run on spark-3.2 and above and are assumed away below it. Before Spark 3.2, TreeNode.withNewChildren is overridable and reaches copyTagsFrom only by way of makeCopy, and the native plan nodes override it with a plain copy of the case class, so no node tag survives a child rebuild there. The converter sets the link on every version, but on 3.0 and 3.1 adaptive execution drops it again as soon as it substitutes a query stage into the tree. That gap sits in the node overrides rather than in the conversion path, it predates this change, and it does not affect the bug being fixed, which is specific to Spark 4.

• the same DataFrame executed twice returns the right rows. This is the regression itself, and it is inert on Spark 3, where the second collect short-circuits
• every native node in the plan carries a logical link
• native nodes in one plan do not all collapse onto a single link, and the node realizing the join is linked to a Join

Each was proved to pin what it claims by mutation, run module-wide against the whole suite rather than a single class:

mutation	which test fails
the fix removed	all three (2 on Spark 3, where the first is inert)
the link stamped onto every descendant, collapsing them	the collapse check
the link stamped one level down, shifting every operator	the identity check, while the collapse check still passes

The third mutation is why both assertions are kept: a shift mislabels every operator while leaving the distinct count intact, so a cardinality check alone would pass it.

Measured on spark-3.5/scala-2.12 and spark-4.0.2/scala-2.13: 150 → 152 passing on 3.5, 149 → 152 on 4.0, with exactly the three new tests changing state and no other test affected in either direction, confirmed by set difference rather than by counting.

The identity assertion matches the three native join base types rather than a concrete class, and was re-run with broadcast joins disabled and sort-merge preferred to confirm it does not depend on a join strategy.

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

• Yes
• No

Generated-by: Claude Code (Opus 5)

[Copilot]

Pull request overview

This PR fixes a Spark 4 regression where native plan conversion drops Spark’s logicalLink, causing repeated execution of the same DataFrame (e.g., multiple collect() calls) to fail an AQE assertion when the final plan is fully native and exchange-free. It does this by explicitly propagating the source node’s logical link onto any newly constructed replacement node during recursive conversion, and adds targeted regression tests across Spark versions.

Changes:

• Propagate Spark logicalLink when convertSparkPlan returns a new (non-eq) replacement node during convertSparkPlanRecursively.
• Add regression coverage ensuring repeated execution succeeds and that native operators retain per-operator logical links (including join linkage).

Reviewed changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated no comments.

File	Description
spark-extension/src/main/scala/org/apache/spark/sql/auron/AuronConverters.scala	When conversion creates a new node instance, stamps the original node’s logicalLink onto the replacement via Shims.get.setLogicalLink.
spark-extension-shims-spark/src/test/scala/org/apache/auron/AuronQuerySuite.scala	Adds three tests covering repeated execution and validating logical link presence and per-operator distinctness (including join linkage).

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[Copilot]

Pull request overview

Copilot reviewed 2 out of 2 changed files in this pull request and generated no new comments.