Skip to content

Conversation

@baibaichen
Copy link
Contributor

What changes were proposed in this pull request?

This PR uses collectFirst to find the first AdaptiveSparkPlanExec node anywhere in the plan tree, instead of assuming the root plan is an AdaptiveSparkPlanExec.

Why are the changes needed?

#52157 introduced the extractShuffleIds method in SQLExecution to find shuffle IDs of SparkPlan. Previously, the method implicitly assumed that if AQE is enabled, the AdaptiveSparkPlanExec would be at the root of the input. Since Spark only inserts AdaptiveSparkPlanExec under Command, this assumption was fine. However, the AdaptiveSparkPlanExec may not be the root node in Gluten. Gluten needs to insert a special physical plan to do column to row transition.

By using collectFirst, we can correctly locate the AdaptiveSparkPlanExec regardless of its position in the plan tree, which improves compatibility.

Does this PR introduce any user-facing change?

No.

How was this patch tested?

Pass GHA.

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

No

@github-actions github-actions bot added the SQL label Dec 26, 2025
Copy link
Member

@yaooqinn yaooqinn left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

+1, LGTM

@yaooqinn
Copy link
Member

@baibaichen SPARK-53413 is shipped with spark 4.1.0, you need re-file a JIRA ticket for the PR.

@baibaichen baibaichen changed the title [SPARK-53413][SQL][FOLLOWUP] Fix extractShuffleIds to find AdaptiveSparkPlanExec anywhere in plan tree [SPARK-54850][SQL][FOLLOWUP] Improve extractShuffleIds to find AdaptiveSparkPlanExec anywhere in plan tree Dec 26, 2025
@baibaichen
Copy link
Contributor Author

@baibaichen SPARK-53413 is shipped with spark 4.1.0, you need re-file a JIRA ticket for the PR.

SPARK-54850 is created for this PR

@yaooqinn yaooqinn changed the title [SPARK-54850][SQL][FOLLOWUP] Improve extractShuffleIds to find AdaptiveSparkPlanExec anywhere in plan tree [SPARK-54850][SQL] Improve extractShuffleIds to find AdaptiveSparkPlanExec anywhere in plan tree Dec 26, 2025
@yaooqinn
Copy link
Member

@baibaichen Can you retrigger the GHA?

@baibaichen baibaichen closed this Dec 29, 2025
@baibaichen baibaichen reopened this Dec 29, 2025
@yaooqinn yaooqinn closed this in 082de7f Dec 29, 2025
yaooqinn pushed a commit that referenced this pull request Dec 29, 2025
…PlanExec` anywhere in plan tree

### What changes were proposed in this pull request?

This PR uses `collectFirst` to find the first `AdaptiveSparkPlanExec` node anywhere in the plan tree, instead of assuming the root plan is an `AdaptiveSparkPlanExec`.

### Why are the changes needed?

#52157 introduced the `extractShuffleIds` method in `SQLExecution` to find shuffle IDs of `SparkPlan`. Previously, the method implicitly assumed that if AQE is enabled, the `AdaptiveSparkPlanExec` would be at the root of the input. Since Spark only inserts `AdaptiveSparkPlanExec` under Command, this assumption was fine. However, the `AdaptiveSparkPlanExec` may not be the root node in Gluten. Gluten needs to insert a special physical plan to do column to row transition.

By using `collectFirst`, we can correctly locate the `AdaptiveSparkPlanExec` regardless of its position in the plan tree, which improves compatibility.

### Does this PR introduce any user-facing change?

No.

### How was this patch tested?

Pass GHA.

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

No

Closes #53620 from baibaichen/feature/extractShuffleIds.

Authored-by: Chang chen <baibaichen@gmail.com>
Signed-off-by: Kent Yao <kentyao@microsoft.com>
(cherry picked from commit 082de7f)
Signed-off-by: Kent Yao <kentyao@microsoft.com>
@yaooqinn
Copy link
Member

Merged to master and 4.1.1.

Thank you @baibaichen @cloud-fan

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

Labels

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants