fix: bound parallel batch sorting in SortExec by geoffreyclaude · Pull Request #26242 · apache/datafusion

geoffreyclaude · 2026-10-12T15:58:45Z

Which issue does this PR close?

Addresses the unbounded local-sort task population discussed in #23174.

Rationale for this change

SortExec sorts 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:

Wait before spawning. Preparations acquire a shared permit before creating a Tokio task. All partitions of one SortExec share 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.
Sort, return, exit. Each job sorts one run, reserves its complete output, and returns a stream holding the batches and their reservations. The job then exits and releases its permit.
Merge the prepared streams. Output stays reserved until its batches are read. The merge receives the streams in their original run order, keeping its output independent of task completion order.

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. SortExec rustdoc describes the limit and its lifetime.

What is the testing strategy for this PR?

Memory and output: a competing-consumer regression checks that prepared output stays reserved while released scratch space can be reused. Tight-workspace tests check complete merging and cleanup when unread output is dropped.
Task lifetime and timing: focused tests cover permit ownership, cancellation, errors, the shared limit and its reset, and compute recording before output consumption.
SQL benchmarks: at 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:

Measurement	Upstream base	PR head
Sampled peak live Tokio tasks	~187k	9
Slowest `SELECT 1` probe	2.83–2.89 s	155–165 ms

Are 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's elapsed_compute now includes local batch sorting and its memory accounting on both inline and spawned paths, while excluding waits for a permit or another task. EXPLAIN ANALYZE can 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.

codecov-commenter · 2026-10-12T18:34:20Z

Codecov Report

❌ Patch coverage is 88.98305% with 26 lines in your changes missing coverage. Please review.
✅ Project coverage is 82.91%. Comparing base (fbf7737) to head (cc90f9e).
⚠️ Report is 2 commits behind head on main.

Files with missing lines	Patch %	Lines
datafusion/physical-plan/src/sorts/sort.rs	88.98%	3 Missing and 23 partials ⚠️

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.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

jayshrivastava · 2026-10-14T18:01:48Z

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


I don't think this comment adds much value.

indeed, removed.

jayshrivastava · 2026-10-14T18:06:26Z

+                        // current batch, returning workspace before the merge
+                        // reserves memory for that batch and its cursor.
+                        let next = stream.next().await;
+                        Some((batch, (next, stream)))


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

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.

jayshrivastava · 2026-10-14T18:14:28Z

