diff --git a/native/vendor/datafusion-physical-plan/src/sorts/sort.rs b/native/vendor/datafusion-physical-plan/src/sorts/sort.rs index 75980ae529a..5f29defc4a7 100644 --- a/native/vendor/datafusion-physical-plan/src/sorts/sort.rs +++ b/native/vendor/datafusion-physical-plan/src/sorts/sort.rs @@ -106,6 +106,9 @@ impl ExternalSorterMetrics { const SMALL_BATCHES_TARGET_BYTES: usize = 4 << 20; const ROW_BYTES_SAMPLES: usize = 8; const ROW_BYTES_SAMPLE_ROWS: usize = 1024; +/// COMET PATCH: the least overcommit an in-memory output merge may take when the pool +/// refuses growth beyond its estimate. See `SpillWorkspace::allow_overcommit`. +const OUTPUT_MERGE_OVERCOMMIT: usize = 16 << 20; /// COMET PATCH: a session config extension. A sort whose input stayed in memory but whose /// reservation exceeds this many bytes spills it before producing output. 0 disables. @@ -493,13 +496,16 @@ impl ExternalSorter { if self.spills_before_output() { self.sort_and_spill(true).await?; } + let mut output_merge = None; if !self.spilled_before() && self.in_mem_batches.len() > 1 && self.reservation.size() >= self.sort_in_place_threshold_bytes && self.runtime.disk_manager.tmp_files_enabled() - && !self.output_merge_fits()? { - self.spill_in_mem_batches_in_groups().await?; + output_merge = self.reserve_output_merge()?; + if output_merge.is_none() { + self.spill_in_mem_batches_in_groups().await?; + } } if self.spilled_before() { // Sort `in_mem_batches` and spill it first. If there are many @@ -542,14 +548,29 @@ impl ExternalSorter { // headroom for the merge's cursors and buffers instead of returning it to the // pool and asking for it again. Released run memory beyond it goes back to the // pool, and growth is charged to the merge. See `SpillWorkspace`. + // + // What `reserve_output_merge` was granted toward the merge's estimate stays + // reserved too, so that the merge does not have to win back from Spark memory + // it just gave up: Spark hands it to a consumer that does not spill for others, + // such as the shuffle writer of a JVM sort in the same task. A small shortfall of + // the estimate is taken as bounded overcommit, because the merge cannot spill. let buffered = self.reservation.size(); let headroom = self.merge_reservation.size(); + let (needed, grant) = match output_merge { + Some(output_merge) => output_merge, + None => (0, self.merge_reservation.new_empty()), + }; let workspace = SpillWorkspace::new(vec![ self.merge_reservation.take(), self.reservation.take(), + grant, ]); - workspace.keep_at_most(headroom); - self.in_mem_sort_stream_in_workspace(&workspace, buffered, true, true) + let stream = + self.in_mem_sort_stream_in_workspace(&workspace, buffered, true, true); + let reserved = MemoryPool::reserved(workspace.as_ref()); + workspace.keep_at_most(headroom.max(needed.saturating_sub(buffered))); + workspace.allow_overcommit(OUTPUT_MERGE_OVERCOMMIT.max(reserved / 10)); + stream } else { // Release the memory reserved for merge back to the pool so // there is some left when `in_mem_sort_stream` requests an @@ -560,12 +581,15 @@ impl ExternalSorter { } } - /// COMET PATCH - fn output_merge_fits(&self) -> Result { + /// COMET PATCH: the memory the output merge of the buffered batches is estimated to + /// need, and a reservation holding what that needs beyond the sorter's reservations. + /// `None` if the pool cannot grant it. + fn reserve_output_merge(&self) -> Result> { + let grant = self.merge_reservation.new_empty(); if self.late_materialization.is_some() && LateMaterialization::applies_to(&self.in_mem_batches) { - return Ok(true); + return Ok(Some((0, grant))); } let row_bytes = self.sampled_row_bytes(&self.in_mem_batches)?; let needed = self @@ -574,7 +598,11 @@ impl ExternalSorter { .map(|batch| self.run_merge_memory(batch, row_bytes)) .sum::>()?; let held = self.reservation.size() + self.merge_reservation.size(); - Ok(self.pool_can_grant(needed.saturating_sub(held))) + let missing = needed.saturating_sub(held); + if missing > 0 && grant.try_grow(missing).is_err() { + return Ok(None); + } + Ok(Some((needed, grant))) } /// COMET PATCH diff --git a/native/vendor/datafusion-physical-plan/src/sorts/sort/comet_memory_tests.rs b/native/vendor/datafusion-physical-plan/src/sorts/sort/comet_memory_tests.rs index 850588cd70a..1dbcb8fd512 100644 --- a/native/vendor/datafusion-physical-plan/src/sorts/sort/comet_memory_tests.rs +++ b/native/vendor/datafusion-physical-plan/src/sorts/sort/comet_memory_tests.rs @@ -870,3 +870,57 @@ async fn output_merge_that_does_not_fit_spills_instead_of_failing() -> Result<() } Ok(()) } + +/// The output merge's precheck asked the pool for what the merge needs beyond the sorter's +/// reservations and gave it straight back. In a Spark task the JVM sort that spilled for +/// that request let its memory go to the task's shuffle writer, which does not spill for +/// other consumers, and the merge failed later with "Additional allocation failed for +/// ExternalSorterMerge[0]". The merge now keeps what the precheck was granted. +#[tokio::test] +async fn output_merge_keeps_what_its_precheck_was_granted() -> Result<()> { + let (schema, ordering, input) = key_dominated_batches(300, 1024)?; + let mut ordering: Vec = ordering.iter().cloned().collect(); + ordering.push(PhysicalSortExpr::new_default(crate::expressions::col( + "v", &schema, + )?)); + let ordering = LexOrdering::new(ordering).unwrap(); + let expected = sort_batch(&concat_batches(&schema, &input)?, &ordering, None)?; + let data: usize = input.iter().map(get_record_batch_memory_size).sum(); + let headroom = 64 * 1024; + let pool_size = 2 * data + headroom; + let stealing = StealingPool::new(pool_size); + let pool: Arc = Arc::clone(&stealing) as _; + let runtime = RuntimeEnvBuilder::new() + .with_memory_pool(Arc::clone(&pool)) + .build_arc()?; + let mut sorter = ExternalSorter::new( + 0, + Arc::clone(&schema), + ordering, + 8192, + headroom, + 1024 * 1024, + SpillCompression::Uncompressed, + &ExecutionPlanMetricsSet::new(), + runtime, + )?; + for batch in input { + sorter.insert_batch(batch).await?; + } + let merge_stream = sorter.sort().await?; + assert_eq!(sorter.spill_count(), 0, "the output merge fits"); + drop(sorter); + + let contender = MemoryConsumer::new("ShuffleExternalSorter").register(&pool); + contender.try_grow(pool_size - pool.reserved())?; + stealing.arm(); + + let output: Vec = merge_stream.try_collect().await?; + assert_eq!(pool.reserved(), contender.size() + stealing.stolen()); + let actual = concat_batches(&schema, &output)?; + assert_eq!(actual.num_rows(), expected.num_rows()); + for column in 0..4 { + assert_eq!(actual.column(column), expected.column(column)); + } + Ok(()) +} diff --git a/native/vendor/datafusion-physical-plan/src/sorts/spill_workspace.rs b/native/vendor/datafusion-physical-plan/src/sorts/spill_workspace.rs index 49956e38519..e9898e7a9cb 100644 --- a/native/vendor/datafusion-physical-plan/src/sorts/spill_workspace.rs +++ b/native/vendor/datafusion-physical-plan/src/sorts/spill_workspace.rs @@ -59,6 +59,9 @@ struct State { used: usize, /// How many unused bytes stay reserved in the parents. keep: usize, + /// How many more bytes a fallible growth may still take with the infallible `grow` + /// when the execution pool refuses them. See [`SpillWorkspace::allow_overcommit`]. + overcommit_left: usize, } impl State { @@ -66,13 +69,18 @@ impl State { self.parents.iter().map(MemoryReservation::size).sum() } - fn cover(&mut self, used: usize, fallible: bool) -> Result<()> { + fn cover(&mut self, used: usize, fallible: bool, overcommit: bool) -> Result<()> { let reserved = self.reserved(); if used > reserved { - if fallible { - self.parents[0].try_grow(used - reserved)?; - } else { - self.parents[0].grow(used - reserved); + let missing = used - reserved; + if !fallible { + self.parents[0].grow(missing); + } else if let Err(e) = self.parents[0].try_grow(missing) { + if !overcommit || missing > self.overcommit_left { + return Err(e); + } + self.overcommit_left -= missing; + self.parents[0].grow(missing); } } self.used = used; @@ -121,6 +129,7 @@ impl SpillWorkspace { parents, used: 0, keep: usize::MAX, + overcommit_left: 0, }), fallible, }) @@ -172,6 +181,14 @@ impl SpillWorkspace { state.trim(); } + /// Lets children take up to `bytes` in total beyond what the execution pool grants, + /// through the infallible `grow`, which the Comet pool records as overcommit and repays + /// from the next grant. For a merge that cannot spill, a small shortfall of its estimate + /// then costs a bounded overcommit rather than the task. + pub(super) fn allow_overcommit(&self, bytes: usize) { + self.state.lock().overcommit_left = bytes; + } + fn release(&self, size: usize) { let mut state = self.state.lock(); state.used = state @@ -197,7 +214,7 @@ impl MemoryPool for SpillWorkspace { let mut state = self.state.lock(); let used = state.used.saturating_add(additional); state - .cover(used, false) + .cover(used, false, false) .expect("an infallible grow cannot fail"); } @@ -214,7 +231,7 @@ impl MemoryPool for SpillWorkspace { let Some(used) = state.used.checked_add(additional) else { return resources_err!("Sort spill workspace overflow"); }; - state.cover(used, self.fallible) + state.cover(used, self.fallible, true) } fn reserved(&self) -> usize { @@ -253,7 +270,7 @@ impl MemoryPool for BoundedWorkspace { let Some(used) = state.used.checked_add(additional) else { return resources_err!("Sort spill workspace overflow"); }; - state.cover(used, true) + state.cover(used, true, false) } fn reserved(&self) -> usize { @@ -301,6 +318,27 @@ mod tests { assert_eq!(parent.reserved(), 20); } + #[test] + fn refused_growth_overcommits_only_within_the_allowance() { + let (parent, workspace, child) = setup(100, 80); + child.try_grow(80).unwrap(); + let contender = MemoryConsumer::new("contender").register(&parent); + contender.try_grow(20).unwrap(); + assert!(child.try_grow(5).is_err()); + + workspace.allow_overcommit(8); + child.try_grow(5).unwrap(); + assert!(child.try_grow(4).is_err()); + child.try_grow(3).unwrap(); + assert!(child.try_grow(1).is_err()); + assert_eq!(parent.reserved(), 108); + assert_eq!(MemoryPool::reserved(workspace.as_ref()), 88); + + drop(child); + workspace.close(); + assert_eq!(parent.reserved(), 20); + } + #[test] fn growth_past_the_workspace_uses_the_first_parent() { let parent: Arc = Arc::new(GreedyMemoryPool::new(100));