andygrove opened a new pull request, #2721: URL: https://github.com/apache/datafusion-ballista/pull/2721
# Which issue does this PR close? Closes #2679. # Rationale for this change When the scheduler splits a stage between tasks, `restrict_plan_to_partitions` rewrites the plan so that each task sees only its own partitions. It could restrict only the leaves `select_output_partitions` knows and the two operators on the hand-kept `as_partition_sliceable` list. Any other node that holds 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. A downstream crate could implement `PartitionSliceable` but had no way to make the scheduler find it, and the trait covered only operators with a single child. Comet Native Executors (apache/datafusion-comet#7288) need this. They work around it today by giving every fragment a fake `MemorySourceConfig` leaf with one row per partition. # What changes are included in this PR? The three things the issue asks for: 1. **`PartitionSliceable` covers any node.** `restrict_to_partitions(children, partitions)` replaces `slice_to_partitions(child, partitions)`. It receives all of the node's children, already restricted, and is used for leaves too. The children are restricted the way an interior node's always are, so an input declared `SinglePartition` arrives whole. `RangeFilterExec` and `PrefixMergeExec` are ported. The walker also checks that the restricted node has `partitions.len()` partitions, because a node that gets that wrong would bring the silent duplication back. 2. **A resolver on `SchedulerConfig`.** `partition_sliceable_resolver` (builder: `with_partition_sliceable_resolver`) maps a plan node to its `PartitionSliceable`, and is asked after the built-in list. `TaskManager` passes it to the rewriter. `DefaultDistributedPlanner` passes it to the check that decides whether a union branch can be restricted to nothing, so a custom node in a union does not always get a stage of its own. 3. **A leaf nothing can restrict fails instead of duplicating output.** It is accepted only in a task that runs all of its partitions, in order, such as the single partition of a one-partition leaf, so plans made of leaves the scheduler knows are unaffected. Anything else is an error that names the node and what to implement. So that `None` from `select_output_partitions` means "unrecognised", `EmptyExec` and `PlaceholderRowExec` are now returned as recognised, and the constructor errors of the readers and the memory source propagate instead of turning into `None`. Smaller changes: `TaskManager::new` now delegates to `with_launcher`, so the config is read in one place. There are entries in the 56.0.0 upgrade guide and a "Custom Physical Plan Nodes" section in the extending-components guide. **The first commit is not part of the feature.** `main`'s Rust CI has been red since #2613: #2633 added a test to `execution_graph.rs` that names `ExecutionPlan`, and #2613 then dropped that import while moving `FailingPlanRewriteExec` out of the module, so the scheduler's test target no longer compiles. This PR cannot pass CI without the one-line fix in `fix(scheduler): import ExecutionPlan in the execution_graph tests`. It can be dropped from here once `main` is fixed, or merged on its own. # Are there any user-facing changes? Yes, and some are breaking for anyone who embeds the scheduler or defines `ExecutionPlan` nodes. All are in `docs/source/upgrading/56.0.0.md`. - `PartitionSliceable::slice_to_partitions` is now `restrict_to_partitions` and takes a `Vec` of children. - `task_builder::restrict_plan_to_partitions` takes a third argument, `Option<&PartitionSliceableResolver>`. - `SchedulerConfig` has a new public field, `partition_sliceable_resolver`. - A job whose stage is split between tasks now fails when the stage holds a leaf the scheduler cannot restrict. It used to return wrong results. A retry of part of a stage counts as split. A `DataSourceExec` over a source that is neither a file scan nor an in-memory source used to log a warning and run, and is now refused when its stage is split. One limitation: under push-staged scheduling the job's status says only "task could not be prepared or was rejected by executor", and the reason is in the scheduler's log. Pull-staged scheduling puts it in the status. Every task-preparation failure behaves this way today, so it is left alone here. ## Testing New unit tests cover the walker (a leaf found through the resolver, inputs restricted first, a `SinglePartition` input left whole, a node below a collapse, the built-in list answering before the resolver, a union branch owned by another task, a wrong partition count, the `EmptyExec` and `PlaceholderRowExec` exemption), the planner's union check, and the config builder. Two `launch_tasks` tests submit a stage that is split between two tasks and holds a custom leaf. Without a resolver the job fails and no task is sent. With one, each task is sent only the payloads of its own partitions, which together are every payload once. Each piece was also broken on purpose, and each break failed the test meant to catch it. `cargo fmt`, `cargo clippy --all-targets --workspace --all-features -- -D warnings`, `cargo doc` with `-D warnings` and prettier are clean, and the `ballista-scheduler` tests pass. The `ballista-core` tests that bind sockets were not run locally and are left to CI. -- 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]
