fix: keep a late-materialized spill's column concatenation within the pool

native/vendor/datafusion-physical-plan/src/sorts/sort.rs

@@ -732,7 +732,13 @@ impl ExternalSorter {
         is_output_stream: bool,
         coalesce_runs: bool,
     ) -> Result<SendableRecordBatchStream> {
-        let pool = Arc::clone(workspace) as Arc<dyn MemoryPool>;
+        let pool = if self.late_materialization.is_some()
+            && LateMaterialization::applies_to(&self.in_mem_batches)
+        {
+            workspace.bounded()
+        } else {
+            Arc::clone(workspace) as Arc<dyn MemoryPool>
+        };
         let runs =
             MemoryConsumer::new(self.reservation.consumer().name()).register(&pool);
         runs.grow(buffered);

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<RecordBatch> {
+        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::<std::result::Result<Vec<_>, _>>()
+            .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<dyn ExecutionPlan>,
+                    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);

native/vendor/datafusion-physical-plan/src/sorts/spill_workspace.rs

@@ -126,6 +126,10 @@ impl SpillWorkspace {
         })
     }
 
+    pub(super) fn bounded(self: &Arc<Self>) -> Arc<dyn MemoryPool> {
+        Arc::new(BoundedWorkspace(Arc::clone(self)))
+    }
+
     /// Lends up to `size` bytes of unused workspace.
     pub(super) fn borrow(self: &Arc<Self>, size: usize) -> WorkspaceLoan {
         let mut state = self.state.lock();
@@ -218,6 +222,45 @@ impl MemoryPool for SpillWorkspace {
     }
 }
 
+#[derive(Debug)]
+struct BoundedWorkspace(Arc<SpillWorkspace>);
+
+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<dyn MemoryPool> = 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);