Repository navigation
Conversation
A physical plan submitted with `submit_physical_plan` or a `PhysicalPlan` query always used the static planner, even with adaptive planning on. Route it through `AdaptiveExecutionGraph` when `ballista.planner.adaptive.enabled` is set, keeping the static planner otherwise. `AdaptivePlanner::try_new_from_physical` runs the plan-preparation rules over the submitted plan and hands it to the adaptive optimizer chain. Before that, `PinSubmittedShufflesRule` turns the plan's hash repartitions into exchanges: DataFusion's `EnforceDistribution` removes any repartition no operator requirement asks for, which would drop a shuffle the submitter placed under an extension operator. `CoalescePartitionsRule` now only coalesces a stage whose shuffles feed operators known to be indifferent to the partition count. An extension operator may depend on it, so such a stage is left alone.
# Conflicts: # ballista/scheduler/src/state/mod.rs
#2633 added a test to this module that names ExecutionPlan, and #2613 then moved FailingPlanRewriteExec out of the module and dropped the imports it had needed, ExecutionPlan among them. Each passed CI against its own base. Together they leave the scheduler's test target uncompilable, so cargo check, Clippy and the test jobs have failed on main since #2613 landed.
…ing it The static planner kept only hash repartitions as stage boundaries and dropped every other RepartitionExec. That is harmless for round-robin, which only spreads rows, but a range repartition is part of what a plan means. A plan that range-partitions and then sorts each partition (a distributed global sort) came back as partitions that were each sorted but not range-disjoint, with no error. A RepartitionExec with Partitioning::Range is now a stage boundary. The producer stage ends in a SortShuffleWriterExec that routes rows with DataFusion's own RangeExpr (the function RepartitionExec applies), so the consumer's partition i reads exactly range i. The writer shares everything after routing with the hash shuffle: buffering, spilling, the consolidated file and its index. The range travels in a new field of SortShuffleWriterExecNode; hash stages keep using the old one. Restricting a range-partitioned ShuffleReaderExec to a task's partitions now slices the split points along with them, instead of failing. A slice that does not ascend (a retried partition queued ahead of a later one) or is empty has no range partitioning to report, so it reports UnknownPartitioning rather than failing the task. A range repartition that must preserve its input order is rejected with an error: a shuffle cannot merge sorted inputs, and planning it anyway would return partitions that are range-disjoint but no longer sorted. RangeShuffleWriterExec and RangeShuffleReaderExec are not used. They need an ordered producer, while the repartition here is followed by the sort. Closes #2680
A range RepartitionExec in a submitted physical plan is now planned as a shuffle instead of being dropped, one that must preserve its order is refused, and SortShuffleWriterExecNode gains a field.
#2723 adds partition_id_routing to the same message as field 10, and two fields cannot share a number. 11 leaves 10 to it, and putting the field next to output_partitioning keeps it clear of the fields that PR appends at the end of the message.
…partitions When a stage is split between tasks, the scheduler restricts each task's plan to the partitions that task runs. It could only restrict the leaves it knew and the two PartitionSliceable operators on a hand-kept list. Any other node holding data per partition was left as it was, so a task that ran partitions 4 and 5 of it executed its partitions 0 and 1, and the stage returned some partitions twice and others not at all, with no error. - PartitionSliceable::restrict_to_partitions takes any number of children, leaves included. It replaces slice_to_partitions, which took one. - SchedulerConfig::partition_sliceable_resolver finds the PartitionSliceable of nodes Ballista does not list. It is asked after the built-in list. - A leaf that nothing can restrict now fails the task unless the task runs all of its partitions in order, instead of running and duplicating output. The scheduler also rejects a restriction that returns the wrong number of partitions.
A hash shuffle always assigns rows to partitions with DataFusion's hash, so it cannot match data that is already partitioned by another function, such as Spark's Murmur3 for bucketed tables. Add a partition-id mode to the sort-based shuffle writer. The plan below the writer computes each row's output partition into an integer column, and the writer sends the row to that partition and drops the column before writing. Output partition i holds exactly the rows whose id is i. An id that is null, negative or out of range fails the task. - ShuffleRouting selects between hashing and partition-id routing in SortShuffleWriterExec. try_new is unchanged; try_new_with_routing takes the routing, and the executor stamps the decoded writer rather than rebuilding it field by field, so the routing is carried over. - PartitionIdRepartitionExec is the marker a client places where a hash RepartitionExec would go. The static distributed planner turns it into a stage boundary ending in the new writer. It cannot be executed. - Both the marker and the writer's routing have proto messages and codec support (a oneof on the sort writer node, wire-compatible with the old field), so plans containing them can be submitted over gRPC. - The adaptive planner rejects the marker: it may coalesce or re-plan an exchange's partitions, which would break the placement. Closes #2681.
`losing_own_shuffle_output_restores_the_cache_reader` (#2633) names `dyn ExecutionPlan` but relied on the test module's import of it, which #2613 removed along with the test helper that was its only other user. Each passed on its own base, but together `ballista-scheduler`'s lib tests no longer compile on main.
# Conflicts: # ballista/scheduler/src/state/aqe/planner.rs
# Conflicts: # ballista/scheduler/src/state/execution_graph.rs
# Conflicts: # ballista/scheduler/src/state/task_builder.rs
# Conflicts: # ballista/core/proto/ballista.proto # ballista/core/src/execution_plans/shuffle_writer.rs # ballista/core/src/execution_plans/sort_shuffle/writer.rs # ballista/core/src/serde/generated/ballista.rs # ballista/core/src/serde/mod.rs # ballista/scheduler/src/planner.rs
…e shuffle PRs Pass the new resolver argument to restrict_plan_to_partitions in range_shuffle_tests and drop the now-unused internal_err import from task_builder.
PinSubmittedShufflesRule pinned only hash repartitions, so EnforceDistribution removed a client's range repartition and nothing put it back: a consumer that sorts each partition then no longer sorted the whole result. It now pins a range repartition too, unless it must preserve the order of its input, which a shuffle cannot.
# Conflicts: # ballista/scheduler/src/state/mod.rs # ballista/scheduler/src/state/task_builder.rs
Cherry-picked from #2743 (adaptive-rules-hook): SchedulerConfig::adaptive_rules takes preparation rules, which run once while a job's plan is created, and replan rules, which run last each time the plan is optimized. ExchangeExec is exported, with accessors and with_inactive_stage, from state::adaptive. On this branch the rules are also passed to the physical-plan path from #2682, where preparation rules run after the submitted shuffles are pinned, so a rule can hold back a stage the submitter placed. Comet Native Executors use it to measure a join's build side before shuffling its probe side (apache/datafusion-comet#7391).
github-actions
Bot
added
the
documentation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Which issue does this PR close?
Not intended to close an issue. This is a draft integration branch and is not meant to be merged as is.
Rationale for this change
I am using this branch (
comet-native) to test Ballista with Comet native executors. It collects the open PRs that Comet native execution depends on, so they can be exercised together before each one lands onmain.What changes are included in this PR?
The branch combines the following open PRs, which are reviewed separately:
ExecutionPlannodes restrict themselves to a task's partitionsRepartitionExecas a shuffle stage instead of dropping itIt also carries a small change on top of #2722:
PinSubmittedShufflesRulekeeps a submitted range repartition as a stage boundary.Already merged to
mainand kept in sync by mergingmaininto this branch: #2682 (run submitted physical plans through the adaptive planner) and #2743 (embedder rules for adaptive query planning).Are there any user-facing changes?
Whatever the PRs above change. This PR adds nothing beyond them.