DerGut commented on code in PR #15692:
URL: https://github.com/apache/datafusion/pull/15692#discussion_r2041254566
##########
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:
This would actually be a good argument for adding a function to
`DataFusionError` that allows you to safely compare against any of its nested
errors (I'm thinking something similar to Go's
[`errors.Is`](https://pkg.go.dev/errors#Is).
Without such an API, it would never be safe to just wrap an error with
context 🤔
--
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
