diff --git a/ballista/core/src/execution_plans/sort_shuffle/spill.rs b/ballista/core/src/execution_plans/sort_shuffle/spill.rs index 8df23af25c..e8c174aea2 100644 --- a/ballista/core/src/execution_plans/sort_shuffle/spill.rs +++ b/ballista/core/src/execution_plans/sort_shuffle/spill.rs @@ -22,46 +22,93 @@ //! because the per-task buffer budget was reached. //! At finalization, the spill bytes are concatenated verbatim into the //! consolidated output file alongside the in-memory remainder. +//! +//! Each spill event appends one complete Arrow IPC stream to each spilled +//! partition's file and closes its writer, so a partition's spill file holds +//! one stream per event and no writer outlives the event. Keeping a writer +//! open per partition would hold a file descriptor and a compression context +//! (about 0.3-0.7 MB each for zstd, outside the memory pool) for every output +//! partition until the task finishes. use crate::JobId; use crate::error::{BallistaError, Result}; +use crate::utils::create_write_options; use datafusion::arrow::datatypes::SchemaRef; -use datafusion::arrow::ipc::reader::StreamReader; -use datafusion::arrow::ipc::writer::StreamWriter; -use datafusion::arrow::ipc::{CompressionType, writer::IpcWriteOptions}; +use datafusion::arrow::ipc::CompressionType; +use datafusion::arrow::ipc::writer::{IpcWriteOptions, StreamWriter}; use datafusion::arrow::record_batch::RecordBatch; use log::debug; use std::collections::HashMap; -use std::fs::File; -use std::io::BufWriter; -use std::path::{Path, PathBuf}; +use std::fs::OpenOptions; +use std::io::{BufWriter, Write}; +use std::path::PathBuf; + +/// Batch, row and byte counts for rows written as IPC. +/// +/// `num_bytes` is the Arrow in-memory size of each batch +/// (`RecordBatch::get_array_memory_size`), not the encoded size. +#[derive(Clone, Copy, Debug, Default)] +pub(super) struct EncodedStats { + pub(super) num_batches: u64, + pub(super) num_rows: u64, + pub(super) num_bytes: u64, +} + +impl std::ops::AddAssign for EncodedStats { + fn add_assign(&mut self, other: Self) { + self.num_batches += other.num_batches; + self.num_rows += other.num_rows; + self.num_bytes += other.num_bytes; + } +} + +/// Writes `batches` to `sink` as one complete Arrow IPC stream and returns +/// the sink with the counts of what was written. +/// +/// Shared by spilling (a spill file) and the final in-memory encode (a +/// `Vec`), so both produce byte-identical streams for the same rows. +pub(super) fn write_ipc_stream( + sink: W, + schema: &SchemaRef, + options: &IpcWriteOptions, + batches: &mut dyn Iterator>, +) -> datafusion::error::Result<(W, EncodedStats)> { + let mut writer = StreamWriter::try_new_with_options(sink, schema, options.clone())?; + let mut stats = EncodedStats::default(); + for batch in batches { + let batch = batch?; + stats.num_batches += 1; + stats.num_rows += batch.num_rows() as u64; + stats.num_bytes += batch.get_array_memory_size() as u64; + writer.write(&batch)?; + } + Ok((writer.into_inner()?, stats)) +} /// Manages spill files for sort-based shuffle. /// /// When partition buffers exceed memory limits, they are spilled to disk -/// as Arrow IPC streams. Each output partition has at most one spill file -/// that is appended to across multiple spill calls. During finalization, -/// the spill file bytes are concatenated directly into the consolidated -/// output file (no decode/re-encode round-trip). +/// as Arrow IPC streams. Each output partition has at most one spill file, +/// and each spill call appends one complete stream to it. During +/// finalization, the spill file bytes are concatenated directly into the +/// consolidated output file (no decode/re-encode round-trip); readers +/// already cross concatenated streams. pub struct SpillManager { /// Base directory for spill files spill_dir: PathBuf, /// Schema shared by all spill writers schema: SchemaRef, - /// Spill file path per output partition: partition_id -> spill_file_path - spill_files: HashMap, - /// Active writers per partition, kept open for appending - active_writers: HashMap>>, - /// Compression codec for spill files - compression: Option, + /// IPC write options (including the compression codec) for spill files + write_options: IpcWriteOptions, /// Total number of batches written across all spill files. One call to /// `spill_all_partitions` typically increments this multiple times (once /// per partition that had buffered rows). total_spilled_batches: u64, /// Total bytes spilled total_bytes_spilled: u64, - /// Per-partition counters: partition_id -> (batches, rows, bytes) - partition_counters: HashMap, + /// Counts per spilled output partition. A partition has an entry exactly + /// when it has a spill file. + partition_counters: HashMap, } impl SpillManager { @@ -97,98 +144,55 @@ impl SpillManager { Ok(Self { spill_dir, schema, - spill_files: HashMap::new(), - active_writers: HashMap::new(), - compression, + write_options: create_write_options(compression)?, total_spilled_batches: 0, total_bytes_spilled: 0, partition_counters: HashMap::new(), }) } - /// Spills a single `batch` for `partition_id` to disk. The first call for - /// a given `partition_id` creates the spill file; subsequent calls append. + /// Spills `batches` for `partition_id` to disk as one complete IPC + /// stream appended to the partition's spill file, which the first call + /// for a given `partition_id` creates. The writer is finished and dropped + /// before returning. /// - /// Returns the number of bytes written (estimated from the batch's array - /// memory size). - pub fn spill(&mut self, partition_id: usize, batch: &RecordBatch) -> Result { - if batch.num_rows() == 0 { - return Ok(0); - } - - if !self.active_writers.contains_key(&partition_id) { - let spill_path = self.spill_file_path(partition_id); - debug!( - "Creating spill file for partition {} at {:?}", - partition_id, spill_path - ); - let file = File::create(&spill_path).map_err(BallistaError::IoError)?; - let buffered = BufWriter::new(file); - - let options = - IpcWriteOptions::default().try_with_compression(self.compression)?; - - let writer = - StreamWriter::try_new_with_options(buffered, &self.schema, options)?; - - self.active_writers.insert(partition_id, writer); - self.spill_files.insert(partition_id, spill_path); - } - - let writer = self.active_writers.get_mut(&partition_id).unwrap(); - let bytes_written = batch.get_array_memory_size() as u64; - writer.write(batch)?; - - let entry = self - .partition_counters - .entry(partition_id) - .or_insert((0, 0, 0)); - entry.0 += 1; - entry.1 += batch.num_rows() as u64; - entry.2 += bytes_written; + /// Returns `(batches, bytes)` written, with bytes estimated from each + /// batch's array memory size. + pub fn spill_batches( + &mut self, + partition_id: usize, + batches: &mut dyn Iterator>, + ) -> Result<(u64, u64)> { + let path = self.spill_file_path(partition_id); + debug!("Spilling partition {partition_id} to {path:?}"); + let file = OpenOptions::new() + .create(true) + .append(true) + .open(&path) + .map_err(BallistaError::IoError)?; + let (sink, stats) = write_ipc_stream( + BufWriter::new(file), + &self.schema, + &self.write_options, + batches, + )?; + sink.into_inner() + .map_err(|e| BallistaError::IoError(e.into_error()))?; - self.total_spilled_batches += 1; - self.total_bytes_spilled += bytes_written; + *self.partition_counters.entry(partition_id).or_default() += stats; + self.total_spilled_batches += stats.num_batches; + self.total_bytes_spilled += stats.num_bytes; - Ok(bytes_written) + Ok((stats.num_batches, stats.num_bytes)) } /// Returns true if the partition has a spill file. pub fn has_spill_files(&self, partition_id: usize) -> bool { - self.spill_files.contains_key(&partition_id) - } - - /// Finishes all active writers so spill files can be read. - /// Must be called before `open_spill_reader`. - pub fn finish_writers(&mut self) -> Result<()> { - for (_, mut writer) in self.active_writers.drain() { - writer.finish()?; - } - Ok(()) - } - - /// Opens the spill file for a partition and returns a streaming - /// reader. `finish_writers` must be called before this method. - pub fn open_spill_reader( - &self, - partition_id: usize, - ) -> Result>> { - match self.spill_files.get(&partition_id) { - Some(spill_path) => { - let file = File::open(spill_path).map_err(BallistaError::IoError)?; - let reader = StreamReader::try_new(file, None)?; - Ok(Some(reader)) - } - None => Ok(None), - } + self.partition_counters.contains_key(&partition_id) } /// Cleans up all spill files. pub fn cleanup(&mut self) -> Result<()> { - // Finish any active writers first - for (_, mut writer) in self.active_writers.drain() { - let _ = writer.finish(); - } if self.spill_dir.exists() { std::fs::remove_dir_all(&self.spill_dir).map_err(BallistaError::IoError)?; } @@ -216,17 +220,19 @@ impl SpillManager { /// at the time of the spill call (`RecordBatch::get_array_memory_size`). /// It is **not** the compressed on-disk size. pub fn partition_stats(&self, partition_id: usize) -> (u64, u64, u64) { - self.partition_counters + let stats = self + .partition_counters .get(&partition_id) .copied() - .unwrap_or((0, 0, 0)) + .unwrap_or_default(); + (stats.num_batches, stats.num_rows, stats.num_bytes) } - /// Returns the path to a partition's spill file, or `None` if no batches - /// were ever spilled for that partition. `finish_writers` must be called - /// before this path is read by another process. - pub fn spill_path(&self, partition_id: usize) -> Option<&Path> { - self.spill_files.get(&partition_id).map(PathBuf::as_path) + /// Returns the path to a partition's spill file, or `None` if the + /// partition never spilled. + pub fn spill_path(&self, partition_id: usize) -> Option { + self.has_spill_files(partition_id) + .then(|| self.spill_file_path(partition_id)) } /// Returns the spill file path for a partition. @@ -239,8 +245,8 @@ impl std::fmt::Debug for SpillManager { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { f.debug_struct("SpillManager") .field("spill_dir", &self.spill_dir) - .field("spill_files", &self.spill_files) - .field("compression", &self.compression) + .field("write_options", &self.write_options) + .field("partition_counters", &self.partition_counters) .field("total_spilled_batches", &self.total_spilled_batches) .field("total_bytes_spilled", &self.total_bytes_spilled) .finish() @@ -261,6 +267,8 @@ mod tests { use super::*; use datafusion::arrow::array::Int32Array; use datafusion::arrow::datatypes::{DataType, Field, Schema}; + use datafusion::arrow::ipc::reader::StreamReader; + use std::io::{Seek, SeekFrom}; use std::sync::Arc; use tempfile::TempDir; @@ -273,6 +281,39 @@ mod tests { RecordBatch::try_new(schema.clone(), vec![Arc::new(array)]).unwrap() } + /// Reads every batch of a partition's spill file, across the one stream + /// each spill event appended. + fn read_spill( + manager: &SpillManager, + partition_id: usize, + ) -> Result> { + let path = manager.spill_path(partition_id).unwrap(); + let len = std::fs::metadata(&path)?.len(); + let mut batches = Vec::new(); + let mut offset = 0; + while offset < len { + let mut file = std::fs::File::open(&path)?; + file.seek(SeekFrom::Start(offset))?; + let mut reader = StreamReader::try_new(file, None)?; + for batch in reader.by_ref() { + batches.push(batch?); + } + offset = reader.get_mut().stream_position()?; + } + Ok(batches) + } + + /// Spills a single batch as its own spill event, returning bytes spilled. + fn spill_one( + manager: &mut SpillManager, + partition_id: usize, + batch: RecordBatch, + ) -> Result { + let (_, bytes) = + manager.spill_batches(partition_id, &mut std::iter::once(Ok(batch)))?; + Ok(bytes) + } + #[test] fn test_spill_and_read() -> Result<()> { let temp_dir = TempDir::new().unwrap(); @@ -290,18 +331,14 @@ mod tests { let b1 = create_test_batch(&schema, vec![1, 2, 3]); let b2 = create_test_batch(&schema, vec![4, 5]); - assert!(manager.spill(0, &b1)? > 0); - assert!(manager.spill(0, &b2)? > 0); + assert!(spill_one(&mut manager, 0, b1)? > 0); + assert!(spill_one(&mut manager, 0, b2)? > 0); assert!(manager.has_spill_files(0)); assert!(!manager.has_spill_files(1)); assert_eq!(manager.total_spilled_batches(), 2); - manager.finish_writers()?; - - let reader = manager.open_spill_reader(0)?.unwrap(); - let read_batches: Vec<_> = - reader.into_iter().collect::>()?; + let read_batches = read_spill(&manager, 0)?; assert_eq!(read_batches.len(), 2); assert_eq!(read_batches[0].num_rows(), 3); assert_eq!(read_batches[1].num_rows(), 2); @@ -324,23 +361,15 @@ mod tests { Some(CompressionType::LZ4_FRAME), )?; - manager.spill(0, &create_test_batch(&schema, vec![1, 2]))?; - manager.spill(1, &create_test_batch(&schema, vec![3, 4]))?; + spill_one(&mut manager, 0, create_test_batch(&schema, vec![1, 2]))?; + spill_one(&mut manager, 1, create_test_batch(&schema, vec![3, 4]))?; assert!(manager.has_spill_files(0)); assert!(manager.has_spill_files(1)); assert_eq!(manager.total_spilled_batches(), 2); - manager.finish_writers()?; - - let r0 = manager - .open_spill_reader(0)? - .unwrap() - .collect::, _>>()?; - let r1 = manager - .open_spill_reader(1)? - .unwrap() - .collect::, _>>()?; + let r0 = read_spill(&manager, 0)?; + let r1 = read_spill(&manager, 1)?; assert_eq!(r0.len(), 1); assert_eq!(r1.len(), 1); @@ -362,9 +391,9 @@ mod tests { Some(CompressionType::LZ4_FRAME), )?; - manager.spill(0, &create_test_batch(&schema, vec![1, 2, 3]))?; - manager.spill(0, &create_test_batch(&schema, vec![4, 5]))?; - manager.spill(1, &create_test_batch(&schema, vec![6]))?; + spill_one(&mut manager, 0, create_test_batch(&schema, vec![1, 2, 3]))?; + spill_one(&mut manager, 0, create_test_batch(&schema, vec![4, 5]))?; + spill_one(&mut manager, 1, create_test_batch(&schema, vec![6]))?; let (b0, r0, bytes0) = manager.partition_stats(0); let (b1, r1, bytes1) = manager.partition_stats(1); @@ -409,7 +438,7 @@ mod tests { Some(CompressionType::LZ4_FRAME), )?; - manager.spill(0, &create_test_batch(&schema, vec![1, 2]))?; + spill_one(&mut manager, 0, create_test_batch(&schema, vec![1, 2]))?; let spill_dir = manager.spill_dir.clone(); assert!(spill_dir.exists()); diff --git a/ballista/core/src/execution_plans/sort_shuffle/writer.rs b/ballista/core/src/execution_plans/sort_shuffle/writer.rs index c6e797dfa2..68de00e490 100644 --- a/ballista/core/src/execution_plans/sort_shuffle/writer.rs +++ b/ballista/core/src/execution_plans/sort_shuffle/writer.rs @@ -39,7 +39,7 @@ use super::super::shuffle_writer_trait::ShuffleWriter; use super::buffer::BufferedBatches; use super::config::SortShuffleConfig; use super::index::ShuffleIndex; -use super::spill::SpillManager; +use super::spill::{EncodedStats, SpillManager, write_ipc_stream}; use crate::JobId; use crate::extension::SessionConfigExt; use crate::serde::protobuf::{ShuffleWritePartition, TaskColumnStats}; @@ -94,14 +94,6 @@ struct InputPartitionOutput { _reservation: MemoryReservation, } -/// What one output partition's encoded rows amount to. -#[derive(Clone, Copy, Default)] -struct EncodedStats { - num_batches: u64, - num_rows: u64, - num_bytes: u64, -} - /// One input partition's rows, encoded into one IPC buffer per output /// partition, with the counts for each. struct EncodedPartitions { @@ -129,17 +121,9 @@ fn encode_buffered_partitions( // One writer open at a time: each holds its own compression context. buffered.drain(|partition, batches| { - let mut writer = - StreamWriter::try_new_with_options(Vec::new(), schema, opts.clone())?; - let stat = &mut stats[partition]; - for batch in batches { - let batch = batch?; - stat.num_rows += batch.num_rows() as u64; - stat.num_bytes += batch.get_array_memory_size() as u64; - stat.num_batches += 1; - writer.write(&batch)?; - } - encoded[partition] = writer.into_inner()?; + let (buf, stat) = write_ipc_stream(Vec::new(), schema, opts, batches)?; + encoded[partition] = buf; + stats[partition] = stat; Ok(()) })?; @@ -688,12 +672,6 @@ impl SortShuffleWriterExec { } } - // Finish spill writers before the coordinator reads them back while - // building the task's consolidated file. - spill_manager - .finish_writers() - .map_err(|e| DataFusionError::Execution(format!("{e:?}")))?; - metrics.spill_count.add(spill_events as usize); metrics .spill_bytes @@ -774,13 +752,11 @@ fn spill_all_partitions( let mut batches_written: u64 = 0; let mut bytes_written: u64 = 0; buffered.drain(|partition_id, batches| { - for batch in batches { - let written = spill_manager - .spill(partition_id, &batch?) - .map_err(|e| DataFusionError::Execution(format!("{e:?}")))?; - batches_written += 1; - bytes_written += written; - } + let (batches, bytes) = spill_manager + .spill_batches(partition_id, batches) + .map_err(|e| DataFusionError::Execution(format!("{e:?}")))?; + batches_written += batches; + bytes_written += bytes; Ok(()) })?; reservation.free(); @@ -860,7 +836,7 @@ fn write_task_consolidated( // Same flush-then-bypass discipline as above: anything already // buffered has to land before the raw copy appends to the file. output.flush()?; - let mut spill_file = File::open(spill_path)?; + let mut spill_file = File::open(&spill_path)?; std::io::copy(&mut spill_file, output.get_mut())?; } num_batches += spill_batches; @@ -1898,6 +1874,27 @@ mod tests { /// reader crossing input boundaries inside a single range. #[tokio::test] async fn multiple_input_partitions_write_one_file() -> Result<()> { + multi_input_round_trip(1, SortShuffleConfig::default(), false).await + } + + /// Every input spills several times and keeps an in-memory remainder, so + /// each output partition's range interleaves spill and encoded spans from + /// every input. + #[tokio::test] + async fn multiple_spilling_input_partitions_round_trip() -> Result<()> { + multi_input_round_trip( + 7, + SortShuffleConfig::default().with_memory_limit_per_task_bytes(4 * 1024), + true, + ) + .await + } + + async fn multi_input_round_trip( + batches_per_input: usize, + config: SortShuffleConfig, + expect_spills: bool, + ) -> Result<()> { use super::super::reader::stream_sort_shuffle_partition; use datafusion::arrow::array::Int64Array; use datafusion::physical_plan::ExecutionPlanProperties; @@ -1913,21 +1910,26 @@ mod tests { let num_partitions = 4; // One input partition per outer Vec, so the writer sees M = 3. + let rows_per_batch = rows_per_input / batches_per_input; + let rows_per_input = rows_per_batch * batches_per_input; let partitions: Vec> = (0..num_inputs) .map(|p| { - let start = (p * rows_per_input) as i64; - let keys: Vec = (start..start + rows_per_input as i64).collect(); - let values: Vec = keys.iter().map(|k| k * 2).collect(); - vec![ - RecordBatch::try_new( - schema.clone(), - vec![ - Arc::new(Int64Array::from(keys)), - Arc::new(Int64Array::from(values)), - ], - ) - .unwrap(), - ] + (0..batches_per_input) + .map(|b| { + let start = (p * rows_per_input + b * rows_per_batch) as i64; + let keys: Vec = + (start..start + rows_per_batch as i64).collect(); + let values: Vec = keys.iter().map(|k| k * 2).collect(); + RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(Int64Array::from(keys)), + Arc::new(Int64Array::from(values)), + ], + ) + .unwrap() + }) + .collect() }) .collect(); @@ -1951,7 +1953,7 @@ mod tests { input, work_dir.path().to_str().unwrap().to_string(), Partitioning::Hash(vec![Arc::new(Column::new("k", 0))], num_partitions), - SortShuffleConfig::default(), + config, )? .with_task_id(7); @@ -1962,6 +1964,13 @@ mod tests { while stream.next().await.is_some() {} } + let spills = writer + .metrics() + .and_then(|m| m.sum_by_name("spill_count")) + .map(|v| v.as_usize()) + .unwrap_or(0); + assert_eq!(spills > 0, expect_spills, "unexpected spill count {spills}"); + // One directory named after the task id, holding one data file. let stage_dir = work_dir.path().join("multi_input_job").join("1"); let file_dirs: Vec<_> = std::fs::read_dir(&stage_dir)? @@ -1989,8 +1998,14 @@ mod tests { .as_any() .downcast_ref::() .unwrap(); - for v in arr.values() { - assert!(seen.insert(*v), "key {v} appeared twice across partitions"); + let values = batch + .column(1) + .as_any() + .downcast_ref::() + .unwrap(); + for (k, v) in arr.values().iter().zip(values.values()) { + assert_eq!(*v, k * 2, "key {k} came back with the wrong value"); + assert!(seen.insert(*k), "key {k} appeared twice across partitions"); } } }