apache / apache/auron

Converted native plans drop logicalLink, breaking repeated DataFrame execution on Spark 4

Open
#2,491 0 comments 0 reactions 0 assignees View on GitHub
Dominant language
Rust
Stars
1.8k
Forks
241
Avg merge
2d 12h
Merged PRs (30d)
21

Description

### 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` — declaration
- `spark-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:

```scala
val df = spark.sql("select c1 from t1 where c2 > (select max(c3) from t2)")
df.collect()
df.collect() // fails
```

```
org.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.

Contributor guide

Open the contributing guide

Research direction

Start in spark-extension/src/main/scala/org/apache/spark/sql/auron/AuronConverters.scala, tracing tryConvert and the conversion path into the 34 handled plan types. Read Shims.scala:234 and ShimsImpl.scala:525, then reproduce the repeated collect on Spark 4.x and run the relevant Spark 3.x and 4.x tests. Done means repeated execution succeeds while ordinary queries remain unaffected.

Written by the indexing model from the issue text.

Assessment

Tech stack
scala
Domain
data-engineering, distributed-systems
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Active
Clarity
Clearly specified
Newbie friendliness
48/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.