perf(scheduler): merge a task's sorted partitions before the shuffle write

ballista/core/src/config.rs

@@ -172,6 +172,11 @@ 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 merging a task's sorted partitions into one output
+/// partition in the passthrough shuffle writer.
+pub const BALLISTA_SHUFFLE_MERGE_ORDERED_PASSTHROUGH: &str =
+    "ballista.shuffle.merge_ordered_passthrough";
+
 /// 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";
@@ -414,6 +419,16 @@ static CONFIG_ENTRIES: LazyLock<HashMap<String, ConfigEntry>> = LazyLock::new(||
             DataType::Utf8,
             Some("lz4".to_string()),
         ),
+        ConfigEntry::new(BALLISTA_SHUFFLE_MERGE_ORDERED_PASSTHROUGH.to_string(),
+                         "Merge each task's sorted partitions into one output partition \
+                          instead of writing one file per partition, where the consumer \
+                          reads the stage back with an ordering-preserving merge. Cuts that \
+                          consumer's fan-in from the stage's partition count to its task \
+                          count, which bounds the merge's working set, at the cost of the \
+                          parallel single-source reads it replaces. Applies to adaptive \
+                          query planning only.".to_string(),
+                         DataType::Boolean,
+                         Some(true.to_string())),
         ConfigEntry::new(
             BALLISTA_SCHEDULER_MAX_PARTITIONS_PER_TASK.to_string(),
             "Upper bound on the number of input partitions packed into a single \
@@ -677,6 +692,12 @@ impl BallistaConfig {
         self.get_usize_setting(BALLISTA_SHUFFLE_WRITER_CHANNEL_CAPACITY)
     }
 
