Conversation
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## dev/chao/codex/pr25172-accounting-base #25428 +/- ##
==========================================================================
- Coverage 81.94% 81.94% -0.01%
==========================================================================
Files 1136 1135 -1
Lines 429614 428452 -1162
Branches 429614 428452 -1162
==========================================================================
- Hits 352067 351106 -961
+ Misses 56491 56386 -105
+ Partials 21056 20960 -96 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
comphead
left a comment
comphead
left a comment
There was a problem hiding this comment.
Reviewed at 6987df2 against the stacked base b0752560. Built and ran the sort/merge suite locally: 162 sorts:: tests pass.
The core idea is right. An intermediate pass feeds a spill writer, not the aggregate, so replay headroom there is pure waste, and keeping the inputs plus the grant alive so a failed wider write can fall back is a clean way to bound the extra peak-disk risk. Scoping is good too: merge_pool is only set by sort.rs and reserve_replay_headroom only by the aggregate streams, so external sort is untouched by this.
Three things I would want resolved before merge:
aggregate_memory_spill.sltnow pinsmax_spill_merge_fan_in = 2for the 4-partition case. Replay headroom also implicitly bounded how much of a shared pool one partition's merge could hold; widening removes that bound for intermediate passes while partitions run concurrently. Details inline.- The flush at every input batch boundary is unconditional and measurably cuts intermediate output batch size roughly in half even when runs are written at
batch_size. Numbers inline. - No isolated benchmark for this change. The only wall-clock data in the description compares the whole prerequisite stack against
mainand shows it 7-15% slower onexternal_aggrQ2 (832 to 892 ms at 8 MiB, 717 to 826 ms at 16 MiB). The base branch is already pinned atb0752560, so an A/B of exactlyb0752560vs this head onexternal_aggrat both pool sizes is cheap and is the only number that attributes the cost. Point 2 is a plausible contributor.
On fuzzing: none is added. datafusion/core/tests/fuzz_cases/spilling_fuzz_in_memory_constrained_env.rs already drives memory-constrained aggregate spilling with varying batch sizes, oversized batches and a competing consumer. Extending it would cover the widened path with randomized inputs and would also replace several of the hand-built fixtures in the new test module.
The rest is reuse and test dedup, inline.
| # Bound merge buffers so each partition can allocate replay state while the | ||
| # other partitions retain aggregate state in the shared greedy memory pool. | ||
| statement ok | ||
| SET datafusion.runtime.max_spill_merge_fan_in = 2 |
There was a problem hiding this comment.
P1. This is the part I would push back on. Reserving replay headroom also implicitly bounded how much of a shared pool a single partition's merge could hold, and at target_partitions = 4 all four partitions merge concurrently against one greedy pool. Widening removes that bound for intermediate passes, and the change here caps fan-in rather than capping the share.
Two costs:
- Case G no longer covers default fan-in at 4 partitions, which is the configuration most likely to hit the new behaviour.
- The condition is timing dependent, so the cap hides it rather than closing it. I removed both statements and ran the file 12 times on a 16-core machine: 12/12 green. It reproduces somewhere but not everywhere, which is the shape of a contention window, not of a deterministic budget error.
Could widening instead be bounded by a share of the pool, or skipped while other consumers hold state, so the SLT can keep default fan-in? If the cap is genuinely the right fix, the commit message should say which failure it addresses, because test: bound merge fan-in for parallel aggregate spilling with an empty body does not record that.
| ); | ||
|
|
||
| drop(timer); | ||
| if self.flush_on_input_batch_boundary { |
There was a problem hiding this comment.
P2. This flushes on every input batch boundary, so a cut happens whenever any one of N inputs exhausts a batch. In steady state the output batch size converges to the input batch size, independent of batch_size.
Measured by driving StreamingMergeBuilder directly, 8 runs x 10 batches, batch_size = 8192:
| input batch | flush off | flush on |
|---|---|---|
| 1000 rows | 10 batches, avg 8000 rows | 80 batches, avg 1000 rows |
| 8192 rows | 80 batches, avg 8192 rows | 150 batches, avg 4369 rows |
So even when runs are written at the full batch_size, the intermediate run ends up with about 1.9x the batches at about 53% the rows, and the final pass re-reads all of them. That is the IPC framing cost the description mentions, quantified, and it is paid on every widened pass.
The invariant you need is "pending output must fit the reserved workspace", not "at most one batch per input". BatchBuilder already tracks batches_mem_used and the per-merge grant is known at construction, so the flush could be conditional on approaching that budget and keep full batches the rest of the time. That would also make the short-wide-batch case the only one that pays.
| .await; | ||
| drop(stream); | ||
| // A successful write drops the backups before the next selection. | ||
| if let Some(retry) = retry |
There was a problem hiding this comment.
The retry fires on any Err, but only a disk-capacity failure motivates it. A memory, I/O or cancellation error pays a full re-read plus re-merge and then surfaces the second error, with the first silently dropped. Worth gating on the quota error specifically, or at least chaining the original error as context so the real cause survives.
| mem::swap(&mut self.reservation, &mut memory_reservation); | ||
| let original_count = spills.len(); | ||
| let original_memory = memory_reservation.size(); | ||
| if self.reserve_replay_headroom |
There was a problem hiding this comment.
This re-implements get_sorted_spill_files_to_merge. effective_spill_merge_fan_in, the get_reserved_bytes_for_record_batch_size(x, x) * buffer_len formula, the fan-in cap and the drain all already exist there. It is the same loop with check_headroom = false plus a "leave one run" stop condition, and that function already carries an allow_minimum_without_headroom knob for exactly this kind of variation. Folding it in keeps admission accounting in one place and removes the second counter, the second drain and the duplicated fan-in lookup.
Style drift worth noting while the two exist side by side: the original loop grows to an absolute total_needed, this one grows by a delta. They agree only because the shrink just above left the reservation at accepted_memory.
Separately: this also runs when allow_minimum_without_headroom is true, that is, immediately after the pool failed to seat two streams with headroom. Trying to grow past the just-admitted minimum in that state is contradictory. Suggest adding && !allow_minimum_without_headroom.
| let reservation = Arc::new(memory_reservation); | ||
| let retry = | ||
| (spills.len() > original_count).then(|| IntermediateMergeRetry { | ||
| spills: spills |
There was a problem hiding this comment.
read_spill_as_stream takes Arc<dyn SpillFile>, so merge_selected_runs can take &[(SortedSpillFile, usize)] and Arc::clone internally. The caller then keeps spills and moves it straight into the guard, and this hand-written clone plus the duplicate Vec both go away. Deriving Clone on SortedSpillFile (an Arc and a usize) would also do it.
Minor, same area: retry is None at four of the five MergeStep::Stream construction sites. Holding it on the builder as self.pending_retry would leave those four unrelated arms untouched.
| self.release_unused_memory(); | ||
| } | ||
|
|
||
| fn retain_current_batches(&mut self, keep_consumed: bool) { |
There was a problem hiding this comment.
retain_current_batches(true) at the pre-existing call site reads as "keep the current batches", which is the opposite of what the flag selects. Naming the parameter for the behaviour it adds (drop_consumed) keeps the old site readable, or split into two small named methods over a shared inner.
| input_bytes: usize, | ||
| } | ||
|
|
||
| fn replay_merge_builder( |
There was a problem hiding this comment.
mod tests in the parent file already provides test_schema, build_spill_manager, make_sorted_spill_file and build_merge_builder, and already holds replay-headroom tests (replay_headroom_splits_an_oversized_first_run, replay_headroom_allows_only_an_indivisible_minimum). This module re-creates all four helpers. Either put the new cases in that module, or make the helpers pub(super) and import them. Two replay-headroom test modules in one file will be confusing to the next person touching this.
| } | ||
|
|
||
| #[tokio::test] | ||
| async fn intermediate_merge_preserves_short_batch_replay() -> Result<()> { |
There was a problem hiding this comment.
This and replay_headroom_preserves_indivisible_run_batch_limits above share a fixture (1-row Utf8 runs of 1024-char values) and differ only in run count and pool size, and the latter in turn overlaps the pre-existing replay_headroom_allows_only_an_indivisible_minimum. One rstest matrix over (run_count, memory_batches, expect_split) would cover all three at a third of the length.
Wider point: datafusion/core/tests/fuzz_cases/spilling_fuzz_in_memory_constrained_env.rs already has RunTestWithLimitedMemoryArgs with MemoryBehavior::TakeAllMemoryAtTheBeginning, varying and oversized record batches, and assert_all_output_batches_roughly_match_batch_size_conf. Those are the same shapes this module hand-builds, including the 90-line HandoffPool in intermediate_merge_keeps_admitted_buffers, and extending it would add randomized coverage this PR currently has none of.
| // rewriting every input row twice before returning the final merge. | ||
| let (spill_count, spilled_rows, spilled_bytes) = | ||
| merge_replay_runs(32, 256, 32).await?; | ||
| println!( |
There was a problem hiding this comment.
Leftover debug output. All three values are already in scope for the asserts below, so folding them into the existing assert message keeps the diagnostic without the noise on every run.
| async fn intermediate_merge_flushes_before_replacing_input_batch( | ||
| #[case] data_type: DataType, | ||
| #[values(false, true)] row_cursor: bool, | ||
| #[values(false, true)] round_robin: bool, |
There was a problem hiding this comment.
round_robin only affects tie-break ordering among equal keys and does not interact with batch-boundary flushing, but it doubles a matrix that is already 3 data types x 2 cursor kinds x 2 fetch values, each running an inner for flush in [false, true], so 48 merges. Dropping it halves that with no coverage loss. The flush == false branch is also asserting pre-existing batching behaviour rather than anything this PR changes.
3 participants
Which issue does this PR close?
Follow-up to the review of #25383, which suggested skipping replay headroom for intermediate multi-level merge passes.
Depends on #25383 and #25172. This draft is stacked on a temporary upstream branch at #25172's exact head,
b0752560ff47f36ce8896574829082e51f79088c, so the diff contains only this follow-up. Merge those prerequisites in that order, then rebase and retarget this PR tomain.Rationale for this change
An intermediate spill merge writes its output back to disk before aggregation resumes. Reserving equal memory for aggregate replay during that pass limits how many runs it can merge, causing more intermediate files and repeated spill I/O.
For example, with 32 equal runs and enough memory for eight buffered inputs, the existing policy admits four inputs per pass. The measured fixture writes ten intermediate files and rewrites every input row twice. Widening intermediate passes reduces this to four files and one rewrite per row. The final merge still needs space for the aggregate consuming its output.
The accounting prerequisite matters: widening intermediate merges on #25383 alone makes both existing legacy aggregate replay tests exceed their 8 KiB reservation limit. Legacy aggregation retains a sibling reservation, which the old
FairSpillPooladmission check overlooks. #25172 accounts for that retained memory. The original peak-memory assertions are unchanged.What changes are included in this PR?
The merge first selects inputs with replay headroom. If that selection requires an intermediate pass, it can admit additional spilled runs using the remaining memory. It keeps the original reservation and read-ahead setting, respects the configured fan-in limit, and leaves at least one run for a subsequent pass. Final selections and split requests retain the existing replay-headroom policy.
Wider passes emit pending rows before replacing an exhausted input batch and discard consumed batches, including their dictionaries. This prevents many short input batches from accumulating into an oversized intermediate batch. Ordinary merges retain their existing batching.
A wider merge can require more peak disk space despite reducing total writes. Its input files and reservation remain available until the intermediate writer finishes. If that write fails, the merge retries its original smaller selection using the retained grant and disables further widening. This also works with compressed spills and custom backends without estimating their output sizes. Widening applies only to spilled inputs, which can be reread.
What is the testing strategy for this PR?
Regression tests cover the three failures reproduced during review: short wide batches exhausting later split workspace, a larger merge exceeding a disk quota that admits the smaller merge, and losing an admitted reservation to a competing consumer. They also cover contention after read-ahead reduction, sufficient and insufficient disk quotas under both Tokio runtime modes, dictionary payload retention, primitive/string/view merge paths, fetch limits, sorted output, and cleanup. Existing tests cover final replay allocations, configured fan-in, split retries, and unchanged aggregate peak-memory limits.
The deterministic fixture uses 32 interleaved Int64 runs, 256 rows per run, a 256-row output target, and a pool of 32 times the largest input batch's recorded memory. It measures actual intermediate spill writes, excluding the original input spills:
b0752560)These counters describe this fixture's cumulative spill work. They do not measure elapsed time or peak disk usage. Input-boundary flushing adds some IPC framing while preserving the reduction in merge passes and rewritten rows.
Fresh local validation on the final source passed:
cargo fmt --all,cargo clippy --all-targets --all-features -- -D warnings, and the complete./dev/rust_lint.shsuite, including private Rust documentation and local Markdown links.Validation used Rust 1.98.1 and the prerequisites' unchanged dependency lockfile. The spill-work comparison used a separate binary at the exact stacked base. Independent rereview covered merge semantics, resource lifetime, fallback ordering, and ordinary-sort compatibility.
A bounded
external_aggrQ2 comparison also passed 24 timed iterations and four result/spill diagnostics on two million synthetic distinct pairs, using one partition, legacy aggregation, FairSpillPool, and 8/16 MiB limits. Both builds used Rust 1.97.0 andrelease-nonlto. The baseline was historical main commit22651d24, this branch's main merge base; the candidate was24df38b0eincluding both prerequisites.All results were exactly 2,000,000 and both revisions stayed within their memory limits. Spill counts and spilled-row totals matched in this workload. The candidate was slower in this small shared-VM sample; these timings compare the entire prerequisite stack and do not isolate this follow-up. No wall-clock speedup is claimed.
Are there any user-facing changes?
Spilling aggregations can perform fewer intermediate merges and spill writes. A failed wider spill retries the smaller merge, which may add I/O before completing or returning an error. Final replay retains its existing headroom policy and indivisible-row fallback. SQL results, configuration, and public APIs are unchanged.