Repository navigation
perf: coalesce small input batches before sorting - #26931
Open
geoffreyclaude wants to merge 1 commit into
Open
geoffreyclaude wants to merge 1 commit into
geoffreyclaude wants to merge 1 commit into
Conversation
github-actions
Bot
added
core
Contributor
|
This is a great PR , @geoffreyclaude are you planning to move on this? |
geoffreyclaude
force-pushed
the
codex/sort-small-batches
branch
2 times, most recently
from
November 2, 2026 10:34
61a0ddf to
b5358e6
Compare
Contributor
Author
|
run benchmark sort_tpch sort_tpch10 |
Contributor
Author
|
run benchmark sort env: |
|
🤖 Benchmark running (GKE) | trigger CPU Details (lscpu)Comparing codex/sort-small-batches (b5358e6) to d4c2a32 (merge-base) diff Run configurationrun benchmark sort_tpchResults will be posted here when complete File an issue against this benchmark runner |
|
🤖 Benchmark running (GKE) | trigger CPU Details (lscpu)Comparing codex/sort-small-batches (b5358e6) to d4c2a32 (merge-base) diff Run configurationrun benchmark sort_tpch10Results will be posted here when complete File an issue against this benchmark runner |
|
🤖 Benchmark completed (GKE) | trigger Instance: Comparing codex/sort-small-batches (b5358e6) to d4c2a32 (merge-base) diff Run configurationrun benchmark sort_tpchCPU Details (lscpu)Details
Resource Usagesort_tpch — base (merge-base)
sort_tpch — branch
File an issue against this benchmark runner |
|
🤖 Benchmark running (GKE) | trigger CPU Details (lscpu)Comparing codex/sort-small-batches (b5358e6) to d4c2a32 (merge-base) diff Run configurationrun benchmark sort
env:
BENCH_FILTER: "small_batches"Results will be posted here when complete File an issue against this benchmark runner |
|
🤖 Benchmark completed (GKE) | trigger Instance: Comparing codex/sort-small-batches (b5358e6) to d4c2a32 (merge-base) diff Run configurationrun benchmark sort_tpch10CPU Details (lscpu)Details
Resource Usagesort_tpch10 — base (merge-base)
sort_tpch10 — branch
File an issue against this benchmark runner |
|
🤖 Benchmark completed (GKE) | trigger Instance: Comparing codex/sort-small-batches (b5358e6) to d4c2a32 (merge-base) diff Run configurationrun benchmark sort
env:
BENCH_FILTER: "small_batches"CPU Details (lscpu)Details
Resource Usagesort — base (merge-base)
sort — branch
File an issue against this benchmark runner |
geoffreyclaude
force-pushed
the
codex/sort-small-batches
branch
from
November 2, 2026 13:27
b5358e6 to
1563663
Compare
Contributor
Author
@comphead Yes definitely, I just need to find time to do a proper self-review of the AI code first. There are also two other follow-up PRs (still very much in AI draft too): |
geoffreyclaude
force-pushed
the
codex/sort-small-batches
branch
from
November 4, 2026 19:43
1563663 to
97f342f
Compare
geoffreyclaude
force-pushed
the
codex/sort-small-batches
branch
from
November 4, 2026 19:59
97f342f to
808bfdb
Compare
geoffreyclaude
marked this pull request as ready for review
November 4, 2026 20:28
Contributor
Author
@comphead I've finally opened the PR for review. I believe it's best to focus on this one first as it brings the clearest benefit. The following two should be #26930 (code simplification, removes 186 LOC) and #26933 (yield every batch of rows). |
comphead
reviewed
Nov 4, 2026
comphead
left a comment
comphead
left a comment
Contributor
There was a problem hiding this comment.
Nice speedup on small_batches/two_keys. Most of my comments are about memory on the spill path, which this PR now coalesces too.
- Benchmarks:
sort_tpchreads 8192-row batches, so the SF1/SF10no changeresults show that the untouched path stays untouched. They don't exercise the new path. The biggest behavior change is coalescing before a spill, so a memory-limited small-batch run (peak RSS and spill count, before and after) would cover it. - Design question: did you consider coalescing at ingest in
insert_batch, for example with theBatchCoalescerthatLimitedBatchCoalesceralready wraps? Each small batch would be copied once as it arrives and charged through the existingreserve_memory_for_batch_and_maybe_spill. That removes the fallback and the realign pass. The costs are an extra copy for inputs that end up on the in-place path, and the shared-slice dedup no longer applying to coalesced batches. - Row target: was a larger target measured for multi-key sorts?
batch_sizesits between the old one-run-per-batch rule and the single-key byte target. The removed comment said many small runs beat a few large lexsorts, and this PR's numbers contradict it, so the right size is an open question. - Nit: the run rule is described in four places (the
SortExecdocs, thein_mem_sort_streamdocs, the comment before the call, and thecoalesce_in_mem_batches_into_runsdocs). Keeping it oncoalesce_in_mem_batches_into_runswould be enough. - Nit: in
sort_spill_reservation, iswith_target_partitions(1)needed? The scenario already has one partition and the plan is asserted. - Nit:
coalesce_small_runs_preserves_rows_and_full_batchesis missing thetest_prefix that the surrounding tests use.
| let target = self.sort_in_place_threshold_bytes.max(1); | ||
| let mut runs: Vec<RecordBatch> = Vec::new(); | ||
| let mut group: Vec<RecordBatch> = Vec::new(); | ||
| let mut group: Vec<&RecordBatch> = Vec::new(); |
Contributor
There was a problem hiding this comment.
group now borrows from batches so that the whole input can be returned on failure. That keeps every original batch alive until the last group has been copied. On main, each group's originals were dropped right after its concat_batches, so the transient peak was one group, not the whole buffer.
The copies are also not reserved before they are made. On the spill path the pool is full by construction. Any growth in footprint (for example the aliased columns in test_small_run_coalescing_falls_back_when_memory_is_full) is only caught by the final try_resize. By then the whole buffer has been copied, and all of it is thrown away.
Owning the group and falling back per group bounds both:
let flush = |group: &mut Vec<RecordBatch>, runs: &mut Vec<RecordBatch>| {
if group.len() > 1
&& let Ok(run) = concat_batches(&schema, group.iter())
{
group.clear();
runs.push(run);
} else {
runs.append(group);
}
};If the group's copy is reserved with try_grow before concatenating, and the group is kept as is when that fails, the final realign becomes a shrink. This also removes both return Ok(batches) paths and the pop().cloned().unwrap().
| if group_rows >= self.batch_size | ||
| && (!coalesce_large_runs || group_bytes.saturating_add(bytes) > target) | ||
| { | ||
| if flush(&mut group, &mut runs, &self.schema).is_err() { |
Contributor
There was a problem hiding this comment.
Two points on this fallback:
is_err()swallows every error, not only the offset and dictionary-key limits named in the comment above. Matching the Arrow overflow variants (DictionaryKeyOverflowError,OffsetOverflowError) would keep real failures visible.- The in-place path in
in_mem_sort_streamstill callsconcat_batcheson all buffered batches. With the defaultsort_in_place_threshold_bytes(1 MiB), the input fromtest_small_run_coalescing_falls_back_on_dictionary_overflowtakes that path and fails withArrow error: Dictionary key bigger than the key type. I checked this onmain, and this PR doesn't change that path. The test only passes because it sets the threshold to 0. Either fall through to the run path when that concat fails, or narrow the dictionary claim in the description and leave the rest as a follow-up.
Comment on lines
+906
to
+907
| if group_rows >= self.batch_size | ||
| && (!coalesce_large_runs || group_bytes.saturating_add(bytes) > target) |
Contributor
There was a problem hiding this comment.
A full batch that follows a short group gets absorbed and copied. With batch_size = 4, [2] [4] becomes one 6-row run, built by copying the 4-row batch. Outside large-run mode, pushing batches with num_rows() >= batch_size straight to runs avoids that copy. Run order doesn't matter, since equal keys have no ordering guarantee.
With that change, the all(..) pre-scan at the top is no longer needed either. An all-full input copies nothing, so a copied flag that skips the realign pass at the end covers the same fast path in one traversal.
adriangb
pushed a commit
to adriangb/datafusion
that referenced
this pull request
Nov 12, 2026
… limit (apache#26953) ## Which issue does this PR close? - Closes apache#26940. - Closes apache#16979. ## 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 apache#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 apache#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.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Which issue does this PR close?
Follow-up to the minimum-run-size suggestion in #26242.
Rationale for this change
A sort processes chunks of input separately, then merges the sorted chunks, called runs. Small input batches create many tiny runs. Processing and merging so many runs adds overhead.
What changes are included in this PR?
Combine adjacent small batches before sorting until each run reaches the configured
batch_size, except for the last run. This applies to sorts with one or several keys, including runs written to spill files. Single-key sorts can still form larger runs based on memory size.Keep the original batches while trying to combine them. If concatenation fails or the combined data cannot fit its memory reservation, sort those original batches separately. For example, two valid dictionaries can contain too many distinct values to fit in one combined dictionary.
Example with multiple sort keys and
batch_size = 4(numbers are rows):What is the testing strategy for this PR?
Tests check batch grouping and row preservation, successful spilling with a full memory pool, and complete sorted output from dictionary batches that cannot be concatenated.
Bot results compare
b5358e61298with baselined4c2a329a8f. The existingsmall_batches/two_keysbenchmark, with two million rows stored in 32-row batches, improved from 1.212 s to 185 ms, about 6.6× faster.The standard SF1 and SF10 SQL sort suites classified all 22 cases as
no change.Are there any user-facing changes?
Rows with equal sort keys may change relative order when batches become one run. Trying concatenation temporarily retains both original and combined arrays.
Prepared with AI assistance.