From 4c9545d0a2e9ef8fe9d04deba7b2bba5bae2d93e Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Sat, 31 Oct 2026 08:52:13 -0600 Subject: [PATCH 1/4] fix(scheduler): import ExecutionPlan in the execution_graph tests #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. --- ballista/scheduler/src/state/execution_graph.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ballista/scheduler/src/state/execution_graph.rs b/ballista/scheduler/src/state/execution_graph.rs index e6c7e5b65d..a2a5f3e2df 100644 --- a/ballista/scheduler/src/state/execution_graph.rs +++ b/ballista/scheduler/src/state/execution_graph.rs @@ -1937,8 +1937,8 @@ mod test { use datafusion::arrow::datatypes::{DataType, Field, Schema}; use datafusion::functions_aggregate::sum::sum; use datafusion::logical_expr::col; - use datafusion::physical_plan::Partitioning; use datafusion::physical_plan::union::UnionExec; + use datafusion::physical_plan::{ExecutionPlan, Partitioning}; use datafusion::prelude::{SessionConfig, SessionContext}; use datafusion::test_util::scan_empty_with_partitions; From 03c13f185f52999fd38f643739ac3d2b8ce121b9 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Sat, 31 Oct 2026 08:42:55 -0600 Subject: [PATCH 2/4] feat: let custom ExecutionPlan nodes restrict themselves to a task's 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. --- .../core/src/execution_plans/plan_algebra.rs | 72 +++- .../core/src/execution_plans/prefix_merge.rs | 7 +- .../core/src/execution_plans/range_filter.rs | 7 +- ballista/scheduler/src/config.rs | 19 + ballista/scheduler/src/planner.rs | 6 +- ballista/scheduler/src/state/mod.rs | 170 +++++++- ballista/scheduler/src/state/task_builder.rs | 407 ++++++++++++++---- ballista/scheduler/src/state/task_manager.rs | 24 +- ballista/scheduler/src/test_utils.rs | 237 ++++++++++ docs/source/upgrading/56.0.0.md | 40 ++ .../source/user-guide/extending-components.md | 40 +- 11 files changed, 917 insertions(+), 112 deletions(-) diff --git a/ballista/core/src/execution_plans/plan_algebra.rs b/ballista/core/src/execution_plans/plan_algebra.rs index 2b0a10711e..92a17fa918 100644 --- a/ballista/core/src/execution_plans/plan_algebra.rs +++ b/ballista/core/src/execution_plans/plan_algebra.rs @@ -74,15 +74,21 @@ pub fn preserves_partitioning(plan: &dyn ExecutionPlan) -> bool { || plan.downcast_ref::().is_some() } -/// Operators carrying data indexed by *global* input partition, which the -/// scheduler must slice when it restricts a stage plan to one task's -/// partition subset. +/// Operators carrying data indexed by *global* partition, which the scheduler +/// must slice when it restricts a stage plan to one task's partition subset. +/// A leaf that owns a partition-indexed source qualifies too. /// /// The executor still never learns its task's global identity — slicing is /// how that stays true. After restriction an operator holds a vec covering /// exactly the partitions its task will run, in local order, so `execute(k)` /// indexes it directly. /// +/// Partition `i` of the node must read partition `i` of each child that is +/// restricted, so a node that maps partitions to children differently, as +/// `UnionExec` does, cannot use this hook. The restricted node is serialized +/// with the scheduler's physical extension codec, so what it carries per +/// partition must survive that round trip. +/// /// Implement this next to the fields being sliced. The alternative — the /// scheduler's task builder knowing each operator's internals — means adding /// a partition-indexed field silently leaves a stale slicer two crates away, @@ -91,25 +97,34 @@ pub fn preserves_partitioning(plan: &dyn ExecutionPlan) -> bool { /// Like the whitelists above, this exists because `ExecutionPlan` has no /// "restrict yourself to a subset of input partitions" method. A defaulted /// upstream one would let each operator answer for itself and retire -/// [`as_partition_sliceable`]. +/// [`as_partition_sliceable`] and [`PartitionSliceableResolver`]. pub trait PartitionSliceable: ExecutionPlan { - /// Rebuild over `child` — already restricted to `partitions` — carrying - /// only the entries for `partitions`, in that order. + /// Rebuild over `children`, carrying only the entries for `partitions`, in + /// that order. + /// + /// The result must have exactly `partitions.len()` output partitions; the + /// scheduler rejects any other count. `partitions` is empty for a branch of + /// a `UnionExec` that belongs to other tasks: return a node with no + /// partitions. /// /// # Arguments /// - /// * `child` - this operator's single input, already restricted - /// * `partitions` - global input partition indices this task will run - fn slice_to_partitions( + /// * `children` - this operator's inputs, restricted to `partitions` except + /// any declared + /// [`SinglePartition`](datafusion::physical_expr::Distribution::SinglePartition), + /// which arrive whole; empty for a leaf + /// * `partitions` - global partition indices this task will run + fn restrict_to_partitions( &self, - child: Arc, + children: Vec>, partitions: &[usize], ) -> Result>; } /// Hand-maintained whitelist of [`PartitionSliceable`] operators, for the /// same reason as the property whitelists above: `dyn ExecutionPlan` can't -/// be downcast to an arbitrary trait. +/// be downcast to an arbitrary trait. Operators defined outside this crate are +/// found through a [`PartitionSliceableResolver`]. pub fn as_partition_sliceable( plan: &Arc, ) -> Option<&dyn PartitionSliceable> { @@ -122,6 +137,41 @@ pub fn as_partition_sliceable( None } +/// Looks up the [`PartitionSliceable`] of a node that [`as_partition_sliceable`] +/// does not list, so the scheduler can restrict `ExecutionPlan` nodes of your +/// own. Set it with `SchedulerConfig::with_partition_sliceable_resolver`. +/// +/// It is asked after the built-in list, and never about a node that is read +/// whole: one below a collapse, or declared `SinglePartition` by its parent. It +/// is called for most nodes of every task's plan as tasks are launched, so keep +/// it cheap, typically a `downcast_ref` to your own types. +pub type PartitionSliceableResolver = Arc< + dyn for<'a> Fn(&'a Arc) -> Option<&'a dyn PartitionSliceable> + + Send + + Sync, +>; + +/// Unwraps the only child that a single-input operator's +/// [`restrict_to_partitions`](PartitionSliceable::restrict_to_partitions) +/// receives. +/// +/// # Arguments +/// +/// * `owner` - operator name, for the error message +/// * `children` - the children `restrict_to_partitions` received +pub fn single_child( + owner: &str, + children: Vec>, +) -> Result> { + let [child] = <[_; 1]>::try_from(children).map_err(|children| { + DataFusionError::Internal(format!( + "{owner}: expected 1 child to restrict, got {}", + children.len() + )) + })?; + Ok(child) +} + /// Take `values[global]` for each global partition index, in task-local /// order. Errors rather than dropping, so a slice that outruns its operator's /// data surfaces as a named failure instead of a silently short vec. diff --git a/ballista/core/src/execution_plans/prefix_merge.rs b/ballista/core/src/execution_plans/prefix_merge.rs index b17c1a1f84..a0127eade8 100644 --- a/ballista/core/src/execution_plans/prefix_merge.rs +++ b/ballista/core/src/execution_plans/prefix_merge.rs @@ -98,7 +98,7 @@ use std::pin::Pin; use std::task::{Context, Poll}; use crate::execution_plans::plan_algebra::{ - PartitionSliceable, slice_by_global_partition, + PartitionSliceable, single_child, slice_by_global_partition, }; /// A single already-prefix-merged window-aggregate state. @@ -397,11 +397,12 @@ impl PartitionSliceable for PrefixMergeExec { /// Both the per-partition state and every [`WindowApply::Scalar`]'s /// offsets are indexed by global input partition, so both slice parallel /// to the input. - fn slice_to_partitions( + fn restrict_to_partitions( &self, - child: Arc, + children: Vec>, partitions: &[usize], ) -> Result> { + let child = single_child("PrefixMergeExec", children)?; let applies = self .applies .iter() diff --git a/ballista/core/src/execution_plans/range_filter.rs b/ballista/core/src/execution_plans/range_filter.rs index 0062dd2cfe..d0252be250 100644 --- a/ballista/core/src/execution_plans/range_filter.rs +++ b/ballista/core/src/execution_plans/range_filter.rs @@ -82,7 +82,7 @@ use futures::{Stream, StreamExt, ready}; use parking_lot::Mutex; use crate::execution_plans::plan_algebra::{ - PartitionSliceable, slice_by_global_partition, + PartitionSliceable, single_child, slice_by_global_partition, }; /// Half-open `[lo, hi)` bound for one input partition. `None` on either side @@ -336,11 +336,12 @@ impl PartitionSliceable for RangeFilterExec { /// `raw_bounds` is indexed by input partition, so it slices parallel to /// the input. Halos and routing carry over verbatim — the fresh operator /// re-widens the unwidened bounds itself. - fn slice_to_partitions( + fn restrict_to_partitions( &self, - child: Arc, + children: Vec>, partitions: &[usize], ) -> Result> { + let child = single_child("RangeFilterExec", children)?; let raw_bounds = self.raw_bounds().ok_or_else(|| { datafusion::common::DataFusionError::Internal( "RangeFilterExec: task-restriction before resolve_bounds()".into(), diff --git a/ballista/scheduler/src/config.rs b/ballista/scheduler/src/config.rs index c86152db68..0ad973af0b 100644 --- a/ballista/scheduler/src/config.rs +++ b/ballista/scheduler/src/config.rs @@ -28,6 +28,7 @@ use crate::SessionBuilder; use crate::cluster::DistributionPolicy; use crate::scheduler_server::JobIdGenerator; +use ballista_core::execution_plans::plan_algebra::PartitionSliceableResolver; use ballista_core::extension::EndpointOverrideFn; use ballista_core::{ ConfigProducer, JobId, config::TaskSchedulingPolicy, ids::new_instance_id, @@ -439,6 +440,13 @@ pub struct SchedulerConfig { pub override_physical_codec: Option>, /// Override function for customizing gRPC client endpoints before they are used pub override_create_grpc_client_endpoint: Option, + /// Finds the [`PartitionSliceable`](ballista_core::execution_plans::plan_algebra::PartitionSliceable) + /// of `ExecutionPlan` nodes of your own that hold data per partition, so + /// the scheduler can restrict them to a task's partitions. Without it, a + /// stage split between tasks that holds such a leaf fails the job. Unlike + /// the `override_*` fields it adds to Ballista's own list rather than + /// replacing it. See [`PartitionSliceableResolver`]. + pub partition_sliceable_resolver: Option, /// Whether to use TLS when connecting to executors (for flight proxy) pub use_tls: bool, /// Minimum number of registered executors before /readyz returns 200. It @@ -523,6 +531,7 @@ impl Default for SchedulerConfig { override_logical_codec: None, override_physical_codec: None, override_create_grpc_client_endpoint: None, + partition_sliceable_resolver: None, use_tls: false, min_ready_executors: 1, task_max_failures: 4, @@ -835,6 +844,15 @@ impl SchedulerConfig { self } + /// Sets [`Self::partition_sliceable_resolver`]. + pub fn with_partition_sliceable_resolver( + mut self, + resolver: PartitionSliceableResolver, + ) -> Self { + self.partition_sliceable_resolver = Some(resolver); + self + } + /// Sets the directory to write per-job event logs to. pub fn with_event_log_dir(mut self, event_log_dir: Option) -> Self { self.event_log_dir = event_log_dir; @@ -980,6 +998,7 @@ impl TryFrom for SchedulerConfig { override_physical_codec: None, override_session_builder: None, override_create_grpc_client_endpoint: None, + partition_sliceable_resolver: None, use_tls: false, min_ready_executors: opt.min_ready_executors, task_max_failures: opt.task_max_failures, diff --git a/ballista/scheduler/src/planner.rs b/ballista/scheduler/src/planner.rs index ddc67a6c24..e92000f4f9 100644 --- a/ballista/scheduler/src/planner.rs +++ b/ballista/scheduler/src/planner.rs @@ -671,9 +671,11 @@ fn can_stay_inline(plan: &Arc) -> bool { /// Whether per-task restriction can reduce `plan` to zero partitions. An empty /// slice means "this task polls nothing here", so a branch that still reports a -/// partition afterwards is one that cannot be restricted away. +/// partition afterwards, or that the rewriter refuses, cannot be restricted +/// away. No resolver is passed, so a node Ballista does not know gets a stage +/// of its own. fn restricts_to_nothing(plan: &Arc) -> bool { - restrict_plan_to_partitions(plan.clone(), &[]) + restrict_plan_to_partitions(plan.clone(), &[], None) .is_ok_and(|p| p.properties().output_partitioning().partition_count() == 0) } diff --git a/ballista/scheduler/src/state/mod.rs b/ballista/scheduler/src/state/mod.rs index 703d1ad9e4..8df4c9add8 100644 --- a/ballista/scheduler/src/state/mod.rs +++ b/ballista/scheduler/src/state/mod.rs @@ -513,14 +513,17 @@ mod tests { use crate::scheduler_server::timestamp_millis; use crate::state::executor_manager::ExecutorManager; use crate::state::task_manager::TaskLauncher; - use crate::test_utils::test_cluster_context; + use crate::test_utils::payload::{ + PayloadCodec, PayloadExec, payload_resolver, payloads_of, + }; + use crate::test_utils::{mock_executor, test_cluster_context}; use ballista_core::extension::SessionConfigExt; - use ballista_core::serde::BallistaCodec; use ballista_core::serde::protobuf::MultiTaskDefinition; use ballista_core::serde::scheduler::{ ExecutorData, ExecutorMetadata, ExecutorOperatingSystemSpecification, ExecutorSpecification, }; + use ballista_core::serde::{BallistaCodec, BallistaLogicalExtensionCodec}; use datafusion::arrow::datatypes::{DataType, Field, Schema}; use datafusion::execution::context::SessionConfig; use datafusion::functions_aggregate::sum::sum; @@ -847,4 +850,167 @@ mod tests { Ok(()) } + + /// Keeps the task definitions it is handed, and rejects none. + #[derive(Default)] + struct CapturingLauncher { + launched: std::sync::Mutex>, + } + + #[async_trait::async_trait] + impl TaskLauncher for CapturingLauncher { + async fn launch_tasks( + &self, + _executor: &ExecutorMetadata, + tasks: Vec, + _executor_manager: &ExecutorManager, + ) -> Result> { + self.launched.lock().unwrap().extend(tasks); + Ok(HashSet::new()) + } + } + + /// Submits a job of one four-partition `PayloadExec` and binds it to two + /// executors of two vcores each, which splits its stage into two tasks. + async fn bind_split_payload_stage( + config: SchedulerConfig, + launcher: Arc, + ) -> Result<( + SchedulerState, + JobId, + Vec, + )> { + let job_id = JobId::from("job-payload"); + let codec = BallistaCodec::new( + Arc::new(BallistaLogicalExtensionCodec::default()), + Arc::new(PayloadCodec::default()), + ); + let state: SchedulerState = + SchedulerState::new_with_task_launcher( + test_cluster_context(), + codec, + "localhost:50050".to_owned(), + Arc::new(config), + launcher, + ); + + for executor in ["executor-1", "executor-2"] { + let metadata = ExecutorMetadata { + specification: ExecutorSpecification::default().with_vcores(2), + ..mock_executor(executor.to_string()) + }; + let data = ExecutorData { + executor_id: executor.to_string(), + total_vcores: 2, + available_vcores: 2, + }; + state + .executor_manager + .register_executor(metadata, data) + .await?; + } + + let ctx = state + .session_manager + .create_or_update_session("session", &SessionConfig::new_with_ballista()) + .await?; + state + .task_manager + .queue_job(&job_id, "", timestamp_millis())?; + state + .task_manager + .submit_physical_plan( + &job_id, + "", + ctx, + PayloadExec::leaf(&["p0", "p1", "p2", "p3"]), + timestamp_millis(), + None, + ) + .await?; + + let bound = state + .executor_manager + .bind_schedulable_tasks(state.task_manager.get_running_job_cache()) + .await?; + let mut slices: Vec<_> = bound + .iter() + .map(|(_, task)| task.global_input_partition_ids.clone()) + .collect(); + slices.sort(); + assert_eq!(slices, [[0, 1], [2, 3]], "the stage is split in two"); + + // As `revive_offers` does between binding and launching. + state.task_manager.record_task_reservations(&bound); + + Ok((state, job_id, bound)) + } + + /// Launches `bound`, returning the vcores refunded and the jobs failed. + async fn launch( + state: &SchedulerState, + bound: Vec, + ) -> Result<(u32, HashSet)> { + let (tx_event, _rx_event) = tokio::sync::mpsc::channel(100); + let sender = EventSender::new(tx_event); + let (unassigned_slots, failed_jobs) = state.launch_tasks(bound, &sender).await?; + Ok((unassigned_slots.iter().map(|(_, n)| *n).sum(), failed_jobs)) + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn launch_tasks_fails_a_job_whose_split_stage_holds_a_leaf_nobody_can_restrict() + -> Result<()> { + let launcher = Arc::new(CapturingLauncher::default()); + let (state, job_id, bound) = + bind_split_payload_stage(SchedulerConfig::default(), launcher.clone()) + .await?; + let vcores: u32 = bound.iter().map(|(_, task)| task.vcores_consumed).sum(); + + let (freed, failed_jobs) = launch(&state, bound).await?; + + assert_eq!(failed_jobs, HashSet::from([job_id])); + assert_eq!(freed, vcores, "the tasks' vcores are refunded"); + assert!(launcher.launched.lock().unwrap().is_empty()); + + Ok(()) + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn launch_tasks_restricts_the_leaf_through_the_configured_resolver() + -> Result<()> { + let launcher = Arc::new(CapturingLauncher::default()); + let config = SchedulerConfig::default() + .with_partition_sliceable_resolver(payload_resolver()); + let (state, _job_id, bound) = + bind_split_payload_stage(config, launcher.clone()).await?; + + let (freed, failed_jobs) = launch(&state, bound).await?; + + assert!(failed_jobs.is_empty()); + assert_eq!(freed, 0, "every task was launched, so no slot is refunded"); + + // Decode what the executors would, with the codec the scheduler used. + let task_ctx = SessionContext::new().task_ctx(); + let mut shipped: Vec> = launcher + .launched + .lock() + .unwrap() + .iter() + .map(|definition| { + let plan = PhysicalPlanNode::try_decode(&definition.plan) + .and_then(|node| { + node.try_into_physical_plan( + &task_ctx, + state.codec.physical_extension_codec(), + ) + }) + .expect("the executor can decode the plan"); + payloads_of(plan.children()[0]) + }) + .collect(); + shipped.sort(); + assert_eq!(shipped, [["p0", "p1"], ["p2", "p3"]]); + + Ok(()) + } } diff --git a/ballista/scheduler/src/state/task_builder.rs b/ballista/scheduler/src/state/task_builder.rs index 5e0a21b3be..c71365201a 100644 --- a/ballista/scheduler/src/state/task_builder.rs +++ b/ballista/scheduler/src/state/task_builder.rs @@ -36,9 +36,11 @@ //! flows from parent to descendants via function arguments, so sibling //! subtrees never share state and there's no traversal-order dependency. -use ballista_core::execution_plans::plan_algebra::as_partition_sliceable; +use ballista_core::execution_plans::plan_algebra::{ + PartitionSliceableResolver, as_partition_sliceable, +}; use ballista_core::execution_plans::{RangeShuffleReaderExec, ShuffleReaderExec}; -use datafusion::common::internal_err; +use datafusion::common::{internal_err, plan_err}; use datafusion::datasource::memory::MemorySourceConfig; use datafusion::datasource::physical_plan::{ FileGroup, FileScanConfig, FileScanConfigBuilder, @@ -47,13 +49,13 @@ use datafusion::datasource::source::DataSourceExec; use datafusion::error::Result; use datafusion::physical_expr::Distribution; use datafusion::physical_plan::coalesce_partitions::CoalescePartitionsExec; +use datafusion::physical_plan::displayable; use datafusion::physical_plan::empty::EmptyExec; use datafusion::physical_plan::placeholder_row::PlaceholderRowExec; use datafusion::physical_plan::replace_children_if_necessary; use datafusion::physical_plan::sorts::sort_preserving_merge::SortPreservingMergeExec; use datafusion::physical_plan::union::UnionExec; use datafusion::physical_plan::{ExecutionPlan, Partitioning}; -use log::warn; use std::any::Any; use std::sync::Arc; @@ -62,11 +64,16 @@ use std::sync::Arc; /// `SortPreservingMergeExec`, or a `SinglePartition`-requiring join build /// side) — those leaves keep the full upstream so the collapse sees every /// partition. +/// +/// Nodes Ballista does not list are restricted through `resolver`. A leaf that +/// nothing can restrict is an error unless `partitions` is all of its +/// partitions, in order. pub fn restrict_plan_to_partitions( plan: Arc, partitions: &[usize], + resolver: Option<&PartitionSliceableResolver>, ) -> Result> { - restrict(plan, partitions, /* under_collect */ false) + restrict(plan, partitions, /* under_collect */ false, resolver) } /// Recursive worker. `under_collect` is the scope inherited from ancestors: @@ -76,6 +83,7 @@ fn restrict( plan: Arc, partitions: &[usize], under_collect: bool, + resolver: Option<&PartitionSliceableResolver>, ) -> Result> { // Leaves: apply the algebraic partition-projection when we're not // under a collapse. Under collapse, leaves keep their full upstream — @@ -86,20 +94,16 @@ fn restrict( return Ok(rewritten); } - // Operators carrying data indexed by global input partition slice it - // parallel to the input restriction. Each one implements the slicing - // beside its own fields; this walker only supplies the restricted child. - if !under_collect && let Some(sliceable) = as_partition_sliceable(&plan) { - let children = plan.children(); - let [child] = children.as_slice() else { - return internal_err!( - "{} is PartitionSliceable but has {} children, expected 1", - plan.name(), - children.len() - ); - }; - let new_child = restrict((*child).clone(), partitions, false)?; - return sliceable.slice_to_partitions(new_child, partitions); + // Operators carrying data indexed by global partition slice it parallel to + // the restriction. Each one implements the slicing beside its own fields; + // this walker only supplies the restricted children. + if !under_collect + && let Some(sliceable) = as_partition_sliceable(&plan) + .or_else(|| resolver.and_then(|resolve| resolve(&plan))) + { + let children = restrict_children(&plan, partitions, false, resolver)?; + let restricted = sliceable.restrict_to_partitions(children, partitions)?; + return checked_restriction(plan.name(), restricted, partitions); } // UnionExec: parent partition `p` maps to exactly one child's local @@ -116,21 +120,83 @@ fn restrict( .into_iter() .cloned() .zip(per_child) - .map(|(child, sub)| restrict(child, &sub, false)) + .map(|(child, sub)| restrict(child, &sub, false, resolver)) .collect::>>()?; return replace_children_if_necessary(plan, new_children); } + // A leaf nothing above could restrict: read whole below a collapse, and + // otherwise only right if the task runs all of it. + if plan.children().is_empty() { + if !under_collect { + ensure_leaf_runs_whole(&plan, partitions)?; + } + return Ok(plan); + } + // Interior: compute per-child scope, recurse, rebuild the node. - let per_child = child_scopes(&plan, under_collect); - let new_children: Vec> = plan - .children() + let new_children = restrict_children(&plan, partitions, under_collect, resolver)?; + replace_children_if_necessary(plan, new_children) +} + +/// Restrict each child of `plan`, in the scope `child_scopes` gives it. +fn restrict_children( + plan: &Arc, + partitions: &[usize], + under_collect: bool, + resolver: Option<&PartitionSliceableResolver>, +) -> Result>> { + let per_child = child_scopes(plan, under_collect); + plan.children() .into_iter() .cloned() .zip(per_child) - .map(|(child, collect)| restrict(child, partitions, collect)) - .collect::>>()?; - replace_children_if_necessary(plan, new_children) + .map(|(child, collect)| restrict(child, partitions, collect, resolver)) + .collect() +} + +/// Hold a `PartitionSliceable` to its contract: one partition for each +/// partition the task runs. Another count would run the wrong partitions. +fn checked_restriction( + name: &str, + restricted: Arc, + partitions: &[usize], +) -> Result> { + let count = restricted + .properties() + .output_partitioning() + .partition_count(); + if count != partitions.len() { + return plan_err!( + "{name}::restrict_to_partitions returned {count} partitions for \ + partitions {partitions:?}; it has to return exactly one for each" + ); + } + Ok(restricted) +} + +/// Accept a leaf that nothing could restrict only if the task runs all of it, +/// in order. Left as it is, it numbers its partitions from zero, so any other +/// slice would run the wrong ones. A one-partition leaf therefore passes in the +/// task that owns it, and fails in a `UnionExec` branch that belongs to another. +fn ensure_leaf_runs_whole( + plan: &Arc, + partitions: &[usize], +) -> Result<()> { + let count = plan.properties().output_partitioning().partition_count(); + if partitions.iter().copied().eq(0..count) { + return Ok(()); + } + let name = plan.name(); + let described = displayable(plan.as_ref()).one_line().to_string(); + let described = described.trim_end(); + plan_err!( + "cannot restrict `{described}` to partitions {partitions:?} of its {count}: \ + the scheduler does not know how to restrict a `{name}`, and left alone it \ + would number its partitions from zero in every task. A node of your own \ + can implement `PartitionSliceable` and be returned from \ + `SchedulerConfig::partition_sliceable_resolver`." + ) } /// Bucket each parent partition of a `UnionExec` into its owning child's @@ -239,10 +305,10 @@ fn child_scopes(plan: &Arc, under_collect: bool) -> Vec /// **When a new leaf type surfaces, add a branch below.** Currently /// covered: `ShuffleReaderExec`, `DataSourceExec` (files, in-memory), /// `PlaceholderRowExec`, `EmptyExec`. `ValuesExec` and other future -/// self-generating leaves will silently fall through as `None`; if one -/// sits under a partition-splitting fan-in, the walker leaves it at its -/// full partition count and the executor duplicates output (see the -/// PlaceholderRow-under-UnionExec bug for the failure mode). +/// self-generating leaves fall through as `None`; if one sits under a +/// partition-splitting fan-in, the walker refuses the task rather than let the +/// executor duplicate output (see the PlaceholderRow-under-UnionExec bug for +/// the failure mode). /// /// # Scope /// @@ -287,14 +353,12 @@ fn select_output_partitions( ); } }; - let Ok(restricted) = ShuffleReaderExec::try_new( + let restricted = ShuffleReaderExec::try_new( reader.stage_id, kept, reader.schema(), partitioning, - ) else { - return Ok(None); - }; + )?; return Ok(Some(Arc::new(restricted))); } @@ -306,14 +370,12 @@ fn select_output_partitions( .iter() .filter_map(|&p| reader.partition.get(p).cloned()) .collect(); - let Ok(restricted) = RangeShuffleReaderExec::try_new( + let restricted = RangeShuffleReaderExec::try_new( reader.stage_id, kept, reader.schema(), reader.merge_ordering().clone(), - ) else { - return Ok(None); - }; + )?; let restricted = restricted.with_fetch_limit(reader.fetch()); // Bounds are per output partition, so they follow the same slice the // locations did. Dropping them here would leave the task reading its @@ -324,10 +386,7 @@ fn select_output_partitions( .iter() .filter_map(|&p| bounds.get(p).cloned()) .collect(); - let Ok(restricted) = restricted.with_bounds(kept) else { - return Ok(None); - }; - restricted + restricted.with_bounds(kept)? } None => restricted, }; @@ -352,20 +411,15 @@ fn select_output_partitions( .iter() .filter_map(|&i| config.partitions().get(i).cloned()) .collect(); - let Ok(restricted) = MemorySourceConfig::try_new( + let restricted = MemorySourceConfig::try_new( &kept, config.original_schema(), config.projection().clone(), - ) else { - return Ok(None); - }; + )?; return Ok(Some(DataSourceExec::from_data_source(restricted))); } - warn!( - "select_output_partitions: unrecognised DataSourceExec source \ - left unrestricted; if it distributes work from a shared queue, \ - tasks would over-read" - ); + // Any other source is left to the walker, which refuses a task that runs + // only some of its partitions. return Ok(None); } @@ -383,18 +437,19 @@ fn select_output_partitions( // // Workaround for the empty-slice case: swap in a 0-partition // `MemorySourceConfig` (which DOES roundtrip its partition list). - // For non-empty slices we return `None` and let the generic recursion - // leave the leaf as-is — currently safe because the only shape emitted - // is `partitions == 1` and UnionExec routing gives sub-slices of `[0]` - // (identity, no restriction needed) or `[]` (handled by this branch). + // For non-empty slices we return the leaf as it is — currently safe because + // the only shape emitted is `partitions == 1` and UnionExec routing gives + // sub-slices of `[0]` (identity, no restriction needed) or `[]` (handled by + // this branch). Returned rather than `None` so the walker's leaf check + // accepts it: an `EmptyExec` under the round-robin `RepartitionExec` that + // `make_empty_exec_serde_safe` adds sees that repartition's partition + // numbers, not its own. if plan.is::() || plan.is::() { if indices.is_empty() { - let Ok(stub) = MemorySourceConfig::try_new(&[], plan.schema(), None) else { - return Ok(None); - }; + let stub = MemorySourceConfig::try_new(&[], plan.schema(), None)?; return Ok(Some(DataSourceExec::from_data_source(stub))); } - return Ok(None); + return Ok(Some(plan.clone())); } Ok(None) @@ -405,7 +460,7 @@ mod tests { use super::*; use ballista_core::JobId; use ballista_core::execution_plans::{ - CoalescePlan, PartitionGroup, ShuffleReaderExec, + CoalescePlan, PartitionGroup, RangeBound, ShuffleReaderExec, }; use ballista_core::serde::scheduler::{ ExecutorMetadata, PartitionId, PartitionLocation, @@ -452,7 +507,7 @@ mod tests { #[test] fn keeps_only_the_wanted_partition() { let plan = shuffle_reader(4); - let pruned = restrict_plan_to_partitions(plan, &[1]).unwrap(); + let pruned = restrict_plan_to_partitions(plan, &[1], None).unwrap(); let r = pruned .downcast_ref::() .expect("expected a ShuffleReaderExec"); @@ -466,7 +521,7 @@ mod tests { #[test] fn keeps_multiple_wanted_partitions_in_request_order() { let plan = shuffle_reader(4); - let pruned = restrict_plan_to_partitions(plan, &[0, 3]).unwrap(); + let pruned = restrict_plan_to_partitions(plan, &[0, 3], None).unwrap(); let r = pruned .downcast_ref::() .expect("expected a ShuffleReaderExec"); @@ -487,7 +542,7 @@ mod tests { ) .unwrap(); let plan: Arc = Arc::new(reader); - let pruned = restrict_plan_to_partitions(plan, &[1]).unwrap(); + let pruned = restrict_plan_to_partitions(plan, &[1], None).unwrap(); let r = pruned .downcast_ref::() .expect("expected a ShuffleReaderExec"); @@ -529,7 +584,7 @@ mod tests { ) .unwrap(); let plan: Arc = Arc::new(reader); - let pruned = restrict_plan_to_partitions(plan, &[1]).unwrap(); + let pruned = restrict_plan_to_partitions(plan, &[1], None).unwrap(); let r = pruned .downcast_ref::() .expect("expected a ShuffleReaderExec"); @@ -541,7 +596,7 @@ mod tests { #[test] fn restricting_to_every_partition_is_a_no_op() { let plan = shuffle_reader(3); - let pruned = restrict_plan_to_partitions(plan, &[0, 1, 2]).unwrap(); + let pruned = restrict_plan_to_partitions(plan, &[0, 1, 2], None).unwrap(); let r = pruned .downcast_ref::() .expect("expected a ShuffleReaderExec"); @@ -565,13 +620,17 @@ mod tests { use crate::test_utils::scan_with_file_groups; use datafusion::physical_plan::union::UnionExec; + /// The file groups of a `DataSourceExec` scan. + fn file_groups(plan: &Arc) -> &[FileGroup] { + let exec = plan.downcast_ref::().unwrap(); + let source: &dyn Any = exec.data_source().as_ref(); + &source.downcast_ref::().unwrap().file_groups + } + /// File counts per file group of a `DataSourceExec` — the shape a union /// test cares about after restriction. fn group_file_counts(plan: &Arc) -> Vec { - let exec = plan.downcast_ref::().unwrap(); - let source: &dyn Any = exec.data_source().as_ref(); - let config = source.downcast_ref::().unwrap(); - config.file_groups.iter().map(|g| g.len()).collect() + file_groups(plan).iter().map(|g| g.len()).collect() } fn union_child_file_counts(plan: &Arc) -> Vec> { @@ -594,7 +653,7 @@ mod tests { ); // Parent partition 1 is served by the left child's local partition 1. - let restricted = restrict_plan_to_partitions(union.clone(), &[1]).unwrap(); + let restricted = restrict_plan_to_partitions(union.clone(), &[1], None).unwrap(); assert_eq!( union_child_file_counts(&restricted), vec![vec![1], vec![]], @@ -605,7 +664,7 @@ mod tests { // Parent partition 6 is served by the right child's local partition 2 // (6 - 4). Before the mapping existed this would have addressed no // group at all in either child. - let restricted = restrict_plan_to_partitions(union, &[6]).unwrap(); + let restricted = restrict_plan_to_partitions(union, &[6], None).unwrap(); assert_eq!( union_child_file_counts(&restricted), vec![vec![], vec![1]], @@ -623,7 +682,8 @@ mod tests { scan_with_file_groups(3), ]) .unwrap(); - let restricted = restrict_plan_to_partitions(union, &[partition]).unwrap(); + let restricted = + restrict_plan_to_partitions(union, &[partition], None).unwrap(); let counts = union_child_file_counts(&restricted); let (serving, idle) = if partition < 3 { (0, 1) } else { (1, 0) }; assert_eq!( @@ -645,7 +705,7 @@ mod tests { #[test] fn scan_without_union_is_pinned_to_its_own_partition() { let plan = scan_with_file_groups(4); - let restricted = restrict_plan_to_partitions(plan, &[2]).unwrap(); + let restricted = restrict_plan_to_partitions(plan, &[2], None).unwrap(); assert_eq!(group_file_counts(&restricted), vec![1]); } @@ -675,7 +735,7 @@ mod tests { "CrossJoinExec must declare SinglePartition on its left input" ); - let restricted = restrict_plan_to_partitions(join, &[1]).unwrap(); + let restricted = restrict_plan_to_partitions(join, &[1], None).unwrap(); let children = restricted.children(); assert_eq!( group_file_counts(children[0]), @@ -689,15 +749,14 @@ mod tests { ); } - /// A `RangeFilterExec` restricted to a subset of partitions must slice - /// its `raw_bounds` by the same indices in the same order as its child. - /// Halos + routing_expr carry over unchanged. - #[test] - fn range_filter_bounds_are_sliced_with_partitions() { + /// A `RangeFilterExec` over a four-partition `ShuffleReaderExec`, with the + /// raw bounds it was built from. Partition `i` of the stage owns bound `i`. + fn range_filter_over_reader() -> (Arc, Vec) { use ballista_core::execution_plans::RangeFilterExec; use datafusion::physical_expr::Partitioning; use datafusion::physical_expr::PhysicalExpr; use datafusion::physical_expr::expressions::Column; + use datafusion::scalar::ScalarValue; // 4 upstream partitions with a global 4-way range partition. let schema = @@ -711,11 +770,10 @@ mod tests { Partitioning::UnknownPartitioning(4), ) .unwrap(); - use datafusion::scalar::ScalarValue; let filter_expr: Arc = Arc::new(Column::new("v", 0)); // K=4 raw bounds derived from cuts [100, 200, 300]. let sv = |v: f64| ScalarValue::Float64(Some(v)); - let raw_bounds: Vec<(Option, Option)> = vec![ + let raw_bounds: Vec = vec![ (None, Some(sv(100.0))), (Some(sv(100.0)), Some(sv(200.0))), (Some(sv(200.0)), Some(sv(300.0))), @@ -732,8 +790,19 @@ mod tests { ) .unwrap(), ); + (plan, raw_bounds) + } + + /// A `RangeFilterExec` restricted to a subset of partitions must slice + /// its `raw_bounds` by the same indices in the same order as its child. + /// Halos + routing_expr carry over unchanged. + #[test] + fn range_filter_bounds_are_sliced_with_partitions() { + use ballista_core::execution_plans::RangeFilterExec; + + let (plan, raw_bounds) = range_filter_over_reader(); - let restricted = restrict_plan_to_partitions(plan, &[1, 3]).unwrap(); + let restricted = restrict_plan_to_partitions(plan, &[1, 3], None).unwrap(); let rf = restricted .downcast_ref::() .expect("top must remain RangeFilterExec"); @@ -751,4 +820,180 @@ mod tests { assert_eq!(reader.partition[0][0].partition_id.partition_id, 1); assert_eq!(reader.partition[1][0].partition_id.partition_id, 3); } + + // --- Nodes the scheduler does not know: `PayloadExec` stands in for one --- + + use crate::test_utils::payload::{PayloadExec, payload_resolver, payloads_of}; + use ballista_core::execution_plans::plan_algebra::PartitionSliceableResolver; + use std::sync::Mutex; + + fn partition_count(plan: &Arc) -> usize { + plan.properties().output_partitioning().partition_count() + } + + /// The file of each group of a scan (`group_file_counts` only counts them). + fn group_files(plan: &Arc) -> Vec { + file_groups(plan) + .iter() + .map(|group| group.files()[0].object_meta.location.to_string()) + .collect() + } + + #[test] + fn resolver_restricts_a_node_the_scheduler_does_not_know() { + let resolver = payload_resolver(); + let leaf = PayloadExec::leaf(&["p0", "p1", "p2", "p3"]); + + let restricted = + restrict_plan_to_partitions(leaf, &[1, 3], Some(&resolver)).unwrap(); + + assert_eq!(payloads_of(&restricted), ["p1", "p3"]); + assert_eq!(partition_count(&restricted), 2); + } + + #[test] + fn a_leaf_nobody_can_restrict_is_refused_unless_the_task_runs_all_of_it() { + let leaf = PayloadExec::leaf(&["p0", "p1"]); + + let kept = restrict_plan_to_partitions(leaf.clone(), &[0, 1], None).unwrap(); + assert!( + Arc::ptr_eq(&kept, &leaf), + "an unrestricted leaf is left alone" + ); + + // A one-partition leaf is right for the task that owns it. + let single = PayloadExec::leaf(&["p0"]); + assert!(restrict_plan_to_partitions(single, &[0], None).is_ok()); + + for partitions in [&[0][..], &[1], &[1, 0], &[]] { + assert!( + restrict_plan_to_partitions(leaf.clone(), partitions, None).is_err(), + "{partitions:?} is not all of the leaf, in order" + ); + } + + let error = restrict_plan_to_partitions(leaf, &[1], None) + .unwrap_err() + .to_string(); + assert!(error.contains("PayloadExec"), "{error}"); + assert!(error.contains("[1]"), "{error}"); + assert!(error.contains("PartitionSliceable"), "{error}"); + } + + #[test] + fn the_resolver_is_asked_only_after_the_built_in_list() { + let asked = Arc::new(Mutex::new(vec![])); + let record = asked.clone(); + let resolver: PartitionSliceableResolver = Arc::new(move |plan| { + record.lock().unwrap().push(plan.name().to_string()); + None + }); + // A `RangeFilterExec` is on the `PartitionSliceable` list, the + // `ShuffleReaderExec` below it is a leaf the scheduler restricts itself. + let (plan, _) = range_filter_over_reader(); + + restrict_plan_to_partitions(plan, &[1, 3], Some(&resolver)).unwrap(); + assert!(asked.lock().unwrap().is_empty()); + + // A node it does not know is put to the resolver. + let leaf = PayloadExec::leaf(&["p0"]); + restrict_plan_to_partitions(leaf, &[0], Some(&resolver)).unwrap(); + assert_eq!(*asked.lock().unwrap(), ["PayloadExec"]); + } + + #[test] + fn inputs_are_restricted_except_one_the_node_reads_whole() { + let resolver = payload_resolver(); + let node = PayloadExec::over( + &["p0", "p1", "p2", "p3"], + vec![scan_with_file_groups(4), scan_with_file_groups(4)], + &[1], + ); + + let restricted = + restrict_plan_to_partitions(node, &[1, 3], Some(&resolver)).unwrap(); + + assert_eq!(payloads_of(&restricted), ["p1", "p3"]); + let inputs = restricted.children(); + assert_eq!(group_files(inputs[0]), ["file1.parquet", "file3.parquet"]); + assert_eq!(group_file_counts(inputs[1]), vec![1, 1, 1, 1]); + } + + #[test] + fn a_node_below_a_collapse_is_read_whole() { + let resolver = payload_resolver(); + let plan: Arc = + Arc::new(CoalescePartitionsExec::new(PayloadExec::leaf(&[ + "p0", "p1", "p2", "p3", + ]))); + + let restricted = + restrict_plan_to_partitions(plan, &[1], Some(&resolver)).unwrap(); + + assert_eq!( + payloads_of(restricted.children()[0]), + ["p0", "p1", "p2", "p3"] + ); + } + + /// A union branch that belongs to another task must restrict to nothing, + /// which needs the resolver for a node Ballista does not know. + #[test] + fn a_union_branch_owned_by_another_task_restricts_to_nothing() { + let resolver = payload_resolver(); + let union = || -> Arc { + UnionExec::try_new(vec![ + scan_with_file_groups(2), + PayloadExec::leaf(&["p0", "p1"]), + ]) + .unwrap() + }; + + // Partition 2 is the second branch's partition 0. + let restricted = + restrict_plan_to_partitions(union(), &[2], Some(&resolver)).unwrap(); + let branches = restricted.children(); + assert_eq!(partition_count(branches[0]), 0); + assert_eq!(payloads_of(branches[1]), ["p0"]); + + // And partition 0 is the first branch's, which leaves the second none. + let restricted = + restrict_plan_to_partitions(union(), &[0], Some(&resolver)).unwrap(); + let branches = restricted.children(); + assert_eq!(group_files(branches[0]), ["file0.parquet"]); + assert_eq!(partition_count(branches[1]), 0); + + assert!(restrict_plan_to_partitions(union(), &[0], None).is_err()); + assert!(restrict_plan_to_partitions(union(), &[2], None).is_err()); + } + + #[test] + fn a_restriction_that_keeps_the_wrong_partition_count_is_rejected() { + let resolver = payload_resolver(); + let leaf = PayloadExec::forgets_to_restrict(&["p0", "p1", "p2", "p3"]); + + let error = restrict_plan_to_partitions(leaf, &[1, 3], Some(&resolver)) + .unwrap_err() + .to_string(); + + assert!(error.contains("PayloadExec"), "{error}"); + assert!(error.contains("returned 4 partitions"), "{error}"); + } + + /// An `EmptyExec` sees the partition numbers of the repartition above it, + /// not its own, so these two leaves are exempt from the leaf check. + #[test] + fn self_generating_leaves_are_exempt_from_the_leaf_check() { + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, true)])); + let empty: Arc = Arc::new(EmptyExec::new(schema.clone())); + assert!(restrict_plan_to_partitions(empty, &[1], None).is_ok()); + + let placeholder: Arc = + Arc::new(PlaceholderRowExec::new(schema)); + assert!(restrict_plan_to_partitions(placeholder.clone(), &[0], None).is_ok()); + + // A task that polls none of it gets a source with no partitions. + let nothing = restrict_plan_to_partitions(placeholder, &[], None).unwrap(); + assert_eq!(partition_count(¬hing), 0); + } } diff --git a/ballista/scheduler/src/state/task_manager.rs b/ballista/scheduler/src/state/task_manager.rs index 49fc0236bc..8469c4b3c3 100644 --- a/ballista/scheduler/src/state/task_manager.rs +++ b/ballista/scheduler/src/state/task_manager.rs @@ -32,6 +32,7 @@ use ballista_core::config::ExecutionMode; use ballista_core::error::BallistaError; use ballista_core::error::Result; use ballista_core::execution_plans::compute_global_output_partition_ids; +use ballista_core::execution_plans::plan_algebra::PartitionSliceableResolver; use ballista_core::extension::{SessionConfigExt, SessionConfigHelperExt}; use ballista_core::serde::BallistaCodec; use ballista_core::serde::protobuf::{ @@ -149,6 +150,9 @@ pub struct TaskManager task_max_failures: usize, /// Maximum number of failure attempts for stage-level retry before the stage is considered failed. stage_max_failures: usize, + /// Resolver for nodes Ballista does not know. See + /// [`SchedulerConfig::partition_sliceable_resolver`]. + partition_sliceable_resolver: Option, } /// Contains the execution graph and cached data to improve performance @@ -200,19 +204,16 @@ impl TaskManager scheduler_endpoint: String, config: Arc, ) -> Self { - Self { + Self::with_launcher( state, codec, - scheduler_id: scheduler_id.clone(), - active_job_cache: Arc::new(DashMap::new()), - outstanding_task_reservations: Arc::new(DashMap::new()), - launcher: Arc::new(DefaultTaskLauncher::new(scheduler_endpoint)), - task_max_failures: config.task_max_failures, - stage_max_failures: config.stage_max_failures, - } + scheduler_id, + Arc::new(DefaultTaskLauncher::new(scheduler_endpoint)), + config, + ) } - #[allow(dead_code)] + /// Creates a new `TaskManager` that launches its tasks with `launcher`. pub(crate) fn with_launcher( state: Arc, codec: BallistaCodec, @@ -229,6 +230,7 @@ impl TaskManager launcher, task_max_failures: config.task_max_failures, stage_max_failures: config.stage_max_failures, + partition_sliceable_resolver: config.partition_sliceable_resolver.clone(), } } @@ -878,6 +880,8 @@ impl TaskManager /// ShuffleReader partitions) to this task's assigned slice, then encoded /// fresh. No shared-plan cache — each task gets its own bytes because /// each task's plan is different. + /// + /// Fails if the plan holds a leaf that cannot be restricted to the slice. #[allow(dead_code)] pub fn prepare_task_definition( &self, @@ -892,6 +896,7 @@ impl TaskManager let restricted = restrict_plan_to_partitions( task.plan.clone(), &task.global_input_partition_ids, + self.partition_sliceable_resolver.as_ref(), )?; let mut plan_buf: Vec = vec![]; let plan_proto = PhysicalPlanNode::try_from_physical_plan( @@ -1013,6 +1018,7 @@ impl TaskManager let restricted = restrict_plan_to_partitions( task.plan.clone(), &task.global_input_partition_ids, + self.partition_sliceable_resolver.as_ref(), )?; let mut plan_buf: Vec = vec![]; let plan_proto = PhysicalPlanNode::try_from_physical_plan(restricted, codec)?; diff --git a/ballista/scheduler/src/test_utils.rs b/ballista/scheduler/src/test_utils.rs index 38c77724d1..77d736d02f 100644 --- a/ballista/scheduler/src/test_utils.rs +++ b/ballista/scheduler/src/test_utils.rs @@ -1394,3 +1394,240 @@ pub fn fail_plan_rewrites(plan: &mut Arc) -> Arc, + inputs: Vec>, + whole: Vec, + forgets_to_restrict: bool, + properties: Arc, + } + + impl PayloadExec { + fn new( + payloads: Vec, + inputs: Vec>, + whole: Vec, + forgets_to_restrict: bool, + ) -> Self { + // The schema of `scan_with_file_groups`, so that the two can be + // branches of one `UnionExec`. + let schema = + Arc::new(Schema::new(vec![Field::new("a", DataType::Int64, false)])); + let properties = Arc::new(PlanProperties::new( + EquivalenceProperties::new(schema), + Partitioning::UnknownPartitioning(payloads.len()), + EmissionType::Incremental, + Boundedness::Bounded, + )); + Self { + payloads, + inputs, + whole, + forgets_to_restrict, + properties, + } + } + + fn strings(payloads: &[&str]) -> Vec { + payloads.iter().map(|p| p.to_string()).collect() + } + + /// A leaf with one partition per payload. + pub fn leaf(payloads: &[&str]) -> Arc { + Arc::new(Self::new(Self::strings(payloads), vec![], vec![], false)) + } + + /// A node over `inputs`, with one partition per payload. + pub fn over( + payloads: &[&str], + inputs: Vec>, + whole: &[usize], + ) -> Arc { + Arc::new(Self::new( + Self::strings(payloads), + inputs, + whole.to_vec(), + false, + )) + } + + /// A leaf whose restriction hands back all of its payloads, as an + /// implementation that forgets to restrict would. + pub fn forgets_to_restrict(payloads: &[&str]) -> Arc { + Arc::new(Self::new(Self::strings(payloads), vec![], vec![], true)) + } + } + + /// The payloads of a [`PayloadExec`], one per partition. + pub fn payloads_of(plan: &Arc) -> Vec { + plan.downcast_ref::() + .expect("expected a PayloadExec") + .payloads + .clone() + } + + impl DisplayAs for PayloadExec { + fn fmt_as(&self, _t: DisplayFormatType, f: &mut Formatter) -> std::fmt::Result { + write!(f, "PayloadExec: payloads={:?}", self.payloads) + } + } + + impl ExecutionPlan for PayloadExec { + fn name(&self) -> &str { + "PayloadExec" + } + + fn properties(&self) -> &Arc { + &self.properties + } + + fn children(&self) -> Vec<&Arc> { + self.inputs.iter().collect() + } + + fn required_input_distribution(&self) -> Vec { + (0..self.inputs.len()) + .map(|i| { + if self.whole.contains(&i) { + Distribution::SinglePartition + } else { + Distribution::UnspecifiedDistribution + } + }) + .collect() + } + + fn apply_expressions( + &self, + _f: &mut dyn FnMut(&Arc) -> Result, + ) -> Result { + Ok(TreeNodeRecursion::Continue) + } + + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + Ok(Arc::new(Self::new( + self.payloads.clone(), + children, + self.whole.clone(), + self.forgets_to_restrict, + ))) + } + + fn execute( + &self, + _partition: usize, + _context: Arc, + ) -> Result { + not_impl_err!("PayloadExec is only ever planned") + } + } + + impl PartitionSliceable for PayloadExec { + fn restrict_to_partitions( + &self, + children: Vec>, + partitions: &[usize], + ) -> Result> { + let payloads = if self.forgets_to_restrict { + self.payloads.clone() + } else { + slice_by_global_partition( + &self.payloads, + partitions, + "PayloadExec", + "payloads", + )? + }; + Ok(Arc::new(Self::new( + payloads, + children, + self.whole.clone(), + self.forgets_to_restrict, + ))) + } + } + + /// A resolver that finds a [`PayloadExec`]. + pub fn payload_resolver() -> PartitionSliceableResolver { + Arc::new(|plan| { + plan.downcast_ref::() + .map(|op| op as &dyn PartitionSliceable) + }) + } + + /// Carries a leaf [`PayloadExec`] and leaves the rest to Ballista's codec. + #[derive(Debug, Default)] + pub struct PayloadCodec { + ballista: BallistaPhysicalExtensionCodec, + } + + impl PhysicalExtensionCodec for PayloadCodec { + fn try_decode( + &self, + buf: &[u8], + inputs: &[Arc], + ctx: &TaskContext, + proto_converter: &dyn PhysicalProtoConverterExtension, + ) -> Result> { + let Some(payloads) = buf.strip_prefix(WIRE_PREFIX) else { + return self.ballista.try_decode(buf, inputs, ctx, proto_converter); + }; + let payloads = std::str::from_utf8(payloads) + .map_err(|e| DataFusionError::Internal(e.to_string()))?; + let payloads: Vec<&str> = + payloads.split(',').filter(|p| !p.is_empty()).collect(); + Ok(PayloadExec::leaf(&payloads)) + } + + fn try_encode( + &self, + node: Arc, + buf: &mut Vec, + proto_converter: &dyn PhysicalProtoConverterExtension, + ) -> Result<()> { + let Some(exec) = node.downcast_ref::() else { + return self.ballista.try_encode(node, buf, proto_converter); + }; + buf.extend_from_slice(WIRE_PREFIX); + buf.extend_from_slice(exec.payloads.join(",").as_bytes()); + Ok(()) + } + } +} diff --git a/docs/source/upgrading/56.0.0.md b/docs/source/upgrading/56.0.0.md index a586a24160..6c4bfdefc6 100644 --- a/docs/source/upgrading/56.0.0.md +++ b/docs/source/upgrading/56.0.0.md @@ -134,6 +134,28 @@ early, so a straggler partition leaves them idle. Set `ballista.scheduler.max_partitions_per_task` to a positive value to dispatch smaller tasks, which release their vcores sooner. +#### Stages split between tasks fail on a leaf the scheduler cannot restrict + +When a stage is split between tasks, each task gets a plan restricted to its +own partitions. A leaf the scheduler did not know how to restrict used to be +left as it was, so a task that ran partitions 4 and 5 of it ran its partitions 0 +and 1, and the stage returned some partitions twice and others not at all, with +no error. + +The scheduler now fails the job instead. Such a leaf is accepted only in a task +that runs all of its partitions, in order, so a retry of part of a stage can +fail a job too. Plans made of Ballista's own leaves (shuffle readers, file and +in-memory scans, `PlaceholderRowExec`, `EmptyExec`) are unaffected, as are +leaves read whole below a `CoalescePartitionsExec` or on a join's build side. A +`DataSourceExec` over any other source used to log a warning and run, and is +now refused when its stage is split; a node of your own that implements +`PartitionSliceable` can replace it. + +The error names the node and what to implement. Under push-staged scheduling +the job's status says only that a task could not be prepared; the reason is in +the scheduler's log. To make nodes of your own restrictable, see +[Custom Physical Plan Nodes](../user-guide/extending-components.md#custom-physical-plan-nodes). + ### API changes #### `sort_shuffle::BufferedBatches` is no longer public @@ -152,6 +174,24 @@ cannot re-run another job's output, so a caller must fail them (the scheduler posts `JobRunningFailed`). A caller that only needs the tasks to cancel can take the first element. +#### `PartitionSliceable` covers any node, and embedders can register their own + +`ballista_core::execution_plans::plan_algebra::PartitionSliceable` is no longer +limited to nodes with one child. `slice_to_partitions(child, partitions)` is now +`restrict_to_partitions(children, partitions)`, which receives all of the node's +children and also applies to leaves. Implementations of your own need the new +signature; `plan_algebra::single_child` unwraps the one child of a single-input +node. + +Nodes of your own are made known to the scheduler with +`SchedulerConfig::with_partition_sliceable_resolver`, which is consulted after +Ballista's own list. + +- `SchedulerConfig` gains a `partition_sliceable_resolver` field, so a struct + literal that doesn't end in `..Default::default()` needs to set it. +- `task_builder::restrict_plan_to_partitions` takes a third argument, + `Option<&PartitionSliceableResolver>`. Pass `None` for none. + ### Python client changes #### Summary diff --git a/docs/source/user-guide/extending-components.md b/docs/source/user-guide/extending-components.md index 3a3b665f9b..41a1010b70 100644 --- a/docs/source/user-guide/extending-components.md +++ b/docs/source/user-guide/extending-components.md @@ -29,10 +29,11 @@ new configuration extensions, object stores, logical and physical codecs ... - `session builder` - function which creates new `SessionState` for each user session - `logical codec` - overrides `LogicalCodec` - `physical codec` - overrides `PhysicalCodec` +- `partition sliceable resolver` - lets the scheduler restrict `ExecutionPlan` nodes of your own to the partitions of a task, see [Custom Physical Plan Nodes](#custom-physical-plan-nodes) Ballista executor can be configured using `ExecutorProcessConfig` which supports overriding `function registry`,`runtime producer`, `config producer`, `logical codec`, `physical codec`. -Ballista scheduler can be tuned using `SchedulerConfig` which supports overriding `config producer`, `session builder`, `logical codec`, `physical codec` +Ballista scheduler can be tuned using `SchedulerConfig` which supports overriding `config producer`, `session builder`, `logical codec`, `physical codec` and setting a `partition sliceable resolver`. ## Example: Custom Object Store Integration @@ -102,3 +103,40 @@ As with the scheduler, to keep the executor's own command line, parse :language: rust :start-after: under the License. ``` + +## Custom Physical Plan Nodes + +A plan with `ExecutionPlan` nodes of your own can run on a Ballista cluster once the scheduler and the executors have a `physical codec` that serializes them. A node that holds data per partition needs one thing more. + +When a stage is split between tasks, each task gets a copy of the plan restricted to its own partitions, and operators see only task-local partition indices: a task that runs partitions 4 and 5 executes partitions 0 and 1 of every node. The scheduler restricts the leaves it knows, such as shuffle readers and file and in-memory scans. A node it doesn't know would read its per-partition data, such as one list of files for each partition, from the wrong index, and the stage would return some partitions twice and others not at all. The scheduler fails a job whose split stage holds such a leaf. It cannot tell when an interior node holds per-partition data, so register those too. + +To make your node restrictable, implement `PartitionSliceable` for it and have the scheduler find it with a `partition sliceable resolver`: + +```rust +use ballista_core::execution_plans::plan_algebra::{PartitionSliceable, slice_by_global_partition}; + +impl PartitionSliceable for MyScanExec { + fn restrict_to_partitions( + &self, + children: Vec>, + partitions: &[usize], + ) -> Result> { + // `partitions` are indices of this node's own partitions, in the order + // the task should number them. + let files = slice_by_global_partition(&self.files, partitions, "MyScanExec", "file lists")?; + Ok(Arc::new(MyScanExec::new(files, children))) + } +} + +let config = SchedulerConfig::default().with_partition_sliceable_resolver(Arc::new(|plan| { + plan.downcast_ref::() + .map(|node| node as &dyn PartitionSliceable) +})); +``` + +- `children` are the node's inputs, none for a scan, restricted to `partitions`. An input the node declares `SinglePartition` arrives whole, as a join's build side does. Partition `i` of the node must read partition `i` of every other input. +- The restricted node must have exactly `partitions.len()` partitions. The scheduler rejects any other count. +- `partitions` is empty for a `UnionExec` branch that belongs to other tasks. Return a node without partitions. +- The restricted node is serialized with the scheduler's `physical codec`, so what it carries per partition has to survive encoding and decoding. +- The resolver is asked after Ballista's own list. A node that is read whole, such as one below a `CoalescePartitionsExec` or on a join's build side, needs none of this. +- The standalone scheduler (`SessionContext::standalone`) builds its own `SchedulerConfig`, so it cannot be given a resolver. From ad87057b4c9157f22f80451be5e7d82e355be1dc Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Sat, 31 Oct 2026 11:01:52 -0600 Subject: [PATCH 3/4] test: look through CooperativeExec for the restricted leaf in the resolver test (cherry picked from commit 2efa802ffd6f454211d516b660a3f9db648bab65) --- ballista/scheduler/src/state/mod.rs | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/ballista/scheduler/src/state/mod.rs b/ballista/scheduler/src/state/mod.rs index 8df4c9add8..a92cd0e61f 100644 --- a/ballista/scheduler/src/state/mod.rs +++ b/ballista/scheduler/src/state/mod.rs @@ -1005,7 +1005,12 @@ mod tests { ) }) .expect("the executor can decode the plan"); - payloads_of(plan.children()[0]) + // The writer's input can come wrapped in a `CooperativeExec`. + let mut input = plan.children()[0]; + while input.name() == "CooperativeExec" { + input = input.children()[0]; + } + payloads_of(input) }) .collect(); shipped.sort(); From aef82ff16e9bee9696491131399a83f706fb7902 Mon Sep 17 00:00:00 2001 From: sim Date: Tue, 10 Nov 2026 12:57:56 -0700 Subject: [PATCH 4/4] refactor: make PartitionSliceableResolver a trait and clarify what is guarded The resolver is now a trait, like the other extension points on SchedulerConfig, so methods can be added later without another breaking change. Fix the missing "can" in the unrestrictable-leaf error and compare the whole message in its test. The upgrade guide and the extending-components guide now say that only leaves are checked, and that whether a stage is split depends on cluster load and retries. Co-Authored-By: Claude Opus 5.5 --- .../core/src/execution_plans/plan_algebra.rs | 45 +++++++++------ ballista/scheduler/src/config.rs | 4 +- ballista/scheduler/src/state/task_builder.rs | 55 +++++++++++-------- ballista/scheduler/src/state/task_manager.rs | 6 +- ballista/scheduler/src/test_utils.rs | 47 +++++++++++++++- docs/source/upgrading/56.0.0.md | 26 ++++++--- .../source/user-guide/extending-components.md | 28 +++++++--- 7 files changed, 150 insertions(+), 61 deletions(-) diff --git a/ballista/core/src/execution_plans/plan_algebra.rs b/ballista/core/src/execution_plans/plan_algebra.rs index 9f42475a95..b09c08a973 100644 --- a/ballista/core/src/execution_plans/plan_algebra.rs +++ b/ballista/core/src/execution_plans/plan_algebra.rs @@ -25,6 +25,7 @@ //! conservative is the safety net: unrecognized node → property assumed //! false → caller falls back to the safer path. +use std::fmt::Debug; use std::sync::Arc; use datafusion::common::{Result, plan_err}; @@ -140,33 +141,41 @@ pub fn as_partition_sliceable( None } -/// Restricts a node that [`as_partition_sliceable`] does not list to a task's +/// Restricts nodes that [`as_partition_sliceable`] does not list to a task's /// partitions, so the scheduler can restrict `ExecutionPlan` nodes of your own, /// including ones you cannot implement [`PartitionSliceable`] on, such as a /// `DataSourceExec` over a `DataSource` of your own. Set it with /// `SchedulerConfig::with_partition_sliceable_resolver`. /// -/// It receives the node, its children already restricted as -/// [`PartitionSliceable::restrict_to_partitions`] receives them, and the global -/// partition indices the task runs. It returns `None` for a node it does not -/// know, and otherwise the node rebuilt over `children` with exactly -/// `partitions.len()` output partitions; any other count is an error. For a node -/// of your own, that is typically a `downcast_ref` to your type followed by a -/// call to its [`PartitionSliceable::restrict_to_partitions`]. -/// /// It is asked after the built-in list, and never about a node that is read /// whole: one below a collapse, or declared `SinglePartition` by its parent. It /// is called for most nodes of every task's plan as tasks are launched, so keep /// the `None` path cheap. -pub type PartitionSliceableResolver = Arc< - dyn Fn( - &Arc, - Vec>, - &[usize], - ) -> Option>> - + Send - + Sync, ->; +/// +/// This is a stopgap for an `ExecutionPlan` method that would let each node +/// restrict itself; see [`PartitionSliceable`]. +pub trait PartitionSliceableResolver: Debug + Send + Sync { + /// Restrict `plan` to `partitions`, or return `None` for a node this + /// resolver does not know. + /// + /// The restricted node is `plan` rebuilt over `children` with exactly + /// `partitions.len()` output partitions; any other count is an error. For a + /// node of your own, that is typically a `downcast_ref` to your type + /// followed by a call to its [`PartitionSliceable::restrict_to_partitions`]. + /// + /// # Arguments + /// + /// * `plan` - the node to restrict + /// * `children` - its children, already restricted as + /// [`PartitionSliceable::restrict_to_partitions`] receives them + /// * `partitions` - global partition indices the task runs + fn restrict_to_partitions( + &self, + plan: &Arc, + children: Vec>, + partitions: &[usize], + ) -> Option>>; +} /// Unwraps the only child that a single-input operator's /// [`restrict_to_partitions`](PartitionSliceable::restrict_to_partitions) diff --git a/ballista/scheduler/src/config.rs b/ballista/scheduler/src/config.rs index db1291fcb9..64f56e3f94 100644 --- a/ballista/scheduler/src/config.rs +++ b/ballista/scheduler/src/config.rs @@ -448,7 +448,7 @@ pub struct SchedulerConfig { /// stage split between tasks that holds such a leaf fails the job. Unlike /// the `override_*` fields it adds to Ballista's own list rather than /// replacing it. See [`PartitionSliceableResolver`]. - pub partition_sliceable_resolver: Option, + pub partition_sliceable_resolver: Option>, /// Rules of your own that adaptive query planning runs, for plans with /// `ExecutionPlan` nodes of your own. Unlike the `override_*` fields it /// adds to Ballista's rules rather than replacing them. Only the adaptive @@ -856,7 +856,7 @@ impl SchedulerConfig { /// Sets [`Self::partition_sliceable_resolver`]. pub fn with_partition_sliceable_resolver( mut self, - resolver: PartitionSliceableResolver, + resolver: Arc, ) -> Self { self.partition_sliceable_resolver = Some(resolver); self diff --git a/ballista/scheduler/src/state/task_builder.rs b/ballista/scheduler/src/state/task_builder.rs index d1925fefe4..db42e01ba7 100644 --- a/ballista/scheduler/src/state/task_builder.rs +++ b/ballista/scheduler/src/state/task_builder.rs @@ -74,7 +74,7 @@ use std::sync::Arc; pub fn restrict_plan_to_partitions( plan: Arc, partitions: &[usize], - resolver: Option<&PartitionSliceableResolver>, + resolver: Option<&dyn PartitionSliceableResolver>, ) -> Result> { restrict(plan, partitions, /* under_collect */ false, resolver) } @@ -116,7 +116,7 @@ fn restrict( plan: Arc, partitions: &[usize], under_collect: bool, - resolver: Option<&PartitionSliceableResolver>, + resolver: Option<&dyn PartitionSliceableResolver>, ) -> Result> { // Leaves: apply the algebraic partition-projection when we're not // under a collapse. Under collapse, leaves keep their full upstream — @@ -159,7 +159,9 @@ fn restrict( Some(sliceable) => { Some(sliceable.restrict_to_partitions(children.clone(), partitions)) } - None => resolver.and_then(|resolve| resolve(&plan, children.clone(), partitions)), + None => resolver.and_then(|resolver| { + resolver.restrict_to_partitions(&plan, children.clone(), partitions) + }), }; if let Some(restricted) = restricted { return checked_restriction(plan.name(), restricted?, partitions); @@ -179,7 +181,7 @@ fn restrict_children( plan: &Arc, partitions: &[usize], under_collect: bool, - resolver: Option<&PartitionSliceableResolver>, + resolver: Option<&dyn PartitionSliceableResolver>, ) -> Result>> { let per_child = child_scopes(plan, under_collect); plan.children() @@ -229,7 +231,7 @@ fn ensure_leaf_runs_whole( "cannot restrict `{described}` to partitions {partitions:?} of its {count}: \ the scheduler does not know how to restrict a `{name}`, and left alone it \ would number its partitions from zero in every task. A node of your own \ - be restricted by `SchedulerConfig::partition_sliceable_resolver`, for instance \ + can be restricted by `SchedulerConfig::partition_sliceable_resolver`, for instance \ by implementing `PartitionSliceable` and calling it from there." ) } @@ -864,7 +866,9 @@ mod tests { // --- Nodes the scheduler does not know: `PayloadExec` stands in for one --- - use crate::test_utils::payload::{PayloadExec, payload_resolver, payloads_of}; + use crate::test_utils::payload::{ + PayloadExec, fn_resolver, payload_resolver, payloads_of, + }; use ballista_core::execution_plans::plan_algebra::PartitionSliceableResolver; use std::sync::Mutex; @@ -886,7 +890,7 @@ mod tests { let leaf = PayloadExec::leaf(&["p0", "p1", "p2", "p3"]); let restricted = - restrict_plan_to_partitions(leaf, &[1, 3], Some(&resolver)).unwrap(); + restrict_plan_to_partitions(leaf, &[1, 3], Some(resolver.as_ref())).unwrap(); assert_eq!(payloads_of(&restricted), ["p1", "p3"]); assert_eq!(partition_count(&restricted), 2); @@ -897,8 +901,8 @@ mod tests { /// resolver building a replacement. #[test] fn resolver_restricts_a_node_without_partition_sliceable() { - let resolver: PartitionSliceableResolver = - Arc::new(|plan, children, partitions| { + let resolver: Arc = + fn_resolver(|plan, children, partitions| { assert!(children.is_empty()); plan.is::().then(|| { Ok(Arc::new( @@ -909,7 +913,7 @@ mod tests { let leaf = PayloadExec::leaf(&["p0", "p1", "p2", "p3"]); let restricted = - restrict_plan_to_partitions(leaf, &[1, 3], Some(&resolver)).unwrap(); + restrict_plan_to_partitions(leaf, &[1, 3], Some(resolver.as_ref())).unwrap(); assert!(restricted.is::()); assert_eq!(partition_count(&restricted), 2); @@ -939,17 +943,24 @@ mod tests { let error = restrict_plan_to_partitions(leaf, &[1], None) .unwrap_err() .to_string(); - assert!(error.contains("PayloadExec"), "{error}"); - assert!(error.contains("[1]"), "{error}"); - assert!(error.contains("PartitionSliceable"), "{error}"); + assert_eq!( + error, + "Error during planning: cannot restrict \ + `PayloadExec: payloads=[\"p0\", \"p1\"]` to partitions [1] of its 2: \ + the scheduler does not know how to restrict a `PayloadExec`, and left \ + alone it would number its partitions from zero in every task. A node \ + of your own can be restricted by \ + `SchedulerConfig::partition_sliceable_resolver`, for instance by \ + implementing `PartitionSliceable` and calling it from there." + ); } #[test] fn the_resolver_is_asked_only_after_the_built_in_list() { let asked = Arc::new(Mutex::new(vec![])); let record = asked.clone(); - let resolver: PartitionSliceableResolver = - Arc::new(move |plan, _children, _partitions| { + let resolver: Arc = + fn_resolver(move |plan, _children, _partitions| { record.lock().unwrap().push(plan.name().to_string()); None }); @@ -957,12 +968,12 @@ mod tests { // `ShuffleReaderExec` below it is a leaf the scheduler restricts itself. let (plan, _) = range_filter_over_reader(); - restrict_plan_to_partitions(plan, &[1, 3], Some(&resolver)).unwrap(); + restrict_plan_to_partitions(plan, &[1, 3], Some(resolver.as_ref())).unwrap(); assert!(asked.lock().unwrap().is_empty()); // A node it does not know is put to the resolver. let leaf = PayloadExec::leaf(&["p0"]); - restrict_plan_to_partitions(leaf, &[0], Some(&resolver)).unwrap(); + restrict_plan_to_partitions(leaf, &[0], Some(resolver.as_ref())).unwrap(); assert_eq!(*asked.lock().unwrap(), ["PayloadExec"]); } @@ -976,7 +987,7 @@ mod tests { ); let restricted = - restrict_plan_to_partitions(node, &[1, 3], Some(&resolver)).unwrap(); + restrict_plan_to_partitions(node, &[1, 3], Some(resolver.as_ref())).unwrap(); assert_eq!(payloads_of(&restricted), ["p1", "p3"]); let inputs = restricted.children(); @@ -993,7 +1004,7 @@ mod tests { ]))); let restricted = - restrict_plan_to_partitions(plan, &[1], Some(&resolver)).unwrap(); + restrict_plan_to_partitions(plan, &[1], Some(resolver.as_ref())).unwrap(); assert_eq!( payloads_of(restricted.children()[0]), @@ -1016,14 +1027,14 @@ mod tests { // Partition 2 is the second branch's partition 0. let restricted = - restrict_plan_to_partitions(union(), &[2], Some(&resolver)).unwrap(); + restrict_plan_to_partitions(union(), &[2], Some(resolver.as_ref())).unwrap(); let branches = restricted.children(); assert_eq!(partition_count(branches[0]), 0); assert_eq!(payloads_of(branches[1]), ["p0"]); // And partition 0 is the first branch's, which leaves the second none. let restricted = - restrict_plan_to_partitions(union(), &[0], Some(&resolver)).unwrap(); + restrict_plan_to_partitions(union(), &[0], Some(resolver.as_ref())).unwrap(); let branches = restricted.children(); assert_eq!(group_files(branches[0]), ["file0.parquet"]); assert_eq!(partition_count(branches[1]), 0); @@ -1037,7 +1048,7 @@ mod tests { let resolver = payload_resolver(); let leaf = PayloadExec::forgets_to_restrict(&["p0", "p1", "p2", "p3"]); - let error = restrict_plan_to_partitions(leaf, &[1, 3], Some(&resolver)) + let error = restrict_plan_to_partitions(leaf, &[1, 3], Some(resolver.as_ref())) .unwrap_err() .to_string(); diff --git a/ballista/scheduler/src/state/task_manager.rs b/ballista/scheduler/src/state/task_manager.rs index e42e2e32c2..b5c6a4c55e 100644 --- a/ballista/scheduler/src/state/task_manager.rs +++ b/ballista/scheduler/src/state/task_manager.rs @@ -158,7 +158,7 @@ pub struct TaskManager stage_max_failures: usize, /// Resolver for nodes Ballista does not know. See /// [`SchedulerConfig::partition_sliceable_resolver`]. - partition_sliceable_resolver: Option, + partition_sliceable_resolver: Option>, /// See [`SchedulerConfig::adaptive_rules`]. adaptive_rules: Option>, } @@ -1023,7 +1023,7 @@ impl TaskManager let restricted = restrict_plan_to_partitions( task.plan.clone(), &task.global_input_partition_ids, - self.partition_sliceable_resolver.as_ref(), + self.partition_sliceable_resolver.as_deref(), )?; let mut plan_buf: Vec = vec![]; let plan_proto = PhysicalPlanNode::try_from_physical_plan( @@ -1145,7 +1145,7 @@ impl TaskManager let restricted = restrict_plan_to_partitions( task.plan.clone(), &task.global_input_partition_ids, - self.partition_sliceable_resolver.as_ref(), + self.partition_sliceable_resolver.as_deref(), )?; let mut plan_buf: Vec = vec![]; let plan_proto = PhysicalPlanNode::try_from_physical_plan(restricted, codec)?; diff --git a/ballista/scheduler/src/test_utils.rs b/ballista/scheduler/src/test_utils.rs index a9f179b4fa..1992fca095 100644 --- a/ballista/scheduler/src/test_utils.rs +++ b/ballista/scheduler/src/test_utils.rs @@ -1774,13 +1774,56 @@ pub mod payload { } /// A resolver that finds a [`PayloadExec`]. - pub fn payload_resolver() -> PartitionSliceableResolver { - Arc::new(|plan, children, partitions| { + pub fn payload_resolver() -> Arc { + fn_resolver(|plan, children, partitions| { plan.downcast_ref::() .map(|op| op.restrict_to_partitions(children, partitions)) }) } + /// A [`PartitionSliceableResolver`] that calls `resolve`. + pub fn fn_resolver(resolve: F) -> Arc + where + F: Fn( + &Arc, + Vec>, + &[usize], + ) -> Option>> + + Send + + Sync + + 'static, + { + Arc::new(FnResolver(resolve)) + } + + struct FnResolver(F); + + impl std::fmt::Debug for FnResolver { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + f.write_str("FnResolver") + } + } + + impl PartitionSliceableResolver for FnResolver + where + F: Fn( + &Arc, + Vec>, + &[usize], + ) -> Option>> + + Send + + Sync, + { + fn restrict_to_partitions( + &self, + plan: &Arc, + children: Vec>, + partitions: &[usize], + ) -> Option>> { + (self.0)(plan, children, partitions) + } + } + /// Carries a leaf [`PayloadExec`] and leaves the rest to Ballista's codec. #[derive(Debug, Default)] pub struct PayloadCodec { diff --git a/docs/source/upgrading/56.0.0.md b/docs/source/upgrading/56.0.0.md index ece9a01a9d..af96fa57cf 100644 --- a/docs/source/upgrading/56.0.0.md +++ b/docs/source/upgrading/56.0.0.md @@ -192,8 +192,19 @@ fail a job too. Plans made of Ballista's own leaves (shuffle readers, file and in-memory scans, `PlaceholderRowExec`, `EmptyExec`) are unaffected, as are leaves read whole below a `CoalescePartitionsExec` or on a join's build side. A `DataSourceExec` over any other source used to log a warning and run, and is -now refused when its stage is split; a node of your own that implements -`PartitionSliceable` can replace it. +now refused when its stage is split unless a resolver restricts it. + +Whether a stage is split depends on how many task slots are free when it is +scheduled, and on whether part of it is retried, so a query that holds such a +leaf can pass on an idle or small cluster and fail on a busy one. To test a plan +that holds nodes of your own, set `ballista.scheduler.max_partitions_per_task` +to a small value, which splits every stage with more partitions than that. + +Only leaves are checked. An interior node the scheduler does not know is left as +it was while its children are restricted, because the scheduler cannot tell one +that holds data per partition from one that does not, such as a `FilterExec`. An +interior node of your own that holds data per partition still returns wrong +results in a split stage unless a resolver restricts it. The error names the node and what to implement. Under push-staged scheduling the job's status says only that a task could not be prepared; the reason is in @@ -228,15 +239,16 @@ signature; `plan_algebra::single_child` unwraps the one child of a single-input node. Nodes of your own are made known to the scheduler with -`SchedulerConfig::with_partition_sliceable_resolver`, which is consulted after -Ballista's own list. The resolver restricts the node itself, given its -restricted children and the task's partitions, so it also covers nodes that -cannot implement `PartitionSliceable`, such as a `DataSourceExec`. +`SchedulerConfig::with_partition_sliceable_resolver`, which takes an +implementation of the new `plan_algebra::PartitionSliceableResolver` trait and +is consulted after Ballista's own list. The resolver restricts the node itself, +given its restricted children and the task's partitions, so it also covers +nodes that cannot implement `PartitionSliceable`, such as a `DataSourceExec`. - `SchedulerConfig` gains a `partition_sliceable_resolver` field, so a struct literal that doesn't end in `..Default::default()` needs to set it. - `task_builder::restrict_plan_to_partitions` takes a third argument, - `Option<&PartitionSliceableResolver>`. Pass `None` for none. + `Option<&dyn PartitionSliceableResolver>`. Pass `None` for none. #### Partition locations hold an `ExecutorEndpoint` diff --git a/docs/source/user-guide/extending-components.md b/docs/source/user-guide/extending-components.md index 477a636d56..2228bf3425 100644 --- a/docs/source/user-guide/extending-components.md +++ b/docs/source/user-guide/extending-components.md @@ -137,12 +137,14 @@ As with the scheduler, to keep the executor's own command line, parse A plan with `ExecutionPlan` nodes of your own can run on a Ballista cluster once the scheduler and the executors have a `physical codec` that serializes them. A node that holds data per partition needs one thing more. -When a stage is split between tasks, each task gets a copy of the plan restricted to its own partitions, and operators see only task-local partition indices: a task that runs partitions 4 and 5 executes partitions 0 and 1 of every node. The scheduler restricts the leaves it knows, such as shuffle readers and file and in-memory scans. A node it doesn't know would read its per-partition data, such as one list of files for each partition, from the wrong index, and the stage would return some partitions twice and others not at all. The scheduler fails a job whose split stage holds such a leaf. It cannot tell when an interior node holds per-partition data, so register those too. +When a stage is split between tasks, each task gets a copy of the plan restricted to its own partitions, and operators see only task-local partition indices: a task that runs partitions 4 and 5 executes partitions 0 and 1 of every node. The scheduler restricts the leaves it knows, such as shuffle readers and file and in-memory scans. A node it doesn't know would read its per-partition data, such as one list of files for each partition, from the wrong index, and the stage would return some partitions twice and others not at all. The scheduler fails a job whose split stage holds such a leaf. It does not check interior nodes: it cannot tell one that holds per-partition data from one that does not, such as a `FilterExec`, so an interior node of yours that holds per-partition data returns wrong results in a split stage unless you register it. Whether a stage is split depends on free task slots and on retries, so test with `ballista.scheduler.max_partitions_per_task` set low rather than relying on a small cluster to show the failure. -To make your node restrictable, give the scheduler a `partition sliceable resolver`. It is called with a node, its children already restricted, and the partitions the task runs, and returns the restricted node, or `None` for a node it doesn't know. A node of your own can implement `PartitionSliceable` to keep the slicing beside its fields: +To make your node restrictable, give the scheduler a `PartitionSliceableResolver`. It is called with a node, its children already restricted, and the partitions the task runs, and returns the restricted node, or `None` for a node it doesn't know. A node of your own can implement `PartitionSliceable` to keep the slicing beside its fields: ```rust -use ballista_core::execution_plans::plan_algebra::{PartitionSliceable, slice_by_global_partition}; +use ballista_core::execution_plans::plan_algebra::{ + PartitionSliceable, PartitionSliceableResolver, slice_by_global_partition, +}; impl PartitionSliceable for MyScanExec { fn restrict_to_partitions( @@ -157,10 +159,22 @@ impl PartitionSliceable for MyScanExec { } } -let config = SchedulerConfig::default().with_partition_sliceable_resolver(Arc::new(|plan, children, partitions| { - plan.downcast_ref::() - .map(|node| node.restrict_to_partitions(children, partitions)) -})); +#[derive(Debug)] +struct MyResolver; + +impl PartitionSliceableResolver for MyResolver { + fn restrict_to_partitions( + &self, + plan: &Arc, + children: Vec>, + partitions: &[usize], + ) -> Option>> { + plan.downcast_ref::() + .map(|node| node.restrict_to_partitions(children, partitions)) + } +} + +let config = SchedulerConfig::default().with_partition_sliceable_resolver(Arc::new(MyResolver)); ``` - `children` are the node's inputs, none for a scan, restricted to `partitions`. An input the node declares `SinglePartition` arrives whole, as a join's build side does. Partition `i` of the node must read partition `i` of every other input.