2010YOUY01 commented on code in PR #15692:
URL: https://github.com/apache/datafusion/pull/15692#discussion_r2041381039
##########
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:
I agree it's better to have a utility to assert an error's root cause, like
`Context(Context(ResourceExhausted, ...`'s root cause is
`ResourceExhanusted`, however, I think this is not required for this PR.
--
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
