From 87fd025e0d18e786ba5e7c92c029e91ccba7434a Mon Sep 17 00:00:00 2001 From: Adrian Garcia Badaracco <1755071+adriangb@users.noreply.github.com> Date: Sun, 25 Oct 2026 15:38:54 -0500 Subject: [PATCH 1/9] test: external sort over wide rows under a memory limit Adds integration tests that sort 512 rows of 64 KiB (Utf8, Utf8View and List) with a 12 MiB pool. Each input batch fits in the pool, but the sort writes batch_size-row spill batches that do not, so the final merge fails with ResourcesExhausted. These tests fail until the next commits. Co-Authored-By: Claude Opus 5.5 --- datafusion/core/tests/memory_limit/mod.rs | 1 + .../tests/memory_limit/wide_row_sort_spill.rs | 147 ++++++++++++++++++ 2 files changed, 148 insertions(+) create mode 100644 datafusion/core/tests/memory_limit/wide_row_sort_spill.rs 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; +} From d50efec810112944fd8f5991ea217a1a6a276546 Mon Sep 17 00:00:00 2001 From: Adrian Garcia Badaracco <1755071+adriangb@users.noreply.github.com> Date: Sun, 25 Oct 2026 15:39:33 -0500 Subject: [PATCH 2/9] refactor: return the row count of the written spill batch SpillManager recorded the row count of the appended batch as the largest batch of the file. The writer now returns the row count of what it wrote, so the next commit can write one appended batch as several smaller ones. No change in behavior. Co-Authored-By: Claude Opus 5.5 --- .../physical-plan/src/spill/in_progress_spill_file.rs | 11 ++++++++++- datafusion/physical-plan/src/spill/spill_manager.rs | 6 ++++-- 2 files changed, 14 insertions(+), 3 deletions(-) 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..0e9e88f560b38 100644 --- a/datafusion/physical-plan/src/spill/in_progress_spill_file.rs +++ b/datafusion/physical-plan/src/spill/in_progress_spill_file.rs @@ -113,6 +113,15 @@ impl InProgressSpillFile { /// 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 batch written. + 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." @@ -159,7 +168,7 @@ impl InProgressSpillFile { )); } } - gc_batch.get_sliced_size() + Ok((gc_batch.get_sliced_size()?, gc_batch.num_rows())) } pub fn flush(&mut self) -> Result<()> { diff --git a/datafusion/physical-plan/src/spill/spill_manager.rs b/datafusion/physical-plan/src/spill/spill_manager.rs index 9c07602fab070..df0cf4da3be03 100644 --- a/datafusion/physical-plan/src/spill/spill_manager.rs +++ b/datafusion/physical-plan/src/spill/spill_manager.rs @@ -157,10 +157,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?; From fa41a998dde548191fe4a1b0441d07ce9ea301c3 Mon Sep 17 00:00:00 2001 From: Adrian Garcia Badaracco <1755071+adriangb@users.noreply.github.com> Date: Sun, 25 Oct 2026 19:17:08 -0500 Subject: [PATCH 3/9] bench: external sort that spills under a memory limit Adds a sort_spill benchmark with three cases: a dictionary column with 100,000 distinct values, short strings, and 1 KiB strings. Each case uses a memory pool smaller than the input, so the sort spills and merges. Co-Authored-By: Claude Opus 5.5 --- datafusion/physical-plan/Cargo.toml | 5 + .../physical-plan/benches/sort_spill.rs | 147 ++++++++++++++++++ 2 files changed, 152 insertions(+) create mode 100644 datafusion/physical-plan/benches/sort_spill.rs diff --git a/datafusion/physical-plan/Cargo.toml b/datafusion/physical-plan/Cargo.toml index 82a2215581f22..62f00464bea27 100644 --- a/datafusion/physical-plan/Cargo.toml +++ b/datafusion/physical-plan/Cargo.toml @@ -142,6 +142,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); From ba86f43a7aba37d4a78e67dcead12bf71b9bf5cd Mon Sep 17 00:00:00 2001 From: Adrian Garcia Badaracco <1755071+adriangb@users.noreply.github.com> Date: Sun, 25 Oct 2026 19:19:19 -0500 Subject: [PATCH 4/9] fix: bound spill file batches by bytes so external sorts over wide rows can merge The sort wrote spill batches of batch_size rows whatever their size. The final merge reserves twice the largest batch of each spill file, so with wide rows it could not start. InProgressSpillFile::with_max_batch_bytes now splits a batch larger than the bound into row ranges. Each range is copied so that it holds only its own rows, and the ranges are written one at a time. A batch within the bound is written as before, so its dictionaries stay shared. Batches with list views, unions, run-end encoded arrays or nested dictionaries are written whole. The bound is set per writer, not on SpillManager, so readers that expect one batch per append (spill_pool, sort-merge join) cannot get a bounded file. The sort sets it for its spill files and for the intermediate spill files of its merge. The bound is the largest batch for which the minimum merge (two files, get_reserved_bytes_for_record_batch_size each) fits in sort_spill_reservation_bytes, and at least 64 KiB. Co-Authored-By: Claude Opus 5.5 --- Cargo.lock | 1 + datafusion/physical-plan/Cargo.toml | 1 + .../src/sorts/multi_level_merge.rs | 18 +- datafusion/physical-plan/src/sorts/sort.rs | 63 ++++++- .../src/sorts/streaming_merge.rs | 15 ++ .../src/spill/in_progress_spill_file.rs | 70 ++++++-- datafusion/physical-plan/src/spill/mod.rs | 1 + .../physical-plan/src/spill/spill_manager.rs | 9 +- .../physical-plan/src/spill/spill_pieces.rs | 170 ++++++++++++++++++ 9 files changed, 327 insertions(+), 21 deletions(-) create mode 100644 datafusion/physical-plan/src/spill/spill_pieces.rs 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/physical-plan/Cargo.toml b/datafusion/physical-plan/Cargo.toml index 62f00464bea27..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 } 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 0e9e88f560b38..addf3494c988b 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,18 +123,23 @@ 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. @@ -117,7 +148,8 @@ impl InProgressSpillFile { } /// Like [`Self::append_batch_async`], but returns both the memory size and the row - /// count of the batch written. + /// 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, @@ -128,7 +160,7 @@ impl InProgressSpillFile { )); } - 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(); @@ -150,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!( @@ -168,7 +206,7 @@ impl InProgressSpillFile { )); } } - Ok((gc_batch.get_sliced_size()?, gc_batch.num_rows())) + Ok((max_size, max_rows)) } pub fn flush(&mut self) -> Result<()> { 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 df0cf4da3be03..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; @@ -457,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)) +} From 84c7bcf6290d04b86e1dbbd0d445bda4cbad23fb Mon Sep 17 00:00:00 2001 From: Adrian Garcia Badaracco <1755071+adriangb@users.noreply.github.com> Date: Sun, 25 Oct 2026 19:19:19 -0500 Subject: [PATCH 5/9] test: the oversized-batch sort fuzz test no longer needs a re-spill The sort now bounds its spill batches, so the merge can seat two streams without the fallback that re-spills a run at half the batch size. The test now checks that the sort spills and finishes. The fallback keeps its unit tests in multi_level_merge.rs. Co-Authored-By: Claude Opus 5.5 --- ...spilling_fuzz_in_memory_constrained_env.rs | 19 ++++++++----------- 1 file changed, 8 insertions(+), 11 deletions(-) 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(()) From 782ef12f49a2815e3366887661f47a3bda157e16 Mon Sep 17 00:00:00 2001 From: Adrian Garcia Badaracco <1755071+adriangb@users.noreply.github.com> Date: Sun, 25 Oct 2026 19:19:19 -0500 Subject: [PATCH 6/9] test: unit tests for bounded spill batches Covers view, list, map, struct and dictionary columns, and the cases that are written whole: one wide row, batches within the bound (dictionary kept), list views with shared children and nested dictionaries. Co-Authored-By: Claude Opus 5.5 --- .../src/spill/in_progress_spill_file.rs | 444 +++++++++++++++++- 1 file changed, 442 insertions(+), 2 deletions(-) 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 addf3494c988b..d3619c67a5f0a 100644 --- a/datafusion/physical-plan/src/spill/in_progress_spill_file.rs +++ b/datafusion/physical-plan/src/spill/in_progress_spill_file.rs @@ -295,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, @@ -364,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(()) + } } From e436f3b4922faa4604fe0778b333c33bb43be931 Mon Sep 17 00:00:00 2001 From: Adrian Garcia Badaracco <1755071+adriangb@users.noreply.github.com> Date: Mon, 26 Oct 2026 17:27:56 -0500 Subject: [PATCH 7/9] fix: prune shared dictionaries once before splitting oversized spill batches Also test the empty-dictionary case with a splittable top-level schema. Co-Authored-By: Claude Opus 5.5 --- .../src/spill/in_progress_spill_file.rs | 57 +++++++++++-------- .../physical-plan/src/spill/spill_pieces.rs | 51 +++++++++++++++-- 2 files changed, 78 insertions(+), 30 deletions(-) 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 d3619c67a5f0a..6a501e68ac05f 100644 --- a/datafusion/physical-plan/src/spill/in_progress_spill_file.rs +++ b/datafusion/physical-plan/src/spill/in_progress_spill_file.rs @@ -297,7 +297,7 @@ mod tests { use super::*; use crate::spill::SPILL_BATCH_MEMORY_MARGIN; use arrow::array::{ - Array, ArrayRef, DictionaryArray, Int32Array, Int64Array, ListBuilder, + Array, ArrayRef, AsArray, DictionaryArray, Int32Array, Int64Array, ListBuilder, MapBuilder, StringArray, StringBuilder, StringViewArray, StructArray, }; use arrow_schema::{DataType, Field, Fields, Schema, SchemaRef}; @@ -773,35 +773,44 @@ mod tests { } /// 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. + /// 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 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]); - } + let schema = Arc::new(Schema::new(vec![Field::new( + "d", + dictionary_type(DataType::Utf8), + true, + )])); + let batch = RecordBatch::try_new(schema, vec![dictionary])?; + // Pieces of at most 4 rows: the first ones have only null keys. + let budget = 4 * WIDE + WIDE / 2; + let pieces = + SpillPieces::new(&batch, Some(budget))?.collect::>>()?; + assert!(pieces.len() > 1); + assert!( + pieces.iter().any(|(piece, _)| piece + .column(0) + .as_any_dictionary() + .values() + .is_empty()), + "expected a piece with an empty dictionary" + ); + SpillPieces::new(&batch.slice(0, 0), Some(budget))? + .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(()) } } diff --git a/datafusion/physical-plan/src/spill/spill_pieces.rs b/datafusion/physical-plan/src/spill/spill_pieces.rs index 6d1ccc45c0487..7e653291a1d30 100644 --- a/datafusion/physical-plan/src/spill/spill_pieces.rs +++ b/datafusion/physical-plan/src/spill/spill_pieces.rs @@ -59,16 +59,25 @@ impl SpillPieces { let batch = gc_view_arrays(batch)?; let size = batch.get_sliced_size()?; let rows = batch.num_rows(); + let whole = |batch, size| Self { + whole: Some((batch, size)), + batch: RecordBatch::new_empty(Arc::new(Schema::empty())), + max_batch_bytes: 0, + pending: vec![], + }; 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![], - }); + return Ok(whole(batch, size)); }; + // A shared dictionary can hold many values that no row of this batch uses. Drop + // them once before choosing the row ranges, so that the ranges are sized by what + // the batch uses and the pieces do not each prune the whole dictionary again. + let batch = prune_dictionaries(&batch)?; + let size = batch.get_sliced_size()?; + if size <= max_batch_bytes { + return Ok(whole(batch, size)); + } // 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); @@ -141,6 +150,36 @@ fn compact(slice: &RecordBatch) -> Result { gc_view_arrays(&piece) } +/// Drops the values of top-level dictionaries that no key uses. +fn prune_dictionaries(batch: &RecordBatch) -> Result { + if !batch + .schema_ref() + .fields() + .iter() + .any(|field| matches!(field.data_type(), DataType::Dictionary(_, _))) + { + return Ok(batch.clone()); + } + let columns = batch + .columns() + .iter() + .map(|column| { + Ok(match column.data_type() { + DataType::Dictionary(_, _) => { + garbage_collect_any_dictionary(column.as_any_dictionary())? + } + _ => Arc::clone(column), + }) + }) + .collect::>>()?; + let options = RecordBatchOptions::new().with_row_count(Some(batch.num_rows())); + Ok(RecordBatch::try_new_with_options( + batch.schema(), + columns, + &options, + )?) +} + /// 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 From 8f46ec431ff3f8bcdfee41a60aab1ad1364ba858 Mon Sep 17 00:00:00 2001 From: Adrian Garcia Badaracco <1755071+adriangb@users.noreply.github.com> Date: Tue, 27 Oct 2026 08:15:31 -0500 Subject: [PATCH 8/9] test: Utf8View and dictionary edge cases for bounded spill batches Co-Authored-By: Claude Opus 5.5 --- .../src/spill/in_progress_spill_file.rs | 170 +++++++++++++++++- 1 file changed, 168 insertions(+), 2 deletions(-) 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 6a501e68ac05f..1c935992eb842 100644 --- a/datafusion/physical-plan/src/spill/in_progress_spill_file.rs +++ b/datafusion/physical-plan/src/spill/in_progress_spill_file.rs @@ -297,8 +297,9 @@ mod tests { use super::*; use crate::spill::SPILL_BATCH_MEMORY_MARGIN; use arrow::array::{ - Array, ArrayRef, AsArray, DictionaryArray, Int32Array, Int64Array, ListBuilder, - MapBuilder, StringArray, StringBuilder, StringViewArray, StructArray, + Array, ArrayRef, AsArray, BinaryViewArray, DictionaryArray, Int32Array, + Int64Array, ListBuilder, MapBuilder, StringArray, StringBuilder, StringViewArray, + StructArray, UInt32Array, }; use arrow_schema::{DataType, Field, Fields, Schema, SchemaRef}; use datafusion_common::utils::memory::get_record_batch_memory_size; @@ -813,4 +814,169 @@ mod tests { 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(()) + } } From 3feccc90a2942676ac0ef7d5b11e3ebf2c57a49d Mon Sep 17 00:00:00 2001 From: Adrian Garcia Badaracco <1755071+adriangb@users.noreply.github.com> Date: Tue, 27 Oct 2026 08:24:10 -0500 Subject: [PATCH 9/9] test: end-to-end memory-limited sorts over view and dictionary edge cases Co-Authored-By: Claude Opus 5.5 --- .../tests/memory_limit/wide_row_sort_spill.rs | 299 ++++++++++++++---- 1 file changed, 242 insertions(+), 57 deletions(-) diff --git a/datafusion/core/tests/memory_limit/wide_row_sort_spill.rs b/datafusion/core/tests/memory_limit/wide_row_sort_spill.rs index 824aa76687379..0e0c42b9e8286 100644 --- a/datafusion/core/tests/memory_limit/wide_row_sort_spill.rs +++ b/datafusion/core/tests/memory_limit/wide_row_sort_spill.rs @@ -26,10 +26,13 @@ use std::sync::Arc; use arrow::array::{ - ArrayRef, Int64Array, ListBuilder, RecordBatch, StringArray, StringBuilder, - StringViewArray, + Array, ArrayRef, BinaryViewArray, DictionaryArray, Int32Array, Int64Array, ListArray, + ListBuilder, RecordBatch, StringArray, StringBuilder, StringViewArray, StructArray, + UInt32Array, }; -use arrow_schema::{DataType, Field, Schema}; +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; @@ -50,57 +53,51 @@ 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) +/// The row ranges of the input batches. +fn input_ranges() -> impl Iterator> { + (0..ROWS) .step_by(INPUT_BATCH_ROWS) - .map(|start| { - let rows = start..start + 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 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() + let mut cols: Vec = vec![Arc::new(ids)]; + cols.extend(columns(rows)); + RecordBatch::try_new(Arc::clone(&schema), cols).unwrap() }) .collect(); + (schema, batches) +} - 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(); +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); + 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 id, payload FROM t ORDER BY id DESC") + .sql("SELECT * FROM t ORDER BY id DESC") .await .unwrap() .create_physical_plan() @@ -109,24 +106,63 @@ async fn sort_wide_rows(data_type: DataType) { 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:?}"); + (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 - .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); + .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] @@ -145,3 +181,152 @@ async fn sort_wide_utf8_view_rows_spills_and_finishes() { 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; +}