Describe the bug
AuronConverters builds fresh plan nodes when it converts a Spark plan to a native one, and it copies only its own four tags onto the replacement. Spark's logicalLink is not among them, so a converted node carries no logical link.
Shims.setLogicalLink exists for exactly this and has had no callers since 441a1a24:
spark-extension/src/main/scala/org/apache/spark/sql/auron/Shims.scala:234 — declarationspark-extension-shims-spark/src/main/scala/org/apache/spark/sql/auron/ShimsImpl.scala:525 — implementation- no call sites anywhere in the repository
On Spark 3.x this is harmless. On Spark 4.x it is not, because of where the link is asserted.
AdaptiveSparkPlanExec.setLogicalLinkForNewQueryStage asserts link.isDefined, and that 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 — getFinalPhysicalPlan() returns early when isFinalPlan is set — so the second collect creates no new query stage. Spark 4.0 re-enters createQueryStages(..., firstRun = true) on every collect; the comment on that branch names the case directly ("e.g, when we do df.collect multiple times"). On that second pass the plan handed to the assertion is the one produced after preColumnarTransitions, that is, after Auron's rewrite, and the link is gone.
The first collect is unaffected because it runs against the pre-columnar plan. QueryStageExec is a LeafExecNode, so link-setting does not recurse into a stage's inner plan either.
To Reproduce
On Spark 4.x, execute the same DataFrame twice where the final plan is fully native and contains no exchange:
val df = spark.sql("select c1 from t1 where c2 > (select max(c3) from t2)")
df.collect()
df.collect() // failsorg.apache.spark.SparkException: [INTERNAL_ERROR] The "collect" action failed.
Cause: java.lang.AssertionError: assertion failed
at org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanExec.setLogicalLinkForNewQueryStage(AdaptiveSparkPlanExec.scala:725)
at org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanExec.createQueryStages(AdaptiveSparkPlanExec.scala:545)
Any shuffle in the tree masks it, since a ShuffleQueryStageExec satisfies the assertion on its own, and so does any partial fallback that leaves a Spark node carrying a link. In the query above the only shuffle sits inside the scalar subquery's own plan, which the walk does not descend into, so the tree carries no link anywhere.
Auron sets ADAPTIVE_EXECUTION_FORCE_APPLY to true, so AQE is always in the path.
Expected behavior
Executing the same DataFrame more than once succeeds, as it does on Spark 3.x and as it does without Auron.
Additional context
Surfaced by a test added in #2490, which collected the same DataFrame twice. That test has been corrected there, since the double collect was not what it meant to check; this issue is about the missing propagation, which is independent of that PR and predates it. Neither AuronConverters.scala nor ShimsImpl.scala is touched by #2490.
The likely fix is to restore the Shims.get.setLogicalLink call in the conversion path. It wants care rather than a one-line patch: tryConvert is the single funnel for the 34 plan types convertSparkPlan handles, so setting the link there changes what AQE observes for every converted node on every supported Spark version. It needs its own test matrix across Spark 3.x and 4.x, covering both the repeated-execution case and the absence of regressions in ordinary queries.
I am not aware of a user report of this; the impact above is derived from the code path rather than observed in a deployment.
I will take a first pass at the fix.
Describe the bug
AuronConvertersbuilds fresh plan nodes when it converts a Spark plan to a native one, and it copies only its own four tags onto the replacement. Spark'slogicalLinkis not among them, so a converted node carries no logical link.Shims.setLogicalLinkexists for exactly this and has had no callers since441a1a24:spark-extension/src/main/scala/org/apache/spark/sql/auron/Shims.scala:234— declarationspark-extension-shims-spark/src/main/scala/org/apache/spark/sql/auron/ShimsImpl.scala:525— implementationOn Spark 3.x this is harmless. On Spark 4.x it is not, because of where the link is asserted.
AdaptiveSparkPlanExec.setLogicalLinkForNewQueryStageassertslink.isDefined, and that 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 —getFinalPhysicalPlan()returns early whenisFinalPlanis set — so the second collect creates no new query stage. Spark 4.0 re-enterscreateQueryStages(..., firstRun = true)on every collect; the comment on that branch names the case directly ("e.g, when we dodf.collectmultiple times"). On that second pass the plan handed to the assertion is the one produced afterpreColumnarTransitions, that is, after Auron's rewrite, and the link is gone.The first collect is unaffected because it runs against the pre-columnar plan.
QueryStageExecis aLeafExecNode, so link-setting does not recurse into a stage's inner plan either.To Reproduce
On Spark 4.x, execute the same
DataFrametwice where the final plan is fully native and contains no exchange:Any shuffle in the tree masks it, since a
ShuffleQueryStageExecsatisfies the assertion on its own, and so does any partial fallback that leaves a Spark node carrying a link. In the query above the only shuffle sits inside the scalar subquery's own plan, which the walk does not descend into, so the tree carries no link anywhere.Auron sets
ADAPTIVE_EXECUTION_FORCE_APPLYto true, so AQE is always in the path.Expected behavior
Executing the same
DataFramemore than once succeeds, as it does on Spark 3.x and as it does without Auron.Additional context
Surfaced by a test added in #2490, which collected the same
DataFrametwice. That test has been corrected there, since the double collect was not what it meant to check; this issue is about the missing propagation, which is independent of that PR and predates it. NeitherAuronConverters.scalanorShimsImpl.scalais touched by #2490.The likely fix is to restore the
Shims.get.setLogicalLinkcall in the conversion path. It wants care rather than a one-line patch:tryConvertis the single funnel for the 34 plan typesconvertSparkPlanhandles, so setting the link there changes what AQE observes for every converted node on every supported Spark version. It needs its own test matrix across Spark 3.x and 4.x, covering both the repeated-execution case and the absence of regressions in ordinary queries.I am not aware of a user report of this; the impact above is derived from the code path rather than observed in a deployment.
I will take a first pass at the fix.