From e9a4878722230892dff4b4425b862dd475e62174 Mon Sep 17 00:00:00 2001 From: msaf Date: Tue, 13 Oct 2026 20:39:56 +0100 Subject: [PATCH] fix: keep a late-materialized spill's column concatenation within the pool A spill sorts its buffered batches inside a SpillWorkspace that grows its first parent with the infallible grow (29b0bb02b), since what the merge's children record there already exists. Late materialization's Gather also asks that workspace, with try_grow, before it concatenates a group of column pieces, and the copy it would make is new memory. In a spill workspace that try_grow never failed, so the concatenation went ahead past Spark's grant and Comet carried the difference as overcommit. When the batches' columns share one allocation, as batches decoded from shuffle do, no piece is freed until the batch's last column is concatenated, so a spill held up to about twice its input. On em32 parse_parquet device_default this put 8 tasks' native reservations at 12.8 GiB against a 9 GiB grant (OVERCOMMIT_DIAG: every overcommit step came from Gather::concatenate under sort_and_spill) and the pods were killed with exit 137. The sorted runs of a late-materialized sort now register with a bounded view of the workspace whose try_grow stays under the execution pool's limits, while grow still records existing memory without limit. A concatenation the pool refuses is skipped and the columns are gathered from their pieces. Gather also reserves the concatenated array's actual size before keeping it, not only its sliced size. The streaming merge of non-late runs keeps the infallible workspace, so the failure 29b0bb02b fixed does not return. Tests: a bounded workspace view refuses growth past the pool and still accepts grow; a late-materialized spill of batches sharing one allocation per batch peaks within the pool for three pool sizes and five orderings (before: 2997665 of 2283968 bytes). Co-Authored-By: Claude Opus 5.5 (1M context) --- .../src/sorts/sort.rs | 8 ++- .../src/sorts/sort/late_materialize.rs | 49 ++++++++++++- .../src/sorts/spill_workspace.rs | 71 +++++++++++++++++++ 3 files changed, 126 insertions(+), 2 deletions(-) diff --git a/native/vendor/datafusion-physical-plan/src/sorts/sort.rs b/native/vendor/datafusion-physical-plan/src/sorts/sort.rs index ce3173f5f73..6ae97abfacd 100644 --- a/native/vendor/datafusion-physical-plan/src/sorts/sort.rs +++ b/native/vendor/datafusion-physical-plan/src/sorts/sort.rs @@ -732,7 +732,13 @@ impl ExternalSorter { is_output_stream: bool, coalesce_runs: bool, ) -> Result { - let pool = Arc::clone(workspace) as Arc; + let pool = if self.late_materialization.is_some() + && LateMaterialization::applies_to(&self.in_mem_batches) + { + workspace.bounded() + } else { + Arc::clone(workspace) as Arc + }; let runs = MemoryConsumer::new(self.reservation.consumer().name()).register(&pool); runs.grow(buffered); diff --git a/native/vendor/datafusion-physical-plan/src/sorts/sort/late_materialize.rs b/native/vendor/datafusion-physical-plan/src/sorts/sort/late_materialize.rs index a0933de1ac6..9d3104945dc 100644 --- a/native/vendor/datafusion-physical-plan/src/sorts/sort/late_materialize.rs +++ b/native/vendor/datafusion-physical-plan/src/sorts/sort/late_materialize.rs @@ -663,7 +663,12 @@ impl Gather { let arrays: Vec<&dyn Array> = arrays.iter().map(|array| array.as_ref()).collect(); let merged = concat(&arrays) .ok() - .filter(|array| take(array, &UInt32Array::from(vec![0u32]), None).is_ok()); + .filter(|array| take(array, &UInt32Array::from(vec![0u32]), None).is_ok()) + .filter(|array| { + let needed = self.needed() + array.get_array_memory_size(); + let size = self.reservation.size(); + size >= needed || self.reservation.try_grow(needed - size).is_ok() + }); if merged.is_none() { self.fit(); } @@ -1567,6 +1572,48 @@ mod tests { Ok(()) } + fn sharing_one_allocation_per_batch(input: &[RecordBatch]) -> Vec { + let mut bytes = vec![]; + let mut writer = + arrow::ipc::writer::StreamWriter::try_new(&mut bytes, &input[0].schema()) + .unwrap(); + for batch in input { + writer.write(batch).unwrap(); + } + writer.finish().unwrap(); + drop(writer); + arrow::ipc::reader::StreamReader::try_new(std::io::Cursor::new(bytes), None) + .unwrap() + .collect::, _>>() + .unwrap() + } + + #[tokio::test] + async fn a_spill_concatenates_only_what_the_pool_grants() -> Result<()> { + let separate = mixed_batches(64, 150, 42); + let bytes: usize = separate + .iter() + .map(RecordBatch::get_array_memory_size) + .sum(); + let input = sharing_one_allocation_per_batch(&separate); + for ordering in mixed_orderings() { + for limit in [bytes / 2, bytes / 3, bytes / 5] { + let pool = PeakPool::new(limit); + let sort = mixed_sort_exec(&input, ordering.clone()); + let output = collect( + Arc::clone(&sort) as Arc, + context(Some(Arc::clone(&pool) as _), 1024, 256 << 10), + ) + .await?; + assert!(sort.metrics().unwrap().spill_count().unwrap() > 0); + assert_matches_reference(&input, &output, &ordering); + assert!(pool.peak() <= limit, "peak {} of {limit}", pool.peak()); + assert_eq!(pool.reserved(), 0); + } + } + Ok(()) + } + #[test] fn concatenates_every_column_of_a_scattered_order() -> Result<()> { let input = mixed_batches(50, 100, 7); 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 271c24f2186..49956e38519 100644 --- a/native/vendor/datafusion-physical-plan/src/sorts/spill_workspace.rs +++ b/native/vendor/datafusion-physical-plan/src/sorts/spill_workspace.rs @@ -126,6 +126,10 @@ impl SpillWorkspace { }) } + pub(super) fn bounded(self: &Arc) -> Arc { + Arc::new(BoundedWorkspace(Arc::clone(self))) + } + /// Lends up to `size` bytes of unused workspace. pub(super) fn borrow(self: &Arc, size: usize) -> WorkspaceLoan { let mut state = self.state.lock(); @@ -218,6 +222,45 @@ impl MemoryPool for SpillWorkspace { } } +#[derive(Debug)] +struct BoundedWorkspace(Arc); + +impl Display for BoundedWorkspace { + fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { + write!(f, "BoundedWorkspace") + } +} + +impl MemoryPool for BoundedWorkspace { + fn name(&self) -> &str { + "BoundedWorkspace" + } + + fn grow(&self, reservation: &MemoryReservation, additional: usize) { + self.0.grow(reservation, additional); + } + + fn shrink(&self, reservation: &MemoryReservation, shrink: usize) { + self.0.shrink(reservation, shrink); + } + + fn try_grow( + &self, + _reservation: &MemoryReservation, + additional: usize, + ) -> Result<()> { + let mut state = self.0.state.lock(); + let Some(used) = state.used.checked_add(additional) else { + return resources_err!("Sort spill workspace overflow"); + }; + state.cover(used, true) + } + + fn reserved(&self) -> usize { + self.0.reserved() + } +} + #[cfg(test)] mod tests { use super::*; @@ -306,6 +349,34 @@ mod tests { assert_eq!(parent.reserved(), 40); } + #[test] + fn a_bounded_view_of_a_spill_workspace_refuses_growth_past_the_pool() { + let parent: Arc = Arc::new(GreedyMemoryPool::new(100)); + let sorter = MemoryConsumer::new("sorter").register(&parent); + sorter.try_grow(60).unwrap(); + let contender = MemoryConsumer::new("contender").register(&parent); + contender.try_grow(30).unwrap(); + let workspace = SpillWorkspace::for_spill(vec![sorter]); + let bounded = workspace.bounded(); + let child = MemoryConsumer::new("child").register(&bounded); + + child.grow(40); + child.try_grow(30).unwrap(); + assert_eq!(parent.reserved(), 100); + assert!(child.try_grow(1).is_err()); + assert_eq!(child.size(), 70); + child.grow(5); + assert_eq!(parent.reserved(), 105); + + child.shrink(35); + child.try_grow(35).unwrap(); + assert_eq!(parent.reserved(), 105); + + workspace.close(); + drop(child); + assert_eq!(parent.reserved(), 30); + } + #[test] fn can_grow_counts_unused_workspace_and_leaves_reservations_unchanged() { let (parent, workspace, child) = setup(100, 40);