Skip to content

fix: split multi-column sort runs into smaller batches under a memory limit - #26953

Open
comphead wants to merge 6 commits into
apache:mainfrom
comphead:chunk-sort-merge-runs
Open

comphead wants to merge 6 commits into
apache:mainfrom
comphead:chunk-sort-merge-runs

Conversation

@comphead

@comphead comphead commented Oct 30, 2026 •

Copy link
Copy Markdown
Contributor

Which issue does this PR close?

Rationale for this change

A sort under a memory limit can fail with ResourcesExhausted for ExternalSorterMerge instead of spilling when its merge uses the row format, as with several sort keys. The merge holds the row-format keys of every sorted run at once. When the keys are most of each row, that needs more than the 2x the sort reserves for each buffered batch. The issue has the details and a reproducer.

What changes are included in this PR?

With a memory limit, ExternalSorter sorts a run that feeds a merge in the row format into batches of an eighth of batch_size, but at least 1024 rows, when the run's encoded keys may be larger than the run, so the merge holds the keys of one such batch per run. Smaller keys fit next to the sorted copy, since the sort reserves about twice the run. That covers several keys and a single key without a specialized cursor, such as FixedSizeBinary. A computed key is sized from its type if that has a fixed width, and otherwise counts as larger. Split runs compact their view arrays, so each batch is charged only for its own strings. Runs whose views share bytes, as dictionary-encoded Parquet strings do, keep batch_size, because compacting them would copy the shared bytes once per row. Spilling also compacts view arrays with several data buffers, since each buffer read back from the spill is a slice of one IPC message and is charged for all of it. The merge output keeps batch_size, and sorts without a memory limit are unchanged. This adds a batch_size parameter to sort_run.

What is the testing strategy for this PR?

test_merge_keys_larger_than_batches (current-thread and multi-thread) sorts in a pool that holds 16 buffered batches, or 8 for the clustered keys. Two Int32 keys and a single FixedSizeBinary(16) key fail on main with ResourcesExhausted and pass here. Two Int32 keys with a Utf8View, a dictionary or a dictionary-backed Utf8View payload check that splitting does not charge each batch for the shared buffers, and two clustered Utf8View keys check the spilled buffers. test_merge_run_batch_size_follows_key_size checks that only runs whose keys outgrow them are split. The reproducer from the issue and two Parquet sorts from review, with datafusion-cli:

main This PR
30M rows, 1 partition, -m 512M fails 2 spills
30M rows, default partitions, -m 1G fails 16 spills
28M rows, 1 partition, -m 1G fails no spill
dictionary-encoded strings, 1 partition, -m 256M 3 spills 3 spills
clustered unique strings, 1 partition, -m 32M 507 spills 507 spills

With dfbench, sort-tpch Q5 from #16979 (SF1, 1 partition, --memory-limit 100M) fails on main and passes here. With the same settings in datafusion-cli, it also fails on main at 50M and 200M and passes here with 6 and 2 spills.

The criterion sort benchmark has no memory limit and is unchanged. With the split forced on at batch_size = 8192, its multi-column in-memory cases took 2 to 12% longer, which is the cost under a memory limit. No-spill Parquet sorts at -m 16G with integer or computed integer keys and a Utf8View payload are within 4% of main in time and peak RSS.

Are there any user-facing changes?

No. Known limitations: runs of at most 1024 rows are not split. That covers sorts with a batch_size of 1024 or less, and small input batches, which #26931 would coalesce into larger runs. Sorts with dictionary, list view or nested view columns, and runs whose views share bytes, are not split, since each batch would be charged for the shared buffers. The key size counts the bytes of variable-length keys without the row format's padding, so string keys that only slightly outgrow their run are not split. Runs sorted by a computed key without a fixed width, such as a string expression, are split under a memory limit even when their keys fit. Each run's existing one-batch read-ahead now holds one of the smaller batches with its reservation already released.

