[comphead]

Which issue does this PR close?

• Closes Multi-column sort fails with ResourcesExhausted instead of spilling when the sort keys are most of each row #26940.

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 the runs of a merge in the row format into batches of an eighth of batch_size, but at least 1024 rows, so the merge holds the keys of one such batch per run. That covers several keys and a single key without a specialized cursor, such as FixedSizeBinary. Split runs compact their view arrays, so each batch is charged only for its own strings. The merge output keeps batch_size, and sorts without a memory limit are unchanged. This adds a batch_size parameter to sort_run, which #26930 also changes.

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. Two Int32 keys and a single FixedSizeBinary(16) key fail on main with ResourcesExhausted and pass here. Two Int32 keys with a Utf8View or a dictionary payload check that splitting does not charge each batch for the shared buffers. The reproducer from the issue, 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

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.

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 are not split, since each batch would be charged for the shared buffers. Each run's existing one-batch read-ahead now holds one of the smaller batches with its reservation already released.

[comphead]

This PR would be good addition to #26931

[jayzhan211]

@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.

[jayzhan211]

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?

[comphead]

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.

[jayzhan211]

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
        ))
}

[comphead]

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.