diff --git a/ballista/core/src/execution_plans/plan_algebra.rs b/ballista/core/src/execution_plans/plan_algebra.rs index 3c8ff03677..b09c08a973 100644 --- a/ballista/core/src/execution_plans/plan_algebra.rs +++ b/ballista/core/src/execution_plans/plan_algebra.rs @@ -25,9 +25,10 @@ //! 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::{DataFusionError, Result}; +use datafusion::common::{Result, plan_err}; use datafusion::physical_plan::ExecutionPlan; use datafusion::physical_plan::filter::FilterExec; use datafusion::physical_plan::projection::ProjectionExec; @@ -77,15 +78,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, @@ -94,25 +101,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> { @@ -125,6 +141,63 @@ pub fn as_partition_sliceable( None } +/// 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 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. +/// +/// 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) +/// 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> { + match <[_; 1]>::try_from(children) { + Ok([child]) => Ok(child), + Err(children) => plan_err!( + "{owner}: expected 1 child to restrict, got {}", + children.len() + ), + } +} + /// 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. @@ -143,13 +216,12 @@ pub fn slice_by_global_partition( ) -> Result> { partitions .iter() - .map(|&global| { - values.get(global).cloned().ok_or_else(|| { - DataFusionError::Internal(format!( - "{owner}: partition index {global} out of bounds ({} {what})", - values.len() - )) - }) + .map(|&global| match values.get(global) { + Some(value) => Ok(value.clone()), + None => plan_err!( + "{owner}: partition index {global} out of bounds ({} {what})", + values.len() + ), }) .collect() } diff --git a/ballista/core/src/execution_plans/prefix_merge.rs b/ballista/core/src/execution_plans/prefix_merge.rs index 8b2cd8c3b3..898a8fb3cb 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 938a338a27..8eb7a5e3cc 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 e5fe28a173..64f56e3f94 100644 --- a/ballista/scheduler/src/config.rs +++ b/ballista/scheduler/src/config.rs @@ -29,6 +29,7 @@ use crate::SessionBuilder; use crate::cluster::DistributionPolicy; use crate::scheduler_server::JobIdGenerator; use crate::state::adaptive::AdaptiveRules; +use ballista_core::execution_plans::plan_algebra::PartitionSliceableResolver; use ballista_core::extension::EndpointOverrideFn; use ballista_core::{ ConfigProducer, JobId, config::TaskSchedulingPolicy, ids::new_instance_id, @@ -440,6 +441,14 @@ 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, + /// Restricts `ExecutionPlan` nodes of your own that hold data per + /// partition to a task's partitions, using + /// [`PartitionSliceable`](ballista_core::execution_plans::plan_algebra::PartitionSliceable) + /// or by building the restricted node itself. 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>, /// 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 @@ -530,6 +539,7 @@ impl Default for SchedulerConfig { override_logical_codec: None, override_physical_codec: None, override_create_grpc_client_endpoint: None, + partition_sliceable_resolver: None, adaptive_rules: None, use_tls: false, min_ready_executors: 1, @@ -843,6 +853,15 @@ impl SchedulerConfig { self } + /// Sets [`Self::partition_sliceable_resolver`]. + pub fn with_partition_sliceable_resolver( + mut self, + resolver: Arc, + ) -> Self { + self.partition_sliceable_resolver = Some(resolver); + self + } + /// Sets [`Self::adaptive_rules`]. pub fn with_adaptive_rules(mut self, rules: Arc) -> Self { self.adaptive_rules = Some(rules); @@ -994,6 +1013,7 @@ impl TryFrom for SchedulerConfig { override_physical_codec: None, override_session_builder: None, override_create_grpc_client_endpoint: None, + partition_sliceable_resolver: None, adaptive_rules: None, use_tls: false, min_ready_executors: opt.min_ready_executors, diff --git a/ballista/scheduler/src/planner.rs b/ballista/scheduler/src/planner.rs index 7fc8036dd9..c8990ee700 100644 --- a/ballista/scheduler/src/planner.rs +++ b/ballista/scheduler/src/planner.rs @@ -669,9 +669,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 e15478ef7c..37c0c15d74 100644 --- a/ballista/scheduler/src/state/mod.rs +++ b/ballista/scheduler/src/state/mod.rs @@ -684,14 +684,21 @@ mod tests { use super::*; use crate::config::SchedulerConfig; use crate::scheduler_server::timestamp_millis; + use crate::state::executor_manager::ExecutorManager; + use crate::test_utils::payload::{ + PayloadCodec, PayloadExec, payload_resolver, payloads_of, + }; use crate::test_utils::{ BlackholeTaskLauncher, RejectingTaskLauncher, aggregation_logical_plan, - register_virtual_executor, test_cluster_context, + mock_executor, register_virtual_executor, test_cluster_context, }; use ballista_core::config::BALLISTA_EXECUTION_MODE; use ballista_core::extension::SessionConfigExt; - use ballista_core::serde::BallistaCodec; - + use ballista_core::serde::protobuf::MultiTaskDefinition; + use ballista_core::serde::scheduler::{ + ExecutorData, ExecutorMetadata, ExecutorSpecification, + }; + use ballista_core::serde::{BallistaCodec, BallistaLogicalExtensionCodec}; use datafusion::execution::context::SessionConfig; use datafusion_proto::protobuf::{LogicalPlanNode, PhysicalPlanNode}; @@ -932,6 +939,174 @@ 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"); + // 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(); + assert_eq!(shipped, [["p0", "p1"], ["p2", "p3"]]); + + Ok(()) + } + /// A pipelined task holds no vcores of its own, so an executor rejecting /// it hands none back: the vcores belong to the job's reservation, which /// is released when the job fails. diff --git a/ballista/scheduler/src/state/task_builder.rs b/ballista/scheduler/src/state/task_builder.rs index a81fcc0c9d..db42e01ba7 100644 --- a/ballista/scheduler/src/state/task_builder.rs +++ b/ballista/scheduler/src/state/task_builder.rs @@ -36,12 +36,14 @@ //! 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::{ PipelinedShuffleReaderExec, RangeShuffleReaderExec, ShuffleReaderExec, UnresolvedShuffleExec, }; -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, @@ -50,13 +52,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; @@ -65,11 +67,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 by `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<&dyn PartitionSliceableResolver>, ) -> Result> { - restrict(plan, partitions, /* under_collect */ false) + restrict(plan, partitions, /* under_collect */ false, resolver) } /// The stages whose output a leaf of the stage plan `plan` reads in full, @@ -88,9 +95,7 @@ pub(crate) fn stages_read_whole(plan: &Arc) -> Vec { return; } // The same scopes `restrict` gives each child. - let scopes = if !under_collect - && (plan.is::() || as_partition_sliceable(plan).is_some()) - { + let scopes = if !under_collect && plan.is::() { vec![false; plan.children().len()] } else { child_scopes(plan, under_collect) @@ -111,6 +116,7 @@ fn restrict( plan: Arc, partitions: &[usize], under_collect: bool, + 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 — @@ -121,22 +127,6 @@ 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); - } - // UnionExec: parent partition `p` maps to exactly one child's local // partition (`p` minus the sum of preceding children's counts). Split // `partitions` into disjoint per-child sub-slices and recurse; a child @@ -151,21 +141,99 @@ 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); } - // Interior: compute per-child scope, recurse, rebuild the node. - let per_child = child_scopes(&plan, under_collect); - let new_children: Vec> = plan - .children() + let children = restrict_children(&plan, partitions, under_collect, resolver)?; + if under_collect { + return replace_children_if_necessary(plan, children); + } + + // 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. A node Ballista does + // not list may be restricted by the resolver. + let restricted = match as_partition_sliceable(&plan) { + Some(sliceable) => { + Some(sliceable.restrict_to_partitions(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); + } + + // A leaf nothing above could restrict is only right if the task runs all + // of it. + if children.is_empty() { + ensure_leaf_runs_whole(&plan, partitions)?; + return Ok(plan); + } + replace_children_if_necessary(plan, 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<&dyn 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 be restricted by `SchedulerConfig::partition_sliceable_resolver`, for instance \ + by implementing `PartitionSliceable` and calling it from there." + ) } /// Bucket each parent partition of a `UnionExec` into its owning child's @@ -274,10 +342,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 /// @@ -331,14 +399,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))); } @@ -350,14 +416,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 @@ -368,10 +432,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, }; @@ -396,20 +457,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); } @@ -427,18 +483,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) @@ -449,7 +506,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::{ ExecutorEndpoint, PartitionId, PartitionLocation, @@ -493,7 +550,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"); @@ -507,7 +564,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"); @@ -528,7 +585,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"); @@ -570,7 +627,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"); @@ -582,7 +639,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"); @@ -606,13 +663,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> { @@ -635,7 +696,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![]], @@ -646,7 +707,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]], @@ -664,7 +725,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!( @@ -686,7 +748,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]); } @@ -716,7 +778,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]), @@ -730,15 +792,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 = @@ -752,11 +813,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))), @@ -773,8 +833,19 @@ mod tests { ) .unwrap(), ); + (plan, raw_bounds) + } - let restricted = restrict_plan_to_partitions(plan, &[1, 3]).unwrap(); + /// 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], None).unwrap(); let rf = restricted .downcast_ref::() .expect("top must remain RangeFilterExec"); @@ -792,4 +863,213 @@ 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, fn_resolver, 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.as_ref())).unwrap(); + + assert_eq!(payloads_of(&restricted), ["p1", "p3"]); + assert_eq!(partition_count(&restricted), 2); + } + + /// A node that cannot implement `PartitionSliceable`, such as a + /// `DataSourceExec` over a source of your own, is restricted by the + /// resolver building a replacement. + #[test] + fn resolver_restricts_a_node_without_partition_sliceable() { + let resolver: Arc = + fn_resolver(|plan, children, partitions| { + assert!(children.is_empty()); + plan.is::().then(|| { + Ok(Arc::new( + EmptyExec::new(plan.schema()).with_partitions(partitions.len()), + ) as Arc) + }) + }); + let leaf = PayloadExec::leaf(&["p0", "p1", "p2", "p3"]); + + let restricted = + restrict_plan_to_partitions(leaf, &[1, 3], Some(resolver.as_ref())).unwrap(); + + assert!(restricted.is::()); + 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_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: Arc = + fn_resolver(move |plan, _children, _partitions| { + 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.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.as_ref())).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.as_ref())).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.as_ref())).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.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.as_ref())).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.as_ref())) + .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 1bc73768fd..b5c6a4c55e 100644 --- a/ballista/scheduler/src/state/task_manager.rs +++ b/ballista/scheduler/src/state/task_manager.rs @@ -38,6 +38,7 @@ use crate::state::task_builder::restrict_plan_to_partitions; 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::{ @@ -155,6 +156,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>, /// See [`SchedulerConfig::adaptive_rules`]. adaptive_rules: Option>, } @@ -217,20 +221,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, - adaptive_rules: config.adaptive_rules.clone(), - } + 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, @@ -247,6 +247,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(), adaptive_rules: config.adaptive_rules.clone(), } } @@ -1006,6 +1007,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, @@ -1020,6 +1023,7 @@ impl TaskManager let restricted = restrict_plan_to_partitions( task.plan.clone(), &task.global_input_partition_ids, + self.partition_sliceable_resolver.as_deref(), )?; let mut plan_buf: Vec = vec![]; let plan_proto = PhysicalPlanNode::try_from_physical_plan( @@ -1141,6 +1145,7 @@ impl TaskManager let restricted = restrict_plan_to_partitions( task.plan.clone(), &task.global_input_partition_ids, + 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 e97d23b810..1992fca095 100644 --- a/ballista/scheduler/src/test_utils.rs +++ b/ballista/scheduler/src/test_utils.rs @@ -1583,3 +1583,283 @@ 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() -> 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 { + 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 77f8957288..af96fa57cf 100644 --- a/docs/source/upgrading/56.0.0.md +++ b/docs/source/upgrading/56.0.0.md @@ -178,6 +178,39 @@ 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 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 +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 @@ -196,6 +229,27 @@ 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 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<&dyn PartitionSliceableResolver>`. Pass `None` for none. + #### Partition locations hold an `ExecutorEndpoint` A partition location only has to say where to fetch from, so diff --git a/docs/source/user-guide/extending-components.md b/docs/source/user-guide/extending-components.md index d4a980adb8..2228bf3425 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`. ## Combining Codecs @@ -131,3 +132,55 @@ 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 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 `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, PartitionSliceableResolver, 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))) + } +} + +#[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. +- 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 does not need `PartitionSliceable`. It can build the restricted node itself, which is the way to restrict a node you cannot implement the trait on, such as a `DataSourceExec` over a `DataSource` of your own: match it with `downcast_ref`, then build a `DataSourceExec` over the restricted source. +- 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.