andygrove opened a new pull request, #2682: URL: https://github.com/apache/datafusion-ballista/pull/2682
# Which issue does this PR close? Closes #2678. # Rationale for this change A client that builds its own physical plan can submit it with `submit_physical_plan` or a `PhysicalPlan` query, but that path always used the static planner, even though adaptive planning is on by default. Those jobs got none of AQE's runtime decisions and skipped `BalanceAndPruneFileGroups`. This matters for clients such as Comet Native Executors, which submit whole queries as physical plans with their own extension operators between DataFusion repartitions. # What changes are included in this PR? - `TaskManager::submit_plan` routes `SubmitPlan::Physical` through `AdaptiveExecutionGraph` when `ballista.planner.adaptive.enabled` is set, and keeps the static planner otherwise, matching logical submissions. - `AdaptivePlanner::try_new_from_physical` and `AdaptiveExecutionGraph::try_new_from_physical` run the plan-preparation rules over the submitted plan and hand it to the adaptive optimizer chain. There is no logical plan, so `EXPLAIN` handling does not apply on this path. - New `PinSubmittedShufflesRule`, run once during plan preparation on this path only. It turns the submitted plan's hash `RepartitionExec`s into `ExchangeExec`s before DataFusion's chain runs. Without it, `EnforceDistribution` removes any repartition that no operator requirement asks for, so a shuffle the submitter placed under an extension operator (which declares no requirement) disappeared. The new test `should_keep_submitted_shuffle_under_extension_operator` reproduces that. - `CoalescePartitionsRule` now checks every operator between the stage root and a leaf exchange against a list of DataFusion and Ballista operators known to produce the same rows over any co-partitioned input. A stage whose shuffles feed an operator outside that list, such as a client's extension operator, is not coalesced, since the operator may depend on the partition count. Existing coalesce tests are unchanged, so the list covers the plans DataFusion produces. - `CoalescePartitionsExec` is not pinned. DataFusion re-derives it from the consumer's `required_input_distribution`, so an extension operator that needs a single partition must declare that. The architecture guide now says so. Tests: `test_submit_physical_plan_uses_adaptive_planner` checks the routing; the two tests in `aqe/test/physical_submission.rs` cover shuffle pinning and the coalesce guard with a pass-through extension operator; the existing standalone end-to-end test `physical_plan_submission` now exercises the adaptive path, since it runs with the default configuration. # Are there any user-facing changes? Physical-plan submissions are planned adaptively when `ballista.planner.adaptive.enabled` is true (the default). Setting it to false restores the previous behaviour. Docs updated in `docs/source/contributors-guide/architecture.md`. No public API changes; the doc comments on `submit_physical_plan` no longer say AQE is unavailable for this path. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
