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