diff --git a/ballista/core/src/execution_plans/sort_shuffle/buffer.rs b/ballista/core/src/execution_plans/sort_shuffle/buffer.rs index b5dc81abb3..f81a72a029 100644 --- a/ballista/core/src/execution_plans/sort_shuffle/buffer.rs +++ b/ballista/core/src/execution_plans/sort_shuffle/buffer.rs @@ -17,42 +17,42 @@ //! In-memory buffering for sort-based shuffle. //! -//! Holds whole input record batches plus per-output-partition row indices -//! (`(batch_idx, row_idx)` pairs). Rows are not copied at insertion time; -//! materialization is deferred to spill or final-write time and performed -//! via `arrow::compute::interleave_record_batch`. +//! Holds whole input record batches plus each output partition's +//! `(batch_idx, row_idx)` pairs. Rows are not copied when they arrive; output +//! batches are gathered with `interleave_record_batch` when the buffer is +//! drained, at spill or final-write time. +use super::partitioned_batch_iterator::PartitionedBatchIterator; use datafusion::arrow::datatypes::SchemaRef; use datafusion::arrow::record_batch::RecordBatch; +use datafusion::error::Result; -/// Return type of [`BufferedBatches::take`]: `(batches, per-partition indices)`. -pub type BufferedTake = (Vec, Vec>); - -/// Holds whole input `RecordBatch`es and per-partition `(batch_idx, row_idx)` -/// index lists. Rows are not copied at insertion time — only the indices are -/// recorded. Materialization happens through `PartitionedBatchIterator` at -/// spill or final-write time, by way of `interleave_record_batch`. +/// Rows buffered for every output partition of one input partition. #[derive(Debug)] pub struct BufferedBatches { schema: SchemaRef, - /// All input batches, in arrival order. Indexed by `batch_idx` in - /// `indices`. Never sliced. + batch_size: usize, + /// Input batches in arrival order, indexed by `batch_idx`. batches: Vec, - /// One entry per output partition. Each `(u32, u32)` is `(batch_idx, - /// row_idx)`, referring to a row inside `batches[batch_idx]`. + /// Total `get_array_memory_size` of `batches`. + batches_bytes: usize, + /// One list of `(batch_idx, row_idx)` pairs per output partition. indices: Vec>, - /// Total rows currently referenced by `indices`. Test diagnostic only. - num_buffered_rows: usize, + /// Total allocated capacity of `indices`, in bytes. + indices_bytes: usize, } impl BufferedBatches { - /// Creates a new buffer for the given partition count and schema. - pub fn new(num_partitions: usize, schema: SchemaRef) -> Self { + /// Creates a buffer for `num_partitions` output partitions that drains + /// into batches of at most `batch_size` rows. + pub fn new(num_partitions: usize, schema: SchemaRef, batch_size: usize) -> Self { Self { schema, + batch_size, batches: Vec::new(), - indices: (0..num_partitions).map(|_| Vec::new()).collect(), - num_buffered_rows: 0, + batches_bytes: 0, + indices: vec![Vec::new(); num_partitions], + indices_bytes: 0, } } @@ -61,35 +61,19 @@ impl BufferedBatches { self.indices.len() } - /// Returns true if no batches have been pushed yet. + /// Returns true if no batches are buffered. pub fn is_empty(&self) -> bool { self.batches.is_empty() } - /// Returns the total number of rows currently referenced by indices. - #[allow(dead_code)] - pub fn num_buffered_rows(&self) -> usize { - self.num_buffered_rows - } - - /// Returns the total heap-allocated size, in bytes, of the per-partition - /// `(batch_idx, row_idx)` index `Vec`s. Uses capacity (not length) so the - /// figure tracks actual heap allocation as `Vec`s grow. - pub fn indices_allocated_size(&self) -> usize { - self.indices - .iter() - .map(|v| v.capacity() * std::mem::size_of::<(u32, u32)>()) - .sum() - } - - /// Returns the row indices for output partition `partition_id`. - pub fn indices_for(&self, partition_id: usize) -> &[(u32, u32)] { - &self.indices[partition_id] + /// Returns the bytes held: the buffered batches plus the allocated + /// capacity of the index lists. + pub fn memory_size(&self) -> usize { + self.batches_bytes + self.indices_bytes } - /// Pushes a whole input `batch` and records, for each output partition - /// `p`, the row indices from `per_partition_rows[p]` as `(batch_idx, r)` - /// pairs in that partition's index list. + /// Buffers `batch`, recording the rows listed in `per_partition_rows[p]` + /// for output partition `p`. /// /// `per_partition_rows.len()` must equal `num_partitions()`. pub fn push_batch(&mut self, batch: RecordBatch, per_partition_rows: &[Vec]) { @@ -99,26 +83,42 @@ impl BufferedBatches { "BufferedBatches::push_batch schema mismatch" ); let batch_idx = self.batches.len() as u32; - for (p, rows) in per_partition_rows.iter().enumerate() { - let dst = &mut self.indices[p]; - dst.reserve(rows.len()); - for &r in rows { - dst.push((batch_idx, r)); - self.num_buffered_rows += 1; - } + for (partition_indices, rows) in self.indices.iter_mut().zip(per_partition_rows) { + let capacity_before = partition_indices.capacity(); + partition_indices.extend(rows.iter().map(|&r| (batch_idx, r))); + self.indices_bytes += (partition_indices.capacity() - capacity_before) + * size_of::<(u32, u32)>(); } + self.batches_bytes += batch.get_array_memory_size(); self.batches.push(batch); } - /// Drains all state, returning the buffered batches and per-partition - /// index lists. After this call the buffer is empty. - pub fn take(&mut self) -> BufferedTake { - self.num_buffered_rows = 0; + /// Empties the buffer, calling `f` once per non-empty partition, in + /// partition order, with an iterator over its output batches of at most + /// `batch_size` rows, in arrival order. Batches are gathered one at a time, + /// as `f` consumes them. The buffer is emptied first, so if `f` or a gather + /// fails, the remaining rows are dropped. + pub fn drain( + &mut self, + mut f: impl FnMut(usize, &mut dyn Iterator>) -> Result<()>, + ) -> Result<()> { + self.batches_bytes = 0; + self.indices_bytes = 0; let batches = std::mem::take(&mut self.batches); - // Preserve the partition count by replacing each inner vec with empty - let indices: Vec> = - self.indices.iter_mut().map(std::mem::take).collect(); - (batches, indices) + let indices: Vec<_> = self.indices.iter_mut().map(std::mem::take).collect(); + let batch_refs: Vec<&RecordBatch> = batches.iter().collect(); + for (partition, partition_indices) in indices.iter().enumerate() { + if partition_indices.is_empty() { + continue; + } + let mut iter = PartitionedBatchIterator::new( + &batch_refs, + partition_indices, + self.batch_size, + ); + f(partition, &mut iter)?; + } + Ok(()) } } @@ -129,58 +129,81 @@ mod tests { use datafusion::arrow::datatypes::{DataType, Field, Schema}; use std::sync::Arc; - fn create_test_schema() -> SchemaRef { + fn schema() -> SchemaRef { Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])) } - fn create_test_batch(schema: &SchemaRef, values: Vec) -> RecordBatch { - let array = Int32Array::from(values); - RecordBatch::try_new(schema.clone(), vec![Arc::new(array)]).unwrap() + fn batch(values: Vec) -> RecordBatch { + RecordBatch::try_new(schema(), vec![Arc::new(Int32Array::from(values))]).unwrap() + } + + /// Each partition `drain` visits, with its batches. + fn drain(bb: &mut BufferedBatches) -> Vec<(usize, Vec)> { + let mut out = vec![]; + bb.drain(|p, batches| { + out.push((p, batches.collect::>()?)); + Ok(()) + }) + .unwrap(); + out + } + + fn index_bytes(bb: &BufferedBatches) -> usize { + bb.indices + .iter() + .map(|v| v.capacity() * size_of::<(u32, u32)>()) + .sum() } #[test] - fn buffered_batches_pushes_and_partitions_indices() { - let schema = create_test_schema(); - let mut bb = BufferedBatches::new(3, schema.clone()); + fn drains_each_partition_in_arrival_order_and_batch_size() { + // With 3 rows a batch, partition 0's first output batch takes rows + // from both inputs, so chunking by `batch_size` is told apart from + // one output batch per input batch. + let mut bb = BufferedBatches::new(3, schema(), 3); assert!(bb.is_empty()); assert_eq!(bb.num_partitions(), 3); - let batch_a = create_test_batch(&schema, vec![10, 20, 30, 40]); - let batch_b = create_test_batch(&schema, vec![50, 60]); - - // Partition 0 gets rows {0, 2} from batch 0; partition 1 gets {1, 3} from - // batch 0; partition 2 gets {0, 1} from batch 1. - let per_partition_a: Vec> = vec![vec![0, 2], vec![1, 3], vec![]]; - let per_partition_b: Vec> = vec![vec![], vec![], vec![0, 1]]; - - bb.push_batch(batch_a, &per_partition_a); - bb.push_batch(batch_b, &per_partition_b); - + // Partition 0 gets rows {0, 2} of the first batch and all of the + // second; partition 1 gets {3, 1} of the first; partition 2, none. + bb.push_batch( + batch(vec![10, 20, 30, 40]), + &[vec![0, 2], vec![3, 1], vec![]], + ); + bb.push_batch(batch(vec![50, 60]), &[vec![0, 1], vec![], vec![]]); assert!(!bb.is_empty()); - // Total rows referenced by indices: 2 + 2 + 2 = 6 - assert_eq!(bb.num_buffered_rows(), 6); - - let p0 = bb.indices_for(0); - assert_eq!(p0, &[(0u32, 0u32), (0, 2)]); - let p1 = bb.indices_for(1); - assert_eq!(p1, &[(0u32, 1u32), (0, 3)]); - let p2 = bb.indices_for(2); - assert_eq!(p2, &[(1u32, 0u32), (1, 1)]); + + assert_eq!( + drain(&mut bb), + vec![ + (0, vec![batch(vec![10, 30, 50]), batch(vec![60])]), + (1, vec![batch(vec![40, 20])]), + ] + ); + assert!(bb.is_empty()); + assert_eq!(drain(&mut bb), vec![]); } #[test] - fn buffered_batches_take_drains_state() { - let schema = create_test_schema(); - let mut bb = BufferedBatches::new(2, schema.clone()); - bb.push_batch(create_test_batch(&schema, vec![1, 2]), &[vec![0], vec![1]]); - - let (batches, indices) = bb.take(); - assert_eq!(batches.len(), 1); - assert_eq!(indices.len(), 2); - assert_eq!(indices[0], vec![(0, 0)]); - assert_eq!(indices[1], vec![(0, 1)]); - - assert!(bb.is_empty()); - assert_eq!(bb.num_buffered_rows(), 0); + fn memory_size_counts_batches_and_index_capacity() { + let mut bb = BufferedBatches::new(2, schema(), 8); + assert_eq!(bb.memory_size(), 0); + + let first = batch(vec![1, 2, 3]); + let mut batch_bytes = first.get_array_memory_size(); + bb.push_batch(first, &[vec![0, 2], vec![1]]); + assert_eq!(bb.memory_size(), batch_bytes + index_bytes(&bb)); + + // More rows for partition 0 than its list has room for, so the + // running total must follow a reallocation. + let capacity_before = bb.indices[0].capacity(); + let second = batch((0..8).collect()); + batch_bytes += second.get_array_memory_size(); + bb.push_batch(second, &[(0..8).collect(), vec![]]); + assert!(bb.indices[0].capacity() > capacity_before); + assert_eq!(bb.memory_size(), batch_bytes + index_bytes(&bb)); + + drain(&mut bb); + assert_eq!(bb.memory_size(), 0); } } diff --git a/ballista/core/src/execution_plans/sort_shuffle/mod.rs b/ballista/core/src/execution_plans/sort_shuffle/mod.rs index c6cde972f4..e7d0112246 100644 --- a/ballista/core/src/execution_plans/sort_shuffle/mod.rs +++ b/ballista/core/src/execution_plans/sort_shuffle/mod.rs @@ -38,7 +38,6 @@ mod reader; mod spill; mod writer; -pub use buffer::BufferedBatches; pub use config::SortShuffleConfig; pub use index::ShuffleIndex; pub use reader::{get_index_path, is_sort_shuffle_output, stream_sort_shuffle_partition}; diff --git a/ballista/core/src/execution_plans/sort_shuffle/partitioned_batch_iterator.rs b/ballista/core/src/execution_plans/sort_shuffle/partitioned_batch_iterator.rs index 2b92f414a9..2f44d10b63 100644 --- a/ballista/core/src/execution_plans/sort_shuffle/partitioned_batch_iterator.rs +++ b/ballista/core/src/execution_plans/sort_shuffle/partitioned_batch_iterator.rs @@ -61,10 +61,10 @@ fn compact_view_columns(batch: RecordBatch) -> Result { /// /// Walks a slice of `(batch_idx, row_idx)` pairs in `indices`, and for each /// chunk of up to `batch_size` pairs gathers the referenced rows from -/// `batches` via `interleave_record_batch`, yielding one `RecordBatch` per +/// `batch_refs` via `interleave_record_batch`, yielding one `RecordBatch` per /// chunk. pub(crate) struct PartitionedBatchIterator<'a> { - batch_refs: Vec<&'a RecordBatch>, + batch_refs: &'a [&'a RecordBatch], indices: &'a [(u32, u32)], batch_size: usize, pos: usize, @@ -76,16 +76,16 @@ pub(crate) struct PartitionedBatchIterator<'a> { impl<'a> PartitionedBatchIterator<'a> { /// Create a new iterator. /// - /// * `batches` — the source record batches to pull rows from. + /// * `batch_refs` — the source record batches to pull rows from, shared + /// by every partition's iterator. /// * `indices` — `(batch_idx, row_idx)` pairs that select which rows to /// emit, in order. /// * `batch_size` — maximum number of rows per output `RecordBatch`. pub(crate) fn new( - batches: &'a [RecordBatch], + batch_refs: &'a [&'a RecordBatch], indices: &'a [(u32, u32)], batch_size: usize, ) -> Self { - let batch_refs = batches.iter().collect(); Self { batch_refs, indices, @@ -112,7 +112,7 @@ impl<'a> Iterator for PartitionedBatchIterator<'a> { ); self.pos = end; - match interleave_record_batch(&self.batch_refs, &self.scratch) { + match interleave_record_batch(self.batch_refs, &self.scratch) { Ok(batch) => Some(compact_view_columns(batch)), Err(e) => Some(Err(DataFusionError::ArrowError( Box::new(e), @@ -142,11 +142,12 @@ mod tests { fn yields_chunks_of_at_most_batch_size_rows() { use super::PartitionedBatchIterator; - let batches = vec![batch(&[1, 2, 3, 4]), batch(&[5, 6, 7])]; + let batches = [batch(&[1, 2, 3, 4]), batch(&[5, 6, 7])]; let indices: Vec<(u32, u32)> = vec![(0, 0), (0, 1), (0, 2), (0, 3), (1, 0), (1, 1), (1, 2)]; - let mut iter = PartitionedBatchIterator::new(&batches, &indices, 4); + let refs: Vec<&RecordBatch> = batches.iter().collect(); + let mut iter = PartitionedBatchIterator::new(&refs, &indices, 4); let first = iter.next().unwrap().unwrap(); assert_eq!(first.num_rows(), 4); @@ -173,9 +174,10 @@ mod tests { fn empty_indices_yields_nothing() { use super::PartitionedBatchIterator; - let batches = vec![batch(&[1, 2, 3])]; + let batches = [batch(&[1, 2, 3])]; let indices: Vec<(u32, u32)> = vec![]; - let mut iter = PartitionedBatchIterator::new(&batches, &indices, 4); + let refs: Vec<&RecordBatch> = batches.iter().collect(); + let mut iter = PartitionedBatchIterator::new(&refs, &indices, 4); assert!(iter.next().is_none()); } @@ -183,12 +185,13 @@ mod tests { fn interleaves_rows_from_multiple_input_batches() { use super::PartitionedBatchIterator; - let batches = vec![batch(&[10, 20, 30]), batch(&[40, 50, 60])]; + let batches = [batch(&[10, 20, 30]), batch(&[40, 50, 60])]; // Pull alternating rows from the two input batches let indices: Vec<(u32, u32)> = vec![(0, 0), (1, 0), (0, 1), (1, 1), (0, 2), (1, 2)]; - let mut iter = PartitionedBatchIterator::new(&batches, &indices, 8192); + let refs: Vec<&RecordBatch> = batches.iter().collect(); + let mut iter = PartitionedBatchIterator::new(&refs, &indices, 8192); let only = iter.next().unwrap().unwrap(); assert_eq!(only.num_rows(), 6); let arr = only diff --git a/ballista/core/src/execution_plans/sort_shuffle/writer.rs b/ballista/core/src/execution_plans/sort_shuffle/writer.rs index 142d2136ac..68b1b93576 100644 --- a/ballista/core/src/execution_plans/sort_shuffle/writer.rs +++ b/ballista/core/src/execution_plans/sort_shuffle/writer.rs @@ -38,7 +38,6 @@ use super::super::shuffle_writer_trait::ShuffleWriter; use super::buffer::BufferedBatches; use super::config::SortShuffleConfig; use super::index::ShuffleIndex; -use super::partitioned_batch_iterator::PartitionedBatchIterator; use super::spill::SpillManager; use crate::JobId; use crate::extension::SessionConfigExt; @@ -124,44 +123,27 @@ struct EncodedPartitions { fn encode_buffered_partitions( buffered: &mut BufferedBatches, schema: &SchemaRef, - config: &SortShuffleConfig, opts: &datafusion::arrow::ipc::writer::IpcWriteOptions, ) -> Result { let num_partitions = buffered.num_partitions(); - let (batches, indices) = buffered.take(); - - let mut encoded = Vec::with_capacity(num_partitions); - let mut stats = Vec::with_capacity(num_partitions); - - for partition_indices in indices.iter() { - let mut buf: Vec = Vec::new(); - let (mut num_batches, mut num_rows, mut num_bytes) = (0u64, 0u64, 0u64); - - if !partition_indices.is_empty() { - let iter = PartitionedBatchIterator::new( - &batches, - partition_indices, - config.batch_size, - ); - let mut writer = - StreamWriter::try_new_with_options(&mut buf, schema, opts.clone())?; - for result in iter { - let batch = result?; - num_rows += batch.num_rows() as u64; - num_bytes += batch.get_array_memory_size() as u64; - num_batches += 1; - writer.write(&batch)?; - } - writer.finish()?; + let mut encoded = vec![Vec::new(); num_partitions]; + let mut stats = vec![EncodedStats::default(); num_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.push(buf); - stats.push(EncodedStats { - num_batches, - num_rows, - num_bytes, - }); - } + encoded[partition] = writer.into_inner()?; + Ok(()) + })?; Ok(EncodedPartitions { encoded, stats }) } @@ -620,8 +602,11 @@ impl SortShuffleWriterExec { ) .map_err(|e| DataFusionError::Execution(format!("{e:?}")))?; - let mut buffered = - BufferedBatches::new(num_output_partitions, schema.clone()); + let mut buffered = BufferedBatches::new( + num_output_partitions, + schema.clone(), + config.batch_size, + ); let mut reservation = MemoryConsumer::new(format!("SortShuffleWriter[{input_partition}]")) @@ -635,13 +620,8 @@ impl SortShuffleWriterExec { let partition_reducer = StrengthReducedU64::new(num_output_partitions as u64); let mut spill_events: u64 = 0; let mut spill_triggers = SpillTriggerCounts::default(); - // Absolute buffered-bytes counter, independent of the runtime - // `MemoryPool`. When `memory_limit` is non-zero it caps this counter - // as a second spill trigger; a `memory_limit` of 0 disables the cap - // so spilling is driven solely by memory-pool pressure. - let mut buffered_bytes: usize = 0; - // A limit of 0 disables the per-task budget, leaving the runtime - // `MemoryPool` as the sole spill trigger. + // `memory_limit` is a second spill trigger, independent of the + // `MemoryPool`; 0 disables it. let memory_limit = config.memory_limit_per_task_bytes; let per_task_budget_enabled = memory_limit > 0; let mut null_counts = vec![0u64; schema.fields().len()]; @@ -664,22 +644,14 @@ impl SortShuffleWriterExec { )?; timer.done(); - // Estimate memory growth: input batch + index Vec growth. - let mut growth = input_batch.get_array_memory_size(); - let before = buffered.indices_allocated_size(); buffered.push_batch(input_batch, &per_partition_rows); - let after = buffered.indices_allocated_size(); - growth += after.saturating_sub(before); - - // Mirror the growth in the runtime pool reservation so the pool - // sees this writer's memory usage. A rejected grow means the - // pool is under pressure and has *not* accounted for the batch - // just buffered, so spill instead of holding memory the pool - // believes is free. Spilling flushes that batch to disk and - // frees the reservation, which is the response a spillable - // consumer owes the pool. - let pool_rejected_growth = reservation.try_grow(growth).is_err(); - buffered_bytes = buffered_bytes.saturating_add(growth); + let buffered_bytes = buffered.memory_size(); + + // Mirror the buffered bytes in the pool. A refused resize + // means the pool has not accounted for the rows just + // buffered, so spill them and free the reservation. + let pool_rejected_growth = + reservation.try_resize(buffered_bytes).is_err(); let budget_reached = per_task_budget_enabled && buffered_bytes >= memory_limit; @@ -691,10 +663,8 @@ impl SortShuffleWriterExec { &mut buffered, &mut spill_manager, &mut reservation, - config.batch_size, )?; spill_timer.done(); - buffered_bytes = 0; if event_batches > 0 { spill_events += 1; @@ -771,7 +741,7 @@ impl SortShuffleWriterExec { // compression cost stays parallel across the task's inputs. let opts = create_write_options(compression_type)?; let EncodedPartitions { encoded, stats } = - encode_buffered_partitions(&mut buffered, &schema, &config, &opts)?; + encode_buffered_partitions(&mut buffered, &schema, &opts)?; Ok(InputPartitionOutput { encoded, @@ -785,8 +755,7 @@ impl SortShuffleWriterExec { } } -/// Spills *all* buffered partitions: for each partition, materializes its -/// indices through `PartitionedBatchIterator` and appends each yielded batch +/// Spills *all* buffered partitions: appends each partition's buffered rows /// to that partition's spill file. After this call, `buffered.is_empty()` is /// true and `reservation.size() == 0`. /// @@ -797,28 +766,22 @@ fn spill_all_partitions( buffered: &mut BufferedBatches, spill_manager: &mut SpillManager, reservation: &mut MemoryReservation, - batch_size: usize, ) -> Result<(u64, u64)> { if buffered.is_empty() { return Ok((0, 0)); } let mut batches_written: u64 = 0; let mut bytes_written: u64 = 0; - let (batches, indices) = buffered.take(); - for (partition_id, partition_indices) in indices.iter().enumerate() { - if partition_indices.is_empty() { - continue; - } - let iter = PartitionedBatchIterator::new(&batches, partition_indices, batch_size); - for result in iter { - let batch = result?; + buffered.drain(|partition_id, batches| { + for batch in batches { let written = spill_manager - .spill(partition_id, &batch) + .spill(partition_id, &batch?) .map_err(|e| DataFusionError::Execution(format!("{e:?}")))?; batches_written += 1; bytes_written += written; } - } + Ok(()) + })?; reservation.free(); Ok((batches_written, bytes_written)) } diff --git a/docs/source/upgrading/56.0.0.md b/docs/source/upgrading/56.0.0.md index 354e9e9d97..f40a7f5ceb 100644 --- a/docs/source/upgrading/56.0.0.md +++ b/docs/source/upgrading/56.0.0.md @@ -134,6 +134,14 @@ early, so a straggler partition leaves them idle. Set `ballista.scheduler.max_partitions_per_task` to a positive value to dispatch smaller tasks, which release their vcores sooner. +### API changes + +#### `sort_shuffle::BufferedBatches` is no longer public + +`ballista_core::execution_plans::sort_shuffle::BufferedBatches`, the sort-shuffle +writer's internal row buffer, is no longer exported. It had no use outside the +writer. Code that used it directly should use `SortShuffleWriterExec` instead. + ### Python client changes #### Summary