Repository navigation
fix: Batch size limit in re-spill compounds - #23286
Merged
Merged
Conversation
github-actions
Bot
added
physical-plan
# Conflicts: # datafusion/physical-plan/src/sorts/multi_level_merge.rs
rluvaton
reviewed
Jul 13, 2026
Comment on lines
+290
to
+299
| let (spill_file, _) = self.sorted_spill_files.remove(0); | ||
|
|
||
| // Not reserving any memory for this disk as we are not holding it in memory | ||
| let output_stream = self | ||
| .spill_manager | ||
| .read_spill_as_stream(spill_file.file, None)?; | ||
|
|
||
| Ok(MergeStep::Stream(self.observe_output(output_stream))) | ||
| Ok(MergeStep::Stream { | ||
| stream: self.observe_output(output_stream), | ||
| batch_size_limit: self.batch_size, |
Member
There was a problem hiding this comment.
the stream should return the spill batch size
rluvaton
reviewed
Jul 13, 2026
Comment on lines
+354
to
+362
| // Cap the merge output at the smallest limit among the runs we're | ||
| // about to merge. Runs that were shrunk for skew carry a smaller limit, | ||
| // if none do, every run carries `self.batch_size` and the merge runs at | ||
| // the full batch size. The output stream is tagged with the same limit | ||
| // (see the `MergeStep::Stream` returns below) so a re-spilled | ||
| // intermediate run stays shrunk and won't rebuild an oversized batch on | ||
| // a later pass. | ||
| let output_batch_size = | ||
| min_batch_size_limit(&sorted_spill_files).unwrap_or(self.batch_size); |
Member
There was a problem hiding this comment.
We are already iterating over the spill files in for (spill, _) in sorted_spill_files {
can you instead inline that?
rluvaton
reviewed
Jul 13, 2026
Comment on lines
+973
to
+1027
| /// Same as [`respill_halves_the_merge_output_batch_size`], but under a budget tight | ||
| /// enough that *both* runs must be re-spilled before the merge fits - the scenario | ||
| /// where the batch-size reduction could compound. Because the reduction is tracked | ||
| /// per-run (each run capped at half) rather than by halving the global batch size on | ||
| /// every split, the merged output is emitted in 4096-row batches - half, not a | ||
| /// quarter. A global-halving implementation would have halved once per re-spill and | ||
| /// emitted 2048-row batches. | ||
| #[tokio::test] | ||
| async fn respilling_two_skewed_runs_halves_the_output_without_compounding() | ||
| -> Result<()> { | ||
| let env = Arc::new(RuntimeEnv::default()); | ||
| let schema = test_schema(); | ||
| let spill_manager = build_spill_manager(&env, &schema); | ||
|
|
||
| let n: i64 = 16384; | ||
| let f0 = make_sorted_spill_file(&spill_manager, &schema, (0..n).collect()); | ||
| let f1 = make_sorted_spill_file(&spill_manager, &schema, (0..n).collect()); | ||
| let m = f0.max_record_batch_memory.max(f1.max_record_batch_memory); | ||
|
|
||
| // 2.5*m is tight enough that even after halving one run the two still don't | ||
| // fit, so *both* runs are re-spilled once before the merge succeeds. (3.5*m, | ||
| // as in the single-split test, would let the pair fit after one split.) This | ||
| // is exactly the scenario where a compounding, global-halving implementation | ||
| // would drive the output batch size down to a quarter. | ||
| let initial_batch_size = 8192; | ||
| let pool: Arc<dyn MemoryPool> = Arc::new(GreedyMemoryPool::new(m * 5 / 2)); | ||
|
|
||
| let builder = build_merge_builder( | ||
| spill_manager, | ||
| Arc::clone(&schema), | ||
| vec![f0, f1], | ||
| &pool, | ||
| initial_batch_size, | ||
| ); | ||
| let stream = builder.create_spillable_merge_stream(); | ||
| let batches: Vec<RecordBatch> = stream.try_collect().await?; | ||
|
|
||
| // All rows are still present. | ||
| let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum(); | ||
| assert_eq!(total_rows, (2 * n) as usize); | ||
|
|
||
| // Each run was re-spilled once, so each is capped at half the original batch | ||
| // size and the merge caps its output at that half - NOT a quarter. A global | ||
| // halving-per-split implementation would have emitted 2048-row batches here. | ||
| let expected_batch_size = initial_batch_size / 2; | ||
| let max_batch_rows = batches.iter().map(|b| b.num_rows()).max().unwrap_or(0); | ||
| assert_eq!( | ||
| max_batch_rows, expected_batch_size, | ||
| "two re-spills must halve (not quarter) the output: expected \ | ||
| {expected_batch_size}-row batches, got a largest batch of \ | ||
| {max_batch_rows} rows" | ||
| ); | ||
|
|
||
| Ok(()) | ||
| } |
Member
There was a problem hiding this comment.
I would prefer more integration test like and not access private functions but I see that all the other tests here are like that so we can do that in a separate pr
github-merge-queue
Bot
removed this pull request from the merge queue due to failed status checks
Jul 13, 2026
kosiew
pushed a commit
to kosiew/datafusion
that referenced
this pull request
Aug 12, 2026
This fixes an issue in my previous design for the re-spill in sort - whenever we half a stream, we half self.batch_size, even if its not the same stream being halved, so for each re-spill we get a smaller and smaller output batch, which really hurts performance in the merge. We can avoid it by just keeping the batch_size_limit in the spilled file struct. then we just use the lowest limit whenever we merge N streams, making it dynamic and much more robust. Has tests. Does not break API --------- Co-authored-by: Raz Luvaton <16746759+rluvaton@users.noreply.github.com>
This was referenced Nov 1, 2026
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.
This fixes an issue in my previous design for the re-spill in sort - whenever we half a stream, we half self.batch_size, even if its not the same stream being halved, so for each re-spill we get a smaller and smaller output batch, which really hurts performance in the merge.
We can avoid it by just keeping the batch_size_limit in the spilled file struct. then we just use the lowest limit whenever we merge N streams, making it dynamic and much more robust.
Has tests.
Does not break API