diff --git a/Cargo.lock b/Cargo.lock index be40df978369d..1ad4559f88862 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2552,6 +2552,7 @@ dependencies = [ "arrow-ipc", "arrow-ord", "arrow-schema", + "arrow-select", "async-trait", "bytes", "criterion", 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 761c492d8dde3..a772ef0610890 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 @@ -273,12 +273,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), @@ -290,12 +290,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 d8814eaa86307..6f12718f6a828 100644 --- a/datafusion/core/tests/memory_limit/mod.rs +++ b/datafusion/core/tests/memory_limit/mod.rs @@ -26,6 +26,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..824aa76687379 --- /dev/null +++ b/datafusion/core/tests/memory_limit/wide_row_sort_spill.rs @@ -0,0 +1,147 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! 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::{ + ArrayRef, Int64Array, ListBuilder, RecordBatch, StringArray, StringBuilder, + StringViewArray, +}; +use arrow_schema::{DataType, Field, Schema}; +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)) +} + +async fn sort_wide_rows(data_type: DataType) { + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int64, false), + Field::new("payload", data_type.clone(), false), + ])); + let batches: Vec = (0..ROWS) + .step_by(INPUT_BATCH_ROWS) + .map(|start| { + let rows = start..start + INPUT_BATCH_ROWS; + let ids = Int64Array::from_iter_values(rows.clone().map(|i| i as i64)); + 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!(), + }; + RecordBatch::try_new(Arc::clone(&schema), vec![Arc::new(ids), payload]) + .unwrap() + }) + .collect(); + + let table = MemTable::try_new(Arc::clone(&schema), vec![batches]).unwrap(); + let runtime = RuntimeEnvBuilder::new() + .with_memory_pool(Arc::new(FairSpillPool::new(POOL_BYTES))) + .build_arc() + .unwrap(); + 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); + ctx.register_table("t", Arc::new(table)).unwrap(); + + let plan = ctx + .sql("SELECT id, payload 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); + assert!(spill_count > 0, "the sort must spill, plan:\n{plan:?}"); + + let ids: Vec = result + .iter() + .flat_map(|batch| { + batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .values() + .to_vec() + }) + .collect(); + let expected: Vec = (0..ROWS as i64).rev().collect(); + assert_eq!(ids, expected); +} + +#[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; +} diff --git a/datafusion/physical-plan/Cargo.toml b/datafusion/physical-plan/Cargo.toml index 82a2215581f22..900af993a02f9 100644 --- a/datafusion/physical-plan/Cargo.toml +++ b/datafusion/physical-plan/Cargo.toml @@ -64,6 +64,7 @@ arrow-data = { workspace = true } arrow-ipc = { workspace = true, features = ["lz4", "zstd"] } arrow-ord = { workspace = true } arrow-schema = { workspace = true } +arrow-select = { workspace = true } async-trait = { workspace = true } bytes = { workspace = true } datafusion-common = { workspace = true } @@ -142,6 +143,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..b5f9ed25b84c2 --- /dev/null +++ b/datafusion/physical-plan/benches/sort_spill.rs @@ -0,0 +1,147 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! 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/sorts/multi_level_merge.rs b/datafusion/physical-plan/src/sorts/multi_level_merge.rs index 28a8000265c7b..55da678ce11f6 100644 --- a/datafusion/physical-plan/src/sorts/multi_level_merge.rs +++ b/datafusion/physical-plan/src/sorts/multi_level_merge.rs @@ -157,6 +157,8 @@ pub(crate) struct MultiLevelMergeBuilder { reservation: MemoryReservation, /// Workspace retained across retries and intermediate spill passes. merge_pool: Option>, + /// Bound on the size of each batch written to intermediate spill files. + spill_max_batch_bytes: Option, /// Leave memory for the aggregate consuming the merged spill rows. reserve_replay_headroom: bool, /// Avoid merging more runs than needed before the final replay pass. @@ -200,6 +202,7 @@ impl MultiLevelMergeBuilder { batch_size, reservation, merge_pool: None, + spill_max_batch_bytes: None, reserve_replay_headroom: false, size_intermediate_merges: false, enable_round_robin_tie_breaker, @@ -207,6 +210,14 @@ impl MultiLevelMergeBuilder { } } + pub(super) fn with_spill_max_batch_bytes( + mut self, + spill_max_batch_bytes: Option, + ) -> Self { + self.spill_max_batch_bytes = spill_max_batch_bytes; + self + } + pub(super) fn with_merge_pool(mut self, pool: Option>) -> Self { self.merge_pool = pool; self @@ -294,6 +305,7 @@ impl MultiLevelMergeBuilder { .spill_record_batch_stream_and_return_max_batch_stats( &mut stream, "MultiLevelMergeBuilder intermediate spill", + self.spill_max_batch_bytes, ) .await? else { @@ -813,6 +825,7 @@ impl MultiLevelMergeBuilder { .spill_record_batch_stream_and_return_max_batch_stats( &mut halved, "MultiLevelMergeBuilder split skewed spill", + self.spill_max_batch_bytes, ) .await?; @@ -907,11 +920,14 @@ fn split_batch_in_half(batch: RecordBatch) -> Vec { vec![batch.slice(0, mid), batch.slice(mid, num_rows - mid)] } +/// The fewest spill files a merge reads at once. +pub(super) const MIN_SPILL_MERGE_FAN_IN: usize = 2; + 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) } } diff --git a/datafusion/physical-plan/src/sorts/sort.rs b/datafusion/physical-plan/src/sorts/sort.rs index da272e2ff3341..297e0b08693fa 100644 --- a/datafusion/physical-plan/src/sorts/sort.rs +++ b/datafusion/physical-plan/src/sorts/sort.rs @@ -42,6 +42,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; @@ -268,6 +269,9 @@ struct ExternalSorter { /// Reservation for in_mem_batches reservation: MemoryReservation, spill_manager: SpillManager, + /// Bound on the size of each batch this sort and its merge write to spill files. See + /// [`spill_max_batch_bytes`]. + spill_max_batch_bytes: usize, /// Reservation for the merging of in-memory batches. If the sort /// might spill, `sort_spill_reservation_bytes` will be @@ -328,6 +332,7 @@ impl ExternalSorter { in_mem_batches: vec![], in_mem_batches_counter, in_progress_spill_file: None, + spill_max_batch_bytes: spill_max_batch_bytes(sort_spill_reservation_bytes), finished_spill_files: vec![], expr, metrics, @@ -390,6 +395,7 @@ impl ExternalSorter { StreamingMergeBuilder::new() .with_sorted_spill_files(std::mem::take(&mut self.finished_spill_files)) .with_spill_manager(self.spill_manager.clone()) + .with_spill_max_batch_bytes(Some(self.spill_max_batch_bytes)) .with_schema(Arc::clone(&self.schema)) .with_expressions(&self.expr.clone()) .with_metrics(self.metrics.baseline.clone()) @@ -445,8 +451,12 @@ impl ExternalSorter { // Lazily initialize the in-progress spill file if self.in_progress_spill_file.is_none() { - self.in_progress_spill_file = - Some((self.spill_manager.create_in_progress_file("Sorting")?, 0)); + self.in_progress_spill_file = Some(( + self.spill_manager + .create_in_progress_file("Sorting")? + .with_max_batch_bytes(Some(self.spill_max_batch_bytes)), + 0, + )); } debug!("Spilling sort data of ExternalSorter to disk whilst inserting"); @@ -998,6 +1008,38 @@ impl ExternalSorter { } } +/// The bound on the in-memory size of each batch that [`ExternalSorter`] spills. +/// +/// The merge of the spill files starts with `sort_spill_reservation_bytes` of headroom. +/// It 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 the headroom, found from the same formula. +/// +/// 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. +fn spill_max_batch_bytes(sort_spill_reservation_bytes: usize) -> usize { + let fits = |batch: usize| { + MIN_SPILL_MERGE_FAN_IN * get_reserved_bytes_for_record_batch_size(batch, batch) + <= sort_spill_reservation_bytes + }; + // Binary search for the largest batch size that fits: `fits` is true up to it. + let (mut fitting, mut too_large) = (0, sort_spill_reservation_bytes + 1); + while too_large - fitting > 1 { + let mid = fitting + (too_large - fitting) / 2; + if fits(mid) { + fitting = mid; + } else { + too_large = mid; + } + } + fitting.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 @@ -2554,6 +2596,23 @@ 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); + } + #[tokio::test] async fn test_in_mem_sort() -> Result<()> { let task_ctx = Arc::new(TaskContext::default()); diff --git a/datafusion/physical-plan/src/sorts/streaming_merge.rs b/datafusion/physical-plan/src/sorts/streaming_merge.rs index 249b20cf49f38..7e95ea9dfb5c3 100644 --- a/datafusion/physical-plan/src/sorts/streaming_merge.rs +++ b/datafusion/physical-plan/src/sorts/streaming_merge.rs @@ -96,6 +96,7 @@ pub struct StreamingMergeBuilder<'a> { fetch: Option, reservation: Option, merge_pool: Option>, + spill_max_batch_bytes: Option, /// Leave memory for the aggregate consuming the merged spill rows. reserve_replay_headroom: bool, min_spill_batch_rows: Option, @@ -158,6 +159,18 @@ impl<'a> StreamingMergeBuilder<'a> { self } + /// Bound the size of each batch that a merge of spill files writes to its + /// intermediate spill files. See [`InProgressSpillFile::with_max_batch_bytes`]. + /// + /// [`InProgressSpillFile::with_max_batch_bytes`]: crate::spill::in_progress_spill_file::InProgressSpillFile::with_max_batch_bytes + pub(super) fn with_spill_max_batch_bytes( + mut self, + spill_max_batch_bytes: Option, + ) -> Self { + self.spill_max_batch_bytes = spill_max_batch_bytes; + self + } + /// Keep spill workspace until the final merge pass selects its buffer budget. pub(crate) fn with_merge_pool(mut self, pool: Arc) -> Self { self.merge_pool = Some(pool); @@ -217,6 +230,7 @@ impl<'a> StreamingMergeBuilder<'a> { batch_size, reservation, merge_pool, + spill_max_batch_bytes, reserve_replay_headroom, min_spill_batch_rows, fetch, @@ -260,6 +274,7 @@ impl<'a> StreamingMergeBuilder<'a> { enable_round_robin_tie_breaker, ) .with_merge_pool(merge_pool) + .with_spill_max_batch_bytes(spill_max_batch_bytes) .with_replay_headroom(reserve_replay_headroom) .with_intermediate_merge_sizing(min_spill_batch_rows) .create_spillable_merge_stream()); 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 2f4ddd3398f8e..d3619c67a5f0a 100644 --- a/datafusion/physical-plan/src/spill/in_progress_spill_file.rs +++ b/datafusion/physical-plan/src/spill/in_progress_spill_file.rs @@ -25,8 +25,8 @@ 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::SpillPieces, }; enum InProgressWriter { @@ -43,6 +43,9 @@ pub struct InProgressSpillFile { writer: Option, /// Lazily initialized in-progress file, it will be moved out when the `finish` method is invoked in_progress_file: Option>, + /// Upper bound on the in-memory size of each written batch. See + /// [`Self::with_max_batch_bytes`]. + max_batch_bytes: Option, } impl InProgressSpillFile { @@ -54,16 +57,39 @@ impl InProgressSpillFile { spill_manager, in_progress_file: Some(in_progress_file), writer: None, + max_batch_bytes: None, } } + /// Bounds the in-memory size of each batch written to this file, 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. + /// Batches within the bound are written as they are. Row order is kept, but one + /// appended batch can come back as several batches when the file is read. + /// + /// The bound is set on the writer and not on the [`SpillManager`], so a reader that + /// expects one batch per append (such as [`crate::spill::spill_pool`]) cannot get a + /// bounded file from a shared manager. + pub(crate) fn with_max_batch_bytes(mut self, max_batch_bytes: Option) -> Self { + self.max_batch_bytes = max_batch_bytes; + self + } + /// Appends a `RecordBatch` to the spill file, initializing the writer if necessary. /// /// Before writing, performs GC on StringView/BinaryView arrays to compact backing /// 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 [`Self::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) @@ -76,7 +102,7 @@ impl InProgressSpillFile { )); } - let gc_batch = gc_view_arrays(batch)?; + let pieces = SpillPieces::new(batch, self.max_batch_bytes)?; if self.writer.is_none() { // Use the SpillManager's declared schema rather than the batch's schema. @@ -97,29 +123,44 @@ impl InProgressSpillFile { self.spill_manager.metrics.spill_file_count.add(1); } } + let mut max_size = 0; 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 (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); + max_size = max_size.max(size); + } } 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(max_size) } /// 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?.0) + } + + /// 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<(usize, usize)> { 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.max_batch_bytes)?; if self.writer.is_none() { let schema = self.spill_manager.schema(); @@ -141,12 +182,18 @@ impl InProgressSpillFile { } } + let (mut max_size, mut max_rows) = (0, 0); match &mut self.writer { Some(InProgressWriter::Async(writer)) => { - let (spilled_rows, delta_bytes) = writer.write(&gc_batch).await?; - - self.spill_manager.metrics.spilled_rows.add(spilled_rows); - self.spill_manager.metrics.spilled_bytes.add(delta_bytes); + for piece in pieces { + let (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); + max_size = max_size.max(size); + max_rows = max_rows.max(piece.num_rows()); + } } Some(InProgressWriter::Sync(_)) => { return Err(exec_datafusion_err!( @@ -159,7 +206,7 @@ impl InProgressSpillFile { )); } } - gc_batch.get_sliced_size() + Ok((max_size, max_rows)) } pub fn flush(&mut self) -> Result<()> { @@ -248,8 +295,13 @@ 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, DictionaryArray, Int32Array, Int64Array, ListBuilder, + MapBuilder, StringArray, StringBuilder, StringViewArray, StructArray, + }; + 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, @@ -317,4 +369,439 @@ mod tests { Ok(()) } + + const WIDE: usize = 64 * 1024; + + fn wide_value(i: usize) -> String { + format!("{i}{}", "x".repeat(WIDE)) + } + + fn manager(schema: SchemaRef) -> Result> { + let runtime = Arc::new(RuntimeEnvBuilder::new().build()?); + let metrics = SpillMetrics::new(&ExecutionPlanMetricsSet::new(), 0); + Ok(Arc::new(SpillManager::new(runtime, metrics, schema))) + } + + /// 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())?; + let mut file = manager + .create_in_progress_file("test")? + .with_max_batch_bytes(Some(max_batch_bytes)); + 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())?; + 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))?.collect::>>()?; + assert_eq!(pieces.len(), 1); + assert_eq!(pieces[0].0, batch); + Ok(()) + } + + /// Values of a dictionary inside a struct are not pruned, so every piece would carry + /// the whole dictionary. The batch is written whole instead. + #[tokio::test] + async fn nested_dictionary_is_written_whole() -> 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 = 1024 * 1024; + let (read, _) = round_trip(std::slice::from_ref(&batch), budget).await?; + assert_eq!(read.len(), 1); + 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))?.collect::>>()?; + assert_eq!(pieces.len(), 1); + assert_eq!(pieces[0].0, 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())?; + let mut file = manager + .create_in_progress_file("test")? + .with_max_batch_bytes(Some(4 * WIDE + WIDE / 2)); + let (size, rows) = file.append_batch_async_with_stats(&batch).await?; + assert!(size <= 4 * WIDE + WIDE / 2); + assert_eq!(rows, 4); + 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. + #[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 fields = + Fields::from(vec![Field::new("d", dictionary_type(DataType::Utf8), true)]); + let structs: ArrayRef = Arc::new(StructArray::try_new( + fields.clone(), + vec![Arc::clone(&dictionary)], + None, + )?); + let schema = Arc::new(Schema::new(vec![ + Field::new("d", dictionary_type(DataType::Utf8), true), + Field::new("s", DataType::Struct(fields), false), + ])); + let batch = RecordBatch::try_new(schema, vec![dictionary, structs])?; + // Pieces of 8 rows: the first has only null keys. + let budget = 20 * WIDE; + for batch in [batch.slice(0, 8), batch.slice(0, 0), batch] { + if batch.num_rows() == 0 { + SpillPieces::new(&batch, Some(budget))?.collect::>>()?; + continue; + } + let (read, _) = round_trip(std::slice::from_ref(&batch), budget).await?; + assert_same_rows(&read, &[batch]); + } + Ok(()) + } } diff --git a/datafusion/physical-plan/src/spill/mod.rs b/datafusion/physical-plan/src/spill/mod.rs index 9eaa04389c1b1..1b4afea19cbd9 100644 --- a/datafusion/physical-plan/src/spill/mod.rs +++ b/datafusion/physical-plan/src/spill/mod.rs @@ -20,6 +20,7 @@ pub(crate) mod in_progress_spill_file; pub(crate) mod replayable_spill_input; pub(crate) mod spill_manager; +pub(crate) mod spill_pieces; pub mod spill_pool; use datafusion_execution::spill_file::{AsyncSpillWriter, SpillWriter}; // Moved for refactor, re-export to keep the public API stable diff --git a/datafusion/physical-plan/src/spill/spill_manager.rs b/datafusion/physical-plan/src/spill/spill_manager.rs index 9c07602fab070..8980dce8453ae 100644 --- a/datafusion/physical-plan/src/spill/spill_manager.rs +++ b/datafusion/physical-plan/src/spill/spill_manager.rs @@ -141,15 +141,19 @@ impl SpillManager { } /// Spill a stream of `RecordBatch`es and return the file, maximum batch memory, - /// and maximum batch row count. + /// and maximum batch row count. See [`InProgressSpillFile::with_max_batch_bytes`] + /// for `max_batch_bytes`. pub(crate) async fn spill_record_batch_stream_and_return_max_batch_stats( &self, stream: &mut SendableRecordBatchStream, request_description: &str, + max_batch_bytes: Option, ) -> Result, usize, usize)>> { use futures::StreamExt; - let mut in_progress_file = self.create_in_progress_file(request_description)?; + let mut in_progress_file = self + .create_in_progress_file(request_description)? + .with_max_batch_bytes(max_batch_bytes); let result = async { let mut max_record_batch_size = 0; @@ -157,10 +161,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 (gc_sliced_size, rows) = 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_batch_rows = max_batch_rows.max(rows); } let file = in_progress_file.finish_async().await?; @@ -455,6 +461,7 @@ mod tests { .spill_record_batch_stream_and_return_max_batch_stats( &mut stream, "abort-test", + None, ) .await .err() 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..6d1ccc45c0487 --- /dev/null +++ b/datafusion/physical-plan/src/spill/spill_pieces.rs @@ -0,0 +1,170 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Splits a batch into byte-bounded pieces before it is written to a spill file. + +use std::ops::Range; +use std::sync::Arc; + +use arrow::array::{ArrayRef, AsArray, RecordBatch, RecordBatchOptions, make_array}; +use arrow::datatypes::{DataType, Schema}; +use arrow_data::transform::MutableArrayData; +use arrow_select::dictionary::garbage_collect_any_dictionary; +use datafusion_common::Result; + +use super::gc_view_arrays; +use super::spill_manager::GetSlicedSize; + +/// 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. Its dictionaries are not +/// changed, so the batches of a spill file can keep sharing one dictionary. +/// +/// When the batch is larger than `max_batch_bytes`, it is cut into equal row ranges by its +/// average row size. Each range is copied out with [`compact`], so that it holds only its +/// own rows, and measured with [`GetSlicedSize`]. A range that is still too large (the rows +/// differ in size) 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 copied. +/// +/// Schemas that [`compact`] cannot narrow (see [`can_split`]) are written whole. +pub(crate) struct SpillPieces { + batch: RecordBatch, + max_batch_bytes: usize, + /// Row ranges still to produce, the next one last. + pending: Vec>, + /// The whole batch, when it is written as one piece. + whole: Option<(RecordBatch, usize)>, +} + +impl SpillPieces { + pub(crate) fn new( + batch: &RecordBatch, + max_batch_bytes: Option, + ) -> Result { + let batch = gc_view_arrays(batch)?; + let size = batch.get_sliced_size()?; + let rows = batch.num_rows(); + let Some(max_batch_bytes) = max_batch_bytes + .filter(|max| size > *max && rows > 1 && can_split(batch.schema_ref())) + else { + return Ok(Self { + whole: Some((batch, size)), + batch: RecordBatch::new_empty(Arc::new(Schema::empty())), + max_batch_bytes: 0, + pending: vec![], + }); + }; + // The sliced size over-counts nested columns, whose children are not narrowed by + // a slice, so this can give more pieces than needed but not fewer. + let pieces = size.div_ceil(max_batch_bytes.max(1)).min(rows); + let rows_per_piece = rows.div_ceil(pieces); + let pending = (0..rows) + .step_by(rows_per_piece) + .map(|start| start..(start + rows_per_piece).min(rows)) + .rev() + .collect(); + Ok(Self { + batch, + max_batch_bytes, + pending, + whole: None, + }) + } +} + +impl Iterator for SpillPieces { + /// A piece and its size, as [`GetSlicedSize`] measures it. + type Item = Result<(RecordBatch, usize)>; + + fn next(&mut self) -> Option { + if let Some(whole) = self.whole.take() { + return Some(Ok(whole)); + } + while let Some(range) = self.pending.pop() { + let piece = match compact(&self.batch.slice(range.start, range.len())) { + Ok(piece) => piece, + Err(e) => return Some(Err(e)), + }; + let size = match piece.get_sliced_size() { + Ok(size) => size, + Err(e) => return 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; + } + return Some(Ok((piece, size))); + } + None + } +} + +/// Copies a slice so that it holds only its own rows: list children are narrowed to the +/// range the slice uses, view arrays are garbage collected, and top-level dictionaries drop +/// the values that no key uses. +fn compact(slice: &RecordBatch) -> Result { + let columns = slice + .columns() + .iter() + .map(|column| { + let data = column.to_data(); + let mut copy = MutableArrayData::new(vec![&data], false, data.len()); + copy.try_extend(0, 0, data.len())?; + let copied = make_array(copy.freeze()); + Ok(match copied.data_type() { + DataType::Dictionary(_, _) => { + garbage_collect_any_dictionary(copied.as_any_dictionary())? + } + _ => copied, + }) + }) + .collect::>>()?; + let options = RecordBatchOptions::new().with_row_count(Some(slice.num_rows())); + let piece = RecordBatch::try_new_with_options(slice.schema(), columns, &options)?; + gc_view_arrays(&piece) +} + +/// Whether [`compact`] narrows every column of `schema` to its rows. It does not for list +/// views and unions, whose children a row range does not bound, run-end encoded arrays, or +/// dictionaries below the top level, whose values it does not prune. A batch with such a +/// column is written whole. +fn can_split(schema: &Schema) -> bool { + fn supported(data_type: &DataType, top_level: bool) -> bool { + match data_type { + DataType::ListView(_) + | DataType::LargeListView(_) + | DataType::Union(_, _) + | DataType::RunEndEncoded(_, _) => false, + DataType::Dictionary(_, values) => top_level && supported(values, false), + DataType::List(field) + | DataType::LargeList(field) + | DataType::FixedSizeList(field, _) + | DataType::Map(field, _) => supported(field.data_type(), false), + DataType::Struct(fields) => fields + .iter() + .all(|field| supported(field.data_type(), false)), + _ => true, + } + } + schema + .fields() + .iter() + .all(|field| supported(field.data_type(), true)) +}