+    /// Whether a passthrough shuffle stage with a sorted input merges the
+    /// task's partitions into one.
+    pub fn shuffle_merge_ordered_passthrough(&self) -> bool {
+        self.get_bool_setting(BALLISTA_SHUFFLE_MERGE_ORDERED_PASSTHROUGH)
+    }
+
     /// Returns the per-task buffered-bytes budget at which the sort shuffle
     /// writer spills its in-memory batches to disk.
     pub fn shuffle_sort_based_memory_limit_per_task_bytes(&self) -> usize {

ballista/core/src/execution_plans/shuffle_writer.rs

@@ -356,6 +356,11 @@ pub struct ShuffleWriterExec {
     /// → 0..K) the writer detects that at path-build time and uses `local`
     /// directly instead of `global_output_partition_ids[local]`.
     global_output_partition_ids: Vec<usize>,
+    /// Whether the scheduler should merge this task's sorted partitions into one
+    /// before the write. Scheduler-side only: `task_builder` consumes it to
+    /// rewrite the per-task plan, after which the executor just runs the plan it
+    /// is handed, so it is deliberately not carried over the wire.
+    merge_per_task: bool,
     /// Execution metrics
     metrics: ExecutionPlanMetricsSet,
     /// Plan properties
@@ -375,6 +380,7 @@ impl Clone for ShuffleWriterExec {
             work_dir: self.work_dir.clone(),
             task_id: self.task_id,
             global_output_partition_ids: self.global_output_partition_ids.clone(),
+            merge_per_task: self.merge_per_task,
             metrics: self.metrics.clone(),
             properties: self.properties.clone(),
             state: self.state.clone(),
@@ -461,6 +467,7 @@ impl ShuffleWriterExec {
             work_dir,
             task_id: 0,
             global_output_partition_ids: default_partition_slice,
+            merge_per_task: false,
             metrics: ExecutionPlanMetricsSet::new(),
             properties,
             state: Arc::new(Mutex::new(WriterState {
@@ -478,6 +485,18 @@ impl ShuffleWriterExec {
         self
     }
 
+    /// Only sound when the consumer merges a partition's locations on the same
+    /// key, so only the caller that plants the consumer's reader may set this.
+    pub fn with_merge_per_task(mut self, merge_per_task: bool) -> Self {
+        self.merge_per_task = merge_per_task;
+        self
+    }
+
+    /// See [`Self::with_merge_per_task`].
+    pub fn merge_per_task(&self) -> bool {
+        self.merge_per_task
+    }
+
     /// Task id (append-order slot within the stage) this writer instance
     /// is bound to.
     pub fn task_id(&self) -> usize {
@@ -688,6 +707,7 @@ impl ExecutionPlan for ShuffleWriterExec {
                     self.work_dir.clone(),
                 )?
                 .with_task_id(self.task_id)
+                .with_merge_per_task(self.merge_per_task)
                 .with_global_output_partition_ids(
                     self.global_output_partition_ids.clone(),
                 ),

ballista/core/src/extension.rs

@@ -245,6 +245,9 @@ pub trait SessionConfigExt {
     /// Is adaptive query planner enabled
     fn ballista_adaptive_query_planner_enabled(&self) -> bool;
 
+    /// Does a sorted passthrough shuffle stage merge each task's partitions
+    fn ballista_shuffle_merge_ordered_passthrough(&self) -> bool;
+
     /// Enables or disables adaptive query planning (enabled by default).
     fn with_ballista_adaptive_query_planner(self, enabled: bool) -> Self;
 
@@ -631,6 +634,16 @@ impl SessionConfigExt for SessionConfig {
             .unwrap_or_else(|| BallistaConfig::default().adaptive_query_planner_enabled())
     }
 
+    fn ballista_shuffle_merge_ordered_passthrough(&self) -> bool {
+        self.options()
+            .extensions
+            .get::<BallistaConfig>()
+            .map(|c| c.shuffle_merge_ordered_passthrough())
+            .unwrap_or_else(|| {
+                BallistaConfig::default().shuffle_merge_ordered_passthrough()
+            })
+    }
+
     fn with_ballista_grpc_metadata(self, metadata: HashMap<String, String>) -> Self {
         let extension = BallistaGrpcMetadataInterceptor::new(metadata);
         self.with_extension(Arc::new(extension))

ballista/scheduler/src/state/aqe/adapter.rs

@@ -22,8 +22,10 @@ use crate::state::aqe::execution_plan::{
 use crate::state::aqe::planner::AdaptiveStageInfo;
 use crate::state::execution_graph::StageOutput;
 use ballista_core::JobId;
+use ballista_core::config::BallistaConfig;
 use ballista_core::execution_plans::{
-    RangeFilterExec, RangeShuffleReaderExec, ShuffleReaderExec,
+    RangeFilterExec, RangeShuffleReaderExec, ShuffleReaderExec, ShuffleWriter,
+    ShuffleWriterExec,
 };
 use datafusion::common::exec_err;
 use datafusion::config::ConfigOptions;
@@ -40,6 +42,32 @@ use datafusion::{
 use std::collections::HashMap;
 use std::sync::Arc;
 
+/// Mark the stage's writer when its output will be read back by an
+/// ordering-preserving reader, so each task can merge its partitions before
+/// writing and leave the consumer one source per task.
+///
+/// Only place that knows both sides: `build_reader` plants
+/// `RangeShuffleReaderExec` for this same exchange under the conditions checked
+/// here. The static planner never calls it, so it is never marked.
+fn mark_merge_per_task(
+    writer: Arc<dyn ShuffleWriter>,
+    exchange: &ExchangeExec,
+    config: &ConfigOptions,
+) -> Arc<dyn ShuffleWriter> {
+    let ballista = config.extensions.get::<BallistaConfig>();
+    let enabled = ballista
+        .map(|c| c.shuffle_merge_ordered_passthrough() && !c.coalesce_enabled())
+        .unwrap_or(false);
+    if !enabled || exchange.broadcast || exchange.input().output_ordering().is_none() {
+        return writer;
+    }
+    let as_plan: &dyn ExecutionPlan = writer.as_ref();
+    match as_plan.downcast_ref::<ShuffleWriterExec>() {
+        Some(w) => Arc::new(w.clone().with_merge_per_task(true)),
+        None => writer,
+    }
+}
+
 #[derive(Debug, Clone, Default)]
 pub(crate) struct BallistaAdapter {
     inputs: HashMap<usize, StageOutput>,
@@ -198,6 +226,7 @@ impl BallistaAdapter {
                 config,
             )
             .map_err(|e| DataFusionError::External(Box::new(e)))?;
+            let writer = mark_merge_per_task(writer, root, config);
 
             Ok(AdaptiveStageInfo {
                 plan: writer,

ballista/scheduler/src/state/task_builder.rs

@@ -37,7 +37,7 @@
 //! subtrees never share state and there's no traversal-order dependency.
 
 use ballista_core::execution_plans::{
-    RangeFilterExec, RangeShuffleReaderExec, ShuffleReaderExec,
+    RangeFilterExec, RangeShuffleReaderExec, ShuffleReaderExec, ShuffleWriterExec,
 };
 use datafusion::common::internal_err;
 use datafusion::datasource::memory::MemorySourceConfig;
@@ -51,9 +51,10 @@ use datafusion::physical_plan::coalesce_partitions::CoalescePartitionsExec;
 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::SortExec;
 use datafusion::physical_plan::sorts::sort_preserving_merge::SortPreservingMergeExec;
 use datafusion::physical_plan::union::UnionExec;
-use datafusion::physical_plan::{ExecutionPlan, Partitioning};
+use datafusion::physical_plan::{ExecutionPlan, ExecutionPlanProperties, Partitioning};
 use log::warn;
 use std::any::Any;
 use std::sync::Arc;
@@ -70,6 +71,56 @@ pub fn restrict_plan_to_partitions(
     restrict(plan, partitions, /* under_collect */ false)
 }
 
+/// Merge a marked writer's sorted partitions into one before the write.
+///
+/// The task writes one file instead of one per partition, and reports it as
+/// output partition 0 — `file_id` already distinguishes producers, which is what
+/// `GlobalPartitionMap::Collapsed` handles. The mark comes from the AQE adapter,
+/// the only place that knows the consumer reads this back with an
+/// ordering-preserving merge.
+///
+/// Must run after [`restrict_plan_to_partitions`], which treats a
+/// `SortPreservingMergeExec` as a collapse and gives leaves below one the full
+/// upstream — merging first would make every task read the whole stage input.
+pub fn merge_task_partitions_before_write(
+    plan: Arc<dyn ExecutionPlan>,
+) -> Result<Arc<dyn ExecutionPlan>> {
+    if let Some(writer) = plan.downcast_ref::<ShuffleWriterExec>()
+        && writer.merge_per_task()
+        && let [input] = plan.children().as_slice()
+        && input.output_partitioning().partition_count() > 1
+        && let Some(ordering) = input.output_ordering()
+    {
+        let merge = Arc::new(
+            SortPreservingMergeExec::new(ordering.clone(), Arc::clone(input))
+                .with_fetch(top_k_fetch(input)),
+        );
+        return replace_children_if_necessary(plan, vec![merge]);
+    }
+    Ok(plan)
+}
+
+/// The row limit of the `SortExec` that established this ordering.
+///
+/// A per-partition `TopK(n)` only exists because a global limit of `n` sits
+/// above the merge that consumes this stage, so the task's merge can stop at `n`
+/// too: any row in the global top `n` that this task holds is in the task's own
+/// top `n`. Stops at the sort — a limit further down belongs to a different
+/// ordering.
+fn top_k_fetch(input: &Arc<dyn ExecutionPlan>) -> Option<usize> {
+    let mut node = Arc::clone(input);
+    loop {
+        if node.is::<SortExec>() {
+            return node.fetch();
+        }
+        let children = node.children();
+        let [child] = children.as_slice() else {
+            return None;
+        };
+        node = Arc::clone(child);
+    }
+}
+
 /// Recursive worker. `under_collect` is the scope inherited from ancestors:
 /// once set, every descendant leaf reads the full upstream. Scope is passed
 /// by value, so sibling subtrees never leak state to each other.
@@ -765,4 +816,94 @@ mod tests {
         assert_eq!(reader.partition[0][0].partition_id.partition_id, 1);
         assert_eq!(reader.partition[1][0].partition_id.partition_id, 3);
     }
+
+    // --- task-local ordered merge ---
+
+    /// A passthrough writer over a 3-partition input — the shape a top-N
+    /// query's aggregate stage has. `marked` is what the AQE adapter stamps.
+    fn passthrough_writer(sorted: bool, marked: bool) -> Arc<dyn ExecutionPlan> {
+        passthrough_writer_with_fetch(sorted, marked, None)
+    }
+
+    fn passthrough_writer_with_fetch(
+        sorted: bool,
+        marked: bool,
+        fetch: Option<usize>,
+    ) -> Arc<dyn ExecutionPlan> {
+        use datafusion::physical_expr::expressions::Column;
+        use datafusion::physical_expr::{LexOrdering, PhysicalSortExpr};
+        use datafusion::physical_plan::sorts::sort::SortExec;
+
+        let source = scan_with_file_groups(3);
+        let input: Arc<dyn ExecutionPlan> = if sorted {
+            let ordering = LexOrdering::new(vec![PhysicalSortExpr::new_default(
+                Arc::new(Column::new("a", 0)),
+            )])
+            .unwrap();
+            Arc::new(
+                SortExec::new(ordering, source)
+                    .with_preserve_partitioning(true)
+                    .with_fetch(fetch),
+            )
+        } else {
+            source
+        };
+        Arc::new(
+            ShuffleWriterExec::try_new("job".to_string().into(), 1, input, String::new())
+                .unwrap()
+                .with_merge_per_task(marked),
+        )
+    }
+
+    #[test]
+    fn task_local_merge_collapses_sorted_partitions() {
+        let plan = passthrough_writer(true, true);
+        assert_eq!(plan.output_partitioning().partition_count(), 3);
+
+        let out = merge_task_partitions_before_write(plan).unwrap();
+        let children = out.children();
+        assert!(
+            children[0]
+                .downcast_ref::<SortPreservingMergeExec>()
+                .is_some(),
+            "expected the merge under the writer, got {}",
+            children[0].name()
+        );
+        assert_eq!(
+            out.output_partitioning().partition_count(),
+            1,
+            "the writer now emits one partition per task"
+        );
+    }
+
+    /// Concatenating files that are not each sorted is not sorted, so an input
+    /// without an ordering has nothing to merge on.
+    #[test]
+    fn task_local_merge_skips_unordered_input() {
+        let out =
+            merge_task_partitions_before_write(passthrough_writer(false, true)).unwrap();
+        assert_eq!(out.output_partitioning().partition_count(), 3);
+    }
+
+    /// Unmarked is the statically planned case: that consumer concatenates a
+    /// partition's locations, so collapsing them would lose the ordering.
+    #[test]
+    fn task_local_merge_skips_an_unmarked_writer() {
+        let out =
+            merge_task_partitions_before_write(passthrough_writer(true, false)).unwrap();
+        assert_eq!(out.output_partitioning().partition_count(), 3);
+    }
+
+    /// The merge must inherit the per-partition `TopK`'s limit — without it the
+    /// task emits `partitions x fetch` rows for the consumer to re-merge.
+    #[test]
+    fn task_local_merge_takes_the_top_k_fetch() {
+        let plan = passthrough_writer_with_fetch(true, true, Some(20));
+        let out = merge_task_partitions_before_write(plan).unwrap();
+        let children = out.children();
+        let merge = children[0]
+            .downcast_ref::<SortPreservingMergeExec>()
+            .expect("expected the merge under the writer");
+        assert_eq!(merge.fetch(), Some(20));
+    }
 }