fix: spill a native sort's buffered batches in groups when their output merge does not fit

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

@@ -26,14 +26,16 @@ use std::pin::Pin;
 use std::sync::Arc;
 use std::task::{Context, Poll};
 
-use arrow::datatypes::{DataType, SchemaRef};
+use arrow::datatypes::SchemaRef;
 use datafusion_common::{Result, internal_err, resources_err};
 use datafusion_execution::memory_pool::{MemoryPool, MemoryReservation};
 
 use crate::sorts::builder::try_grow_reservation_to_at_least;
 use crate::sorts::sort::get_reserved_bytes_for_record_batch_size;
 use crate::sorts::spill_workspace::SpillWorkspace;
-use crate::sorts::streaming_merge::{SortedSpillFile, StreamingMergeBuilder};
+use crate::sorts::streaming_merge::{
+    SortedSpillFile, StreamingMergeBuilder, merge_uses_row_cursor,
+};
 use crate::stream::{ObservedStream, RecordBatchStreamAdapter};
 use datafusion_execution::{RecordBatchStream, SendableRecordBatchStream};
 use datafusion_physical_expr_common::sort_expr::LexOrdering;
@@ -713,23 +715,9 @@ impl MultiLevelMergeBuilder {
         }
     }
 
-    /// COMET PATCH: whether the merge compares rows, which `StreamingMergeBuilder` does
-    /// unless it sorts on one primitive or string/binary column.
+    /// COMET PATCH
     fn merge_uses_row_cursor(&self) -> bool {
-        let [sort] = self.expr.as_ref() else {
-            return true;
-        };
-        !sort.expr.data_type(&self.schema).is_ok_and(|data_type| {
-            data_type.is_primitive()
-                || matches!(
-                    data_type,
-                    DataType::Utf8
-                        | DataType::Utf8View
-                        | DataType::LargeUtf8
-                        | DataType::Binary
-                        | DataType::LargeBinary
-                )
-        })
+        merge_uses_row_cursor(&self.expr, &self.schema)
     }
 
     /// Re-spill the spill file at `index` with half its batch size, putting it back

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

