diff --git a/datafusion/core/tests/fuzz_cases/spilling_fuzz_in_memory_constrained_env.rs b/datafusion/core/tests/fuzz_cases/spilling_fuzz_in_memory_constrained_env.rs index ba566d597bb27..bfac2bbf71f54 100644 --- a/datafusion/core/tests/fuzz_cases/spilling_fuzz_in_memory_constrained_env.rs +++ b/datafusion/core/tests/fuzz_cases/spilling_fuzz_in_memory_constrained_env.rs @@ -260,12 +260,12 @@ async fn test_sort_with_limited_memory_and_oversized_record_batch() -> Result<() let number_of_record_batches = 100; - // Each spilled run's largest batch is so big that two merge streams cannot be - // reserved at once even at the smallest read-buffer size (`2 * (2 * batch) > - // pool`), yet a single stream still fits (`2 * batch < pool`). Reducing the - // buffer size therefore cannot help, the multi-level merge has to re-spill a - // run with a smaller batch size to make progress instead of failing with - // `ResourcesExhausted`. + // Each input batch is so big that two merge streams over spill files of such + // batches cannot be reserved at once (`2 * (2 * batch) > pool`), yet a single + // stream fits (`2 * batch < pool`). The sort must not fail with + // `ResourcesExhausted`: it bounds the batches it spills by bytes, so the merge + // can seat two streams. (The merge's fallback of re-spilling a run at half the + // batch size is covered by the unit tests in `multi_level_merge.rs`.) let metrics = run_sort_test_with_limited_memory(RunTestWithLimitedMemoryArgs { pool_size, task_ctx: Arc::new(task_ctx), @@ -277,12 +277,9 @@ async fn test_sort_with_limited_memory_and_oversized_record_batch() -> Result<() }) .await?; - let output_batches = get_output_batches_from_metrics(&metrics); - - // minimum 2 batches more assert!( - output_batches >= number_of_record_batches + 2, - "output_batches {output_batches} should be greater than number_of_record_batches ({number_of_record_batches}) + 2" + metrics.spill_count().unwrap_or(0) > 0, + "the sort must spill" ); Ok(()) diff --git a/datafusion/core/tests/memory_limit/mod.rs b/datafusion/core/tests/memory_limit/mod.rs index 8d67c27e96859..69fdc074bde07 100644 --- a/datafusion/core/tests/memory_limit/mod.rs +++ b/datafusion/core/tests/memory_limit/mod.rs @@ -11,6 +11,7 @@ mod nlj_spill_unmatched; mod repartition_mem_limit; mod union_nullable_spill; mod view_spill_compaction; +mod wide_row_sort_spill; use arrow::array::{ ArrayRef, DictionaryArray, Int32Array, Int64Array, Int64Builder, ListBuilder, RecordBatch, StringArray, StringViewArray, StructArray, diff --git a/datafusion/core/tests/memory_limit/wide_row_sort_spill.rs b/datafusion/core/tests/memory_limit/wide_row_sort_spill.rs new file mode 100644 index 0000000000000..893ee25167027 --- /dev/null +++ b/datafusion/core/tests/memory_limit/wide_row_sort_spill.rs @@ -0,0 +1,319 @@ +// 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 + +//! An external sort over wide rows must finish under a memory limit that holds a few rows, +//! even though `batch_size` rows do not fit. +//! +//! The input arrives in small batches, so every input batch fits. Before the fix, the sort +//! wrote its spill files in batches of up to `batch_size` rows. With wide rows, one spilled +//! batch was larger than the memory pool, and the final merge of the spill files, which +//! reserves memory for the largest batch of each file, failed with `ResourcesExhausted`. + +use std::sync::Arc; + +use arrow::array::{ + Array, ArrayRef, BinaryViewArray, DictionaryArray, Int32Array, Int64Array, ListArray, + ListBuilder, RecordBatch, StringArray, StringBuilder, StringViewArray, StructArray, + UInt32Array, +}; +use arrow::buffer::OffsetBuffer; +use arrow::compute::{concat_batches, take}; +use arrow_schema::{DataType, Field, Fields, Schema, SchemaRef}; +use datafusion::datasource::MemTable; +use datafusion::execution::memory_pool::FairSpillPool; +use datafusion::execution::runtime_env::RuntimeEnvBuilder; +use datafusion::physical_plan::collect; +use datafusion::prelude::{SessionConfig, SessionContext}; + +/// 64 KiB per row. +const ROW_BYTES: usize = 64 * 1024; +/// 32 MiB of rows in total. +const ROWS: usize = 512; +/// Rows per input batch: 256 KiB, which fits in the pool. +const INPUT_BATCH_ROWS: usize = 4; +/// Holds about 190 rows, less than half of `ROWS`, so the sort must spill. +const POOL_BYTES: usize = 12 * 1024 * 1024; +const SPILL_RESERVATION_BYTES: usize = 1024 * 1024; + +fn payload(i: usize) -> String { + format!("{i:08}{}", "x".repeat(ROW_BYTES)) +} + +/// The row ranges of the input batches. +fn input_ranges() -> impl Iterator> { + (0..ROWS) + .step_by(INPUT_BATCH_ROWS) + .map(|start| start..start + INPUT_BATCH_ROWS) +} + +/// Builds the input batches: an `id` column followed by the given columns for each range. +fn input( + fields: Vec, + mut columns: impl FnMut(std::ops::Range) -> Vec, +) -> (SchemaRef, Vec) { + let mut all = vec![Field::new("id", DataType::Int64, false)]; + all.extend(fields); + let schema = Arc::new(Schema::new(all)); + let batches = input_ranges() + .map(|rows| { + let ids = Int64Array::from_iter_values(rows.clone().map(|i| i as i64)); + let mut cols: Vec = vec![Arc::new(ids)]; + cols.extend(columns(rows)); + RecordBatch::try_new(Arc::clone(&schema), cols).unwrap() + }) + .collect(); + (schema, batches) +} + +async fn run_sort( + schema: &SchemaRef, + batches: &[RecordBatch], + pool_bytes: Option, +) -> (Vec, usize) { + let table = MemTable::try_new(Arc::clone(schema), vec![batches.to_vec()]).unwrap(); + let mut runtime = RuntimeEnvBuilder::new(); + if let Some(pool_bytes) = pool_bytes { + runtime = runtime.with_memory_pool(Arc::new(FairSpillPool::new(pool_bytes))); + } + let config = SessionConfig::new() + .with_target_partitions(1) + // 512 rows of 64 KiB is 32 MiB, more than the pool. + .with_batch_size(ROWS) + .with_sort_spill_reservation_bytes(SPILL_RESERVATION_BYTES); + let ctx = SessionContext::new_with_config_rt(config, runtime.build_arc().unwrap()); + ctx.register_table("t", Arc::new(table)).unwrap(); + let plan = ctx + .sql("SELECT * FROM t ORDER BY id DESC") + .await + .unwrap() + .create_physical_plan() + .await + .unwrap(); + let result = collect(Arc::clone(&plan), ctx.task_ctx()) + .await + .expect("the sort must spill and finish"); + let spill_count = plan.metrics().and_then(|m| m.spill_count()).unwrap_or(0); + (result, spill_count) +} + +/// Sorts the input under the memory limit, checks that it spilled, and checks that the +/// result equals a sort of the same input without a limit, column by column. +async fn sort_spills_and_matches(schema: SchemaRef, batches: Vec) { + let (result, spill_count) = run_sort(&schema, &batches, Some(POOL_BYTES)).await; + assert!(spill_count > 0, "the sort must spill"); + let (expected, _) = run_sort(&schema, &batches, None).await; + + let result = concat_batches(&schema, &result).unwrap(); + let expected = concat_batches(&schema, &expected).unwrap(); + assert_eq!(result.num_rows(), ROWS); + let ids: Vec = result + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .values() + .to_vec(); + assert_eq!(ids, (0..ROWS as i64).rev().collect::>()); + for (i, field) in schema.fields().iter().enumerate() { + assert_eq!( + result.column(i).as_ref(), + expected.column(i).as_ref(), + "column {} differs", + field.name() + ); + } +} + +async fn sort_wide_rows(data_type: DataType) { + let field = Field::new("payload", data_type.clone(), false); + let (schema, batches) = input(vec![field], |rows| { + let payload: ArrayRef = match data_type { + DataType::Utf8 => Arc::new(StringArray::from_iter_values(rows.map(payload))), + DataType::Utf8View => { + Arc::new(StringViewArray::from_iter_values(rows.map(payload))) + } + DataType::List(_) => { + // Four elements of a quarter row each. + let mut builder = ListBuilder::new(StringBuilder::new()); + for i in rows { + let quarter = format!("{i:08}{}", "x".repeat(ROW_BYTES / 4)); + for _ in 0..4 { + builder.values().append_value(&quarter); + } + builder.append(true); + } + Arc::new(builder.finish()) + } + _ => unreachable!(), + }; + vec![payload] + }); + sort_spills_and_matches(schema, batches).await; +} + +#[tokio::test] +async fn sort_wide_utf8_rows_spills_and_finishes() { + sort_wide_rows(DataType::Utf8).await; +} + +#[tokio::test] +async fn sort_wide_utf8_view_rows_spills_and_finishes() { + sort_wide_rows(DataType::Utf8View).await; +} + +/// A list column keeps its whole child array alive when it is sliced, so the spill writer +/// must measure and copy only the child range that each piece uses. +#[tokio::test] +async fn sort_wide_list_rows_spills_and_finishes() { + sort_wide_rows(DataType::new_list(DataType::Utf8, true)).await; +} + +fn view_field(name: &str) -> Field { + Field::new(name, DataType::Utf8View, true) +} + +/// Input batches that are slices of a few shared view arrays, as a scan or a `take` over a +/// larger batch emits. Each group of 4 input batches shares one parent's buffers (1 MiB of +/// values), so the sort's in-memory batches and its spilled batches reference the same +/// buffers. The sort charges each input slice for its parent's buffers, so a parent must +/// fit in the pool. +#[tokio::test] +async fn sort_view_slices_sharing_buffers() { + let group = 4 * INPUT_BATCH_ROWS; + let parents: Vec = (0..ROWS) + .step_by(group) + .map(|start| { + StringViewArray::from_iter_values((start..start + group).map(payload)) + }) + .collect(); + let (schema, batches) = input(vec![view_field("v")], |rows| { + let parent = &parents[rows.start / group]; + vec![Arc::new(parent.slice(rows.start % group, rows.len()))] + }); + sort_spills_and_matches(schema, batches).await; +} + +/// Input batches that are sparse `take`s of shared view arrays: not slices, so their +/// buffers hold values of other batches' rows that only garbage collection drops. +#[tokio::test] +async fn sort_sparse_taken_views_sharing_buffers() { + // Each parent holds the rows of 4 input batches, interleaved. + let group = 4 * INPUT_BATCH_ROWS; + let parents: Vec = (0..ROWS) + .step_by(group) + .map(|start| { + Arc::new(StringViewArray::from_iter_values( + (start..start + group).map(payload), + )) as ArrayRef + }) + .collect(); + let (schema, batches) = input(vec![view_field("v")], |rows| { + let offset = rows.start % group; + let indices = UInt32Array::from_iter_values( + (0..INPUT_BATCH_ROWS).map(|i| (offset + i) as u32), + ); + vec![take(&parents[rows.start / group], &indices, None).unwrap()] + }); + sort_spills_and_matches(schema, batches).await; +} + +/// Two view columns of each batch share one set of buffers, in different row orders, +/// with short inline values and nulls mixed in. +#[tokio::test] +async fn sort_view_columns_sharing_buffers_with_inline_and_null_values() { + let (schema, batches) = input(vec![view_field("a"), view_field("b")], |rows| { + let a: ArrayRef = + Arc::new(StringViewArray::from_iter(rows.clone().map(|i| { + match i % 3 { + 0 => None, + 1 => Some(format!("s{i}")), + _ => Some(payload(i)), + } + }))); + let reversed = UInt32Array::from_iter_values((0..rows.len() as u32).rev()); + let b = take(&a, &reversed, None).unwrap(); + vec![a, b] + }); + sort_spills_and_matches(schema, batches).await; +} + +#[tokio::test] +async fn sort_wide_binary_view_rows() { + let field = Field::new("b", DataType::BinaryView, false); + let (schema, batches) = input(vec![field], |rows| { + vec![Arc::new(BinaryViewArray::from_iter_values( + rows.map(|i| payload(i).into_bytes()), + ))] + }); + sort_spills_and_matches(schema, batches).await; +} + +fn dictionary_field(values: DataType) -> Field { + Field::new( + "d", + DataType::Dictionary(Box::new(DataType::Int32), Box::new(values)), + false, + ) +} + +/// Each input batch has its own dictionary, with one wide value per row plus unused +/// values. The sort concatenates and slices these, so the spilled batches carry merged +/// dictionaries that the spill writer prunes. +async fn sort_wide_dictionary_rows(values: DataType) { + let (schema, batches) = input(vec![dictionary_field(values.clone())], |rows| { + // Used values first, then as many unused ones. + let strings = rows + .clone() + .chain(rows.clone().map(|i| i + ROWS)) + .map(payload); + let values: ArrayRef = match values { + DataType::Utf8 => Arc::new(StringArray::from_iter_values(strings)), + DataType::Utf8View => Arc::new(StringViewArray::from_iter_values(strings)), + _ => unreachable!(), + }; + let keys = Int32Array::from_iter_values(0..rows.len() as i32); + vec![Arc::new(DictionaryArray::try_new(keys, values).unwrap())] + }); + sort_spills_and_matches(schema, batches).await; +} + +#[tokio::test] +async fn sort_wide_dictionary_rows_spills_and_finishes() { + sort_wide_dictionary_rows(DataType::Utf8).await; +} + +#[tokio::test] +async fn sort_wide_dictionary_of_views_rows_spills_and_finishes() { + sort_wide_dictionary_rows(DataType::Utf8View).await; +} + +/// Views below the top level, in a struct and in a list. +#[tokio::test] +async fn sort_wide_nested_view_rows() { + let struct_fields = Fields::from(vec![view_field("v")]); + let item = Arc::new(Field::new_list_field(DataType::Utf8View, true)); + let fields = vec![ + Field::new("s", DataType::Struct(struct_fields.clone()), false), + Field::new("l", DataType::List(Arc::clone(&item)), false), + ]; + let (schema, batches) = input(fields, |rows| { + let views = |rows: std::ops::Range| -> ArrayRef { + Arc::new(StringViewArray::from_iter_values( + rows.map(|i| format!("{i:08}{}", "x".repeat(ROW_BYTES / 2))), + )) + }; + let structs = + StructArray::try_new(struct_fields.clone(), vec![views(rows.clone())], None) + .unwrap(); + let lists = ListArray::try_new( + Arc::clone(&item), + OffsetBuffer::from_lengths(vec![1; rows.len()]), + views(rows), + None, + ) + .unwrap(); + vec![Arc::new(structs), Arc::new(lists)] + }); + sort_spills_and_matches(schema, batches).await; +} diff --git a/datafusion/physical-plan/Cargo.toml b/datafusion/physical-plan/Cargo.toml index 25f8e437407a1..ef5f7a70763eb 100644 --- a/datafusion/physical-plan/Cargo.toml +++ b/datafusion/physical-plan/Cargo.toml @@ -130,6 +130,11 @@ name = "invariant_check" harness = false name = "union_schema" +[[bench]] +harness = false +name = "sort_spill" +required-features = ["test_utils"] + [[bench]] harness = false name = "spill_io" diff --git a/datafusion/physical-plan/benches/sort_spill.rs b/datafusion/physical-plan/benches/sort_spill.rs new file mode 100644 index 0000000000000..0e7da7329dc21 --- /dev/null +++ b/datafusion/physical-plan/benches/sort_spill.rs @@ -0,0 +1,134 @@ +// 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 + +//! Benchmarks for an external sort that spills under a memory limit. +//! +//! Each case sorts one partition of `NUM_BATCHES` batches of `BATCH_SIZE` rows with a +//! memory pool smaller than the input, so the sort writes spill files and merges them. + +use std::sync::Arc; + +use arrow::array::{ArrayRef, DictionaryArray, Int64Array, StringArray}; +use arrow::datatypes::{DataType, Field, Int32Type, Schema}; +use arrow::record_batch::RecordBatch; +use criterion::{BenchmarkId, Criterion, criterion_group, criterion_main}; +use datafusion_execution::TaskContext; +use datafusion_execution::config::SessionConfig; +use datafusion_execution::memory_pool::FairSpillPool; +use datafusion_execution::runtime_env::RuntimeEnvBuilder; +use datafusion_physical_expr::{LexOrdering, PhysicalSortExpr, expressions::col}; +use datafusion_physical_plan::sorts::sort::SortExec; +use datafusion_physical_plan::test::TestMemoryExec; +use datafusion_physical_plan::{ExecutionPlan, collect}; +use rand::rngs::StdRng; +use rand::{Rng, SeedableRng}; +use tokio::runtime::Runtime; + +const NUM_BATCHES: usize = 16; +const BATCH_SIZE: usize = 8192; + +/// A sort key and one payload column, in random key order. +fn make_batches(payload: impl Fn(&mut StdRng, usize) -> ArrayRef) -> Vec { + let mut rng = StdRng::seed_from_u64(42); + (0..NUM_BATCHES) + .map(|_| { + let keys: ArrayRef = Arc::new(Int64Array::from_iter_values( + (0..BATCH_SIZE).map(|_| rng.random::()), + )); + let payload = payload(&mut rng, BATCH_SIZE); + let schema = Arc::new(Schema::new(vec![ + Field::new("key", DataType::Int64, false), + Field::new("payload", payload.data_type().clone(), false), + ])); + RecordBatch::try_new(schema, vec![keys, payload]).unwrap() + }) + .collect() +} + +/// Dictionary with many distinct values, shared by all input batches. +fn dictionary_batches() -> Vec { + const DISTINCT: usize = 100_000; + let values: ArrayRef = Arc::new(StringArray::from_iter_values( + (0..DISTINCT).map(|i| format!("value number {i:08}")), + )); + make_batches(|rng, rows| { + let keys = (0..rows).map(|_| rng.random_range(0..DISTINCT as i32)); + Arc::new( + DictionaryArray::::try_new(keys.collect(), Arc::clone(&values)) + .unwrap(), + ) + }) +} + +/// Short strings, so no spilled batch is split. +fn narrow_utf8_batches() -> Vec { + make_batches(|rng, rows| { + Arc::new(StringArray::from_iter_values( + (0..rows).map(|_| format!("string {:016}", rng.random::())), + )) + }) +} + +/// 1 KiB strings: 8 MiB per batch, larger than the spill batch bound. +fn wide_utf8_batches() -> Vec { + make_batches(|rng, rows| { + Arc::new(StringArray::from_iter_values((0..rows).map(|_| { + format!("{:016}{}", rng.random::(), "x".repeat(1008)) + }))) + }) +} + +fn sort_spilled(rt: &Runtime, batches: &[RecordBatch], pool_bytes: usize) { + let schema = batches[0].schema(); + let input = + TestMemoryExec::try_new_exec(&[batches.to_vec()], Arc::clone(&schema), None) + .unwrap(); + let ordering = + LexOrdering::new([PhysicalSortExpr::new_default(col("key", &schema).unwrap())]) + .unwrap(); + let sort = Arc::new(SortExec::new(ordering, input)); + let runtime = RuntimeEnvBuilder::new() + .with_memory_pool(Arc::new(FairSpillPool::new(pool_bytes))) + .build_arc() + .unwrap(); + let task_ctx = Arc::new( + TaskContext::default() + .with_session_config(SessionConfig::new().with_batch_size(BATCH_SIZE)) + .with_runtime(runtime), + ); + let output = rt + .block_on(collect( + Arc::clone(&sort) as Arc, + task_ctx, + )) + .unwrap(); + assert_eq!( + output.iter().map(|b| b.num_rows()).sum::(), + NUM_BATCHES * BATCH_SIZE + ); + let spills = sort.metrics().and_then(|m| m.spill_count()).unwrap_or(0); + assert!(spills > 0, "the sort must spill"); +} + +fn criterion_benchmark(c: &mut Criterion) { + let rt = Runtime::new().unwrap(); + let cases: [(&str, Vec, usize); 3] = [ + ("dictionary", dictionary_batches(), 32 * 1024 * 1024), + ("narrow_utf8", narrow_utf8_batches(), 16 * 1024 * 1024), + ("wide_utf8", wide_utf8_batches(), 64 * 1024 * 1024), + ]; + let mut group = c.benchmark_group("sort_spill"); + group.sample_size(10); + for (name, batches, pool_bytes) in &cases { + group.bench_with_input( + BenchmarkId::from_parameter(name), + batches, + |b, batches| b.iter(|| sort_spilled(&rt, batches, *pool_bytes)), + ); + } + group.finish(); +} + +criterion_group!(benches, criterion_benchmark); +criterion_main!(benches); diff --git a/datafusion/physical-plan/src/aggregates/grouped_hash_stream.rs b/datafusion/physical-plan/src/aggregates/grouped_hash_stream.rs index 1ee6d02d13e8e..0ffad45e7e9f8 100644 --- a/datafusion/physical-plan/src/aggregates/grouped_hash_stream.rs +++ b/datafusion/physical-plan/src/aggregates/grouped_hash_stream.rs @@ -24,6 +24,7 @@ use std::vec; use super::aggregate_hash_table::{accumulator_phases, create_group_accumulator}; use super::order::GroupClustering; use super::skip_partial::SkipAggregationProbe; +use super::spill::bound_spill_batches_by_first_spill; use super::{AggregateExec, format_human_display}; use crate::aggregates::group_values::{ AccumulatorPhase, AggregateAccumulatorMetrics, AggregateArgumentMetrics, @@ -1250,6 +1251,12 @@ impl GroupedHashAggregateStream { ) })?; + if self.spill_state.spills.is_empty() { + self.spill_state.spill_manager = bound_spill_batches_by_first_spill( + &self.spill_state.spill_manager, + &emit, + )?; + } let sorted_iter = IncrementalSortIterator::new( emit, self.spill_state.spill_expr.clone(), diff --git a/datafusion/physical-plan/src/aggregates/mod.rs b/datafusion/physical-plan/src/aggregates/mod.rs index 39389ad77df9b..827c140ab6422 100644 --- a/datafusion/physical-plan/src/aggregates/mod.rs +++ b/datafusion/physical-plan/src/aggregates/mod.rs @@ -3927,7 +3927,13 @@ mod tests { } } assert_eq!(seen.len(), 2 * KEYS_PER_PREFIX); - assert!(pool.peak_reserved() <= MEMORY_LIMIT); + // The replay merge grows a second reservation of the aggregate's consumer + // while the first still holds a few bytes of group state. `FairSpillPool` + // checks each reservation against the fair share on its own, so the two + // together can exceed the share by what the first one holds (16 bytes + // here). That is a pool bug, not a spill bug, so the test allows it. + const SIBLING_RESERVATION_SLACK: usize = 1024; + assert!(pool.peak_reserved() <= MEMORY_LIMIT + SIBLING_RESERVATION_SLACK); let metrics = aggregate.metrics().unwrap(); assert!(metrics.spill_count().unwrap() > 1); assert!(metrics.spilled_rows().unwrap() > 0); diff --git a/datafusion/physical-plan/src/aggregates/spill.rs b/datafusion/physical-plan/src/aggregates/spill.rs index c08a76adc2153..e4fc47f4feac1 100644 --- a/datafusion/physical-plan/src/aggregates/spill.rs +++ b/datafusion/physical-plan/src/aggregates/spill.rs @@ -23,9 +23,11 @@ use super::{AggregateExec, AggregateMode}; use crate::SendableRecordBatchStream; use crate::metrics::{BaselineMetrics, SpillMetrics}; use crate::sorts::IncrementalSortIterator; +use crate::sorts::sort::spill_max_batch_bytes; use crate::sorts::streaming_merge::{SortedSpillFile, StreamingMergeBuilder}; use crate::spill::gc_nested_arrays; -use crate::spill::spill_manager::SpillManager; +use crate::spill::in_progress_spill_file::WrittenBatchStats; +use crate::spill::spill_manager::{GetSlicedSize, SpillManager}; use crate::stream::RecordBatchStreamAdapter; /// Target size in bytes of the batches that spilling writes and replay @@ -193,7 +195,8 @@ impl AggregateSpill { spill_metrics, Arc::clone(spill_schema), ) - .with_compression_type(context.session_config().spill_compression()); + .with_compression_type(context.session_config().spill_compression()) + .with_max_batch_bytes(Some(SPILL_BATCH_TARGET_BYTES)); Ok(Self { replay_agg, @@ -222,29 +225,34 @@ impl AggregateSpill { let Some(state_batch) = state_batch else { return Ok(()); }; + if self.spills.is_empty() { + self.spill_manager = + bound_spill_batches_by_first_spill(&self.spill_manager, &state_batch)?; + } - 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()); + .with_max_batch_bytes(SPILL_BATCH_TARGET_BYTES); + let mut in_progress_file = + self.spill_manager.create_in_progress_file(self.label)?; + let mut written = WrittenBatchStats::default(); + for batch in sorted_iter { + let batch = batch?; + if batch.num_rows() == 0 { + continue; } - }); - let spill_file = self - .spill_manager - .spill_record_batch_iter_and_return_max_batch_memory( - sorted_iter, - self.label, - )?; - - let Some((file, max_record_batch_memory)) = spill_file else { + let stats = in_progress_file.append_batch_with_stats(&batch)?; + written.memory_size = written.memory_size.max(stats.memory_size); + written.rows = written.rows.max(stats.rows); + } + let Some(file) = in_progress_file.finish()? else { return internal_err!("{}: produced an empty spill", self.label); }; + let (max_record_batch_memory, max_batch_rows) = + (written.memory_size, written.rows); self.spills.push(SortedSpillFile { file, @@ -297,6 +305,7 @@ impl AggregateSpill { } = self; let spill_schema = Arc::clone(spill_manager.schema()); + let max_batch_bytes = spill_manager.max_batch_bytes(); // The merge and replay table are two components of the same aggregate // operator. Keep them under one consumer registration so a fair memory // pool does not divide this operator's quota between its own phases. @@ -308,7 +317,7 @@ 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_max_batch_bytes(max_batch_bytes) .with_reservation(merge_reservation) .with_replay_headroom() .with_intermediate_merge_sizing(Some(min_spill_batch_rows)) @@ -336,3 +345,21 @@ impl AggregateSpill { Ok(replay.into_stream()) } } + +/// Lowers the bound on the batches of an aggregate's sorted spill runs (see +/// [`SpillManager::with_max_batch_bytes`]) to fit the state of its first spill. +/// +/// Unlike the sort, the merge of an aggregate's runs holds no workspace of its +/// own. The batches are bounded so that the smallest merge, plus as much again +/// for the replay, fits in what the aggregate held when memory first ran out. +/// Later spills, such as the remainder spilled at the end of the input, can +/// hold much less, so they do not lower the bound. +pub(super) fn bound_spill_batches_by_first_spill( + spill_manager: &SpillManager, + first_spill: &RecordBatch, +) -> Result { + let held = first_spill.get_sliced_size()?; + let bound = spill_max_batch_bytes(held / 2) + .min(spill_manager.max_batch_bytes().unwrap_or(usize::MAX)); + Ok(spill_manager.clone().with_max_batch_bytes(Some(bound))) +} diff --git a/datafusion/physical-plan/src/sorts/multi_level_merge.rs b/datafusion/physical-plan/src/sorts/multi_level_merge.rs index ebdecb90f38fa..cc5900599cfe2 100644 --- a/datafusion/physical-plan/src/sorts/multi_level_merge.rs +++ b/datafusion/physical-plan/src/sorts/multi_level_merge.rs @@ -14,14 +14,12 @@ use std::sync::Arc; use std::task::{Context, Poll}; use arrow::datatypes::SchemaRef; -use datafusion_common::{Result, internal_err, resources_err}; +use datafusion_common::{Result, resources_err}; use datafusion_execution::memory_pool::{MemoryReservation, MergeMemoryPool}; use crate::sorts::builder::try_grow_reservation_to_at_least; use crate::sorts::sort::get_reserved_bytes_for_record_batch_size; use crate::sorts::streaming_merge::{SortedSpillFile, StreamingMergeBuilder}; -use crate::spill::gc_view_arrays; -use crate::spill::spill_manager::GetSlicedSize; use crate::stream::{ObservedStream, RecordBatchStreamAdapter}; use datafusion_execution::{RecordBatchStream, SendableRecordBatchStream}; use datafusion_physical_expr_common::sort_expr::LexOrdering; @@ -107,36 +105,32 @@ use futures::{Stream, StreamExt}; /// /// ## Memory Management Strategy /// -/// This multi-level merge make sure that we can handle any amount of data to sort as long as -/// we have enough memory to merge at least 2 streams at a time, even when individual record -/// batches are skewed (very wide). +/// The merge relies on one property of its inputs: every batch of a sorted spill +/// run is at most the bound its [`SpillManager`] sets, unless it holds a single row +/// or a column that cannot be split by rows (see [`SpillManager::with_max_batch_bytes`]). +/// The writers of sorted runs (the sort, the aggregate, and this merge for its +/// intermediate runs) get it from their [`SpillManager`], which splits larger +/// batches as it writes them. /// -/// 1. **Worst-Case Memory Reservation**: Reserves memory based on the largest -/// batch size encountered in each spill file to merge, ensuring sufficient memory is always -/// available during merge operations. -/// 2. **Adaptive Buffer Sizing**: Reduces buffer sizes when memory is constrained -/// 3. **Spill-to-Disk**: Spill to disk when we cannot merge all files in memory -/// 4. **Re-spilling Skewed Runs**: If even at the smallest read-buffer size we still cannot -/// reserve memory for the minimum of 2 streams - because a single run's largest batch is so -/// wide that two streams' worth of reservation exceeds the budget - the larger of the two -/// runs is re-spilled with each batch sliced in half. This shrinks its largest batch, -/// lowering the per-stream reservation, and the merge pass is retried. The re-spilled run -/// is tracked alongside a per-run batch-size limit equal to half the batch size it was -/// written with, so any later merge that includes it caps its output batch size to match - -/// otherwise the merged run could rebuild a full-size batch and reintroduce the skew. -/// Crucially the global merge batch size is *not* lowered, so re-spilling more than one run -/// does not compound the reduction. If a batch cannot be split any further (a single row -/// wider than the budget), the merge surfaces `ResourcesExhausted` instead of looping -/// forever. +/// 1. **Worst-Case Memory Reservation**: Reserves memory for the largest batch of +/// each run it reads, times the number of batches read ahead. +/// 2. **Adaptive Buffer Sizing**: Reads fewer batches ahead when memory is +/// constrained. +/// 3. **Replay Headroom**: When the caller asks for it (the aggregate replaying +/// its spills), a run is admitted only if the pool could give as much again. +/// Only the two runs a pass needs, with one batch read ahead, may use the +/// whole pool. +/// 4. **Spill-to-Disk**: Spill to disk when we cannot merge all runs in memory. +/// When the caller asks for it, an intermediate pass merges only as many runs +/// as needed for the remaining runs to fit one final pass. +/// +/// If two runs do not fit even with one batch read ahead (single rows or columns +/// that cannot be split, or a pool smaller than the bound was chosen for), the +/// merge returns `ResourcesExhausted`. pub(crate) struct MultiLevelMergeBuilder { spill_manager: SpillManager, schema: SchemaRef, - /// Sorted runs still to be merged. Each run is paired with the batch-size limit a - /// merge consuming it must cap its output at. Runs written at the full batch size - /// carry `batch_size`. A run re-spilled smaller to resolve skew carries its halved - /// limit (see [`Self::split_spill_file_in_half`]). Tracking it here keeps this limit - /// out of the public [`SortedSpillFile`], so no external caller has to set it. - sorted_spill_files: Vec<(SortedSpillFile, usize)>, + sorted_spill_files: Vec, sorted_streams: Vec, expr: LexOrdering, metrics: BaselineMetrics, @@ -144,7 +138,7 @@ pub(crate) struct MultiLevelMergeBuilder { /// See [`StreamingMergeBuilder::with_max_batch_bytes`] max_batch_bytes: Option, reservation: MemoryReservation, - /// Workspace retained across retries and intermediate spill passes. + /// Workspace retained across intermediate spill passes. merge_pool: Option>, /// Leave memory for the aggregate consuming the merged spill rows. reserve_replay_headroom: bool, @@ -154,6 +148,9 @@ pub(crate) struct MultiLevelMergeBuilder { enable_round_robin_tie_breaker: bool, } +/// The fewest spill files a merge reads at once. +pub(super) const MIN_SPILL_MERGE_FAN_IN: usize = 2; + impl Debug for MultiLevelMergeBuilder { fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { write!(f, "MultiLevelMergeBuilder") @@ -177,12 +174,7 @@ impl MultiLevelMergeBuilder { Self { spill_manager, schema, - // Initial runs are written at the full batch size, so they impose no cap - // on later merges - record `batch_size` as their (unconstrained) limit. - sorted_spill_files: sorted_spill_files - .into_iter() - .map(|file| (file, batch_size)) - .collect(), + sorted_spill_files, sorted_streams, expr, metrics, @@ -209,8 +201,7 @@ impl MultiLevelMergeBuilder { self } - /// Leave replay headroom while selecting merge buffers. Temporary splitting - /// workspace can still use the full pool because replay has not started. + /// Leave replay headroom while selecting merge buffers. pub(super) fn with_replay_headroom(mut self, reserve: bool) -> Self { self.reserve_replay_headroom = reserve; self @@ -235,38 +226,8 @@ impl MultiLevelMergeBuilder { } async fn create_stream(mut self) -> Result { - let mut allow_minimum_without_headroom = false; loop { - let (mut stream, batch_size_limit) = match self - .merge_sorted_runs_within_mem_limit(allow_minimum_without_headroom)? - { - MergeStep::Stream { - stream, - batch_size_limit, - } => (stream, batch_size_limit), - MergeStep::SplitThenRetry(index) => { - // Couldn't reserve memory for the minimum of 2 streams. Re-spill - // the larger of the two we're trying to merge with half its batch - // size so its largest batch shrinks, lowering the per-stream - // reservation, then retry. Makes the merge resilient to skewed - // (very wide) rows. - let retry_unsplittable = - self.reserve_replay_headroom && !allow_minimum_without_headroom; - if !self - .split_spill_file_in_half(index, retry_unsplittable) - .await? - { - // A single row may prevent leaving replay headroom while - // the minimum merge still fits the actual shared pool. - allow_minimum_without_headroom = true; - } - continue; - } - }; - allow_minimum_without_headroom = false; - - // TODO - add a threshold for number of files to disk even if empty and reading from disk so - // we can avoid the memory reservation + let mut stream = self.merge_sorted_runs_within_mem_limit()?; // If no spill files are left, we can return the stream as this is the last sorted run // TODO - We can write to disk before reading it back to avoid having multiple streams in memory @@ -286,7 +247,7 @@ impl MultiLevelMergeBuilder { } // Need to sort to a spill file - let Some((spill_file, max_record_batch_memory, max_batch_rows)) = self + let Some((file, max_record_batch_memory, max_batch_rows)) = self .spill_manager .spill_record_batch_stream_and_return_max_batch_stats( &mut stream, @@ -297,21 +258,14 @@ impl MultiLevelMergeBuilder { continue; }; - // Offset overflow can flush partial batches even without a fetch limit. - // Keep sizing disabled once a run cannot fill its output batch limit. - self.size_intermediate_merges &= max_batch_rows >= batch_size_limit; - - // Add the spill file paired with the batch-size limit of the merge that - // produced it: if that merge consumed a shrunk (skew-resolved) run, its - // output was capped and this intermediate run is likewise capped, so a - // later pass that re-merges it won't rebuild an oversized batch. - self.sorted_spill_files.push(( - SortedSpillFile { - file: spill_file, - max_record_batch_memory, - }, - batch_size_limit, - )); + // Offset overflow or the byte bound can end batches early. Keep + // sizing disabled once a run cannot fill its output batch. + self.size_intermediate_merges &= max_batch_rows >= self.batch_size; + + self.sorted_spill_files.push(SortedSpillFile { + file, + max_record_batch_memory, + }); } } @@ -319,59 +273,42 @@ impl MultiLevelMergeBuilder { /// as possible within the memory limit. fn merge_sorted_runs_within_mem_limit( &mut self, - allow_minimum_without_headroom: bool, - ) -> Result { + ) -> Result { match (self.sorted_spill_files.len(), self.sorted_streams.len()) { // No data so empty batch (0, 0) => { let empty_stream = Box::pin(EmptyRecordBatchStream::new(Arc::clone(&self.schema))); - Ok(MergeStep::Stream { - stream: self.observe_output(empty_stream), - batch_size_limit: self.batch_size, - }) + Ok(self.observe_output(empty_stream)) } // Only in-memory stream, return that (0, 1) => { let output_stream = self.sorted_streams.remove(0); - Ok(MergeStep::Stream { - stream: self.observe_output(output_stream), - batch_size_limit: self.batch_size, - }) + Ok(self.observe_output(output_stream)) } // Only single sorted spill file so return it (1, 0) => { - let (spill_file, batch_size) = self.sorted_spill_files.remove(0); + let spill_file = self.sorted_spill_files.remove(0); // Not reserving any memory for this disk as we are not holding it in memory let output_stream = self .spill_manager .read_spill_as_stream(spill_file.file, None)?; - - Ok(MergeStep::Stream { - stream: self.observe_output(output_stream), - batch_size_limit: batch_size, - }) + Ok(self.observe_output(output_stream)) } - // Only in memory streams, so merge them all in a single pass. In-memory - // runs are never shrunk for skew, so this merge runs at the full batch - // size and its output carries no limit. + // Only in memory streams, so merge them all in a single pass (0, _) => { let sorted_stream = mem::take(&mut self.sorted_streams); // No need to wrap with observed stream since merge sort will update the observed metrics - Ok(MergeStep::Stream { - stream: self.create_new_merge_sort( - sorted_stream, - // If we have no sorted spill files left, this is the last run - true, - true, - self.batch_size, - )?, - batch_size_limit: self.batch_size, - }) + self.create_new_merge_sort( + sorted_stream, + // If we have no sorted spill files left, this is the last run + true, + true, + ) } // Need to merge multiple streams @@ -383,30 +320,13 @@ impl MultiLevelMergeBuilder { // allocation. let mut memory_reservation = self.reservation.take(); - // Compute the minimum before taking the in-memory streams so that, if we - // need to re-spill and retry, `self.sorted_streams` is left untouched. - let minimum_number_of_required_streams = - 2_usize.saturating_sub(self.sorted_streams.len()); - - let (mut sorted_spill_files, buffer_size) = match self + let (mut sorted_spill_files, buffer_size) = self .get_sorted_spill_files_to_merge( 2, // we must have at least 2 streams to merge - minimum_number_of_required_streams, + MIN_SPILL_MERGE_FAN_IN.saturating_sub(self.sorted_streams.len()), &mut memory_reservation, - allow_minimum_without_headroom, - )? { - SpillFilesToMerge::Ready(sorted_spill_files, buffer_size) => { - (sorted_spill_files, buffer_size) - } - // Not enough memory to seat 2 streams. Re-spill the blocking file - // smaller and retry. `get_sorted_spill_files_to_merge` already freed - // the reservation and `self.sorted_streams` is untouched, so the - // retry starts clean. - SpillFilesToMerge::SplitThenRetry(index) => { - return Ok(MergeStep::SplitThenRetry(index)); - } - }; + )?; if self.size_intermediate_merges && self.sorted_streams.is_empty() @@ -422,21 +342,20 @@ impl MultiLevelMergeBuilder { } else { final_fan_in }; - let needed = self.sorted_spill_files.len() + sorted_spill_files.len() - - final_fan_in - + 1; + let needed = + (self.sorted_spill_files.len() + sorted_spill_files.len() + 1) + .saturating_sub(final_fan_in); if needed < sorted_spill_files.len() { - let (first, first_limit) = &sorted_spill_files[0]; - // A wider pending run or a smaller output batch limit can - // change the next pass's fan-in and add intermediate work. - // Apply this sizing heuristic only to uniform run budgets. + let first = &sorted_spill_files[0]; + // A wider pending run can change the next pass's + // fan-in and add intermediate work. Apply this sizing + // heuristic only to uniform run budgets. let uniform_runs = sorted_spill_files .iter() .chain(&self.sorted_spill_files) - .all(|(file, limit)| { + .all(|file| { file.max_record_batch_memory == first.max_record_batch_memory - && limit == first_limit }); if uniform_runs { let unused = sorted_spill_files.split_off(needed); @@ -464,15 +383,7 @@ impl MultiLevelMergeBuilder { mem::swap(&mut self.reservation, &mut memory_reservation); } - // Cap the merge output at the smallest limit among the runs we're - // about to merge. Runs that were shrunk for skew carry a smaller limit, - // if none do, every run carries `self.batch_size` and the merge runs at - // the full batch size. The output stream is tagged with the same limit - // (see the `MergeStep::Stream` returns below) so a re-spilled - // intermediate run stays shrunk and won't rebuild an oversized batch on - // a later pass. - let mut output_batch_size = self.batch_size; - for (spill, batch_size_limit) in sorted_spill_files { + for spill in sorted_spill_files { let stream = self .spill_manager .clone() @@ -481,7 +392,6 @@ impl MultiLevelMergeBuilder { spill.file, Some(spill.max_record_batch_memory), )?; - output_batch_size = output_batch_size.min(batch_size_limit); sorted_streams.push(stream); } let merge_sort_stream = self.create_new_merge_sort( @@ -489,7 +399,6 @@ impl MultiLevelMergeBuilder { // If we have no sorted spill files left, this is the last run self.sorted_spill_files.is_empty(), is_only_merging_memory_streams, - output_batch_size, )?; // If we're only merging memory streams, we don't need to attach the memory reservation @@ -501,20 +410,14 @@ impl MultiLevelMergeBuilder { "when only merging memory streams, we should not have any memory reservation and let the merge sort handle the memory" ); - Ok(MergeStep::Stream { - stream: merge_sort_stream, - batch_size_limit: output_batch_size, - }) + Ok(merge_sort_stream) } else { // Attach the memory reservation to the stream to make sure we have enough memory // throughout the merge process as we bypassed the memory pool for the merge sort stream - Ok(MergeStep::Stream { - stream: Box::pin(StreamAttachedReservation::new( - merge_sort_stream, - memory_reservation, - )), - batch_size_limit: output_batch_size, - }) + Ok(Box::pin(StreamAttachedReservation::new( + merge_sort_stream, + memory_reservation, + ))) } } } @@ -525,12 +428,12 @@ impl MultiLevelMergeBuilder { streams: Vec, is_output: bool, all_in_memory: bool, - output_batch_size: usize, ) -> Result { let mut builder = StreamingMergeBuilder::new() .with_schema(Arc::clone(&self.schema)) .with_expressions(&self.expr) - .with_batch_size(output_batch_size) + .with_batch_size(self.batch_size) + .with_max_batch_bytes(self.max_batch_bytes) .with_fetch(self.fetch) .with_metrics(if is_output { // Only add the metrics to the last run @@ -539,8 +442,7 @@ impl MultiLevelMergeBuilder { self.metrics.intermediate() }) .with_round_robin_tie_breaker(self.enable_round_robin_tie_breaker) - .with_streams(streams) - .with_max_batch_bytes(self.max_batch_bytes); + .with_streams(streams); if !all_in_memory { // Don't track memory used by this stream as we reserve that memory by worst case sceneries @@ -559,56 +461,41 @@ impl MultiLevelMergeBuilder { builder.build() } - /// Return the sorted spill files to use for the next phase, and the buffer size - /// This will try to get as many spill files as possible to merge, and if we don't have enough streams - /// it will try to reduce the buffer size until we have enough streams to merge - /// otherwise it will return an error + /// Returns the sorted spill files to read in the next pass, and how many + /// batches to read ahead from each. + /// + /// Takes as many runs as fit, at least `minimum_number_of_required_streams`, + /// leaving replay headroom if the caller asked for it (see the type docs). If + /// the minimum does not fit, retries with fewer batches read ahead; with one + /// batch read ahead the minimum only has to fit the pool, or the pool's error + /// is returned. fn get_sorted_spill_files_to_merge( &mut self, buffer_len: usize, minimum_number_of_required_streams: usize, reservation: &mut MemoryReservation, - allow_minimum_without_headroom: bool, - ) -> Result { + ) -> Result<(Vec, usize)> { assert_ne!(buffer_len, 0, "Buffer length must be greater than 0"); - let mut number_of_spills_to_read_for_current_phase = 0; - let configured_fan_in = self - .spill_manager - .env() - .disk_manager - .max_spill_merge_fan_in(); - let max_spill_files = effective_spill_merge_fan_in(configured_fan_in); - // Track total memory needed for spill file buffers. When the - // reservation has pre-reserved bytes (from sort_spill_reservation_bytes), - // those bytes cover the first N spill files without additional pool - // allocation, preventing starvation under memory pressure. - let mut total_needed: usize = 0; + let max_spill_files = self.max_fan_in(); + let mut number_of_spills = 0; let mut accepted_memory: usize = 0; - for (spill, _) in &self.sorted_spill_files { - if number_of_spills_to_read_for_current_phase >= max_spill_files - || (allow_minimum_without_headroom - && number_of_spills_to_read_for_current_phase - >= minimum_number_of_required_streams) - { + for spill in &self.sorted_spill_files { + if number_of_spills >= max_spill_files { break; } - let per_spill = get_reserved_bytes_for_record_batch_size( - spill.max_record_batch_memory, - // Size will be the same as the sliced size, bc it is a spilled batch. - spill.max_record_batch_memory, - ) * buffer_len; - total_needed += per_spill; - - // If a run cannot shrink, allow only the minimum merge without - // replay headroom. Disable read-ahead and still ask the pool for - // every byte used by the merge. - let skip_headroom = allow_minimum_without_headroom - && buffer_len == 1 - && number_of_spills_to_read_for_current_phase - < minimum_number_of_required_streams; - let check_headroom = self.reserve_replay_headroom && !skip_headroom; + let total_needed = accepted_memory + + get_reserved_bytes_for_record_batch_size( + spill.max_record_batch_memory, + // Size will be the same as the sliced size, bc it is a spilled batch. + spill.max_record_batch_memory, + ) * buffer_len; + // Only the minimum merge with one batch read ahead may use the + // whole pool: without it, the merge cannot make progress at all. + let check_headroom = self.reserve_replay_headroom + && !(buffer_len == 1 + && number_of_spills < minimum_number_of_required_streams); let admission = if check_headroom { // Ask the pool for merge buffers plus equal replay space, then // return the spare bytes before exposing the merge stream. @@ -622,62 +509,36 @@ impl MultiLevelMergeBuilder { try_grow_reservation_to_at_least(reservation, total_needed) }; match admission { - Ok(_) => { - number_of_spills_to_read_for_current_phase += 1; + Ok(()) => { + number_of_spills += 1; accepted_memory = total_needed; } - // If we can't grow the reservation, we need to stop - Err(err) => { - // We must have at least 2 streams to merge, so if we don't have enough memory - // fail - if minimum_number_of_required_streams - > number_of_spills_to_read_for_current_phase - { - // Free the memory we reserved for this merge as we either try again or fail - reservation.free(); - if buffer_len > 1 { - // Try again with smaller buffer size, it will be slower but at least we can merge - return self.get_sorted_spill_files_to_merge( - buffer_len - 1, - minimum_number_of_required_streams, - reservation, - allow_minimum_without_headroom, - ); - } - - // buffer_len == 1 and we still can't seat the minimum of 2 streams. - if number_of_spills_to_read_for_current_phase == 0 { - if check_headroom { - // Replay has not started, so splitting may use the - // full pool to leave room for replay afterward. - return Ok(SpillFilesToMerge::SplitThenRetry(0)); - } - // We couldn't even reserve a single stream - one record batch - // is larger than the whole merge budget. That's the lone-batch - // case, not the 2-stream merge skew we rescue here - surface it. - return Err(err); - } - - // We seated one stream (index 0) but not the second (index 1, the - // batch that just failed to reserve). Those are by definition the - // only two streams we are trying to merge, so re-spill the larger - // of them with a smaller batch size and retry, the smaller max - // batch lowers the per-stream reservation enough to seat both. - let split_index = usize::from( - self.sorted_spill_files[1].0.max_record_batch_memory - > self.sorted_spill_files[0].0.max_record_batch_memory, + Err(err) if number_of_spills < minimum_number_of_required_streams => { + // Free the memory we reserved for this merge as we either try again or fail + reservation.free(); + if buffer_len > 1 { + // Try again with smaller buffer size, it will be slower but at least we can merge + return self.get_sorted_spill_files_to_merge( + buffer_len - 1, + minimum_number_of_required_streams, + reservation, ); - return Ok(SpillFilesToMerge::SplitThenRetry(split_index)); } - - // We reached the maximum amount of memory we can use - // for this merge - break; + return Err(err.context(format!( + "Cannot merge two sorted spill runs: each needs {} bytes \ + to read one batch; consider increasing the memory limit", + get_reserved_bytes_for_record_batch_size( + spill.max_record_batch_memory, + spill.max_record_batch_memory + ) + ))); } + // We reached the maximum amount of memory we can use for this merge + Err(_) => break, } } - if self.reserve_replay_headroom { + if self.reserve_replay_headroom && reservation.size() > accepted_memory { // `total_needed` may include a rejected candidate. Keep only the // buffers that were admitted, releasing temporary replay headroom. reservation.shrink(reservation.size() - accepted_memory); @@ -685,180 +546,20 @@ impl MultiLevelMergeBuilder { let spills = self .sorted_spill_files - .drain(..number_of_spills_to_read_for_current_phase) + .drain(..number_of_spills) .collect::>(); - Ok(SpillFilesToMerge::Ready(spills, buffer_len)) + Ok((spills, buffer_len)) } - /// Re-spill the spill file at `index` with half its batch size, putting it back - /// at the same position. We read the file back and re-spill it through the normal - /// spill API (which owns batch layout), slicing every batch in two, which halves - /// the largest written batch and so lowers the per-stream merge reservation enough - /// for the next attempt to seat both streams. One stream's worth of memory is - /// reserved for the duration and freed afterwards. Makes the merge resilient to skew. - /// - /// Instead of halving the *global* merge batch size (which would compound when more - /// than one run is re-spilled), the shrunk run records its own smaller batch-size - /// limit (tracked alongside the run in `sorted_spill_files`), so only merges that - /// actually consume it pay the reduced batch size. - /// - /// Returns whether the largest batch shrank. If `retry_unsplittable` is true, - /// restore an unchanged run for one minimum-merge admission attempt. - async fn split_spill_file_in_half( - &mut self, - index: usize, - retry_unsplittable: bool, - ) -> Result { - log::debug!( - "2 spilled streams could not be loaded into memory for merge \ - (requires 2x of the largest batch from both), re-spilling the larger of the two with half \ - the batch size to reduce memory needs for the next merge attempt. the shrunk run carries \ - a halved batch-size limit so only merges consuming it use the smaller batch size" - ); - - // Extract the target in O(1) instead of `remove(index)`, which would shift - // every following spill file. Swap it to the back and pop it; the matching - // swap after re-spilling restores the original order, so the vec ends up - // exactly as it started, just with the target file shrunk. - // `old_batch_size` is the batch size this run was written with (the full merge - // batch size unless it was already shrunk once). Halving it caps the next merge - // that reads this run so the merged output can't rebuild a full-size batch. - let last = self.sorted_spill_files.len() - 1; - self.sorted_spill_files.swap(index, last); - let (mut target, old_batch_size) = self - .sorted_spill_files - .pop() - .expect("index is in bounds, so the vec is non-empty"); - let old_max = target.max_record_batch_memory; - let mut max_batch_rows = old_batch_size; - - // Reserve enough to hold a single stream of this file while we re-spill it. - let reservation = self.reservation.new_empty(); - reservation - .try_grow(get_reserved_bytes_for_record_batch_size(old_max, old_max))?; - - if self.reserve_replay_headroom { - // A maximum-sized singleton cannot shrink. Find it without writing - // another file: the original runs may already fill the disk quota. - // Use an unbuffered reader so no background read outlives this guard. - let mut source = self.spill_manager.read_spill_as_stream_unbuffered( - Arc::clone(&target.file), - Some(old_max), - )?; - let mut all_singletons = true; - let mut max_is_singleton = false; - let mut decoded_max = 0; - max_batch_rows = 0; - while let Some(batch) = source.next().await { - let batch = batch?; - max_batch_rows = max_batch_rows.max(batch.num_rows()); - all_singletons &= batch.num_rows() == 1; - decoded_max = decoded_max.max(batch.get_sliced_size()?); - if batch.num_rows() == 1 - && gc_view_arrays(&batch)?.get_sliced_size()? >= old_max - { - max_is_singleton = true; - break; - } - } - // IPC can discard spare view-buffer capacity included in `old_max`. - // A complete scan can correct that estimate without rewriting the - // file. Use decoded buffers without GC, as the merge retains them. - let shrank = !max_is_singleton && decoded_max < old_max; - if shrank || all_singletons || max_is_singleton { - if !shrank && !retry_unsplittable { - return resources_err!( - "Cannot merge sorted runs: a single record batch of {old_max} bytes \ - exceeds the available merge memory and cannot be split further" - ); - } - if shrank { - target.max_record_batch_memory = decoded_max; - } - let batch_size_limit = if all_singletons || max_is_singleton { - 1 - } else { - old_batch_size.min(max_batch_rows).max(1) - }; - self.sorted_spill_files.push((target, batch_size_limit)); - self.sorted_spill_files.swap(index, last); - return Ok(shrank); - } - } - - let source = self - .spill_manager - .read_spill_as_stream(target.file, Some(old_max))?; - // Re-spill with half the batch size: slice every batch in two. The spill - // writer owns the batch layout, we only change how many rows per batch. - let mut halved: SendableRecordBatchStream = - Box::pin(RecordBatchStreamAdapter::new( - Arc::clone(&self.schema), - source.flat_map(|batch| { - futures::stream::iter(match batch { - Ok(batch) => split_batch_in_half(batch) - .into_iter() - .map(Ok) - .collect::>(), - Err(e) => vec![Err(e)], - }) - }), - )); - - let result = self - .spill_manager - .spill_record_batch_stream_and_return_max_batch_stats( - &mut halved, - "MultiLevelMergeBuilder split skewed spill", - ) - .await?; - - reservation.free(); - - let Some((file, new_max, new_max_batch_rows)) = result else { - return internal_err!("re-spilling a skewed spill file produced no data"); - }; - - // If halving cannot reduce the largest batch, only a requested retry - // against the actual pool can make progress. The caller permits that - // retry once before surfacing this error. - let shrank = new_max < old_max; - if !shrank && !retry_unsplittable { - return resources_err!( - "Cannot merge sorted runs: a single record batch of {old_max} bytes \ - exceeds the available merge memory and cannot be split further" - ); - } - - // Record the halved batch size as a *per-run* limit rather than lowering the - // global batch size. Merges that don't touch this run keep the full batch - // size. a merge that reads it caps its output at this limit so the merged run - // can't rebuild a full-size batch and reintroduce the skew. Actual batches - // can be shorter than the configured limit; use their size so the merge - // cannot accumulate extra split batches outside its reservation. - let new_batch_size_limit = if shrank { - (old_batch_size / 2).min(max_batch_rows.div_ceil(2)).max(1) - } else { - // Skipping headroom only covers indivisible input rows, not a - // larger output batch formed by concatenating those rows. - 1 - }; - self.size_intermediate_merges &= new_max_batch_rows >= new_batch_size_limit; - - // Push the re-spilled (smaller) file and swap it back into `index`, undoing - // the swap-to-back above so the order is preserved. - self.sorted_spill_files.push(( - SortedSpillFile { - file, - max_record_batch_memory: new_max, - }, - new_batch_size_limit, - )); - let last = self.sorted_spill_files.len() - 1; - self.sorted_spill_files.swap(index, last); - - Ok(shrank) + /// The most runs one pass may merge (`max_spill_merge_fan_in`). + fn max_fan_in(&self) -> usize { + effective_spill_merge_fan_in( + self.spill_manager + .env() + .disk_manager + .max_spill_merge_fan_in(), + ) } fn observe_output( @@ -869,47 +570,11 @@ impl MultiLevelMergeBuilder { } } -/// Outcome of trying to reserve memory for one multi-level merge pass. -enum SpillFilesToMerge { - /// Enough memory: the spill files to read this pass (each paired with its - /// batch-size limit) and the read-ahead buffer size. - Ready(Vec<(SortedSpillFile, usize)>, usize), - /// Could not seat the minimum of 2 streams. Re-spill the spill file at this index - /// with a smaller (halved) batch size, then retry the pass. - SplitThenRetry(usize), -} - -/// What one iteration of the multi-level merge loop should do next. -enum MergeStep { - /// A merged stream is ready to be consumed (and possibly spilled back). - Stream { - stream: SendableRecordBatchStream, - /// The batch-size limit to stamp on the run if this stream is re-spilled as an - /// intermediate result: the batch size its merge ran at. It equals the full - /// merge batch size unless the merge consumed a skew-resolved run, in which - /// case it is that run's smaller limit so the re-spilled result stays capped - /// and can't rebuild an oversized batch. - batch_size_limit: usize, - }, - /// Re-spill the spill file at this index smaller, then retry the merge step. - SplitThenRetry(usize), -} - -/// Slice `batch` into two row-halves so a re-spill writes batches half the size. -fn split_batch_in_half(batch: RecordBatch) -> Vec { - let num_rows = batch.num_rows(); - if num_rows <= 1 { - return vec![batch]; - } - let mid = num_rows / 2; - vec![batch.slice(0, mid), batch.slice(mid, num_rows - mid)] -} - fn effective_spill_merge_fan_in(configured_fan_in: usize) -> usize { if configured_fan_in == 0 { usize::MAX } else { - configured_fan_in.max(2) + configured_fan_in.max(MIN_SPILL_MERGE_FAN_IN) } } @@ -1194,122 +859,28 @@ mod tests { Ok(()) } - /// Two sorted runs whose largest batches are too big to both - /// be seated in the merge budget at once are re-spilled (halved) until they - /// fit, and the merge then completes with fully sorted, complete output. - #[tokio::test] - async fn skewed_runs_reuse_retained_workspace_across_retries() -> Result<()> { - // These budgets require one and two re-spills, respectively. - for (budget_halves, expected_splits) in [(7, 1), (5, 2)] { - let capacity = 1024 * 1024; - let parent: Arc = Arc::new(GreedyMemoryPool::new(capacity)); - let env = RuntimeEnvBuilder::new() - .with_memory_pool(Arc::clone(&parent)) - .build_arc()?; - let schema = test_schema(); - let spill_manager = build_spill_manager(&env, &schema); - let spill_count = spill_manager.metrics.spill_file_count.clone(); - let n: i64 = 16384; - let f0 = make_sorted_spill_file(&spill_manager, &schema, (0..n).collect()); - let f1 = make_sorted_spill_file(&spill_manager, &schema, (0..n).collect()); - let m = f0.max_record_batch_memory.max(f1.max_record_batch_memory); - let workspace = m * budget_halves / 2; - assert!(workspace < capacity); - let merge_pool = Arc::new(MergeMemoryPool::new( - Arc::clone(&parent), - MemoryConsumer::new("merge workspace"), - )); - let pool: Arc = - Arc::clone(&merge_pool) as Arc; - let batch_size = 8192; - let mut builder = build_merge_builder( - spill_manager, - Arc::clone(&schema), - vec![f0, f1], - &pool, - batch_size, - ) - .with_merge_pool(Some(Arc::clone(&merge_pool))); - builder.reservation.try_grow(workspace)?; - merge_pool.retain(workspace); - let contender = MemoryConsumer::new("contender").register(&parent); - contender.try_grow(capacity - workspace)?; - - // Drive each retry so its freed child reservation cannot silently - // release the workspace and reacquire it from the parent pool. - for _ in 0..expected_splits { - let MergeStep::SplitThenRetry(index) = - builder.merge_sorted_runs_within_mem_limit(false)? - else { - panic!("the merge must re-spill a skewed run"); - }; - assert_eq!(parent.reserved(), capacity); - assert!(contender.try_grow(1).is_err()); - builder.split_spill_file_in_half(index, false).await?; - assert_eq!(parent.reserved(), capacity); - assert!(contender.try_grow(1).is_err()); - } - - let final_bytes: usize = builder - .sorted_spill_files - .iter() - .map(|(file, _)| { - get_reserved_bytes_for_record_batch_size( - file.max_record_batch_memory, - file.max_record_batch_memory, - ) - }) - .sum(); - assert!(final_bytes < workspace && workspace < 2 * final_bytes); - - let mut stream = builder.create_spillable_merge_stream(); - let first = stream.try_next().await?.expect("nonempty merge"); - assert_eq!(spill_count.value(), 2 + expected_splits); - assert_eq!(merge_pool.reserved(), final_bytes); - assert_eq!(parent.reserved(), contender.size() + final_bytes); - - let mut batches = vec![first]; - while let Some(batch) = stream.try_next().await? { - batches.push(batch); - } - assert_eq!( - batches.iter().map(|batch| batch.num_rows()).max(), - Some(batch_size / 2), - "splitting both runs must halve, not quarter, the output batch size" - ); - let merged = concat_batches(&schema, &batches)?; - let expected = - Int64Array::from_iter_values((0..n).flat_map(|value| [value, value])); - assert_eq!(merged.column(0).as_primitive::(), &expected); - - // The pool stays alive: EOF must release live bytes and the idle floor. - assert_eq!(merge_pool.reserved(), 0); - assert_eq!(parent.reserved(), contender.size()); - drop(stream); - assert_eq!(env.disk_manager.spilling_progress().active_files_count, 0); - assert_eq!(env.disk_manager.used_disk_space(), 0); - drop(contender); - assert_eq!(parent.reserved(), 0); - } - Ok(()) - } - + /// Runs written through a manager with a byte bound merge in a pool that + /// could not hold two of their batches had they been written whole. #[tokio::test] - async fn skewed_runs_are_respilled_so_the_merge_fits() -> Result<()> { + async fn bounded_runs_merge_in_a_pool_smaller_than_two_whole_batches() -> Result<()> { let env = Arc::new(RuntimeEnv::default()); let schema = test_schema(); - let spill_manager = build_spill_manager(&env, &schema); - let n: i64 = 16384; + let whole = build_spill_manager(&env, &schema); + let whole_batch = make_sorted_spill_file(&whole, &schema, (0..n).collect()) + .max_record_batch_memory; + let bound = whole_batch / 4; + let spill_manager = + build_spill_manager(&env, &schema).with_max_batch_bytes(Some(bound)); let f0 = make_sorted_spill_file(&spill_manager, &schema, (0..n).collect()); let f1 = make_sorted_spill_file(&spill_manager, &schema, (0..n).collect()); - let m = f0.max_record_batch_memory.max(f1.max_record_batch_memory); - - // Seating two streams needs ~4*m (2*m each), which does NOT fit, but the - // budget is large enough once a run is halved. The rescue keeps halving - // the blocking run until two streams fit (here, after one halving). - let pool: Arc = Arc::new(GreedyMemoryPool::new(m * 7 / 2)); + assert!(f0.max_record_batch_memory <= bound); + assert!(f1.max_record_batch_memory <= bound); + // Merging the whole batches needs 4 * whole_batch with one batch read + // ahead, more than this pool. + let pool: Arc = + Arc::new(GreedyMemoryPool::new(whole_batch * 7 / 2)); let builder = build_merge_builder( spill_manager, Arc::clone(&schema), @@ -1317,83 +888,23 @@ mod tests { &pool, 8192, ); - let stream = builder.create_spillable_merge_stream(); - let batches: Vec = stream.try_collect().await?; - - let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum(); - assert_eq!( - total_rows, - (2 * n) as usize, - "the merge must emit every input row" - ); - + let batches: Vec = builder + .create_spillable_merge_stream() + .try_collect() + .await?; let merged = concat_batches(&schema, &batches)?; let col = merged.column(0).as_primitive::(); - for i in 1..col.len() { - assert!( - col.value(i - 1) <= col.value(i), - "merge output must be sorted: {} > {} at {i}", - col.value(i - 1), - col.value(i), - ); - } - - Ok(()) - } - - /// A run can fit the pool during splitting but leave no room for replay. - /// Splitting must shrink the merge buffers to half the available memory. - #[rstest::rstest] - #[tokio::test] - async fn replay_headroom_splits_an_oversized_first_run( - #[values(128, 129, 4096)] rows: i64, - ) -> Result<()> { - let env = Arc::new(RuntimeEnv::default()); - let schema = test_schema(); - let spill_manager = build_spill_manager(&env, &schema); - let first = make_sorted_spill_file(&spill_manager, &schema, (0..rows).collect()); - let second = - make_sorted_spill_file(&spill_manager, &schema, (rows..2 * rows).collect()); - let merge_limit = first.max_record_batch_memory; - // One original run fills the pool. Splitting may use that space before - // replay starts, but the final merge must leave half the pool for replay. - let pool_size = 2 * merge_limit; - let pool: Arc = Arc::new(GreedyMemoryPool::new(pool_size)); - let builder = build_merge_builder( - spill_manager, - Arc::clone(&schema), - vec![first, second], - &pool, - 8192, - ) - .with_replay_headroom(true); - let mut stream = builder.create_spillable_merge_stream(); - let mut batches = vec![]; - while let Some(batch) = stream.next().await { - let batch = batch?; - // Short input runs must not rebuild a larger output batch than - // the attached merge reservation can hold. - assert!( - crate::spill::get_record_batch_memory_size(&batch) <= pool.reserved() - ); - batches.push(batch); - assert!(pool.reserved() <= merge_limit); - } - let merged = concat_batches(&schema, &batches)?; - let values = merged.column(0).as_primitive::(); - assert_eq!(values.len(), (2 * rows) as usize); - for (expected, value) in values.values().iter().enumerate() { - assert_eq!(*value, expected as i64); - } - drop(stream); + assert_eq!(col.len(), (2 * n) as usize); + assert!(col.values().windows(2).all(|w| w[0] <= w[1])); + drop(batches); assert_eq!(pool.reserved(), 0); - assert_eq!(env.disk_manager.spilling_progress().current_bytes, 0); - assert_eq!(env.disk_manager.spilling_progress().active_files_count, 0); Ok(()) } + /// With one batch read ahead, the two runs a pass needs only have to fit + /// the pool, not the replay headroom. #[tokio::test] - async fn replay_headroom_allows_only_an_indivisible_minimum() -> Result<()> { + async fn replay_headroom_allows_the_minimum_without_headroom() -> Result<()> { let env = Arc::new(RuntimeEnv::default()); let schema = test_schema(); let spill_manager = build_spill_manager(&env, &schema); @@ -1407,16 +918,9 @@ mod tests { build_merge_builder(spill_manager, Arc::clone(&schema), spills, &pool, 8192) .with_replay_headroom(true); - assert!(!builder.split_spill_file_in_half(0, true).await?); - assert_eq!(builder.sorted_spill_files.len(), 3); - // Actual batches contain one row even though the nominal size is 8192. - assert_eq!(builder.sorted_spill_files[0].1, 1); let mut reservation = builder.reservation.new_empty(); - let SpillFilesToMerge::Ready(spills, buffer_len) = - builder.get_sorted_spill_files_to_merge(2, 2, &mut reservation, true)? - else { - panic!("minimum merge should fit the pool"); - }; + let (spills, buffer_len) = + builder.get_sorted_spill_files_to_merge(2, 2, &mut reservation)?; assert_eq!(buffer_len, 1); assert_eq!(spills.len(), 2); assert_eq!(builder.sorted_spill_files.len(), 1); @@ -1425,11 +929,10 @@ mod tests { reservation.free(); let mut stream = builder.create_spillable_merge_stream(); - let mut values = Vec::new(); + let mut values: Vec = Vec::new(); while let Some(batch) = stream.next().await { let batch = batch?; - assert_eq!(batch.num_rows(), 1); - values.push(batch.column(0).as_primitive::().value(0)); + values.extend(batch.column(0).as_primitive::().values()); } assert_eq!(values, vec![0, 1, 2]); drop(stream); @@ -1439,6 +942,32 @@ mod tests { Ok(()) } + /// With replay headroom, reading two batches ahead must also leave as much + /// again for the replay; otherwise the merge reads one batch ahead. + #[test] + fn read_ahead_needs_headroom() -> Result<()> { + let env = Arc::new(RuntimeEnv::default()); + let schema = test_schema(); + let spill_manager = build_spill_manager(&env, &schema); + let spills = (0..2) + .map(|value| make_sorted_spill_file(&spill_manager, &schema, vec![value])) + .collect::>(); + let batch_memory = spills[0].max_record_batch_memory; + // Two batches read ahead from two runs take 8 batches of the pool's 9. + let pool: Arc = Arc::new(GreedyMemoryPool::new(9 * batch_memory)); + let mut builder = build_merge_builder(spill_manager, schema, spills, &pool, 1) + .with_replay_headroom(true); + let mut reservation = builder.reservation.new_empty(); + let (spills, buffer_len) = + builder.get_sorted_spill_files_to_merge(2, 2, &mut reservation)?; + assert_eq!(buffer_len, 1); + assert_eq!(spills.len(), 2); + assert_eq!(reservation.size(), 4 * batch_memory); + drop((reservation, spills, builder)); + assert_eq!(pool.reserved(), 0); + Ok(()) + } + #[test] fn replay_headroom_is_released_after_rejected_candidate() -> Result<()> { let env = Arc::new(RuntimeEnv::default()); @@ -1453,11 +982,8 @@ mod tests { let mut builder = build_merge_builder(spill_manager, schema, spills, &pool, 1) .with_replay_headroom(true); let mut reservation = builder.reservation.new_empty(); - let SpillFilesToMerge::Ready(spills, buffer_len) = - builder.get_sorted_spill_files_to_merge(1, 2, &mut reservation, false)? - else { - panic!("two streams and replay headroom should fit the pool"); - }; + let (spills, buffer_len) = + builder.get_sorted_spill_files_to_merge(1, 2, &mut reservation)?; assert_eq!(buffer_len, 1); assert_eq!(spills.len(), 2); assert_eq!(builder.sorted_spill_files.len(), 1); @@ -1488,7 +1014,13 @@ mod tests { .with_replay_headroom(true); let mut stream = builder.create_spillable_merge_stream(); let error = stream.next().await.unwrap().unwrap_err(); - assert!(error.to_string().contains("cannot be split further")); + assert!( + matches!( + error.find_root(), + datafusion_common::DataFusionError::ResourcesExhausted(_) + ), + "{error}" + ); drop(stream); assert_eq!(pool.reserved(), 0); assert_eq!(env.disk_manager.spilling_progress().current_bytes, 0); @@ -1496,330 +1028,6 @@ mod tests { Ok(()) } - #[rstest::rstest] - #[case(1, false, DataType::Utf8)] - #[case(8192, true, DataType::Utf8)] - #[case(8192, true, DataType::Utf8View)] - #[tokio::test] - async fn indivisible_spill_merge_needs_no_extra_disk_space( - #[case] batch_size: usize, - #[case] mixed_batches: bool, - #[case] data_type: DataType, - ) -> Result<()> { - use arrow::array::{ArrayRef, StringArray, StringViewArray}; - - const KEY_BYTES: usize = 300_000; - const POOL_BYTES: usize = 2 * 1024 * 1024; - let schema = Arc::new(Schema::new(vec![Field::new("x", data_type, false)])); - let make_runs = |manager: &SpillManager| -> Result> { - (0..2) - .map(|run| { - let batches = (0..8).map(|batch| { - // A small two-row batch before the large singleton - // prevents treating the first batch as representative. - let keys = if mixed_batches && batch == 0 { - vec![format!("{run:04}"), format!("{:04}", run + 2)] - } else { - let key = run + 2 * (batch + usize::from(mixed_batches)); - vec![format!("{key:04}{}", "x".repeat(KEY_BYTES - 4))] - }; - let values: ArrayRef = match schema.field(0).data_type() { - DataType::Utf8 => Arc::new(StringArray::from(keys)), - DataType::Utf8View => Arc::new(StringViewArray::from(keys)), - _ => unreachable!(), - }; - RecordBatch::try_new(Arc::clone(&schema), vec![values]) - .map_err(Into::into) - }); - let (file, max_record_batch_memory) = manager - .spill_record_batch_iter_and_return_max_batch_memory( - batches, - "indivisible input run", - )? - .unwrap(); - Ok(SortedSpillFile { - file, - max_record_batch_memory, - }) - }) - .collect() - }; - - // Calibrate the quota to exactly the original IPC files, with no room - // for even a replacement header. Keep several batches in each run so - // read-ahead cannot retire the original before the unnecessary write. - let calibration = Arc::new(RuntimeEnv::default()); - let runs = make_runs(&build_spill_manager(&calibration, &schema))?; - let quota = runs.iter().map(|run| run.file.size().unwrap()).sum(); - drop(runs); - let env = RuntimeEnvBuilder::new() - .with_max_temp_directory_size(quota) - .build_arc()?; - let manager = build_spill_manager(&env, &schema); - let runs = make_runs(&manager)?; - assert_eq!(env.disk_manager.spilling_progress().current_bytes, quota); - assert!(4 * runs[0].max_record_batch_memory > POOL_BYTES / 2); - assert!(4 * runs[0].max_record_batch_memory <= POOL_BYTES); - let pool: Arc = Arc::new(GreedyMemoryPool::new(POOL_BYTES)); - let builder = build_merge_builder(manager, schema, runs, &pool, batch_size) - .with_replay_headroom(true); - let mut stream = builder.create_spillable_merge_stream(); - let mut expected = 0; - while let Some(batch) = stream.next().await { - let batch = batch?; - assert_eq!(batch.num_rows(), 1); - let values = arrow::compute::cast(batch.column(0), &DataType::Utf8)?; - let key = values.as_string::().value(0); - assert_eq!(key[..4].parse::().unwrap(), expected); - expected += 1; - } - assert_eq!(expected, if mixed_batches { 18 } else { 16 }); - drop(stream); - assert_eq!(pool.reserved(), 0); - let progress = env.disk_manager.spilling_progress(); - assert_eq!(progress.current_bytes, 0); - assert_eq!(progress.active_files_count, 0); - Ok(()) - } - - #[rstest::rstest] - #[tokio::test] - async fn small_view_spill_merge_needs_no_extra_disk_space( - #[values(3, 6)] pool_batches: usize, - #[values(false, true)] mixed_batches: bool, - ) -> Result<()> { - use arrow::array::StringViewBuilder; - - let schema = Arc::new(Schema::new(vec![Field::new( - "x", - DataType::Utf8View, - false, - )])); - let make_runs = |manager: &SpillManager| -> Result> { - (0..2) - .map(|run| { - let batches = (0..8).map(|batch| { - // A non-inline value retains the builder's 8 KiB block, - // but IPC stores only its used bytes. Small view buffers - // are below the spill writer's compaction threshold. - let key = format!("{:04}xxxxxxxxx", run + 2 * batch); - let mut values = StringViewBuilder::with_capacity(1) - .with_fixed_block_size(8192); - values.append_value(key); - // A larger final batch requires scanning the entire run - // before replacing its maximum decoded size. - if mixed_batches && batch == 7 { - values.append_value(format!( - "{:04}xxxxxxxxx", - run + 2 * (batch + 1) - )); - } - RecordBatch::try_new( - Arc::clone(&schema), - vec![Arc::new(values.finish())], - ) - .map_err(Into::into) - }); - let (file, max_record_batch_memory) = manager - .spill_record_batch_iter_and_return_max_batch_memory( - batches, - "small view singleton run", - )? - .unwrap(); - Ok(SortedSpillFile { - file, - max_record_batch_memory, - }) - }) - .collect() - }; - - // Fill the quota with the original runs. Multiple batches prevent - // read-ahead from retiring a run before an unnecessary replacement write. - let calibration = Arc::new(RuntimeEnv::default()); - let runs = make_runs(&build_spill_manager(&calibration, &schema))?; - let quota = runs.iter().map(|run| run.file.size().unwrap()).sum(); - drop(runs); - let env = RuntimeEnvBuilder::new() - .with_max_temp_directory_size(quota) - .build_arc()?; - let manager = build_spill_manager(&env, &schema); - let runs = make_runs(&manager)?; - assert_eq!(env.disk_manager.spilling_progress().current_bytes, quota); - let stored_max = runs[0].max_record_batch_memory; - let mut source = manager.read_spill_as_stream_unbuffered( - Arc::clone(&runs[0].file), - Some(stored_max), - )?; - let batch = source.next().await.unwrap()?; - assert_eq!(batch.num_rows(), 1); - assert!(gc_view_arrays(&batch)?.get_sliced_size()? < stored_max); - drop((source, batch)); - - // The stale estimate needs four batches for the minimum merge. Test - // limits below and above that estimate: both can fit the decoded rows. - let pool_bytes = pool_batches * stored_max; - let pool: Arc = Arc::new(GreedyMemoryPool::new(pool_bytes)); - let builder = build_merge_builder(manager, schema, runs, &pool, 8192) - .with_replay_headroom(true); - let mut stream = builder.create_spillable_merge_stream(); - let mut expected = 0; - while let Some(batch) = stream.next().await { - let batch = batch?; - if mixed_batches { - assert!(batch.num_rows() <= 2); - } else { - assert_eq!(batch.num_rows(), 1); - } - for key in batch.column(0).as_string_view().iter() { - assert_eq!(key.unwrap(), format!("{expected:04}xxxxxxxxx")); - expected += 1; - } - } - assert_eq!(expected, if mixed_batches { 18 } else { 16 }); - drop(stream); - assert_eq!(pool.reserved(), 0); - let progress = env.disk_manager.spilling_progress(); - assert_eq!(progress.current_bytes, 0); - assert_eq!(progress.active_files_count, 0); - Ok(()) - } - - /// Tests the `new_max >= old_max` guard: a single-row run cannot be split - /// any smaller, so re-spilling it does not shrink the largest batch and the - /// rescue surfaces `ResourcesExhausted` rather than looping forever. - #[tokio::test] - async fn respilling_an_unsplittable_run_surfaces_resources_exhausted() -> Result<()> { - let env = Arc::new(RuntimeEnv::default()); - let schema = test_schema(); - let spill_manager = build_spill_manager(&env, &schema); - - // A one-row run: `split_batch_in_half` returns it unchanged, so the - // re-spilled file's largest batch cannot drop below the original. - let f0 = make_sorted_spill_file(&spill_manager, &schema, vec![42]); - - // Ample budget so the only possible failure is the un-splittable guard, - // not the single-stream reservation itself. - let pool: Arc = Arc::new(GreedyMemoryPool::new(1024 * 1024)); - let mut builder = - build_merge_builder(spill_manager, schema, vec![f0], &pool, 1024); - - let err = builder - .split_spill_file_in_half(0, false) - .await - .expect_err("re-spilling a one-row run cannot shrink it"); - assert!( - err.to_string().contains("cannot be split further"), - "expected the un-splittable guard error, got: {err}" - ); - - Ok(()) - } - - /// Proves the re-spill also halves the merge output batch size: after one - /// re-spill the merged run is emitted in 4096-row batches (not the original - /// 8192), so it cannot rebuild a full-size batch and reintroduce the skew. - #[tokio::test] - async fn respill_halves_the_merge_output_batch_size() -> Result<()> { - let env = Arc::new(RuntimeEnv::default()); - let schema = test_schema(); - let spill_manager = build_spill_manager(&env, &schema); - - let n: i64 = 16384; - let f0 = make_sorted_spill_file(&spill_manager, &schema, (0..n).collect()); - let f1 = make_sorted_spill_file(&spill_manager, &schema, (0..n).collect()); - let m = f0.max_record_batch_memory.max(f1.max_record_batch_memory); - - // 3.5*m forces exactly one re-spill (split one run, then both fit), which - // halves the merge output batch size. - let initial_batch_size = 8192; - let pool: Arc = Arc::new(GreedyMemoryPool::new(m * 7 / 2)); - - let builder = build_merge_builder( - spill_manager, - Arc::clone(&schema), - vec![f0, f1], - &pool, - initial_batch_size, - ); - let stream = builder.create_spillable_merge_stream(); - let batches: Vec = stream.try_collect().await?; - - // All rows are still present. - let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum(); - assert_eq!(total_rows, (2 * n) as usize); - - // The largest emitted batch is the halved size, not the original 8192: the - // shrunk run carries a halved batch-size limit, and the final pass consumes - // it, so the merge output is capped there. Without the per-run limit the merge - // would rebuild 8192-row batches. - let expected_batch_size = initial_batch_size / 2; - let max_batch_rows = batches.iter().map(|b| b.num_rows()).max().unwrap_or(0); - assert_eq!( - max_batch_rows, expected_batch_size, - "after one re-spill the merge must emit {expected_batch_size}-row \ - batches, got a largest batch of {max_batch_rows} rows" - ); - - Ok(()) - } - - /// Same as [`respill_halves_the_merge_output_batch_size`], but under a budget tight - /// enough that *both* runs must be re-spilled before the merge fits - the scenario - /// where the batch-size reduction could compound. Because the reduction is tracked - /// per-run (each run capped at half) rather than by halving the global batch size on - /// every split, the merged output is emitted in 4096-row batches - half, not a - /// quarter. A global-halving implementation would have halved once per re-spill and - /// emitted 2048-row batches. - #[tokio::test] - async fn respilling_two_skewed_runs_halves_the_output_without_compounding() - -> Result<()> { - let env = Arc::new(RuntimeEnv::default()); - let schema = test_schema(); - let spill_manager = build_spill_manager(&env, &schema); - - let n: i64 = 16384; - let f0 = make_sorted_spill_file(&spill_manager, &schema, (0..n).collect()); - let f1 = make_sorted_spill_file(&spill_manager, &schema, (0..n).collect()); - let m = f0.max_record_batch_memory.max(f1.max_record_batch_memory); - - // 2.5*m is tight enough that even after halving one run the two still don't - // fit, so *both* runs are re-spilled once before the merge succeeds. (3.5*m, - // as in the single-split test, would let the pair fit after one split.) This - // is exactly the scenario where a compounding, global-halving implementation - // would drive the output batch size down to a quarter. - let initial_batch_size = 8192; - let pool: Arc = Arc::new(GreedyMemoryPool::new(m * 5 / 2)); - - let builder = build_merge_builder( - spill_manager, - Arc::clone(&schema), - vec![f0, f1], - &pool, - initial_batch_size, - ); - let stream = builder.create_spillable_merge_stream(); - let batches: Vec = stream.try_collect().await?; - - // All rows are still present. - let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum(); - assert_eq!(total_rows, (2 * n) as usize); - - // Each run was re-spilled once, so each is capped at half the original batch - // size and the merge caps its output at that half - NOT a quarter. A global - // halving-per-split implementation would have emitted 2048-row batches here. - let expected_batch_size = initial_batch_size / 2; - let max_batch_rows = batches.iter().map(|b| b.num_rows()).max().unwrap_or(0); - assert_eq!( - max_batch_rows, expected_batch_size, - "two re-spills must halve (not quarter) the output: expected \ - {expected_batch_size}-row batches, got a largest batch of \ - {max_batch_rows} rows" - ); - - Ok(()) - } - #[test] fn spill_merge_fan_in_is_unlimited_by_default() { assert_eq!(effective_spill_merge_fan_in(0), usize::MAX); @@ -1874,17 +1082,8 @@ mod tests { let mut merge_reservation = MemoryConsumer::new("spill_merge_fan_in_phase") .register(&runtime.memory_pool); - let (spills, buffer_len) = match builder.get_sorted_spill_files_to_merge( - 1, - 2, - &mut merge_reservation, - false, - )? { - SpillFilesToMerge::Ready(spills, buffer_len) => (spills, buffer_len), - SpillFilesToMerge::SplitThenRetry(index) => { - panic!("expected ready spill files, got retry for index {index}") - } - }; + let (spills, buffer_len) = + builder.get_sorted_spill_files_to_merge(1, 2, &mut merge_reservation)?; assert_eq!(spills.len(), 2); assert_eq!(buffer_len, 1); diff --git a/datafusion/physical-plan/src/sorts/multi_level_merge/intermediate_sizing_tests.rs b/datafusion/physical-plan/src/sorts/multi_level_merge/intermediate_sizing_tests.rs index 35b2f6486a6a5..03c726b8a2974 100644 --- a/datafusion/physical-plan/src/sorts/multi_level_merge/intermediate_sizing_tests.rs +++ b/datafusion/physical-plan/src/sorts/multi_level_merge/intermediate_sizing_tests.rs @@ -213,11 +213,7 @@ async fn sized_merge_retains_reservation_until_builder_drop() -> Result<()> { build_merge_builder(spill_manager, Arc::clone(&schema), files, &pool, ROWS) .with_replay_headroom(true) .with_intermediate_merge_sizing(Some(ROWS)); - let MergeStep::Stream { stream, .. } = - builder.merge_sorted_runs_within_mem_limit(false)? - else { - panic!("six inputs should be selected for the intermediate merge"); - }; + let stream = builder.merge_sorted_runs_within_mem_limit()?; assert_eq!(builder.sorted_spill_files.len(), 2); let retained = 4 * 7 * batch_memory; assert_eq!(builder.reservation.size(), retained); diff --git a/datafusion/physical-plan/src/sorts/sort.rs b/datafusion/physical-plan/src/sorts/sort.rs index 96e673e4043bc..3a73f39e000a3 100644 --- a/datafusion/physical-plan/src/sorts/sort.rs +++ b/datafusion/physical-plan/src/sorts/sort.rs @@ -30,6 +30,7 @@ use crate::metrics::{ }; use crate::projection::{ProjectionExec, make_with_child, update_ordering}; use crate::sorts::IncrementalSortIterator; +use crate::sorts::multi_level_merge::MIN_SPILL_MERGE_FAN_IN; use crate::sorts::streaming_merge::{SortedSpillFile, StreamingMergeBuilder}; use crate::spill::get_record_batch_memory_size; use crate::spill::in_progress_spill_file::InProgressSpillFile; @@ -347,7 +348,8 @@ impl ExternalSorter { metrics.spill_metrics.clone(), Arc::clone(&schema), ) - .with_compression_type(spill_compression); + .with_compression_type(spill_compression) + .with_max_batch_bytes(Some(spill_max_batch_bytes(sort_spill_reservation_bytes))); let in_mem_batches_counter = (!schema .fields() @@ -1149,6 +1151,27 @@ impl ExternalSorter { } } +/// The bound on the in-memory size of each batch of a sorted spill run, such that the +/// smallest merge of the runs fits in `merge_memory`. +/// +/// The merge reads at least [`MIN_SPILL_MERGE_FAN_IN`] files, and reserves +/// [`get_reserved_bytes_for_record_batch_size`] of the largest batch for each, with +/// read-ahead reduced to one batch under pressure. The bound is the largest batch size for +/// which that minimum merge fits in `merge_memory`, found from the same formula. +/// [`ExternalSorter`] passes `sort_spill_reservation_bytes`, the headroom its merge starts +/// with. An aggregate passes part of the memory it held when it first spilled. +/// +/// The bound is at least [`MIN_SPILL_MAX_BATCH_BYTES`], so that a small or zero headroom +/// does not write, read and merge tiny batches. Below that headroom the merge cannot fit +/// in it anyway, and takes the rest from the memory pool, as it does without a bound. +pub(crate) fn spill_max_batch_bytes(merge_memory: usize) -> usize { + // Each batch reserves `get_reserved_bytes_for_record_batch_size(b, b)`, which is `2 * b`. + (merge_memory / (2 * MIN_SPILL_MERGE_FAN_IN)).max(MIN_SPILL_MAX_BATCH_BYTES) +} + +/// The smallest bound [`spill_max_batch_bytes`] returns. +const MIN_SPILL_MAX_BATCH_BYTES: usize = 64 * 1024; + /// Estimate how much memory is needed to sort a `RecordBatch`. /// /// This is used to pre-reserve memory for the sort/merge. The sort/merge process involves @@ -2934,6 +2957,31 @@ mod tests { } } + #[test] + fn test_spill_max_batch_bytes_fits_minimum_merge() { + let minimum_merge = |bound| { + MIN_SPILL_MERGE_FAN_IN + * get_reserved_bytes_for_record_batch_size(bound, bound) + }; + for headroom in [1024 * 1024, 10 * 1024 * 1024, 10 * 1024 * 1024 + 7] { + let bound = spill_max_batch_bytes(headroom); + // The largest bound whose minimum merge fits in the headroom. + assert!(minimum_merge(bound) <= headroom, "headroom {headroom}"); + assert!(minimum_merge(bound + 1) > headroom, "headroom {headroom}"); + } + // A small or zero headroom gets the floor. + assert_eq!(spill_max_batch_bytes(0), MIN_SPILL_MAX_BATCH_BYTES); + assert_eq!(spill_max_batch_bytes(1024), MIN_SPILL_MAX_BATCH_BYTES); + // The largest headroom does not overflow. + let bound = spill_max_batch_bytes(usize::MAX); + assert_eq!(bound, usize::MAX / (2 * MIN_SPILL_MERGE_FAN_IN)); + assert!( + MIN_SPILL_MERGE_FAN_IN + .checked_mul(get_reserved_bytes_for_record_batch_size(bound, bound)) + .is_some() + ); + } + #[tokio::test] async fn test_in_mem_sort() -> Result<()> { let task_ctx = Arc::new(TaskContext::default()); diff --git a/datafusion/physical-plan/src/sorts/sort/spill_tests.rs b/datafusion/physical-plan/src/sorts/sort/spill_tests.rs index 625f96be682cb..b89db6c5bb620 100644 --- a/datafusion/physical-plan/src/sorts/sort/spill_tests.rs +++ b/datafusion/physical-plan/src/sorts/sort/spill_tests.rs @@ -7,7 +7,7 @@ use super::{ }; use crate::metrics::ExecutionPlanMetricsSet; use crate::spill::get_record_batch_memory_size; -use crate::spill::spill_manager::GetSlicedSize; +use crate::spill::spill_pieces::recorded_size; use arrow::array::{ ArrayRef, Decimal128Array, DictionaryArray, Int8Array, Int64Array, StringArray, StringViewArray, @@ -1017,7 +1017,7 @@ async fn check_final_spilled_merge_releases_unused_workspace( .iter() .all(|spill| spill.max_record_batch_memory == large_spill_bytes) ); - let tail_spill_bytes = sort_batch(&tail, &ordering, None)?.get_sliced_size()?; + let tail_spill_bytes = recorded_size(&sort_batch(&tail, &ordering, None)?, false)?; let single_buffer_bytes = if intermediate { 4 * large_spill_bytes // Two intermediate files remain for the final merge. } else { diff --git a/datafusion/physical-plan/src/spill/in_progress_spill_file.rs b/datafusion/physical-plan/src/spill/in_progress_spill_file.rs index f91fba0e54641..685718703b2fc 100644 --- a/datafusion/physical-plan/src/spill/in_progress_spill_file.rs +++ b/datafusion/physical-plan/src/spill/in_progress_spill_file.rs @@ -4,6 +4,7 @@ //! Define the `InProgressSpillFile` struct, which represents an in-progress spill file used for writing `RecordBatch`es to disk, created by `SpillManager`. +use datafusion_common::config::SpillCompression; use datafusion_common::{Result, internal_datafusion_err}; use std::sync::Arc; @@ -12,8 +13,9 @@ use datafusion_common::exec_datafusion_err; use datafusion_execution::spill_file::SpillFile; use super::{ - AsyncIPCStreamWriter, IPCStreamEncoder, IPCStreamWriter, gc_view_arrays, - spill_manager::{GetSlicedSize, SpillManager}, + AsyncIPCStreamWriter, IPCStreamEncoder, IPCStreamWriter, + spill_manager::SpillManager, + spill_pieces::{SpillPiece, SpillPieces}, }; enum InProgressWriter { @@ -21,6 +23,15 @@ enum InProgressWriter { Async(AsyncIPCStreamWriter), } +/// What one append wrote, for the reader to size its reservation. +#[derive(Debug, Default, Clone, Copy, PartialEq, Eq)] +pub(crate) struct WrittenBatchStats { + /// The largest recorded size of a piece. + pub(crate) memory_size: usize, + /// The most rows in a piece. + pub(crate) rows: usize, +} + /// Represents an in-progress spill file used for writing `RecordBatch`es to disk, created by `SpillManager`. /// Caller is able to use this struct to incrementally append in-memory batches to /// the file, and then finalize the file by calling the `finish` method. @@ -50,20 +61,37 @@ impl InProgressSpillFile { /// buffers. When a view array is sliced, it still references the original full buffers, /// causing massive spill files without GC (see issue #19414: 820MB → 33MB after GC). /// - /// Returns the post-GC sliced memory size of the batch for memory accounting. + /// Returns the post-GC sliced memory size of the batch for memory accounting. When the + /// batch size is bounded and the batch is split (see [`SpillManager::with_max_batch_bytes`]), + /// this is the size of the largest piece, + /// which is the largest batch a reader of the file decodes. /// /// # Errors /// - Returns an error if the file is not active (has been finalized) /// - Returns an error if appending would exceed the disk usage limit configured /// by `max_temp_directory_size` in `DiskManager` pub fn append_batch(&mut self, batch: &RecordBatch) -> Result { + Ok(self.append_batch_with_stats(batch)?.memory_size) + } + + /// Like [`Self::append_batch`], but returns both the memory size and the row count + /// of the largest batch written. When the batch is split, both describe the pieces a + /// reader decodes, not the appended batch. + pub(crate) fn append_batch_with_stats( + &mut self, + batch: &RecordBatch, + ) -> Result { if self.in_progress_file.is_none() { return Err(exec_datafusion_err!( "Append operation failed: No active in-progress file. The file may have already been finalized." )); } - let gc_batch = gc_view_arrays(batch)?; + let pieces = SpillPieces::new( + batch, + self.spill_manager.max_batch_bytes(), + self.spill_manager.compression != SpillCompression::Uncompressed, + )?; if self.writer.is_none() { // Use the SpillManager's declared schema rather than the batch's schema. @@ -78,36 +106,56 @@ impl InProgressSpillFile { spill_writer, schema.as_ref(), self.spill_manager.compression, - gc_batch.num_rows(), + batch.num_rows(), )?)); // Update metrics self.spill_manager.metrics.spill_file_count.add(1); } } + let mut stats = WrittenBatchStats::default(); if let Some(InProgressWriter::Sync(writer)) = &mut self.writer { - // The writer calculates how many serialized bytes were emitted - let (spilled_rows, delta_bytes) = writer.write(&gc_batch)?; + for piece in pieces { + let SpillPiece { batch: piece, size } = piece?; + // The writer calculates how many serialized bytes were emitted + let (spilled_rows, delta_bytes) = writer.write(&piece)?; - self.spill_manager.metrics.spilled_rows.add(spilled_rows); - self.spill_manager.metrics.spilled_bytes.add(delta_bytes); + self.spill_manager.metrics.spilled_rows.add(spilled_rows); + self.spill_manager.metrics.spilled_bytes.add(delta_bytes); + stats.memory_size = stats.memory_size.max(size); + stats.rows = stats.rows.max(piece.num_rows()); + } } else if self.writer.is_some() { return Err(exec_datafusion_err!( "Cannot use synchronous append after asynchronous spill writing has started" )); } - gc_batch.get_sliced_size() + Ok(stats) } /// Appends a `RecordBatch` using the asynchronous spill writer. pub async fn append_batch_async(&mut self, batch: &RecordBatch) -> Result { + Ok(self.append_batch_async_with_stats(batch).await?.memory_size) + } + + /// Like [`Self::append_batch_async`], but returns both the memory size and the row + /// count of the largest batch written. When the batch is split, both describe the + /// pieces a reader decodes, not the appended batch. + pub(crate) async fn append_batch_async_with_stats( + &mut self, + batch: &RecordBatch, + ) -> Result { if self.in_progress_file.is_none() { return Err(exec_datafusion_err!( "Append operation failed: No active in-progress file. The file may have already been finalized." )); } - let gc_batch = gc_view_arrays(batch)?; + let pieces = SpillPieces::new( + batch, + self.spill_manager.max_batch_bytes(), + self.spill_manager.compression != SpillCompression::Uncompressed, + )?; if self.writer.is_none() { let schema = self.spill_manager.schema(); @@ -117,7 +165,7 @@ impl InProgressSpillFile { let encoder = IPCStreamEncoder::new( schema.as_ref(), self.spill_manager.compression, - gc_batch.num_rows(), + batch.num_rows(), )?; let spill_writer = in_progress_file.open_async_writer().await?; @@ -130,12 +178,18 @@ impl InProgressSpillFile { } } + let mut stats = WrittenBatchStats::default(); match &mut self.writer { Some(InProgressWriter::Async(writer)) => { - let (spilled_rows, delta_bytes) = writer.write(&gc_batch).await?; + for piece in pieces { + let SpillPiece { batch: piece, size } = piece?; + let (spilled_rows, delta_bytes) = writer.write(&piece).await?; - self.spill_manager.metrics.spilled_rows.add(spilled_rows); - self.spill_manager.metrics.spilled_bytes.add(delta_bytes); + self.spill_manager.metrics.spilled_rows.add(spilled_rows); + self.spill_manager.metrics.spilled_bytes.add(delta_bytes); + stats.memory_size = stats.memory_size.max(size); + stats.rows = stats.rows.max(piece.num_rows()); + } } Some(InProgressWriter::Sync(_)) => { return Err(exec_datafusion_err!( @@ -148,7 +202,7 @@ impl InProgressSpillFile { )); } } - gc_batch.get_sliced_size() + Ok(stats) } pub fn flush(&mut self) -> Result<()> { @@ -237,8 +291,14 @@ impl InProgressSpillFile { #[cfg(test)] mod tests { use super::*; - use arrow::array::Int64Array; - use arrow_schema::{DataType, Field, Schema}; + use crate::spill::SPILL_BATCH_MEMORY_MARGIN; + use arrow::array::{ + Array, ArrayRef, AsArray, BinaryViewArray, DictionaryArray, Int32Array, + Int64Array, ListBuilder, MapBuilder, StringArray, StringBuilder, StringViewArray, + StructArray, UInt8Array, UInt32Array, + }; + use arrow_schema::{DataType, Field, Fields, Schema, SchemaRef}; + use datafusion_common::utils::memory::get_record_batch_memory_size; use datafusion_execution::runtime_env::RuntimeEnvBuilder; use datafusion_physical_expr_common::metrics::{ ExecutionPlanMetricsSet, SpillMetrics, @@ -306,4 +366,735 @@ mod tests { Ok(()) } + + const WIDE: usize = 64 * 1024; + + fn wide_value(i: usize) -> String { + format!("{i}{}", "x".repeat(WIDE)) + } + + fn manager( + schema: SchemaRef, + max_batch_bytes: Option, + ) -> Result> { + let runtime = Arc::new(RuntimeEnvBuilder::new().build()?); + let metrics = SpillMetrics::new(&ExecutionPlanMetricsSet::new(), 0); + Ok(Arc::new( + SpillManager::new(runtime, metrics, schema) + .with_max_batch_bytes(max_batch_bytes), + )) + } + + /// Writes `batches` through a bounded spill file, reads the file back with the reader's + /// size check on, and returns what was read and the recorded largest batch. + async fn round_trip( + batches: &[RecordBatch], + max_batch_bytes: usize, + ) -> Result<(Vec, usize)> { + let manager = manager(batches[0].schema(), Some(max_batch_bytes))?; + let mut file = manager.create_in_progress_file("test")?; + let mut max_memory = 0; + for batch in batches { + max_memory = max_memory.max(file.append_batch(batch)?); + } + let file = file.finish()?.unwrap(); + let read = manager + .read_spill_as_stream(file, Some(max_memory))? + .try_collect::>() + .await?; + // The reader only logs a batch larger than the recorded size plus its margin, so + // check it here: the merge reserves memory from the recorded size. + for batch in &read { + let size = get_record_batch_memory_size(batch); + assert!( + size <= max_memory + SPILL_BATCH_MEMORY_MARGIN, + "read back a batch of {size} bytes, recorded {max_memory}" + ); + } + Ok((read, max_memory)) + } + + fn assert_same_rows(read: &[RecordBatch], written: &[RecordBatch]) { + let schema = written[0].schema(); + let read = arrow::compute::concat_batches(&schema, read).unwrap(); + let written = arrow::compute::concat_batches(&schema, written).unwrap(); + assert_eq!(read, written); + } + + /// Slices of one large view array keep the whole parent's data buffers alive. The split + /// must look at what each slice points at, or every half reports the whole parent and the + /// batch is written unsplit. + #[tokio::test] + async fn view_slices_of_a_large_parent_are_split() -> Result<()> { + let parent = StringViewArray::from_iter_values((0..256).map(wide_value)); + let schema = Arc::new(Schema::new(vec![Field::new( + "v", + DataType::Utf8View, + false, + )])); + // Two slices that each share the parent's buffers (about 16 MiB). + let batches: Vec = [(0, 128), (128, 128)] + .into_iter() + .map(|(offset, len)| { + RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(parent.slice(offset, len))], + ) + }) + .collect::>()?; + + let budget = 1024 * 1024; + let (read, max_memory) = round_trip(&batches, budget).await?; + assert!( + max_memory <= budget, + "largest spilled batch is {max_memory} bytes" + ); + assert!( + read.len() >= 16, + "expected at least 16 pieces, got {}", + read.len() + ); + assert_same_rows(&read, &batches); + Ok(()) + } + + /// Many rows share each long value, so the pieces reference overlapping parts of the + /// parent's data buffers. The bound and the row order must still hold. + #[tokio::test] + async fn repeated_view_values_are_split_and_kept_in_order() -> Result<()> { + let values = (0..256).map(|i| wide_value(i % 4)); + let batch = RecordBatch::try_new( + Arc::new(Schema::new(vec![Field::new( + "v", + DataType::Utf8View, + false, + )])), + vec![Arc::new(StringViewArray::from_iter_values(values))], + )?; + let budget = 256 * 1024; + let (read, max_memory) = round_trip(std::slice::from_ref(&batch), budget).await?; + assert!( + max_memory <= budget, + "largest spilled batch is {max_memory} bytes" + ); + assert!(read.len() > 1); + assert_same_rows(&read, &[batch]); + Ok(()) + } + + /// Each key points at its own wide value, so the dictionary values are most of the batch. + /// Split pieces drop the values they do not use, so the dictionary is divided between + /// them, and each piece passes the reader's size check. + #[tokio::test] + async fn dictionary_with_large_values_is_divided() -> Result<()> { + let values = StringArray::from_iter_values((0..128).map(wide_value)); + let keys = Int32Array::from_iter_values(0..128); + let dictionary = DictionaryArray::try_new(keys, Arc::new(values))?; + let batch = RecordBatch::try_new( + Arc::new(Schema::new(vec![Field::new( + "d", + DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)), + false, + )])), + vec![Arc::new(dictionary)], + )?; + let budget = 1024 * 1024; + let (read, max_memory) = round_trip(std::slice::from_ref(&batch), budget).await?; + assert!( + max_memory <= budget, + "largest spilled batch is {max_memory} bytes" + ); + assert!( + read.len() >= 8, + "expected at least 8 pieces, got {}", + read.len() + ); + let cast = + |b: &RecordBatch| arrow::compute::cast(b.column(0), &DataType::Utf8).unwrap(); + let read_values: Vec = read.iter().map(cast).collect(); + let read_values = arrow::compute::concat( + &read_values.iter().map(|a| a.as_ref()).collect::>(), + )?; + assert_eq!(&read_values, &cast(&batch)); + Ok(()) + } + + /// One row larger than the budget cannot be split. It is written whole and recorded, so + /// the reader's check still passes. + #[tokio::test] + async fn single_wide_row_is_written_whole() -> Result<()> { + let batch = RecordBatch::try_new( + Arc::new(Schema::new(vec![Field::new("v", DataType::Utf8, false)])), + vec![Arc::new(StringArray::from_iter_values([wide_value(0)]))], + )?; + let (read, max_memory) = round_trip(std::slice::from_ref(&batch), 1024).await?; + assert!(max_memory > 1024); + assert_eq!(read, vec![batch]); + Ok(()) + } + + /// Without a bound, one appended batch is one written batch, as before. + #[tokio::test] + async fn unbounded_manager_keeps_batches_whole() -> Result<()> { + let batch = RecordBatch::try_new( + Arc::new(Schema::new(vec![Field::new( + "v", + DataType::Utf8View, + false, + )])), + vec![Arc::new(StringViewArray::from_iter_values( + (0..64).map(wide_value), + ))], + )?; + let manager = manager(batch.schema(), None)?; + let file = manager + .spill_record_batch_and_finish(std::slice::from_ref(&batch), "test")? + .unwrap(); + let read = manager + .read_spill_as_stream(file, None)? + .try_collect::>() + .await?; + assert_eq!(read.len(), 1); + assert_same_rows(&read, &[batch]); + Ok(()) + } + + /// A `List` column: `rows` rows of four quarter-`WIDE` strings each. + fn wide_list(rows: usize) -> ArrayRef { + let mut builder = ListBuilder::new(StringBuilder::new()); + for i in 0..rows { + for _ in 0..4 { + builder + .values() + .append_value(format!("{i}{}", "x".repeat(WIDE / 4))); + } + builder.append(true); + } + Arc::new(builder.finish()) + } + + /// A sliced list keeps its whole child array. The split must measure and copy only the + /// child range that each piece uses, or both halves report the parent and the batch is + /// written unsplit. + #[tokio::test] + async fn list_column_is_split() -> Result<()> { + let list = wide_list(256); + let schema = Arc::new(Schema::new(vec![Field::new( + "l", + list.data_type().clone(), + false, + )])); + // A slice of a larger parent, as a sort or aggregate emits. + let batch = RecordBatch::try_new(schema, vec![list.slice(64, 128)])?; + let budget = 1024 * 1024; + let (read, max_memory) = round_trip(std::slice::from_ref(&batch), budget).await?; + assert!( + max_memory <= budget, + "largest spilled batch is {max_memory} bytes" + ); + assert!( + read.len() >= 8, + "expected at least 8 pieces, got {}", + read.len() + ); + assert_same_rows(&read, &[batch]); + Ok(()) + } + + /// A map keeps its whole entries array when it is sliced, like a list. + #[tokio::test] + async fn map_column_is_split() -> Result<()> { + let mut builder = + MapBuilder::new(None, StringBuilder::new(), StringBuilder::new()); + for i in 0..256 { + builder.keys().append_value(format!("k{i}")); + builder.values().append_value(wide_value(i)); + builder.append(true)?; + } + let map: ArrayRef = Arc::new(builder.finish()); + let schema = Arc::new(Schema::new(vec![Field::new( + "m", + map.data_type().clone(), + false, + )])); + let batch = RecordBatch::try_new(schema, vec![map.slice(64, 128)])?; + let budget = 1024 * 1024; + let (read, max_memory) = round_trip(std::slice::from_ref(&batch), budget).await?; + assert!( + max_memory <= budget, + "largest spilled batch is {max_memory} bytes" + ); + assert!( + read.len() >= 8, + "expected at least 8 pieces, got {}", + read.len() + ); + assert_same_rows(&read, &[batch]); + Ok(()) + } + + /// A struct narrows its children when it is sliced, but a list inside it still keeps its + /// whole child array. + #[tokio::test] + async fn struct_with_list_child_is_split() -> Result<()> { + let list = wide_list(256); + let ids: ArrayRef = Arc::new(Int64Array::from_iter_values(0..256)); + let fields = Fields::from(vec![ + Field::new("id", DataType::Int64, false), + Field::new("l", list.data_type().clone(), false), + ]); + let structs: ArrayRef = + Arc::new(StructArray::try_new(fields.clone(), vec![ids, list], None)?); + let schema = Arc::new(Schema::new(vec![Field::new( + "s", + DataType::Struct(fields), + false, + )])); + let batch = RecordBatch::try_new(schema, vec![structs.slice(64, 128)])?; + let budget = 1024 * 1024; + let (read, max_memory) = round_trip(std::slice::from_ref(&batch), budget).await?; + assert!( + max_memory <= budget, + "largest spilled batch is {max_memory} bytes" + ); + assert!( + read.len() >= 8, + "expected at least 8 pieces, got {}", + read.len() + ); + assert_same_rows(&read, &[batch]); + Ok(()) + } + + fn dictionary_type(values: DataType) -> DataType { + DataType::Dictionary(Box::new(DataType::Int32), Box::new(values)) + } + + /// A small batch whose keys use a few values of a large dictionary, as `take` produces + /// when a sort emits a chunk. It is within the budget, so it is written as it is: the + /// batches of a spill file keep sharing one dictionary, which the reader decodes once. + #[tokio::test] + async fn unsplit_batch_keeps_its_dictionary() -> Result<()> { + let values = StringArray::from_iter_values((0..4).map(|i| format!("v{i}"))); + let keys = Int32Array::from_iter_values([3, 1, 1, 0]); + let dictionary = DictionaryArray::try_new(keys, Arc::new(values))?; + let batch = RecordBatch::try_new( + Arc::new(Schema::new(vec![Field::new( + "d", + dictionary_type(DataType::Utf8), + false, + )])), + vec![Arc::new(dictionary)], + )?; + let pieces = SpillPieces::new(&batch, Some(1024 * 1024), false)? + .collect::>>()?; + assert_eq!(pieces.len(), 1); + assert_eq!(pieces[0].batch, batch); + Ok(()) + } + + /// A dictionary inside a struct is compacted like a top-level one, so each piece + /// holds only the values of its rows. + #[tokio::test] + async fn nested_dictionary_is_split() -> Result<()> { + let values = StringArray::from_iter_values((0..128).map(wide_value)); + let keys = Int32Array::from_iter_values(0..128); + let dictionary: ArrayRef = + Arc::new(DictionaryArray::try_new(keys, Arc::new(values))?); + let fields = Fields::from(vec![Field::new( + "d", + dictionary_type(DataType::Utf8), + false, + )]); + let structs: ArrayRef = Arc::new(StructArray::try_new( + fields.clone(), + vec![dictionary], + None, + )?); + let batch = RecordBatch::try_new( + Arc::new(Schema::new(vec![Field::new( + "s", + DataType::Struct(fields), + false, + )])), + vec![structs], + )?; + let budget = 16 * WIDE; + let (read, max_memory) = round_trip(std::slice::from_ref(&batch), budget).await?; + assert!(read.len() > 1); + assert!(max_memory <= budget, "recorded {max_memory}"); + assert_same_rows(&read, &[batch]); + Ok(()) + } + + /// Rows of a list view can share child ranges, so copying each piece's ranges out can + /// multiply the data. A batch with a list view is written whole. + #[tokio::test] + async fn list_view_with_shared_child_is_written_whole() -> Result<()> { + let child: ArrayRef = Arc::new(StringArray::from_iter_values([wide_value(0)])); + let field = Arc::new(Field::new_list_field(DataType::Utf8, false)); + let rows = 1024; + let list_view: ArrayRef = Arc::new(arrow::array::ListViewArray::try_new( + Arc::clone(&field), + vec![0; rows].into(), + vec![1; rows].into(), + child, + None, + )?); + let schema = Arc::new(Schema::new(vec![Field::new( + "l", + DataType::ListView(field), + false, + )])); + let batch = RecordBatch::try_new(schema, vec![list_view])?; + let pieces = SpillPieces::new(&batch, Some(16 * 1024), false)? + .collect::>>()?; + assert_eq!(pieces.len(), 1); + assert_eq!(pieces[0].batch, batch); + Ok(()) + } + + /// The largest row count returned with the size describes the pieces, not the + /// appended batch. + #[tokio::test] + async fn stats_count_rows_per_piece() -> Result<()> { + let batch = RecordBatch::try_new( + Arc::new(Schema::new(vec![Field::new("v", DataType::Utf8, false)])), + vec![Arc::new(StringArray::from_iter_values( + (0..64).map(wide_value), + ))], + )?; + let manager = manager(batch.schema(), Some(4 * WIDE + WIDE / 2))?; + let mut file = manager.create_in_progress_file("test")?; + let WrittenBatchStats { + memory_size: size, + rows, + } = file.append_batch_async_with_stats(&batch).await?; + assert!(size <= 4 * WIDE + WIDE / 2); + assert_eq!(rows, 4); + Ok(()) + } + + /// With many narrow columns the IPC padding of each buffer is larger than the data, so + /// the recorded size must include it or a reader decodes more than was reserved. + #[tokio::test] + async fn recorded_size_includes_ipc_padding() -> Result<()> { + let schema = Arc::new(Schema::new( + (0..4096) + .map(|i| Field::new(format!("c{i}"), DataType::UInt8, false)) + .collect::>(), + )); + let columns = (0..4096) + .map(|_| Arc::new(UInt8Array::from_iter_values(0..17)) as ArrayRef) + .collect(); + let batch = RecordBatch::try_new(schema, columns)?; + // `round_trip` checks every batch read back against the recorded size. + let (read, _) = round_trip(std::slice::from_ref(&batch), 64 * 1024).await?; + assert_same_rows(&read, std::slice::from_ref(&batch)); + Ok(()) + } + + /// A piece whose dictionary keys are all null has no used values. After the unused + /// values are dropped the dictionary is empty, which must not panic when measured, and + /// must read back as nulls. + #[tokio::test] + async fn dictionary_with_only_null_keys() -> Result<()> { + let values = StringArray::from_iter_values((0..64).map(wide_value)); + let keys = Int32Array::from_iter((0..64).map(|i| (i >= 8).then_some(i))); + let dictionary: ArrayRef = + Arc::new(DictionaryArray::try_new(keys, Arc::new(values))?); + let schema = Arc::new(Schema::new(vec![Field::new( + "d", + dictionary_type(DataType::Utf8), + true, + )])); + let batch = RecordBatch::try_new(schema, vec![dictionary])?; + let budget = 4 * WIDE + WIDE / 2; + // The first 8 rows have only null keys. With a budget of one byte each is its own + // piece. + let pieces = SpillPieces::new(&batch.slice(0, 8), Some(1), false)? + .collect::>>()?; + assert_eq!(pieces.len(), 8); + assert!( + pieces.iter().all(|piece| piece + .batch + .column(0) + .as_any_dictionary() + .values() + .is_empty()), + "expected only pieces with an empty dictionary" + ); + SpillPieces::new(&batch.slice(0, 0), Some(budget), false)? + .collect::>>()?; + + let (read, _) = round_trip(std::slice::from_ref(&batch), budget).await?; + let cast = + |b: &RecordBatch| arrow::compute::cast(b.column(0), &DataType::Utf8).unwrap(); + let read_values: Vec = read.iter().map(cast).collect(); + let read_values = arrow::compute::concat( + &read_values.iter().map(|a| a.as_ref()).collect::>(), + )?; + assert_eq!(&read_values, &cast(&batch)); + Ok(()) + } + fn view_schema(fields: &[(&str, DataType)]) -> SchemaRef { + Arc::new(Schema::new( + fields + .iter() + .map(|(name, t)| Field::new(*name, t.clone(), true)) + .collect::>(), + )) + } + + fn assert_bounded(read: &[RecordBatch], max_memory: usize, budget: usize) { + assert!( + max_memory <= budget, + "largest spilled batch is {max_memory} bytes" + ); + assert!(read.len() > 1, "expected the batch to be split"); + } + + /// `take` of every 16th row of a large view array keeps all the parent's buffers, but + /// is not a slice. Garbage collection must see the few values it uses, so the batch + /// fits the budget and is written as one small piece. + #[tokio::test] + async fn sparse_taken_view_array_is_compacted() -> Result<()> { + let parent = StringViewArray::from_iter_values((0..256).map(wide_value)); + let indices = UInt32Array::from_iter_values((0..256).step_by(16)); + let taken = arrow::compute::take(&parent, &indices, None)?; + let batch = + RecordBatch::try_new(view_schema(&[("v", DataType::Utf8View)]), vec![taken])?; + let budget = 2 * 1024 * 1024; + assert!(get_record_batch_memory_size(&batch) > budget); + let (read, max_memory) = round_trip(std::slice::from_ref(&batch), budget).await?; + assert_eq!(read.len(), 1); + assert!(max_memory <= budget, "spilled batch is {max_memory} bytes"); + assert_same_rows(&read, &[batch]); + Ok(()) + } + + /// Sparse `take` outputs of one parent in several batches, each larger than the + /// budget once compacted: every piece holds only its own values. + #[tokio::test] + async fn sparse_taken_views_across_batches_are_split() -> Result<()> { + let parent = StringViewArray::from_iter_values((0..512).map(wide_value)); + let schema = view_schema(&[("v", DataType::Utf8View)]); + let batches = [0u32, 1] + .into_iter() + .map(|start| { + let indices = UInt32Array::from_iter_values((start..512).step_by(2)); + let taken = arrow::compute::take(&parent, &indices, None)?; + Ok(RecordBatch::try_new(Arc::clone(&schema), vec![taken])?) + }) + .collect::>>()?; + let budget = 1024 * 1024; + let (read, max_memory) = round_trip(&batches, budget).await?; + assert_bounded(&read, max_memory, budget); + assert!(read.len() >= 32, "got {} pieces", read.len()); + assert_same_rows(&read, &batches); + Ok(()) + } + + /// Two columns of one batch share the same view buffers: one is the parent, the other + /// a reversed `take` of it. Each piece must count and copy each column's own values. + #[tokio::test] + async fn columns_sharing_view_buffers_are_split() -> Result<()> { + let parent: ArrayRef = + Arc::new(StringViewArray::from_iter_values((0..128).map(wide_value))); + let reversed = arrow::compute::take( + &parent, + &UInt32Array::from_iter_values((0..128).rev()), + None, + )?; + let batch = RecordBatch::try_new( + view_schema(&[("a", DataType::Utf8View), ("b", DataType::Utf8View)]), + vec![Arc::clone(&parent), reversed], + )?; + let budget = 1024 * 1024; + let (read, max_memory) = round_trip(std::slice::from_ref(&batch), budget).await?; + assert_bounded(&read, max_memory, budget); + assert_same_rows(&read, &[batch]); + Ok(()) + } + + /// An unsliced view array mixing inline (at most 12 bytes) values, long values and + /// nulls is split, and every kind of value reads back in order. + #[tokio::test] + async fn mixed_inline_long_and_null_views_are_split() -> Result<()> { + let values = (0..256).map(|i| match i % 3 { + 0 => None, + 1 => Some(format!("s{i}")), + _ => Some(wide_value(i)), + }); + let batch = RecordBatch::try_new( + view_schema(&[("v", DataType::Utf8View)]), + vec![Arc::new(StringViewArray::from_iter(values))], + )?; + let budget = 1024 * 1024; + let (read, max_memory) = round_trip(std::slice::from_ref(&batch), budget).await?; + assert_bounded(&read, max_memory, budget); + assert_same_rows(&read, &[batch]); + Ok(()) + } + + /// Binary views are split like string views. + #[tokio::test] + async fn binary_view_column_is_split() -> Result<()> { + let values: Vec> = (0..128).map(|i| wide_value(i).into_bytes()).collect(); + let batch = RecordBatch::try_new( + view_schema(&[("b", DataType::BinaryView)]), + vec![Arc::new(BinaryViewArray::from_iter_values(values))], + )?; + let budget = 1024 * 1024; + let (read, max_memory) = round_trip(std::slice::from_ref(&batch), budget).await?; + assert_bounded(&read, max_memory, budget); + assert_same_rows(&read, &[batch]); + Ok(()) + } + + /// A dictionary whose values are views: pruning the dictionary must also drop the + /// view buffers that only the unused values pointed at. + #[tokio::test] + async fn dictionary_of_views_is_divided() -> Result<()> { + let values = StringViewArray::from_iter_values((0..128).map(wide_value)); + let keys = Int32Array::from_iter_values(0..128); + let dictionary = DictionaryArray::try_new(keys, Arc::new(values))?; + let batch = RecordBatch::try_new( + view_schema(&[("d", dictionary_type(DataType::Utf8View))]), + vec![Arc::new(dictionary)], + )?; + let budget = 1024 * 1024; + let (read, max_memory) = round_trip(std::slice::from_ref(&batch), budget).await?; + assert_bounded(&read, max_memory, budget); + let cast = + |b: &RecordBatch| arrow::compute::cast(b.column(0), &DataType::Utf8).unwrap(); + let read_values: Vec = read.iter().map(cast).collect(); + let read_values = arrow::compute::concat( + &read_values.iter().map(|a| a.as_ref()).collect::>(), + )?; + assert_eq!(&read_values, &cast(&batch)); + Ok(()) + } + + /// Views inside a struct and a list are below the top level. Each piece must still + /// hold only the view data of its own rows. + #[tokio::test] + async fn nested_views_are_split() -> Result<()> { + let views: ArrayRef = + Arc::new(StringViewArray::from_iter_values((0..256).map(wide_value))); + let fields = Fields::from(vec![Field::new("v", DataType::Utf8View, false)]); + let structs: ArrayRef = + Arc::new(StructArray::try_new(fields.clone(), vec![views], None)?); + let item = Arc::new(Field::new_list_field(DataType::Utf8View, false)); + let lists: ArrayRef = Arc::new(arrow::array::ListArray::try_new( + Arc::clone(&item), + arrow::buffer::OffsetBuffer::from_lengths(vec![1; 256]), + Arc::new(StringViewArray::from_iter_values((0..256).map(wide_value))), + None, + )?); + let batch = RecordBatch::try_new( + view_schema(&[("s", DataType::Struct(fields)), ("l", DataType::List(item))]), + vec![structs, lists], + )?; + let budget = 2 * 1024 * 1024; + let (read, max_memory) = round_trip(std::slice::from_ref(&batch), budget).await?; + assert_bounded(&read, max_memory, budget); + assert_same_rows(&read, &[batch]); + Ok(()) + } +} + +#[cfg(test)] +/// The recorded size bounds what a reader decodes, for every codec and for buffer types +/// whose IPC padding or decompressed allocation is larger than their data. +mod decoded_size_bound_tests { + use super::*; + use arrow::array::*; + use arrow_schema::{DataType, Field, Schema}; + use datafusion_common::config::SpillCompression; + use datafusion_common::utils::memory::get_record_batch_memory_size; + use datafusion_execution::runtime_env::RuntimeEnvBuilder; + use datafusion_physical_expr_common::metrics::{ + ExecutionPlanMetricsSet, SpillMetrics, + }; + use futures::TryStreamExt; + + fn cols(n: usize, rows: usize) -> Vec<(String, RecordBatch)> { + let mk = |name: &str, dt: DataType, f: &dyn Fn() -> ArrayRef| { + let schema = Arc::new(Schema::new( + (0..n) + .map(|i| Field::new(format!("c{i}"), dt.clone(), true)) + .collect::>(), + )); + ( + name.to_string(), + RecordBatch::try_new(schema, (0..n).map(|_| f()).collect()).unwrap(), + ) + }; + vec![ + mk("u8", DataType::UInt8, &|| { + Arc::new(UInt8Array::from_iter_values((0..rows).map(|i| i as u8))) + }), + mk("i64_nulls", DataType::Int64, &|| { + Arc::new(Int64Array::from_iter( + (0..rows).map(|i| (i % 2 == 0).then_some(i as i64)), + )) + }), + mk("bool", DataType::Boolean, &|| { + Arc::new(BooleanArray::from_iter((0..rows).map(|i| Some(i % 2 == 0)))) + }), + mk("utf8", DataType::Utf8, &|| { + Arc::new(StringArray::from_iter_values( + (0..rows).map(|i| format!("v{i}")), + )) + }), + mk("utf8view", DataType::Utf8View, &|| { + Arc::new(StringViewArray::from_iter_values( + (0..rows).map(|i| format!("a-long-string-value-{i}")), + )) + }), + mk("dec256", DataType::Decimal256(40, 2), &|| { + Arc::new( + Decimal256Array::from_iter_values( + (0..rows).map(|i| arrow::datatypes::i256::from_i128(i as i128)), + ) + .with_precision_and_scale(40, 2) + .unwrap(), + ) + }), + ] + } + + #[tokio::test] + async fn recorded_size_bounds_decoded_batch() -> Result<()> { + for comp in [ + SpillCompression::Uncompressed, + SpillCompression::Lz4Frame, + SpillCompression::Zstd, + ] { + for (n, rows) in [(256usize, 17usize), (4, 4000)] { + for (name, batch) in cols(n, rows) { + let runtime = Arc::new(RuntimeEnvBuilder::new().build()?); + let metrics = SpillMetrics::new(&ExecutionPlanMetricsSet::new(), 0); + let m = Arc::new( + SpillManager::new(runtime, metrics, batch.schema()) + .with_compression_type(comp) + .with_max_batch_bytes(Some(1 << 30)), + ); + let mut f = m.create_in_progress_file("t")?; + let rec = f.append_batch(&batch)?; + let file = f.finish()?.unwrap(); + let read: Vec = + m.read_spill_as_stream(file, None)?.try_collect().await?; + let actual = + read.iter().map(get_record_batch_memory_size).max().unwrap(); + assert!( + actual <= rec + crate::spill::SPILL_BATCH_MEMORY_MARGIN, + "{comp:?} {name} cols={n} rows={rows}: recorded {rec}, decoded {actual}" + ); + } + } + } + Ok(()) + } } diff --git a/datafusion/physical-plan/src/spill/mod.rs b/datafusion/physical-plan/src/spill/mod.rs index f2ec70c6f3544..bb0c2d047c637 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; mod lengths; pub(crate) mod replayable_spill_input; pub(crate) mod spill_manager; +pub(crate) mod spill_pieces; pub mod spill_pool; pub(crate) mod spilled_row_sizes; use datafusion_execution::spill_file::{AsyncSpillWriter, SpillWriter}; @@ -780,7 +781,7 @@ fn buffer_to_bytes(buffer: Buffer) -> bytes::Bytes { } // Returns the maximum byte alignment required by any field in the schema (>= 8), derived from Arrow buffer layouts. -fn get_max_alignment_for_schema(schema: &Schema) -> usize { +pub(crate) fn get_max_alignment_for_schema(schema: &Schema) -> usize { let minimum_alignment = 8; let mut max_alignment = minimum_alignment; for field in schema.fields() { diff --git a/datafusion/physical-plan/src/spill/spill_manager.rs b/datafusion/physical-plan/src/spill/spill_manager.rs index 57e819b1408c4..98df5d681798a 100644 --- a/datafusion/physical-plan/src/spill/spill_manager.rs +++ b/datafusion/physical-plan/src/spill/spill_manager.rs @@ -32,6 +32,8 @@ pub struct SpillManager { batch_read_buffer_capacity: usize, /// general-purpose compression options pub(crate) compression: SpillCompression, + /// See [`Self::with_max_batch_bytes`] + max_batch_bytes: Option, } impl SpillManager { @@ -42,6 +44,7 @@ impl SpillManager { schema, batch_read_buffer_capacity: 2, compression: SpillCompression::default(), + max_batch_bytes: None, } } @@ -58,6 +61,27 @@ impl SpillManager { self } + /// Bounds the in-memory size of each batch written through this manager, or `None` + /// for no bound. + /// + /// A batch is bounded by `batch_size` rows, so with wide rows it can be very large in + /// bytes. A reader must decode each spilled batch whole, and a merge of spill files + /// reserves memory for the largest batch of each file (`max_record_batch_memory`). + /// With a bound, a larger batch is split into row ranges of at most this size (best + /// effort: one row can be larger), and each range is written as its own IPC message. + /// Row order is kept, but one appended batch can come back as several batches, so a + /// reader that expects one batch per append (such as [`crate::spill::spill_pool`]) + /// must not use a bounded manager. + pub(crate) fn with_max_batch_bytes(mut self, max_batch_bytes: Option) -> Self { + self.max_batch_bytes = max_batch_bytes; + self + } + + /// See [`Self::with_max_batch_bytes`] + pub(crate) fn max_batch_bytes(&self) -> Option { + self.max_batch_bytes + } + /// Returns the schema for batches managed by this SpillManager pub fn schema(&self) -> &SchemaRef { &self.schema @@ -144,10 +168,12 @@ impl SpillManager { while let Some(batch) = stream.next().await { let batch = batch?; - let gc_sliced_size = in_progress_file.append_batch_async(&batch).await?; + let stats = in_progress_file + .append_batch_async_with_stats(&batch) + .await?; - max_record_batch_size = max_record_batch_size.max(gc_sliced_size); - max_batch_rows = max_batch_rows.max(batch.num_rows()); + max_record_batch_size = max_record_batch_size.max(stats.memory_size); + max_batch_rows = max_batch_rows.max(stats.rows); } let file = in_progress_file.finish_async().await?; diff --git a/datafusion/physical-plan/src/spill/spill_pieces.rs b/datafusion/physical-plan/src/spill/spill_pieces.rs new file mode 100644 index 0000000000000..1128f019c37aa --- /dev/null +++ b/datafusion/physical-plan/src/spill/spill_pieces.rs @@ -0,0 +1,271 @@ +// 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 + +//! Splits a batch into byte-bounded pieces before it is written to a spill file. + +use std::ops::Range; + +use arrow::array::{ArrayRef, RecordBatch, RecordBatchOptions, make_array}; +use arrow::datatypes::{DataType, Schema}; +use arrow_data::ArrayData; +use arrow_data::transform::MutableArrayData; +use datafusion_common::Result; + +use super::spill_manager::GetSlicedSize; +use super::spilled_row_sizes::SpilledRowSizes; +use super::{gc_columns, gc_view_arrays, get_max_alignment_for_schema}; + +/// One piece of a batch, to be written as one IPC message. +#[derive(Debug)] +pub(crate) struct SpillPiece { + pub(crate) batch: RecordBatch, + /// The size recorded for the piece, see [`recorded_size`]. + pub(crate) size: usize, +} + +/// The pieces that one appended batch is written as, in row order. +/// +/// Without a bound, or when the batch is already within it, there is one piece: the batch +/// after [`gc_view_arrays`], as it is written without a bound. +/// +/// Otherwise the batch is cut into consecutive row ranges by the size of each row once +/// spilled ([`SpilledRowSizes`]), less the padding that every piece adds. Each range is +/// sliced, its list columns are copied (see [`narrow_lists`]), and it is compacted like the +/// chunks of a bounded sort (see [`SpilledRowSizes::compacted_columns`]), so a dictionary whose values many keys share +/// stays shared. That size counts only the keys of such a dictionary, so a piece can hold +/// more than it estimates. A piece whose [`recorded_size`] is still larger than the bound +/// is halved until it fits or has one row. Pieces are produced one at a time, so that the +/// caller can write each one before the next is compacted. +/// +/// Schemas that a slice does not narrow (see [`can_split`]) are written whole. +pub(crate) struct SpillPieces { + batch: RecordBatch, + max_batch_bytes: usize, + compressed: bool, + /// The bytes left for rows in each piece, once its padding is counted. + row_budget: usize, + /// `None` when the batch is written whole. + row_sizes: Option, + /// The first row not yet in a piece. + next_row: usize, + /// Ranges that were too large, to produce in halves, the next one last. + pending: Vec>, + /// The whole batch, when it is written as one piece. + whole: Option, +} + +impl SpillPieces { + pub(crate) fn new( + batch: &RecordBatch, + max_batch_bytes: Option, + compressed: bool, + ) -> Result { + let batch = gc_view_arrays(batch)?; + let size = recorded_size(&batch, compressed)?; + let max_batch_bytes = max_batch_bytes.filter(|max| { + size > *max && batch.num_rows() > 1 && can_split(batch.schema_ref()) + }); + let Some(max_batch_bytes) = max_batch_bytes else { + return Ok(Self { + whole: Some(SpillPiece { + batch: batch.clone(), + size, + }), + batch, + max_batch_bytes: 0, + compressed, + row_budget: 0, + row_sizes: None, + next_row: usize::MAX, + pending: vec![], + }); + }; + let padding = padding(&batch, padding_alignment(batch.schema_ref(), compressed)); + Ok(Self { + row_sizes: Some(SpilledRowSizes::new(&batch)), + batch, + max_batch_bytes, + compressed, + row_budget: max_batch_bytes.saturating_sub(padding), + next_row: 0, + pending: vec![], + whole: None, + }) + } + + /// Returns the next range of rows, from `self.next_row`, that fits the row budget. + fn next_range(&mut self, row_sizes: &SpilledRowSizes) -> Option> { + let (start, rows) = (self.next_row, self.batch.num_rows()); + if start >= rows { + return None; + } + let end = if let Some(row_bytes) = row_sizes.fixed() { + start + (self.row_budget / row_bytes.max(1)).clamp(1, rows - start) + } else { + let mut bytes = 0; + let mut end = start; + while end < rows { + bytes += row_sizes.row(end); + if end > start && bytes > self.row_budget { + break; + } + end += 1; + } + end + }; + self.next_row = end; + Some(start..end) + } +} + +impl Iterator for SpillPieces { + type Item = Result; + + fn next(&mut self) -> Option { + if let Some(whole) = self.whole.take() { + return Some(Ok(whole)); + } + let row_sizes = self.row_sizes.take()?; + let piece = loop { + let Some(range) = self.pending.pop().or_else(|| self.next_range(&row_sizes)) + else { + break None; + }; + let piece = self.batch.slice(range.start, range.len()); + let piece = match narrow_lists(&piece) + .and_then(|piece| gc_columns(&piece, row_sizes.compacted_columns())) + .and_then(|piece| gc_view_arrays(&piece)) + { + Ok(piece) => piece, + Err(e) => break Some(Err(e)), + }; + let size = match recorded_size(&piece, self.compressed) { + Ok(size) => size, + Err(e) => break Some(Err(e)), + }; + if size > self.max_batch_bytes && range.len() > 1 { + let mid = range.start + range.len() / 2; + self.pending.push(mid..range.end); + self.pending.push(range.start..mid); + continue; + } + break Some(Ok(SpillPiece { batch: piece, size })); + }; + self.row_sizes = Some(row_sizes); + piece + } +} + +/// The capacity granularity of arrow's `MutableBuffer`, which a compressed buffer is +/// decompressed into. +const DECOMPRESSED_ALIGNMENT: usize = 64; + +/// The size recorded for a batch written to a spill file: an upper bound on the memory a +/// reader uses for it once decoded from IPC. Use it wherever a batch read back is compared +/// with the recorded size. +/// +/// It is the batch's [`GetSlicedSize`] plus the padding of every buffer. Without +/// compression the buffers are slices of the message body, where the IPC writer pads each +/// one to the schema's alignment. With compression each is decompressed into its own +/// allocation, rounded up to [`DECOMPRESSED_ALIGNMENT`]. For many narrow columns the +/// padding can be larger than the data. +pub(crate) fn recorded_size(batch: &RecordBatch, compressed: bool) -> Result { + let alignment = padding_alignment(batch.schema_ref(), compressed); + Ok(batch.get_sliced_size()? + padding(batch, alignment)) +} + +/// The padding a reader can add to each buffer of `schema`. +fn padding_alignment(schema: &Schema, compressed: bool) -> usize { + let alignment = get_max_alignment_for_schema(schema); + if compressed { + alignment.max(DECOMPRESSED_ALIGNMENT) + } else { + alignment + } +} + +/// The padding of every buffer of `batch`, at `alignment` bytes each. +fn padding(batch: &RecordBatch, alignment: usize) -> usize { + fn buffers(data: &ArrayData) -> usize { + // One validity buffer, which the writer emits even when it is empty. + 1 + data.buffers().len() + data.child_data().iter().map(buffers).sum::() + } + let buffers: usize = batch.columns().iter().map(|c| buffers(&c.to_data())).sum(); + buffers * alignment +} + +/// Copies the columns of `slice` that hold lists, so that their children hold only the +/// elements of its rows. The IPC writer narrows them the same way, but +/// [`GetSlicedSize`] counts the whole children of a sliced list, so without the copy +/// [`recorded_size`] would count every piece as the whole batch. +fn narrow_lists(slice: &RecordBatch) -> Result { + fn has_list(data_type: &DataType) -> bool { + match data_type { + DataType::List(_) + | DataType::LargeList(_) + | DataType::FixedSizeList(_, _) + | DataType::Map(_, _) => true, + DataType::Struct(fields) => { + fields.iter().any(|field| has_list(field.data_type())) + } + DataType::Dictionary(_, values) => has_list(values), + _ => false, + } + } + if !slice + .schema_ref() + .fields() + .iter() + .any(|field| has_list(field.data_type())) + { + return Ok(slice.clone()); + } + let columns = slice + .columns() + .iter() + .map(|column| { + if !has_list(column.data_type()) { + return Ok(ArrayRef::clone(column)); + } + let data = column.to_data(); + let mut copy = MutableArrayData::new(vec![&data], false, data.len()); + copy.try_extend(0, 0, data.len())?; + Ok(make_array(copy.freeze())) + }) + .collect::>>()?; + let options = RecordBatchOptions::new().with_row_count(Some(slice.num_rows())); + Ok(RecordBatch::try_new_with_options( + slice.schema(), + columns, + &options, + )?) +} + +/// Whether slicing, with [`narrow_lists`], narrows every column of `schema` to its rows. +/// It does not for list +/// views and unions, whose children a row range does not bound, or for run-end encoded +/// arrays, whose values a slice keeps whole. A batch with such a column is written whole. +fn can_split(schema: &Schema) -> bool { + fn supported(data_type: &DataType) -> bool { + match data_type { + DataType::ListView(_) + | DataType::LargeListView(_) + | DataType::Union(_, _) + | DataType::RunEndEncoded(_, _) => false, + DataType::Dictionary(_, values) => supported(values), + DataType::List(field) + | DataType::LargeList(field) + | DataType::FixedSizeList(field, _) + | DataType::Map(field, _) => supported(field.data_type()), + DataType::Struct(fields) => { + fields.iter().all(|field| supported(field.data_type())) + } + _ => true, + } + } + schema + .fields() + .iter() + .all(|field| supported(field.data_type())) +} diff --git a/datafusion/physical-plan/tests/cases/aggregate_memory_call_sequence.rs b/datafusion/physical-plan/tests/cases/aggregate_memory_call_sequence.rs index 853cbabf3ce50..ac45a5d795f50 100644 --- a/datafusion/physical-plan/tests/cases/aggregate_memory_call_sequence.rs +++ b/datafusion/physical-plan/tests/cases/aggregate_memory_call_sequence.rs @@ -524,7 +524,6 @@ async fn spills_twice_then_replays() -> Result<()> { FinalHashAggregateStream[0] shrink [bytes] FinalHashAggregateStream[0] try_grow [bytes] FinalHashAggregateStream[0] try_grow [bytes] - FinalHashAggregateStream[0] shrink [bytes] FinalHashAggregateStream[0] try_grow [bytes] FinalHashAggregateStream[0] try_grow [bytes] FinalHashAggregateStream[0] shrink [bytes] @@ -787,7 +786,6 @@ async fn clustered_spills_then_replays() -> Result<()> { ClusteredFinalAggregateStream[0] try_grow [bytes] -> refused ClusteredFinalAggregateStream[0] try_grow [bytes] ClusteredFinalAggregateStream[0] try_grow [bytes] - ClusteredFinalAggregateStream[0] shrink [bytes] ClusteredFinalAggregateStream[0] try_grow [bytes] ClusteredFinalAggregateStream[0] try_grow [bytes] ClusteredFinalAggregateStream[0] shrink [bytes] diff --git a/docs/source/library-user-guide/upgrading/56.0.0.md b/docs/source/library-user-guide/upgrading/56.0.0.md index 07714dcd5589c..6c73d31ab01cb 100644 --- a/docs/source/library-user-guide/upgrading/56.0.0.md +++ b/docs/source/library-user-guide/upgrading/56.0.0.md @@ -1621,3 +1621,22 @@ let return_type: Option = definition `CreateFunction::try_new` takes `Option` for the return type. See [issue #18235](https://github.com/apache/datafusion/issues/18235) for details. + +### Spill merges no longer re-spill runs with halved batches + +`StreamingMergeBuilder::with_sorted_spill_files` merges sorted spill runs by +reserving memory for the largest batch of each run. When two runs did not fit, +the merge used to re-spill the larger run with every batch halved and retry. +It now returns `ResourcesExhausted` instead. + +DataFusion's sort and aggregate write their runs in batches of a bounded size, +so that two runs fit the memory their merges reserve. Only a single row, or a +column that cannot be split by rows (such as a list view or a union), can be +larger. + +**Migration guide:** + +If you write your own runs for `with_sorted_spill_files`, keep each batch small +enough that two runs fit the memory you expect the merge to have. The merge +charges each batch twice its size and reads at least one batch ahead from each +run, so batches of at most a quarter of that memory fit.