perf: spill byte arrays as lengths and coalesce spill writes

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?;
 

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
+//! <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. 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<SchemaRef> {
+    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<RecordBatch> {
+    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::<i32>())?
+            }
+            (DataType::LargeUtf8, DataType::Struct(struct_fields)) => {
+                encode_bytes::<_, Int64Type>(struct_fields, column.as_string::<i64>())?
+            }
+            (DataType::Binary, DataType::Struct(struct_fields)) => {
+                encode_bytes::<_, Int32Type>(struct_fields, column.as_binary::<i32>())?
+            }
+            (DataType::LargeBinary, DataType::Struct(struct_fields)) => {
+                encode_bytes::<_, Int64Type>(struct_fields, column.as_binary::<i64>())?
+            }
+            _ => {
+                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<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 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>(fields: &Fields, array: &GenericByteArray<T>) -> Result<ArrayRef>
+where
+    T: ByteArrayType,
+    L: ArrowPrimitiveType<Native = T::Offset>,
+{
+    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::<L>::new(lengths, array.nulls().cloned());
+
+    let (start, len) = offset_span(offsets);
+    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.clone(),
+        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 offsets = OffsetBuffer::<T::Offset>::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::<Int64Type>().values();
+    let bytes = bytes.as_binary::<T::Offset>();
+    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::<T>::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(())
+    }
+}