@github-actions github-actions Bot added the physical-plan Changes to the physical-plan crate label Oct 30, 2026
@comphead
comphead marked this pull request as ready for review October 30, 2026 20:42
@comphead

Copy link
Copy Markdown
Contributor Author

This PR would be good addition to #26931

@comphead
comphead requested review from andygrove and sunchao October 30, 2026 20:59

@jayzhan211 jayzhan211 left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@comphead thanks for the fix! I found two things, both inline: a case where the split now runs out of memory where main succeeds, and single-key sorts that the gate misses.

/// be larger than the batches. With a memory limit, runs therefore use an eighth
/// of `batch_size`, but at least 1024 rows to keep the per-batch merge cost
/// small. The merge output still uses `batch_size`.
fn merge_run_batch_size(&self) -> usize {

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

With a non-inlined Utf8View column, each of the 8 chunks from sort_batch_chunked keeps the input's full view data buffer, and sort_run adds up get_record_batch_memory_size per chunk. With logging added to sort_run, one run had input=1228800 reserved=2457600 total_sorted=8454144, so it asks for about 8x the payload. This test passes on main and fails here with Failed to allocate additional 5.7 MB for ExternalSorter[0]:

let schema = Arc::new(Schema::new(vec![
    Field::new("a", DataType::Int32, false),
    Field::new("b", DataType::Int32, false),
    Field::new("c", DataType::Utf8View, false),
]));
// 64 batches of 8192 rows; c = format!("{i:0>64}"); sort by (a, b)
let capacity = MERGE_BYTES + 16 * get_reserved_bytes_for_record_batch(&batches[0])?;
let runtime = RuntimeEnvBuilder::new()
    .with_memory_pool(Arc::new(GreedyMemoryPool::new(capacity)))
    .build_arc()?;

Running crate::spill::gc_view_arrays on the chunks when sorted_batches.len() > 1 fixed it locally. Three exact-accounting tests in spill_tests.rs then need new expected numbers. Could you add this case as a test with the fix?

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks, confirmed and fixed in 571f7f9. Split runs now go through gc_view_arrays, and a Utf8View payload is a case of test_merge_keys_larger_than_batches. I limited the compaction to runs split for the merge. Runs chunked only because they exceed batch_size keep their current charge, and the three workspace-loan tests stay unchanged, since they create the loan with those unsplit chunks.

A high-cardinality dictionary payload failed the same way, and gc_view_arrays cannot compact shared dictionary values. So schemas whose split batches would still share buffers keep batch_size, and that payload is a test case too.

self.runtime.memory_pool.memory_limit(),
MemoryLimit::Infinite
);
if limited && self.expr.len() > 1 {

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

expr.len() > 1 doesn't match when the merge uses the row format: StreamingMergeBuilder::build also uses RowCursorStream for a single key that isn't primitive or a plain string or binary type. A single FixedSizeBinary(16) key (UUIDs) still fails the same way. Could the gate follow the merge's choice, ideally as a helper next to build so the two stay in sync, with a test for a single FixedSizeBinary key? With this, that test passes and the existing sort tests still pass:

fn merges_in_row_format(expr: &LexOrdering, schema: &Schema) -> bool {
    if expr.len() > 1 {
        return true;
    }
    let Ok(data_type) = expr[0].expr.data_type(schema) else {
        return true;
    };
    !(data_type.is_primitive()
        || matches!(
            data_type,
            DataType::Utf8
                | DataType::Utf8View
                | DataType::LargeUtf8
                | DataType::Binary
                | DataType::LargeBinary
        ))
}

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Good catch, thanks. merges_in_row_format now sits next to StreamingMergeBuilder::build, with a note in build that the two must match, and merge_run_batch_size uses it. A single FixedSizeBinary(16) key is a case of test_merge_keys_larger_than_batches. It fails on main and on the previous head, and passes now.

@comphead

comphead commented Nov 5, 2026

Copy link
Copy Markdown
Contributor Author

@jayzhan211 PTAL

@jayzhan211 jayzhan211 left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@comphead thanks for the update! Both earlier asks look good, and the issue's reproducer passes. My gc_view_arrays suggestion had two gaps, though: two real Parquet sorts now run out of memory where main succeeds. Both are inline with a fix.

Comment on lines +1040 to +1041
// Set when `merge_run_batch_size` split this run
let split = batch_size < self.batch_size;

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Sorry, my gc suggestion misses views that share bytes, which is what dictionary-encoded Parquet strings produce. A chunk that would write at least as much as the buffer holds skips gc, so all 8 chunks keep the buffer and each is charged for it. This fails at -m 256M, 512M and 1G with Not enough memory to continue external sort, and passes on main:

COPY (SELECT value % 1000003 AS a, value AS b,
             concat(repeat('x', 240), (value % 1000)::varchar) AS c
      FROM range(5000000)) TO 'shared.parquet' STORED AS PARQUET;
SET datafusion.execution.target_partitions = 1;
SELECT * FROM 'shared.parquet' ORDER BY a, b;

Could these runs keep batch_size? The helper below goes next to gc_view_arrays in spill/mod.rs, and sort.rs needs to import it.

Suggested change
// Set when `merge_run_batch_size` split this run
let split = batch_size < self.batch_size;
// Split batches of views sharing bytes would each be charged for them.
let batch_size = if views_share_bytes(&batch) {
self.batch_size
} else {
batch_size
};
// Set when `merge_run_batch_size` split this run
let split = batch_size < self.batch_size;
pub(crate) fn views_share_bytes(batch: &RecordBatch) -> bool {
    use arrow::array::AsArray;
    fn shares<T: ByteViewType>(array: &GenericByteViewArray<T>) -> bool {
        let (_, unique) = dedup_buffer_entries(array.data_buffers());
        array.total_buffer_bytes_used() > unique.iter().map(|b| b.len()).sum::<usize>()
    }
    batch.columns().iter().any(|array| match array.data_type() {
        DataType::Utf8View => shares(array.as_string_view()),
        DataType::BinaryView => shares(array.as_binary_view()),
        _ => false,
    })
}

Nested views already keep batch_size through sorted_copy_keeps_buffers, so only these two top-level types need the check. Could you add a test with a payload of cast(&dict, &DataType::Utf8View), using 8192 keys over 1024 values of 64 bytes? It should succeed with correct output. Skip the assert_released peak check there, because the grow in try_sort_and_spill_in_mem_batches goes over the pool on main too. Please also unit test the boundary: views_share_bytes is true for that array and false for 8192 values each used once.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks, applied in 0938167. views_share_bytes sits next to gc_view_arrays, and sort_run only checks it when the run would be split, so sorts without a memory limit don't scan the views. The dictionary-backed Utf8View payload is a case of test_merge_keys_larger_than_batches, without the pool checks as you suggested, and without the change it fails with Not enough memory to continue external sort. The boundary test is test_views_share_bytes in spill/mod.rs. Your query now passes at 256M, 512M and 1G with the same spill counts as main.

sort_batch_chunked(&batch, &expressions, batch_size)?;
// `take` keeps the data buffers of view arrays, so splitting would
// charge every batch for all of them.
if split && sorted_batches.len() > 1 {

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The per-chunk gc also causes a second failure, on clustered input. An output batch taken from one run holds 8 chunk buffers with no unused bytes, so gc skips it. After the spill is read back, each buffer is a slice of one IPC message and is charged for the whole message, so a batch's max_record_batch_memory reached 22.1 MB against 3.0 MB on main. This fails at -m 32M and 64M, and passes on main:

COPY (SELECT value / 1000 AS a, value AS b,
             concat(repeat('x', 60), value::varchar) AS c,
             concat(repeat('y', 60), value::varchar) AS d,
             concat(repeat('z', 60), value::varchar) AS e
      FROM range(5000000)) TO 'clustered3.parquet' STORED AS PARQUET;
SET datafusion.execution.target_partitions = 1;
SELECT * FROM 'clustered3.parquet' ORDER BY a, b;

Could spilling compact these arrays to one buffer? With this and the views_share_bytes change, all of my probes match or beat main, and all sorts:: and spill:: tests pass, including yours. The issue's reproducer still passes too.

--- a/datafusion/physical-plan/src/spill/mod.rs
+++ b/datafusion/physical-plan/src/spill/mod.rs
@@ fn compact_view_array
-    if should_gc_view_array(array, unique_size) {
+    // Several buffers read back as slices of one IPC message, each charged for all of it
+    if should_gc_view_array(array, unique_size)
+        || (unique.len() > 1 && array.total_buffer_bytes_used() <= unique_size)
+    {
         return Some(array.gc());
     }

The <= keeps gc away from views that share bytes, where it would copy once per row (#23564). Could you add a test with two Utf8View keys of unique 66-byte strings, prefixed by batch number so the runs don't overlap, in a pool of 8 buffered batches? It should succeed; today it fails in ExternalSorterMerge.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Applied as suggested in 0938167, and the doc comment of compact_view_array now covers it. The two clustered Utf8View keys with room for 8 batches fail without it and pass now. Your query passes at 32M and 64M with the same 507 and 113 spills as main. Since this also changes spilling for other operators, I checked a GROUP BY on a Utf8View column at 64M and 128M, which spills the same as on main. A BinaryView payload, Utf8View keys with overlapping runs, shared and unique views in one table, nine Int64 keys and a single FixedSizeBinary(16) key also match or beat main.

@comphead
comphead requested a review from jayzhan211 November 6, 2026 18:10

@jayzhan211 jayzhan211 left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@comphead thanks for the update! All of the earlier failures are fixed, and I found no case that fails where main succeeds. I do have one performance ask inline, which I should have raised earlier, and a small test request.

One optional doc fix: in spill/mod.rs around lines 762 and 770, the gc_view_arrays docs still say small buffers skip gc. Arrays with several buffers now compact at any size.

/// when their batches would share buffers other than the data of top-level
/// view columns, such as dictionary values, since every batch is charged for
/// them. `sort_run` also keeps it for a run whose views share bytes.
fn merge_run_batch_size(&self) -> usize {

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Sorry, I should have raised this last round: splitting also costs runs that don't need it. With integer keys and a Utf8View payload, the per-batch gc copies every string, so Parquet sorts at -m 16G with no spill took 9–17% longer than main, and peak RSS rose 13–38%. Their encoded keys are 26 bytes a row against about 100 for the run, so they already fit in the 2x reservation.

Could the split apply only when the keys outgrow the run? With the diff below those sorts are within 3% of main, and the issue's reproducer still passes. The trade-off is that the Utf8View payload in your test goes back to main's pool overshoot, hence within_pool. Repro:

COPY (SELECT value % 1000003 AS a, value AS b,
             concat(repeat('u', 60), value::varchar) AS c
      FROM range(5000000)) TO 'unique.parquet' STORED AS PARQUET;
SET datafusion.execution.target_partitions = 1;
EXPLAIN ANALYZE SELECT * FROM 'unique.parquet' ORDER BY a, b;
--- a/datafusion/physical-plan/src/sorts/sort.rs
+++ b/datafusion/physical-plan/src/sorts/sort.rs
@@ imports
 use datafusion_physical_expr::PhysicalExpr;
+use datafusion_physical_expr::expressions::Column;
 use datafusion_physical_expr::expressions::{DynamicFilterPhysicalExpr, lit};
@@ fn prepare_sort_runs
-        let batch_size = self.merge_run_batch_size();
         let jobs =
             runs.into_iter()
                 .zip(sizes)
                 .enumerate()
                 .map(move |(index, (batch, size))| {
                     let reservation = self.reservation.split(size);
+                    let batch_size = self.merge_run_batch_size(&batch)?;
                     let sort = self.sort_run(
@@ fn sort_run_streams
-        let batch_size = self.merge_run_batch_size();
         runs.into_iter()
             .zip(sizes)
             .map(|(run, size)| {
                 let reservation = self.reservation.split(size);
+                let batch_size = self.merge_run_batch_size(&run)?;
                 self.sort_batch_stream(run, reservation, shared_input.clone(), batch_size)
@@ fn sort_run
-        // Split batches of views sharing bytes would each be charged for them
-        let batch_size = if batch_size < self.batch_size && views_share_bytes(&batch) {
-            self.batch_size
-        } else {
-            batch_size
-        };
         // Set when `merge_run_batch_size` split this run
         let split = batch_size < self.batch_size;
@@ fn merge_run_batch_size
-    /// Batch size of the sorted runs that feed a merge. A merge in the row format
-    /// holds the keys of every run's current batch at once, and these keys can be
-    /// larger than the batches. With a memory limit, runs therefore use an eighth
-    /// of `batch_size`, but at least 1024 rows to keep the per-batch merge cost
-    /// small. The merge output still uses `batch_size`. Runs keep `batch_size`
-    /// when their batches would share buffers other than the data of top-level
-    /// view columns, such as dictionary values, since every batch is charged for
-    /// them. `sort_run` also keeps it for a run whose views share bytes.
-    fn merge_run_batch_size(&self) -> usize {
+    /// Batch size of the sorted run `run` when it feeds a merge. A merge in the
+    /// row format holds the keys of every run's current batch at once. The sort
+    /// reserves about twice the run, so the keys fit next to the sorted copy if
+    /// they are at most as large as the run. Otherwise, with a memory limit, the
+    /// run uses an eighth of `batch_size`, but at least 1024 rows to keep the
+    /// per-batch merge cost small. The merge output still uses `batch_size`.
+    ///
+    /// Runs keep `batch_size` when their batches would share buffers other than
+    /// the data of top-level view columns, such as dictionary values, or when the
+    /// views of `run` share bytes, since every batch is charged for them.
+    fn merge_run_batch_size(&self, run: &RecordBatch) -> Result<usize> {
         let limited = !matches!(
             self.runtime.memory_pool.memory_limit(),
             MemoryLimit::Infinite
         );
         let shares_buffers = self.schema.fields().iter().any(|field| {
             let data_type = field.data_type();
             sorted_copy_keeps_buffers(data_type)
                 && !matches!(data_type, DataType::Utf8View | DataType::BinaryView)
         });
-        if limited && !shares_buffers && merges_in_row_format(&self.expr, &self.schema) {
-            (self.batch_size / 8).max(1024).min(self.batch_size)
+        if limited
+            && !shares_buffers
+            && merges_in_row_format(&self.expr, &self.schema)
+            && !views_share_bytes(run)
+            && self.encoded_keys_may_exceed(run)?
+        {
+            Ok((self.batch_size / 8).max(1024).min(self.batch_size))
         } else {
-            self.batch_size
+            Ok(self.batch_size)
+        }
+    }
+
+    /// Returns whether the row-format keys of `run` may be larger than `run`. A
+    /// computed key counts as larger, since sizing it would evaluate it again.
+    fn encoded_keys_may_exceed(&self, run: &RecordBatch) -> Result<bool> {
+        // An 8-byte offset per row, then a null byte and the bytes of each key
+        let mut encoded = run.num_rows() * size_of::<usize>();
+        for sort_expr in self.expr.iter() {
+            let Some(column) = sort_expr.expr.downcast_ref::<Column>() else {
+                return Ok(true);
+            };
+            encoded += run.num_rows()
+                + run
+                    .column(column.index())
+                    .to_data()
+                    .get_slice_memory_size()?;
         }
+        Ok(encoded > run.get_sliced_size()?)
     }
--- a/datafusion/physical-plan/src/sorts/sort/spill_tests.rs
+++ b/datafusion/physical-plan/src/sorts/sort/spill_tests.rs
@@ fn check_merge_keys_larger_than_batches
-        check_merge_keys(&fields, 2, 16, true, |batch| {
+        // Runs of this `Utf8View` payload are not split, as on `main`
+        let within_pool = payload != DataType::Utf8View;
+        check_merge_keys(&fields, 2, 16, within_pool, |batch| {

Please also add this boundary test. It separates the rule from "any view payload": a short inlined payload must not split either.

/// Only runs whose encoded keys are larger than the run are split.
#[test]
fn test_merge_run_batch_size_follows_key_size() -> Result<()> {
    let batch_size = 8192;
    let rows = batch_size as i32;
    let check = |payload: Option<ArrayRef>, expected: usize| -> Result<()> {
        let mut columns: Vec<(&str, ArrayRef)> = vec![
            ("a", Arc::new(Int64Array::from_iter_values(0..rows as i64))),
            ("b", Arc::new(Int64Array::from_iter_values(0..rows as i64))),
        ];
        columns.extend(payload.map(|payload| ("c", payload)));
        let batch = RecordBatch::try_from_iter(columns)?;
        let ordering = LexOrdering::new((0..2).map(|i| {
            PhysicalSortExpr::new_default(Arc::new(Column::new(
                batch.schema().field(i).name(),
                i,
            )))
        }))
        .expect("two sort keys");
        let runtime = RuntimeEnvBuilder::new()
            .with_memory_limit(1024 * 1024 * 1024, 1.0)
            .build_arc()?;
        let sorter = ExternalSorter::new(
            0,
            batch.schema(),
            ordering,
            batch_size,
            0,
            0,
            SpillCompression::Uncompressed,
            &ExecutionPlanMetricsSet::new(),
            runtime,
        )?;
        assert_eq!(sorter.merge_run_batch_size(&batch)?, expected);
        Ok(())
    };
    // Encoded keys: 26 bytes a row, against 16 bytes a row
    check(None, batch_size / 8)?;
    // Against 16 + 16 bytes a row: inlined views add more than the keys need
    let short: StringViewArray = (0..rows).map(|i| Some(format!("{i}"))).collect();
    check(Some(Arc::new(short)), batch_size)?;
    // Against 16 + 16 + 64 bytes a row
    let long: StringViewArray = (0..rows).map(|i| Some(format!("{i:0>64}"))).collect();
    check(Some(Arc::new(long)), batch_size)
}

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks, this is a good catch. Applied in 8eb93eb, with one change: merge_run_batch_size checks the key size before views_share_bytes, so runs that keep batch_size skip the scan of their views. With your repro at -m 16G, the previous version was 14% slower than main with 35% more peak RSS, and this one is within 4% on both. The issue's reproducer and the memory-limited probes from the earlier rounds still pass with the same spill counts. Your boundary test is in as test_merge_run_batch_size_follows_key_size.

For variable-length keys the estimate counts their bytes without the row format's padding, so string keys that only slightly outgrow a run are not split, as on main. I noted that under the known limitations.

let unique_size: usize = unique.iter().map(|b| b.len()).sum();
if should_gc_view_array(array, unique_size) {
if should_gc_view_array(array, unique_size)
|| (unique.len() > 1 && array.total_buffer_bytes_used() <= unique_size)

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Nothing tests the <= here: changing this to unique.len() > 1 still passes every sorts:: and spill:: test. It's what keeps gc from copying shared bytes once per row when views span several distinct buffers. Could you add this next to test_gc_keeps_shared_dictionary_buffer? It passes today and fails without the <=.

#[test]
fn test_gc_keeps_shared_views_in_two_buffers() -> Result<()> {
    use arrow::array::AsArray;
    use arrow::compute::concat;

    let (_, first) = shared_dictionary_views(1000, 8, 1024)?;
    let (_, second) = shared_dictionary_views(1000, 8, 1024)?;
    let merged = concat(&[&first as &dyn Array, &second])?;
    let merged = merged.as_string_view();

    // `gc()` would copy 16384 values, 8x what the two dictionaries hold
    let compacted = compact_view_array(merged).expect("buffer entries repeat");
    assert_eq!(compacted.data_buffers().len(), 2);
    assert_eq!(&compacted, merged);

    Ok(())
}

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Added as test_gc_keeps_shared_views_in_two_buffers in 8eb93eb. It fails when the condition is only unique.len() > 1. Thanks!

… limit

A multi-column merge holds the row-format keys of every sorted run's current
batch. When the sort keys are most of each row, those keys are larger than the
batches, and the merge needs more memory than the sort reserved while
buffering them, so the sort fails with ResourcesExhausted instead of spilling.

With a memory limit, sort the runs of a multi-column merge into batches of an
eighth of `batch_size`, but at least 1024 rows. The merge output keeps
`batch_size`, and sorts without a memory limit are unchanged.

Closes apache#26940
- Split runs whenever the merge uses the row format, including a single
  key without a specialized cursor, such as `FixedSizeBinary`.
- Compact the view arrays of split runs, so each batch is charged for its
  own strings rather than for the run's whole data buffers.
- Keep `batch_size` when split batches would still share buffers, such as
  dictionary values.
- Keep `batch_size` for a run whose views share bytes, as dictionary-encoded
  Parquet strings do. Compacting such views copies the shared bytes once per
  row, so every split batch would keep and be charged for the whole buffer.
- Compact view arrays with several data buffers before spilling unless their
  views share bytes. Read back, each buffer is a slice of one IPC message and
  is charged for all of it.
The sort reserves about twice each run, so a merge in the row format fits
when the keys are at most as large as the run. Splitting smaller keys only
added work: with integer keys and a `Utf8View` payload, the compaction of
split batches copied every string. `merge_run_batch_size` now checks each
run, from the bytes of its key columns, and also covers views that share
bytes.

Also test that spill compaction leaves views sharing bytes across several
buffers alone, and update the `gc_view_arrays` docs.
@comphead
comphead force-pushed the chunk-sort-merge-runs branch from 0938167 to 8eb93eb Compare November 9, 2026 17:15
@comphead

comphead commented Nov 9, 2026

Copy link
Copy Markdown
Contributor Author

@jayzhan211 thank you for the careful reviews, your probes made this PR much better. I rebased onto main, resolved the conflict with #27015 in spill/mod.rs, and applied your latest suggestions, including the gc_view_arrays doc fix. Could you take another look?

@comphead
comphead requested a review from jayzhan211 November 9, 2026 17:16

@jayzhan211 jayzhan211 left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@comphead thanks for applying all of the last round's changes. One gap is left: a test for the views_share_bytes gate, details inline.

// dictionary-encoded Parquet column produces. Sorting the buffered batches
// before a spill grows past the pool, on `main` too.
let fields = [&int_fields[..], &[("c", DataType::Utf8View)]].concat();
check_merge_keys(&fields, 2, 16, false, |batch| {

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Since the key-size check went in, no test reaches views_share_bytes: removing && !views_share_bytes(run) from merge_run_batch_size still passes every sorts:: and spill:: test. It's still needed for computed keys, which always count as larger than the run. With round 2's shared.parquet, this passes on main and here at -m 256M, 512M and 1G, and fails without the gate with Not enough memory to continue external sort:

SET datafusion.execution.target_partitions = 1;
SELECT * FROM 'shared.parquet' ORDER BY a + 1, b;

Could you add a computed-key case? The diff below passes here and fails without the gate.

Test diff for spill_tests.rs
-use datafusion_physical_expr::expressions::Column;
-use datafusion_physical_expr::{LexOrdering, PhysicalSortExpr};
+use datafusion_physical_expr::expressions::{CastExpr, Column};
+use datafusion_physical_expr::{LexOrdering, PhysicalExpr, PhysicalSortExpr};
@@ async fn check_merge_keys_larger_than_batches
     let fields = [&int_fields[..], &[("c", DataType::Utf8View)]].concat();
-    check_merge_keys(&fields, 2, 16, false, |batch| {
+    let shared_views = |batch: i32| -> Result<Vec<ArrayRef>> {
         let values: StringArray = (0..1024)
@@
         Ok([int_keys(batch), vec![views]].concat())
-    })
-    .await?;
+    };
+    check_merge_keys(&fields, 2, 16, false, shared_views).await?;
+    // A computed key counts as larger than the run, so only the shared views
+    // keep this run at `batch_size`.
+    let a: Arc<dyn PhysicalExpr> = Arc::new(Column::new("a", 0));
+    let keys = vec![
+        Arc::new(CastExpr::new(a, DataType::Int64, None)) as Arc<dyn PhysicalExpr>,
+        Arc::new(Column::new("b", 1)),
+    ];
+    check_merge_exprs(&fields, keys, 16, false, shared_views).await?;
@@ async fn check_merge_keys(
     columns: impl Fn(i32) -> Result<Vec<ArrayRef>>,
+) -> Result<()> {
+    let keys = (0..keys)
+        .map(|i| Arc::new(Column::new(fields[i].0, i)) as Arc<dyn PhysicalExpr>)
+        .collect();
+    check_merge_exprs(fields, keys, buffered, within_pool, columns).await
+}
+
+/// Like [`check_merge_keys`], sorting by `keys`.
+async fn check_merge_exprs(
+    fields: &[(&str, DataType)],
+    keys: Vec<Arc<dyn PhysicalExpr>>,
+    buffered: usize,
+    within_pool: bool,
+    columns: impl Fn(i32) -> Result<Vec<ArrayRef>>,
 ) -> Result<()> {
@@
-    let ordering =
-        LexOrdering::new((0..keys).map(|i| {
-            PhysicalSortExpr::new_default(Arc::new(Column::new(fields[i].0, i)))
-        }))
+    let ordering = LexOrdering::new(keys.into_iter().map(PhysicalSortExpr::new_default))
         .expect("at least one sort key");

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks, you're right on both points. With the gate removed every sorts:: and spill:: test still passed, and your ORDER BY a + 1, b on shared.parquet fails at 256M, 512M and 1G while main and this PR pass with the same spill counts.

Rather than add a second shared-views case, b29d5b0 switches the existing one to CAST(a AS Int64), b through your check_merge_exprs. With column keys that case no longer reached the gate, and the Utf8View payload case already covers keys that fit. It now fails without the gate with Not enough memory to continue external sort.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

One more change in 0e71291. Computed keys counted as larger, so a computed integer key with a Utf8View payload still split, and ORDER BY a + 1, b on your unique.parquet at -m 16G took 11% longer than main with 35% more peak RSS. Keys with a fixed-width type are now sized from it, which brings that sort back to main. A computed-key version of the issue's reproducer, ORDER BY a + 1, b over range(30000000) at 512M, still fails on main and passes here.

The shared-views case now sorts by CAST(a AS Utf8), b, which still counts as larger and still fails without the gate, and test_merge_run_batch_size_follows_key_size now checks both kinds of computed key.

Since runs whose keys fit keep `batch_size`, the shared-views case sorted
by column keys never reached `views_share_bytes`. Sort it by a computed key,
which counts as larger than the run, so the case fails without the gate.
@comphead
comphead requested a review from jayzhan211 November 11, 2026 16:25
A computed key counted as larger than the run, so its runs were always
split under a memory limit. With a computed integer key and a `Utf8View`
payload, the compaction of split batches copied every string. Size such
keys from their type when it has a fixed width. The shared-views case now
sorts by a string cast, which still counts as larger.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

physical-plan Changes to the physical-plan crate

Projects

None yet

2 participants