dwsmith1983 commented on code in PR #6336:
URL: https://github.com/apache/datafusion-comet/pull/6336#discussion_r4380872689
##########
native/shuffle/src/writers/local/spill.rs:
##########
@@ -214,17 +217,63 @@ impl PartitionedSpill {
fn ensure_spill_file_created(
&mut self,
runtime: &RuntimeEnv,
+ buffers: Option<&MemoryReservation>,
) -> datafusion::common::Result<()> {
if self.spill_file.is_none() {
let temp_file = runtime
.disk_manager
.create_tmp_file("shuffle writer spill")?;
- let writer = BufWriter::with_capacity(self.write_buffer_size,
temp_file.open_writer()?);
+ // Open before charging, so a failed open leaves nothing on the
reservation.
+ let file = temp_file.open_writer()?;
+ let capacity = buffers.map_or(self.write_buffer_size,
|reservation| {
+ super::reserve_buffer(reservation, self.write_buffer_size)
+ });
+ let writer = BufWriter::with_capacity(capacity, file);
self.spill_file = Some(ActiveSpillFile { temp_file, writer });
}
Ok(())
}
+ /// Grows a fallback-sized buffer back to the configured size if `buffers`
grants it.
+ /// Buffered bytes move to the new buffer; nothing is written to the file.
+ pub(crate) fn restore_buffer(&mut self, buffers: &MemoryReservation) {
+ if self.failed {
+ return;
+ }
+ let Some(spill_file) = self.spill_file.take() else {
+ return;
+ };
+ let capacity = spill_file.writer.capacity();
+ if capacity >= self.write_buffer_size
+ || buffers.try_grow(self.write_buffer_size - capacity).is_err()
Review Comment:
After a spill the writer now grows the spill buffer to at most `min(write
buffer size, freed / 2)`, where `freed` is what the spill released, and never
shrinks it. That leaves at least half of the freed memory to the next batches.
`spill_buffer_restore_leaves_room_for_the_input` runs your sweep (8 hash
partitions, LZ4, no buffer limit, 1 MiB write buffer, 256 batches of 1,024
integers) at 2,112 KiB and 1,088 KiB. It checks that no restore grows the
buffer past half of what that spill freed, that the writer spills at most about
twice as often as a writer that holds no buffers, and that all rows decode.
One edge remains: a spill buffer the pool grants in full when the first
spill opens the file keeps that size, since the restore never shrinks it.
--
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]
