fix: keep what an in-memory output merge's precheck was granted, and overcommit a small shortfall

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<bool> {
+    /// 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<Option<(usize, MemoryReservation)>> {
+        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::<Result<usize>>()?;
         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

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<PhysicalSortExpr> = 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<dyn MemoryPool> = 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<RecordBatch> = 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(())
+}

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

@@ -59,20 +59,28 @@ 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 {
     fn reserved(&self) -> usize {
         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<dyn MemoryPool> = Arc::new(GreedyMemoryPool::new(100));