jayzhan211 commented on code in PR #25428:
URL: https://github.com/apache/datafusion/pull/25428#discussion_r4100377018
##########
datafusion/physical-plan/src/sorts/multi_level_merge.rs:
##########
@@ -393,6 +401,35 @@ impl MultiLevelMergeBuilder {
}
};
+ if self.size_intermediate_merges
+ && self.sorted_streams.is_empty()
+ && !self.sorted_spill_files.is_empty()
+ {
+ // If one intermediate merge can leave a final pass with
the
+ // admitted fan-in, merge only the runs needed to get
there.
+ // Keep the admitted reservation so sizing does not change
+ // reservation lifetime or admission under contention.
+ let needed = self.sorted_spill_files.len() + 1;
+ if needed < sorted_spill_files.len() {
+ let (first, first_limit) = &sorted_spill_files[0];
+ // A wider pending run or a smaller output batch limit
can
+ // change the next pass's fan-in and add intermediate
work.
+ // Apply this sizing heuristic only to uniform run
budgets.
+ let uniform_runs = sorted_spill_files
+ .iter()
+ .chain(&self.sorted_spill_files)
+ .all(|(file, limit)| {
+ file.max_record_batch_memory
+ == first.max_record_batch_memory
+ && limit == first_limit
Review Comment:
The uniformity guard only looks at the input runs, but the intermediate
output is rebatched up to `batch_size`. When the runs are shorter than
`batch_size`, the new run is wider than its inputs, the final pass no longer
fits in F slots, and sizing adds a pass instead of removing one. With 9 runs of
64 rows at `batch_size = 128` (same pool formula as
`intermediate_merge_sizing_preserves_final_output`), sizing off rewrites 448
rows (7 runs) and sizing on rewrites 576 rows (3 runs, then 6)
```suggestion
&& limit == first_limit
&& file.max_batch_rows == *limit
```
Please also add the short-run case as a test that asserts sizing never
rewrites more rows than the unsized path
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
