diff --git a/datafusion/physical-plan/src/sorts/sort.rs b/datafusion/physical-plan/src/sorts/sort.rs index 53c2929d4aca9..7f9c2eab6c431 100644 --- a/datafusion/physical-plan/src/sorts/sort.rs +++ b/datafusion/physical-plan/src/sorts/sort.rs @@ -30,10 +30,12 @@ use crate::metrics::{ }; use crate::projection::{ProjectionExec, make_with_child, update_ordering}; use crate::sorts::IncrementalSortIterator; -use crate::sorts::streaming_merge::{SortedSpillFile, StreamingMergeBuilder}; -use crate::spill::get_record_batch_memory_size; +use crate::sorts::streaming_merge::{ + SortedSpillFile, StreamingMergeBuilder, merges_in_row_format, +}; use crate::spill::in_progress_spill_file::InProgressSpillFile; use crate::spill::spill_manager::{GetSlicedSize, SpillManager}; +use crate::spill::{gc_view_arrays, get_record_batch_memory_size}; use crate::statistics::{ChildStats, StatisticsArgs, with_per_partition_fetch}; use crate::stream::ReservationStream; use crate::stream::{ObservedStream, RecordBatchStreamAdapter}; @@ -60,7 +62,7 @@ use datafusion_common::{ use datafusion_common_runtime::SpawnedTask; use datafusion_execution::TaskContext; use datafusion_execution::memory_pool::{ - MemoryConsumer, MemoryPool, MemoryReservation, MergeMemoryPool, + MemoryConsumer, MemoryLimit, MemoryPool, MemoryReservation, MergeMemoryPool, }; use datafusion_execution::runtime_env::RuntimeEnv; use datafusion_physical_expr::LexOrdering; @@ -708,7 +710,8 @@ impl ExternalSorter { if self.in_mem_batches.len() == 1 { let batch = self.in_mem_batches.swap_remove(0); let reservation = self.reservation.take(); - let sorted_stream = self.sort_batch_stream(batch, reservation, None)?; + let sorted_stream = + self.sort_batch_stream(batch, reservation, None, self.batch_size)?; return Ok(self.observe_if_output(sorted_stream, is_output_stream)); } @@ -725,7 +728,8 @@ impl ExternalSorter { .try_resize(get_reserved_bytes_for_record_batch(&batch)?) .map_err(Self::err_with_oom_context)?; let reservation = self.reservation.take(); - let sorted_stream = self.sort_batch_stream(batch, reservation, None)?; + let sorted_stream = + self.sort_batch_stream(batch, reservation, None, self.batch_size)?; return Ok(self.observe_if_output(sorted_stream, is_output_stream)); } @@ -781,13 +785,19 @@ impl ExternalSorter { let (sizes, shared_input_size) = self.sort_run_sizes(&runs)?; let shared_input = shared_input_size.map(|size| Arc::new(self.reservation.split(size))); + let batch_size = self.merge_run_batch_size(); let jobs = runs.into_iter() .zip(sizes) .enumerate() .map(move |(index, (batch, size))| { let reservation = self.reservation.split(size); - let sort = self.sort_run(batch, reservation, shared_input.clone())?; + let sort = self.sort_run( + batch, + reservation, + shared_input.clone(), + batch_size, + )?; let permits = Arc::clone(&concurrency.permits); Ok::<_, DataFusionError>(async move { let stream = Self::prepare_sort_run(sort, permits).await?; @@ -969,11 +979,12 @@ impl ExternalSorter { let (sizes, shared_input_size) = self.sort_run_sizes(&runs)?; let shared_input = shared_input_size.map(|size| Arc::new(self.reservation.split(size))); + let batch_size = self.merge_run_batch_size(); runs.into_iter() .zip(sizes) .map(|(run, size)| { let reservation = self.reservation.split(size); - self.sort_batch_stream(run, reservation, shared_input.clone()) + self.sort_batch_stream(run, reservation, shared_input.clone(), batch_size) }) .collect() } @@ -996,11 +1007,11 @@ impl ExternalSorter { batch: RecordBatch, reservation: MemoryReservation, shared_input: Option>, + batch_size: usize, ) -> Result { let schema = batch.schema(); - let stream = - futures::stream::once(self.sort_run(batch, reservation, shared_input)?) - .try_flatten(); + let sort = self.sort_run(batch, reservation, shared_input, batch_size)?; + let stream = futures::stream::once(sort).try_flatten(); Ok(Box::pin(RecordBatchStreamAdapter::new(schema, stream))) } @@ -1013,6 +1024,7 @@ impl ExternalSorter { batch: RecordBatch, reservation: MemoryReservation, shared_input: Option>, + batch_size: usize, ) -> Result> + Send + use<>> { let reserved_bytes = get_reserved_bytes_for_record_batch(&batch)?; @@ -1023,9 +1035,10 @@ impl ExternalSorter { } let expressions = self.expr.clone(); - let batch_size = self.batch_size; let merge_pool = Arc::clone(&self.merge_pool); let elapsed_compute = self.metrics.baseline.elapsed_compute().clone(); + // Set when `merge_run_batch_size` split this run + let split = batch_size < self.batch_size; Ok(async move { // Held until `batch` is sorted, when this block ends. The last run to @@ -1036,7 +1049,16 @@ impl ExternalSorter { let schema = batch.schema(); // Sort the batch immediately and get all output batches - let sorted_batches = sort_batch_chunked(&batch, &expressions, batch_size)?; + let mut sorted_batches = + sort_batch_chunked(&batch, &expressions, batch_size)?; + // `take` keeps the data buffers of view arrays, so splitting would + // charge every batch for all of them. + if split && sorted_batches.len() > 1 { + sorted_batches = sorted_batches + .iter() + .map(gc_view_arrays) + .collect::>()?; + } // Chunked output can retain shared buffers in every batch and // exceed the input estimate. Borrow only already-reserved spill @@ -1078,6 +1100,31 @@ impl ExternalSorter { }) } + /// Batch size of the sorted runs that feed a merge. A merge in the row format + /// holds the keys of every run's current batch at once, and these keys can be + /// larger than the batches. With a memory limit, runs therefore use an eighth + /// of `batch_size`, but at least 1024 rows to keep the per-batch merge cost + /// small. The merge output still uses `batch_size`. Runs keep `batch_size` + /// when their batches would share buffers other than the data of top-level + /// view columns, such as dictionary values, since every batch is charged for + /// them. + fn merge_run_batch_size(&self) -> usize { + let limited = !matches!( + self.runtime.memory_pool.memory_limit(), + MemoryLimit::Infinite + ); + let shares_buffers = self.schema.fields().iter().any(|field| { + let data_type = field.data_type(); + sorted_copy_keeps_buffers(data_type) + && !matches!(data_type, DataType::Utf8View | DataType::BinaryView) + }); + if limited && !shares_buffers && merges_in_row_format(&self.expr, &self.schema) { + (self.batch_size / 8).max(1024).min(self.batch_size) + } else { + self.batch_size + } + } + /// If this sort may spill, pre-allocates /// `sort_spill_reservation_bytes` of memory to guarantee memory /// left for the in memory sort/merge. @@ -2557,7 +2604,8 @@ mod tests { let reservation = sorter.reservation.new_empty(); reservation.try_grow(get_reserved_bytes_for_record_batch(&batch)?)?; let before = elapsed_compute.value(); - let sort = sorter.sort_run(batch.clone(), reservation, None)?; + let sort = + sorter.sort_run(batch.clone(), reservation, None, sorter.batch_size)?; assert_eq!(elapsed_compute.value(), before); let stream = if spawned { diff --git a/datafusion/physical-plan/src/sorts/sort/spill_tests.rs b/datafusion/physical-plan/src/sorts/sort/spill_tests.rs index 625f96be682cb..090d9972805f5 100644 --- a/datafusion/physical-plan/src/sorts/sort/spill_tests.rs +++ b/datafusion/physical-plan/src/sorts/sort/spill_tests.rs @@ -9,10 +9,10 @@ use crate::metrics::ExecutionPlanMetricsSet; use crate::spill::get_record_batch_memory_size; use crate::spill::spill_manager::GetSlicedSize; use arrow::array::{ - ArrayRef, Decimal128Array, DictionaryArray, Int8Array, Int64Array, StringArray, - StringViewArray, + ArrayRef, Decimal128Array, DictionaryArray, FixedSizeBinaryArray, Int8Array, + Int32Array, Int64Array, StringArray, StringViewArray, }; -use arrow::compute::{SortOptions, concat_batches}; +use arrow::compute::{SortOptions, cast, concat_batches}; use arrow::datatypes::{DataType, Field, Int8Type, Schema}; use arrow::record_batch::RecordBatch; use datafusion_common::config::{ExecutionOptions, SpillCompression}; @@ -679,6 +679,107 @@ async fn test_sort_preparation_keeps_output_reserved() -> Result<()> { Ok(()) } +/// Merging sorted runs in the row format holds the keys of every run's current +/// batch. Two `Int32` keys encode to 18 bytes per row against 8 bytes in the +/// batches, and a single `FixedSizeBinary(16)` key also merges in the row +/// format, so merging full-size runs needs more memory than the sort reserves +/// while buffering them. The `Utf8View` and dictionary payloads keep buffers +/// shared after sorting, which must not be charged once per split batch. +async fn check_merge_keys_larger_than_batches() -> Result<()> { + let rows = 8192; + let int_keys = |batch: i32| -> Vec { + vec![ + Arc::new(Int32Array::from_iter_values( + (0..rows).map(|i| (i * 31 + batch) % 1000), + )), + Arc::new(Int32Array::from_iter_values((0..rows).rev())), + ] + }; + let int_fields = [("a", DataType::Int32), ("b", DataType::Int32)]; + check_merge_keys(&int_fields, 2, |batch| Ok(int_keys(batch))).await?; + // Longer than 12 bytes, so views keep them in data buffers. + let strings = |batch: i32| -> StringArray { + (0..rows) + .map(|i| Some(format!("{batch}-{i:0>64}"))) + .collect() + }; + let dictionary = + DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)); + for payload in [DataType::Utf8View, dictionary] { + let fields = [&int_fields[..], &[("c", payload.clone())]].concat(); + check_merge_keys(&fields, 2, |batch| { + let payload = cast(&strings(batch), &payload)?; + Ok([int_keys(batch), vec![payload]].concat()) + }) + .await?; + } + check_merge_keys(&[("k", DataType::FixedSizeBinary(16))], 1, |batch| { + let keys = (0..rows).map(|i| u128::from(((i * 31 + batch) % 1000) as u32)); + let keys = FixedSizeBinaryArray::try_from_iter(keys.map(u128::to_be_bytes))?; + Ok(vec![Arc::new(keys)]) + }) + .await +} + +/// Sorts 64 batches of `columns` by the first `keys` fields with room to buffer +/// 16 of them, then checks the spilled output and that all memory is released. +async fn check_merge_keys( + fields: &[(&str, DataType)], + keys: usize, + columns: impl Fn(i32) -> Result>, +) -> Result<()> { + let schema = Arc::new(Schema::new( + fields + .iter() + .map(|(name, data_type)| Field::new(*name, data_type.clone(), false)) + .collect::>(), + )); + let batches = (0..64) + .map(|batch| Ok(RecordBatch::try_new(Arc::clone(&schema), columns(batch)?)?)) + .collect::>>()?; + let ordering = + LexOrdering::new((0..keys).map(|i| { + PhysicalSortExpr::new_default(Arc::new(Column::new(fields[i].0, i))) + })) + .expect("at least one sort key"); + let capacity = MERGE_BYTES + 16 * get_reserved_bytes_for_record_batch(&batches[0])?; + let pool = AdjustablePool::new(capacity); + let runtime = RuntimeEnvBuilder::new() + .with_memory_pool(Arc::clone(&pool) as Arc) + .build_arc()?; + let mut sorter = ExternalSorter::new( + 0, + Arc::clone(&schema), + ordering.clone(), + batches[0].num_rows(), + MERGE_BYTES, + 0, + SpillCompression::Uncompressed, + &ExecutionPlanMetricsSet::new(), + Arc::clone(&runtime), + )?; + for batch in &batches { + sorter.insert_batch(batch.clone()).await?; + } + assert!(sorter.spilled_before()); + let output: Vec = sorter.sort().await?.try_collect().await?; + let expected = sort_batch(&concat_batches(&schema, &batches)?, &ordering, None)?; + assert_eq!(concat_batches(&schema, &output)?, expected); + drop(sorter); + assert_released(&pool, &runtime).await; + Ok(()) +} + +#[tokio::test] +async fn test_merge_keys_larger_than_batches() -> Result<()> { + check_merge_keys_larger_than_batches().await +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn test_merge_keys_larger_than_batches_multithreaded() -> Result<()> { + check_merge_keys_larger_than_batches().await +} + #[tokio::test] async fn test_single_batch_spill_preserves_workspace_after_limit_decreases() -> Result<()> { diff --git a/datafusion/physical-plan/src/sorts/streaming_merge.rs b/datafusion/physical-plan/src/sorts/streaming_merge.rs index da341cf0cf047..d08f387702cd0 100644 --- a/datafusion/physical-plan/src/sorts/streaming_merge.rs +++ b/datafusion/physical-plan/src/sorts/streaming_merge.rs @@ -13,7 +13,7 @@ use crate::sorts::{ }; use crate::{EmptyRecordBatchStream, SendableRecordBatchStream, SpillManager}; use arrow::array::*; -use arrow::datatypes::{DataType, SchemaRef}; +use arrow::datatypes::{DataType, Schema, SchemaRef}; use datafusion_common::human_readable_size; use datafusion_common::{Result, assert_or_internal_err, internal_err}; use datafusion_execution::SpillFile; @@ -263,7 +263,8 @@ impl<'a> StreamingMergeBuilder<'a> { let reservation = reservation.expect("Reservation cannot be empty for streaming merge"); - // Special case single column comparisons with optimized cursor implementations + // Special case single column comparisons with optimized cursor implementations. + // `merges_in_row_format` must match this choice. if expressions.len() == 1 { let sort = expressions[0].clone(); let data_type = sort.expr.data_type(schema.as_ref())?; @@ -297,6 +298,26 @@ impl<'a> StreamingMergeBuilder<'a> { } } +/// Returns whether [`StreamingMergeBuilder::build`] merges on `expressions` in +/// the row format rather than with a cursor for a single column. +pub(crate) fn merges_in_row_format(expressions: &LexOrdering, schema: &Schema) -> bool { + if expressions.len() > 1 { + return true; + } + let Ok(data_type) = expressions[0].expr.data_type(schema) else { + return true; + }; + !(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};