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);