diff --git a/ballista/core/src/config.rs b/ballista/core/src/config.rs index 6f2a5fd12..c22db05a6 100644 --- a/ballista/core/src/config.rs +++ b/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> = 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 { diff --git a/ballista/core/src/execution_plans/shuffle_writer.rs b/ballista/core/src/execution_plans/shuffle_writer.rs index df3cdcab9..451e365b6 100644 --- a/ballista/core/src/execution_plans/shuffle_writer.rs +++ b/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, + /// 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(), ), diff --git a/ballista/core/src/extension.rs b/ballista/core/src/extension.rs index a4f19c1de..1642d0786 100644 --- a/ballista/core/src/extension.rs +++ b/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::() + .map(|c| c.shuffle_merge_ordered_passthrough()) + .unwrap_or_else(|| { + BallistaConfig::default().shuffle_merge_ordered_passthrough() + }) + } + fn with_ballista_grpc_metadata(self, metadata: HashMap) -> Self { let extension = BallistaGrpcMetadataInterceptor::new(metadata); self.with_extension(Arc::new(extension)) diff --git a/ballista/scheduler/src/state/aqe/adapter.rs b/ballista/scheduler/src/state/aqe/adapter.rs index 9e39c5ffb..78073ad50 100644 --- a/ballista/scheduler/src/state/aqe/adapter.rs +++ b/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, + exchange: &ExchangeExec, + config: &ConfigOptions, +) -> Arc { + let ballista = config.extensions.get::(); + 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::() { + Some(w) => Arc::new(w.clone().with_merge_per_task(true)), + None => writer, + } +} + #[derive(Debug, Clone, Default)] pub(crate) struct BallistaAdapter { inputs: HashMap, @@ -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, diff --git a/ballista/scheduler/src/state/task_builder.rs b/ballista/scheduler/src/state/task_builder.rs index a9ab5e502..4fab3bff4 100644 --- a/ballista/scheduler/src/state/task_builder.rs +++ b/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, +) -> Result> { + if let Some(writer) = plan.downcast_ref::() + && 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) -> Option { + let mut node = Arc::clone(input); + loop { + if node.is::() { + 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 { + passthrough_writer_with_fetch(sorted, marked, None) + } + + fn passthrough_writer_with_fetch( + sorted: bool, + marked: bool, + fetch: Option, + ) -> Arc { + 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 = 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::() + .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::() + .expect("expected the merge under the writer"); + assert_eq!(merge.fetch(), Some(20)); + } } diff --git a/ballista/scheduler/src/state/task_manager.rs b/ballista/scheduler/src/state/task_manager.rs index 60e533047..c22926cca 100644 --- a/ballista/scheduler/src/state/task_manager.rs +++ b/ballista/scheduler/src/state/task_manager.rs @@ -25,7 +25,9 @@ use crate::state::execution_graph::{ ExecutionGraphBox, RunningTaskInfo, StaticExecutionGraph, TaskDescription, }; use crate::state::executor_manager::ExecutorManager; -use crate::state::task_builder::restrict_plan_to_partitions; +use crate::state::task_builder::{ + merge_task_partitions_before_write, restrict_plan_to_partitions, +}; use ballista_core::error::BallistaError; use ballista_core::error::Result; use ballista_core::execution_plans::compute_global_output_partition_ids; @@ -177,6 +179,15 @@ pub struct UpdatedStages { pub resubmit_successful_stages: HashSet, } +/// The per-task plan rewrites applied before encoding: restrict the plan's +/// leaves to the task's partition slice, then merge the task's sorted +/// partitions before its writer. +fn rewrite_plan_for_task(task: &TaskDescription) -> Result> { + let restricted = + restrict_plan_to_partitions(task.plan.clone(), &task.global_input_partition_ids)?; + Ok(merge_task_partitions_before_write(restricted)?) +} + impl TaskManager { /// Creates a new `TaskManager` with the default task launcher. pub fn new( @@ -773,10 +784,7 @@ impl TaskManager let stage_id = task.key.stage_id; if self.active_job_cache.get(&job_id).is_some() { - let restricted = restrict_plan_to_partitions( - task.plan.clone(), - &task.global_input_partition_ids, - )?; + let restricted = rewrite_plan_for_task(&task)?; let mut plan_buf: Vec = vec![]; let plan_proto = PhysicalPlanNode::try_from_physical_plan( restricted, @@ -879,10 +887,7 @@ impl TaskManager let mut multi_tasks = Vec::with_capacity(tasks.len()); for task in tasks { - let restricted = restrict_plan_to_partitions( - task.plan.clone(), - &task.global_input_partition_ids, - )?; + let restricted = rewrite_plan_for_task(&task)?; let mut plan_buf: Vec = vec![]; let plan_proto = PhysicalPlanNode::try_from_physical_plan(restricted, codec)?; plan_proto.try_encode(&mut plan_buf)?; diff --git a/docs/source/user-guide/configs.md b/docs/source/user-guide/configs.md index 91b068356..9d8e8bc88 100644 --- a/docs/source/user-guide/configs.md +++ b/docs/source/user-guide/configs.md @@ -150,6 +150,7 @@ explanation of the sort-based shuffle writer. | ballista.shuffle.compression.codec | Utf8 | lz4 | Compression codec specification used in the shuffle process. Possible values: none, lz4, zstd. Defaults to lz4 to preserve current behaviour | | ballista.shuffle.force_remote_read | Boolean | false | Forces the shuffle reader to always read partitions via the Arrow Flight client, even when partitions are local to the node. | | ballista.shuffle.max_concurrent_read_requests | UInt64 | 64 | Maximum concurrent requests shuffle reader can process | +| ballista.shuffle.merge_ordered_passthrough | Boolean | true | 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. | | ballista.shuffle.reader.default_block_size_bytes | UInt64 | 1048576 | Assumed per-partition byte size charged to the shuffle governor when partition stats carry no byte count. | | ballista.shuffle.reader.max_blocks_in_flight_per_address | UInt64 | 128 | Reduce-side shuffle governor: maximum concurrent in-flight partition fetches to a single executor address. | | ballista.shuffle.reader.max_bytes_in_flight | UInt64 | 50331648 | Reduce-side shuffle governor: maximum total in-flight bytes across concurrent remote partition fetches. Mirrors Spark's spark.reducer.maxSizeInFlight. Values above 4 GiB are clamped to 4 GiB (u32 semaphore limit). |