diff --git a/ballista/client/tests/context_checks.rs b/ballista/client/tests/context_checks.rs index dd67cc1032..5d9cf83ea5 100644 --- a/ballista/client/tests/context_checks.rs +++ b/ballista/client/tests/context_checks.rs @@ -1192,7 +1192,7 @@ mod supported { "| | PlaceholderRowExec, metrics=[...] |", "| | |", "| | =========SuccessfulStage[stage_id=2, partitions=16]========= |", - "| | ShuffleWriterExec: partitioning: None, metrics=[output_rows=..., input_rows=..., repart_time=..., write_time=...] |", + "| | ShuffleWriterExec: partitioning: None, metrics=[output_rows=..., input_rows=..., write_time=...] |", "| | ProjectionExec: expr=[count(Int64(1))@1 as count(*), id@0 as id], metrics=[output_rows=..., elapsed_compute=..., output_bytes=..., output_batches=..., expr_0_eval_time=..., expr_1_eval_time=...] |", "| | AggregateExec: mode=FinalPartitioned, gby=[id@0 as id], aggr=[count(Int64(1))], metrics=[output_rows=..., elapsed_compute=..., output_bytes=..., output_batches=..., spill_count=..., spilled_bytes=..., spilled_rows=..., peak_mem_used=..., aggregate_arguments_time=..., aggregation_time=..., emitting_time=..., time_calculating_group_ids=...] |", "| | ShuffleReaderExec: upstream_stage: 1, partitioning: Hash([id@0], 16), metrics=[output_rows=..., elapsed_compute=..., output_bytes=..., output_batches=..., decoded_bytes=..., fetch_requests=..., fetch_retries=..., local_partitions=..., remote_partitions=..., fetch_time=..., local_read_time=..., permit_wait_time=...] |", diff --git a/ballista/client/tests/sort_shuffle.rs b/ballista/client/tests/sort_shuffle.rs index 0a9cc4bf72..fbeb7f052c 100644 --- a/ballista/client/tests/sort_shuffle.rs +++ b/ballista/client/tests/sort_shuffle.rs @@ -37,14 +37,12 @@ mod sort_shuffle_tests { BALLISTA_SHUFFLE_READER_MAX_BLOCKS_PER_ADDRESS, BALLISTA_SHUFFLE_READER_MAX_BYTES_IN_FLIGHT, BALLISTA_SHUFFLE_READER_REMOTE_PREFER_FLIGHT, - BALLISTA_SHUFFLE_SORT_BASED_ENABLED, }; use datafusion::arrow::util::pretty::pretty_format_batches; use datafusion::common::Result; use datafusion::execution::SessionStateBuilder; use datafusion::prelude::{ParquetReadOptions, SessionConfig, SessionContext}; use rstest::rstest; - use std::collections::HashSet; /// Read mode for shuffle data #[derive(Debug, Clone, Copy)] @@ -69,7 +67,6 @@ mod sort_shuffle_tests { aqe_enabled: bool, ) -> SessionContext { let mut config = SessionConfig::new_with_ballista() - .set_str(BALLISTA_SHUFFLE_SORT_BASED_ENABLED, "true") .set_bool(BALLISTA_ADAPTIVE_PLANNER_ENABLED, aqe_enabled); // Configure read mode @@ -103,7 +100,6 @@ mod sort_shuffle_tests { /// complete and return correct results. async fn create_tiny_budget_remote_context() -> SessionContext { let config = SessionConfig::new_with_ballista() - .set_str(BALLISTA_SHUFFLE_SORT_BASED_ENABLED, "true") .set_str(BALLISTA_SHUFFLE_READER_FORCE_REMOTE_READ, "true") .set_str(BALLISTA_SHUFFLE_READER_REMOTE_PREFER_FLIGHT, "true") .set_str(BALLISTA_SHUFFLE_READER_MAX_BYTES_IN_FLIGHT, "65536") @@ -117,17 +113,6 @@ mod sort_shuffle_tests { SessionContext::standalone_with_state(state).await.unwrap() } - /// Creates a standalone session context with hash-based shuffle. - async fn create_hash_shuffle_context() -> SessionContext { - let config = SessionConfig::new_with_ballista() - .set_str(BALLISTA_SHUFFLE_SORT_BASED_ENABLED, "false"); - let state = SessionStateBuilder::new() - .with_config(config) - .with_default_features() - .build(); - SessionContext::standalone_with_state(state).await.unwrap() - } - /// Registers test data in the context. async fn register_test_data(ctx: &SessionContext) { ctx.register_parquet( @@ -154,21 +139,6 @@ mod sort_shuffle_tests { ); } - /// Extracts values from a result set, ignoring order. - fn extract_values_unordered( - results: &[datafusion::arrow::record_batch::RecordBatch], - ) -> HashSet { - pretty_format_batches(results) - .unwrap() - .to_string() - .trim() - .lines() - .skip(3) // Skip header lines - .filter(|line| !line.starts_with('+')) - .map(|s| s.to_string()) - .collect() - } - // ==================== Basic Aggregation Tests ==================== #[rstest] @@ -327,62 +297,6 @@ mod sort_shuffle_tests { Ok(()) } - // ==================== Comparison with Hash Shuffle ==================== - - #[tokio::test] - async fn test_sort_vs_hash_shuffle_group_by() -> Result<()> { - // Test with sort shuffle (local read is sufficient for comparison) - let sort_ctx = create_sort_shuffle_context(ReadMode::Local).await; - register_test_data(&sort_ctx).await; - let sort_results = sort_ctx - .sql("SELECT bool_col, SUM(id) as total FROM test GROUP BY bool_col") - .await? - .collect() - .await?; - - // Test with hash shuffle - let hash_ctx = create_hash_shuffle_context().await; - register_test_data(&hash_ctx).await; - let hash_results = hash_ctx - .sql("SELECT bool_col, SUM(id) as total FROM test GROUP BY bool_col") - .await? - .collect() - .await?; - - // Results should be equivalent (order may differ) - let sort_values = extract_values_unordered(&sort_results); - let hash_values = extract_values_unordered(&hash_results); - assert_eq!(sort_values, hash_values); - Ok(()) - } - - #[tokio::test] - async fn test_sort_vs_hash_shuffle_distinct() -> Result<()> { - // Test with sort shuffle (local read is sufficient for comparison) - let sort_ctx = create_sort_shuffle_context(ReadMode::Local).await; - register_test_data(&sort_ctx).await; - let sort_results = sort_ctx - .sql("SELECT DISTINCT bool_col FROM test") - .await? - .collect() - .await?; - - // Test with hash shuffle - let hash_ctx = create_hash_shuffle_context().await; - register_test_data(&hash_ctx).await; - let hash_results = hash_ctx - .sql("SELECT DISTINCT bool_col FROM test") - .await? - .collect() - .await?; - - // Results should be equivalent - let sort_values = extract_values_unordered(&sort_results); - let hash_values = extract_values_unordered(&hash_results); - assert_eq!(sort_values, hash_values); - Ok(()) - } - // ==================== Edge Cases ==================== #[rstest] diff --git a/ballista/core/src/config.rs b/ballista/core/src/config.rs index 716597ff6b..533c1f2b2d 100644 --- a/ballista/core/src/config.rs +++ b/ballista/core/src/config.rs @@ -84,9 +84,6 @@ pub const BALLISTA_CLIENT_IO_RETRY_WAIT_TIME_MS: &str = "ballista.client.io_retry_wait_time_ms"; /// Enables adaptive query planning pub const BALLISTA_ADAPTIVE_PLANNER_ENABLED: &str = "ballista.planner.adaptive.enabled"; -/// Configuration key for enabling sort-based shuffle. -pub const BALLISTA_SHUFFLE_SORT_BASED_ENABLED: &str = - "ballista.shuffle.sort_based.enabled"; /// Configuration key for sort shuffle target batch size in rows. pub const BALLISTA_SHUFFLE_SORT_BASED_BATCH_SIZE: &str = "ballista.shuffle.sort_based.batch_size"; @@ -226,10 +223,6 @@ static CONFIG_ENTRIES: LazyLock> = LazyLock::new(|| "Enables Adaptive Query Planning (EXPERIMENTAL)".to_string(), DataType::Boolean, Some(false.to_string())), - ConfigEntry::new(BALLISTA_SHUFFLE_SORT_BASED_ENABLED.to_string(), - "Enable sort-based shuffle which writes consolidated files with index".to_string(), - DataType::Boolean, - Some(true.to_string())), ConfigEntry::new(BALLISTA_SHUFFLE_SORT_BASED_BATCH_SIZE.to_string(), "Target batch size in rows for coalescing small batches in sort shuffle".to_string(), DataType::UInt64, @@ -576,14 +569,6 @@ impl BallistaConfig { self.get_bool_setting(BALLISTA_ADAPTIVE_PLANNER_ENABLED) } - /// Returns whether sort-based shuffle is enabled. - /// - /// When enabled, shuffle writes produce a single consolidated file per input - /// partition with an index file, rather than one file per output partition. - pub fn shuffle_sort_based_enabled(&self) -> bool { - self.get_bool_setting(BALLISTA_SHUFFLE_SORT_BASED_ENABLED) - } - /// Returns the target batch size for sort-based shuffle. pub fn shuffle_sort_based_batch_size(&self) -> usize { self.get_usize_setting(BALLISTA_SHUFFLE_SORT_BASED_BATCH_SIZE) diff --git a/ballista/core/src/execution_plans/shuffle_reader.rs b/ballista/core/src/execution_plans/shuffle_reader.rs index c4d7b511b1..d85f1978fa 100644 --- a/ballista/core/src/execution_plans/shuffle_reader.rs +++ b/ballista/core/src/execution_plans/shuffle_reader.rs @@ -772,7 +772,6 @@ fn send_fetch_partitions( let max_blocks_per_addr = ballista_config.shuffle_reader_max_blocks_in_flight_per_address(); let default_block_size = ballista_config.shuffle_reader_default_block_size_bytes(); - let sort_shuffle_enabled = config.ballista_sort_shuffle_enabled(); let (response_sender, response_receiver) = mpsc::channel(max_reqs.max(1)); @@ -816,7 +815,7 @@ fn send_fetch_partitions( for p in local_locations { let r = { let _timer = local_read_time.timer(); - fetch_partition_local(&work_dir, &p, sort_shuffle_enabled) + fetch_partition_local(&work_dir, &p) }; if let Err(e) = response_sender_c.blocking_send(r) { error!("Fail to send response event to the channel due to {e}"); @@ -1112,7 +1111,6 @@ async fn fetch_partition_remote( fn fetch_partition_local( work_dir: &str, location: &PartitionLocation, - sort_shuffle_enabled: bool, ) -> result::Result { let path = &location.path(work_dir)?; let metadata = &location.executor_meta; @@ -1123,10 +1121,10 @@ fn fetch_partition_local( // replace this check with open, and check for error // // Check if this is a sort-based shuffle output (has index file) - if sort_shuffle_enabled && is_sort_shuffle_output(data_path) { - // note: in some cases sort shuffle is not going to be used - // even its enabled. thus we need to check if there is - // sort shuffle file index + if is_sort_shuffle_output(data_path) { + // A stage's on-disk layout is authoritative: sort-shuffle outputs have a + // companion index file. Standard single-partition outputs do not, so a + // missing index means this is a plain Arrow IPC file. debug!( "Reading sort-based shuffle for partition {} from {:?}", partition_id.partition_id, data_path @@ -1147,7 +1145,7 @@ fn fetch_partition_local( }); } debug!("fetch local partition file: {data_path:?} "); - // Standard hash-based shuffle - read the file directly + // Standard single-file shuffle output - read the file directly let reader = fetch_partition_local_inner(path).map_err(|e| { // return BallistaError::FetchFailed may let scheduler retry this task. BallistaError::FetchFailed( @@ -1290,7 +1288,6 @@ mod tests { use datafusion::common::DataFusionError; use datafusion::datasource::memory::MemorySourceConfig; use datafusion::datasource::source::DataSourceExec; - use datafusion::physical_expr::expressions::Column; use datafusion::physical_plan::common; use datafusion::prelude::SessionContext; use tempfile::{TempDir, tempdir}; @@ -1744,7 +1741,6 @@ mod tests { 1, create_test_data_plan().unwrap(), work_dir.path().to_str().unwrap().to_owned(), - Some(Partitioning::Hash(vec![Arc::new(Column::new("a", 0))], 1)), ) .unwrap(); @@ -1772,10 +1768,9 @@ mod tests { .map_err(|e| DataFusionError::Execution(format!("{e:?}"))) .unwrap(); - // 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); + // With single-partition (None) output, executing input partition 0 + // writes just that partition's 2 batches to a single output file. + assert_eq!(result.len(), 2); for b in result { assert_eq!(b, create_test_batch()) } diff --git a/ballista/core/src/execution_plans/shuffle_writer.rs b/ballista/core/src/execution_plans/shuffle_writer.rs index 56d3787406..0742dbf213 100644 --- a/ballista/core/src/execution_plans/shuffle_writer.rs +++ b/ballista/core/src/execution_plans/shuffle_writer.rs @@ -20,14 +20,9 @@ //! partition is re-partitioned and streamed to disk in Arrow IPC format. Future stages of the query //! will use the ShuffleReaderExec to read these results. -use datafusion::arrow::ipc::writer::StreamWriter; use std::fmt::Debug; -use std::fs; -use std::fs::File; use std::future::Future; -use std::io::BufWriter; use std::iter::Iterator; -use std::path::PathBuf; use std::sync::Arc; use std::time::Instant; @@ -39,7 +34,6 @@ use crate::utils; use crate::serde::protobuf::ShuffleWritePartition; use crate::serde::scheduler::PartitionStats; -use crate::utils::create_write_options; use datafusion::arrow::array::{ ArrayBuilder, ArrayRef, StringBuilder, StructBuilder, UInt32Builder, UInt64Builder, }; @@ -57,14 +51,14 @@ use datafusion::physical_plan::{ DisplayAs, DisplayFormatType, ExecutionPlan, Partitioning, PlanProperties, SendableRecordBatchStream, Statistics, }; -use futures::{StreamExt, TryStreamExt}; +use futures::TryStreamExt; use datafusion::arrow::error::ArrowError; use datafusion::execution::context::TaskContext; -use datafusion::physical_plan::repartition::{BatchPartitioner, RepartitionExec}; +use datafusion::physical_plan::repartition::RepartitionExec; use datafusion::physical_plan::sorts::sort_preserving_merge::SortPreservingMergeExec; use datafusion::physical_plan::stream::RecordBatchStreamAdapter; -use log::{debug, error}; +use log::debug; use std::sync::Mutex; use tokio::sync::oneshot; @@ -148,12 +142,11 @@ fn walk_child_partition_mapping( /// the wire. The executor-side writer then uses these ids directly rather /// than re-walking the plan. /// -/// Three cases: +/// Two cases: /// /// - `SortShuffleWriter(Hash(K))` — HashSpace by construction; the K-space /// `[0..K-1]` is intrinsic. Input ids are irrelevant. -/// - `ShuffleWriter(Hash(K)|RoundRobin(K))` — same K-space semantics. -/// - `ShuffleWriter(None)` — passthrough; the child's plan shape decides: +/// - `ShuffleWriter` — always passthrough; the child's plan shape decides: /// - `SortPreservingMergeExec` in the child chain → `[0]` (collapse). /// - `RepartitionExec(Hash|RoundRobin)` in the child chain → `[0..K-1]`. /// - Otherwise (leaf-preserved partitioning) → the input ids as-is. @@ -171,36 +164,19 @@ pub fn compute_global_output_partition_ids( }; return (0..*k).collect(); } - if let Some(w) = stage_plan.downcast_ref::() { - match w.shuffle_output_partitioning() { - Some(Partitioning::Hash(_, k)) | Some(Partitioning::RoundRobinBatch(k)) => { - return (0..*k).collect(); - } - None => { - let children = stage_plan.children(); - let [child] = children.as_slice() else { - unreachable!("ShuffleWriterExec always has exactly one child"); - }; - return match walk_child_partition_mapping( - child, - global_input_partition_ids, - ) { - GlobalPartitionMap::Collapsed => vec![0], - GlobalPartitionMap::HashSpace => { - let k = - child.properties().output_partitioning().partition_count(); - (0..k).collect() - } - GlobalPartitionMap::PassThrough(ids) => ids, - }; - } - Some(Partitioning::UnknownPartitioning(_)) => { - panic!( - "ShuffleWriterExec::UnknownPartitioning has no meaningful output-id \ - space; upstream planner should never emit this shape" - ); + if stage_plan.downcast_ref::().is_some() { + let children = stage_plan.children(); + let [child] = children.as_slice() else { + unreachable!("ShuffleWriterExec always has exactly one child"); + }; + return match walk_child_partition_mapping(child, global_input_partition_ids) { + GlobalPartitionMap::Collapsed => vec![0], + GlobalPartitionMap::HashSpace => { + let k = child.properties().output_partitioning().partition_count(); + (0..k).collect() } - } + GlobalPartitionMap::PassThrough(ids) => ids, + }; } global_input_partition_ids.to_vec() } @@ -240,13 +216,13 @@ impl Debug for WriterState { /// partition is re-partitioned and streamed to disk in Arrow IPC format. Future stages of the query /// will use the ShuffleReaderExec to read these results. /// -/// # Threading (passthrough / `None` branch) +/// # Threading /// /// One Ballista task holds one `ShuffleWriterExec`. The first `execute(k)` /// call from the executor initializes K oneshot handoffs and spawns -/// `run_coordinator`, which invokes `execute_shuffle_write`. In the -/// passthrough branch (`shuffle_output_partitioning == None`, i.e. the -/// child already has the target K partitions), `execute_shuffle_write` +/// `run_coordinator`, which invokes `execute_shuffle_write`. The writer only +/// ever passes its child's partitioning through (the child already has the +/// target K partitions), so `execute_shuffle_write` /// spawns K concurrent tokio tasks — one per output partition — each /// pulling `child.execute(k)` and streaming directly to /// `data-{task_id}.arrow` for partition k. Each drain emits one summary; @@ -263,8 +239,8 @@ impl Debug for WriterState { /// deadlock the scatter side. /// /// Data flows bottom → top (child produces, executor consumes), matching -/// DataFusion's convention. Under passthrough M = K (child preserves -/// partition count). Example: K=3. +/// DataFusion's convention. M = K (the child preserves partition count). +/// Example: K=3. /// /// ```text /// K=3 output partitions (pulled by executor) @@ -315,18 +291,16 @@ impl Debug for WriterState { /// what differs, and downstream `ShuffleReaderExec` uses the summaries to /// open whichever set of files is right. /// -/// The `Hash` branch (`shuffle_output_partitioning == Some(Hash)`) is a -/// legacy shape kept for correctness while the planner transitions to -/// `RepartitionExec(Hash) → ShuffleWriterExec(None)`; not diagrammed here. +/// This writer never repartitions: hash-repartition stages use +/// [`SortShuffleWriterExec`] instead, so the only scheme left here is +/// passthrough and the writer carries no output-partitioning of its own. /// /// The coordinator + oneshot plumbing is a shared idiom with -/// [`SortShuffleWriterExec`] and -/// with the `Hash` branch; passthrough alone doesn't structurally need it — -/// only the K concurrent drains, which are the real deadlock guard. Once -/// `Hash` is retired in favor of DataFusion's `RepartitionExec(Hash)`, this -/// branch could collapse to per-`execute(k)` eager K-spawn with -/// `JoinHandle` handoff, keeping the coordinator idiom only where it does -/// real work (SortShuffle's M×K summary re-bucketing). +/// [`SortShuffleWriterExec`]; passthrough alone doesn't structurally need +/// it — only the K concurrent drains, which are the real deadlock guard. +/// This could collapse to per-`execute(k)` eager K-spawn with `JoinHandle` +/// handoff, keeping the coordinator idiom only where it does real work +/// (SortShuffle's M×K summary re-bucketing). #[derive(Debug)] pub struct ShuffleWriterExec { /// Unique ID for the job (query) that this stage is a part of @@ -337,9 +311,6 @@ pub struct ShuffleWriterExec { plan: Arc, /// Path to write output streams to work_dir: String, - /// Optional shuffle output partitioning. - /// If it's none, it means there's no need to do repartitioning. - shuffle_output_partitioning: Option, /// Task id (the task's append-order slot in `RunningStage.task_infos`) /// used as `file_id` in shuffle paths so files from different tasks /// (including retries) don't collide. Seeded by the executor's @@ -351,11 +322,10 @@ pub struct ShuffleWriterExec { /// `global_output_partition_ids[i]` globally. /// /// Consumed by the passthrough (`None`) branch when the plan is a straight - /// pass-through of the slice. The `Hash` branch ignores it — the hash - /// K-space is its own global identity — and even for `None`, if the child - /// plan contains a partitioning-resetting operator (SPM → 1 output, - /// RepartitionExec::Hash → 0..K) the writer detects that at path-build - /// time and uses `local` directly instead of `global_output_partition_ids[local]`. + /// pass-through of the slice. If the child plan contains a + /// partitioning-resetting operator (SPM → 1 output, RepartitionExec::Hash + /// → 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, /// Execution metrics metrics: ExecutionPlanMetricsSet, @@ -374,7 +344,6 @@ impl Clone for ShuffleWriterExec { stage_id: self.stage_id, plan: self.plan.clone(), work_dir: self.work_dir.clone(), - shuffle_output_partitioning: self.shuffle_output_partitioning.clone(), task_id: self.task_id, global_output_partition_ids: self.global_output_partition_ids.clone(), metrics: self.metrics.clone(), @@ -391,28 +360,16 @@ impl std::fmt::Display for ShuffleWriterExec { .indent(false); write!( f, - "ShuffleWriterExec: job={} stage={} work_dir={} partitioning={:?} plan: \n {}", - self.job_id, - self.stage_id, - self.work_dir, - self.shuffle_output_partitioning, - printable_plan + "ShuffleWriterExec: job={} stage={} work_dir={} plan: \n {}", + self.job_id, self.stage_id, self.work_dir, printable_plan ) } } -pub struct WriteTracker { - pub num_batches: usize, - pub num_rows: usize, - pub writer: StreamWriter>, - pub path: PathBuf, -} - #[derive(Debug, Clone)] struct ShuffleWriteMetrics { /// Time spend writing batches to shuffle files write_time: metrics::Time, - repart_time: metrics::Time, input_rows: metrics::Count, output_rows: metrics::Count, } @@ -428,8 +385,6 @@ impl ShuffleWriteMetrics { fn new(input_partition: usize, metrics: &ExecutionPlanMetricsSet) -> Self { let write_time = MetricBuilder::new(metrics).subset_time("write_time", input_partition); - let repart_time = - MetricBuilder::new(metrics).subset_time("repart_time", input_partition); let input_rows = MetricBuilder::new(metrics).counter("input_rows", input_partition); @@ -438,7 +393,6 @@ impl ShuffleWriteMetrics { Self { write_time, - repart_time, input_rows, output_rows, } @@ -454,13 +408,10 @@ impl ShuffleWriterExec { stage_id: usize, plan: Arc, work_dir: String, - shuffle_output_partitioning: Option, ) -> Result { - // If [`shuffle_output_partitioning`] is none, then there's no need to do repartitioning. - // Therefore, the partition is the same as its input plan's. - let partitioning = shuffle_output_partitioning - .clone() - .unwrap_or_else(|| plan.properties().output_partitioning().clone()); + // This writer never repartitions, so its output partitioning is + // exactly its input plan's. + let partitioning = plan.properties().output_partitioning().clone(); let output_partition_count = partitioning.partition_count(); let properties = Arc::new(PlanProperties::new( datafusion::physical_expr::EquivalenceProperties::new(plan.schema()), @@ -473,15 +424,12 @@ impl ShuffleWriterExec { // unit tests that don't set a slice. Once every construction path // calls `with_global_output_partition_ids`, drop the default and fold it into // `try_new`'s signature. - let child_partition_count = - plan.properties().output_partitioning().partition_count(); - let default_partition_slice: Vec = (0..child_partition_count).collect(); + let default_partition_slice: Vec = (0..output_partition_count).collect(); Ok(Self { job_id, stage_id, plan, work_dir, - shuffle_output_partitioning, task_id: 0, global_output_partition_ids: default_partition_slice, metrics: ExecutionPlanMetricsSet::new(), @@ -539,28 +487,20 @@ impl ShuffleWriterExec { .partition_count() } - /// Get the true output partitioning - pub fn shuffle_output_partitioning(&self) -> Option<&Partitioning> { - self.shuffle_output_partitioning.as_ref() - } - /// Executes the shuffle write operation for this task. /// /// Returns `(handoff_idx, summary)` pairs. `handoff_idx` indexes into the /// coordinator's per-output-partition oneshot slots (0..K in this /// operator's output_partitioning). `summary.partition_id` is the - /// **global** output partition id downstream will address: - /// - /// - Passthrough (`None`): computed via `walk_child_partition_mapping` - /// over the child plan — either `global_output_partition_ids[local]`, `local` (hash - /// K-space), or `0` (collapsed / SPM). - /// - Hash: hash bucket space is global; global = local. + /// **global** output partition id downstream will address, computed via + /// `walk_child_partition_mapping` over the child plan — either + /// `global_output_partition_ids[local]`, `local` (hash K-space), or `0` + /// (collapsed / SPM). pub fn execute_shuffle_write( self, context: Arc, ) -> impl Future>> { let task_id = self.task_id; - let output_partitioning = self.shuffle_output_partitioning.clone(); let plan = self.plan.clone(); let partition_map = walk_child_partition_mapping(&plan, &self.global_output_partition_ids); @@ -572,272 +512,83 @@ impl ShuffleWriterExec { let compression_type = config.shuffle_compression_codec()?; let channel_capacity = config.shuffle_writer_channel_capacity(); - match output_partitioning { - None => { - // Passthrough shuffle: drain each of the child's output - // partitions into its own file. All K must drain - // CONCURRENTLY, not sequentially — coordinating operators - // below (DynamicRangeRepartitionExec) push to all K - // senders from shared scatter tasks; draining one to EOF - // before the next starts fills up the undrained channel - // and deadlocks the scatter side. - let num_partitions = - plan.properties().output_partitioning().partition_count(); - let mut handles = Vec::with_capacity(num_partitions); - for local_input_partition in 0..num_partitions { - // Each drain owns its own metric bucket, keyed by the - // operator-local input partition it drains. Passthrough - // is 1:1 so local input == local output here. - let write_metrics = - ShuffleWriteMetrics::new(local_input_partition, &metrics); - let global_partition = - partition_map.resolve(local_input_partition) as usize; - let path = create_shuffle_path( - &self.work_dir, - &self.job_id, - self.stage_id, - global_partition, - Some(task_id as u64), - false, - )?; - - if let Some(parent) = path.parent() { - std::fs::create_dir_all(parent)?; - } - - debug!("Writing results to {path:?}"); - - let mut stream = - plan.execute(local_input_partition, context.clone())?; - handles.push(tokio::spawn(async move { - let stats = utils::write_stream_to_disk( - &mut stream, - path.as_path(), - &write_metrics.write_time, - channel_capacity, - compression_type, - ) - .await - .map_err(|e| DataFusionError::Execution(format!("{e:?}")))?; - let rows = stats.num_rows.unwrap_or(0) as usize; - write_metrics.input_rows.add(rows); - write_metrics.output_rows.add(rows); - Ok::<_, DataFusionError>((local_input_partition, stats)) - })); - } - - let mut results = Vec::with_capacity(num_partitions); - for handle in handles { - let (local_input_partition, stats) = - handle.await.map_err(|e| { - DataFusionError::Execution(format!( - "shuffle-write drain task panicked: {e}" - )) - })??; - results.push(( - local_input_partition, - ShuffleWritePartition { - partition_id: partition_map - .resolve(local_input_partition), - num_batches: stats.num_batches.unwrap_or(0), - num_rows: stats.num_rows.unwrap_or(0), - num_bytes: stats.num_bytes.unwrap_or(0), - file_id: Some(task_id as u64), - is_sort_shuffle: false, - }, - )); - } - debug!( - "task_id {} drained {} partitions in {}s", - task_id, - num_partitions, - now.elapsed().as_secs() - ); - Ok(results) + // Passthrough shuffle: drain each of the child's output + // partitions into its own file. All K must drain + // CONCURRENTLY, not sequentially — coordinating operators + // below (DynamicRangeRepartitionExec) push to all K + // senders from shared scatter tasks; draining one to EOF + // before the next starts fills up the undrained channel + // and deadlocks the scatter side. + let num_partitions = + plan.properties().output_partitioning().partition_count(); + let mut handles = Vec::with_capacity(num_partitions); + for local_input_partition in 0..num_partitions { + // Each drain owns its own metric bucket, keyed by the + // operator-local input partition it drains. Passthrough + // is 1:1 so local input == local output here. + let write_metrics = + ShuffleWriteMetrics::new(local_input_partition, &metrics); + let global_partition = + partition_map.resolve(local_input_partition) as usize; + let path = create_shuffle_path( + &self.work_dir, + &self.job_id, + self.stage_id, + global_partition, + Some(task_id as u64), + false, + )?; + + if let Some(parent) = path.parent() { + std::fs::create_dir_all(parent)?; } - // TODO: just use Datafusion's hash partition operator and remove this branch - Some(Partitioning::Hash(exprs, num_output_partitions)) => { - // Task drains ALL of the child's (already slice-restricted) - // input partitions and routes rows into K writers by hash. - // file_id = task_id so files from different tasks - // (including retries) don't collide. Hash buckets 0..K are - // a global K-space and aren't further sliced. - // - // One metric bucket per operator-local input partition, - // matching the pre-K-drain model (before: N tasks × 1 - // bucket; now: 1 task × N buckets). One `BatchPartitioner` - // per input partition so `repart_time` disaggregates the - // same way — total partitioner count across the stage is - // unchanged. - let num_input_partitions = - plan.properties().output_partitioning().partition_count(); - let write_metrics_per_input: Vec = (0 - ..num_input_partitions) - .map(|i| ShuffleWriteMetrics::new(i, &metrics)) - .collect(); - let schema = plan.schema(); - let (tx, mut rx) = tokio::sync::mpsc::channel::<(usize, RecordBatch)>( + + debug!("Writing results to {path:?}"); + + let mut stream = plan.execute(local_input_partition, context.clone())?; + handles.push(tokio::spawn(async move { + let stats = utils::write_stream_to_disk( + &mut stream, + path.as_path(), + &write_metrics.write_time, channel_capacity, - ); - let write_times: Vec = write_metrics_per_input - .iter() - .map(|m| m.write_time.clone()) - .collect(); - let output_rows_per_input: Vec = - write_metrics_per_input - .iter() - .map(|m| m.output_rows.clone()) - .collect(); - let partitioners: Vec = write_metrics_per_input - .iter() - .map(|m| { - BatchPartitioner::new_hash_partitioner( - exprs.clone(), - num_output_partitions, - m.repart_time.clone(), - ) - }) - .collect::, _>>()?; - let work_dir = self.work_dir.clone(); - let job_id = self.job_id.clone(); - let stage_id = self.stage_id; - - let handle = tokio::task::spawn_blocking(move || { - let mut writers: Vec> = - (0..num_output_partitions).map(|_| None).collect(); - let mut partitioners = partitioners; - - while let Some((local_input_partition, input_batch)) = - rx.blocking_recv() - { - let write_time = &write_times[local_input_partition]; - let output_rows = - &output_rows_per_input[local_input_partition]; - partitioners[local_input_partition].partition( - input_batch, - |output_partition, output_batch| { - let timer = write_time.timer(); - match &mut writers[output_partition] { - Some(w) => { - w.num_batches += 1; - w.num_rows += output_batch.num_rows(); - w.writer.write(&output_batch)?; - } - None => { - let p = create_shuffle_path( - &work_dir, - &job_id, - stage_id, - output_partition, - Some(task_id as u64), - false, - )?; - - if let Some(parent) = p.parent() { - std::fs::create_dir_all(parent)?; - } - - debug!("Writing results to {p:?}"); - - let options = - create_write_options(compression_type)?; - - let file = - BufWriter::new(File::create(p.clone())?); - let mut writer = - StreamWriter::try_new_with_options( - file, - schema.as_ref(), - options, - )?; - writer.write(&output_batch)?; - writers[output_partition] = - Some(WriteTracker { - num_batches: 1, - num_rows: output_batch.num_rows(), - writer, - path: p, - }); - } - } - output_rows.add(output_batch.num_rows()); - timer.done(); - Ok(()) - }, - )?; - } - - let mut part_locs = vec![]; - for (i, w) in writers.iter_mut().enumerate() { - if let Some(w) = w { - w.writer.finish()?; - let num_bytes = fs::metadata(&w.path)?.len(); - debug!( - "Finished writing shuffle partition {} at {:?}. Batches: {}. Rows: {}. Bytes: {}.", - i, w.path, w.num_batches, w.num_rows, num_bytes - ); - part_locs.push(( - i, - ShuffleWritePartition { - partition_id: i as u64, - num_batches: w.num_batches as u64, - num_rows: w.num_rows as u64, - num_bytes, - file_id: Some(task_id as u64), - is_sort_shuffle: false, - }, - )); - } - } - Ok(part_locs) - }); - - let mut stream_err = None; - 'outer: for (local_input_partition, per_input_metrics) in - write_metrics_per_input.iter().enumerate() - { - let mut stream = - plan.execute(local_input_partition, context.clone())?; - loop { - match stream.next().await { - Some(Ok(batch)) => { - per_input_metrics.input_rows.add(batch.num_rows()); - if tx - .send((local_input_partition, batch)) - .await - .is_err() - { - break 'outer; - } - } - Some(Err(e)) => { - stream_err = Some(e); - break 'outer; - } - None => break, - } - } - } - drop(tx); - - let write_result = handle.await.map_err(|e| { - DataFusionError::Execution(format!( - "Shuffle writer task failed: {e}" - )) - })?; - if let Some(e) = stream_err { - if let Err(write_err) = &write_result { - error!("Shuffle writer also failed: {write_err}"); - } - return Err(e); - } - write_result - } + compression_type, + ) + .await + .map_err(|e| DataFusionError::Execution(format!("{e:?}")))?; + let rows = stats.num_rows.unwrap_or(0) as usize; + write_metrics.input_rows.add(rows); + write_metrics.output_rows.add(rows); + Ok::<_, DataFusionError>((local_input_partition, stats)) + })); + } - _ => Err(DataFusionError::Execution( - "Invalid shuffle partitioning scheme".to_owned(), - )), + let mut results = Vec::with_capacity(num_partitions); + for handle in handles { + let (local_input_partition, stats) = handle.await.map_err(|e| { + DataFusionError::Execution(format!( + "shuffle-write drain task panicked: {e}" + )) + })??; + results.push(( + local_input_partition, + ShuffleWritePartition { + partition_id: partition_map.resolve(local_input_partition), + num_batches: stats.num_batches.unwrap_or(0), + num_rows: stats.num_rows.unwrap_or(0), + num_bytes: stats.num_bytes.unwrap_or(0), + file_id: Some(task_id as u64), + is_sort_shuffle: false, + }, + )); } + debug!( + "task_id {} drained {} partitions in {}s", + task_id, + num_partitions, + now.elapsed().as_secs() + ); + Ok(results) } } } @@ -849,25 +600,13 @@ impl DisplayAs for ShuffleWriterExec { f: &mut std::fmt::Formatter, ) -> std::fmt::Result { match t { + // "None" is retained for plan-shape stability: this writer never + // repartitions, so the value can only ever be None. DisplayFormatType::Default | DisplayFormatType::Verbose => { - write!( - f, - "ShuffleWriterExec: partitioning: {}", - self.shuffle_output_partitioning - .as_ref() - .map(|p| p.to_string()) - .unwrap_or("None".to_string()) - ) + write!(f, "ShuffleWriterExec: partitioning: None") } DisplayFormatType::TreeRender => { - write!( - f, - "partitioning={}", - self.shuffle_output_partitioning - .as_ref() - .map(|p| p.to_string()) - .unwrap_or("None".to_string()) - ) + write!(f, "partitioning=None") } } } @@ -907,7 +646,6 @@ impl ExecutionPlan for ShuffleWriterExec { self.stage_id, input, self.work_dir.clone(), - self.shuffle_output_partitioning.clone(), )? .with_task_id(self.task_id) .with_global_output_partition_ids( @@ -1020,8 +758,9 @@ impl ShuffleWriter for ShuffleWriterExec { self.stage_id } + /// Always `None`: this writer preserves its input partitioning. fn shuffle_output_partitioning(&self) -> Option<&Partitioning> { - self.shuffle_output_partitioning.as_ref() + None } fn input_partition_count(&self) -> usize { @@ -1217,25 +956,23 @@ mod tests { } #[tokio::test] - // number of rows in each partition is a function of the hash output, so don't test here - #[cfg(not(feature = "force_hash_collisions"))] async fn test() -> Result<()> { let session_ctx = SessionContext::new(); let task_ctx = session_ctx.task_ctx(); - let input_plan = Arc::new(CoalescePartitionsExec::new(create_input_plan()?)); + // No output partitioning: passthrough writer, one file per one of + // the input plan's 2 partitions. + let input_plan = create_input_plan()?; let work_dir = TempDir::new()?; let query_stage = Arc::new(ShuffleWriterExec::try_new( JobId::new("jobOne"), 1, input_plan, work_dir.path().to_str().unwrap().to_owned(), - Some(Partitioning::Hash(vec![Arc::new(Column::new("a", 0))], 2)), )?); let batches = drive_all_partitions(query_stage, task_ctx).await?; // K=2 output partitions -> one metadata batch per execute(N) call - // (drive_all_partitions collects one per K). Empty output slots - // now emit 0-row batches (no summary), non-empty ones emit 1-row. + // (drive_all_partitions collects one per K). assert_eq!(2, batches.len()); for batch in &batches { assert_eq!(4, batch.num_columns()); @@ -1275,7 +1012,6 @@ mod tests { } #[tokio::test] - #[cfg(not(feature = "force_hash_collisions"))] async fn display_renders_child_operator_metrics() -> Result<()> { let session_ctx = SessionContext::new(); let task_ctx = session_ctx.task_ctx(); @@ -1287,7 +1023,6 @@ mod tests { 1, input_plan, work_dir.path().to_str().unwrap().to_owned(), - Some(Partitioning::Hash(vec![Arc::new(Column::new("a", 0))], 2)), )?; let mut stream = query_stage.execute(0, task_ctx)?; let _ = utils::collect_stream(&mut stream) @@ -1306,45 +1041,6 @@ mod tests { Ok(()) } - #[tokio::test] - // number of rows in each partition is a function of the hash output, so don't test here - #[cfg(not(feature = "force_hash_collisions"))] - async fn test_partitioned() -> Result<()> { - let session_ctx = SessionContext::new(); - let task_ctx = session_ctx.task_ctx(); - - let input_plan = create_input_plan()?; - let work_dir = TempDir::new()?; - let query_stage = Arc::new(ShuffleWriterExec::try_new( - JobId::new("jobOne"), - 1, - input_plan, - work_dir.path().to_str().unwrap().to_owned(), - Some(Partitioning::Hash(vec![Arc::new(Column::new("a", 0))], 2)), - )?); - let batches = drive_all_partitions(query_stage, task_ctx).await?; - assert_eq!(2, batches.len()); - let total: u64 = batches - .iter() - .flat_map(|b| { - let stats = b.column(3).as_any().downcast_ref::().unwrap(); - let num_rows = stats - .column_by_name("num_rows") - .unwrap() - .as_any() - .downcast_ref::() - .unwrap() - .clone(); - (0..b.num_rows()).map(move |i| num_rows.value(i)) - }) - .sum(); - // Row conservation across the K output partitions: writer reads its - // full input slice (2 partitions × 2 batches × 2 rows = 8). - assert_eq!(8, total); - - Ok(()) - } - #[tokio::test] async fn test_no_repart_write_failure_propagates() -> Result<()> { let session_ctx = SessionContext::new(); @@ -1370,7 +1066,6 @@ mod tests { 1, input_plan, work_dir, - None, )?); let result = drive_all_partitions(query_stage, task_ctx).await; assert!( @@ -1381,13 +1076,13 @@ mod tests { } #[tokio::test] - async fn test_hash_repart_write_failure_propagates() -> Result<()> { + async fn test_create_dir_failure_propagates() -> Result<()> { let session_ctx = SessionContext::new(); let task_ctx = session_ctx.task_ctx(); // Place a regular file at the stage_id path component so - // create_dir_all fails when the writer tries to create - // output partition subdirectories underneath it. + // create_dir_all fails when the writer tries to create the + // output partition subdirectory underneath it. // Path structure: work_dir / job_id / stage_id / ... let tmp = tempfile::TempDir::new().unwrap(); let job_dir = tmp.path().join("jobOne"); @@ -1402,12 +1097,11 @@ mod tests { 1, input_plan, work_dir, - Some(Partitioning::Hash(vec![Arc::new(Column::new("a", 0))], 2)), )?); let result = drive_all_partitions(query_stage, task_ctx).await; assert!( result.is_err(), - "expected create_dir_all failure in hash writer to propagate" + "expected create_dir_all failure in writer to propagate" ); Ok(()) } diff --git a/ballista/core/src/extension.rs b/ballista/core/src/extension.rs index b0a58289b9..ff78e7c0a8 100644 --- a/ballista/core/src/extension.rs +++ b/ballista/core/src/extension.rs @@ -257,9 +257,6 @@ pub trait SessionConfigExt { /// Get whether to use TLS for executor connections fn ballista_use_tls(&self) -> bool; - /// Is short shuffle used - fn ballista_sort_shuffle_enabled(&self) -> bool; - /// Returns whether the AQE coalesce-shuffle-partitions rule is enabled. fn ballista_coalesce_enabled(&self) -> bool; /// Sets whether the AQE coalesce-shuffle-partitions rule is enabled. @@ -523,14 +520,6 @@ impl SessionConfigExt for SessionConfig { }) } - fn ballista_sort_shuffle_enabled(&self) -> bool { - self.options() - .extensions - .get::() - .map(|c| c.shuffle_sort_based_enabled()) - .unwrap_or_else(|| BallistaConfig::default().shuffle_sort_based_enabled()) - } - fn with_ballista_shuffle_reader_maximum_concurrent_requests( self, max_requests: usize, diff --git a/ballista/core/src/serde/mod.rs b/ballista/core/src/serde/mod.rs index 213a600a78..e3f4365dae 100644 --- a/ballista/core/src/serde/mod.rs +++ b/ballista/core/src/serde/mod.rs @@ -387,19 +387,22 @@ impl PhysicalExtensionCodec for BallistaPhysicalExtensionCodec { PhysicalPlanType::ShuffleWriter(shuffle_writer) => { let input = inputs[0].clone(); - let shuffle_output_partitioning = parse_protobuf_hash_partitioning( - shuffle_writer.output_partitioning.as_ref(), - &decode_ctx, - input.schema().as_ref(), - &converter, - )?; + // ShuffleWriterExec never repartitions. A plan that still + // carries an output partitioning is a legacy hash-shuffle + // plan, which this writer no longer implements. + if shuffle_writer.output_partitioning.is_some() { + return Err(DataFusionError::Internal( + "hash-partitioned ShuffleWriterExec is no longer supported; \ + hash-repartition stages use SortShuffleWriterExec" + .to_string(), + )); + } Ok(Arc::new(ShuffleWriterExec::try_new( shuffle_writer.job_id.clone().into(), shuffle_writer.stage_id as usize, input, "".to_string(), // this is intentional but hacky - the executor will fill this in - shuffle_output_partitioning, )?)) } PhysicalPlanType::SortShuffleWriter(sort_shuffle_writer) => { @@ -559,33 +562,15 @@ impl PhysicalExtensionCodec for BallistaPhysicalExtensionCodec { buf: &mut Vec, ) -> Result<(), DataFusionError> { if let Some(exec) = node.downcast_ref::() { - // note that we use shuffle_output_partitioning() rather than output_partitioning() - // to get the true output partitioning - let output_partitioning = match exec.shuffle_output_partitioning() { - Some(Partitioning::Hash(exprs, partition_count)) => { - Some(datafusion_proto::protobuf::PhysicalHashRepartition { - hash_expr: exprs - .iter() - .map(|expr|datafusion_proto::physical_plan::to_proto::serialize_physical_expr(&expr.clone(), self.default_codec.as_ref())) - .collect::, DataFusionError>>()?, - partition_count: *partition_count as u64, - }) - } - None => None, - other => { - return Err(DataFusionError::Internal(format!( - "physical_plan::to_proto() invalid partitioning for ShuffleWriterExec: {other:?}" - ))); - } - }; - let proto = protobuf::BallistaPhysicalPlanNode { physical_plan_type: Some(PhysicalPlanType::ShuffleWriter( protobuf::ShuffleWriterExecNode { job_id: exec.job_id().to_string(), stage_id: exec.stage_id() as u32, input: None, - output_partitioning, + // This writer preserves its input partitioning, so + // there is never a shuffle partitioning to encode. + output_partitioning: None, }, )), }; diff --git a/ballista/executor/src/execution_engine.rs b/ballista/executor/src/execution_engine.rs index 5633a6ddeb..cc55c7aa1a 100644 --- a/ballista/executor/src/execution_engine.rs +++ b/ballista/executor/src/execution_engine.rs @@ -149,18 +149,17 @@ impl ExecutionEngine for DefaultExecutionEngine { // the query plan created by the scheduler always starts with a shuffle writer // (either ShuffleWriterExec or SortShuffleWriterExec) - if let Some(shuffle_writer) = plan.downcast_ref::() { + if plan.downcast_ref::().is_some() { let exec = ShuffleWriterExec::try_new( job_id, stage_id, plan.children()[0].clone(), work_dir.to_string(), - shuffle_writer.shuffle_output_partitioning().cloned(), )? .with_task_id(task_id) .with_global_output_partition_ids(global_output_partition_ids); Ok(Arc::new(DefaultQueryStageExec::new( - ShuffleWriterVariant::Hash(exec), + ShuffleWriterVariant::Passthrough(exec), ))) } else if let Some(sort_shuffle_writer) = plan.downcast_ref::() @@ -190,8 +189,9 @@ impl ExecutionEngine for DefaultExecutionEngine { /// Enum representing the different shuffle writer implementations. #[derive(Debug, Clone)] pub enum ShuffleWriterVariant { - /// Hash-based shuffle writer (original implementation). - Hash(ShuffleWriterExec), + /// Passthrough shuffle writer: preserves its input partitioning, + /// one file per output partition. + Passthrough(ShuffleWriterExec), /// Sort-based shuffle writer. Sort(SortShuffleWriterExec), } @@ -216,7 +216,7 @@ impl DefaultQueryStageExec { impl Display for DefaultQueryStageExec { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { match &self.shuffle_writer { - ShuffleWriterVariant::Hash(writer) => { + ShuffleWriterVariant::Passthrough(writer) => { let stage_metrics: Vec = writer .metrics() .unwrap_or_default() @@ -225,7 +225,7 @@ impl Display for DefaultQueryStageExec { .collect(); write!( f, - "DefaultQueryStageExec(Hash): ({})\n{}", + "DefaultQueryStageExec(Passthrough): ({})\n{}", stage_metrics.join(", "), writer ) @@ -257,7 +257,9 @@ impl QueryStageExecutor for DefaultQueryStageExec { ) -> Result> { let (plan_arc, is_sort_shuffle): (Arc, bool) = match &self.shuffle_writer { - ShuffleWriterVariant::Hash(writer) => (Arc::new(writer.clone()), false), + ShuffleWriterVariant::Passthrough(writer) => { + (Arc::new(writer.clone()), false) + } ShuffleWriterVariant::Sort(writer) => (Arc::new(writer.clone()), true), }; info!( @@ -282,7 +284,9 @@ impl QueryStageExecutor for DefaultQueryStageExec { fn collect_plan_metrics(&self) -> Vec { match &self.shuffle_writer { - ShuffleWriterVariant::Hash(writer) => utils::collect_plan_metrics(writer), + ShuffleWriterVariant::Passthrough(writer) => { + utils::collect_plan_metrics(writer) + } ShuffleWriterVariant::Sort(writer) => utils::collect_plan_metrics(writer), } } diff --git a/ballista/executor/src/executor.rs b/ballista/executor/src/executor.rs index 3a1c103be8..dce01ae7f5 100644 --- a/ballista/executor/src/executor.rs +++ b/ballista/executor/src/executor.rs @@ -424,12 +424,11 @@ mod test { 1, Arc::new(NeverendingOperator::new()), work_dir.clone(), - None, ) .expect("creating shuffle writer"); let query_stage_exec = - DefaultQueryStageExec::new(ShuffleWriterVariant::Hash(shuffle_write)); + DefaultQueryStageExec::new(ShuffleWriterVariant::Passthrough(shuffle_write)); let executor_registration = ExecutorRegistration { id: "executor".to_string(), diff --git a/ballista/executor/src/flight_service.rs b/ballista/executor/src/flight_service.rs index 32a0db4047..3e3c8b079b 100644 --- a/ballista/executor/src/flight_service.rs +++ b/ballista/executor/src/flight_service.rs @@ -143,7 +143,7 @@ impl FlightService for BallistaFlightService { )); } - // Standard hash-based shuffle - read the entire file + // Standard single-file shuffle output - read the entire file let file = File::open(&path) .map_err(|e| { BallistaError::General(format!( diff --git a/ballista/scheduler/src/planner.rs b/ballista/scheduler/src/planner.rs index 0128540c5b..7283558991 100644 --- a/ballista/scheduler/src/planner.rs +++ b/ballista/scheduler/src/planner.rs @@ -700,16 +700,17 @@ pub(crate) fn create_shuffle_writer_with_config( ) -> Result> { let plan = make_empty_exec_serde_safe(plan)?; - // Check if sort-based shuffle is enabled + // Sort-based shuffle is the only shuffle writer for hash-repartition + // stages. Its tuning values still come from the session config. let ballista_config = config .extensions .get::() .cloned() .unwrap_or_default(); - if ballista_config.shuffle_sort_based_enabled() { - // Sort shuffle requires hash partitioning - if let Some(Partitioning::Hash(exprs, partition_count)) = partitioning { + // Sort shuffle requires hash partitioning. + match partitioning { + Some(Partitioning::Hash(exprs, partition_count)) => { let sort_config = SortShuffleConfig::new( true, ballista_config.shuffle_sort_based_batch_size(), @@ -718,25 +719,27 @@ pub(crate) fn create_shuffle_writer_with_config( ballista_config.shuffle_sort_based_memory_limit_per_task_bytes(), ); - return Ok(Arc::new(SortShuffleWriterExec::try_new( + Ok(Arc::new(SortShuffleWriterExec::try_new( job_id.to_owned(), stage_id, plan, "".to_owned(), Partitioning::Hash(exprs, partition_count), sort_config, - )?)); + )?)) } + // Stages that don't repartition write their input partitioning + // through: one file per output partition. + None => Ok(Arc::new(ShuffleWriterExec::try_new( + job_id.to_owned(), + stage_id, + plan, + "".to_owned(), + )?)), + Some(other) => Err(BallistaError::General(format!( + "unsupported shuffle output partitioning: {other}" + ))), } - - // Fall back to standard shuffle writer - Ok(Arc::new(ShuffleWriterExec::try_new( - job_id.to_owned(), - stage_id, - plan, - "".to_owned(), - partitioning, - )?)) } #[cfg(test)] diff --git a/ballista/scheduler/src/state/aqe/test/mod.rs b/ballista/scheduler/src/state/aqe/test/mod.rs index a8640d3e30..7c99abe83f 100644 --- a/ballista/scheduler/src/state/aqe/test/mod.rs +++ b/ballista/scheduler/src/state/aqe/test/mod.rs @@ -26,8 +26,6 @@ mod join_selection; /// Tests if plan is going to be split to stages correctly mod plan_to_stages; -use ballista_core::config::BALLISTA_SHUFFLE_SORT_BASED_ENABLED; -use ballista_core::extension::SessionConfigExt; use ballista_core::serde::scheduler::{ ExecutorMetadata, ExecutorOperatingSystemSpecification, ExecutorSpecification, PartitionId, PartitionLocation, PartitionStats, @@ -131,17 +129,3 @@ pub(crate) fn mock_context() -> SessionContext { SessionContext::new_with_state(state) } - -pub(crate) fn mock_context_sort_shuffle() -> SessionContext { - let config = SessionConfig::new_with_ballista() - .set_str(BALLISTA_SHUFFLE_SORT_BASED_ENABLED, "true") - .with_target_partitions(2) - .with_round_robin_repartition(false); - - let state = SessionStateBuilder::new() - .with_config(config) - .with_default_features() - .build(); - - SessionContext::new_with_state(state) -} diff --git a/ballista/scheduler/src/state/aqe/test/plan_to_stages.rs b/ballista/scheduler/src/state/aqe/test/plan_to_stages.rs index b115f273c5..993996af49 100644 --- a/ballista/scheduler/src/state/aqe/test/plan_to_stages.rs +++ b/ballista/scheduler/src/state/aqe/test/plan_to_stages.rs @@ -19,8 +19,7 @@ use crate::assert_plan; use crate::state::aqe::execution_plan::ExchangeExec; use crate::state::aqe::planner::AdaptivePlanner; use crate::state::aqe::test::{ - mock_batch, mock_context, mock_context_sort_shuffle, mock_memory_table, - mock_partitions_with_statistics, + mock_batch, mock_context, mock_memory_table, mock_partitions_with_statistics, }; use ballista_core::execution_plans::SortShuffleWriterExec; use datafusion::arrow::datatypes::{DataType, Field, Schema}; @@ -455,36 +454,6 @@ async fn should_ignore_inactive_stages() -> datafusion::error::Result<()> { Ok(()) } -#[tokio::test] -async fn should_use_sort_shuffle_when_enabled() -> datafusion::error::Result<()> { - let ctx = mock_context_sort_shuffle(); - ctx.register_batch("t", mock_batch()?)?; - - let q = r#" - select min(a) as c0, max(b) as c1, c as c2 from t group by c - "#; - - let plan = ctx.sql(q).await?.create_physical_plan().await?; - let mut planner = AdaptivePlanner::try_from_plan( - ctx.state().config(), - plan, - "test_job".to_owned(), - )?; - - let stages = planner.runnable_stages()?.unwrap(); - assert_eq!(1, stages.len()); - - let plan = stages.first().unwrap().plan.as_ref(); - assert!( - (plan as &dyn ExecutionPlan) - .downcast_ref::() - .is_some(), - "expected SortShuffleWriterExec when sort shuffle is enabled, got plan: {plan:?}" - ); - - Ok(()) -} - #[tokio::test] async fn should_use_sort_shuffle_by_default() -> datafusion::error::Result<()> { let ctx = mock_context(); diff --git a/benchmarks/src/bin/shuffle_bench.rs b/benchmarks/src/bin/shuffle_bench.rs index 8fd9184d28..e2bdbcbf35 100644 --- a/benchmarks/src/bin/shuffle_bench.rs +++ b/benchmarks/src/bin/shuffle_bench.rs @@ -17,14 +17,13 @@ //! Standalone shuffle benchmark for profiling Ballista shuffle write //! performance outside of a cluster. Streams input from Parquet files and -//! drives either the hash-based or sort-based shuffle writer end-to-end. +//! drives the sort-based shuffle writer end-to-end. //! //! # Usage //! //! ```sh //! cargo run --release --bin shuffle_bench -- \ //! --input /data/tpch-sf100/lineitem/ \ -//! --writer sort \ //! --partitions 200 \ //! --hash-columns 0,3 //! ``` @@ -33,10 +32,9 @@ //! ```sh //! cargo flamegraph --release --bin shuffle_bench -- \ //! --input /data/tpch-sf100/lineitem/ \ -//! --writer sort --partitions 200 +//! --partitions 200 //! ``` -use ballista_core::execution_plans::ShuffleWriterExec; use ballista_core::execution_plans::sort_shuffle::{ SortShuffleConfig, SortShuffleWriterExec, }; @@ -70,12 +68,8 @@ struct Args { #[arg(long)] input: PathBuf, - /// Shuffle writer to drive: `hash` (default) or `sort`. - #[arg(long, default_value = "hash")] - writer: String, - /// Partitioning scheme: `hash`, `single`, or `round-robin`. Currently - /// both writers only support `hash`; other values are rejected. + /// the writer only supports `hash`; other values are rejected. #[arg(long, default_value = "hash")] partitioning: String, @@ -117,27 +111,11 @@ struct Args { vcores: usize, } -#[derive(Clone, Copy, Debug)] -enum WriterKind { - Hash, - Sort, -} - #[derive(Clone, Copy, Debug)] enum PartitioningKind { Hash, } -fn parse_writer(s: &str) -> Result { - match s.to_lowercase().as_str() { - "hash" => Ok(WriterKind::Hash), - "sort" => Ok(WriterKind::Sort), - other => Err(format!( - "unknown writer: {other} (expected 'hash' or 'sort')" - )), - } -} - fn parse_partitioning(s: &str) -> Result { match s.to_lowercase().as_str() { "hash" => Ok(PartitioningKind::Hash), @@ -227,7 +205,6 @@ fn build_partitioning( async fn execute_shuffle_write( args: &Args, - writer_kind: WriterKind, partitioning_kind: PartitioningKind, hash_col_indices: &[usize], work_dir: PathBuf, @@ -270,43 +247,25 @@ async fn execute_shuffle_write( let work_dir_str = work_dir.to_str().unwrap().to_string(); fs::create_dir_all(&work_dir).expect("create work dir"); - let metrics: MetricsSet = match writer_kind { - WriterKind::Hash => { - let exec = ShuffleWriterExec::try_new( - format!("bench_job_{task_id}").into(), - 1, - input, - work_dir_str, - Some(partitioning), - )?; - let task_ctx = ctx.task_ctx(); - let mut stream = exec.execute(0, task_ctx)?; - let _ = utils::collect_stream(&mut stream).await; - exec.metrics().unwrap_or_default() - } - WriterKind::Sort => { - let cfg = SortShuffleConfig::new(true, args.batch_size); - let exec = SortShuffleWriterExec::try_new( - format!("bench_job_{task_id}").into(), - 1, - input, - work_dir_str, - partitioning, - cfg, - )?; - let task_ctx = ctx.task_ctx(); - let mut stream = exec.execute(0, task_ctx)?; - let _ = utils::collect_stream(&mut stream).await; - exec.metrics().unwrap_or_default() - } - }; + let cfg = SortShuffleConfig::new(true, args.batch_size); + let exec = SortShuffleWriterExec::try_new( + format!("bench_job_{task_id}").into(), + 1, + input, + work_dir_str, + partitioning, + cfg, + )?; + let task_ctx = ctx.task_ctx(); + let mut stream = exec.execute(0, task_ctx)?; + let _ = utils::collect_stream(&mut stream).await; + let metrics: MetricsSet = exec.metrics().unwrap_or_default(); Ok(metrics) } fn run_iteration( args: &Args, - writer_kind: WriterKind, partitioning_kind: PartitioningKind, hash_col_indices: &[usize], ) -> (f64, Option) { @@ -317,7 +276,6 @@ fn run_iteration( let work_dir = args.output_dir.join("task_0"); let metrics = execute_shuffle_write( args, - writer_kind, partitioning_kind, hash_col_indices, work_dir.clone(), @@ -337,7 +295,6 @@ fn run_iteration( handles.push(tokio::spawn(async move { let m = execute_shuffle_write( &args, - writer_kind, partitioning_kind, &hash_col_indices, work_dir.clone(), @@ -433,10 +390,6 @@ fn main() { .init(); let args = Args::parse(); - let writer_kind = parse_writer(&args.writer).unwrap_or_else(|e| { - eprintln!("error: {e}"); - std::process::exit(2); - }); let partitioning_kind = parse_partitioning(&args.partitioning).unwrap_or_else(|e| { eprintln!("error: {e}"); std::process::exit(2); @@ -448,7 +401,6 @@ fn main() { let (schema, total_rows) = read_parquet_metadata(&args.input, args.limit); println!("=== Ballista Shuffle Benchmark ==="); - println!("Writer: {writer_kind:?}"); println!("Partitioning: {partitioning_kind:?}"); println!("Input: {}", args.input.display()); println!( @@ -483,7 +435,7 @@ fn main() { format!("iter {}/{}", i - args.warmup + 1, args.iterations) }; let (elapsed, metrics) = - run_iteration(&args, writer_kind, partitioning_kind, &hash_col_indices); + run_iteration(&args, partitioning_kind, &hash_col_indices); if !is_warmup { times.push(elapsed); if metrics.is_some() { diff --git a/benchmarks/src/bin/tpch.rs b/benchmarks/src/bin/tpch.rs index 5c561b171f..8932dc29db 100644 --- a/benchmarks/src/bin/tpch.rs +++ b/benchmarks/src/bin/tpch.rs @@ -127,7 +127,7 @@ struct BallistaBenchmarkOpt { /// Configuration overrides in key=value format. /// Can be specified multiple times, e.g. - /// -c ballista.shuffle.sort_based.enabled=true + /// -c ballista.shuffle.sort_based.batch_size=8192 /// -c datafusion.execution.target_partitions=16 #[structopt(short = "c", long = "config", number_of_values = 1)] config_overrides: Vec, diff --git a/docs/source/user-guide/configs.md b/docs/source/user-guide/configs.md index 83f0351db5..8bc7276f4b 100644 --- a/docs/source/user-guide/configs.md +++ b/docs/source/user-guide/configs.md @@ -92,14 +92,13 @@ let expected = [ The following session-level keys control Ballista's shuffle behavior. See the [tuning guide](tuning-guide.md#shuffle-implementation) for an -explanation of the sort-based (default) and hash-based shuffle writers. +explanation of the sort-based shuffle writer. | key | type | default | description | | ------------------------------------------------------- | ------- | --------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------ | | ballista.shuffle.max_concurrent_read_requests | UInt64 | 64 | Maximum number of concurrent fetch requests the shuffle reader will issue. | | ballista.shuffle.force_remote_read | Boolean | false | Forces the shuffle reader to fetch every partition through Arrow Flight, even when the data is local. Intended for testing. | | ballista.shuffle.remote_read_prefer_flight | Boolean | false | For remote reads, prefer the Arrow Flight reader over the block reader. The block reader is generally faster. | -| ballista.shuffle.sort_based.enabled | Boolean | true | Enables the sort-based shuffle writer (consolidated data file per input partition with an index, instead of one file per (input partition, output partition) pair). | | ballista.shuffle.sort_based.batch_size | UInt64 | 8192 | Target row count when coalescing buffered batches before they are written or spilled. | | ballista.shuffle.sort_based.memory_limit_per_task_bytes | UInt64 | 268435456 | Per-task buffered-bytes budget at which the sort-based writer spills to disk (256 MiB default), counted independently of the runtime memory pool. Set to `0` to disable the per-task budget and spill only under memory-pool pressure — safe only with a bounded memory pool, otherwise the writer never spills and may run out of memory. | diff --git a/docs/source/user-guide/tuning-guide.md b/docs/source/user-guide/tuning-guide.md index b9f3d5d7f1..5391c399dd 100644 --- a/docs/source/user-guide/tuning-guide.md +++ b/docs/source/user-guide/tuning-guide.md @@ -127,11 +127,12 @@ scheduler or executors. Ballista exchanges data between query stages by writing the output of each upstream task to local files, which downstream tasks read either from disk -(when co-located) or over Arrow Flight. Two shuffle implementations are -available, with different trade-offs around file count, memory use, and -write latency. +(when co-located) or over Arrow Flight. Ballista uses a sort-based shuffle +writer that bounds file count and memory use. Single-partition stages (a +query's final output, coalesce, and broadcast-build stages) are written +directly to one file. -### Sort-based shuffle (default) +### Sort-based shuffle The sort-based writer accumulates incoming batches in memory, tracking each row's output partition. It spills the buffered batches to disk when either of @@ -160,7 +161,7 @@ output partition. This produces `2 × N` files instead of `N × M`, coalesces small batches to a target size before writing, and bounds shuffle memory use via -spilling — at the cost of higher write latency than the hash writer. +spilling. Worst-case sort-shuffle memory per executor is approximately `vcores × memory_limit_per_task_bytes`, since one writer task can run per @@ -169,31 +170,12 @@ sooner, or raise it to keep more data in memory and reduce spill I/O. Setting it to `0` removes the budget entirely and is safe only with a bounded memory pool (see the warning above). -### Hash-based shuffle (opt-in) - -The hash-based writer hashes each incoming `RecordBatch` and immediately -encodes the per-partition slices to Arrow IPC, streaming them into one -file per `(input_partition, output_partition)` pair. Nothing is buffered -in memory across batches. - -This is simple and low latency, but for `N` input partitions and `M` -output partitions it produces `N × M` files. Wide shuffles can therefore -generate a large number of small files. Consider switching to the -hash-based writer for narrow shuffles where the additional buffering -and merging of the sort-based writer is unnecessary overhead: - -```rust -let session_config = SessionConfig::new_with_ballista() - .set_bool("ballista.shuffle.sort_based.enabled", false); -``` - The following session-level keys tune its behavior: -| key | type | default | description | -| ------------------------------------------------------- | ------- | --------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------ | -| ballista.shuffle.sort_based.enabled | Boolean | true | Enables the sort-based shuffle writer. | -| ballista.shuffle.sort_based.batch_size | UInt64 | 8192 | Target row count when coalescing buffered batches before they are written or spilled. | -| ballista.shuffle.sort_based.memory_limit_per_task_bytes | UInt64 | 268435456 | Per-task buffered-bytes budget at which the writer spills to disk (256 MiB default). Counted independently of the runtime memory pool. Set to `0` to spill only under memory pressure — safe only with a bounded memory pool, otherwise the writer never spills and may run out of memory. | +| key | type | default | description | +| ------------------------------------------------------- | ------ | --------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------ | +| ballista.shuffle.sort_based.batch_size | UInt64 | 8192 | Target row count when coalescing buffered batches before they are written or spilled. | +| ballista.shuffle.sort_based.memory_limit_per_task_bytes | UInt64 | 268435456 | Per-task buffered-bytes budget at which the writer spills to disk (256 MiB default). Counted independently of the runtime memory pool. Set to `0` to spill only under memory pressure — safe only with a bounded memory pool, otherwise the writer never spills and may run out of memory. | ## Adaptive Query Execution (Experimental) diff --git a/python/python/tests/test_context.py b/python/python/tests/test_context.py index 047c8986e7..6e51803784 100644 --- a/python/python/tests/test_context.py +++ b/python/python/tests/test_context.py @@ -95,7 +95,7 @@ def test_cluster_config_propagates_to_distributed_dataframe(): def test_cluster_config_accepts_ballista_namespaced_keys(): - """Ballista-namespaced keys (e.g. ``ballista.shuffle.sort_based.enabled``) + """Ballista-namespaced keys (e.g. ``ballista.shuffle.sort_based.batch_size``) are not understood by the local DataFusion ``SessionConfig`` and used to panic when applied to it. They are forwarded to the scheduler only and must be ignored locally rather than crashing context construction. @@ -103,7 +103,7 @@ def test_cluster_config_accepts_ballista_namespaced_keys(): (address, port) = setup_test_cluster() overrides = { "datafusion.execution.target_partitions": "8", - "ballista.shuffle.sort_based.enabled": "true", + "ballista.shuffle.sort_based.batch_size": "8192", } ctx = BallistaSessionContext( address=f"df://{address}:{port}",