viirya commented on code in PR #25428:
URL: https://github.com/apache/datafusion/pull/25428#discussion_r4109889402
##########
datafusion/physical-plan/src/sorts/multi_level_merge.rs:
##########
@@ -393,6 +411,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
Review Comment:
This keeps the full `k`-run grant while the pass merges only `R + 1` runs,
then releases it when the intermediate stream is dropped after the write. The
final pass then re-admits from an empty reservation at line 387. So the grant
is held during the pass that doesn't need all of it, and released just before
the pass that does.
Since intermediate spill merges use `with_bypass_mempool()`, the grant is
only accounting. For a sized pass it could stay in `self.reservation` instead
of being attached to the intermediate stream. The pool would see exactly the
same reserved bytes during the merge, and the next `self.reservation.take()`
would hand them to the final admission, which already treats pre-reserved bytes
as covering the first runs via `try_grow_reservation_to_at_least`. The final
pass would then get its `k` slots whenever the new run is no wider than its
inputs, which the uniform guard already expects. The other option is to leave
one run of slack in `needed`, at the cost of some of the saving.
##########
datafusion/physical-plan/src/sorts/multi_level_merge.rs:
##########
@@ -985,6 +1033,148 @@ mod tests {
)
}
+ #[rstest::rstest]
+ #[case::full_batches(9, 128, 3)]
+ #[case::unchanged_selection(13, 128, 7)]
+ #[case::short_batches(9, 64, 7)]
+ #[tokio::test]
+ async fn intermediate_merge_sizing_preserves_final_output(
Review Comment:
These cases run against a fresh pool with no other consumer, so the final
admission always seats the same fan-in as the first. Could you add a case where
a competing reservation takes part of the released grant between the
intermediate write and the final admission, asserting that sizing rewrites no
more rows than the unsized path? The `HandoffPool` probe from the earlier
revision had the right shape: a `shrink` hook that lets a contender grow when
the merge releases memory.
--
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]