@@ -48,7 +48,9 @@ use crate::metrics::{
 use crate::projection::{ProjectionExec, make_with_child, update_ordering};
 use crate::sorts::IncrementalSortIterator;
 use crate::sorts::spill_workspace::SpillWorkspace;
-use crate::sorts::streaming_merge::{SortedSpillFile, StreamingMergeBuilder};
+use crate::sorts::streaming_merge::{
+    SortedSpillFile, StreamingMergeBuilder, merge_uses_row_cursor,
+};
 use crate::spill::get_record_batch_memory_size;
 use crate::spill::in_progress_spill_file::InProgressSpillFile;
 use crate::spill::spill_manager::{GetSlicedSize, SpillManager};
@@ -65,6 +67,7 @@ use crate::{
 use arrow::array::{RecordBatch, RecordBatchOptions};
 use arrow::compute::{concat_batches, lexsort_to_indices, take_arrays};
 use arrow::datatypes::{DataType, SchemaRef};
+use arrow::row::{RowConverter, SortField};
 use datafusion_common::config::SpillCompression;
 use datafusion_common::tree_node::TreeNodeRecursion;
 use datafusion_common::utils::memory::RecordBatchMemoryCounter;
@@ -78,6 +81,7 @@ use datafusion_execution::runtime_env::RuntimeEnv;
 use datafusion_physical_expr::LexOrdering;
 use datafusion_physical_expr::PhysicalExpr;
 use datafusion_physical_expr::expressions::{DynamicFilterPhysicalExpr, lit};
+use datafusion_physical_expr_common::utils::evaluate_expressions_to_arrays;
 
 use futures::{StreamExt, TryStreamExt};
 use log::{debug, trace};
@@ -100,6 +104,8 @@ impl ExternalSorterMetrics {
 
 /// COMET PATCH
 const SMALL_BATCHES_TARGET_BYTES: usize = 4 << 20;
+const ROW_BYTES_SAMPLES: usize = 8;
+const ROW_BYTES_SAMPLE_ROWS: usize = 1024;
 
 /// 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.
@@ -487,6 +493,14 @@ impl ExternalSorter {
         if self.spills_before_output() {
             self.sort_and_spill(true).await?;
         }
+        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?;
+        }
         if self.spilled_before() {
             // Sort `in_mem_batches` and spill it first. If there are many
             // `in_mem_batches` and the memory limit is almost reached, merging
@@ -546,6 +560,112 @@ impl ExternalSorter {
         }
     }
 
+    /// COMET PATCH
+    fn output_merge_fits(&self) -> Result<bool> {
+        if self.late_materialization.is_some()
+            && LateMaterialization::applies_to(&self.in_mem_batches)
+        {
+            return Ok(true);
+        }
+        let row_bytes = self.sampled_row_bytes(&self.in_mem_batches)?;
+        let needed = self
+            .in_mem_batches
+            .iter()
+            .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)))
+    }
+
+    /// COMET PATCH
+    fn run_merge_memory(&self, batch: &RecordBatch, row_bytes: usize) -> Result<usize> {
+        let rows = batch.num_rows();
+        let buffers = if rows > self.batch_size { 2 } else { 1 };
+        let encoded = rows.min(self.batch_size) * buffers * row_bytes;
+        Ok(batch.get_sliced_size()? + encoded + encoded / 4)
+    }
+
+    /// COMET PATCH
+    fn sampled_row_bytes(&self, batches: &[RecordBatch]) -> Result<usize> {
+        if batches.is_empty() || !merge_uses_row_cursor(&self.expr, &self.schema) {
+            return Ok(0);
+        }
+        let fields = self
+            .expr
+            .iter()
+            .map(|sort| {
+                Ok(SortField::new_with_options(
+                    sort.expr.data_type(&self.schema)?,
+                    sort.options,
+                ))
+            })
+            .collect::<Result<Vec<_>>>()?;
+        let converter = RowConverter::new(fields)?;
+        let exprs: Vec<_> = self.expr.iter().map(|sort| Arc::clone(&sort.expr)).collect();
+        let step = batches.len().div_ceil(ROW_BYTES_SAMPLES).max(1);
+        let mut row_bytes = 0;
+        for batch in batches.iter().step_by(step) {
+            let rows = batch.num_rows().min(ROW_BYTES_SAMPLE_ROWS);
+            if rows == 0 {
+                continue;
+            }
+            let columns = evaluate_expressions_to_arrays(&exprs, &batch.slice(0, rows))?;
+            let encoded = converter.convert_columns(&columns)?;
+            row_bytes = row_bytes.max(encoded.size().div_ceil(rows));
+        }
+        Ok(row_bytes)
+    }
+
+    /// COMET PATCH
+    fn pool_can_grant(&self, bytes: usize) -> bool {
+        bytes == 0 || self.merge_reservation.new_empty().try_grow(bytes).is_ok()
+    }
+
+    /// COMET PATCH
+    async fn spill_in_mem_batches_in_groups(&mut self) -> Result<()> {
+        let mut batches = std::mem::take(&mut self.in_mem_batches);
+        let row_bytes = self.sampled_row_bytes(&batches)?;
+        let mut counter = RecordBatchMemoryCounter::new();
+        let mut reserved = batches
+            .iter()
+            .map(|batch| reserved_bytes_counting_shared_buffers(batch, &mut counter))
+            .collect::<Result<Vec<_>>>()?;
+        while !batches.is_empty() {
+            let headroom = self.merge_reservation.size();
+            let mut len = 0;
+            let mut needed = 0;
+            let mut held = 0;
+            for (batch, bytes) in batches.iter().zip(&reserved) {
+                needed += self.run_merge_memory(batch, row_bytes)?;
+                held += bytes;
+                if len > 0 && !self.pool_can_grant(needed.saturating_sub(held + headroom))
+                {
+                    break;
+                }
+                len += 1;
+            }
+            let group: Vec<RecordBatch> = batches.drain(..len).collect();
+            let bytes: usize = reserved.drain(..len).sum();
+            let bytes = if batches.is_empty() {
+                self.reservation.size()
+            } else {
+                bytes.min(self.reservation.size())
+            };
+            self.in_mem_batches = group;
+            let workspace = SpillWorkspace::for_spill(vec![
+                self.reservation.split(bytes),
+                self.merge_reservation.take(),
+            ]);
+            let result = self
+                .merge_and_spill_in_mem_batches(&workspace, bytes, true)
+                .await;
+            workspace.close();
+            result?;
+            let _ = self.reserve_memory_for_merge();
+        }
+        Ok(())
+    }
+
     /// How much memory is buffered in this `ExternalSorter`?
     fn used(&self) -> usize {
         self.reservation.size()

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

@@ -749,3 +749,124 @@ async fn spill_merge_of_key_dominated_rows_does_not_fail() -> Result<()> {
     }
     Ok(())
 }
