Multi-partition tasks: partition_slice plumbing + SortShuffleWriter refactor

.gitignore

@@ -116,3 +116,4 @@ CLAUDE.md
 .worktrees/
 # ignore insta captures
 *-snap
+.local/

ballista/core/proto/ballista.proto

@@ -132,7 +132,6 @@ message ExecutionGraph {
   uint64 output_partitions = 5;
   repeated PartitionLocation output_locations = 6;
   string scheduler_id = 7;
-  uint32 task_id_gen = 8;
   repeated StageAttempts failed_attempts = 9;
   string job_name = 10;
   uint64 start_time = 11;
@@ -289,10 +288,30 @@ message PartitionId {
   uint32 partition_id = 4;
 }
 
+// Within-stage task identity (paired with job_id + stage_id on the parent
+// MultiTaskDefinition). One task processes a slice of partitions;
+// `task_id` names the task within the stage (it's the task's append-order
+// slot in `RunningStage.task_infos`, so (job_id, stage_id, task_id) is
+// globally unique). `global_output_partition_ids` gives the concrete
+// global partition ids the task's restricted plan is covering. Writers
+// use these to name shuffle files with global identity (via
+// `create_shuffle_path`) so downstream reads have a stable, canonical
+// address regardless of how the scheduler split the work.
 message TaskId {
   uint32 task_id = 1;
   uint32 task_attempt_num = 2;
-  uint32 partition_id = 3;
+  // Global partition ids covered by this task, in slice order. Position i in
+  // the restricted plan corresponds to `global_output_partition_ids[i]` globally.
+  // `[packed=false]` keeps the single-partition case wire-compatible with the
+  // legacy `uint32 partition_id = 3` field: one element writes as a single
+  // varint at tag 3, exactly matching the old scalar encoding.
+  repeated uint32 global_output_partition_ids = 3 [packed=false];
+  // Vcores this task consumed from the executor's budget at bind time.
+  // The executor uses this to size the task's memory pool proportionally
+  // (`pool = total_memory * vcores_consumed / total_vcores`), so a
+  // multi-partition task claiming N vcores gets N/total of the executor's
+  // memory budget rather than a single per-vcore share.
+  uint32 vcores_consumed = 5;
 }
 
 message PartitionStats {
@@ -359,6 +378,10 @@ message OperatorMetric {
     NamedRatio ratio = 14;
     uint64 output_batches = 15;
   }
+  // Local partition index within the reporting task's plan (0..slice_len-1).
+  // The scheduler maps this to a global partition id via the task's
+  // `global_input_partition_ids` when folding metrics into stage metrics.
+  optional uint32 partition = 16;
 }
 
 // Used by scheduler
@@ -526,11 +549,14 @@ message ShuffleWritePartition {
 }
 
 message TaskStatus {
+  // Task's append-order slot in `RunningStage.task_infos`; (job_id,
+  // stage_id, task_id) is globally unique. Under the multi-partition-task
+  // model one task processes a partition slice, so this names the task,
+  // not a single partition.
   uint32 task_id = 1;
   string job_id = 2;
   uint32 stage_id = 3;
   uint32 stage_attempt_num = 4;
-  uint32 partition_id = 5;
   uint64 launch_time = 6;
   uint64 start_exec_time = 7;
   uint64 end_exec_time = 8;
@@ -549,17 +575,30 @@ message PollWorkParams {
   repeated TaskStatus task_status = 3;
 }
 
+// Pull-based single-task dispatch. One task processes a slice of
+// partitions; `task_id` names the task within the stage (its append-order
+// slot in `RunningStage.task_infos`), `global_output_partition_ids` gives
+// the concrete global partition ids it covers. The task's plan is
+// scheduler-side shrink-restricted so its leaves report `slice.len()`
+// partitions; the writer uses `global_output_partition_ids` to attach
+// global identity to shuffle files (except in cases where the plan itself
+// resets partitioning: the writer walks its child plan to detect
+// SortPreservingMergeExec or RepartitionExec::Hash and picks the right
+// global mapping).
 message TaskDefinition {
   uint32 task_id = 1;
   uint32 task_attempt_num = 2;
   string job_id = 3;
   uint32 stage_id = 4;
   uint32 stage_attempt_num = 5;
-  uint32 partition_id = 6;
   bytes plan = 7;
   string session_id = 9;
   uint64 launch_time = 10;
   repeated KeyValuePair props = 11;
+  // Global partition ids covered by this task, in slice order.
+  repeated uint32 global_output_partition_ids = 12;
+  // Vcores this task consumed at bind time — see TaskId.vcores_consumed.
+  uint32 vcores_consumed = 13;
 }
 
 // A set of tasks in the same stage
@@ -840,10 +879,11 @@ message RemoveJobDataResult {
 }
 
 message RunningTaskInfo {
+  // Per-stage append-order slot; (job_id, stage_id, task_id) is globally
+  // unique (see TaskStatus).
   uint32 task_id = 1;
   string job_id = 2;
   uint32 stage_id = 3;
-  uint32 partition_id = 4;;
 }
 
 service SchedulerGrpc {

ballista/core/src/config.rs

@@ -142,6 +142,10 @@ pub const BALLISTA_CHAOS_EXECUTION_SEED: &str = "ballista.testing.chaos_executio
 /// Valid values are: none, lz4, zstd
 pub const BALLISTA_SHUFFLE_COMPRESSION_CODEC: &str = "ballista.shuffle.compression.codec";
 
+/// Configuration key for the scheduler's per-task partition-slice cap.
+pub const BALLISTA_SCHEDULER_MAX_PARTITIONS_PER_TASK: &str =
+    "ballista.scheduler.max_partitions_per_task";
+
 /// Result type for configuration parsing operations.
 pub type ParseResult<T> = result::Result<T, String>;
 use std::sync::LazyLock;
@@ -335,6 +339,20 @@ static CONFIG_ENTRIES: LazyLock<HashMap<String, ConfigEntry>> = LazyLock::new(||
             DataType::Utf8,
             Some("lz4".to_string()),
         ),
+        ConfigEntry::new(
+            BALLISTA_SCHEDULER_MAX_PARTITIONS_PER_TASK.to_string(),
+            "Upper bound on the number of input partitions packed into a single \
+             task's `partition_slice`. `1` (default) is the pre-multi-partition-tasks \
+             execution model — one task per partition, matching the Spark user \
+             persona and preserving master's behaviour on merge. Raise to enable \
+             multi-partition tasks (fewer tasks, parallel-sort / parallel-join wins); \
+             `0` means unbounded — the scheduler fills each task up to the executor's \
+             free vcore count. Does not apply to collapse stages, which must pack \
+             their full pending queue into a single task for correctness."
+                .to_string(),
+            DataType::UInt64,
+            Some(1.to_string()),
+        ),
     ];
     entries
         .into_iter()
@@ -619,6 +637,11 @@ impl BallistaConfig {
         self.get_bool_setting(BALLISTA_ADAPTIVE_JOIN_ENABLED)
     }
 
+    /// Returns the scheduler's per-task partition-slice cap. `0` means unbounded.
+    pub fn max_partitions_per_task(&self) -> usize {
+        self.get_usize_setting(BALLISTA_SCHEDULER_MAX_PARTITIONS_PER_TASK)
+    }
+
     /// Returns whether chaos-monkey execution injection is enabled.
     pub fn chaos_execution_enabled(&self) -> bool {
         self.get_bool_setting(BALLISTA_CHAOS_EXECUTION_ENABLED)

ballista/core/src/execution_plans/distributed_explain_analyze.rs

@@ -326,11 +326,13 @@ mod tests {
                         metrics: vec![
                             OperatorMetric {
                                 metric: Some(operator_metric::Metric::OutputRows(4)),
+                                partition: None,
                             },
                             OperatorMetric {
                                 metric: Some(operator_metric::Metric::ElapseTime(
                                     15_000_000,
                                 )),
+                                partition: None,
                             },
                         ],
                     }],
@@ -344,6 +346,7 @@ mod tests {
                         operator_desc: "FilterExec: a@0 > 1".to_string(),
                         metrics: vec![OperatorMetric {
                             metric: Some(operator_metric::Metric::OutputRows(2)),
+                            partition: None,
                         }],
                     }],
                 },

ballista/core/src/execution_plans/mod.rs

@@ -37,6 +37,7 @@ pub use shuffle_reader::{CoalescePlan, PartitionGroup, ShuffleReaderExec};
 pub use shuffle_reader::{stats_for_partition, stats_for_partitions};
 pub use shuffle_writer::DEFAULT_SHUFFLE_CHANNEL_CAPACITY;
 pub use shuffle_writer::ShuffleWriterExec;
+pub use shuffle_writer::compute_global_output_partition_ids;
 pub use shuffle_writer_trait::ShuffleWriter;
 pub use sort_shuffle::SortShuffleWriterExec;
 pub use unresolved_shuffle::UnresolvedShuffleExec;

ballista/core/src/execution_plans/shuffle_reader.rs

@@ -1772,7 +1772,10 @@ mod tests {
             .map_err(|e| DataFusionError::Execution(format!("{e:?}")))
             .unwrap();
 
-        assert_eq!(result.len(), 2);
+        // Writer's K=1 hash routes every row to partition 0. Its coordinator
+        // reads its full input slice (2 partitions × 2 batches = 4 batches),
+        // all routed to the single output file.
+        assert_eq!(result.len(), 4);
         for b in result {
             assert_eq!(b, create_test_batch())
         }