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 new file mode 100644 index 0000000000000..142296d6cddfa --- /dev/null +++ b/datafusion/physical-plan/src/spill/lengths.rs @@ -0,0 +1,307 @@ +// 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. Files whose first batch +//! has fewer than [`MIN_LENGTHS_ROWS`](super::MIN_LENGTHS_ROWS) rows keep their +//! offsets. + +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::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 { + 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(); + encoded + .then(|| Arc::new(Schema::new_with_metadata(fields, schema.metadata().clone()))) +} + +/// 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() + .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` +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 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(fields: &Fields, array: &GenericByteArray) -> Result +where + T: ByteArrayType, + L: ArrowPrimitiveType, +{ + if array.is_empty() { + 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, len) = offset_span(offsets); + 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.clone(), + 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 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 (start, len) = offset_span(bytes.offsets()); + let end = offsets.last().as_usize(); + if end != len { + return internal_err!( + "Spilled {} lengths add up to {end}, but there are {len} value bytes", + T::DATA_TYPE + ); + } + + let array = GenericByteArray::::try_new( + offsets, + 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::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) + } + + /// 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 batch = RecordBatch::try_from_iter(vec![ + ("large_utf8", cast(&utf8, &DataType::LargeUtf8)?), + ("binary", cast(&utf8, &DataType::Binary)?), + ("large_binary", cast(&utf8, &DataType::LargeBinary)?), + ("utf8", utf8), + ])?; + 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")? + .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()?; + assert!( + spilled * 4 < ipc.len(), + "spilled {spilled} bytes, against {} bytes of IPC", + ipc.len() + ); + 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 842105c99da6b..15fe7d851085d 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,12 @@ 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; + 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!( @@ -379,6 +389,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 { @@ -465,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)?, }) } @@ -645,9 +667,15 @@ 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, + /// The schema written when byte arrays are spilled as lengths, see [`lengths`] + lengths_schema: Option, } impl IPCStreamEncoder { @@ -663,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, @@ -677,25 +709,62 @@ 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 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, + )?; Ok(Self { encoder: Some(encoder), + lengths_schema, }) } 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 = 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, 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); impl AsRef<[u8]> for ArrowBufferOwner { @@ -1123,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?; @@ -1132,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)); @@ -1165,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 { @@ -1199,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 { @@ -1236,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), @@ -1278,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 dc6af9b1e4e47..8178a8388a117 100644 --- a/datafusion/sqllogictest/test_files/aggregate_memory_spill.slt +++ b/datafusion/sqllogictest/test_files/aggregate_memory_spill.slt @@ -247,6 +247,61 @@ FROM ( ---- 8192 8192 +# --- Case I: byte array keys in compressed spill files --- +# Compressed spill files store `Utf8`, `LargeUtf8`, `Binary` and `LargeBinary` +# 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 + 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