+
+#[derive(Debug)]
+struct PeakPool {
+    inner: GreedyMemoryPool,
+    peak: AtomicUsize,
+}
+
+impl PeakPool {
+    fn new(size: usize) -> Arc<Self> {
+        Arc::new(Self {
+            inner: GreedyMemoryPool::new(size),
+            peak: AtomicUsize::new(0),
+        })
+    }
+
+    fn peak(&self) -> usize {
+        self.peak.load(Ordering::Relaxed)
+    }
+
+    fn record(&self) {
+        self.peak.fetch_max(self.inner.reserved(), Ordering::Relaxed);
+    }
+}
+
+impl fmt::Display for PeakPool {
+    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
+        write!(f, "peak({})", self.inner)
+    }
+}
+
+impl MemoryPool for PeakPool {
+    fn name(&self) -> &str {
+        "peak"
+    }
+
+    fn grow(&self, reservation: &MemoryReservation, additional: usize) {
+        self.inner.grow(reservation, additional);
+        self.record();
+    }
+
+    fn shrink(&self, reservation: &MemoryReservation, shrink: usize) {
+        self.inner.shrink(reservation, shrink)
+    }
+
+    fn try_grow(&self, reservation: &MemoryReservation, additional: usize) -> Result<()> {
+        self.inner.try_grow(reservation, additional)?;
+        self.record();
+        Ok(())
+    }
+
+    fn reserved(&self) -> usize {
+        self.inner.reserved()
+    }
+
+    fn memory_limit(&self) -> MemoryLimit {
+        self.inner.memory_limit()
+    }
+}
+
+/// The output merge of batches that stayed in memory holds the sorted batches and the
+/// rows it encodes their sort keys into. With a key that is the whole row these rows
+/// outgrow what sorting frees, and the merge failed with "Additional allocation failed
+/// for ExternalSorterMerge[0]" in a pool the buffered batches filled. It now spills the
+/// batches in groups that fit and merges the spill files.
+#[tokio::test]
+async fn output_merge_that_does_not_fit_spills_instead_of_failing() -> 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 headroom = 64 * 1024;
+    let sort = |pool: Arc<dyn MemoryPool>| {
+        let runtime = RuntimeEnvBuilder::new().with_memory_pool(pool).build_arc();
+        let schema = Arc::clone(&schema);
+        let ordering = ordering.clone();
+        let input = input.clone();
+        async move {
+            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 buffered = sorter.reservation.size();
+            let spilled_before = sorter.spill_count();
+            let output: Vec<RecordBatch> = sorter.sort().await?.try_collect().await?;
+            Ok::<_, DataFusionError>((buffered, spilled_before, sorter.spill_count(), output))
+        }
+    };
+    let (buffered, _, spills, _) =
+        sort(Arc::new(GreedyMemoryPool::new(1 << 30)) as _).await?;
+    assert_eq!(spills, 0);
+    let data: usize = input.iter().map(get_record_batch_memory_size).sum();
+    for extra in [0, data / 100, data / 50, data / 25] {
+        let limit = buffered + headroom + extra;
+        let peak = PeakPool::new(limit);
+        let pool: Arc<dyn MemoryPool> = Arc::clone(&peak) as _;
+        let (_, spilled_before, spills, output) = sort(Arc::clone(&pool)).await?;
+        assert_eq!(spilled_before, 0, "the batches stay in memory until the output");
+        assert!(spills > 1, "the batches spill in groups");
+        assert!(peak.peak() <= limit, "peak {} of {limit}", peak.peak());
+        assert_eq!(pool.reserved(), 0);
+        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/streaming_merge.rs

@@ -286,6 +286,24 @@ impl<'a> StreamingMergeBuilder<'a> {
     }
 }
 
+/// COMET PATCH
+pub(super) fn merge_uses_row_cursor(expr: &LexOrdering, schema: &SchemaRef) -> bool {
+    let [sort] = expr.as_ref() else {
+        return true;
+    };
+    !sort.expr.data_type(schema).is_ok_and(|data_type| {
+        data_type.is_primitive()
+            || matches!(
+                data_type,
+                DataType::Utf8
+                    | DataType::Utf8View
+                    | DataType::LargeUtf8
+                    | DataType::Binary
+                    | DataType::LargeBinary
+            )
+    })
+}
+
 #[cfg(test)]
 mod tests {
     use crate::{common::collect, stream::RecordBatchStreamAdapter};