From ab0b083f840c80b4b0ba6ae8d549eee10e7fbee0 Mon Sep 17 00:00:00 2001 From: sim Date: Sat, 7 Nov 2026 14:26:23 -0700 Subject: [PATCH 1/3] test: round-trip sort shuffle output from several spilling inputs The existing multi-input test never spilled, so nothing checked that each output partition's range stitches every input's spill and in-memory bytes together in order. Run it again with a tiny buffer budget so every input spills several times and keeps an in-memory remainder. --- .../execution_plans/sort_shuffle/writer.rs | 71 ++++++++++++++----- 1 file changed, 55 insertions(+), 16 deletions(-) diff --git a/ballista/core/src/execution_plans/sort_shuffle/writer.rs b/ballista/core/src/execution_plans/sort_shuffle/writer.rs index c6e797dfa2..552cf6e85b 100644 --- a/ballista/core/src/execution_plans/sort_shuffle/writer.rs +++ b/ballista/core/src/execution_plans/sort_shuffle/writer.rs @@ -1898,6 +1898,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 +1934,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 +1977,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 +1988,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 +2022,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"); } } } From b63bceb7a75fe46d0acbfe501bda01c9d1e60b32 Mon Sep 17 00:00:00 2001 From: sim Date: Sat, 7 Nov 2026 14:11:47 -0700 Subject: [PATCH 2/3] perf: close sort shuffle spill writers after each spill event The spill manager kept one IPC writer open per spilled output partition until the task finished, each holding a file descriptor and its own compression context outside the memory pool. With hundreds of output partitions that is hundreds of open files per input partition and, with zstd, hundreds of megabytes of compressor state. Write each spill event as one complete stream appended to the partition's file and close it; readers already cross concatenated streams. --- .../src/execution_plans/sort_shuffle/spill.rs | 203 ++++++++++-------- .../execution_plans/sort_shuffle/writer.rs | 18 +- 2 files changed, 115 insertions(+), 106 deletions(-) diff --git a/ballista/core/src/execution_plans/sort_shuffle/spill.rs b/ballista/core/src/execution_plans/sort_shuffle/spill.rs index 8df23af25c..ca0cc1a140 100644 --- a/ballista/core/src/execution_plans/sort_shuffle/spill.rs +++ b/ballista/core/src/execution_plans/sort_shuffle/spill.rs @@ -22,27 +22,34 @@ //! 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 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::record_batch::RecordBatch; use log::debug; use std::collections::HashMap; -use std::fs::File; +use std::fs::OpenOptions; use std::io::BufWriter; use std::path::{Path, PathBuf}; /// 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, @@ -50,8 +57,6 @@ pub struct SpillManager { 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, /// Total number of batches written across all spill files. One call to @@ -98,7 +103,6 @@ impl SpillManager { spill_dir, schema, spill_files: HashMap::new(), - active_writers: HashMap::new(), compression, total_spilled_batches: 0, total_bytes_spilled: 0, @@ -106,51 +110,82 @@ impl SpillManager { }) } - /// 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); + /// 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 spill_path = self.spill_file_path(partition_id); + let mut writer: Option>> = None; + let mut event_batches: u64 = 0; + let mut event_bytes: u64 = 0; + for batch in batches { + let batch = batch?; + if batch.num_rows() == 0 { + continue; + } + let writer = match &mut writer { + Some(writer) => writer, + None => { + debug!("Spilling partition {} to {:?}", partition_id, spill_path); + let file = OpenOptions::new() + .create(true) + .append(true) + .open(&spill_path) + .map_err(BallistaError::IoError)?; + let options = IpcWriteOptions::default() + .try_with_compression(self.compression)?; + writer.insert(StreamWriter::try_new_with_options( + BufWriter::new(file), + &self.schema, + options, + )?) + } + }; + let bytes = 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; + event_batches += 1; + event_bytes += bytes; } - - 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); + if let Some(mut writer) = writer { + writer.finish()?; + writer + .into_inner()? + .into_inner() + .map_err(|e| BallistaError::IoError(e.into_error()))?; 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; + self.total_spilled_batches += event_batches; + self.total_bytes_spilled += event_bytes; - self.total_spilled_batches += 1; - self.total_bytes_spilled += bytes_written; + Ok((event_batches, event_bytes)) + } - Ok(bytes_written) + /// Spills a single `batch` for `partition_id` to disk; see + /// [`Self::spill_batches`]. + /// + /// 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 { + let (_, bytes) = + self.spill_batches(partition_id, &mut std::iter::once(Ok(batch.clone())))?; + Ok(bytes) } /// Returns true if the partition has a spill file. @@ -158,37 +193,8 @@ impl SpillManager { 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), - } - } - /// 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)?; } @@ -223,8 +229,7 @@ impl SpillManager { } /// 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. + /// were ever spilled for that partition. pub fn spill_path(&self, partition_id: usize) -> Option<&Path> { self.spill_files.get(&partition_id).map(PathBuf::as_path) } @@ -273,6 +278,30 @@ 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> { + use datafusion::arrow::ipc::reader::StreamReader; + use std::io::{Seek, SeekFrom}; + 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) + } + #[test] fn test_spill_and_read() -> Result<()> { let temp_dir = TempDir::new().unwrap(); @@ -297,11 +326,7 @@ mod tests { 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); @@ -331,16 +356,8 @@ mod tests { 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); diff --git a/ballista/core/src/execution_plans/sort_shuffle/writer.rs b/ballista/core/src/execution_plans/sort_shuffle/writer.rs index 552cf6e85b..24b5d02de5 100644 --- a/ballista/core/src/execution_plans/sort_shuffle/writer.rs +++ b/ballista/core/src/execution_plans/sort_shuffle/writer.rs @@ -688,12 +688,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 +768,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(); From 7d51a949c38f411bb33c4ba38b5ce0b99be492ce Mon Sep 17 00:00:00 2001 From: sim Date: Sun, 8 Nov 2026 02:28:09 -0700 Subject: [PATCH 3/3] refactor: share IPC stream encoding between spill and final write Address review feedback: spill and the in-memory encode now call one write_ipc_stream helper, SpillManager builds its IpcWriteOptions once, derives spill paths from partition_counters instead of a duplicate map, and the single-batch spill() wrapper is removed (tests call spill_batches). --- .../src/execution_plans/sort_shuffle/spill.rs | 202 ++++++++++-------- .../execution_plans/sort_shuffle/writer.rs | 26 +-- 2 files changed, 112 insertions(+), 116 deletions(-) diff --git a/ballista/core/src/execution_plans/sort_shuffle/spill.rs b/ballista/core/src/execution_plans/sort_shuffle/spill.rs index ca0cc1a140..e8c174aea2 100644 --- a/ballista/core/src/execution_plans/sort_shuffle/spill.rs +++ b/ballista/core/src/execution_plans/sort_shuffle/spill.rs @@ -32,15 +32,58 @@ use crate::JobId; use crate::error::{BallistaError, Result}; +use crate::utils::create_write_options; use datafusion::arrow::datatypes::SchemaRef; -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::OpenOptions; -use std::io::BufWriter; -use std::path::{Path, PathBuf}; +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. /// @@ -55,18 +98,17 @@ pub struct SpillManager { 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, - /// 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 { @@ -102,8 +144,7 @@ impl SpillManager { Ok(Self { spill_dir, schema, - spill_files: HashMap::new(), - compression, + write_options: create_write_options(compression)?, total_spilled_batches: 0, total_bytes_spilled: 0, partition_counters: HashMap::new(), @@ -122,75 +163,32 @@ impl SpillManager { partition_id: usize, batches: &mut dyn Iterator>, ) -> Result<(u64, u64)> { - let spill_path = self.spill_file_path(partition_id); - let mut writer: Option>> = None; - let mut event_batches: u64 = 0; - let mut event_bytes: u64 = 0; - for batch in batches { - let batch = batch?; - if batch.num_rows() == 0 { - continue; - } - let writer = match &mut writer { - Some(writer) => writer, - None => { - debug!("Spilling partition {} to {:?}", partition_id, spill_path); - let file = OpenOptions::new() - .create(true) - .append(true) - .open(&spill_path) - .map_err(BallistaError::IoError)?; - let options = IpcWriteOptions::default() - .try_with_compression(self.compression)?; - writer.insert(StreamWriter::try_new_with_options( - BufWriter::new(file), - &self.schema, - options, - )?) - } - }; - let bytes = 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; - event_batches += 1; - event_bytes += bytes; - } - if let Some(mut writer) = writer { - writer.finish()?; - writer - .into_inner()? - .into_inner() - .map_err(|e| BallistaError::IoError(e.into_error()))?; - self.spill_files.insert(partition_id, spill_path); - } - - self.total_spilled_batches += event_batches; - self.total_bytes_spilled += event_bytes; + 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()))?; - Ok((event_batches, event_bytes)) - } + *self.partition_counters.entry(partition_id).or_default() += stats; + self.total_spilled_batches += stats.num_batches; + self.total_bytes_spilled += stats.num_bytes; - /// Spills a single `batch` for `partition_id` to disk; see - /// [`Self::spill_batches`]. - /// - /// 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 { - let (_, bytes) = - self.spill_batches(partition_id, &mut std::iter::once(Ok(batch.clone())))?; - Ok(bytes) + 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) + self.partition_counters.contains_key(&partition_id) } /// Cleans up all spill files. @@ -222,16 +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. - 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. @@ -244,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() @@ -266,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; @@ -284,14 +287,12 @@ mod tests { manager: &SpillManager, partition_id: usize, ) -> Result> { - use datafusion::arrow::ipc::reader::StreamReader; - use std::io::{Seek, SeekFrom}; let path = manager.spill_path(partition_id).unwrap(); - let len = std::fs::metadata(path)?.len(); + 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)?; + 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() { @@ -302,6 +303,17 @@ mod tests { 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(); @@ -319,8 +331,8 @@ 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)); @@ -349,8 +361,8 @@ 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)); @@ -379,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); @@ -426,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 24b5d02de5..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(()) })?; @@ -852,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;