comphead commented on code in PR #26960:
URL: https://github.com/apache/datafusion/pull/26960#discussion_r4429723558
##########
datafusion/physical-plan/src/spill/mod.rs:
##########
@@ -677,23 +699,62 @@ impl IPCStreamEncoder {
let compression_type =
Option::<CompressionType>::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<Buffer>)> {
- 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, buffers))
+ Ok((batch.num_rows(), bytes, coalesce_buffers(buffers)))
}
fn finish(&mut self) -> Result<(usize, Vec<Buffer>)> {
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<Buffer>) -> Vec<Buffer> {
+ let mut coalesced = Vec::new();
+ let mut small: Vec<u8> = 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
}
Review Comment:
The cost sits in the unbuffered `FileSpillWriter`, and a `BufWriter` around
its file does the same job. Writing one 8192-row batch of the reproducer data,
with compressed columns in this PR's layout, took:
| | uncompressed | lz4 | zstd |
|---|---:|---:|---:|
| one `write_all` per buffer | 125-137 µs | 127-139 µs | 122-125 µs |
| `coalesce_buffers` | 87-95 µs | 11-12 µs | 8-9 µs |
| 64 KB `BufWriter`, flushed per batch | 101-105 µs | 11 µs | 7 µs |
The `BufWriter` reuses one bounded buffer and passes large buffers through
without copying. `coalesce_buffers` allocates a `Vec` per run that grows by
doubling, keeps up to twice its bytes alive across the async `write_all`, and
copies every small buffer of uncompressed spills. `spill_pool` flushes after
each batch and both finish paths flush, so readers of in-progress files are
unaffected. Custom `AsyncSpillWriter` backends would keep getting separate
buffers, which the object store example copies into its parts anyway.
If coalescing stays here, sizing `small` from the lengths summed in the same
pass as `bytes` avoids the regrowth.
##########
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
+//! <https://github.com/apache/datafusion/issues/26627>. 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<SchemaRef> {
+ 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<RecordBatch> {
+ let Some(schema) = encode_schema(batch.schema_ref()) else {
+ return Ok(batch.clone());
+ };
+ let columns = batch
+ .columns()
+ .iter()
+ .map(encode_array)
+ .collect::<Result<_>>()?;
+ Ok(RecordBatch::try_new(schema, columns)?)
+}
+
+/// Restores a batch that [`encode`] wrote for `schema`
+pub(super) fn decode(batch: &RecordBatch, schema: &SchemaRef) ->
Result<RecordBatch> {
+ let columns = schema
+ .fields()
+ .iter()
+ .zip(batch.columns())
+ .map(|(field, column)| decode_array(field.data_type(), column))
+ .collect::<Result<_>>()?;
+ Ok(RecordBatch::try_new(Arc::clone(schema), columns)?)
+}
+
+fn encoded_type(data_type: &DataType) -> Option<DataType> {
+ let fields = match data_type {
+ DataType::Utf8 | DataType::Binary => encoded_fields::<Int32Type>(),
+ DataType::LargeUtf8 | DataType::LargeBinary =>
encoded_fields::<Int64Type>(),
+ _ => return None,
+ };
+ Some(DataType::Struct(fields))
+}
+
+/// The fields of the struct a byte array with `L` offsets is spilled as
+fn encoded_fields<L>() -> Fields
+where
+ L: ArrowPrimitiveType,
+ L::Native: OffsetSizeTrait,
+{
+ let run_ends = Field::new("run_ends", DataType::Int64, false);
+ let bytes = Field::new("values",
GenericBinaryType::<L::Native>::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<ArrayRef> {
+ match array.data_type() {
+ DataType::Utf8 => encode_bytes::<_,
Int32Type>(array.as_string::<i32>()),
+ DataType::LargeUtf8 => encode_bytes::<_,
Int64Type>(array.as_string::<i64>()),
+ DataType::Binary => encode_bytes::<_,
Int32Type>(array.as_binary::<i32>()),
+ DataType::LargeBinary => encode_bytes::<_,
Int64Type>(array.as_binary::<i64>()),
+ _ => Ok(Arc::clone(array)),
+ }
+}
+
+fn decode_array(data_type: &DataType, array: &ArrayRef) -> Result<ArrayRef> {
+ match data_type {
+ DataType::Utf8 => decode_bytes::<Utf8Type, Int32Type>(array),
+ DataType::LargeUtf8 => decode_bytes::<LargeUtf8Type, Int64Type>(array),
+ DataType::Binary => decode_bytes::<BinaryType, Int32Type>(array),
+ DataType::LargeBinary => decode_bytes::<LargeBinaryType,
Int64Type>(array),
+ _ => Ok(Arc::clone(array)),
+ }
+}
+
+fn encode_bytes<T, L>(array: &GenericByteArray<T>) -> Result<ArrayRef>
+where
+ T: ByteArrayType,
+ L: ArrowPrimitiveType<Native = T::Offset>,
+{
+ let fields = encoded_fields::<L>();
+ 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::<L>::new(lengths, array.nulls().cloned());
+
+ let start = offsets[0].as_usize();
+ let len = offsets[array.len()].as_usize() - start;
+ let bytes = GenericBinaryArray::<T::Offset>::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<T, L>(array: &ArrayRef) -> Result<ArrayRef>
+where
+ T: ByteArrayType,
+ L: ArrowPrimitiveType<Native = T::Offset>,
+{
+ let array = array.as_struct();
+ let lengths = array.column(0).as_primitive::<L>();
+ 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::<Int64Type>().values();
+ let bytes = bytes.as_binary::<T::Offset>();
+ 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::<T>::new_unchecked(
+ OffsetBuffer::new_unchecked(offsets.into()),
+ bytes.values().slice_with_length(start, len),
+ lengths.nulls().cloned(),
+ )
+ };
Review Comment:
This skips checks that spill reads still run today. arrow-ipc 60 validates
record batches despite `with_skip_validation(true)`, so on `main` corrupted
offsets or invalid UTF-8 in a spill file fail with an error. Here the sum check
misses negative lengths that cancel out, and the bytes are never checked as
UTF-8, so a corrupted compressed file can decode into an invalid `StringArray`
instead.
- `OffsetBuffer::try_from_lengths` replaces the hand-written sum and the
`unsafe` `OffsetBuffer::new_unchecked`, and rejects negative or overflowing
lengths. It costs about 2-3 µs more per 8192-row batch of five `Utf8` columns.
- `GenericByteArray::try_new` keeps the UTF-8 check for about 22 µs per such
batch, roughly 7% of encode plus decode.
##########
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
+//! <https://github.com/apache/datafusion/issues/26627>. 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<SchemaRef> {
+ 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<RecordBatch> {
+ let Some(schema) = encode_schema(batch.schema_ref()) else {
+ return Ok(batch.clone());
+ };
Review Comment:
`IPCStreamEncoder::new` builds this schema but keeps only `encode_lengths`,
so it is rebuilt here for every batch. `encode_schema` also builds each struct
type twice, once only for `is_some`, and `encode_bytes` builds the same
`Fields` again for every column.
Keeping the encoded `SchemaRef` in `IPCStreamEncoder` instead of the bool
and passing it in, like `decode(batch, schema)`, removes the rebuild and this
unreachable `batch.clone()` branch. The struct `Fields` can come from that
schema. One catch: `RecordBatch::try_new` would then compare pass-through
columns with the spill schema, which rejects union batches that differ in
nested nullability (see the comment in `InProgressSpillFile::append_batch`).
##########
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
+//! <https://github.com/apache/datafusion/issues/26627>. 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<SchemaRef> {
+ 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<RecordBatch> {
+ let Some(schema) = encode_schema(batch.schema_ref()) else {
+ return Ok(batch.clone());
+ };
+ let columns = batch
+ .columns()
+ .iter()
+ .map(encode_array)
+ .collect::<Result<_>>()?;
+ Ok(RecordBatch::try_new(schema, columns)?)
+}
+
+/// Restores a batch that [`encode`] wrote for `schema`
+pub(super) fn decode(batch: &RecordBatch, schema: &SchemaRef) ->
Result<RecordBatch> {
+ let columns = schema
+ .fields()
+ .iter()
+ .zip(batch.columns())
+ .map(|(field, column)| decode_array(field.data_type(), column))
+ .collect::<Result<_>>()?;
+ Ok(RecordBatch::try_new(Arc::clone(schema), columns)?)
+}
+
+fn encoded_type(data_type: &DataType) -> Option<DataType> {
+ let fields = match data_type {
+ DataType::Utf8 | DataType::Binary => encoded_fields::<Int32Type>(),
+ DataType::LargeUtf8 | DataType::LargeBinary =>
encoded_fields::<Int64Type>(),
+ _ => return None,
+ };
+ Some(DataType::Struct(fields))
+}
+
+/// The fields of the struct a byte array with `L` offsets is spilled as
+fn encoded_fields<L>() -> Fields
+where
+ L: ArrowPrimitiveType,
+ L::Native: OffsetSizeTrait,
+{
+ let run_ends = Field::new("run_ends", DataType::Int64, false);
+ let bytes = Field::new("values",
GenericBinaryType::<L::Native>::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<ArrayRef> {
+ match array.data_type() {
+ DataType::Utf8 => encode_bytes::<_,
Int32Type>(array.as_string::<i32>()),
+ DataType::LargeUtf8 => encode_bytes::<_,
Int64Type>(array.as_string::<i64>()),
+ DataType::Binary => encode_bytes::<_,
Int32Type>(array.as_binary::<i32>()),
+ DataType::LargeBinary => encode_bytes::<_,
Int64Type>(array.as_binary::<i64>()),
+ _ => Ok(Arc::clone(array)),
+ }
+}
+
+fn decode_array(data_type: &DataType, array: &ArrayRef) -> Result<ArrayRef> {
+ match data_type {
+ DataType::Utf8 => decode_bytes::<Utf8Type, Int32Type>(array),
+ DataType::LargeUtf8 => decode_bytes::<LargeUtf8Type, Int64Type>(array),
+ DataType::Binary => decode_bytes::<BinaryType, Int32Type>(array),
+ DataType::LargeBinary => decode_bytes::<LargeBinaryType,
Int64Type>(array),
+ _ => Ok(Arc::clone(array)),
+ }
+}
+
+fn encode_bytes<T, L>(array: &GenericByteArray<T>) -> Result<ArrayRef>
+where
+ T: ByteArrayType,
+ L: ArrowPrimitiveType<Native = T::Offset>,
+{
+ let fields = encoded_fields::<L>();
+ 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::<L>::new(lengths, array.nulls().cloned());
+
+ let start = offsets[0].as_usize();
+ let len = offsets[array.len()].as_usize() - start;
+ let bytes = GenericBinaryArray::<T::Offset>::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<T, L>(array: &ArrayRef) -> Result<ArrayRef>
+where
+ T: ByteArrayType,
+ L: ArrowPrimitiveType<Native = T::Offset>,
+{
+ let array = array.as_struct();
+ let lengths = array.column(0).as_primitive::<L>();
+ 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::<Int64Type>().values();
+ let bytes = bytes.as_binary::<T::Offset>();
+ 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::<T>::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<()> {
Review Comment:
No `.slt` reaches this code. `spill_compression` defaults to `uncompressed`
and no test file sets it, and `varchar` columns spill as `Utf8View` because
`map_string_types_to_utf8view` defaults to true. Re-running a spilling query
from `aggregate_memory_spill.slt` or `sort_merge_join_spill.slt` with
`spill_compression` set to `zstd` and to `lz4_frame`, on columns cast to
`Utf8`, `LargeUtf8`, `Binary` and `LargeBinary` with `arrow_cast`, would cover
the real operator paths. This test could then keep only the sliced and empty
batches.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
