rluvaton commented on code in PR #15692:
URL: https://github.com/apache/datafusion/pull/15692#discussion_r2041239380
##########
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:
If I match on the result of this function will it still be resources
exhausted or different type, if different type this can break usages that
depend on that, no?
--
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
