2010YOUY01 commented on code in PR #14644:
URL: https://github.com/apache/datafusion/pull/14644#discussion_r1957240155
##########
datafusion/physical-plan/src/sorts/sort.rs:
##########
@@ -408,50 +395,100 @@ impl ExternalSorter {
debug!("Spilling sort data of ExternalSorter to disk whilst
inserting");
- self.in_mem_sort().await?;
-
let spill_file = self.runtime.disk_manager.create_tmp_file("Sorting")?;
let batches = std::mem::take(&mut self.in_mem_batches);
- let spilled_rows = spill_record_batches(
+ let (spilled_rows, spilled_bytes) = spill_record_batches(
batches,
spill_file.path().into(),
Arc::clone(&self.schema),
)?;
let used = self.reservation.free();
self.metrics.spill_count.add(1);
- self.metrics.spilled_bytes.add(used);
+ self.metrics.spilled_bytes.add(spilled_bytes);
self.metrics.spilled_rows.add(spilled_rows);
self.spills.push(spill_file);
Ok(used)
}
/// Sorts the in_mem_batches in place
- async fn in_mem_sort(&mut self) -> Result<()> {
- if self.in_mem_batches_sorted {
- return Ok(());
- }
-
+ ///
+ /// Sorting may have freed memory, especially if fetch is `Some`. If
+ /// the memory usage has dropped by a factor of 2, then we don't have
+ /// to spill. Otherwise, we spill to free up memory for inserting
+ /// more batches.
+ ///
+ /// The factor of 2 aims to avoid a degenerate case where the
+ /// memory required for `fetch` is just under the memory available,
+ // causing repeated re-sorting of data
+ async fn sort_or_spill_in_mem_batches(&mut self) -> Result<()> {
// Release the memory reserved for merge back to the pool so
- // there is some left when `in_memo_sort_stream` requests an
+ // there is some left when `in_mem_sort_stream` requests an
// allocation.
Review Comment:
```suggestion
// allocation. At the end of this function, memory will be reserved
again for the next spill.
```
I found this pre-reserved memory confusing when reading this code for the
first time, so I would like to make it more clear.
##########
datafusion/physical-plan/src/sorts/sort.rs:
##########
@@ -408,50 +395,100 @@ impl ExternalSorter {
debug!("Spilling sort data of ExternalSorter to disk whilst
inserting");
- self.in_mem_sort().await?;
-
let spill_file = self.runtime.disk_manager.create_tmp_file("Sorting")?;
let batches = std::mem::take(&mut self.in_mem_batches);
- let spilled_rows = spill_record_batches(
+ let (spilled_rows, spilled_bytes) = spill_record_batches(
batches,
spill_file.path().into(),
Arc::clone(&self.schema),
)?;
let used = self.reservation.free();
self.metrics.spill_count.add(1);
- self.metrics.spilled_bytes.add(used);
+ self.metrics.spilled_bytes.add(spilled_bytes);
self.metrics.spilled_rows.add(spilled_rows);
self.spills.push(spill_file);
Ok(used)
}
/// Sorts the in_mem_batches in place
- async fn in_mem_sort(&mut self) -> Result<()> {
- if self.in_mem_batches_sorted {
- return Ok(());
- }
-
+ ///
+ /// Sorting may have freed memory, especially if fetch is `Some`. If
+ /// the memory usage has dropped by a factor of 2, then we don't have
+ /// to spill. Otherwise, we spill to free up memory for inserting
+ /// more batches.
+ ///
+ /// The factor of 2 aims to avoid a degenerate case where the
+ /// memory required for `fetch` is just under the memory available,
+ // causing repeated re-sorting of data
+ async fn sort_or_spill_in_mem_batches(&mut self) -> Result<()> {
// Release the memory reserved for merge back to the pool so
- // there is some left when `in_memo_sort_stream` requests an
+ // there is some left when `in_mem_sort_stream` requests an
// allocation.
self.merge_reservation.free();
- self.in_mem_batches = self
- .in_mem_sort_stream(self.metrics.baseline.intermediate())?
- .try_collect()
- .await?;
+ let before = self.reservation.size();
+
+ let mut sorted_stream =
+ self.in_mem_sort_stream(self.metrics.baseline.intermediate())?;
+
+ // `self.in_mem_batches` is already taken away by the sort_stream, now
it is empty.
+ // We'll gradually collect the sorted stream into self.in_mem_batches,
or directly
+ // write sorted batches to disk when the memory is insufficient.
+ let mut spill_writer: Option<IPCWriter> = None;
+ while let Some(batch) = sorted_stream.next().await {
+ let batch = batch?;
+ match &mut spill_writer {
+ None => {
+ let sorted_size =
get_reserved_byte_for_record_batch(&batch);
+ if self.reservation.try_grow(sorted_size).is_err() {
+ // Directly write in_mem_batches as well as all the
remaining batches in
+ // sorted_stream to disk. Further batches fetched from
`sorted_stream` will
+ // be handled by the `Some(writer)` matching arm.
+ let spill_file =
+
self.runtime.disk_manager.create_tmp_file("Sorting")?;
+ let mut writer = IPCWriter::new(spill_file.path(),
&self.schema)?;
+ // Flush everything in memory to the spill file
+ for batch in self.in_mem_batches.drain(..) {
+ writer.write(&batch)?;
+ }
+ // as well as the newly sorted batch
+ writer.write(&batch)?;
+ spill_writer = Some(writer);
+ self.reservation.free();
+ self.spills.push(spill_file);
+ } else {
+ self.in_mem_batches.push(batch);
Review Comment:
I think adding another field `in_mem_batches_sorted` for this purpose can
slightly improve readability
##########
datafusion/physical-plan/src/sorts/sort.rs:
##########
@@ -612,6 +659,20 @@ impl ExternalSorter {
}
}
+/// Estimate how much memory is needed to sort a `RecordBatch`.
Review Comment:
I suggest to explicitly point out the extra memory is reserved for potential
`Row` conversion of the original array, which is used to speed up comparison in
sorting and merging.
--
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: github-unsubscr...@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
us...@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscr...@datafusion.apache.org
For additional commands, e-mail: github-h...@datafusion.apache.org
