From 5fb670cb0c2456b7dc215ae1d1a4ec174534079b Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Sun, 8 Nov 2026 16:46:13 +0000 Subject: [PATCH] fix: size aggregate spill batches by bytes as well as rows (#25853) - Closes #25851. A grouped aggregate whose groups carry large state, such as `array_agg` over a low-cardinality key, fails with `ResourcesExhausted` after it spills, even though each group fits in memory and the same data spread over many small groups succeeds. Spilling writes the aggregate state in batches of at most `batch_size` rows, whatever their size in bytes. With fewer groups than `batch_size`, every group lands in one batch, and everything after the spill has to hold that batch at once: - With a single spill file, replay reads the batch directly and its table must hold every group at once. In the reproducer from the issue, that is 225 MB of state in a 128 MB pool, although each group is about 28 MB. - With several spill files, the merge reserves twice the largest batch of each file. When that does not fit, it tries to split the file, which first reserves twice the batch it is about to split. Even with small spilled batches, the merge combines rows until it has `batch_size` of them, rebuilding a batch of large groups. Aggregate spilling and replay now bound their batches by bytes as well as rows, with a target of 1 MiB (`SPILL_BATCH_TARGET_BYTES`): - `IncrementalSortIterator::with_max_batch_bytes` cuts the sorted state into chunks by the size each row takes once spilled, so a batch only exceeds the target when it holds a single larger group. Row sizes come from the new `SpilledRowSizes`, which counts view values as they are written after compaction and nested values by their elements. - `StreamingMergeBuilder::with_max_batch_bytes` makes the merge emit a batch once its rows reach the target, estimated from the average row size of the batches they come from, as well as at `batch_size` rows. It applies to every pass of the multi-level merge, so rows of small groups still merge in batches of `batch_size` rows. - Dictionaries whose values are mostly unique (fewer than two keys per value) count the values their keys reference, and each chunk holds only those values, since `take` keeps every value of a dictionary. Dictionaries whose values many keys share count only their keys and stay shared, as on `main`, so that each chunk does not write its own copy of them. - Replay compacts the views and dictionaries of the nested columns of merged batches (`gc_nested_arrays`). A merged batch references the buffers of every batch its rows come from, and `array_agg` keeps the lists it merges until its group completes. Both new options are crate-private and only used by aggregate spilling and replay; other merges are unchanged. The intermediate merge sizing from #25428 now receives the largest row count of the batches each run actually wrote. This applies to every stream built on `AggregateSpill` (hash and ordered, final and single-stage). The legacy `GroupedHashAggregateStream`, used only when `datafusion.execution.enable_migration_aggregate` is `false`, is unchanged. - New Case I in `aggregate_memory_spill.slt`, with a 24 MB limit, where `EXPLAIN ANALYZE` assertions check that each query spilled. All of these fail on `main` with `ResourcesExhausted`: - `array_agg` over 32 groups of about 2 MB each, with two partitions (one spill file per partition) and with one partition (several spill files, merged). - 14 large groups, adjacent in key order, spilled together with 62,500 small groups. - Runs of large groups followed by a run of small groups. - `array_agg` over dictionary-encoded values, with two partitions and with one. - `string_agg` with one partition. - `array_agg` over dictionary values that each repeat within one group, at 8 MB. - Unit tests for `SpilledRowSizes` (including views that share bytes, nested views, fixed-width rows, and dictionaries of unique, shared and no values), for byte-limited chunks in `IncrementalSortIterator` and their dictionaries, for byte-limited merges of spill files, and for dictionary compaction. With `datafusion-cli -m 128M`, the reproducer from the issue now returns 16 groups and 3,000,000 values, and so does the variant from the issue's additional context. The reproducer with 3,000 one-value groups added, `array_agg` over dictionary-encoded Parquet strings, and the reproducer with its values cast to `Dictionary(Int32, Utf8)`, which fail on `main`, also succeed. Performance compared with `main`: - `external_aggr` (TPC-H SF1, integer keys): no case is slower; 7 are unchanged and 1 is 1.06x faster. - A spilling `GROUP BY` over 5M string keys at 256 MB is unchanged (0.61 s against 0.62 s). - 8M small groups plus 200 large `array_agg` groups take 0.80 s at 128 MB, against 2.99 s on `main`. There `main` spills 26M rows (1 GB) while merging its runs, and this PR spills each of the 8M rows once. At 512 MB, where both spill twice, this PR takes 1.00 s against 0.96 s, from measuring the rows of the spilled batch that holds the large groups. - `sort_preserving_merge` shows no change beyond run-to-run noise. Merges without a byte limit only check that none is set. - `array_agg` over a dictionary of 100 distinct values (2M rows, 100K groups, 24 MB) spills the same 10.5 MB (20-byte values) and 14.7 MB (200-byte values) as `main`, in the same time. Aggregations whose groups have large state can now complete under memory limits where they previously failed after spilling. No API or configuration changes. (cherry picked from commit 32dd1d2c5caf8563cdfd471ee5cbce77e9031376) --- .../src/aggregates/clustered_final_stream.rs | 4 +- .../physical-plan/src/aggregates/spill.rs | 33 +- datafusion/physical-plan/src/sorts/builder.rs | 6 +- datafusion/physical-plan/src/sorts/merge.rs | 43 +- .../src/sorts/multi_level_merge.rs | 13 +- datafusion/physical-plan/src/sorts/stream.rs | 177 +++++- .../src/sorts/streaming_merge.rs | 128 ++++- datafusion/physical-plan/src/spill/mod.rs | 183 +++++- .../src/spill/spilled_row_sizes.rs | 524 ++++++++++++++++++ .../test_files/aggregate_memory_spill.slt | 256 ++++++++- 10 files changed, 1332 insertions(+), 35 deletions(-) create mode 100644 datafusion/physical-plan/src/spill/spilled_row_sizes.rs 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 262bb243c053a..53897bec5d8d9 100644 --- a/datafusion/physical-plan/src/spill/mod.rs +++ b/datafusion/physical-plan/src/spill/mod.rs @@ -8,6 +8,7 @@ pub(crate) mod in_progress_spill_file; 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; @@ -734,7 +735,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. @@ -779,11 +779,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); } @@ -795,7 +832,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 @@ -817,16 +856,27 @@ 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)); + } + Ok((gc.with_values(values), 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), } } @@ -911,7 +961,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)); @@ -921,7 +971,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()); } @@ -2712,6 +2762,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 dc6af9b1e4e47..71d82ddebca10 100644 --- a/datafusion/sqllogictest/test_files/aggregate_memory_spill.slt +++ b/datafusion/sqllogictest/test_files/aggregate_memory_spill.slt @@ -247,8 +247,257 @@ FROM ( ---- 8192 8192 +# --- Case I: 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 @@ -260,5 +509,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