From 64d148ecf75d831f8ccd2d29f3e396464592e7be Mon Sep 17 00:00:00 2001 From: comphead Date: Fri, 30 Oct 2026 15:25:44 -0700 Subject: [PATCH 1/4] perf: coalesce small IPC buffers into one spill write `StreamEncoder` returns every buffer of a batch and its padding as a separate `Buffer`, and the spill writers wrote them one by one. For file backed spills that is one `write` system call per buffer, about 35 for a batch of six string and integer columns. Copy runs of buffers under 64 KB into one, so a spilled batch typically takes a single write. --- datafusion/physical-plan/src/spill/mod.rs | 34 +++++++++++++++++++++-- 1 file changed, 31 insertions(+), 3 deletions(-) diff --git a/datafusion/physical-plan/src/spill/mod.rs b/datafusion/physical-plan/src/spill/mod.rs index ab57d87dceecb..bc6c72d9c400e 100644 --- a/datafusion/physical-plan/src/spill/mod.rs +++ b/datafusion/physical-plan/src/spill/mod.rs @@ -686,16 +686,43 @@ impl IPCStreamEncoder { fn encode(&mut self, batch: &RecordBatch) -> Result<(usize, usize, Vec)> { let buffers = self.encoder.as_mut().unwrap().encode(batch)?; let bytes = buffers.iter().map(Buffer::len).sum(); - Ok((batch.num_rows(), bytes, buffers)) + Ok((batch.num_rows(), bytes, coalesce_buffers(buffers))) } fn finish(&mut self) -> Result<(usize, Vec)> { let buffers = self.encoder.take().unwrap().finish()?; let bytes = buffers.iter().map(Buffer::len).sum(); - Ok((bytes, buffers)) + Ok((bytes, coalesce_buffers(buffers))) } } +/// Buffers at least this large are written on their own, see [`coalesce_buffers`] +const COALESCE_BUFFERS_BELOW: usize = 64 * 1024; + +/// Copies runs of small buffers into one, so that writing the IPC encoding of +/// a batch takes a few writes instead of one per buffer and padding. +/// +/// [`StreamEncoder`] returns every buffer of a batch and its padding separately, +/// and file backed spills would write each with its own system call. +fn coalesce_buffers(buffers: Vec) -> Vec { + let mut coalesced = Vec::new(); + let mut small: Vec = Vec::new(); + for buffer in buffers { + if buffer.len() < COALESCE_BUFFERS_BELOW { + small.extend_from_slice(buffer.as_slice()); + continue; + } + if !small.is_empty() { + coalesced.push(Buffer::from_vec(std::mem::take(&mut small))); + } + coalesced.push(buffer); + } + if !small.is_empty() { + coalesced.push(Buffer::from_vec(small)); + } + coalesced +} + struct ArrowBufferOwner(Buffer); impl AsRef<[u8]> for ArrowBufferOwner { @@ -1132,7 +1159,8 @@ mod tests { let trailing_bytes = writer.finish().await?; assert!(finished.load(Ordering::Relaxed)); - assert!(writes.load(Ordering::Relaxed) > 2); + // One write per batch and one for the end of the stream + assert_eq!(writes.load(Ordering::Relaxed), 3); drop(writer); assert!(!aborted.load(Ordering::Relaxed)); From 75e617205fc757bd04e36284ba64ec26905c43d1 Mon Sep 17 00:00:00 2001 From: comphead Date: Fri, 30 Oct 2026 14:41:22 -0700 Subject: [PATCH 2/4] perf: spill byte arrays as lengths when compressed Closes #26627. Offsets are a running total, so zstd keeps about 2.5 of every 4 bytes and lz4 all of them, while the lengths of similar values repeat. Compressed spill files now store each top level `Utf8`, `LargeUtf8`, `Binary` and `LargeBinary` column as a struct of its lengths and of its value bytes as a single run. Both are valid arrays, so the IPC reader decodes them as usual and the spill reader rebuilds the offsets with a running sum. --- datafusion/physical-plan/src/spill/lengths.rs | 298 ++++++++++++++++++ datafusion/physical-plan/src/spill/mod.rs | 40 ++- 2 files changed, 335 insertions(+), 3 deletions(-) create mode 100644 datafusion/physical-plan/src/spill/lengths.rs diff --git a/datafusion/physical-plan/src/spill/lengths.rs b/datafusion/physical-plan/src/spill/lengths.rs new file mode 100644 index 0000000000000..888a71f7188d4 --- /dev/null +++ b/datafusion/physical-plan/src/spill/lengths.rs @@ -0,0 +1,298 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: See the NOTICE file distributed with this work for additional information regarding copyright ownership +// SPDX-FileContributor: Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements + +//! Spills byte arrays as the lengths of their values. +//! +//! The offsets of a `Utf8` array are a running total that zstd and lz4 barely +//! compress, while the lengths of similar values repeat, see +//! . Compressed spill files +//! therefore store each top level `Utf8`, `LargeUtf8`, `Binary` and +//! `LargeBinary` column as a struct of its lengths, which keep its nulls, and +//! of all its value bytes as a single run: +//! +//! ```text +//! Utf8 ["ab", null, "cde"] +//! +//! Struct { +//! lengths: Int32 [2, null, 3] +//! values: RunEndEncoded { run_ends: [3], values: Binary ["abcde"] } +//! } +//! ``` +//! +//! Both are valid arrays, so the IPC reader decodes them as usual before +//! [`decode`] rebuilds the offsets with a running sum. + +use std::sync::Arc; + +use arrow::array::{ + Array, ArrayRef, ArrowPrimitiveType, AsArray, GenericBinaryArray, GenericByteArray, + Int64Array, OffsetSizeTrait, PrimitiveArray, RecordBatch, RunArray, StructArray, + new_empty_array, +}; +use arrow::buffer::OffsetBuffer; +use arrow::datatypes::{ + ArrowNativeType, BinaryType, ByteArrayType, DataType, Field, Fields, + GenericBinaryType, Int32Type, Int64Type, LargeBinaryType, LargeUtf8Type, Schema, + SchemaRef, Utf8Type, +}; +use datafusion_common::{Result, internal_err}; + +/// Returns the schema that `schema` is spilled as, or `None` if it has no byte +/// array columns +pub(super) fn encode_schema(schema: &Schema) -> Option { + if !schema + .fields() + .iter() + .any(|field| encoded_type(field.data_type()).is_some()) + { + return None; + } + let fields: Vec<_> = schema + .fields() + .iter() + .map(|field| match encoded_type(field.data_type()) { + Some(data_type) => { + Arc::new(Field::new(field.name(), data_type, field.is_nullable())) + } + None => Arc::clone(field), + }) + .collect(); + Some(Arc::new(Schema::new_with_metadata( + fields, + schema.metadata().clone(), + ))) +} + +/// Encodes the byte array columns of `batch`, see the [module docs](self) +pub(super) fn encode(batch: &RecordBatch) -> Result { + let Some(schema) = encode_schema(batch.schema_ref()) else { + return Ok(batch.clone()); + }; + let columns = batch + .columns() + .iter() + .map(encode_array) + .collect::>()?; + Ok(RecordBatch::try_new(schema, columns)?) +} + +/// Restores a batch that [`encode`] wrote for `schema` +pub(super) fn decode(batch: &RecordBatch, schema: &SchemaRef) -> Result { + let columns = schema + .fields() + .iter() + .zip(batch.columns()) + .map(|(field, column)| decode_array(field.data_type(), column)) + .collect::>()?; + Ok(RecordBatch::try_new(Arc::clone(schema), columns)?) +} + +fn encoded_type(data_type: &DataType) -> Option { + let fields = match data_type { + DataType::Utf8 | DataType::Binary => encoded_fields::(), + DataType::LargeUtf8 | DataType::LargeBinary => encoded_fields::(), + _ => return None, + }; + Some(DataType::Struct(fields)) +} + +/// The fields of the struct a byte array with `L` offsets is spilled as +fn encoded_fields() -> Fields +where + L: ArrowPrimitiveType, + L::Native: OffsetSizeTrait, +{ + let run_ends = Field::new("run_ends", DataType::Int64, false); + let bytes = Field::new("values", GenericBinaryType::::DATA_TYPE, true); + let values = DataType::RunEndEncoded(Arc::new(run_ends), Arc::new(bytes)); + Fields::from(vec![ + Field::new("lengths", L::DATA_TYPE, true), + Field::new("values", values, false), + ]) +} + +fn encode_array(array: &ArrayRef) -> Result { + match array.data_type() { + DataType::Utf8 => encode_bytes::<_, Int32Type>(array.as_string::()), + DataType::LargeUtf8 => encode_bytes::<_, Int64Type>(array.as_string::()), + DataType::Binary => encode_bytes::<_, Int32Type>(array.as_binary::()), + DataType::LargeBinary => encode_bytes::<_, Int64Type>(array.as_binary::()), + _ => Ok(Arc::clone(array)), + } +} + +fn decode_array(data_type: &DataType, array: &ArrayRef) -> Result { + match data_type { + DataType::Utf8 => decode_bytes::(array), + DataType::LargeUtf8 => decode_bytes::(array), + DataType::Binary => decode_bytes::(array), + DataType::LargeBinary => decode_bytes::(array), + _ => Ok(Arc::clone(array)), + } +} + +fn encode_bytes(array: &GenericByteArray) -> Result +where + T: ByteArrayType, + L: ArrowPrimitiveType, +{ + let fields = encoded_fields::(); + if array.is_empty() { + return Ok(new_empty_array(&DataType::Struct(fields))); + } + let offsets = array.offsets(); + let lengths = offsets.windows(2).map(|w| w[1] - w[0]).collect(); + let lengths = PrimitiveArray::::new(lengths, array.nulls().cloned()); + + let start = offsets[0].as_usize(); + let len = offsets[array.len()].as_usize() - start; + let bytes = GenericBinaryArray::::new( + OffsetBuffer::from_lengths([len]), + array.values().slice_with_length(start, len), + None, + ); + let run_ends = Int64Array::from(vec![array.len() as i64]); + let values = RunArray::try_new(&run_ends, &bytes)?; + + let array = + StructArray::try_new(fields, vec![Arc::new(lengths), Arc::new(values)], None)?; + Ok(Arc::new(array)) +} + +fn decode_bytes(array: &ArrayRef) -> Result +where + T: ByteArrayType, + L: ArrowPrimitiveType, +{ + let array = array.as_struct(); + let lengths = array.column(0).as_primitive::(); + let mut end = T::Offset::usize_as(0); + let mut offsets = Vec::with_capacity(lengths.len() + 1); + offsets.push(end); + offsets.extend(lengths.values().iter().map(|&len| { + end += len; + end + })); + + let bytes = array.column(1).as_run::().values(); + let bytes = bytes.as_binary::(); + let bytes_offsets = bytes.value_offsets(); + let start = bytes_offsets[0].as_usize(); + let len = bytes_offsets[bytes_offsets.len() - 1].as_usize() - start; + if end.as_usize() != len { + return internal_err!( + "Spilled {} lengths add up to {}, but there are {len} value bytes", + T::DATA_TYPE, + end.as_usize() + ); + } + + // SAFETY: `encode_bytes` took the lengths and value bytes from a valid + // array of type `T`, so their running sum gives back its offsets + let array = unsafe { + GenericByteArray::::new_unchecked( + OffsetBuffer::new_unchecked(offsets.into()), + bytes.values().slice_with_length(start, len), + lengths.nulls().cloned(), + ) + }; + Ok(Arc::new(array)) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::common::collect; + use crate::metrics::{ExecutionPlanMetricsSet, SpillMetrics}; + use crate::spill::spill_manager::SpillManager; + use arrow::array::{Int32Array, ListArray, StringArray}; + use arrow::compute::cast; + use arrow_ipc::CompressionType; + use arrow_ipc::writer::{IpcWriteOptions, StreamWriter}; + use datafusion_common::config::SpillCompression; + use datafusion_execution::runtime_env::RuntimeEnv; + + fn spill_manager(schema: SchemaRef, compression: SpillCompression) -> SpillManager { + let env = Arc::new(RuntimeEnv::default()); + let metrics = SpillMetrics::new(&ExecutionPlanMetricsSet::new(), 0); + SpillManager::new(env, metrics, schema).with_compression_type(compression) + } + + /// Byte arrays read back as written, sliced, empty or with nulls, next to + /// the columns that are spilled as they are + #[tokio::test] + async fn test_spill_lengths_round_trip() -> Result<()> { + let utf8: ArrayRef = Arc::new(StringArray::from_iter( + (0..1000).map(|i| (i % 7 != 0).then(|| "x".repeat(i % 5))), + )); + let list = ListArray::new( + Arc::new(Field::new("item", DataType::Utf8, true)), + OffsetBuffer::from_lengths([1; 1000]), + Arc::clone(&utf8), + None, + ); + let batch = RecordBatch::try_from_iter(vec![ + ( + "int32", + Arc::new(Int32Array::from_iter_values(0..1000)) as _, + ), + ("large_utf8", cast(&utf8, &DataType::LargeUtf8)?), + ("binary", cast(&utf8, &DataType::Binary)?), + ("large_binary", cast(&utf8, &DataType::LargeBinary)?), + ("utf8_view", cast(&utf8, &DataType::Utf8View)?), + ("list", Arc::new(list)), + ("utf8", utf8), + ])?; + let batches = [ + batch.clone(), + batch.slice(3, 500), + batch.slice(998, 2), + batch.slice(0, 0), + ]; + for compression in [ + SpillCompression::Uncompressed, + SpillCompression::Lz4Frame, + SpillCompression::Zstd, + ] { + let spill_manager = spill_manager(batch.schema(), compression); + let spill_file = spill_manager + .spill_record_batch_and_finish(&batches, "Test")? + .unwrap(); + let stream = spill_manager.read_spill_as_stream(spill_file, None)?; + assert_eq!(collect(stream).await?, batches, "{compression}"); + } + Ok(()) + } + + /// The lengths of similar values compress to almost nothing, unlike their + /// offsets + #[test] + fn test_spill_lengths_compress() -> Result<()> { + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Utf8, false)])); + let values = (0..8192).map(|i| format!("item-{}", i / 64)); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(StringArray::from_iter_values(values))], + )?; + + let spill_manager = spill_manager(Arc::clone(&schema), SpillCompression::Zstd); + spill_manager + .spill_record_batch_and_finish(std::slice::from_ref(&batch), "Test")?; + let spilled = spill_manager.metrics.spilled_bytes.value(); + + let options = IpcWriteOptions::default() + .try_with_compression(Some(CompressionType::ZSTD))?; + let mut ipc = vec![]; + let mut writer = StreamWriter::try_new_with_options(&mut ipc, &schema, options)?; + writer.write(&batch)?; + writer.finish()?; + drop(writer); + assert!( + spilled * 4 < ipc.len(), + "spilled {spilled} bytes, against {} bytes of IPC", + ipc.len() + ); + Ok(()) + } +} diff --git a/datafusion/physical-plan/src/spill/mod.rs b/datafusion/physical-plan/src/spill/mod.rs index bc6c72d9c400e..0ba6687020047 100644 --- a/datafusion/physical-plan/src/spill/mod.rs +++ b/datafusion/physical-plan/src/spill/mod.rs @@ -5,6 +5,7 @@ //! Defines the spilling functions pub(crate) mod in_progress_spill_file; +mod lengths; pub(crate) mod replayable_spill_input; pub(crate) mod spill_manager; pub mod spill_pool; @@ -75,6 +76,9 @@ struct SpillReaderStream { _spill_file: Arc, schema_validated: bool, + + /// Whether the file stores byte arrays as lengths, see [`lengths`] + decode_lengths: bool, } // Small margin allowed to accommodate slight memory accounting variation @@ -335,6 +339,7 @@ impl SpillReaderStream { pending: VecDeque::new(), _spill_file: spill_file, schema_validated: false, + decode_lengths: false, }) } } @@ -366,7 +371,11 @@ impl Stream for SpillReaderStream { if !this.schema_validated { this.schema_validated = true; let actual = batch.schema(); - if actual != this.schema { + // Compressed spill files store byte arrays as lengths + this.decode_lengths = actual != this.schema + && lengths::encode_schema(&this.schema) + .is_some_and(|encoded| encoded == actual); + if actual != this.schema && !this.decode_lengths { this.is_done = true; return Poll::Ready(Some(Err( datafusion_common::exec_datafusion_err!( @@ -379,6 +388,17 @@ impl Stream for SpillReaderStream { ))); } } + let batch = if this.decode_lengths { + match lengths::decode(&batch, &this.schema) { + Ok(batch) => batch, + Err(e) => { + this.is_done = true; + return Poll::Ready(Some(Err(e))); + } + } + } else { + batch + }; if let Some(max_record_batch_memory) = this.max_record_batch_memory { @@ -648,6 +668,8 @@ impl Drop for AsyncIPCStreamWriter { /// Encodes Arrow IPC streams without performing I/O. struct IPCStreamEncoder { encoder: Option, + /// Whether byte arrays are written as lengths, see [`lengths`] + encode_lengths: bool, } impl IPCStreamEncoder { @@ -677,14 +699,26 @@ impl IPCStreamEncoder { let compression_type = Option::::from(spill_compression); write_options = write_options.try_with_compression(compression_type)?; - let encoder = StreamEncoder::try_new_with_options(schema, write_options)?; + // Lengths take as much space as offsets when uncompressed + let encoded_schema = + compression_type.and_then(|_| lengths::encode_schema(schema)); + let encoder = StreamEncoder::try_new_with_options( + encoded_schema.as_deref().unwrap_or(schema), + write_options, + )?; Ok(Self { encoder: Some(encoder), + encode_lengths: encoded_schema.is_some(), }) } fn encode(&mut self, batch: &RecordBatch) -> Result<(usize, usize, Vec)> { - let buffers = self.encoder.as_mut().unwrap().encode(batch)?; + let encoder = self.encoder.as_mut().unwrap(); + let buffers = if self.encode_lengths { + encoder.encode(&lengths::encode(batch)?)? + } else { + encoder.encode(batch)? + }; let bytes = buffers.iter().map(Buffer::len).sum(); Ok((batch.num_rows(), bytes, coalesce_buffers(buffers))) } From 338baf953e60a9219feb17fed140120a95d82857 Mon Sep 17 00:00:00 2001 From: comphead Date: Sat, 31 Oct 2026 10:31:38 -0700 Subject: [PATCH 3/4] perf: buffer spill file writes and validate restored byte arrays - `FileSpillWriter` writes through a 256 KB `BufWriter` instead of the IPC encoder copying small buffers together. One bounded buffer is reused, large buffers are not copied, and uncompressed spills benefit too. - Byte arrays read back from compressed spill files are rebuilt with `OffsetBuffer::try_from_lengths` and `GenericByteArray::try_new`, so a corrupted file fails instead of producing invalid arrays. - `IPCStreamEncoder` keeps the encoded schema and reuses its struct fields for every batch. - `aggregate_memory_spill.slt` spills byte array keys with zstd and lz4. --- datafusion/execution/src/disk_manager.rs | 14 +- datafusion/physical-plan/src/spill/lengths.rs | 161 ++++++++---------- datafusion/physical-plan/src/spill/mod.rs | 60 ++----- .../test_files/aggregate_memory_spill.slt | 51 ++++++ 4 files changed, 149 insertions(+), 137 deletions(-) diff --git a/datafusion/execution/src/disk_manager.rs b/datafusion/execution/src/disk_manager.rs index d9a1d2f596937..3465402b6585c 100644 --- a/datafusion/execution/src/disk_manager.rs +++ b/datafusion/execution/src/disk_manager.rs @@ -14,6 +14,7 @@ use log::debug; use parking_lot::Mutex; use rand::{Rng, rng}; use std::fmt::Debug; +use std::io::{BufWriter, Write}; use std::path::{Path, PathBuf}; use std::sync::Arc; use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering}; @@ -21,6 +22,11 @@ use tempfile::{Builder, NamedTempFile, TempDir}; pub const DEFAULT_MAX_TEMP_DIRECTORY_SIZE: u64 = 100 * 1024 * 1024 * 1024; // 100GB pub const DEFAULT_MAX_SPILL_MERGE_FAN_IN: usize = 0; +/// Capacity of the buffer in front of a spill file. The IPC writer writes every +/// buffer and padding of a batch separately, which would otherwise take one +/// system call each. +const SPILL_WRITE_BUFFER_SIZE: usize = 256 * 1024; + /// Builder pattern for the [DiskManager] structure #[derive(Clone)] pub struct DiskManagerBuilder { @@ -469,12 +475,12 @@ fn create_local_dirs(local_dirs: &[PathBuf]) -> Result>> { } pub struct FileSpillWriter { - file: std::fs::File, + file: BufWriter, disk_manager: Arc, current_file_disk_usage: Arc, } -impl std::io::Write for FileSpillWriter { +impl Write for FileSpillWriter { fn write(&mut self, buf: &[u8]) -> std::io::Result { let len = buf.len() as u64; if len == 0 { @@ -516,7 +522,7 @@ impl std::io::Write for FileSpillWriter { impl SpillWriter for FileSpillWriter { fn finish(&mut self) -> Result<()> { - // flush() is handled by Arrow, nothing left to do here + self.file.flush()?; Ok(()) } } @@ -578,7 +584,7 @@ impl SpillFile for RefCountedTempFile { .try_clone() .map_err(DataFusionError::IoError)?; Ok(Box::new(FileSpillWriter { - file, + file: BufWriter::with_capacity(SPILL_WRITE_BUFFER_SIZE, file), disk_manager: Arc::clone(&self.disk_manager), current_file_disk_usage: Arc::clone(&self.current_file_disk_usage), })) diff --git a/datafusion/physical-plan/src/spill/lengths.rs b/datafusion/physical-plan/src/spill/lengths.rs index 888a71f7188d4..13e0f5fe298f6 100644 --- a/datafusion/physical-plan/src/spill/lengths.rs +++ b/datafusion/physical-plan/src/spill/lengths.rs @@ -36,45 +36,67 @@ use arrow::datatypes::{ GenericBinaryType, Int32Type, Int64Type, LargeBinaryType, LargeUtf8Type, Schema, SchemaRef, Utf8Type, }; -use datafusion_common::{Result, internal_err}; +use datafusion_common::utils::offset_span; +use datafusion_common::{Result, internal_datafusion_err, internal_err}; /// Returns the schema that `schema` is spilled as, or `None` if it has no byte /// array columns pub(super) fn encode_schema(schema: &Schema) -> Option { - if !schema - .fields() - .iter() - .any(|field| encoded_type(field.data_type()).is_some()) - { - return None; - } + let mut encoded = false; let fields: Vec<_> = schema .fields() .iter() .map(|field| match encoded_type(field.data_type()) { Some(data_type) => { + encoded = true; Arc::new(Field::new(field.name(), data_type, field.is_nullable())) } None => Arc::clone(field), }) .collect(); - Some(Arc::new(Schema::new_with_metadata( - fields, - schema.metadata().clone(), - ))) + encoded + .then(|| Arc::new(Schema::new_with_metadata(fields, schema.metadata().clone()))) } -/// Encodes the byte array columns of `batch`, see the [module docs](self) -pub(super) fn encode(batch: &RecordBatch) -> Result { - let Some(schema) = encode_schema(batch.schema_ref()) else { - return Ok(batch.clone()); - }; - let columns = batch - .columns() +/// Encodes the byte array columns of `batch` as [`encode_schema`] returned +/// `schema`, see the [module docs](self) +/// +/// The other columns keep the fields of `batch`, which can differ from the +/// spill schema in nested nullability, see `InProgressSpillFile::append_batch`. +pub(super) fn encode(batch: &RecordBatch, schema: &SchemaRef) -> Result { + let mut fields = Vec::with_capacity(batch.num_columns()); + let mut columns = Vec::with_capacity(batch.num_columns()); + for ((field, spilled), column) in batch + .schema_ref() + .fields() .iter() - .map(encode_array) - .collect::>()?; - Ok(RecordBatch::try_new(schema, columns)?) + .zip(schema.fields()) + .zip(batch.columns()) + { + let encoded = match (column.data_type(), spilled.data_type()) { + (DataType::Utf8, DataType::Struct(struct_fields)) => { + encode_bytes::<_, Int32Type>(struct_fields, column.as_string::())? + } + (DataType::LargeUtf8, DataType::Struct(struct_fields)) => { + encode_bytes::<_, Int64Type>(struct_fields, column.as_string::())? + } + (DataType::Binary, DataType::Struct(struct_fields)) => { + encode_bytes::<_, Int32Type>(struct_fields, column.as_binary::())? + } + (DataType::LargeBinary, DataType::Struct(struct_fields)) => { + encode_bytes::<_, Int64Type>(struct_fields, column.as_binary::())? + } + _ => { + fields.push(Arc::clone(field)); + columns.push(Arc::clone(column)); + continue; + } + }; + fields.push(Arc::clone(spilled)); + columns.push(encoded); + } + let schema = Schema::new_with_metadata(fields, batch.schema_ref().metadata().clone()); + Ok(RecordBatch::try_new(Arc::new(schema), columns)?) } /// Restores a batch that [`encode`] wrote for `schema` @@ -112,16 +134,6 @@ where ]) } -fn encode_array(array: &ArrayRef) -> Result { - match array.data_type() { - DataType::Utf8 => encode_bytes::<_, Int32Type>(array.as_string::()), - DataType::LargeUtf8 => encode_bytes::<_, Int64Type>(array.as_string::()), - DataType::Binary => encode_bytes::<_, Int32Type>(array.as_binary::()), - DataType::LargeBinary => encode_bytes::<_, Int64Type>(array.as_binary::()), - _ => Ok(Arc::clone(array)), - } -} - fn decode_array(data_type: &DataType, array: &ArrayRef) -> Result { match data_type { DataType::Utf8 => decode_bytes::(array), @@ -132,21 +144,19 @@ fn decode_array(data_type: &DataType, array: &ArrayRef) -> Result { } } -fn encode_bytes(array: &GenericByteArray) -> Result +fn encode_bytes(fields: &Fields, array: &GenericByteArray) -> Result where T: ByteArrayType, L: ArrowPrimitiveType, { - let fields = encoded_fields::(); if array.is_empty() { - return Ok(new_empty_array(&DataType::Struct(fields))); + return Ok(new_empty_array(&DataType::Struct(fields.clone()))); } let offsets = array.offsets(); let lengths = offsets.windows(2).map(|w| w[1] - w[0]).collect(); let lengths = PrimitiveArray::::new(lengths, array.nulls().cloned()); - let start = offsets[0].as_usize(); - let len = offsets[array.len()].as_usize() - start; + let (start, len) = offset_span(offsets); let bytes = GenericBinaryArray::::new( OffsetBuffer::from_lengths([len]), array.values().slice_with_length(start, len), @@ -155,8 +165,11 @@ where let run_ends = Int64Array::from(vec![array.len() as i64]); let values = RunArray::try_new(&run_ends, &bytes)?; - let array = - StructArray::try_new(fields, vec![Arc::new(lengths), Arc::new(values)], None)?; + let array = StructArray::try_new( + fields.clone(), + vec![Arc::new(lengths), Arc::new(values)], + None, + )?; Ok(Arc::new(array)) } @@ -167,36 +180,29 @@ where { let array = array.as_struct(); let lengths = array.column(0).as_primitive::(); - let mut end = T::Offset::usize_as(0); - let mut offsets = Vec::with_capacity(lengths.len() + 1); - offsets.push(end); - offsets.extend(lengths.values().iter().map(|&len| { - end += len; - end - })); + let offsets = OffsetBuffer::::try_from_lengths( + lengths.values().iter().map(|len| len.as_usize()), + ) + .map_err(|e| { + internal_datafusion_err!("Invalid spilled {} lengths: {e}", T::DATA_TYPE) + })?; let bytes = array.column(1).as_run::().values(); let bytes = bytes.as_binary::(); - let bytes_offsets = bytes.value_offsets(); - let start = bytes_offsets[0].as_usize(); - let len = bytes_offsets[bytes_offsets.len() - 1].as_usize() - start; - if end.as_usize() != len { + let (start, len) = offset_span(bytes.offsets()); + let end = offsets.last().as_usize(); + if end != len { return internal_err!( - "Spilled {} lengths add up to {}, but there are {len} value bytes", - T::DATA_TYPE, - end.as_usize() + "Spilled {} lengths add up to {end}, but there are {len} value bytes", + T::DATA_TYPE ); } - // SAFETY: `encode_bytes` took the lengths and value bytes from a valid - // array of type `T`, so their running sum gives back its offsets - let array = unsafe { - GenericByteArray::::new_unchecked( - OffsetBuffer::new_unchecked(offsets.into()), - bytes.values().slice_with_length(start, len), - lengths.nulls().cloned(), - ) - }; + let array = GenericByteArray::::try_new( + offsets, + bytes.values().slice_with_length(start, len), + lengths.nulls().cloned(), + )?; Ok(Arc::new(array)) } @@ -206,7 +212,7 @@ mod tests { use crate::common::collect; use crate::metrics::{ExecutionPlanMetricsSet, SpillMetrics}; use crate::spill::spill_manager::SpillManager; - use arrow::array::{Int32Array, ListArray, StringArray}; + use arrow::array::StringArray; use arrow::compute::cast; use arrow_ipc::CompressionType; use arrow_ipc::writer::{IpcWriteOptions, StreamWriter}; @@ -219,42 +225,20 @@ mod tests { SpillManager::new(env, metrics, schema).with_compression_type(compression) } - /// Byte arrays read back as written, sliced, empty or with nulls, next to - /// the columns that are spilled as they are + /// Sliced and empty byte arrays with nulls read back as written #[tokio::test] async fn test_spill_lengths_round_trip() -> Result<()> { let utf8: ArrayRef = Arc::new(StringArray::from_iter( (0..1000).map(|i| (i % 7 != 0).then(|| "x".repeat(i % 5))), )); - let list = ListArray::new( - Arc::new(Field::new("item", DataType::Utf8, true)), - OffsetBuffer::from_lengths([1; 1000]), - Arc::clone(&utf8), - None, - ); let batch = RecordBatch::try_from_iter(vec![ - ( - "int32", - Arc::new(Int32Array::from_iter_values(0..1000)) as _, - ), ("large_utf8", cast(&utf8, &DataType::LargeUtf8)?), ("binary", cast(&utf8, &DataType::Binary)?), ("large_binary", cast(&utf8, &DataType::LargeBinary)?), - ("utf8_view", cast(&utf8, &DataType::Utf8View)?), - ("list", Arc::new(list)), ("utf8", utf8), ])?; - let batches = [ - batch.clone(), - batch.slice(3, 500), - batch.slice(998, 2), - batch.slice(0, 0), - ]; - for compression in [ - SpillCompression::Uncompressed, - SpillCompression::Lz4Frame, - SpillCompression::Zstd, - ] { + let batches = [batch.slice(3, 500), batch.slice(998, 2), batch.slice(0, 0)]; + for compression in [SpillCompression::Lz4Frame, SpillCompression::Zstd] { let spill_manager = spill_manager(batch.schema(), compression); let spill_file = spill_manager .spill_record_batch_and_finish(&batches, "Test")? @@ -287,7 +271,6 @@ mod tests { let mut writer = StreamWriter::try_new_with_options(&mut ipc, &schema, options)?; writer.write(&batch)?; writer.finish()?; - drop(writer); assert!( spilled * 4 < ipc.len(), "spilled {spilled} bytes, against {} bytes of IPC", diff --git a/datafusion/physical-plan/src/spill/mod.rs b/datafusion/physical-plan/src/spill/mod.rs index 0ba6687020047..1e73dea06ade3 100644 --- a/datafusion/physical-plan/src/spill/mod.rs +++ b/datafusion/physical-plan/src/spill/mod.rs @@ -372,10 +372,11 @@ impl Stream for SpillReaderStream { this.schema_validated = true; let actual = batch.schema(); // Compressed spill files store byte arrays as lengths - this.decode_lengths = actual != this.schema - && lengths::encode_schema(&this.schema) - .is_some_and(|encoded| encoded == actual); - if actual != this.schema && !this.decode_lengths { + this.decode_lengths = actual != this.schema; + if this.decode_lengths + && !lengths::encode_schema(&this.schema) + .is_some_and(|encoded| encoded == actual) + { this.is_done = true; return Poll::Ready(Some(Err( datafusion_common::exec_datafusion_err!( @@ -668,8 +669,8 @@ impl Drop for AsyncIPCStreamWriter { /// Encodes Arrow IPC streams without performing I/O. struct IPCStreamEncoder { encoder: Option, - /// Whether byte arrays are written as lengths, see [`lengths`] - encode_lengths: bool, + /// The schema written when byte arrays are spilled as lengths, see [`lengths`] + lengths_schema: Option, } impl IPCStreamEncoder { @@ -700,61 +701,33 @@ impl IPCStreamEncoder { write_options = write_options.try_with_compression(compression_type)?; // Lengths take as much space as offsets when uncompressed - let encoded_schema = + let lengths_schema = compression_type.and_then(|_| lengths::encode_schema(schema)); let encoder = StreamEncoder::try_new_with_options( - encoded_schema.as_deref().unwrap_or(schema), + lengths_schema.as_deref().unwrap_or(schema), write_options, )?; Ok(Self { encoder: Some(encoder), - encode_lengths: encoded_schema.is_some(), + lengths_schema, }) } fn encode(&mut self, batch: &RecordBatch) -> Result<(usize, usize, Vec)> { let encoder = self.encoder.as_mut().unwrap(); - let buffers = if self.encode_lengths { - encoder.encode(&lengths::encode(batch)?)? - } else { - encoder.encode(batch)? + let buffers = match &self.lengths_schema { + Some(schema) => encoder.encode(&lengths::encode(batch, schema)?)?, + None => encoder.encode(batch)?, }; let bytes = buffers.iter().map(Buffer::len).sum(); - Ok((batch.num_rows(), bytes, coalesce_buffers(buffers))) + Ok((batch.num_rows(), bytes, buffers)) } fn finish(&mut self) -> Result<(usize, Vec)> { let buffers = self.encoder.take().unwrap().finish()?; let bytes = buffers.iter().map(Buffer::len).sum(); - Ok((bytes, coalesce_buffers(buffers))) - } -} - -/// Buffers at least this large are written on their own, see [`coalesce_buffers`] -const COALESCE_BUFFERS_BELOW: usize = 64 * 1024; - -/// Copies runs of small buffers into one, so that writing the IPC encoding of -/// a batch takes a few writes instead of one per buffer and padding. -/// -/// [`StreamEncoder`] returns every buffer of a batch and its padding separately, -/// and file backed spills would write each with its own system call. -fn coalesce_buffers(buffers: Vec) -> Vec { - let mut coalesced = Vec::new(); - let mut small: Vec = Vec::new(); - for buffer in buffers { - if buffer.len() < COALESCE_BUFFERS_BELOW { - small.extend_from_slice(buffer.as_slice()); - continue; - } - if !small.is_empty() { - coalesced.push(Buffer::from_vec(std::mem::take(&mut small))); - } - coalesced.push(buffer); - } - if !small.is_empty() { - coalesced.push(Buffer::from_vec(small)); + Ok((bytes, buffers)) } - coalesced } struct ArrowBufferOwner(Buffer); @@ -1193,8 +1166,7 @@ mod tests { let trailing_bytes = writer.finish().await?; assert!(finished.load(Ordering::Relaxed)); - // One write per batch and one for the end of the stream - assert_eq!(writes.load(Ordering::Relaxed), 3); + assert!(writes.load(Ordering::Relaxed) > 2); drop(writer); assert!(!aborted.load(Ordering::Relaxed)); diff --git a/datafusion/sqllogictest/test_files/aggregate_memory_spill.slt b/datafusion/sqllogictest/test_files/aggregate_memory_spill.slt index dc6af9b1e4e47..14a5f093a9cad 100644 --- a/datafusion/sqllogictest/test_files/aggregate_memory_spill.slt +++ b/datafusion/sqllogictest/test_files/aggregate_memory_spill.slt @@ -247,6 +247,57 @@ FROM ( ---- 8192 8192 +# --- Case I: byte array keys in compressed spill files --- +# Compressed spill files store `Utf8`, `LargeUtf8`, `Binary` and `LargeBinary` +# columns as lengths. The keys are cast because `varchar` maps to `Utf8View`. +statement ok +CREATE VIEW byte_keys AS +SELECT + arrow_cast(CASE WHEN v % 10 = 0 THEN NULL ELSE concat('k', (v * 7) % 100000) END, 'Utf8') AS s, + arrow_cast(concat('k', (v * 7) % 100000), 'LargeUtf8') AS ls, + arrow_cast(concat('k', (v * 7) % 100000), 'Binary') AS b, + arrow_cast(concat('k', (v * 7) % 100000), 'LargeBinary') AS lb +FROM generate_series(1, 100000) AS t(v) + +statement ok +SET datafusion.execution.spill_compression = 'zstd' + +query IITTII +SELECT count(*), count(s), min(s), max(ls), sum(length(s)), count(*) FILTER ( + WHERE arrow_cast(b, 'LargeUtf8') = ls AND arrow_cast(lb, 'LargeUtf8') = ls + AND (s IS NULL OR s = ls) +) +FROM (SELECT s, ls, b, lb FROM byte_keys GROUP BY s, ls, b, lb) +---- +100000 90000 k1 k99999 530001 100000 + +statement ok +SET datafusion.execution.spill_compression = 'lz4_frame' + +query IITTII +SELECT count(*), count(s), min(s), max(ls), sum(length(s)), count(*) FILTER ( + WHERE arrow_cast(b, 'LargeUtf8') = ls AND arrow_cast(lb, 'LargeUtf8') = ls + AND (s IS NULL OR s = ls) +) +FROM (SELECT s, ls, b, lb FROM byte_keys GROUP BY s, ls, b, lb) +---- +100000 90000 k1 k99999 530001 100000 + +# Assert spill happened +query TT +EXPLAIN ANALYZE +SELECT count(*) FROM (SELECT s, ls, b, lb FROM byte_keys GROUP BY s, ls, b, lb) +---- + +04)------AggregateExec: mode=Single, gby=[s@0 as s, ls@1 as ls, b@2 as b, lb@3 as lb], aggr=[], metrics=[spill_count=, spilled_bytes=, spilled_rows= K,] + + +statement ok +RESET datafusion.execution.spill_compression + +statement ok +DROP VIEW byte_keys + statement ok SET datafusion.execution.target_partitions = 4 From 03a6c93f4ba9923c4d9bc923e85db3b60e41a25f Mon Sep 17 00:00:00 2001 From: comphead Date: Tue, 3 Nov 2026 08:30:01 -0800 Subject: [PATCH 4/4] perf: coalesce spill writes per batch, skip lengths for small batches - Copy the small IPC buffers of each batch into one write again, instead of keeping a 256 KB `BufWriter` per spill file. Those buffers were outside the memory pool, and a preserve-order `RepartitionExec` keeps one spill file open per input and output partition. - Write lengths only when the first batch of a file has at least 128 rows. Below that, the struct and run-end metadata cost more than the lengths save. --- datafusion/execution/src/disk_manager.rs | 14 ++--- .../src/spill/in_progress_spill_file.rs | 2 + datafusion/physical-plan/src/spill/lengths.rs | 28 ++++++++- datafusion/physical-plan/src/spill/mod.rs | 63 ++++++++++++++++--- .../test_files/aggregate_memory_spill.slt | 6 +- 5 files changed, 91 insertions(+), 22 deletions(-) diff --git a/datafusion/execution/src/disk_manager.rs b/datafusion/execution/src/disk_manager.rs index 3465402b6585c..d9a1d2f596937 100644 --- a/datafusion/execution/src/disk_manager.rs +++ b/datafusion/execution/src/disk_manager.rs @@ -14,7 +14,6 @@ use log::debug; use parking_lot::Mutex; use rand::{Rng, rng}; use std::fmt::Debug; -use std::io::{BufWriter, Write}; use std::path::{Path, PathBuf}; use std::sync::Arc; use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering}; @@ -22,11 +21,6 @@ use tempfile::{Builder, NamedTempFile, TempDir}; pub const DEFAULT_MAX_TEMP_DIRECTORY_SIZE: u64 = 100 * 1024 * 1024 * 1024; // 100GB pub const DEFAULT_MAX_SPILL_MERGE_FAN_IN: usize = 0; -/// Capacity of the buffer in front of a spill file. The IPC writer writes every -/// buffer and padding of a batch separately, which would otherwise take one -/// system call each. -const SPILL_WRITE_BUFFER_SIZE: usize = 256 * 1024; - /// Builder pattern for the [DiskManager] structure #[derive(Clone)] pub struct DiskManagerBuilder { @@ -475,12 +469,12 @@ fn create_local_dirs(local_dirs: &[PathBuf]) -> Result>> { } pub struct FileSpillWriter { - file: BufWriter, + file: std::fs::File, disk_manager: Arc, current_file_disk_usage: Arc, } -impl Write for FileSpillWriter { +impl std::io::Write for FileSpillWriter { fn write(&mut self, buf: &[u8]) -> std::io::Result { let len = buf.len() as u64; if len == 0 { @@ -522,7 +516,7 @@ impl Write for FileSpillWriter { impl SpillWriter for FileSpillWriter { fn finish(&mut self) -> Result<()> { - self.file.flush()?; + // flush() is handled by Arrow, nothing left to do here Ok(()) } } @@ -584,7 +578,7 @@ impl SpillFile for RefCountedTempFile { .try_clone() .map_err(DataFusionError::IoError)?; Ok(Box::new(FileSpillWriter { - file: BufWriter::with_capacity(SPILL_WRITE_BUFFER_SIZE, file), + file, disk_manager: Arc::clone(&self.disk_manager), current_file_disk_usage: Arc::clone(&self.current_file_disk_usage), })) 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 0f83ef1a4b441..f91fba0e54641 100644 --- a/datafusion/physical-plan/src/spill/in_progress_spill_file.rs +++ b/datafusion/physical-plan/src/spill/in_progress_spill_file.rs @@ -78,6 +78,7 @@ impl InProgressSpillFile { spill_writer, schema.as_ref(), self.spill_manager.compression, + gc_batch.num_rows(), )?)); // Update metrics @@ -116,6 +117,7 @@ impl InProgressSpillFile { let encoder = IPCStreamEncoder::new( schema.as_ref(), self.spill_manager.compression, + gc_batch.num_rows(), )?; let spill_writer = in_progress_file.open_async_writer().await?; diff --git a/datafusion/physical-plan/src/spill/lengths.rs b/datafusion/physical-plan/src/spill/lengths.rs index 13e0f5fe298f6..142296d6cddfa 100644 --- a/datafusion/physical-plan/src/spill/lengths.rs +++ b/datafusion/physical-plan/src/spill/lengths.rs @@ -21,7 +21,9 @@ //! ``` //! //! Both are valid arrays, so the IPC reader decodes them as usual before -//! [`decode`] rebuilds the offsets with a running sum. +//! [`decode`] rebuilds the offsets with a running sum. Files whose first batch +//! has fewer than [`MIN_LENGTHS_ROWS`](super::MIN_LENGTHS_ROWS) rows keep their +//! offsets. use std::sync::Arc; @@ -278,4 +280,28 @@ mod tests { ); Ok(()) } + + /// Files whose first batch is too small for the lengths to pay off keep + /// their offsets + #[test] + fn test_spill_lengths_skipped_for_small_batches() -> Result<()> { + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Utf8, false)])); + for (rows, encoded) in [(127, false), (128, true)] { + let values = (0..rows).map(|i| format!("item-{}", i / 64)); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(StringArray::from_iter_values(values))], + )?; + let spill_manager = + spill_manager(Arc::clone(&schema), SpillCompression::Zstd); + let spill_file = spill_manager + .spill_record_batch_and_finish(std::slice::from_ref(&batch), "Test")? + .unwrap(); + let file = std::fs::File::open(spill_file.path().unwrap())?; + let reader = arrow_ipc::reader::StreamReader::try_new(file, None)?; + let spilled_type = reader.schema().field(0).data_type().clone(); + assert_eq!(spilled_type != DataType::Utf8, encoded, "{rows} rows"); + } + Ok(()) + } } diff --git a/datafusion/physical-plan/src/spill/mod.rs b/datafusion/physical-plan/src/spill/mod.rs index 1e73dea06ade3..a15584892c628 100644 --- a/datafusion/physical-plan/src/spill/mod.rs +++ b/datafusion/physical-plan/src/spill/mod.rs @@ -486,10 +486,11 @@ impl IPCStreamWriter { spill_writer: Box, schema: &Schema, spill_compression: SpillCompression, + first_batch_rows: usize, ) -> Result { Ok(Self { writer: spill_writer, - encoder: IPCStreamEncoder::new(schema, spill_compression)?, + encoder: IPCStreamEncoder::new(schema, spill_compression, first_batch_rows)?, }) } @@ -666,6 +667,10 @@ impl Drop for AsyncIPCStreamWriter { } } +/// Fewer rows per batch than this and the struct and run-end metadata of a +/// lengths column outweighs what the lengths save, see [`lengths`] +const MIN_LENGTHS_ROWS: usize = 128; + /// Encodes Arrow IPC streams without performing I/O. struct IPCStreamEncoder { encoder: Option, @@ -686,7 +691,11 @@ impl IPCStreamEncoder { /// contract local and build-visible during Cargo feature resolution, /// rather than relying solely on workspace-level feature unification; /// see #21917. - fn new(schema: &Schema, spill_compression: SpillCompression) -> Result { + fn new( + schema: &Schema, + spill_compression: SpillCompression, + first_batch_rows: usize, + ) -> Result { let metadata_version = MetadataVersion::V5; // Depending on the schema, some array types such as StringViewArray require larger (16 byte in this case) alignment. // If the actual buffer layout after IPC read does not satisfy the alignment requirement, @@ -701,8 +710,9 @@ impl IPCStreamEncoder { write_options = write_options.try_with_compression(compression_type)?; // Lengths take as much space as offsets when uncompressed - let lengths_schema = - compression_type.and_then(|_| lengths::encode_schema(schema)); + let lengths_schema = compression_type + .filter(|_| first_batch_rows >= MIN_LENGTHS_ROWS) + .and_then(|_| lengths::encode_schema(schema)); let encoder = StreamEncoder::try_new_with_options( lengths_schema.as_deref().unwrap_or(schema), write_options, @@ -720,14 +730,39 @@ impl IPCStreamEncoder { None => encoder.encode(batch)?, }; let bytes = buffers.iter().map(Buffer::len).sum(); - Ok((batch.num_rows(), bytes, buffers)) + Ok((batch.num_rows(), bytes, coalesce_buffers(buffers))) } fn finish(&mut self) -> Result<(usize, Vec)> { let buffers = self.encoder.take().unwrap().finish()?; let bytes = buffers.iter().map(Buffer::len).sum(); - Ok((bytes, buffers)) + Ok((bytes, coalesce_buffers(buffers))) + } +} + +/// Buffers at least this large are written on their own, see [`coalesce_buffers`] +const COALESCE_BUFFERS_BELOW: usize = 64 * 1024; + +/// Copies runs of small buffers into one, so that writing a batch takes a few +/// writes instead of one per IPC buffer and padding. The copy lives only for +/// the write, unlike a buffer kept by every open spill file. +fn coalesce_buffers(buffers: Vec) -> Vec { + let mut coalesced = Vec::new(); + let mut small: Vec = Vec::new(); + for buffer in buffers { + if buffer.len() < COALESCE_BUFFERS_BELOW { + small.extend_from_slice(buffer.as_slice()); + continue; + } + if !small.is_empty() { + coalesced.push(Buffer::from_vec(std::mem::take(&mut small))); + } + coalesced.push(buffer); } + if !small.is_empty() { + coalesced.push(Buffer::from_vec(small)); + } + coalesced } struct ArrowBufferOwner(Buffer); @@ -1157,8 +1192,11 @@ mod tests { finished: Arc::clone(&finished), aborted: Arc::clone(&aborted), }; - let encoder = - IPCStreamEncoder::new(schema.as_ref(), SpillCompression::Uncompressed)?; + let encoder = IPCStreamEncoder::new( + schema.as_ref(), + SpillCompression::Uncompressed, + usize::MAX, + )?; let mut writer = AsyncIPCStreamWriter::new(Box::new(spill_writer), encoder); let (_, first_bytes) = writer.write(&batch1).await?; @@ -1166,7 +1204,8 @@ mod tests { let trailing_bytes = writer.finish().await?; assert!(finished.load(Ordering::Relaxed)); - assert!(writes.load(Ordering::Relaxed) > 2); + // One write per batch and one for the end of the stream + assert_eq!(writes.load(Ordering::Relaxed), 3); drop(writer); assert!(!aborted.load(Ordering::Relaxed)); @@ -1199,6 +1238,7 @@ mod tests { let encoder = IPCStreamEncoder::new( batch.schema().as_ref(), SpillCompression::Uncompressed, + usize::MAX, )?; let mut writer = AsyncIPCStreamWriter::new( Box::new(FailingAsyncWriter { @@ -1233,6 +1273,7 @@ mod tests { let encoder = IPCStreamEncoder::new( batch.schema().as_ref(), SpillCompression::Uncompressed, + usize::MAX, )?; let mut writer = AsyncIPCStreamWriter::new( Box::new(PendingFinishAsyncWriter { @@ -1270,7 +1311,8 @@ mod tests { let schema = Schema::new(vec![Field::new("a", DataType::Int32, false)]); let abort_count = Arc::new(AtomicUsize::new(0)); let cleanup_done = Arc::new(Notify::new()); - let encoder = IPCStreamEncoder::new(&schema, SpillCompression::Uncompressed)?; + let encoder = + IPCStreamEncoder::new(&schema, SpillCompression::Uncompressed, usize::MAX)?; let mut writer = AsyncIPCStreamWriter::new( Box::new(FailOnceAbortWriter { abort_count: Arc::clone(&abort_count), @@ -1312,6 +1354,7 @@ mod tests { let encoder = IPCStreamEncoder::new( batch.schema().as_ref(), SpillCompression::Uncompressed, + usize::MAX, )?; AsyncIPCStreamWriter::new(Box::new(spill_writer), encoder) }; diff --git a/datafusion/sqllogictest/test_files/aggregate_memory_spill.slt b/datafusion/sqllogictest/test_files/aggregate_memory_spill.slt index 14a5f093a9cad..8178a8388a117 100644 --- a/datafusion/sqllogictest/test_files/aggregate_memory_spill.slt +++ b/datafusion/sqllogictest/test_files/aggregate_memory_spill.slt @@ -249,7 +249,11 @@ FROM ( # --- Case I: byte array keys in compressed spill files --- # Compressed spill files store `Utf8`, `LargeUtf8`, `Binary` and `LargeBinary` -# columns as lengths. The keys are cast because `varchar` maps to `Utf8View`. +# columns as lengths, unless their batches are too small to benefit. The keys +# are cast because `varchar` maps to `Utf8View`. +statement ok +SET datafusion.execution.batch_size = 8192 + statement ok CREATE VIEW byte_keys AS SELECT