[SPARK-47148] Avoid to materialize AQE ExchangeQueryStageExec on the cancellation

• Type: Bug
• Status: Resolved
• Priority: Major
• Resolution: Fixed
• Affects Version/s: 4.0.0
• Fix Version/s: 4.0.0
• Component/s: Shuffle, SQL
• Labels: pull-request-available
• Assignee: Eren Avsarogullari
• Reporter: Eren Avsarogullari
• Created: 23/Feb/24 18:00
• Updated: 23/Jul/24 02:00
• Resolved: 29/Apr/24 14:12

Description

AQE can materialize both ShuffleQueryStage and BroadcastQueryStage on the cancellation. This causes unnecessary stage materialization by submitting Shuffle Job and Broadcast Job. Under normal circumstances, if the stage is already non-materialized (a.k.a ShuffleQueryStage.shuffleFuture or BroadcastQueryStage.broadcastFuture is not initialized yet), it should just be skipped without materializing it.

Please find sample use-case:
1- Stage Materialization Steps:
When stage materialization is failed:

1.1- ShuffleQueryStage1 - is materialized successfully,
1.2- ShuffleQueryStage2 - materialization is failed,
1.3- ShuffleQueryStage3 - Not materialized yet so ShuffleQueryStage3.shuffleFuture is not initialized yet

2- Stage Cancellation Steps:

2.1- ShuffleQueryStage1 - is canceled due to already materialized,
2.2- ShuffleQueryStage2 - is earlyFailedStage so currently, it is skipped as default by AQE because it could not be materialized,
2.3- ShuffleQueryStage3 - Problem is here: This stage is not materialized yet but currently, it is also tried to cancel and this stage requires to be materialized first.

Issue Links

• links to: GitHub Pull Request #45234
• links to: GitHub Pull Request #47007