diff --git a/datafusion/physical-plan/src/aggregates/clustered_final_stream.rs b/datafusion/physical-plan/src/aggregates/clustered_final_stream.rs index 38cb99e71318d..6eeb37a475ffe 100644 --- a/datafusion/physical-plan/src/aggregates/clustered_final_stream.rs +++ b/datafusion/physical-plan/src/aggregates/clustered_final_stream.rs @@ -40,8 +40,8 @@ use crate::stream::{ObservedStream, RecordBatchStreamAdapter}; /// extra index array is used in later sorting before spilling. /// - On memory pressure, materialize all group states into one batch. /// - Use [`IncrementalSortIterator`](crate::sorts::IncrementalSortIterator) to compute the full-batch index, then -/// materialize and write one sorted `batch_size` slice at a time. The original -/// batch and full index remain live until the run is written. +/// materialize and write one sorted slice at a time. The original batch and +/// full index remain live until the run is written. /// - After input ends, merge the sorted runs and replay them through a fully /// ordered final aggregate stream. pub(crate) struct ClusteredFinalAggregateStream { diff --git a/datafusion/physical-plan/src/aggregates/spill.rs b/datafusion/physical-plan/src/aggregates/spill.rs index 9da8da2eae62a..c08a76adc2153 100644 --- a/datafusion/physical-plan/src/aggregates/spill.rs +++ b/datafusion/physical-plan/src/aggregates/spill.rs @@ -14,6 +14,7 @@ use datafusion_execution::memory_pool::MemoryReservation; use datafusion_physical_expr::PhysicalSortExpr; use datafusion_physical_expr::expressions::Column; use datafusion_physical_expr_common::sort_expr::LexOrdering; +use futures::StreamExt; use super::aggregate_hash_table::ClusteredAggregateTableMetrics; use super::clustered_final_stream::ClusteredFinalAggregateStream; @@ -23,7 +24,19 @@ use crate::SendableRecordBatchStream; use crate::metrics::{BaselineMetrics, SpillMetrics}; use crate::sorts::IncrementalSortIterator; use crate::sorts::streaming_merge::{SortedSpillFile, StreamingMergeBuilder}; +use crate::spill::gc_nested_arrays; use crate::spill::spill_manager::SpillManager; +use crate::stream::RecordBatchStreamAdapter; + +/// Target size in bytes of the batches that spilling writes and replay +/// merges. A batch can be larger when it holds a larger group. +/// +/// Replay holds the groups of the batch it merges in memory, and merging +/// several spill files reserves memory for a few batches of each. A limit of +/// `batch_size` rows alone does not bound the size of a batch: a few groups +/// with large state, such as `array_agg` over a low-cardinality key, would fit +/// in a single batch. +const SPILL_BATCH_TARGET_BYTES: usize = 1024 * 1024; /// Spill configuration and accumulated runs of one grouped aggregation stream. /// @@ -210,12 +223,18 @@ impl AggregateSpill { return Ok(()); }; - let max_batch_rows = state_batch.num_rows().min(self.batch_size); + let mut max_batch_rows = 0; let sorted_iter = IncrementalSortIterator::new( state_batch, self.spill_expr.clone(), self.batch_size, - ); + ) + .with_max_batch_bytes(SPILL_BATCH_TARGET_BYTES) + .inspect(|batch| { + if let Ok(batch) = batch { + max_batch_rows = max_batch_rows.max(batch.num_rows()); + } + }); let spill_file = self .spill_manager .spill_record_batch_iter_and_return_max_batch_memory( @@ -289,10 +308,20 @@ impl AggregateSpill { .with_expressions(&spill_expr) .with_metrics(baseline_metrics.intermediate()) .with_batch_size(batch_size) + .with_max_batch_bytes(Some(SPILL_BATCH_TARGET_BYTES)) .with_reservation(merge_reservation) .with_replay_headroom() .with_intermediate_merge_sizing(Some(min_spill_batch_rows)) .build()?; + // A merged batch references the whole data buffers of the spilled + // batches its rows come from, and can hold all their dictionary + // values. Replay keeps the list states of a group, such as those of + // `array_agg`, until the group completes, so compact them to hold only + // the bytes of their own rows. + let merged: SendableRecordBatchStream = Box::pin(RecordBatchStreamAdapter::new( + merged.schema(), + merged.map(|batch| gc_nested_arrays(&batch?)), + )); let replay = ClusteredFinalAggregateStream::new_with_input_and_metrics( &replay_agg, &context, diff --git a/datafusion/physical-plan/src/sorts/builder.rs b/datafusion/physical-plan/src/sorts/builder.rs index 93b7622f40e49..97abe7b62ef5a 100644 --- a/datafusion/physical-plan/src/sorts/builder.rs +++ b/datafusion/physical-plan/src/sorts/builder.rs @@ -85,8 +85,8 @@ impl BatchBuilder { } } - /// Append a new batch in `stream_idx` - pub fn push_batch(&mut self, stream_idx: usize, batch: RecordBatch) -> Result<()> { + /// Append a new batch in `stream_idx`, returning its memory size + pub fn push_batch(&mut self, stream_idx: usize, batch: RecordBatch) -> Result { let size = get_record_batch_memory_size(&batch); self.batches_mem_used += size; // Only request additional memory from the pool when actual batch @@ -99,7 +99,7 @@ impl BatchBuilder { batch_idx, row_idx: 0, }; - Ok(()) + Ok(size) } /// Append the next row from `stream_idx` diff --git a/datafusion/physical-plan/src/sorts/merge.rs b/datafusion/physical-plan/src/sorts/merge.rs index 8e900372121bd..5bc5bd348af26 100644 --- a/datafusion/physical-plan/src/sorts/merge.rs +++ b/datafusion/physical-plan/src/sorts/merge.rs @@ -73,6 +73,16 @@ pub(crate) struct SortPreservingMergeStream { /// Target batch size batch_size: usize, + /// See [`Self::with_max_batch_bytes`] + max_batch_bytes: Option, + + /// The average memory size of the rows of each input's current batch + avg_row_bytes: Vec, + + /// The estimated memory size of the in-progress rows, tracked when + /// `max_batch_bytes` is set + in_progress_bytes: usize, + /// Cursors for each input partition. `None` means the input is exhausted cursors: Vec>>, @@ -133,11 +143,22 @@ impl SortPreservingMergeStream { poll_reset_epochs: vec![0; stream_count], loser_tree: vec![], batch_size, + max_batch_bytes: None, + avg_row_bytes: vec![0; stream_count], + in_progress_bytes: 0, fetch, produced: 0, } } + /// Also emits a batch once its rows take about `max_batch_bytes`, as + /// estimated from the average row size of the batches they come from, + /// before it has `batch_size` rows. + pub(crate) fn with_max_batch_bytes(mut self, max_batch_bytes: Option) -> Self { + self.max_batch_bytes = max_batch_bytes; + self + } + pub(crate) fn into_stream(self) -> SendableRecordBatchStream where C: 'static, @@ -171,7 +192,13 @@ impl SortPreservingMergeStream { Some(Err(e)) => Poll::Ready(Err(e)), Some(Ok((cursor, batch))) => { self.cursors[idx] = Some(Cursor::new(cursor)); - Poll::Ready(self.in_progress.push_batch(idx, batch)) + let num_rows = batch.num_rows(); + let size = match self.in_progress.push_batch(idx, batch) { + Ok(size) => size, + Err(e) => return Poll::Ready(Err(e)), + }; + self.avg_row_bytes[idx] = size.div_ceil(num_rows.max(1)); + Poll::Ready(Ok(())) } } } @@ -179,7 +206,10 @@ impl SortPreservingMergeStream { fn emit_in_progress_batch(&mut self) -> Result> { let rows_before = self.in_progress.len(); let result = self.in_progress.build_record_batch(); - self.produced += rows_before - self.in_progress.len(); + let rows_after = self.in_progress.len(); + self.produced += rows_before - rows_after; + // Rows kept after an offset overflow are assumed to be of average size + self.in_progress_bytes = self.in_progress_bytes * rows_after / rows_before.max(1); result } @@ -237,6 +267,9 @@ impl SortPreservingMergeStream { // 3.1. add loser_tree[0] (minimum) stream to pending record batch let winner_stream = self.loser_tree[0]; self.in_progress.push_row(winner_stream); + if self.max_batch_bytes.is_some() { + self.in_progress_bytes += self.avg_row_bytes[winner_stream]; + } // 3.2. If the new row reached the limit if self.fetch_reached() { @@ -244,7 +277,11 @@ impl SortPreservingMergeStream { } // 3.3. if there is enough to emit for a full record batch - if self.in_progress.len() >= self.batch_size { + if self.in_progress.len() >= self.batch_size + || self + .max_batch_bytes + .is_some_and(|max_bytes| self.in_progress_bytes >= max_bytes) + { // 3.3.1 build pending record batch and reset builder let Some(batch) = self.emit_in_progress_batch()? else { return internal_err!("must have batch in progress to emit"); diff --git a/datafusion/physical-plan/src/sorts/multi_level_merge.rs b/datafusion/physical-plan/src/sorts/multi_level_merge.rs index 424d1faffcf33..ebdecb90f38fa 100644 --- a/datafusion/physical-plan/src/sorts/multi_level_merge.rs +++ b/datafusion/physical-plan/src/sorts/multi_level_merge.rs @@ -141,6 +141,8 @@ pub(crate) struct MultiLevelMergeBuilder { expr: LexOrdering, metrics: BaselineMetrics, batch_size: usize, + /// See [`StreamingMergeBuilder::with_max_batch_bytes`] + max_batch_bytes: Option, reservation: MemoryReservation, /// Workspace retained across retries and intermediate spill passes. merge_pool: Option>, @@ -185,6 +187,7 @@ impl MultiLevelMergeBuilder { expr, metrics, batch_size, + max_batch_bytes: None, reservation, merge_pool: None, reserve_replay_headroom: false, @@ -199,6 +202,13 @@ impl MultiLevelMergeBuilder { self } + /// Applies [`StreamingMergeBuilder::with_max_batch_bytes`] to every pass, + /// including the passes whose output is spilled again. + pub(super) fn with_max_batch_bytes(mut self, max_batch_bytes: Option) -> Self { + self.max_batch_bytes = max_batch_bytes; + self + } + /// Leave replay headroom while selecting merge buffers. Temporary splitting /// workspace can still use the full pool because replay has not started. pub(super) fn with_replay_headroom(mut self, reserve: bool) -> Self { @@ -529,7 +539,8 @@ impl MultiLevelMergeBuilder { self.metrics.intermediate() }) .with_round_robin_tie_breaker(self.enable_round_robin_tie_breaker) - .with_streams(streams); + .with_streams(streams) + .with_max_batch_bytes(self.max_batch_bytes); if !all_in_memory { // Don't track memory used by this stream as we reserve that memory by worst case sceneries diff --git a/datafusion/physical-plan/src/sorts/stream.rs b/datafusion/physical-plan/src/sorts/stream.rs index bdd3cb16dee6e..246237fed85b1 100644 --- a/datafusion/physical-plan/src/sorts/stream.rs +++ b/datafusion/physical-plan/src/sorts/stream.rs @@ -3,6 +3,8 @@ // SPDX-FileContributor: Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements use crate::sorts::cursor::{ArrayValues, CursorArray, RowValues}; +use crate::spill::gc_columns; +use crate::spill::spilled_row_sizes::SpilledRowSizes; use crate::{EmptyRecordBatchStream, SendableRecordBatchStream}; use crate::{PhysicalExpr, PhysicalSortExpr}; use arrow::array::{Array, UInt32Array}; @@ -328,12 +330,13 @@ impl PartitionedStream for FieldCursorStream { /// peak memory to hold both the unsorted and sorted copies simultaneously). /// /// On the first call to `next()`, a sorted index array (`UInt32Array`) is -/// computed via `lexsort_to_indices`. Subsequent calls yield chunks of -/// `batch_size` rows by `take`-ing from the original batch using slices of -/// this index array. Each `take` copies data for the chunk (not zero-copy), -/// but only one chunk is live at a time since the caller consumes it before -/// requesting the next. Once all rows have been yielded, the original batch -/// and index array are dropped to free memory. +/// computed via `lexsort_to_indices`. Subsequent calls yield chunks of up to +/// `batch_size` rows, fewer if [`Self::with_max_batch_bytes`] is set, by +/// `take`-ing from the original batch using slices of this index array. Each +/// `take` copies data for the chunk (not zero-copy), but only one chunk is +/// live at a time since the caller consumes it before requesting the next. +/// Once all rows have been yielded, the original batch and index array are +/// dropped to free memory. /// /// The caller must reserve `sizeof(batch) + sizeof(one chunk)` for this iterator, /// and free the reservation once the iterator is depleted. @@ -341,7 +344,11 @@ pub(crate) struct IncrementalSortIterator { batch: RecordBatch, expressions: LexOrdering, batch_size: usize, + /// See [`Self::with_max_batch_bytes`]. + max_batch_bytes: Option, indices: Option, + /// Row sizes of `batch`, measured with `indices` when `max_batch_bytes` is set + row_sizes: Option, cursor: usize, } @@ -355,10 +362,47 @@ impl IncrementalSortIterator { batch, expressions, batch_size, + max_batch_bytes: None, cursor: 0, indices: None, + row_sizes: None, } } + + /// Also limits each chunk to `max_batch_bytes`, by the size its rows take + /// once spilled (see [`SpilledRowSizes`]). A row larger than that is + /// returned on its own. Chunks hold only the dictionary values of their + /// own rows in the columns where that size counts them (see + /// [`SpilledRowSizes::compacted_columns`]). + pub(crate) fn with_max_batch_bytes(mut self, max_batch_bytes: usize) -> Self { + self.max_batch_bytes = Some(max_batch_bytes); + self + } + + /// Returns how many of the rows at `indices` fit in `max_batch_bytes`. + fn rows_within(&self, indices: &[u32]) -> usize { + let (Some(max_batch_bytes), Some(row_sizes)) = + (self.max_batch_bytes, &self.row_sizes) + else { + return indices.len(); + }; + if let Some(row_bytes) = row_sizes.fixed() { + return (max_batch_bytes / row_bytes.max(1)) + .max(1) + .min(indices.len()); + } + if row_sizes.max_row().saturating_mul(indices.len()) <= max_batch_bytes { + return indices.len(); + } + let mut bytes = 0; + for (rows, &row) in indices.iter().enumerate() { + bytes += row_sizes.row(row as usize); + if rows > 0 && bytes > max_batch_bytes { + return rows; + } + } + indices.len() + } } impl Iterator for IncrementalSortIterator { @@ -386,12 +430,18 @@ impl Iterator for IncrementalSortIterator { Err(e) => return Some(Err(e.into())), }; self.indices = Some(indices); + if self.max_batch_bytes.is_some() { + self.row_sizes = Some(SpilledRowSizes::new(&self.batch)); + } // Call again, this time it will hit the Some(indices) branch and return the first batch self.next() } Some(indices) => { let batch_size = self.batch_size.min(self.batch.num_rows() - self.cursor); + let batch_size = self.rows_within( + &indices.values()[self.cursor..self.cursor + batch_size], + ); // Perform the take to produce the next batch let new_batch_indices = indices.slice(self.cursor, batch_size); @@ -399,6 +449,17 @@ impl Iterator for IncrementalSortIterator { Ok(batch) => batch, Err(e) => return Some(Err(e.into())), }; + // `take` keeps every dictionary value, while the row sizes + // count only the values of the rows of `compacted_columns` + let new_batch = match &self.row_sizes { + Some(row_sizes) if !row_sizes.compacted_columns().is_empty() => { + match gc_columns(&new_batch, row_sizes.compacted_columns()) { + Ok(batch) => batch, + Err(e) => return Some(Err(e)), + } + } + _ => new_batch, + }; self.cursor += batch_size; @@ -407,6 +468,7 @@ impl Iterator for IncrementalSortIterator { let schema = self.batch.schema(); let _ = mem::replace(&mut self.batch, RecordBatch::new_empty(schema)); self.indices = None; + self.row_sizes = None; } // Return the new batch @@ -421,7 +483,13 @@ impl Iterator for IncrementalSortIterator { let num_rows = self.batch.num_rows().saturating_sub(self.cursor); let batch_size = self.batch_size; let num_batches = num_rows.div_ceil(batch_size); - (num_batches, Some(num_batches)) + // Limiting bytes can leave as few as one row per batch + let max_batches = if self.max_batch_bytes.is_some() { + num_rows + } else { + num_batches + }; + (num_batches, Some(max_batches)) } } @@ -431,7 +499,7 @@ impl FusedIterator for IncrementalSortIterator {} mod tests { use super::*; use crate::memory::MemoryStream; - use arrow::array::{AsArray, Int32Array, StringArray}; + use arrow::array::{ArrayRef, AsArray, DictionaryArray, Int32Array, StringArray}; use arrow::datatypes::{DataType, Field, Int32Type}; use arrow_schema::SchemaRef; use datafusion_common::DataFusionError; @@ -591,6 +659,99 @@ mod tests { } } + #[test] + fn incremental_sort_iterator_limits_batch_bytes() -> Result<()> { + // Keys 100 down to 0. Keys 0..5 have 1000 byte values, the others 10 + // bytes, so a row takes 1008 or 18 bytes with the Int32 key and the + // 4 byte offset. + let keys: Vec = (0..=100).rev().collect(); + let values: Vec = keys + .iter() + .map(|&key| "x".repeat(if key < 5 { 1000 } else { 10 })) + .collect(); + let batch = RecordBatch::try_from_iter(vec![ + ("key", Arc::new(Int32Array::from(keys)) as ArrayRef), + ("value", Arc::new(StringArray::from(values)) as ArrayRef), + ])?; + let expressions = LexOrdering::new(vec![PhysicalSortExpr::new_default(col( + "key", + &batch.schema(), + )?)]) + .unwrap(); + + let chunk_lens = |iter: IncrementalSortIterator| -> Result> { + let mut next_key = 0; + iter.map(|chunk| { + let chunk = chunk?; + let keys = chunk.column(0).as_primitive::(); + for key in keys.values() { + assert_eq!(*key, next_key, "rows must stay sorted"); + next_key += 1; + } + Ok(chunk.num_rows()) + }) + .collect() + }; + + let iter = IncrementalSortIterator::new(batch.clone(), expressions.clone(), 50); + assert_eq!(chunk_lens(iter)?, vec![50, 50, 1]); + + // Two large rows fit in 2500 bytes. The last large row leaves room for + // many small rows, up to 50 rows per batch. + let iter = IncrementalSortIterator::new(batch.clone(), expressions.clone(), 50) + .with_max_batch_bytes(2500); + assert_eq!(chunk_lens(iter)?, vec![2, 2, 50, 47]); + + // Without the values, every row takes the 4 bytes of its key + let keys = batch.project(&[0])?; + let iter = + IncrementalSortIterator::new(keys, expressions, 50).with_max_batch_bytes(100); + assert_eq!(chunk_lens(iter)?, vec![25, 25, 25, 25, 1]); + Ok(()) + } + + #[test] + fn incremental_sort_iterator_compacts_unique_dictionaries() -> Result<()> { + // Keys 99 down to 0, each with its own 100 byte value in `unique`, and + // one of four values, each of 25 keys, in `shared` + let keys: Vec = (0..100).rev().collect(); + let values: Vec = keys.iter().map(|key| format!("{key:0>100}")).collect(); + let unique: DictionaryArray = + values.iter().map(String::as_str).collect(); + let shared: DictionaryArray = keys + .iter() + .map(|key| ["a", "b", "c", "d"][*key as usize / 25]) + .collect(); + let shared_values = shared.values().to_data(); + let batch = RecordBatch::try_from_iter(vec![ + ("key", Arc::new(Int32Array::from(keys)) as ArrayRef), + ("unique", Arc::new(unique) as ArrayRef), + ("shared", Arc::new(shared) as ArrayRef), + ])?; + let expressions = LexOrdering::new(vec![PhysicalSortExpr::new_default(col( + "key", + &batch.schema(), + )?)]) + .unwrap(); + + // A row takes 4 bytes of key and 4 of `shared` key, plus a 4 byte key, + // a 4 byte offset and 100 bytes of value in `unique`, so 10 rows fit + // in 1200 bytes + let chunks = IncrementalSortIterator::new(batch, expressions, 50) + .with_max_batch_bytes(1200) + .collect::>>()?; + assert_eq!(chunks.len(), 10); + for chunk in &chunks { + // Each chunk holds the unique values of its own rows, and shares + // the values of the other dictionary + let unique = chunk.column(1).as_any_dictionary(); + assert_eq!(unique.values().len(), chunk.num_rows()); + let shared = chunk.column(2).as_any_dictionary(); + assert!(shared.values().to_data().ptr_eq(&shared_values)); + } + Ok(()) + } + fn assert_iterator_size_hint(iter: &IncrementalSortIterator, expected_len: usize) { assert_eq!(iter.size_hint(), (expected_len, Some(expected_len))); } diff --git a/datafusion/physical-plan/src/sorts/streaming_merge.rs b/datafusion/physical-plan/src/sorts/streaming_merge.rs index da341cf0cf047..2e5e0790accec 100644 --- a/datafusion/physical-plan/src/sorts/streaming_merge.rs +++ b/datafusion/physical-plan/src/sorts/streaming_merge.rs @@ -30,7 +30,7 @@ macro_rules! primitive_merge_helper { } macro_rules! merge_helper { - ($t:ty, $sort:ident, $streams:ident, $schema:ident, $tracking_metrics:ident, $batch_size:ident, $fetch:ident, $reservation:ident, $enable_round_robin_tie_breaker:ident) => {{ + ($t:ty, $sort:ident, $streams:ident, $schema:ident, $tracking_metrics:ident, $batch_size:ident, $max_batch_bytes:ident, $fetch:ident, $reservation:ident, $enable_round_robin_tie_breaker:ident) => {{ let streams = FieldCursorStream::<$t>::new($sort, $streams, $reservation.new_empty()); return Ok(SortPreservingMergeStream::new( @@ -42,6 +42,7 @@ macro_rules! merge_helper { $reservation, $enable_round_robin_tie_breaker, ) + .with_max_batch_bytes($max_batch_bytes) .into_stream()); }}; } @@ -80,6 +81,7 @@ pub struct StreamingMergeBuilder<'a> { expressions: Option<&'a LexOrdering>, metrics: Option, batch_size: Option, + max_batch_bytes: Option, fetch: Option, reservation: Option, merge_pool: Option>, @@ -135,6 +137,12 @@ impl<'a> StreamingMergeBuilder<'a> { self } + /// See [`SortPreservingMergeStream::with_max_batch_bytes`] + pub(crate) fn with_max_batch_bytes(mut self, max_batch_bytes: Option) -> Self { + self.max_batch_bytes = max_batch_bytes; + self + } + pub fn with_fetch(mut self, fetch: Option) -> Self { self.fetch = fetch; self @@ -202,6 +210,7 @@ impl<'a> StreamingMergeBuilder<'a> { schema, metrics, batch_size, + max_batch_bytes, reservation, merge_pool, reserve_replay_headroom, @@ -249,6 +258,7 @@ impl<'a> StreamingMergeBuilder<'a> { .with_merge_pool(merge_pool) .with_replay_headroom(reserve_replay_headroom) .with_intermediate_merge_sizing(min_spill_batch_rows) + .with_max_batch_bytes(max_batch_bytes) .create_spillable_merge_stream()); } @@ -268,12 +278,12 @@ impl<'a> StreamingMergeBuilder<'a> { let sort = expressions[0].clone(); let data_type = sort.expr.data_type(schema.as_ref())?; downcast_primitive! { - data_type => (primitive_merge_helper, sort, streams, schema, metrics, batch_size, fetch, reservation, enable_round_robin_tie_breaker), - DataType::Utf8 => merge_helper!(StringArray, sort, streams, schema, metrics, batch_size, fetch, reservation, enable_round_robin_tie_breaker) - DataType::Utf8View => merge_helper!(StringViewArray, sort, streams, schema, metrics, batch_size, fetch, reservation, enable_round_robin_tie_breaker) - DataType::LargeUtf8 => merge_helper!(LargeStringArray, sort, streams, schema, metrics, batch_size, fetch, reservation, enable_round_robin_tie_breaker) - DataType::Binary => merge_helper!(BinaryArray, sort, streams, schema, metrics, batch_size, fetch, reservation, enable_round_robin_tie_breaker) - DataType::LargeBinary => merge_helper!(LargeBinaryArray, sort, streams, schema, metrics, batch_size, fetch, reservation, enable_round_robin_tie_breaker) + data_type => (primitive_merge_helper, sort, streams, schema, metrics, batch_size, max_batch_bytes, fetch, reservation, enable_round_robin_tie_breaker), + DataType::Utf8 => merge_helper!(StringArray, sort, streams, schema, metrics, batch_size, max_batch_bytes, fetch, reservation, enable_round_robin_tie_breaker) + DataType::Utf8View => merge_helper!(StringViewArray, sort, streams, schema, metrics, batch_size, max_batch_bytes, fetch, reservation, enable_round_robin_tie_breaker) + DataType::LargeUtf8 => merge_helper!(LargeStringArray, sort, streams, schema, metrics, batch_size, max_batch_bytes, fetch, reservation, enable_round_robin_tie_breaker) + DataType::Binary => merge_helper!(BinaryArray, sort, streams, schema, metrics, batch_size, max_batch_bytes, fetch, reservation, enable_round_robin_tie_breaker) + DataType::LargeBinary => merge_helper!(LargeBinaryArray, sort, streams, schema, metrics, batch_size, max_batch_bytes, fetch, reservation, enable_round_robin_tie_breaker) _ => {} } } @@ -293,6 +303,7 @@ impl<'a> StreamingMergeBuilder<'a> { reservation, enable_round_robin_tie_breaker, ) + .with_max_batch_bytes(max_batch_bytes) .into_stream()) } } @@ -305,6 +316,7 @@ mod tests { use super::*; use arrow::array::{ArrayRef, RecordBatch}; + use arrow::datatypes::{Field, Int32Type, Schema}; use arrow_schema::SortOptions; use datafusion_common::Result; use datafusion_execution::TaskContext; @@ -313,6 +325,108 @@ mod tests { ExecutionPlanMetricsSet, SpillMetrics, }; + /// Merges a spill file of five 1000 byte rows, one per batch, with a spill + /// file of 1000 small rows in batches of 100. + async fn merge_wide_and_narrow_rows( + max_batch_bytes: Option, + ) -> Result> { + let schema = Arc::new(Schema::new(vec![ + Field::new("key", DataType::Int32, false), + Field::new("value", DataType::Utf8, false), + ])); + let batch = |keys: std::ops::Range, value_len: usize| { + let values = keys.clone().map(|_| "x".repeat(value_len)); + RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int32Array::from_iter_values(keys)), + Arc::new(StringArray::from_iter_values(values)), + ], + ) + .unwrap() + }; + let wide: Vec<_> = (0..5).map(|key| batch(key..key + 1, 1000)).collect(); + let narrow: Vec<_> = (0..10) + .map(|i| batch(10 + i * 100..10 + (i + 1) * 100, 1)) + .collect(); + + let sort: LexOrdering = + [PhysicalSortExpr::new_default(col("key", &schema)?)].into(); + let spill_manager = SpillManager::new( + Arc::new(TaskContext::default()).runtime_env(), + SpillMetrics::new(&ExecutionPlanMetricsSet::new(), 0), + Arc::clone(&schema), + ); + let spill = |batches: &[RecordBatch]| SortedSpillFile { + file: spill_manager + .spill_record_batch_and_finish(batches, "test") + .unwrap() + .unwrap(), + max_record_batch_memory: batches + .iter() + .map(RecordBatch::get_array_memory_size) + .max() + .unwrap(), + }; + let sorted_spill_files = vec![spill(&wide), spill(&narrow)]; + let merged = StreamingMergeBuilder::new() + .with_schema(Arc::clone(&schema)) + .with_expressions(&sort) + .with_metrics(BaselineMetrics::new(&ExecutionPlanMetricsSet::new(), 0)) + .with_batch_size(100) + .with_max_batch_bytes(max_batch_bytes) + .with_bypass_mempool() + .with_sorted_spill_files(sorted_spill_files) + .with_spill_manager(spill_manager) + .build()?; + let batches = collect(merged).await?; + + let keys: Vec = batches + .iter() + .flat_map(|batch| { + batch + .column(0) + .as_primitive::() + .values() + .to_vec() + }) + .collect(); + let expected: Vec = (0..5).chain(10..1010).collect(); + assert_eq!(keys, expected); + Ok(batches) + } + + #[tokio::test] + async fn test_max_batch_bytes_limits_merged_batches() -> Result<()> { + let num_rows = |batches: &[RecordBatch]| -> Vec { + batches.iter().map(RecordBatch::num_rows).collect() + }; + // By rows only, the five large rows end up in the first batch + let batches = merge_wide_and_narrow_rows(None).await?; + assert_eq!(num_rows(&batches), [vec![100; 10], vec![5]].concat()); + + // Three large rows exceed 3000 bytes, so the first batch holds fewer. + // Once the large rows are merged, batches of small rows still hold + // `batch_size` rows. + let batches = merge_wide_and_narrow_rows(Some(3000)).await?; + let rows = num_rows(&batches); + assert!(rows[0] < 5, "{rows:?}"); + let last_wide = batches + .iter() + .position(|batch| { + batch + .column(0) + .as_primitive::() + .values() + .contains(&4) + }) + .unwrap(); + let narrow_only = &rows[last_wide + 1..rows.len() - 1]; + assert!(narrow_only.len() > 5, "{rows:?}"); + assert!(narrow_only.iter().all(|&rows| rows == 100), "{rows:?}"); + Ok(()) + } + #[tokio::test] async fn test_sort_merge_fetch_zero_with_only_1_stream() { test_fetch_0_should_output_0_rows(1, 0).await.unwrap(); diff --git a/datafusion/physical-plan/src/spill/mod.rs b/datafusion/physical-plan/src/spill/mod.rs index c6b63ca287161..f2ec70c6f3544 100644 --- a/datafusion/physical-plan/src/spill/mod.rs +++ b/datafusion/physical-plan/src/spill/mod.rs @@ -9,6 +9,7 @@ mod lengths; pub(crate) mod replayable_spill_input; pub(crate) mod spill_manager; pub mod spill_pool; +pub(crate) mod spilled_row_sizes; use datafusion_execution::spill_file::{AsyncSpillWriter, SpillWriter}; // Moved for refactor, re-export to keep the public API stable pub use datafusion_common::utils::memory::get_record_batch_memory_size; @@ -803,7 +804,6 @@ fn get_max_alignment_for_schema(schema: &Schema) -> usize { /// Size of a single view structure in StringView/BinaryView arrays (in bytes). /// Each view is 16 bytes: 4 bytes length + 4 bytes prefix + 8 bytes buffer ID/offset. -#[cfg(test)] const VIEW_SIZE_BYTES: usize = 16; /// Compacts view buffers and oversized dictionaries before spilling. @@ -848,11 +848,48 @@ const VIEW_SIZE_BYTES: usize = 16; /// unused ListView child entries can still retain dictionary values /// - The Arrow `gc()` method itself is optimized and only copies referenced data pub(crate) fn gc_view_arrays(batch: &RecordBatch) -> Result { + gc_arrays_of_columns(batch, |_, _| true, false) +} + +/// Like [`gc_view_arrays`], but only compacts `columns`, and also removes the +/// dictionary values that no key of theirs references. +/// +/// `take` and `slice` keep every value of a dictionary, so a few rows taken +/// from a large batch carry the values of all its rows. +pub(crate) fn gc_columns(batch: &RecordBatch, columns: &[usize]) -> Result { + gc_arrays_of_columns(batch, |index, _| columns.contains(&index), true) +} + +/// Like [`gc_view_arrays`], but only compacts nested columns, such as lists of +/// strings, and also removes the dictionary values that no key of theirs +/// references. +/// +/// A batch built by `interleave`, such as the output of a sort-preserving +/// merge, references the whole data buffers of the batches it takes rows +/// from, and can hold all their dictionary values. An accumulator that keeps +/// the arrays it merges, such as `array_agg`, keeps them alive with each +/// group. Top-level columns are left as they are, since group keys and the +/// accumulators of top-level states typically copy the values they keep. +pub(crate) fn gc_nested_arrays(batch: &RecordBatch) -> Result { + gc_arrays_of_columns(batch, |_, array| array.data_type().is_nested(), true) +} + +/// Compacts the view arrays of the columns of `batch` that `include` selects +/// by index, and their dictionary arrays if `dictionaries` is set. +fn gc_arrays_of_columns( + batch: &RecordBatch, + include: impl Fn(usize, &ArrayRef) -> bool, + dictionaries: bool, +) -> Result { let mut mutated = false; let mut new_columns: Vec> = Vec::with_capacity(batch.num_columns()); - for array in batch.columns() { - let (gc_array, array_mutated) = gc_array(array)?; + for (index, array) in batch.columns().iter().enumerate() { + let (gc_array, array_mutated) = if include(index, array) { + gc_array(array, dictionaries)? + } else { + (Arc::clone(array), false) + }; mutated |= array_mutated; new_columns.push(gc_array); } @@ -864,7 +901,9 @@ pub(crate) fn gc_view_arrays(batch: &RecordBatch) -> Result { } } -fn gc_array(array: &ArrayRef) -> Result<(ArrayRef, bool)> { +/// Compacts the view arrays in `array`, and its dictionary arrays if +/// `dictionaries` is set. Also returns whether anything was compacted. +fn gc_array(array: &ArrayRef, dictionaries: bool) -> Result<(ArrayRef, bool)> { match array.data_type() { DataType::Utf8View => { let string_view = array @@ -886,16 +925,31 @@ fn gc_array(array: &ArrayRef) -> Result<(ArrayRef, bool)> { None => (Arc::clone(array), false), }) } + DataType::Dictionary(_, _) if dictionaries => { + let dictionary = array.as_any_dictionary(); + let gc = garbage_collect_any_dictionary(dictionary)?; + let gc = gc.as_any_dictionary(); + let shrunk = gc.values().len() < dictionary.values().len(); + let (values, values_mutated) = gc_array(gc.values(), dictionaries)?; + if !shrunk && !values_mutated { + return Ok((Arc::clone(array), false)); + } + // `gc` holds only the values its rows reference, so there are no more + // than it has rows. + #[expect(clippy::disallowed_methods)] + let compacted = gc.with_values(values); + Ok((compacted, true)) + } DataType::Dictionary(_, _) if array.as_any_dictionary().values().len() > array.len() => { // Struct-valued merge dictionaries can concatenate whole value tables. // Prune the outer dictionary before compacting its retained children. let compacted = garbage_collect_any_dictionary(array.as_any_dictionary())?; - let (compacted, _) = gc_array_children(&compacted)?; + let (compacted, _) = gc_array_children(&compacted, dictionaries)?; Ok((compacted, true)) } - _ => gc_array_children(array), + _ => gc_array_children(array, dictionaries), } } @@ -980,7 +1034,7 @@ fn dedup_buffer_entries(buffers: &[Buffer]) -> (Vec, Vec) { (remap, unique) } -fn gc_array_children(array: &ArrayRef) -> Result<(ArrayRef, bool)> { +fn gc_array_children(array: &ArrayRef, dictionaries: bool) -> Result<(ArrayRef, bool)> { let data = array.to_data(); if data.child_data().is_empty() { return Ok((Arc::clone(array), false)); @@ -990,7 +1044,7 @@ fn gc_array_children(array: &ArrayRef) -> Result<(ArrayRef, bool)> { let mut child_data = Vec::with_capacity(data.child_data().len()); for child in data.child_data() { let child_array = make_array(child.clone()); - let (gc_child, child_mutated) = gc_array(&child_array)?; + let (gc_child, child_mutated) = gc_array(&child_array, dictionaries)?; mutated |= child_mutated; child_data.push(gc_child.to_data()); } @@ -2789,6 +2843,123 @@ mod tests { Ok(()) } + #[test] + fn test_gc_nested_arrays_skips_top_level_columns() -> Result<()> { + use arrow::buffer::Buffer; + + let strings: Vec = (0..400) + .map(|i| format!("http://example.com/nested/path/that/is/not/inlined/{i}")) + .collect(); + let string_values = Arc::new(StringViewArray::from(strings)) as ArrayRef; + + // Both columns hold 20 of the strings, but reference the data of all 400 + let top_level = string_values.slice(0, 20); + let list_data = ArrayDataBuilder::new(DataType::List(Arc::new( + Field::new_list_field(DataType::Utf8View, true), + ))) + .len(20) + .buffers(vec![Buffer::from_iter(0..=20_i32)]) + .child_data(vec![string_values.slice(20, 20).to_data()]) + .build()?; + let list_array = make_array(list_data); + + let schema = Arc::new(Schema::new(vec![ + Field::new("strings", DataType::Utf8View, false), + Field::new( + "list_strings", + DataType::List(Arc::new(Field::new_list_field(DataType::Utf8View, true))), + false, + ), + ])); + let batch = + RecordBatch::try_new(schema, vec![Arc::clone(&top_level), list_array])?; + + // `gc_view_arrays` compacts the top-level column, `gc_nested_arrays` does not + assert!(!Arc::ptr_eq(gc_view_arrays(&batch)?.column(0), &top_level)); + let gc_batch = gc_nested_arrays(&batch)?; + assert!(Arc::ptr_eq(gc_batch.column(0), &top_level)); + + let gc_list_values = + make_array(gc_batch.column(1).to_data().child_data()[0].clone()); + let gc_list_values = gc_list_values + .as_any() + .downcast_ref::() + .unwrap(); + assert!( + calculate_string_view_waste_ratio(gc_list_values) < 0.2, + "GC should compact nested List child views" + ); + + Ok(()) + } + + #[test] + fn test_gc_columns_compacts_dictionaries() -> Result<()> { + use arrow::array::{DictionaryArray, ListArray}; + use arrow::buffer::OffsetBuffer; + use arrow::datatypes::Int32Type; + + // Every column references a few of the same 100 values + let values: ArrayRef = Arc::new(StringArray::from_iter_values( + (0..100).map(|i| format!("value {i}")), + )); + let dictionary = |keys: Vec| { + DictionaryArray::new(Int32Array::from(keys), Arc::clone(&values)) + }; + let decode = |array: &ArrayRef| -> Vec { + let dictionary = array.as_dictionary::(); + let strings = dictionary.downcast_dict::().unwrap(); + strings + .into_iter() + .map(|s| s.unwrap().to_string()) + .collect() + }; + let list = ListArray::new( + Arc::new(Field::new_list_field( + DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)), + true, + )), + OffsetBuffer::from_lengths([2, 1]), + Arc::new(dictionary(vec![7, 3, 7])), + None, + ); + let batch = RecordBatch::try_from_iter(vec![ + ("compacted", Arc::new(dictionary(vec![5, 9])) as ArrayRef), + ("other", Arc::new(dictionary(vec![1, 2])) as ArrayRef), + ("list", Arc::new(list) as ArrayRef), + ])?; + + // `gc_view_arrays` prunes only dictionaries with more values than keys, + // so it leaves one whose keys reference every value as it is + let dense = RecordBatch::try_from_iter(vec![( + "dense", + Arc::new(DictionaryArray::new( + Int32Array::from(vec![0, 1, 0]), + Arc::new(StringArray::from(vec!["a", "b"])) as ArrayRef, + )) as ArrayRef, + )])?; + let gc_dense = gc_view_arrays(&dense)?; + assert!(Arc::ptr_eq(dense.column(0), gc_dense.column(0))); + + // `gc_columns` keeps only the values that the given columns reference + let gc_batch = gc_columns(&batch, &[0, 2])?; + let compacted = gc_batch.column(0); + assert_eq!(compacted.as_any_dictionary().values().len(), 2); + assert_eq!(decode(compacted), ["value 5", "value 9"]); + assert!(Arc::ptr_eq(gc_batch.column(1), batch.column(1))); + let list_values = gc_batch.column(2).as_list::().values(); + assert_eq!(list_values.as_any_dictionary().values().len(), 2); + assert_eq!(decode(list_values), ["value 7", "value 3", "value 7"]); + + // `gc_nested_arrays` compacts only the nested column + let gc_batch = gc_nested_arrays(&batch)?; + assert!(Arc::ptr_eq(gc_batch.column(0), batch.column(0))); + let list_values = gc_batch.column(2).as_list::().values(); + assert_eq!(list_values.as_any_dictionary().values().len(), 2); + + Ok(()) + } + #[rstest::rstest] #[case(SpillCompression::Uncompressed)] #[case(SpillCompression::Lz4Frame)] diff --git a/datafusion/physical-plan/src/spill/spilled_row_sizes.rs b/datafusion/physical-plan/src/spill/spilled_row_sizes.rs new file mode 100644 index 0000000000000..fe9fc03d1c6c9 --- /dev/null +++ b/datafusion/physical-plan/src/spill/spilled_row_sizes.rs @@ -0,0 +1,524 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: See the NOTICE file distributed with this work for additional information regarding copyright ownership +// SPDX-FileContributor: Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements + +//! Per-row sizes of a batch once it is written to a spill file. + +use arrow::array::{ + AnyDictionaryArray, Array, ArrayRef, AsArray, RecordBatch, downcast_integer_array, + make_array, +}; +use arrow::buffer::{OffsetBuffer, ScalarBuffer}; +use arrow::datatypes::{ArrowNativeType, DataType}; +use arrow_data::MAX_INLINE_VIEW_LEN; + +use super::VIEW_SIZE_BYTES; + +/// Estimates how many bytes each row of a batch takes once written to a spill +/// file, where view arrays are compacted (see [`super::gc_view_arrays`]), and +/// so are the dictionaries of [`Self::compacted_columns`]. +/// +/// A row counts its fixed-width values, the offsets and bytes of its +/// variable-width values, the elements of its nested values, and the values +/// that the keys of compacted dictionaries reference. Buffers that every row +/// shares, such as validity bitmaps and the values of other dictionaries, are +/// not counted. Each view counts the bytes of its value, and each key the +/// value it references, so rows that share bytes are counted as if compaction +/// copied them, although it can write shared bytes only once. +pub(crate) struct SpilledRowSizes { + /// Bytes that every row takes in the fixed-width columns + fixed: usize, + /// The other columns + variable: Vec, + /// See [`Self::max_row`] + max_row: usize, + /// See [`Self::compacted_columns`] + compacted_columns: Vec, +} + +impl SpilledRowSizes { + pub(crate) fn new(batch: &RecordBatch) -> Self { + let mut fixed = 0; + let mut variable = vec![]; + let mut compacted_columns = vec![]; + for (index, column) in batch.columns().iter().enumerate() { + let column = ColumnSize::new(column.as_ref()); + if column.compacted() { + compacted_columns.push(index); + } + match column.fixed() { + Some(width) => fixed += width, + None => variable.push(column), + } + } + let max_row = variable.iter().fold(fixed, |bytes, column| { + bytes.saturating_add(column.max_row()) + }); + Self { + fixed, + variable, + max_row, + compacted_columns, + } + } + + /// Returns the size in bytes of row `row`. + pub(crate) fn row(&self, row: usize) -> usize { + self.variable.iter().fold(self.fixed, |bytes, column| { + bytes + column.rows(row, row + 1) + }) + } + + /// Returns the size in bytes of every row, if they all have the same size. + pub(crate) fn fixed(&self) -> Option { + self.variable.is_empty().then_some(self.fixed) + } + + /// Returns an upper bound on the size in bytes of any row. + pub(crate) fn max_row(&self) -> usize { + self.max_row + } + + /// Returns the columns with dictionaries whose values are mostly unique. + /// Their rows count the values their keys reference, so batches taken + /// from them must be compacted to hold only those values (see + /// [`super::gc_columns`]). + /// + /// `take` keeps every value of a dictionary, so a batch of a few rows + /// would otherwise hold the values of all rows, such as the `array_agg` + /// states of every group. Compacting writes a value once per batch that + /// references it, while a dictionary that consecutive batches share is + /// written only once, so values that many keys reference stay shared. + pub(crate) fn compacted_columns(&self) -> &[usize] { + &self.compacted_columns + } +} + +/// Size model of one column, see [`SpilledRowSizes`]. +enum ColumnSize { + /// Every row takes the same number of bytes. + Fixed(usize), + /// An offset per row, plus the bytes between consecutive offsets. + Bytes(Offsets), + /// A view per row, plus the bytes of values too long to be inlined. + Views(ScalarBuffer), + /// An offset per row, plus the elements between consecutive offsets. + List(Offsets, Box), + /// A fixed number of elements per row. + FixedSizeList(usize, Box), + /// A row of each field. + Struct(Vec), + /// A key of the given width per row, plus the value it references unless + /// the key is null. Only dictionaries that are compacted are sized this + /// way, the others take the bytes of their keys only. + Dictionary(usize, ArrayRef, Box), +} + +impl ColumnSize { + fn new(array: &dyn Array) -> Self { + match array.data_type() { + DataType::Utf8 => Self::Bytes(array.as_string::().offsets().into()), + DataType::LargeUtf8 => Self::Bytes(array.as_string::().offsets().into()), + DataType::Binary => Self::Bytes(array.as_binary::().offsets().into()), + DataType::LargeBinary => { + Self::Bytes(array.as_binary::().offsets().into()) + } + DataType::Utf8View => Self::Views(array.as_string_view().views().clone()), + DataType::BinaryView => Self::Views(array.as_binary_view().views().clone()), + DataType::List(_) => { + let list = array.as_list::(); + Self::List(list.offsets().into(), Box::new(Self::new(list.values()))) + } + DataType::LargeList(_) => { + let list = array.as_list::(); + Self::List(list.offsets().into(), Box::new(Self::new(list.values()))) + } + DataType::Map(_, _) => { + let map = array.as_map(); + Self::List(map.offsets().into(), Box::new(Self::new(map.entries()))) + } + DataType::FixedSizeList(_, size) => Self::FixedSizeList( + *size as usize, + Box::new(Self::new(array.as_fixed_size_list().values())), + ), + DataType::Struct(_) => Self::Struct( + array + .as_struct() + .columns() + .iter() + .map(|field| Self::new(field.as_ref())) + .collect(), + ), + DataType::Dictionary(key, _) => { + let width = key.primitive_width().unwrap_or_default(); + let dictionary = array.as_any_dictionary(); + if should_compact(dictionary) { + Self::Dictionary( + width, + make_array(dictionary.keys().to_data()), + Box::new(Self::new(dictionary.values().as_ref())), + ) + } else { + Self::Fixed(width) + } + } + DataType::FixedSizeBinary(size) => Self::Fixed(*size as usize), + DataType::Null | DataType::Boolean => Self::Fixed(0), + data_type => Self::Fixed(data_type.primitive_width().unwrap_or_else(|| { + // Other layouts are spread evenly over their rows + let bytes = array.to_data().get_slice_memory_size().unwrap_or_default(); + bytes / array.len().max(1) + })), + } + } + + /// Returns whether this column has a dictionary that is compacted. + fn compacted(&self) -> bool { + match self { + Self::Dictionary(_, _, _) => true, + Self::List(_, values) | Self::FixedSizeList(_, values) => values.compacted(), + Self::Struct(fields) => fields.iter().any(Self::compacted), + Self::Fixed(_) | Self::Bytes(_) | Self::Views(_) => false, + } + } + + /// Returns the size in bytes of every row, if they all have the same size. + fn fixed(&self) -> Option { + match self { + Self::Fixed(width) => Some(*width), + Self::FixedSizeList(size, values) => values.fixed().map(|width| size * width), + Self::Struct(fields) => fields.iter().map(Self::fixed).sum(), + Self::Bytes(_) + | Self::Views(_) + | Self::List(_, _) + | Self::Dictionary(_, _, _) => None, + } + } + + /// Returns an upper bound on the size in bytes of a row. + fn max_row(&self) -> usize { + match self { + Self::Fixed(width) => *width, + Self::Bytes(offsets) => offsets.width() + offsets.max_len(), + Self::Views(views) => views + .iter() + .copied() + .map(view_bytes) + .max() + .unwrap_or_default(), + Self::List(offsets, values) => offsets + .width() + .saturating_add(offsets.max_len().saturating_mul(values.max_row())), + Self::FixedSizeList(size, values) => size.saturating_mul(values.max_row()), + Self::Struct(fields) => fields + .iter() + .fold(0, |bytes, field| bytes.saturating_add(field.max_row())), + Self::Dictionary(width, _, values) => width.saturating_add(values.max_row()), + } + } + + /// Returns the size in bytes of rows `start..end`. + fn rows(&self, start: usize, end: usize) -> usize { + let num_rows = end - start; + match self { + Self::Fixed(width) => num_rows * width, + Self::Bytes(offsets) => { + num_rows * offsets.width() + offsets.get(end) - offsets.get(start) + } + Self::Views(views) => views[start..end].iter().copied().map(view_bytes).sum(), + Self::List(offsets, values) => { + num_rows * offsets.width() + + values.rows(offsets.get(start), offsets.get(end)) + } + Self::FixedSizeList(size, values) => values.rows(start * size, end * size), + Self::Struct(fields) => { + fields.iter().map(|field| field.rows(start, end)).sum() + } + Self::Dictionary(width, keys, values) => { + let value_bytes = downcast_integer_array!( + keys => (start..end) + .filter(|&row| keys.is_valid(row)) + .map(|row| { + let key = keys.value(row).as_usize(); + values.rows(key, key + 1) + }) + .sum::(), + data_type => unreachable!("dictionary keys of type {data_type}") + ); + num_rows * width + value_bytes + } + } + } +} + +/// Returns whether to compact the batches of `dictionary`, which is when fewer +/// than two keys reference each of its values on average (see +/// [`SpilledRowSizes::compacted_columns`]). +fn should_compact(dictionary: &dyn AnyDictionaryArray) -> bool { + let keys = dictionary.keys(); + keys.len() - keys.null_count() < 2 * dictionary.values().len() +} + +/// Returns the size in bytes of a view, plus its value unless it is inlined. +fn view_bytes(view: u128) -> usize { + let len = view as u32; + if len > MAX_INLINE_VIEW_LEN { + VIEW_SIZE_BYTES + len as usize + } else { + VIEW_SIZE_BYTES + } +} + +/// The offsets of a variable-width or list column. +enum Offsets { + Small(ScalarBuffer), + Large(ScalarBuffer), +} + +impl Offsets { + /// Returns the size in bytes of one offset. + fn width(&self) -> usize { + match self { + Self::Small(_) => size_of::(), + Self::Large(_) => size_of::(), + } + } + + fn get(&self, index: usize) -> usize { + match self { + Self::Small(offsets) => offsets[index] as usize, + Self::Large(offsets) => offsets[index] as usize, + } + } + + /// Returns the largest distance between consecutive offsets. + fn max_len(&self) -> usize { + match self { + Self::Small(offsets) => offsets + .windows(2) + .map(|pair| (pair[1] - pair[0]) as usize) + .max(), + Self::Large(offsets) => offsets + .windows(2) + .map(|pair| (pair[1] - pair[0]) as usize) + .max(), + } + .unwrap_or_default() + } +} + +impl From<&OffsetBuffer> for Offsets { + fn from(offsets: &OffsetBuffer) -> Self { + Self::Small(offsets.inner().clone()) + } +} + +impl From<&OffsetBuffer> for Offsets { + fn from(offsets: &OffsetBuffer) -> Self { + Self::Large(offsets.inner().clone()) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use arrow::array::{ + ArrayRef, BooleanArray, DictionaryArray, FixedSizeBinaryArray, Int32Array, + Int64Array, ListArray, StringArray, StringViewArray, StructArray, + }; + use arrow::buffer::Buffer; + use arrow::datatypes::{Field, Int32Type, Int64Type}; + use std::sync::Arc; + + fn batch(columns: Vec) -> RecordBatch { + RecordBatch::try_from_iter( + columns + .into_iter() + .enumerate() + .map(|(i, column)| (format!("c{i}"), column)), + ) + .unwrap() + } + + /// Returns the size of each row of a batch of `columns`, and the upper + /// bound on the size of any row. + fn row_sizes(columns: Vec) -> (Vec, usize) { + let batch = batch(columns); + let sizes = SpilledRowSizes::new(&batch); + let rows = (0..batch.num_rows()).map(|row| sizes.row(row)).collect(); + (rows, sizes.max_row()) + } + + #[test] + fn fixed_width_and_offsets() { + let ints: ArrayRef = Arc::new(Int64Array::from(vec![1, 2])); + let strings: ArrayRef = Arc::new(StringArray::from(vec!["a", "bbbb"])); + // 8 bytes of Int64, plus a 4 byte offset and the string bytes + assert_eq!(row_sizes(vec![ints, strings]), (vec![13, 16], 16)); + } + + #[test] + fn rows_of_fixed_width_columns_have_one_size() { + let ints: ArrayRef = Arc::new(Int64Array::from(vec![1, 2])); + let fixed: ArrayRef = Arc::new(FixedSizeBinaryArray::new( + 3, + Buffer::from(vec![0u8; 6]), + None, + )); + let booleans: ArrayRef = Arc::new(BooleanArray::from(vec![true, false])); + let strings: ArrayRef = Arc::new(StringArray::from(vec!["a", "bbbb"])); + // 8 bytes of Int64 and 3 of FixedSizeBinary(3). The bits of booleans + // are not counted. + let columns = vec![Arc::clone(&ints), fixed, booleans]; + assert_eq!(row_sizes(columns.clone()), (vec![11, 11], 11)); + assert_eq!(SpilledRowSizes::new(&batch(columns)).fixed(), Some(11)); + assert_eq!( + SpilledRowSizes::new(&batch(vec![ints, strings])).fixed(), + None + ); + } + + #[test] + fn views_count_their_bytes_unless_inlined() { + let long = "a string longer than twelve"; + let views: ArrayRef = Arc::new(StringViewArray::from(vec!["short", long])); + let long_row = 16 + long.len(); + assert_eq!(row_sizes(vec![views]), (vec![16, long_row], long_row)); + } + + #[test] + fn views_sharing_bytes_count_them_for_every_row() { + // Three views of the same 100 bytes, stored once + let value = "x".repeat(100); + let one = StringViewArray::from(vec![value.as_str()]); + let view = one.views()[0]; + let shared = StringViewArray::new( + ScalarBuffer::from(vec![view; 3]), + one.data_buffers().to_vec(), + None, + ); + // Compacting before spilling copies the bytes for each view + assert_eq!(row_sizes(vec![Arc::new(shared)]), (vec![116; 3], 116)); + } + + #[test] + fn lists_count_their_elements() { + let list: ArrayRef = + Arc::new(ListArray::from_iter_primitive::(vec![ + Some(vec![Some(1), Some(2), Some(3)]), + Some(vec![]), + Some(vec![Some(4)]), + ])); + // An offset plus the elements of the row. The bound is an offset plus + // the most elements in a row, each of the largest size. + assert_eq!(row_sizes(vec![Arc::clone(&list)]), (vec![28, 4, 12], 28)); + // A slice keeps the sizes of its own rows + assert_eq!(row_sizes(vec![list.slice(1, 2)]), (vec![4, 12], 12)); + } + + #[test] + fn nested_views_count_their_bytes() { + let values = + StringViewArray::from(vec!["x".repeat(20), "y".repeat(5), "z".repeat(30)]); + let list: ArrayRef = Arc::new(ListArray::new( + Arc::new(Field::new_list_field(DataType::Utf8View, true)), + OffsetBuffer::new(vec![0, 2, 3].into()), + Arc::new(values), + None, + )); + // The bound takes two elements, the most in a row, of the largest size + assert_eq!( + row_sizes(vec![list]), + (vec![4 + 36 + 16, 4 + 46], 4 + 2 * 46) + ); + } + + #[test] + fn struct_fields_add_up() { + let fields: ArrayRef = Arc::new(StructArray::from(vec![ + ( + Arc::new(Field::new("a", DataType::Int32, false)), + Arc::new(Int32Array::from(vec![1, 2])) as ArrayRef, + ), + ( + Arc::new(Field::new("b", DataType::Utf8, false)), + Arc::new(StringArray::from(vec!["", "abc"])) as ArrayRef, + ), + ])); + assert_eq!(row_sizes(vec![fields]), (vec![8, 11], 11)); + } + + #[test] + fn dictionaries_of_unique_values_count_them() { + let values = StringArray::from(vec!["a long dictionary value", "b", "cc"]); + let dictionary: ArrayRef = Arc::new(DictionaryArray::new( + Int32Array::from(vec![Some(2), None, Some(0)]), + Arc::new(values), + )); + // A 4 byte key, plus the 4 byte offset and the bytes of the value it + // references. A null key references none. + let long_row = 4 + 4 + 23; + assert_eq!( + row_sizes(vec![Arc::clone(&dictionary)]), + (vec![4 + 4 + 2, 4, long_row], long_row) + ); + let sizes = SpilledRowSizes::new(&batch(vec![dictionary])); + assert_eq!(sizes.compacted_columns(), [0]); + } + + #[test] + fn dictionaries_of_shared_values_count_their_keys() { + // Two keys reference each value, which every batch keeps + let dictionary: ArrayRef = Arc::new( + vec![ + "a long dictionary value", + "b", + "a long dictionary value", + "b", + ] + .into_iter() + .collect::>(), + ); + assert_eq!(row_sizes(vec![Arc::clone(&dictionary)]), (vec![4; 4], 4)); + let sizes = SpilledRowSizes::new(&batch(vec![dictionary])); + assert!(sizes.compacted_columns().is_empty()); + } + + #[test] + fn nested_dictionaries_count_their_values() { + let values = DictionaryArray::new( + Int32Array::from(vec![0, 1, 2]), + Arc::new(StringArray::from(vec![ + "x".repeat(20), + "y".repeat(5), + "z".repeat(30), + ])), + ); + let list: ArrayRef = Arc::new(ListArray::new( + Arc::new(Field::new_list_field(values.data_type().clone(), true)), + OffsetBuffer::new(vec![0, 2, 3, 3].into()), + Arc::new(values), + None, + )); + // An offset, plus a key, an offset and the bytes of the value of each + // element. The bound takes two elements of the largest size. + assert_eq!( + row_sizes(vec![Arc::clone(&list)]), + (vec![4 + 28 + 13, 4 + 38, 4], 4 + 2 * 38) + ); + let sizes = SpilledRowSizes::new(&batch(vec![list])); + assert_eq!(sizes.compacted_columns(), [0]); + } + + #[test] + fn empty_nested_dictionaries_count_nothing() { + // The values of null lists are an empty dictionary, without values + let dictionary = + DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)); + let list: ArrayRef = Arc::new(ListArray::new_null( + Arc::new(Field::new_list_field(dictionary, true)), + 2, + )); + assert_eq!(row_sizes(vec![list]), (vec![4, 4], 4)); + } +} diff --git a/datafusion/sqllogictest/test_files/aggregate_memory_spill.slt b/datafusion/sqllogictest/test_files/aggregate_memory_spill.slt index 915e69642e260..e8b9cdae130e8 100644 --- a/datafusion/sqllogictest/test_files/aggregate_memory_spill.slt +++ b/datafusion/sqllogictest/test_files/aggregate_memory_spill.slt @@ -344,8 +344,257 @@ RESET datafusion.execution.spill_compression statement ok DROP VIEW byte_keys +# --- Case J: a few groups with large state --- +# The array_agg states of all 32 groups do not fit in the memory limit, but a +# single group does. Spilling must write the groups in batches that can be +# read back within the limit, instead of one batch holding every group. +# See https://github.com/apache/datafusion/issues/25851 statement ok -SET datafusion.execution.target_partitions = 4 +RESET datafusion.runtime.max_spill_merge_fan_in + +statement ok +SET datafusion.execution.batch_size = 8192 + +statement ok +SET datafusion.runtime.memory_limit = '24M' + +# Partial/final aggregation: the final aggregation spills the partial states of +# its partition. +statement ok +SET datafusion.execution.target_partitions = 2 + +query II +SELECT count(*), sum(n) +FROM ( + SELECT v % 32 AS k, cardinality(array_agg(concat(CAST(v AS VARCHAR), repeat('x', 100)))) AS n + FROM generate_series(1, 500000) AS t(v) + GROUP BY v % 32 +) +---- +32 500000 + +# Assert the final aggregation spilled, the `spilled_bytes` metric must be in MB +query TT +EXPLAIN ANALYZE +SELECT count(*), sum(n) +FROM ( + SELECT v % 32 AS k, cardinality(array_agg(concat(CAST(v AS VARCHAR), repeat('x', 100)))) AS n + FROM generate_series(1, 500000) AS t(v) + GROUP BY v % 32 +) +---- + +06)----------AggregateExec: mode=FinalPartitioned, gby=[t.v % Int64(32)@0 as t.v % Int64(32)], aggr=[], metrics=[spilled_bytes= MB, spilled_rows=] +07)------------RepartitionExec: + +# Single-stage aggregation: spills several runs as input arrives, which replay +# merges. +statement ok +SET datafusion.execution.target_partitions = 1 + +query II +SELECT count(*), sum(n) +FROM ( + SELECT v % 32 AS k, cardinality(array_agg(concat(CAST(v AS VARCHAR), repeat('x', 100)))) AS n + FROM generate_series(1, 500000) AS t(v) + GROUP BY v % 32 +) +---- +32 500000 + +# Assert spill happened, the `spilled_bytes` metric must be in MB +query TT +EXPLAIN ANALYZE +SELECT count(*), sum(n) +FROM ( + SELECT v % 32 AS k, cardinality(array_agg(concat(CAST(v AS VARCHAR), repeat('x', 100)))) AS n + FROM generate_series(1, 500000) AS t(v) + GROUP BY v % 32 +) +---- + +04)------AggregateExec: mode=Single, gby=[v@0 % 32 as t.v % Int64(32)], aggr=[], metrics=[spilled_bytes= MB, spilled_rows=] +05)--------ProjectionExec: + +# 14 large groups, adjacent in key order, spilled together with 62,500 small +# groups. The small groups must not make a batch take several large groups. +query II +SELECT count(*), sum(n) +FROM ( + SELECT k, cardinality(array_agg(p)) AS n + FROM ( + SELECT CASE WHEN v % 8 = 0 THEN v + 1000000 ELSE v % 16 END AS k, + concat(CAST(v AS VARCHAR), repeat('x', 100)) AS p + FROM generate_series(1, 500000) AS t(v) + ) + GROUP BY k +) +---- +62514 500000 + +# Assert spill happened, the `spilled_bytes` metric must be in MB +query TT +EXPLAIN ANALYZE +SELECT count(*), sum(n) +FROM ( + SELECT k, cardinality(array_agg(p)) AS n + FROM ( + SELECT CASE WHEN v % 8 = 0 THEN v + 1000000 ELSE v % 16 END AS k, + concat(CAST(v AS VARCHAR), repeat('x', 100)) AS p + FROM generate_series(1, 500000) AS t(v) + ) + GROUP BY k +) +---- + +04)------AggregateExec: mode=Single, gby=[k@0 as k], aggr=[array_agg(p)], metrics=[spilled_bytes= MB, spilled_rows=] +05)--------ProjectionExec: + +# Runs of large groups followed by a run of small groups. Replay must not +# merge the large rows into one batch. +query II +SELECT count(*), sum(n) +FROM ( + SELECT k, cardinality(array_agg(s)) AS n + FROM ( + SELECT CASE WHEN v <= 360000 THEN v % 32 ELSE v END AS k, + CASE WHEN v <= 360000 THEN concat(CAST(v AS VARCHAR), repeat('x', 100)) ELSE CAST(v AS VARCHAR) END AS s + FROM generate_series(1, 700000) AS t(v) + ) + GROUP BY k +) +---- +340032 700000 + +# Assert spill happened, the `spilled_bytes` metric must be in MB +query TT +EXPLAIN ANALYZE +SELECT count(*), sum(n) +FROM ( + SELECT k, cardinality(array_agg(s)) AS n + FROM ( + SELECT CASE WHEN v <= 360000 THEN v % 32 ELSE v END AS k, + CASE WHEN v <= 360000 THEN concat(CAST(v AS VARCHAR), repeat('x', 100)) ELSE CAST(v AS VARCHAR) END AS s + FROM generate_series(1, 700000) AS t(v) + ) + GROUP BY k +) +---- + +04)------AggregateExec: mode=Single, gby=[k@0 as k], aggr=[array_agg(s)], metrics=[spilled_bytes= MB, spilled_rows=] +05)--------ProjectionExec: + +# Dictionary-encoded values, each referenced once. The array_agg states of all +# groups share one dictionary, so each spilled batch must hold the values of +# its own groups only. +statement ok +SET datafusion.execution.target_partitions = 2 + +query II +SELECT count(*), sum(n) +FROM ( + SELECT v % 32 AS k, cardinality(array_agg(arrow_cast(concat(CAST(v AS VARCHAR), repeat('x', 100)), 'Dictionary(Int32, Utf8)'))) AS n + FROM generate_series(1, 500000) AS t(v) + GROUP BY v % 32 +) +---- +32 500000 + +# Assert the final aggregation spilled, the `spilled_bytes` metric must be in MB +query TT +EXPLAIN ANALYZE +SELECT count(*), sum(n) +FROM ( + SELECT v % 32 AS k, cardinality(array_agg(arrow_cast(concat(CAST(v AS VARCHAR), repeat('x', 100)), 'Dictionary(Int32, Utf8)'))) AS n + FROM generate_series(1, 500000) AS t(v) + GROUP BY v % 32 +) +---- + +06)----------AggregateExec: mode=FinalPartitioned, gby=[t.v % Int64(32)@0 as t.v % Int64(32)], aggr=[], metrics=[spilled_bytes= MB, spilled_rows=] +07)------------RepartitionExec: + +statement ok +SET datafusion.execution.target_partitions = 1 + +query II +SELECT count(*), sum(n) +FROM ( + SELECT v % 32 AS k, cardinality(array_agg(arrow_cast(concat(CAST(v AS VARCHAR), repeat('x', 100)), 'Dictionary(Int32, Utf8)'))) AS n + FROM generate_series(1, 500000) AS t(v) + GROUP BY v % 32 +) +---- +32 500000 + +# Assert spill happened, the `spilled_bytes` metric must be in MB +query TT +EXPLAIN ANALYZE +SELECT count(*), sum(n) +FROM ( + SELECT v % 32 AS k, cardinality(array_agg(arrow_cast(concat(CAST(v AS VARCHAR), repeat('x', 100)), 'Dictionary(Int32, Utf8)'))) AS n + FROM generate_series(1, 500000) AS t(v) + GROUP BY v % 32 +) +---- + +04)------AggregateExec: mode=Single, gby=[v@0 % 32 as t.v % Int64(32)], aggr=[], metrics=[spilled_bytes= MB, spilled_rows=] +05)--------ProjectionExec: + +# string_agg states are top-level strings +query II +SELECT count(*), sum(n) +FROM ( + SELECT v % 32 AS k, length(string_agg(concat(CAST(v AS VARCHAR), repeat('x', 100)), ',')) AS n + FROM generate_series(1, 500000) AS t(v) + GROUP BY v % 32 +) +---- +32 53388863 + +# Assert spill happened, the `spilled_bytes` metric must be in MB +query TT +EXPLAIN ANALYZE +SELECT count(*), sum(n) +FROM ( + SELECT v % 32 AS k, length(string_agg(concat(CAST(v AS VARCHAR), repeat('x', 100)), ',')) AS n + FROM generate_series(1, 500000) AS t(v) + GROUP BY v % 32 +) +---- + +04)------AggregateExec: mode=Single, gby=[v@0 % 32 as t.v % Int64(32)], aggr=[], metrics=[spilled_bytes= MB, spilled_rows=] +05)--------ProjectionExec: + +# Each dictionary value repeats 50 times within one group. Spilled batches +# share the whole dictionary, so replay must compact the lists it keeps to the +# values of their own group. +statement ok +SET datafusion.runtime.memory_limit = '8M' + +query II +SELECT count(*), sum(n) +FROM ( + SELECT (v / 50) % 32 AS k, cardinality(array_agg(arrow_cast(concat(CAST(v / 50 AS VARCHAR), repeat('x', 100)), 'Dictionary(Int32, Utf8)'))) AS n + FROM generate_series(1, 2000000) AS t(v) + GROUP BY (v / 50) % 32 +) +---- +32 2000000 + +# Assert spill happened, the `spilled_bytes` metric must be in MB +query TT +EXPLAIN ANALYZE +SELECT count(*), sum(n) +FROM ( + SELECT (v / 50) % 32 AS k, cardinality(array_agg(arrow_cast(concat(CAST(v / 50 AS VARCHAR), repeat('x', 100)), 'Dictionary(Int32, Utf8)'))) AS n + FROM generate_series(1, 2000000) AS t(v) + GROUP BY (v / 50) % 32 +) +---- + +04)------AggregateExec: mode=Single, gby=[__common_expr_1@0 % 32 as t.v / Int64(50) % Int64(32)], aggr=[], metrics=[spilled_bytes= MB, spilled_rows=] +05)--------ProjectionExec: # Restore settings to slt runner defaults statement ok @@ -357,5 +606,10 @@ RESET datafusion.runtime.max_spill_merge_fan_in statement ok RESET datafusion.execution.batch_size +# The SLT runner sets `target_partitions` to 4 instead of using the default, so +# reset it explicitly. +statement ok +SET datafusion.execution.target_partitions = 4 + statement ok RESET datafusion.catalog.create_default_catalog_and_schema