From 238f2db337a9b26c260650f498285b678874707e Mon Sep 17 00:00:00 2001 From: msaf Date: Sun, 18 Oct 2026 11:49:11 +0100 Subject: [PATCH] fix: keep what an in-memory output merge's precheck was granted, and overcommit a small shortfall The precheck added in #22 asked the pool for the memory the output merge of the buffered batches needs beyond the sorter's reservations, and gave it straight back. The workspace then kept only the merge headroom. In a Spark task that also runs a JVM sort and an UnsafeShuffleWriter (parse_product_default, stage 275 on joom-1.1-12), the JVM sort spilled for the precheck's request, the memory it freed went to the shuffle writer, which does not spill for other consumers, and the merge failed a few KB later with "Additional allocation failed for ExternalSorterMerge[0] ... can spill: false" (7 task failures; 0 on 1.1-9..1.1-11). - reserve_output_merge keeps the precheck's grant as a reservation, and the in-memory merge takes it into its workspace as a third parent. The workspace keeps max(headroom, needed - buffered) unused bytes, so the merge grows into memory it already holds. keep_at_most now runs after the runs are registered, so the sorter's reservation is no longer returned to the pool and grown back with the infallible grow. - A fallible growth the pool refuses is taken with the infallible grow while the total of such growth stays within max(16 MiB, 10% of the workspace), which the Comet pool carries as overcommit and repays from its next grant. Only the in-memory output merge allows it; spill merges keep their fan-in and read-ahead degradation, and the bounded view stays strict. Co-Authored-By: Claude Opus 5.5 (1M context) --- .../src/sorts/sort.rs | 44 ++++++++++++--- .../src/sorts/sort/comet_memory_tests.rs | 54 +++++++++++++++++++ .../src/sorts/spill_workspace.rs | 54 ++++++++++++++++--- 3 files changed, 136 insertions(+), 16 deletions(-) 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));