From b7288bbbb58eea0b34b571b7e76bc013a7355e14 Mon Sep 17 00:00:00 2001 From: comphead Date: Fri, 23 Oct 2026 09:33:03 -0700 Subject: [PATCH 1/2] fix: skip validating spilled record batches again #15454 stopped re-validating spill files on read. #21882 moved the reader to `StreamDecoder`, whose `with_skip_validation` only applies to dictionary batches, so every offset and string of a spilled record batch is validated again. Decode the framed messages with `RecordBatchDecoder` instead, which honors the flag. --- datafusion/physical-plan/src/spill/mod.rs | 179 ++++++++++++++++------ 1 file changed, 133 insertions(+), 46 deletions(-) diff --git a/datafusion/physical-plan/src/spill/mod.rs b/datafusion/physical-plan/src/spill/mod.rs index 9eaa04389c1b1..92911833b0544 100644 --- a/datafusion/physical-plan/src/spill/mod.rs +++ b/datafusion/physical-plan/src/spill/mod.rs @@ -28,7 +28,7 @@ pub use datafusion_common::utils::memory::get_record_batch_memory_size; #[doc(hidden)] pub use spill_manager::SpillManager; -use std::collections::{HashMap, VecDeque}; +use std::collections::HashMap; use std::pin::Pin; use std::sync::{Arc, LazyLock}; use std::task::{Context, Poll}; @@ -42,11 +42,12 @@ use arrow::datatypes::DataType; use arrow::datatypes::{ByteViewType, Schema, SchemaRef}; use arrow::ipc::{ MetadataVersion, - reader::StreamDecoder, + convert::try_fb_to_schema, + reader::{RecordBatchDecoder, read_dictionary_impl}, writer::{IpcWriteOptions, StreamEncoder}, }; use arrow::record_batch::RecordBatch; -use arrow_data::{ArrayDataBuilder, ByteView, MAX_INLINE_VIEW_LEN}; +use arrow_data::{ArrayDataBuilder, ByteView, MAX_INLINE_VIEW_LEN, UnsafeFlag}; #[cfg(test)] use arrow_ipc::writer::StreamWriter; use arrow_ipc::{CompressionType, root_as_message}; @@ -59,13 +60,13 @@ use futures::Stream; use log::debug; /// Stream that reads spill files from a [`SpillFile`] backend as a stream of [`RecordBatch`]es. -/// Uses [`StreamDecoder`] to decode IPC bytes received from the backend's async byte stream. +/// Uses [`MessageFramer`] and [`MessageDecoder`] to decode IPC bytes received from the backend's async byte stream. /// Backends handle their own threading concerns internally - OS files use /// `tokio::fs::File` which performs blocking IO per-syscall without holding a thread /// for the file's lifetime, avoiding deadlocks when concurrent reads exceed thread pool limits. struct SpillReaderStream { schema: SchemaRef, - decoder: StreamDecoder, + decoder: MessageDecoder, byte_stream: Pin> + Send>>, is_done: bool, @@ -81,9 +82,6 @@ struct SpillReaderStream { /// Assembles the chunks of `byte_stream` into exactly sized messages framer: MessageFramer, - /// Framed message buffers not yet consumed by `decoder` - pending: VecDeque, - /// Keeps the file alive until the stream is dropped _spill_file: Arc, @@ -94,7 +92,7 @@ struct SpillReaderStream { const SPILL_BATCH_MEMORY_MARGIN: usize = 4096; /// Reassembles an IPC stream read in arbitrary chunks into one exactly sized -/// allocation per message, so that the zero-copy [`StreamDecoder`] produces +/// allocation per message, so that the zero-copy [`MessageDecoder`] produces /// batches whose buffers pin only their own message. /// /// # Why @@ -295,10 +293,11 @@ impl MessageFramer { /// /// This mirrors [`StreamDecoder::finish`]: a stream is complete either /// when the end-of-stream marker was read, or when it ends right after a - /// message. Framing moves the partial bytes of a truncated file out of - /// the decoder's own scratch space and into the framer, so without this - /// check a truncated spill file would silently decode as a shorter - /// stream instead of erroring. + /// message. Only the framer holds the partial bytes of a truncated file, + /// so without this check a truncated spill file would silently decode as + /// a shorter stream instead of erroring. + /// + /// [`StreamDecoder::finish`]: arrow::ipc::reader::StreamDecoder::finish fn finish(&self) -> Result<()> { match &self.state { FramerState::Finished @@ -327,6 +326,81 @@ impl FramerState { /// Marks a length prefix in the IPC stream format, see `arrow_ipc`. const CONTINUATION_MARKER: [u8; 4] = [0xff; 4]; +/// Decodes the messages framed by [`MessageFramer`] without validating them. +/// +/// DataFusion controls what it writes so it can trust its own IPC output. +/// [`StreamDecoder::with_skip_validation`] would only skip validating the +/// dictionaries, and check every offset and string of the record batches. +/// +/// [`StreamDecoder::with_skip_validation`]: arrow::ipc::reader::StreamDecoder::with_skip_validation +#[derive(Default)] +struct MessageDecoder { + /// The schema of the stream, from its first message + schema: Option, + dictionaries: HashMap, +} + +impl MessageDecoder { + /// Decodes a message framed into `buffers`, returning its batch if it is a + /// record batch. + fn decode(&mut self, buffers: &[Buffer]) -> Result> { + let head = &buffers[0]; + let empty = Buffer::from(&[]); + let body = buffers.get(1).unwrap_or(&empty); + let prefix_len = if head.starts_with(&CONTINUATION_MARKER) { + 8 + } else { + 4 + }; + if head.len() == prefix_len { + // The end-of-stream marker + return Ok(None); + } + let message = root_as_message(&head[prefix_len..]).map_err(|e| { + datafusion_common::exec_datafusion_err!( + "Invalid IPC message in spill file: {e}" + ) + })?; + let version = message.version(); + let mut skip_validation = UnsafeFlag::new(); + // SAFETY: spill files are only read by the process that wrote them + unsafe { skip_validation.set(true) }; + + if let Some(schema) = message.header_as_schema() { + self.schema = Some(Arc::new(try_fb_to_schema(schema)?)); + return Ok(None); + } + let schema = self.schema.as_ref().ok_or_else(|| { + datafusion_common::exec_datafusion_err!("Missing schema in spill file") + })?; + if let Some(dictionary) = message.header_as_dictionary_batch() { + read_dictionary_impl( + body, + dictionary, + schema, + &mut self.dictionaries, + &version, + false, + skip_validation, + )?; + return Ok(None); + } + let Some(batch) = message.header_as_record_batch() else { + return Ok(None); + }; + let batch = RecordBatchDecoder::try_new( + body, + batch, + Arc::clone(schema), + &self.dictionaries, + &version, + )? + .with_skip_validation(skip_validation) + .read_record_batch()?; + Ok(Some(batch)) + } +} + impl SpillReaderStream { fn new( schema: SchemaRef, @@ -334,18 +408,14 @@ impl SpillReaderStream { max_record_batch_memory: Option, ) -> Result { let byte_stream = spill_file.read_stream()?; - // DataFusion controls what it writes so it can trust its own IPC output, - // matching the behavior of the previous StreamReader-based implementation. - let decoder = unsafe { StreamDecoder::new().with_skip_validation(true) }; Ok(Self { schema, - decoder, + decoder: MessageDecoder::default(), byte_stream, max_record_batch_memory, is_done: false, current_buffer: Buffer::from(&[]), framer: MessageFramer::new(), - pending: VecDeque::new(), _spill_file: spill_file, schema_validated: false, }) @@ -363,17 +433,18 @@ impl Stream for SpillReaderStream { } loop { - // Decode the framed messages first - if let Some(buffer) = this.pending.front_mut() { - if buffer.is_empty() { - this.pending.pop_front(); - continue; - } - match this.decoder.decode(buffer) { + // Frame the next message out of the current chunk and decode it + if !this.current_buffer.is_empty() { + let decoded = match this.framer.push(&mut this.current_buffer) { + Ok(Some(buffers)) => this.decoder.decode(&buffers), + Ok(None) => Ok(None), + Err(e) => Err(e), + }; + match decoded { Ok(Some(batch)) => { // One-time schema validation on the first decoded batch. // The IPC stream embeds the writer's schema in its header; - // StreamDecoder surfaces it via the first batch's schema. + // the decoder surfaces it via the first batch's schema. // We check here rather than in new() because schema bytes // only arrive after decoding the IPC header from the stream. if !this.schema_validated { @@ -409,22 +480,9 @@ impl Stream for SpillReaderStream { return Poll::Ready(Some(Ok(batch))); } Ok(None) => { - // A schema or dictionary message, or the buffer was - // only part of a message. Carry on with the next one. + // A schema or dictionary message, or the chunk ended + // inside a message. Carry on with the next one. } - Err(e) => { - this.is_done = true; - return Poll::Ready(Some(Err(e.into()))); - } - } - continue; - } - - // Then frame the next message out of the current chunk - if !this.current_buffer.is_empty() { - match this.framer.push(&mut this.current_buffer) { - Ok(Some(buffers)) => this.pending.extend(buffers), - Ok(None) => {} Err(e) => { this.is_done = true; return Poll::Ready(Some(Err(e))); @@ -433,7 +491,7 @@ impl Stream for SpillReaderStream { continue; } - // Finally fetch another chunk + // Then fetch another chunk match futures::ready!(this.byte_stream.as_mut().poll_next(cx)) { Some(Ok(chunk)) => { this.current_buffer = Buffer::from(chunk); @@ -450,9 +508,6 @@ impl Stream for SpillReaderStream { if let Err(e) = this.framer.finish() { return Poll::Ready(Some(Err(e))); } - if let Err(e) = this.decoder.finish() { - return Poll::Ready(Some(Err(e.into()))); - } return Poll::Ready(None); } } @@ -1008,7 +1063,7 @@ mod tests { use arrow::array::{ArrayRef, Int32Array, StringArray}; use arrow::compute::cast; use arrow::datatypes::{DataType, Field}; - use arrow::ipc::reader::StreamReader; + use arrow::ipc::reader::{StreamDecoder, StreamReader}; use datafusion_execution::runtime_env::RuntimeEnv; use futures::StreamExt as _; use std::io::Cursor; @@ -1705,6 +1760,38 @@ mod tests { Ok(()) } + /// Spill files are read back without validating them again, see + /// [`MessageDecoder`]. Validation would reject these bytes, which are not + /// UTF-8. + #[tokio::test] + async fn test_spill_read_skips_validation() -> Result<()> { + use arrow::array::BinaryArray; + + let binary = BinaryArray::from_iter_values([b"\xa0\xa1".as_slice(), b"valid"]); + // SAFETY: the test only compares the bytes and never reads them as `str` + let strings = unsafe { + StringArray::new_unchecked( + binary.offsets().clone(), + binary.values().clone(), + None, + ) + }; + let batch = RecordBatch::try_from_iter([("s", Arc::new(strings) as ArrayRef)])?; + + let env = Arc::new(RuntimeEnv::default()); + let metrics = SpillMetrics::new(&ExecutionPlanMetricsSet::new(), 0); + let spill_manager = SpillManager::new(env, metrics, batch.schema()); + let spill_file = spill_manager + .spill_record_batch_and_finish(std::slice::from_ref(&batch), "Test")? + .unwrap(); + let stream = spill_manager.read_spill_as_stream(spill_file, None)?; + let read_back = collect(stream).await?; + assert_eq!(read_back.len(), 1); + // `ArrayData` formats the bytes, not the strings, if this fails + assert_eq!(read_back[0].column(0).to_data(), batch.column(0).to_data()); + Ok(()) + } + #[tokio::test] async fn test_batch_spill_and_read_dictionary_arrays() -> Result<()> { // See https://github.com/apache/datafusion/issues/4658 From e113431e28fc54f53dc0b880217be3d33dd4c790 Mon Sep 17 00:00:00 2001 From: comphead Date: Fri, 23 Oct 2026 09:40:41 -0700 Subject: [PATCH 2/2] perf: write offsets to compressed spill files as lengths 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 hold the lengths of `Utf8`, `Binary`, `List` and `Map` arrays, their large variants included, at any nesting depth through `Struct` and `FixedSizeList`. The reader rebuilds the offsets with a running sum, in place in the decompressed buffer. --- datafusion/physical-plan/src/spill/mod.rs | 179 +++++++++++- .../physical-plan/src/spill/offset_lengths.rs | 264 ++++++++++++++++++ 2 files changed, 440 insertions(+), 3 deletions(-) create mode 100644 datafusion/physical-plan/src/spill/offset_lengths.rs diff --git a/datafusion/physical-plan/src/spill/mod.rs b/datafusion/physical-plan/src/spill/mod.rs index 92911833b0544..9a5a5197fd3f7 100644 --- a/datafusion/physical-plan/src/spill/mod.rs +++ b/datafusion/physical-plan/src/spill/mod.rs @@ -18,6 +18,7 @@ //! Defines the spilling functions pub(crate) mod in_progress_spill_file; +mod offset_lengths; pub(crate) mod replayable_spill_input; pub(crate) mod spill_manager; pub mod spill_pool; @@ -326,7 +327,9 @@ impl FramerState { /// Marks a length prefix in the IPC stream format, see `arrow_ipc`. const CONTINUATION_MARKER: [u8; 4] = [0xff; 4]; -/// Decodes the messages framed by [`MessageFramer`] without validating them. +/// Decodes the messages framed by [`MessageFramer`] without validating them, +/// and restores the offsets of compressed record batches, see +/// [`offset_lengths`]. /// /// DataFusion controls what it writes so it can trust its own IPC output. /// [`StreamDecoder::with_skip_validation`] would only skip validating the @@ -388,6 +391,7 @@ impl MessageDecoder { let Some(batch) = message.header_as_record_batch() else { return Ok(None); }; + let compressed = batch.compression().is_some(); let batch = RecordBatchDecoder::try_new( body, batch, @@ -397,7 +401,11 @@ impl MessageDecoder { )? .with_skip_validation(skip_validation) .read_record_batch()?; - Ok(Some(batch)) + if compressed { + offset_lengths::lengths_to_offsets(batch).map(Some) + } else { + Ok(Some(batch)) + } } } @@ -716,6 +724,8 @@ impl Drop for AsyncIPCStreamWriter { /// Encodes Arrow IPC streams without performing I/O. struct IPCStreamEncoder { encoder: Option, + /// Whether offsets are written as lengths, see [`offset_lengths`] + offsets_as_lengths: bool, } impl IPCStreamEncoder { @@ -745,14 +755,28 @@ impl IPCStreamEncoder { let compression_type = Option::::from(spill_compression); write_options = write_options.try_with_compression(compression_type)?; + // Lengths compress better than offsets, but take the same space + // uncompressed, and the reader only restores them for compressed batches + let offsets_as_lengths = compression_type.is_some() + && schema + .fields() + .iter() + .any(|field| offset_lengths::has_offsets(field.data_type())); + let encoder = StreamEncoder::try_new_with_options(schema, write_options)?; Ok(Self { encoder: Some(encoder), + offsets_as_lengths, }) } 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.offsets_as_lengths { + encoder.encode(&offset_lengths::offsets_to_lengths(batch)?)? + } else { + encoder.encode(batch)? + }; let bytes = buffers.iter().map(Buffer::len).sum(); Ok((batch.num_rows(), bytes, buffers)) } @@ -1995,6 +2019,155 @@ mod tests { Ok(()) } + /// A batch with every type whose offsets are spilled as lengths, at the + /// top level and nested, next to types that keep their offsets. + fn build_offsets_batch(n: usize) -> Result { + use arrow::array::{ + FixedSizeListArray, LargeListArray, ListArray, MapArray, StructArray, + }; + use arrow::buffer::{NullBuffer, OffsetBuffer}; + + let strings = |len: usize| -> ArrayRef { + Arc::new(StringArray::from_iter( + (0..len).map(|i| (i % 7 != 0).then(|| "x".repeat(i % 5))), + )) + }; + let nulls = Some(NullBuffer::from_iter((0..n).map(|i| i % 11 != 0))); + let offsets = OffsetBuffer::::from_lengths((0..n).map(|i| i % 3)); + let values = offsets.last() as usize; + let item = Arc::new(Field::new("item", DataType::Utf8, true)); + + let list = ListArray::new( + Arc::clone(&item), + offsets.clone(), + strings(values), + nulls.clone(), + ); + let large_list = LargeListArray::new( + Arc::new(Field::new("item", DataType::Int32, true)), + OffsetBuffer::::from_lengths((0..n).map(|i| i % 3)), + Arc::new(Int32Array::from_iter_values(0..values as i32)), + None, + ); + let keys = StringArray::from_iter_values((0..values).map(|i| format!("k{i}"))); + let entries = StructArray::from(vec![ + ( + Arc::new(Field::new("key", DataType::Utf8, false)), + Arc::new(keys) as ArrayRef, + ), + ( + Arc::new(Field::new("value", DataType::Utf8, true)), + strings(values), + ), + ]); + let map = MapArray::new( + Arc::new(Field::new("entries", entries.data_type().clone(), false)), + offsets, + entries, + nulls.clone(), + false, + ); + let fixed_size_list = + FixedSizeListArray::new(item, 2, strings(2 * n), nulls.clone()); + let list: ArrayRef = Arc::new(list); + let r#struct = StructArray::new( + vec![ + Field::new("s", DataType::Utf8, true), + Field::new("l", list.data_type().clone(), true), + ] + .into(), + vec![strings(n), Arc::clone(&list)], + nulls, + ); + + let utf8 = strings(n); + let dictionary = + DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)); + Ok(RecordBatch::try_from_iter(vec![ + ( + "int32", + Arc::new(Int32Array::from_iter_values(0..n as i32)) as _, + ), + ("large_utf8", cast(&utf8, &DataType::LargeUtf8)?), + ("binary", cast(&utf8, &DataType::Binary)?), + ("large_binary", cast(&utf8, &DataType::LargeBinary)?), + ("utf8_view", cast(&utf8, &DataType::Utf8View)?), + ("dictionary", cast(&utf8, &dictionary)?), + ("utf8", utf8), + ("list", list), + ("large_list", Arc::new(large_list)), + ("map", Arc::new(map)), + ("fixed_size_list", Arc::new(fixed_size_list)), + ("struct", Arc::new(r#struct)), + ])?) + } + + /// Offsets spilled as lengths read back as the arrays that were written, + /// sliced or not, see [`offset_lengths`]. The two row slice has buffers + /// small enough that the IPC writer stores them uncompressed. + #[tokio::test] + async fn test_spill_offsets_as_lengths_round_trip() -> Result<()> { + let batch = build_offsets_batch(1000)?; + 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 env = Arc::new(RuntimeEnv::default()); + let metrics = SpillMetrics::new(&ExecutionPlanMetricsSet::new(), 0); + let spill_manager = SpillManager::new(env, metrics, batch.schema()) + .with_compression_type(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_offsets_as_lengths_compress() -> Result<()> { + use arrow_ipc::writer::IpcWriteOptions; + + 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 env = Arc::new(RuntimeEnv::default()); + let metrics = SpillMetrics::new(&ExecutionPlanMetricsSet::new(), 0); + let spill_manager = SpillManager::new(env, metrics, Arc::clone(&schema)) + .with_compression_type(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(()) + } + // ==== Spill manager tests ==== #[test] diff --git a/datafusion/physical-plan/src/spill/offset_lengths.rs b/datafusion/physical-plan/src/spill/offset_lengths.rs new file mode 100644 index 0000000000000..19a763aa7b739 --- /dev/null +++ b/datafusion/physical-plan/src/spill/offset_lengths.rs @@ -0,0 +1,264 @@ +// 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. + +//! Writes offsets to compressed spill files as lengths, see +//! [`offsets_to_lengths`]. + +use std::ops::Range; +use std::sync::Arc; + +use arrow::array::{ + Array, ArrayRef, OffsetSizeTrait, RecordBatch, RecordBatchOptions, make_array, +}; +use arrow::buffer::Buffer; +use arrow::datatypes::DataType; +use arrow_data::{ArrayData, ArrayDataBuilder}; +use datafusion_common::Result; + +/// Returns whether [`offsets_to_lengths`] rewrites arrays of `data_type`. +pub(super) fn has_offsets(data_type: &DataType) -> bool { + match data_type { + DataType::Utf8 + | DataType::LargeUtf8 + | DataType::Binary + | DataType::LargeBinary + | DataType::List(_) + | DataType::LargeList(_) + | DataType::Map(_, _) => true, + DataType::Struct(fields) => fields.iter().any(|f| has_offsets(f.data_type())), + DataType::FixedSizeList(field, _) => has_offsets(field.data_type()), + _ => false, + } +} + +/// Rewrites every offsets buffer of `batch` as the lengths it encodes, before +/// the batch is written to a compressed spill file. +/// +/// Offsets are a running total that never repeats, so zstd and lz4 barely +/// compress them, while the lengths of similar values repeat +/// (). For `n` values the +/// buffer holds `[0, len_0, ..., len_(n-2), total]`. The IPC writer only reads +/// the first and last entries, to slice the values, so it writes the rewritten +/// arrays as they are. [`lengths_to_offsets`] restores the offsets. +/// +/// The returned batch breaks the offsets invariant, so it must only be passed +/// to the IPC writer. +pub(super) fn offsets_to_lengths(batch: &RecordBatch) -> Result { + let columns = batch + .columns() + .iter() + .map(|column| { + if has_offsets(column.data_type()) { + make_array(encode(&column.to_data())) + } else { + 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, + )?) +} + +/// Restores the offsets of a batch written by [`offsets_to_lengths`], in +/// place unless a buffer is shared. +pub(super) fn lengths_to_offsets(batch: RecordBatch) -> Result { + let (schema, columns, num_rows) = batch.into_parts(); + let columns: Vec = columns + .into_iter() + .map(|column| { + if has_offsets(column.data_type()) { + // `into_data` drops `column`, leaving its buffers unshared + make_array(decode(column.into_data())) + } else { + column + } + }) + .collect(); + let options = RecordBatchOptions::new().with_row_count(Some(num_rows)); + Ok(RecordBatch::try_new_with_options( + schema, columns, &options, + )?) +} + +fn encode(data: &ArrayData) -> ArrayData { + let (buffers, child_data) = match data.data_type() { + DataType::Utf8 | DataType::Binary | DataType::List(_) | DataType::Map(_, _) => { + encode_offsets::(data) + } + DataType::LargeUtf8 | DataType::LargeBinary | DataType::LargeList(_) => { + encode_offsets::(data) + } + DataType::Struct(_) => ( + vec![], + data.child_data() + .iter() + .map(|child| encode_child(&child.slice(data.offset(), data.len()))) + .collect(), + ), + DataType::FixedSizeList(_, size) => { + let size = *size as usize; + let child = + data.child_data()[0].slice(data.offset() * size, data.len() * size); + (vec![], vec![encode_child(&child)]) + } + _ => return data.clone(), + }; + // SAFETY: only the IPC writer reads the result, see `offsets_to_lengths` + unsafe { + ArrayDataBuilder::new(data.data_type().clone()) + .len(data.len()) + .nulls(data.nulls().cloned()) + .buffers(buffers) + .child_data(child_data) + .build_unchecked() + } +} + +fn encode_child(child: &ArrayData) -> ArrayData { + if has_offsets(child.data_type()) { + encode(child) + } else { + child.clone() + } +} + +/// Returns the buffers and children of an array with offsets, its offsets +/// rewritten as lengths and its values sliced to the ones they cover. +fn encode_offsets(data: &ArrayData) -> (Vec, Vec) { + // The IPC writer writes empty arrays without reading their buffers + if data.is_empty() { + return (data.buffers().to_vec(), data.child_data().to_vec()); + } + let (lengths, values) = lengths::(data); + match data.child_data().first() { + Some(child) => ( + vec![lengths], + vec![encode_child(&child.slice(values.start, values.len()))], + ), + None => ( + vec![ + lengths, + data.buffers()[1].slice_with_length(values.start, values.len()), + ], + vec![], + ), + } +} + +/// Returns the lengths of the values of `data`, laid out as described in +/// [`offsets_to_lengths`], and the range of values they cover. +fn lengths(data: &ArrayData) -> (Buffer, Range) { + let len = data.len(); + let offsets = &data.buffers()[0].typed_data::()[data.offset()..][..=len]; + let (first, last) = (offsets[0], offsets[len]); + let mut lengths = Vec::with_capacity(len + 1); + lengths.push(O::default()); + lengths.extend( + offsets[1..len] + .iter() + .zip(&offsets[..len - 1]) + .map(|(end, start)| *end - *start), + ); + lengths.push(last - first); + (lengths.into(), first.as_usize()..last.as_usize()) +} + +fn decode(data: ArrayData) -> ArrayData { + let (data_type, len, nulls, offset, mut buffers, child_data) = data.into_parts(); + match &data_type { + DataType::Utf8 | DataType::Binary | DataType::List(_) | DataType::Map(_, _) => { + restore_offsets::(&mut buffers[0]) + } + DataType::LargeUtf8 | DataType::LargeBinary | DataType::LargeList(_) => { + restore_offsets::(&mut buffers[0]) + } + _ => {} + } + let child_data = child_data + .into_iter() + .map(|child| { + if has_offsets(child.data_type()) { + decode(child) + } else { + child + } + }) + .collect(); + // SAFETY: the offsets are the ones of the array that was spilled + unsafe { + ArrayDataBuilder::new(data_type) + .len(len) + .offset(offset) + .nulls(nulls) + .buffers(buffers) + .child_data(child_data) + .build_unchecked() + } +} + +/// Turns the lengths written by [`offsets_to_lengths`] back into offsets. +fn restore_offsets(buffer: &mut Buffer) { + fn running_sum(lengths: &mut [O]) { + // The last entry already holds the total + let len = lengths.len().saturating_sub(1); + let mut sum = O::default(); + for length in &mut lengths[..len] { + sum += *length; + *length = sum; + } + } + *buffer = match std::mem::take(buffer).into_mutable() { + Ok(mut lengths) => { + running_sum(lengths.typed_data_mut::()); + lengths.into() + } + // The IPC writer stores a buffer uncompressed when compressing it does + // not help, and the reader then slices it out of the message body + Err(shared) => { + let mut lengths = shared.typed_data::().to_vec(); + running_sum(&mut lengths); + lengths.into() + } + }; +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_restore_offsets_in_place_unless_shared() { + // Values of lengths 3, 4 and 5 + let lengths = Buffer::from_vec(vec![0_i32, 3, 4, 12]); + let offsets = [0, 3, 7, 12]; + + let mut shared = lengths.clone(); + restore_offsets::(&mut shared); + assert_eq!(shared.typed_data::(), offsets); + assert_ne!(shared.as_ptr(), lengths.as_ptr()); + + let ptr = lengths.as_ptr(); + let mut unshared = lengths; + restore_offsets::(&mut unshared); + assert_eq!(unshared.typed_data::(), offsets); + assert_eq!(unshared.as_ptr(), ptr); + } +}