From 542d344eb52547ef178fc78599d93828cdc42c2a Mon Sep 17 00:00:00 2001 From: msaf Date: Sat, 17 Oct 2026 10:23:11 +0100 Subject: [PATCH] fix: spill a native sort's buffered batches in groups when their output merge does not fit The in-memory output merge of a sort that never spilled holds the sorted batches and the rows its cursor encodes the sort key of each run's batch into. With a multi-column key every buffered batch is a run, so these rows can outgrow what sorting frees, and the merge failed with "Additional allocation failed for ExternalSorterMerge[...]" once the sorter had filled its share of the pool, with nothing left to spill. Before that merge, estimate its memory from the batches and a sampled row encoding. If the pool cannot grant what exceeds the sorter's and the merge's reservations, spill the batches in consecutive groups whose merges fit, and merge the spill files with the multi-level merge, which already reduces its fan-in and read-ahead to the memory it gets. Co-Authored-By: Claude Opus 5.5 (1M context) --- .../src/sorts/multi_level_merge.rs | 24 +--- .../src/sorts/sort.rs | 122 +++++++++++++++++- .../src/sorts/sort/comet_memory_tests.rs | 121 +++++++++++++++++ .../src/sorts/streaming_merge.rs | 18 +++ 4 files changed, 266 insertions(+), 19 deletions(-) diff --git a/native/vendor/datafusion-physical-plan/src/sorts/multi_level_merge.rs b/native/vendor/datafusion-physical-plan/src/sorts/multi_level_merge.rs index 2dc8e40668c..51710d60a9f 100644 --- a/native/vendor/datafusion-physical-plan/src/sorts/multi_level_merge.rs +++ b/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 diff --git a/native/vendor/datafusion-physical-plan/src/sorts/sort.rs b/native/vendor/datafusion-physical-plan/src/sorts/sort.rs index 6ae97abfacd..75980ae529a 100644 --- a/native/vendor/datafusion-physical-plan/src/sorts/sort.rs +++ b/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 { + 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::>()?; + 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 { + 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 { + 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::>>()?; + 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::>>()?; + 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 = 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() diff --git a/native/vendor/datafusion-physical-plan/src/sorts/sort/comet_memory_tests.rs b/native/vendor/datafusion-physical-plan/src/sorts/sort/comet_memory_tests.rs index 8a4897a4ebe..850588cd70a 100644 --- a/native/vendor/datafusion-physical-plan/src/sorts/sort/comet_memory_tests.rs +++ b/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 { + 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 = 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| { + 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 = 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 = 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(()) +} diff --git a/native/vendor/datafusion-physical-plan/src/sorts/streaming_merge.rs b/native/vendor/datafusion-physical-plan/src/sorts/streaming_merge.rs index d4287411ed0..3dc4f2a42a9 100644 --- a/native/vendor/datafusion-physical-plan/src/sorts/streaming_merge.rs +++ b/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};