Skip to content

fix: preserve current AQE logical-stage links on Comet operators - #5483

Open
sunchao wants to merge 1 commit into
apache:mainfrom
sunchao:fix/aqe-logical-stage-links-20260826
Open

fix: preserve current AQE logical-stage links on Comet operators#5483
sunchao wants to merge 1 commit into
apache:mainfrom
sunchao:fix/aqe-logical-stage-links-20260826

Conversation

@sunchao

@sunchao sunchao commented Aug 26, 2026

Copy link
Copy Markdown
Member

Which issue does this PR close?

Closes #5482.

Rationale for this change

AQE can reuse a native final aggregate above a shuffle stage and assign its physical root a direct link to the current LogicalQueryStage. Comet's logical-link repair restores the older originalPlan.logicalLink, or clears the link when the original has none.

That older logical aggregate is hidden inside the logical-stage leaf. A later exchange can inherit the stale link, preventing Spark's identity-based replacement from inserting it into the active logical tree. The generated SQL regression demonstrates missed native broadcast adaptation on OSS Spark.

What changes are included in this PR?

  • Preserve an existing direct LogicalQueryStage tag when revisiting a CometExec.
  • Retain the existing repair and clearing behavior for ordinary and inherited links. Exchange handling, including the shuffle empty-link invariant from CometShuffleExchangeExec logical link is different to Spark ShuffleExchangeExec #323, is unchanged.
  • Add planner regressions for repeated replanning with direct, inherited, and absent original links.
  • Add a generated-data execution regression that verifies adaptive native broadcasts, results, direct stage links, output row counts, and native compute metrics.

No Spark changes, native execution changes, or configuration changes are required.

How are these changes tested?

Local validation on Linux x86_64 with JDK 17, Spark 4.1.3, and Scala 2.13.17:

  • Built the current OSS native source with make core (debug build, debug symbols disabled).
  • Ran Maven from the reactor root: all 31 CometExecRuleSuite tests, the new AQE broadcast test, and the existing shuffle logical-link test passed (33 tests).
  • Removed only the production fix and reran both new regressions: both failed. The planner test lost the current direct stage identity; the SQL query returned the correct result but had zero native broadcast hash joins instead of two. Restored the fix and reran the full focused selection: all 33 tests passed again.
  • With the fix, the SQL result is 48738816; both broadcast aggregate producers output 64 rows, record positive native compute time, and retain direct logical-stage links.
  • Scalastyle, Spotless, and git diff --check passed.

Local runtime coverage is Spark 4.1.3; the other Spark profiles and full suite remain for CI.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

CometExecRule overwrites direct AQE LogicalQueryStage links during replanning

1 participant