feat: Improve sort memory resilience #19494
Open
+1,145
−92
Conversation
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
EmilyMatt
commented
Dec 25, 2025
| let indices = lexsort_to_indices(&sort_columns, fetch)?; | ||
| let mut columns = take_arrays(batch.columns(), &indices, None)?; | ||
|
|
||
| // The columns may be larger than the unsorted columns in `batch` especially for variable length |
Contributor
Author
There was a problem hiding this comment.
This is not really needed as take actually uses the offsets for variable length columns.
At worst the buffers themselves will be rounded up to the 64 byte multiplier to optimize SIMD operations, but that is actually desirable and not something we want to mess with.
Member
|
run benchmarks |
|
🤖 |
Member
|
run benchmark sort |
|
🤖 Hi @rluvaton, thanks for the request (#19494 (comment)).
Please choose one or more of these with |
|
🤖: Benchmark completed Details
|
Contributor
|
run benchmark clickbench_partitioned |
|
🤖 |
|
🤖: Benchmark completed Details
|
Member
|
run benchmark sort |
|
🤖 |
Weijun-H
reviewed
Dec 29, 2025
| // Spilling to disk and reading back also ensures batch size is consistent | ||
| // rather than potentially having one significantly larger last batch. | ||
| self.spill()?; | ||
| self.spill()?; // TODO: use sort_batch_chunked instead |
Member
There was a problem hiding this comment.
Will you address this in this pr?
Contributor
Author
There was a problem hiding this comment.
I don't think so, I initially thought about this but then I saw there's an immediate spill now, rather than returning an in memory stream, which feels like a deliberate decision, and I didn't want to add side-effects to this PR
Comment on lines
832
to
836
| /// This is calculated by adding the record batch's memory size | ||
| /// (which can be much larger than expected for sliced record batches) | ||
| /// with the sliced buffer sizes, as that is the amount that will be needed to create the new buffer. | ||
| /// The latter is rounded up to the nearest multiple of 64 based on the architecture, | ||
| /// as this is how arrow creates buffers. |
Member
There was a problem hiding this comment.
/// Estimate how much memory is needed to sort a `RecordBatch`.
///
/// For sliced batches, `get_record_batch_memory_size` returns the size of the
/// underlying shared buffers (which may be larger than the logical data).
/// We add `get_sliced_size()` (the actual logical data size, rounded to 64 bytes)
/// because sorting will create new buffers containing only the referenced data.
///
/// Total = existing buffer size + new sorted buffer size
Contributor
Author
There was a problem hiding this comment.
Fine by me, updated.
Weijun-H
reviewed
Dec 29, 2025
Member
Weijun-H
left a comment
Weijun-H
left a comment
There was a problem hiding this comment.
Very impressive pr 👍 I left some suggestion for your reference
|
🤖: Benchmark completed Details
|
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
5 participants
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?
Closes #19493 .
Rationale for this change
Greatly reduces the memory requested by ExternalSorter to perform sorts, adds much more granularity to the reservations, and Tries to do this with minimal overhead by merging the splitting and sorting processes.
What changes are included in this PR?
The sort stream will calculate the indices once, but the take will be done in batches, so we create batch_size sized RecordBatches, whose get_record_batch_size results return info that is very close to their sliced sizes(if not completely the same), this means there is no need for the precaution of reserving a huge amount of memory in order to do the merge sort, meaning we can merge more streams at the same time, and so on and so forth.
Are these changes tested?
Yes
Are there any user-facing changes?
There is a new sort_batch_chunked function, which returns a Vec of RecordBatch, based on the provided batch_size.
Some docs are updated.