fix: bound spill file batches by bytes so external sorts over wide rows can merge

datafusion/core/tests/fuzz_cases/spilling_fuzz_in_memory_constrained_env.rs

@@ -273,12 +273,12 @@ async fn test_sort_with_limited_memory_and_oversized_record_batch() -> Result<()
 
     let number_of_record_batches = 100;
 
-    // Each spilled run's largest batch is so big that two merge streams cannot be
-    // reserved at once even at the smallest read-buffer size (`2 * (2 * batch) >
-    // pool`), yet a single stream still fits (`2 * batch < pool`). Reducing the
-    // buffer size therefore cannot help, the multi-level merge has to re-spill a
-    // run with a smaller batch size to make progress instead of failing with
-    // `ResourcesExhausted`.
+    // Each input batch is so big that two merge streams over spill files of such
+    // batches cannot be reserved at once (`2 * (2 * batch) > pool`), yet a single
+    // stream fits (`2 * batch < pool`). The sort must not fail with
+    // `ResourcesExhausted`: it bounds the batches it spills by bytes, so the merge
+    // can seat two streams. (The merge's fallback of re-spilling a run at half the
+    // batch size is covered by the unit tests in `multi_level_merge.rs`.)
     let metrics = run_sort_test_with_limited_memory(RunTestWithLimitedMemoryArgs {
         pool_size,
         task_ctx: Arc::new(task_ctx),
@@ -290,12 +290,9 @@ async fn test_sort_with_limited_memory_and_oversized_record_batch() -> Result<()
     })
     .await?;
 
-    let output_batches = get_output_batches_from_metrics(&metrics);
-
-    // minimum 2 batches more
     assert!(
-        output_batches >= number_of_record_batches + 2,
-        "output_batches {output_batches} should be greater than number_of_record_batches ({number_of_record_batches}) + 2"
+        metrics.spill_count().unwrap_or(0) > 0,
+        "the sort must spill"
     );
 
     Ok(())

datafusion/core/tests/memory_limit/mod.rs

@@ -26,6 +26,7 @@ mod nlj_spill_unmatched;
 mod repartition_mem_limit;
 mod union_nullable_spill;
 mod view_spill_compaction;
