Repository navigation
fix: bound parallel batch sorting in SortExec - #26242
Open
geoffreyclaude wants to merge 12 commits into
Open
geoffreyclaude wants to merge 12 commits into
geoffreyclaude wants to merge 12 commits into
Conversation
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #26242 +/- ##
========================================
Coverage 82.91% 82.91%
========================================
Files 1155 1155
Lines 458091 458301 +210
Branches 458091 458301 +210
========================================
+ Hits 379823 380008 +185
Misses 55229 55229
- Partials 23039 23064 +25 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
geoffreyclaude
force-pushed
the
codex/bound-sort-preparation
branch
13 times, most recently
from
October 13, 2026 15:34
e2e4c5e to
9d0897a
Compare
geoffreyclaude
force-pushed
the
codex/bound-sort-preparation
branch
4 times, most recently
from
October 14, 2026 09:49
67093a0 to
1eacb12
Compare
geoffreyclaude
marked this pull request as ready for review
October 14, 2026 14:07
| .map_err(|e| internal_datafusion_err!("Sort concurrency closed: {e}"))?; | ||
| SpawnedTask::spawn(async move { | ||
| // Match the buffered producer's first-batch handoff so spill | ||
| // workspace can be reused while the remaining runs are sorted. |
Contributor
There was a problem hiding this comment.
I don't think this comment adds much value.
Contributor
Author
There was a problem hiding this comment.
indeed, removed.
| // current batch, returning workspace before the merge | ||
| // reserves memory for that batch and its cursor. | ||
| let next = stream.next().await; | ||
| Some((batch, (next, stream))) |
Contributor
There was a problem hiding this comment.
Haven't read the consumer side of this stream yet but here's one observation -
Since the batch is unaccounted, why not re-reserve the memory here in self.merge_pool?
- In total across all permits, there's
batch_size * semaphore size, so that's a lot of unaccounted memory right? - The consumed stream returns all the reserved memory (to the workspace or otherwise) before yielding a batch
Contributor
Author
There was a problem hiding this comment.
Good catch. Preparation now keeps all sorted output reserved until merging starts. During merging, the one-batch lookahead still temporarily releases the prefetched batch’s reservation so the merge can reuse that workspace.
| #[derive(Debug)] | ||
| struct SortConcurrency { | ||
| /// Bounds preparation futures buffered by each partition. | ||
| limit: usize, |
Contributor
There was a problem hiding this comment.
Hmm I guess you need this bc there's no API for total semaphore size?
Contributor
Author
There was a problem hiding this comment.
Yes, you can only easily get the still available permits, but we want the total size. Renamed the field to per_partition_limit to clarify its purpose.
| .prepare_sort_runs(runs, concurrency) | ||
| .with_elapsed_compute(elapsed_compute.clone()) | ||
| .await?; | ||
| timer.restart(); |
Contributor
There was a problem hiding this comment.
So pausing the timer actually fails to capture the time spent in the sort_batch_stream stream which calls sort_batch_chunked, which is elapsed compute we want to capture. Maybe pass the timer as a param? Or something better, a timer param feels like a smell.
Contributor
Author
There was a problem hiding this comment.
Fixed: sort_run now measures sorting and memory accounting on the task doing the work, for both inline and spawned execution. Scheduling waits remain excluded.
Contributor
There was a problem hiding this comment.
Oh neat, with_elapsed_compute is quite handy
| .sum::<usize>() | ||
| - input_bytes; | ||
| let first_bytes = get_record_batch_memory_size(&sorted_batches[0]); | ||
| assert!(loan > first_bytes); |
Contributor
There was a problem hiding this comment.
Can you comment why it's bigger? I guess it's because some 4096 size buffers are re-used but not filled?
Contributor
Author
There was a problem hiding this comment.
Each sorted StringView chunk retains the shared string buffers, and summing the chunks' memory sizes counts those buffers repeatedly. I’ve added a comment explaining this.
Related: #25820 fixes shared-buffer counting for sort inputs, but explicitly leaves StringView accounting for follow-up work (I ran into a conflict with that PR when rebasing over the latest main.)
| assert!(alive <= 1, "{alive} sort tasks remain before merging"); | ||
| assert_eq!( | ||
| sorter.merge_pool.borrow(usize::MAX).size(), | ||
| workspace - 2 * (loan - first_bytes) |
Contributor
There was a problem hiding this comment.
Adding comments around the math would make it easier to reason about
Contributor
Author
There was a problem hiding this comment.
Should hopefully be a tiny bit clearer now!
geoffreyclaude
force-pushed
the
codex/bound-sort-preparation
branch
2 times, most recently
from
October 14, 2026 22:25
4ba6d18 to
96f3f8a
Compare
geoffreyclaude
force-pushed
the
codex/bound-sort-preparation
branch
from
October 15, 2026 06:47
96f3f8a to
d1a59c1
Compare
Contributor
Author
|
run benchmark sort_tpch sort_tpch10 |
|
🤖 Benchmark running (GKE) | trigger CPU Details (lscpu)Comparing codex/bound-sort-preparation (d1a59c1) to e732212 (merge-base) diff Run configurationrun benchmark sort_tpchResults will be posted here when complete File an issue against this benchmark runner |
geoffreyclaude
force-pushed
the
codex/bound-sort-preparation
branch
from
October 19, 2026 15:14
5d19187 to
cc90f9e
Compare
gabotechs
self-requested a review
October 19, 2026 15:37
Contributor
|
Thanks @geoffreyclaude I'm doing first automated review and then check manually |
comphead
reviewed
Oct 20, 2026
comphead
left a comment
comphead
left a comment
Contributor
There was a problem hiding this comment.
Thanks for the fix. The permit-before-spawn design and the ordered reassembly look sound to me. Notes below, most important first.
- Coverage of the new path.
sort_fuzz.rsandspilling_fuzz_in_memory_constrained_env.rsuse#[tokio::test](current-thread), so they only run the inline path. The defaultsort_in_place_threshold_bytesis 1 MiB, so SLT inputs mostly collapse into one run and skip the spawned path too. The spawned path is then covered only by the unit and spill tests here. A multi-thread variant of the sort fuzz with a small threshold, asserting the same output as the inline path, would cover run counts, batch sizes and spill together. - Benchmark evidence. The description gives the before numbers for the 32-row reproducer but no after numbers (alive task count,
SELECT 1probe latency). The bot runs were ond1a59c1, which is earlier than the current headcc90f9e, and TPC-H sort inputs are not fragmented, so "no change" says little about the changed path. Please post the reproducer numbers on the current head, or add a small-batch case to an existing sort bench. - Budget scope. The
OnceLocklives on the plan and is shared by every clone, so two concurrent executions of one plan instance share the permits, and a re-run on another runtime keeps the first runtime's worker count untilreset_state. Probably acceptable, but worth one sentence in theSortExecrustdoc if intended.
| /// reservation when that batch is pulled. Reading one batch ahead returns | ||
| /// extra spill workspace for the merge's current batch and cursor. | ||
| /// Each run can hold one prefetched batch whose reservation is already released. | ||
| fn with_lookahead(stream: SendableRecordBatchStream) -> SendableRecordBatchStream { |
Contributor
There was a problem hiding this comment.
with_lookahead looks equivalent to stream.map(std::future::ready).buffered(2) wrapped in RecordBatchStreamAdapter. Buffered fills its queue to 2 before yielding, so each yielded batch has exactly one successor already pulled, same as the unfold state here. That would drop the hand-written state machine. I have not run it, so please confirm the reservation release order in test_sort_merge_lookahead_returns_cursor_workspace.
Contributor
Author
There was a problem hiding this comment.
You're right: this was premature optimization (I was concerned about extra per-batch allocations and bookkeeping). I've switched to buffered(2): it's much clearer, and the query-level benchmarks showed no performance difference. I also ran test_sort_merge_lookahead_returns_cursor_workspace and can confirm it passes with buffered(2).
| concurrency: &SortConcurrency, | ||
| ) -> Result<Vec<SendableRecordBatchStream>> { | ||
| let mut streams: Vec<_> = (0..runs.len()).map(|_| None).collect(); | ||
| let (sizes, shared_input_size) = self.sort_run_sizes(&runs)?; |
Contributor
There was a problem hiding this comment.
Two small points. (1) These three lines (sort_run_sizes, shared-input split, per-run split) are repeated in sort_run_streams. A shared helper returning the per-run reservations plus the optional shared one would keep the two paths from drifting. (2) prepare_sort_runs takes &mut self but only calls reservation.split(&self), like sort_run_streams(&self). &self is enough if the future stays Send (not checked).
Contributor
Author
There was a problem hiding this comment.
(1) I’d prefer keeping these few lines explicit, following @jayshrivastava's earlier feedback. A shared helper would add ownership complexity to preserve lazy splitting.
(2) Unfortunately, switching to &self would make the future non-Send because ExternalSorter isn’t Sync.
| } | ||
|
|
||
| #[tokio::test(flavor = "multi_thread", worker_threads = 2)] | ||
| async fn sort_execution_initializes_and_resets_concurrency() -> Result<()> { |
Contributor
There was a problem hiding this comment.
This test collects but never checks the output. Since the claim is that the spawned path matches the inline one, please assert the collected rows against the same plan run on a current-thread runtime. Use a payload column so equal keys are distinguishable.
Contributor
Author
There was a problem hiding this comment.
Added distinct payloads and checked each partition against explicit expected rows, including equal-key ordering, in b87cde6. I used explicit expectations rather than a second runtime to keep the test simple.
|
|
||
| /// Prepared runs retain their workspace loans until merging starts. Looking | ||
| /// ahead during the merge returns enough workspace for its batch and cursors. | ||
| async fn check_sort_preparation_workspace(consume_output: bool) -> Result<()> { |
Contributor
There was a problem hiding this comment.
The description says the merge receives runs in original order so equal-key ties are preserved, but nothing pins that. Both runs here are the same batch and only have the key column, so any run order gives identical output. A test with a payload column, where the earlier run finishes last (for example a much larger first run), would cover the try_buffer_unordered plus streams[index] reassembly.
Contributor
Author
There was a problem hiding this comment.
Added sort_preparation_restores_run_order in 00cf3cd. It supplies completed runs in reverse order and checks their original order is restored, with distinct payloads for equal keys. This avoids relying on batch sizes or timing.
geoffreyclaude
force-pushed
the
codex/bound-sort-preparation
branch
from
October 21, 2026 08:59
29b0afc to
b87cde6
Compare
geoffreyclaude
force-pushed
the
codex/bound-sort-preparation
branch
from
October 21, 2026 09:48
9e14937 to
00cf3cd
Compare
Keep the randomized sort configurations and compare each non-TopK query against forced inline and bounded parallel sorting, preserving small input batches in both arms. Compare complete results in row order after concatenating batches. Every projected column is an ORDER BY key, so equal keys have identical output. Retain the current-thread determinism test. This changes internal test coverage only; no public documentation update is needed.
Contributor
Author
|
run benchmark sort_tpch sort_tpch10 |
|
🤖 Benchmark running (GKE) | trigger CPU Details (lscpu)Comparing codex/bound-sort-preparation (416c2d6) to fbf7737 (merge-base) diff Run configurationrun benchmark sort_tpch10Results will be posted here when complete File an issue against this benchmark runner |
|
🤖 Benchmark running (GKE) | trigger CPU Details (lscpu)Comparing codex/bound-sort-preparation (416c2d6) to fbf7737 (merge-base) diff Run configurationrun benchmark sort_tpchResults will be posted here when complete File an issue against this benchmark runner |
|
🤖 Benchmark completed (GKE) | trigger Instance: Comparing codex/bound-sort-preparation (416c2d6) to fbf7737 (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 completed (GKE) | trigger Instance: Comparing codex/bound-sort-preparation (416c2d6) to fbf7737 (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 |
Contributor
Author
@comphead Thanks for the detailed review! Addressed all three of the above (as well as the inline comments):
|
Contributor
|
run benchmark tpch tpcds |
|
🤖 Benchmark running (GKE) | trigger CPU Details (lscpu)Comparing codex/bound-sort-preparation (416c2d6) to fbf7737 (merge-base) diff Run configurationrun benchmark tpchResults will be posted here when complete File an issue against this benchmark runner |
|
🤖 Benchmark running (GKE) | trigger CPU Details (lscpu)Comparing codex/bound-sort-preparation (416c2d6) to fbf7737 (merge-base) diff Run configurationrun benchmark tpcdsResults will be posted here when complete File an issue against this benchmark runner |
|
🤖 Benchmark completed (GKE) | trigger Instance: Comparing codex/bound-sort-preparation (416c2d6) to fbf7737 (merge-base) diff Run configurationrun benchmark tpchCPU Details (lscpu)Details
Resource Usagetpch — base (merge-base)
tpch — branch
File an issue against this benchmark runner |
|
🤖 Benchmark completed (GKE) | trigger Instance: Comparing codex/bound-sort-preparation (416c2d6) to fbf7737 (merge-base) diff Run configurationrun benchmark tpcdsCPU Details (lscpu)Details
Resource Usagetpcds — base (merge-base)
tpcds — branch
File an issue against this benchmark runner |
comphead
reviewed
Oct 21, 2026
comphead
left a comment
comphead
left a comment
Contributor
There was a problem hiding this comment.
Thanks for addressing the earlier points. A few design questions after a second pass.
| .into_iter() | ||
| .map(|stream| spawn_buffered(stream, 1)) | ||
| .collect(); | ||
| let concurrency = SortConcurrency::for_current_runtime(&self.sort_concurrency); |
Contributor
There was a problem hiding this comment.
The bound caps the task count, but I think the root cause is the run count. 187k runs is SF1 lineitem at one run per 32-row batch. With this PR the merge still has 187k streams, each with its own lookahead batch (the gap tracked in #26411).
The run count stays at one per input batch because coalesce_in_mem_batches_into_runs only runs on the output path and only for single-column sorts (coalesce_runs && self.expr.len() == 1 just above). Multi-column sorts and every spill keep one run per input batch.
Would a minimum run size, for example coalescing runs below batch_size rows regardless of column count, fix the reported case more directly? It would cut task count, merge fan-in and lookahead memory together. For comparison, the description reports at most 123 ms for 8192-row batches on main and 155 to 165 ms for this PR on 32-row batches. Those are separate runs, so this is only indicative. The semaphore would still make sense as a backstop for very large inputs. This could also be a follow-up if you prefer to keep this PR focused.
Contributor
Author
There was a problem hiding this comment.
I tested this on my fork, in [#22](geoffreyclaude#22). Coalescing made fragmented-input sorts 19 to 79% faster (depending on the benched scenario), with essentially unchanged performance on the ordinary preloaded wide sort. This definitely looks worth pursuing separately: I can open a follow-up PR as the gains are massive on small batches, and the code change is pretty simple.
To be clear though, our production case involved very large inputs with 8,192-row batches. The 32-row batches simply stress run count with a smaller dataset and they aren’t intended to represent typical workloads. Coalescing helps fragmented inputs, while the task bound remains necessary for the production case.
| // Ordered buffering would keep completed runs behind a slow earlier run, | ||
| // filling the window and preventing new jobs from starting. | ||
| let jobs = futures::stream::iter(jobs) | ||
| .try_buffer_unordered(concurrency.per_partition_limit); |
Contributor
There was a problem hiding this comment.
One job per run is what makes the per-partition window, the unordered buffering, the (index, stream) reassembly and collect_prepared_runs necessary. An alternative is to split the runs into at most per_partition_limit contiguous, byte-balanced chunks, with one job sorting each chunk sequentially.
- All jobs are in flight at once, so ordered
try_bufferedhas no head-of-line blocking and run order is preserved without index bookkeeping. - Per-job overhead (permit, spawn, join, boxed stream) is paid per chunk rather than per run. With 32-row runs that overhead is likely comparable to the sort itself. I have not measured this.
- The shared semaphore still bounds jobs across partitions.
Contributor
Author
There was a problem hiding this comment.
I also tested this on my fork, in #21. It reduced unrelated SELECT 1 probe delays (because fewer workers were busy initializing merges simultaneously), but made wide and uneven-work sorts 6 to 13% slower (because, even with byte-balanced chunks, some workers still finish early while others have substantial sorting left).
It does simplify result collection, but adds group balancing and more involved memory-reservation ownership.
Given those tradeoffs, I’d prefer keeping the current per-run design here.
| Self::collect_prepared_runs(run_count, jobs).await | ||
| } | ||
|
|
||
| /// Restore input order because the merge uses it to break equal-key ties. |
Contributor
There was a problem hiding this comment.
Small wording point. The merge does not preserve run order for equal keys, for two reasons.
- Each run is sorted with
lexsort_to_indices, which arrow documents as an unstable sort. StreamingMergeBuilderdefaults toenable_round_robin_tie_breaker: trueandin_mem_sort_streamdoes not override it, so equal keys are interleaved across runs by poll count.
What the reassembly does guarantee is deterministic output regardless of job completion order, which is still worth keeping. I would phrase this comment, the one on prepare_sort_runs and the description that way. This also corrects my earlier comment on spill_tests.rs, which framed it as tie stability. The new tests are fine as determinism checks.
Contributor
Author
There was a problem hiding this comment.
Agreed: I've updated the relevant comments in docs: clarify sort run ordering guarantee.
Contributor
Author
@comphead Thanks for your very pertinent suggestions! I tested both on my fork (thank you Codex!), then investigated where the scheduling gains came from:
My proposed path is to keep bounded per-run preparation in this PR, then pursue coalescing and cooperative merge initialization as separate improvements in two new follow-up PRs. I also corrected the ordering comments and description: restoring run order keeps output independent of task completion order. |
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
6 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?
Addresses the unbounded local-sort task population discussed in #23174.
Rationale for this change
SortExecsorts pieces of buffered input into independent runs, then merges them. On multithreaded Tokio runtimes, it currently creates one buffered producer task per run.A one-batch channel limits each producer's output, but the number of tasks still grows with the number of runs. With many input batches, this task buildup can stall unrelated queries on the same runtime. Large inputs can trigger this buildup even with normal-sized batches.
What changes are included in this PR?
Use bounded preparation for both final in-memory sorting and sorting before a spill:
SortExecshare a budget capped by both the runtime's worker count and DataFusion's estimate of available CPU parallelism. Each partition also bounds its pending preparations.Memory handoff: during merging, one-batch lookahead returns workspace for the merge's batches and cursors. Each run may hold one prefetched batch whose sort-side reservation has already been released, preserving the previous buffering bound.
Cancellation aborts outstanding jobs. Resetting execution state gives the new plan a fresh concurrency limit.
Single-batch and concatenate-and-sort paths continue inline, as do current-thread runtimes and execution outside Tokio. TopK is unchanged.
SortExecrustdoc describes the limit and its lifetime.What is the testing strategy for this PR?
416c2d6, the benchmark bot classified all 11 queries in each of sort_tpch (SF1) and sort_tpch10 (SF10) as "no change". These Parquet-based suites use the default 8,192-row batch target.A local TPC-H SF1 test uses 32-row batches to stress task growth with a modest dataset. With 4 Tokio workers and 4 partitions, two runs per revision gave:
SELECT 1probeAre there any user-facing changes?
Sorts with many runs have a bounded population of batch-sort jobs across partitions. Public configuration and API signatures are unchanged.
SortExec'selapsed_computenow includes local batch sorting and its memory accounting on both inline and spawned paths, while excluding waits for a permit or another task.EXPLAIN ANALYZEcan therefore report larger values.Durations from concurrent sort jobs are added together and can exceed the query's wall-clock duration.
This PR was prepared with AI assistance.