+#[derive(Debug)]
+struct SortConcurrency {
+    /// Bounds preparation futures buffered by each partition.
+    limit: usize,


Hmm I guess you need this bc there's no API for total semaphore size?

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.

jayshrivastava · 2026-10-14T18:20:52Z

+                .prepare_sort_runs(runs, concurrency)
+                .with_elapsed_compute(elapsed_compute.clone())
+                .await?;
+            timer.restart();


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.

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.

Oh neat, with_elapsed_compute is quite handy

jayshrivastava · 2026-10-14T18:28:25Z

+        .sum::<usize>()
+        - input_bytes;
+    let first_bytes = get_record_batch_memory_size(&sorted_batches[0]);
+    assert!(loan > first_bytes);


Can you comment why it's bigger? I guess it's because some 4096 size buffers are re-used but not filled?

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

jayshrivastava · 2026-10-14T18:29:46Z

+        assert!(alive <= 1, "{alive} sort tasks remain before merging");
+        assert_eq!(
+            sorter.merge_pool.borrow(usize::MAX).size(),
+            workspace - 2 * (loan - first_bytes)


Adding comments around the math would make it easier to reason about

Should hopefully be a tiny bit clearer now!

geoffreyclaude · 2026-10-15T06:49:31Z

run benchmark sort_tpch sort_tpch10

adriangbot · 2026-10-15T06:52:23Z

🤖 Benchmark running (GKE) | trigger
Instance: c4a-highmem-16 (12 vCPU / 65 GiB) | Linux bench-c6186268738-3434-bwhsq 6.12.94+ #1 SMP Fri Aug 21 08:00:16 UTC 2026 aarch64 GNU/Linux

CPU Details (lscpu)

Architecture:                            aarch64
CPU op-mode(s):                          64-bit
Byte Order:                              Little Endian
CPU(s):                                  16
On-line CPU(s) list:                     0-15
Vendor ID:                               ARM
Model name:                              Neoverse-V2
Model:                                   1
Thread(s) per core:                      1
Core(s) per cluster:                     16
Socket(s):                               -
Cluster(s):                              1
Stepping:                                r0p1
BogoMIPS:                                2000.00
Flags:                                   fp asimd evtstrm aes pmull sha1 sha2 crc32 atomics fphp asimdhp cpuid asimdrdm jscvt fcma lrcpc dcpop sha3 sm3 sm4 asimddp sha512 sve asimdfhm dit uscat ilrcpc flagm sb paca pacg dcpodp sve2 sveaes svepmull svebitperm svesha3 svesm4 flagm2 frint svei8mm svebf16 i8mm bf16 dgh rng bti
L1d cache:                               1 MiB (16 instances)
L1i cache:                               1 MiB (16 instances)
L2 cache:                                32 MiB (16 instances)
L3 cache:                                80 MiB (1 instance)
NUMA node(s):                            1
NUMA node0 CPU(s):                       0-15
Vulnerability Gather data sampling:      Not affected
Vulnerability Indirect target selection: Not affected
Vulnerability Itlb multihit:             Not affected
Vulnerability L1tf:                      Not affected
Vulnerability Mds:                       Not affected
Vulnerability Meltdown:                  Not affected
Vulnerability Mmio stale data:           Not affected
Vulnerability Reg file data sampling:    Not affected
Vulnerability Retbleed:                  Not affected
Vulnerability Spec rstack overflow:      Not affected
Vulnerability Spec store bypass:         Mitigation; Speculative Store Bypass disabled via prctl
Vulnerability Spectre v1:                Mitigation; __user pointer sanitization
Vulnerability Spectre v2:                Mitigation; CSV2, BHB
Vulnerability Srbds:                     Not affected
Vulnerability Tsa:                       Not affected
Vulnerability Tsx async abort:           Not affected
Vulnerability Vmscape:                   Not affected

Comparing codex/bound-sort-preparation (d1a59c1) to e732212 (merge-base) diff

Run configuration

run benchmark sort_tpch

Results will be posted here when complete

File an issue against this benchmark runner

geoffreyclaude · 2026-10-22T12:00:18Z

Thanks for addressing the earlier points. A few design questions after a second pass.

@comphead Thanks for your very pertinent suggestions! I tested both on my fork (thank you Codex!), then investigated where the scheduling gains came from:

#22: coalesce small batches. Fragmented-input sorts improved by 19 to 79%. This looks worth pursuing separately. Large inputs with normal-sized batches still need the task bound here.
#21: fixed groups of runs. Unrelated query delays improved substantially, but wide and uneven-work sorts became 6 to 13% slower. It simplifies result collection while adding group balancing and more involved reservation ownership. I’d prefer keeping the current per-run design.
#23: yield during merge row-key initialization. Traces pointed to merge initialization occupying all workers. Yielding directly there preserved the responsiveness gains without grouping runs: on wide input with 8,192-row batches, the median worst probe delay per run fell from 49.6 to 4.5 ms, with sort time changing by +0.7% (within noise).

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.

Share a worker-sized permit budget across sort partitions and acquire permits before spawning. Each job sorts and accounts for one run, keeping its output reserved until the merge consumes it. Preserve run order, cancellation, and merge-time workspace reuse.

Check permit ownership, cancellation, errors, shared-limit wiring and reset. Verify inline and spawned sorting record compute before output consumption. Keep prepared output reserved against a competing consumer and exercise tight-workspace merging and unread-output cleanup. Preserve spill coverage on both current-thread and multithreaded runtimes.

Use the smaller of the Tokio worker count and DataFusion's available parallelism estimate for the shared sort-preparation budget. Explain why CPU-bound jobs need both limits and update the SortExec documentation.

Return per-run and shared input sizes from sort_run_sizes, and create reservations at both call sites. Preserve sizing, lazy run construction, and shared-input lifetimes.

Document the independent-input and common-input cases explicitly, including which budget covers private input buffers and when input charges are released.

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.

comphead

Thanks for testing both alternatives on the fork and for the numbers. Keeping the per-run design here makes sense to me given the 6 to 13% slowdown from chunking. The ordering comments now match what the merge guarantees, and the tpch, tpcds and sort_tpch runs show nothing beyond noise.

I think the PR is good to go

comphead

Thanks @geoffreyclaude lets give it a try

alamb · 2026-10-27T21:01:49Z

🚀

geoffreyclaude · 2026-10-29T14:32:07Z

Thanks @geoffreyclaude lets give it a try

Thanks @comphead for your extremely useful reviews! I'll open the follow ups once this is merged and tag you.

2010YOUY01 · 2026-11-01T03:29:26Z

-                        check_equality_of_batches(expected, &results).unwrap();
+                        // Every projected column is an ORDER BY key, so tied
+                        // rows have identical output and can be compared in order.
+                        assert_eq!(expected, &results);


This line change causes false positives for this test, it's getting fixed in

fix: Fix flaky sort fuzzer test #26983

Just want to double check it's only for cleanup here, and there is no other reason requires assert_eq

github-actions Bot added the physical-plan Changes to the physical-plan crate label Oct 12, 2026

geoffreyclaude force-pushed the codex/bound-sort-preparation branch 13 times, most recently from e2e4c5e to 9d0897a Compare October 13, 2026 15:34

github-actions Bot added the optimizer Optimizer rules label Oct 13, 2026

geoffreyclaude force-pushed the codex/bound-sort-preparation branch 4 times, most recently from 67093a0 to 1eacb12 Compare October 14, 2026 09:49

github-actions Bot removed the optimizer Optimizer rules label Oct 14, 2026

geoffreyclaude marked this pull request as ready for review October 14, 2026 14:07

jayshrivastava reviewed Oct 14, 2026

View reviewed changes

geoffreyclaude force-pushed the codex/bound-sort-preparation branch 2 times, most recently from 4ba6d18 to 96f3f8a Compare October 14, 2026 22:25

github-actions Bot added the auto detected api change Auto detected API change label Oct 14, 2026

geoffreyclaude force-pushed the codex/bound-sort-preparation branch from 96f3f8a to d1a59c1 Compare October 15, 2026 06:47

This was referenced Oct 22, 2026

Experiment: coalesce small sort batches into larger runs geoffreyclaude/datafusion#22

Open

Experiment: yield during merge row-key initialization geoffreyclaude/datafusion#23

Open

geoffreyclaude and others added 12 commits October 27, 2026 10:00

fix: cap sort preparation by available CPU parallelism

f0ce6ef

Use the smaller of the Tokio worker count and DataFusion's available parallelism estimate for the shared sort-preparation budget. Explain why CPU-bound jobs need both limits and update the SortExec documentation.

refactor: keep sort reservation ownership with callers

69fb4e0

Return per-run and shared input sizes from sort_run_sizes, and create reservations at both call sites. Preserve sizing, lazy run construction, and shared-input lifetimes.

docs: explain shared and per-run sort reservation sizes

8178f21

Document the independent-input and common-input cases explicitly, including which budget covers private input buffers and when input charges are released.

docs: clarify sort buffer deduplication boundary

42309e4

refactor: simplify sort run lookahead with buffered

c1a35a7

test: assert sorted output when initializing sort concurrency

f25c93e

test: verify sort run ordering after out-of-order preparation

bb1ba6d

docs: clarify the scope of the shared sort task limit

e015ef7

docs: clarify sort run ordering guarantee

aba15cd

geoffreyclaude force-pushed the codex/bound-sort-preparation branch from 8d5b96d to aba15cd Compare October 27, 2026 09:04

comphead reviewed Oct 27, 2026

View reviewed changes

comphead approved these changes Oct 27, 2026

View reviewed changes

geoffreyclaude mentioned this pull request Oct 29, 2026

Cherry-pick apache/datafusion#26242 DataDog/datafusion#192

Open

comphead added this pull request to the merge queue Oct 30, 2026

Merged via the queue into apache:main with commit 676c003 Oct 30, 2026
45 checks passed

github-actions Bot added the v56.0.0 label Oct 30, 2026

This was referenced Oct 30, 2026

refactor: use iterator lookahead for materialized sort runs #26930

Draft

perf: coalesce small input batches before sorting #26931

Draft

perf: yield during sorted merge initialization #26933

Draft

2010YOUY01 mentioned this pull request Nov 1, 2026

fix: Fix flaky sort fuzzer test #26983

Open

2010YOUY01 reviewed Nov 1, 2026

View reviewed changes

Conversation

geoffreyclaude commented Oct 12, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Which issue does this PR close?

Rationale for this change

What changes are included in this PR?

What is the testing strategy for this PR?

Are there any user-facing changes?

Uh oh!

codecov-commenter commented Oct 12, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

geoffreyclaude Oct 15, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

geoffreyclaude commented Oct 15, 2026

Uh oh!

adriangbot commented Oct 15, 2026

Uh oh!

geoffreyclaude commented Oct 22, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

comphead left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

comphead left a comment

Choose a reason for hiding this comment

Uh oh!

alamb commented Oct 27, 2026

Uh oh!

geoffreyclaude commented Oct 29, 2026

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

8 participants

geoffreyclaude commented Oct 12, 2026 •

edited

Loading

codecov-commenter commented Oct 12, 2026 •

edited

Loading

geoffreyclaude Oct 15, 2026 •

edited

Loading

geoffreyclaude commented Oct 22, 2026 •

edited

Loading

comphead left a comment •

edited

Loading