DerGut commented on code in PR #15692:
URL: https://github.com/apache/datafusion/pull/15692#discussion_r2041251862
##########
datafusion/physical-plan/src/sorts/sort.rs:
##########
@@ -759,12 +761,51 @@ impl ExternalSorter {
if self.runtime.disk_manager.tmp_files_enabled() {
let size = self.sort_spill_reservation_bytes;
if self.merge_reservation.size() != size {
- self.merge_reservation.try_resize(size)?;
+ self.merge_reservation
+ .try_resize(size)
+ .map_err(Self::err_with_oom_context)?;
}
}
Ok(())
}
+
+ /// Reserves memory to be able to accommodate the given batch.
+ /// If memory is scarce, tries to spill current in-memory batches to disk
first.
+ async fn reserve_memory_for_batch_and_maybe_spill(
+ &mut self,
+ input: &RecordBatch,
+ ) -> Result<()> {
+ let size = get_reserved_byte_for_record_batch(input);
+
+ match self.reservation.try_grow(size) {
+ Ok(_) => Ok(()),
+ Err(e) => {
+ if self.in_mem_batches.is_empty() {
+ return Err(Self::err_with_oom_context(e));
+ }
+
+ // Spill and try again.
+ self.sort_and_spill_in_mem_batches().await?;
+ self.reservation
+ .try_grow(size)
+ .map_err(Self::err_with_oom_context)
+ }
+ }
+ }
+
+ /// Wraps the error with a context message suggesting settings to tweak.
+ /// This is meant to be used with DataFusionError::ResourcesExhausted only.
+ fn err_with_oom_context(e: DataFusionError) -> DataFusionError {
+ match e {
+ DataFusionError::ResourcesExhausted(_) => e.context(format!(
+ "Not enough memory to continue external sort. \
+ Consider increasing the memory limit, or decreasing
sort_spill_reservation_bytes"
+ )),
Review Comment:
That's a good point! It will be a different type (or enum variant), that's
also why the test is currently failing (see
https://github.com/apache/datafusion/pull/15692/files#r2041251481)
--
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
