feat: let custom ExecutionPlan nodes restrict themselves to a task's partitions

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::<RuntimeStatsExec>().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<dyn ExecutionPlan>,
+        children: Vec<Arc<dyn ExecutionPlan>>,
         partitions: &[usize],
     ) -> Result<Arc<dyn ExecutionPlan>>;
 }
 
 /// 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<dyn ExecutionPlan>,
 ) -> 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<dyn ExecutionPlan>,
+        children: Vec<Arc<dyn ExecutionPlan>>,
+        partitions: &[usize],
+    ) -> Option<Result<Arc<dyn ExecutionPlan>>>;
+}
+
+/// 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<Arc<dyn ExecutionPlan>>,
+) -> Result<Arc<dyn ExecutionPlan>> {
+    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<T: Clone>(
 ) -> Result<Vec<T>> {
     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()
 }

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<dyn ExecutionPlan>,
+        children: Vec<Arc<dyn ExecutionPlan>>,
         partitions: &[usize],
     ) -> Result<Arc<dyn ExecutionPlan>> {
+        let child = single_child("PrefixMergeExec", children)?;
         let applies = self
             .applies
             .iter()

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<dyn ExecutionPlan>,
+        children: Vec<Arc<dyn ExecutionPlan>>,
         partitions: &[usize],
     ) -> Result<Arc<dyn ExecutionPlan>> {
+        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(),

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<Arc<dyn PhysicalExtensionCodec>>,
     /// Override function for customizing gRPC client endpoints before they are used
     pub override_create_grpc_client_endpoint: Option<EndpointOverrideFn>,
+    /// 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<Arc<dyn PartitionSliceableResolver>>,
     /// 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<dyn PartitionSliceableResolver>,
+    ) -> Self {
+        self.partition_sliceable_resolver = Some(resolver);
+        self
+    }
+
     /// Sets [`Self::adaptive_rules`].
     pub fn with_adaptive_rules(mut self, rules: Arc<dyn AdaptiveRules>) -> Self {
         self.adaptive_rules = Some(rules);
@@ -994,6 +1013,7 @@ impl TryFrom<Config> 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,

ballista/scheduler/src/planner.rs

@@ -669,9 +669,11 @@ fn can_stay_inline(plan: &Arc<dyn ExecutionPlan>) -> 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<dyn ExecutionPlan>) -> 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)
 }