+mod wide_row_sort_spill;
 use arrow::array::{
     ArrayRef, DictionaryArray, Int32Array, Int64Array, Int64Builder, ListBuilder,
     RecordBatch, StringArray, StringViewArray, StructArray,

datafusion/core/tests/memory_limit/wide_row_sort_spill.rs

@@ -0,0 +1,147 @@
+// 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.
+
+//! An external sort over wide rows must finish under a memory limit that holds a few rows,
+//! even though `batch_size` rows do not fit.
+//!
+//! The input arrives in small batches, so every input batch fits. Before the fix, the sort
+//! wrote its spill files in batches of up to `batch_size` rows. With wide rows, one spilled
+//! batch was larger than the memory pool, and the final merge of the spill files, which
+//! reserves memory for the largest batch of each file, failed with `ResourcesExhausted`.
+
+use std::sync::Arc;
+
+use arrow::array::{
+    ArrayRef, Int64Array, ListBuilder, RecordBatch, StringArray, StringBuilder,
+    StringViewArray,
+};
+use arrow_schema::{DataType, Field, Schema};
+use datafusion::datasource::MemTable;
+use datafusion::execution::memory_pool::FairSpillPool;
+use datafusion::execution::runtime_env::RuntimeEnvBuilder;
+use datafusion::physical_plan::collect;
+use datafusion::prelude::{SessionConfig, SessionContext};
+
+/// 64 KiB per row.
+const ROW_BYTES: usize = 64 * 1024;
+/// 32 MiB of rows in total.
+const ROWS: usize = 512;
+/// Rows per input batch: 256 KiB, which fits in the pool.
+const INPUT_BATCH_ROWS: usize = 4;
+/// Holds about 190 rows, less than half of `ROWS`, so the sort must spill.
+const POOL_BYTES: usize = 12 * 1024 * 1024;
+const SPILL_RESERVATION_BYTES: usize = 1024 * 1024;
+
+fn payload(i: usize) -> String {
+    format!("{i:08}{}", "x".repeat(ROW_BYTES))
+}
+
+async fn sort_wide_rows(data_type: DataType) {
+    let schema = Arc::new(Schema::new(vec![
+        Field::new("id", DataType::Int64, false),
+        Field::new("payload", data_type.clone(), false),
+    ]));
+    let batches: Vec<RecordBatch> = (0..ROWS)
+        .step_by(INPUT_BATCH_ROWS)
+        .map(|start| {
+            let rows = start..start + INPUT_BATCH_ROWS;
+            let ids = Int64Array::from_iter_values(rows.clone().map(|i| i as i64));
+            let payload: ArrayRef = match data_type {
+                DataType::Utf8 => {
+                    Arc::new(StringArray::from_iter_values(rows.map(payload)))
+                }
+                DataType::Utf8View => {
+                    Arc::new(StringViewArray::from_iter_values(rows.map(payload)))
+                }
+                DataType::List(_) => {
+                    // Four elements of a quarter row each.
+                    let mut builder = ListBuilder::new(StringBuilder::new());
+                    for i in rows {
+                        let quarter = format!("{i:08}{}", "x".repeat(ROW_BYTES / 4));
+                        for _ in 0..4 {
+                            builder.values().append_value(&quarter);
+                        }
+                        builder.append(true);
+                    }
+                    Arc::new(builder.finish())
+                }
+                _ => unreachable!(),
+            };
+            RecordBatch::try_new(Arc::clone(&schema), vec![Arc::new(ids), payload])
+                .unwrap()
+        })
+        .collect();
+
+    let table = MemTable::try_new(Arc::clone(&schema), vec![batches]).unwrap();
+    let runtime = RuntimeEnvBuilder::new()
+        .with_memory_pool(Arc::new(FairSpillPool::new(POOL_BYTES)))
+        .build_arc()
+        .unwrap();
+    let config = SessionConfig::new()
+        .with_target_partitions(1)
+        // 512 rows of 64 KiB is 32 MiB, more than the pool.
+        .with_batch_size(ROWS)
+        .with_sort_spill_reservation_bytes(SPILL_RESERVATION_BYTES);
+    let ctx = SessionContext::new_with_config_rt(config, runtime);
+    ctx.register_table("t", Arc::new(table)).unwrap();
+
+    let plan = ctx
+        .sql("SELECT id, payload FROM t ORDER BY id DESC")
+        .await
+        .unwrap()
+        .create_physical_plan()
+        .await
+        .unwrap();
+    let result = collect(Arc::clone(&plan), ctx.task_ctx())
+        .await
+        .expect("the sort must spill and finish");
+
+    let spill_count = plan.metrics().and_then(|m| m.spill_count()).unwrap_or(0);
+    assert!(spill_count > 0, "the sort must spill, plan:\n{plan:?}");
+
+    let ids: Vec<i64> = result
+        .iter()
+        .flat_map(|batch| {
+            batch
+                .column(0)
+                .as_any()
+                .downcast_ref::<Int64Array>()
+                .unwrap()
+                .values()
+                .to_vec()
+        })
+        .collect();
+    let expected: Vec<i64> = (0..ROWS as i64).rev().collect();
+    assert_eq!(ids, expected);
+}
+
+#[tokio::test]
+async fn sort_wide_utf8_rows_spills_and_finishes() {
+    sort_wide_rows(DataType::Utf8).await;
+}
+
+#[tokio::test]
+async fn sort_wide_utf8_view_rows_spills_and_finishes() {
+    sort_wide_rows(DataType::Utf8View).await;
+}
+
+/// A list column keeps its whole child array alive when it is sliced, so the spill writer
+/// must measure and copy only the child range that each piece uses.
+#[tokio::test]
+async fn sort_wide_list_rows_spills_and_finishes() {
+    sort_wide_rows(DataType::new_list(DataType::Utf8, true)).await;
+}

datafusion/physical-plan/Cargo.toml

@@ -64,6 +64,7 @@ arrow-data = { workspace = true }
 arrow-ipc = { workspace = true, features = ["lz4", "zstd"] }
 arrow-ord = { workspace = true }
 arrow-schema = { workspace = true }
+arrow-select = { workspace = true }
 async-trait = { workspace = true }
 bytes = { workspace = true }
 datafusion-common = { workspace = true }
@@ -142,6 +143,11 @@ name = "invariant_check"
 harness = false
 name = "union_schema"
 
+[[bench]]
+harness = false
+name = "sort_spill"
+required-features = ["test_utils"]
+
 [[bench]]
 harness = false
 name = "spill_io"

datafusion/physical-plan/benches/sort_spill.rs

@@ -0,0 +1,147 @@
+// 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.
+
+//! Benchmarks for an external sort that spills under a memory limit.
+//!
+//! Each case sorts one partition of `NUM_BATCHES` batches of `BATCH_SIZE` rows with a
+//! memory pool smaller than the input, so the sort writes spill files and merges them.
+
+use std::sync::Arc;
+
+use arrow::array::{ArrayRef, DictionaryArray, Int64Array, StringArray};
+use arrow::datatypes::{DataType, Field, Int32Type, Schema};
+use arrow::record_batch::RecordBatch;
+use criterion::{BenchmarkId, Criterion, criterion_group, criterion_main};
+use datafusion_execution::TaskContext;
+use datafusion_execution::config::SessionConfig;
+use datafusion_execution::memory_pool::FairSpillPool;
+use datafusion_execution::runtime_env::RuntimeEnvBuilder;
+use datafusion_physical_expr::{LexOrdering, PhysicalSortExpr, expressions::col};
+use datafusion_physical_plan::sorts::sort::SortExec;
+use datafusion_physical_plan::test::TestMemoryExec;
+use datafusion_physical_plan::{ExecutionPlan, collect};
+use rand::rngs::StdRng;
+use rand::{Rng, SeedableRng};
+use tokio::runtime::Runtime;
+
+const NUM_BATCHES: usize = 16;
+const BATCH_SIZE: usize = 8192;
+
+/// A sort key and one payload column, in random key order.
+fn make_batches(payload: impl Fn(&mut StdRng, usize) -> ArrayRef) -> Vec<RecordBatch> {
+    let mut rng = StdRng::seed_from_u64(42);
+    (0..NUM_BATCHES)
+        .map(|_| {
+            let keys: ArrayRef = Arc::new(Int64Array::from_iter_values(
+                (0..BATCH_SIZE).map(|_| rng.random::<i64>()),
+            ));
+            let payload = payload(&mut rng, BATCH_SIZE);
+            let schema = Arc::new(Schema::new(vec![
+                Field::new("key", DataType::Int64, false),
+                Field::new("payload", payload.data_type().clone(), false),
+            ]));
+            RecordBatch::try_new(schema, vec![keys, payload]).unwrap()
+        })
+        .collect()
+}
+
+/// Dictionary with many distinct values, shared by all input batches.
+fn dictionary_batches() -> Vec<RecordBatch> {
+    const DISTINCT: usize = 100_000;
+    let values: ArrayRef = Arc::new(StringArray::from_iter_values(
+        (0..DISTINCT).map(|i| format!("value number {i:08}")),
+    ));
+    make_batches(|rng, rows| {
+        let keys = (0..rows).map(|_| rng.random_range(0..DISTINCT as i32));
+        Arc::new(
+            DictionaryArray::<Int32Type>::try_new(keys.collect(), Arc::clone(&values))
+                .unwrap(),
+        )
+    })
+}
+
+/// Short strings, so no spilled batch is split.
+fn narrow_utf8_batches() -> Vec<RecordBatch> {
+    make_batches(|rng, rows| {
+        Arc::new(StringArray::from_iter_values(
+            (0..rows).map(|_| format!("string {:016}", rng.random::<u32>())),
+        ))
+    })
+}
+
+/// 1 KiB strings: 8 MiB per batch, larger than the spill batch bound.
+fn wide_utf8_batches() -> Vec<RecordBatch> {
+    make_batches(|rng, rows| {
+        Arc::new(StringArray::from_iter_values((0..rows).map(|_| {
+            format!("{:016}{}", rng.random::<u32>(), "x".repeat(1008))
+        })))
+    })
+}
+
+fn sort_spilled(rt: &Runtime, batches: &[RecordBatch], pool_bytes: usize) {
+    let schema = batches[0].schema();
+    let input =
+        TestMemoryExec::try_new_exec(&[batches.to_vec()], Arc::clone(&schema), None)
+            .unwrap();
+    let ordering =
+        LexOrdering::new([PhysicalSortExpr::new_default(col("key", &schema).unwrap())])
+            .unwrap();
+    let sort = Arc::new(SortExec::new(ordering, input));
+    let runtime = RuntimeEnvBuilder::new()
+        .with_memory_pool(Arc::new(FairSpillPool::new(pool_bytes)))
+        .build_arc()
+        .unwrap();
+    let task_ctx = Arc::new(
+        TaskContext::default()
+            .with_session_config(SessionConfig::new().with_batch_size(BATCH_SIZE))
+            .with_runtime(runtime),
+    );
+    let output = rt
+        .block_on(collect(
+            Arc::clone(&sort) as Arc<dyn ExecutionPlan>,
+            task_ctx,
+        ))
+        .unwrap();
+    assert_eq!(
+        output.iter().map(|b| b.num_rows()).sum::<usize>(),
+        NUM_BATCHES * BATCH_SIZE
+    );
+    let spills = sort.metrics().and_then(|m| m.spill_count()).unwrap_or(0);
+    assert!(spills > 0, "the sort must spill");
+}
+
+fn criterion_benchmark(c: &mut Criterion) {
+    let rt = Runtime::new().unwrap();
+    let cases: [(&str, Vec<RecordBatch>, usize); 3] = [
+        ("dictionary", dictionary_batches(), 32 * 1024 * 1024),
+        ("narrow_utf8", narrow_utf8_batches(), 16 * 1024 * 1024),
+        ("wide_utf8", wide_utf8_batches(), 64 * 1024 * 1024),
+    ];
+    let mut group = c.benchmark_group("sort_spill");
+    group.sample_size(10);
+    for (name, batches, pool_bytes) in &cases {
+        group.bench_with_input(
+            BenchmarkId::from_parameter(name),
+            batches,
+            |b, batches| b.iter(|| sort_spilled(&rt, batches, *pool_bytes)),
+        );
+    }
+    group.finish();
+}
+
+criterion_group!(benches, criterion_benchmark);
+criterion_main!(benches);