diff --git a/datafusion/common/src/config.rs b/datafusion/common/src/config.rs index dec26ec8f969b..1fdefa63c13a2 100644 --- a/datafusion/common/src/config.rs +++ b/datafusion/common/src/config.rs @@ -1006,6 +1006,14 @@ config_namespace! { /// batches and merged. pub sort_in_place_threshold_bytes: usize, default = 1024 * 1024 + /// When greater than 0, a sort sorts every input batch in a background task as + /// soon as it is buffered, and merges every this many sorted batches into one + /// sorted run in a background task, so sorting overlaps reading the input and + /// little is left to merge once it ends. Each merge only reads that many batches, + /// so keep it small enough for their sort keys and rows to fit in the CPU cache. + /// 0 sorts the buffered batches once the input ends or memory runs out. + pub sort_background_merge_batches: usize, default = 8 + /// Maximum buffer capacity (in bytes) per partition for BufferExec /// inserted during sort pushdown optimization. /// diff --git a/datafusion/execution/src/config.rs b/datafusion/execution/src/config.rs index dbfd9bf1d8442..c8438c5ba0e32 100644 --- a/datafusion/execution/src/config.rs +++ b/datafusion/execution/src/config.rs @@ -534,6 +534,18 @@ impl SessionConfig { self } + /// Set [`sort_background_merge_batches`] + /// + /// [`sort_background_merge_batches`]: datafusion_common::config::ExecutionOptions::sort_background_merge_batches + pub fn with_sort_background_merge_batches( + mut self, + sort_background_merge_batches: usize, + ) -> Self { + self.options_mut().execution.sort_background_merge_batches = + sort_background_merge_batches; + self + } + /// Set [`enforce_batch_size_in_joins`] /// /// [`enforce_batch_size_in_joins`]: datafusion_common::config::ExecutionOptions::enforce_batch_size_in_joins diff --git a/datafusion/physical-plan/src/sorts/sort.rs b/datafusion/physical-plan/src/sorts/sort.rs index 13065bed078c0..08a022f57e88f 100644 --- a/datafusion/physical-plan/src/sorts/sort.rs +++ b/datafusion/physical-plan/src/sorts/sort.rs @@ -23,7 +23,7 @@ use std::fmt; use std::fmt::{Debug, Formatter}; use std::sync::Arc; -use parking_lot::RwLock; +use parking_lot::{Mutex, RwLock}; use crate::common::spawn_buffered; use crate::execution_plan::{ @@ -38,7 +38,8 @@ use crate::filter_pushdown::{ }; use crate::limit::LimitStream; use crate::metrics::{ - BaselineMetrics, ExecutionPlanMetricsSet, MetricsSet, SpillMetrics, + BaselineMetrics, Count, ExecutionPlanMetricsSet, MetricBuilder, MetricsSet, + SpillMetrics, Time, }; use crate::projection::{ProjectionExec, make_with_child, update_ordering}; use crate::sorts::IncrementalSortIterator; @@ -61,12 +62,14 @@ use arrow::array::{RecordBatch, RecordBatchOptions}; use arrow::compute::{concat_batches, lexsort_to_indices, take_arrays}; use arrow::datatypes::{DataType, SchemaRef}; use datafusion_common::config::SpillCompression; +use datafusion_common::instant::Instant; use datafusion_common::tree_node::TreeNodeRecursion; use datafusion_common::utils::memory::RecordBatchMemoryCounter; use datafusion_common::{ DataFusionError, Result, assert_or_internal_err, internal_datafusion_err, - unwrap_or_internal_err, + internal_err, unwrap_or_internal_err, }; +use datafusion_common_runtime::JoinSet; use datafusion_execution::TaskContext; use datafusion_execution::memory_pool::{ MemoryConsumer, MemoryPool, MemoryReservation, MergeMemoryPool, @@ -87,17 +90,86 @@ struct ExternalSorterMetrics { baseline: BaselineMetrics, spill_metrics: SpillMetrics, + + background: BackgroundSortMetrics, } impl ExternalSorterMetrics { fn new(metrics: &ExecutionPlanMetricsSet, partition: usize) -> Self { + let baseline = BaselineMetrics::new(metrics, partition); + let background = BackgroundSortMetrics::new( + metrics, + partition, + baseline.elapsed_compute().clone(), + ); Self { - baseline: BaselineMetrics::new(metrics, partition), + baseline, spill_metrics: SpillMetrics::new(metrics, partition), + background, + } + } +} + +/// Metrics of the [`BackgroundSorts`]. +/// +/// The sort's `elapsed_compute` covers all of its work, including what runs in the +/// background, which `background_sort_time` breaks out. The sort's own thread spent +/// `elapsed_compute` - `background_sort_time` working, and `background_sort_wait_time` +/// idle waiting for the background. +#[derive(Clone)] +struct BackgroundSortMetrics { + /// The sort's `elapsed_compute`, which background work is added to as well. + elapsed_compute: Time, + /// How many groups of sorted batches were merged into one run. + merges: Count, + /// Time the background tasks spent sorting each batch and merging them, part of the + /// sort's `elapsed_compute`. + sort_time: Time, + /// Time the sort's own thread waited for the background tasks: when memory ran out, + /// and before the final merge or a spill. Not part of `elapsed_compute`, the work it + /// waits for is. + wait_time: Time, +} + +impl BackgroundSortMetrics { + fn new( + metrics: &ExecutionPlanMetricsSet, + partition: usize, + elapsed_compute: Time, + ) -> Self { + Self { + elapsed_compute, + merges: MetricBuilder::new(metrics).counter("background_merges", partition), + sort_time: MetricBuilder::new(metrics) + .subset_time("background_sort_time", partition), + wait_time: MetricBuilder::new(metrics) + .subset_time("background_sort_wait_time", partition), + } + } + + /// Times background work into both `sort_time` and `elapsed_compute` until dropped. + fn timer(&self) -> BackgroundTimer<'_> { + BackgroundTimer { + metrics: self, + start: Instant::now(), } } } +/// See [`BackgroundSortMetrics::timer`]. +struct BackgroundTimer<'a> { + metrics: &'a BackgroundSortMetrics, + start: Instant, +} + +impl Drop for BackgroundTimer<'_> { + fn drop(&mut self) { + let elapsed = self.start.elapsed(); + self.metrics.sort_time.add_duration(elapsed); + self.metrics.elapsed_compute.add_duration(elapsed); + } +} + /// Sorts an arbitrary sized, unsorted, stream of [`RecordBatch`]es to /// a total order. Depending on the input size and memory manager /// configuration, writes intermediate results to disk ("spills") @@ -279,6 +351,13 @@ struct ExternalSorter { /// How much memory to reserve for performing in-memory sort/merges /// prior to spilling. sort_spill_reservation_bytes: usize, + + /// When non zero, every inserted batch is sorted in the background as it arrives, and + /// every this many sorted batches are merged into one sorted run, see + /// [`BackgroundSorts`]. + background_merge_batches: usize, + /// The background work on the batches inserted since the last spill. + background: BackgroundSorts, } impl ExternalSorter { @@ -339,9 +418,19 @@ impl ExternalSorter { batch_size, sort_spill_reservation_bytes, sort_in_place_threshold_bytes, + background_merge_batches: 0, + background: BackgroundSorts::default(), }) } + /// Sorts every inserted batch in the background as it arrives, and merges every + /// `background_merge_batches` sorted batches into one sorted run, see + /// [`BackgroundSorts`]. `0` disables it. + fn with_background_merge_batches(mut self, background_merge_batches: usize) -> Self { + self.background_merge_batches = background_merge_batches; + self + } + /// Appends an unsorted [`RecordBatch`] to `in_mem_batches` /// /// Updates memory usage metrics, and possibly triggers spilling to disk @@ -351,9 +440,15 @@ impl ExternalSorter { } self.reserve_memory_for_merge()?; - self.reserve_memory_for_batch_and_maybe_spill(&input) + let reserved = self + .reserve_memory_for_batch_and_maybe_spill(&input) .await?; + if self.background_merge_batches > 0 { + self.send_to_background_sort(input, reserved)?; + return self.collect_finished_background_tasks(); + } + self.in_mem_batches.push(input); Ok(()) } @@ -362,6 +457,309 @@ impl ExternalSorter { !self.finished_spill_files.is_empty() } + /// Whether there are inserted batches not yet sorted into the output or a spill file. + fn has_buffered_batches(&self) -> bool { + !self.in_mem_batches.is_empty() || !self.background.is_empty() + } + + /// Queues `batch`, for which `reserved` was just added to `self.reservation`, to be + /// sorted in the background, starting a sort task if none is running, see + /// [`BackgroundSorts`]. + fn send_to_background_sort( + &mut self, + batch: RecordBatch, + reserved: usize, + ) -> Result<()> { + // Without a counter, `reserved` covers this batch's buffers on their own, so they + // go with it. With one, buffers it shares with other buffered batches may be + // counted with one of them, so the input part stays here until every buffered + // batch is sorted, and only the sorted copy's part goes with it. + let task_share = match &self.in_mem_batches_counter { + None => reserved, + Some(_) => { + // The counter tells buffers apart by their address, which the allocator + // can give to a new buffer once this one is freed, so it is kept alive for + // as long as the counter knows it. + self.background.counted_inputs.push(batch.clone()); + batch.get_sliced_size()?.min(reserved) + } + }; + self.background.inputs += reserved - task_share; + let start_sorting = { + let mut queue = self.background.sort_queue.lock(); + queue.batches.push_back((batch, task_share)); + !std::mem::replace(&mut queue.sorting, true) + }; + self.background.pending_sorts += 1; + self.background.pending += 1; + if !start_sorting { + return Ok(()); + } + + let queue = Arc::clone(&self.background.sort_queue); + let results = self.background.results_sender.clone(); + let expr = self.expr.clone(); + let batch_size = self.batch_size; + let metrics = self.metrics.background.clone(); + self.background.tasks.spawn(async move { + loop { + // Ends the task once nothing is queued, under the same lock the next batch + // is queued with, so that batch starts a new one. + let Some((batch, reserved)) = ({ + let mut queue = queue.lock(); + let next = queue.batches.pop_front(); + queue.sorting = next.is_some(); + next + }) else { + break; + }; + let sorted = { + let _timer = metrics.timer(); + sort_batch_chunked(&batch, &expr, batch_size).and_then(|batches| { + // Only the sorted copy is kept. + drop(batch); + SortedRun::new(batches, reserved, &expr) + }) + }; + let failed = sorted.is_err(); + // Sending only fails when the sort is being dropped. After an error, the + // sort fails too, so nothing more is sorted. + if results.send(sorted.map(BackgroundOutput::Sorted)).is_err() || failed { + queue.lock().sorting = false; + break; + } + } + }); + Ok(()) + } + + /// Takes in what the background tasks produced so far, without waiting for the + /// others, and starts the merge that is due. + fn collect_finished_background_tasks(&mut self) -> Result<()> { + while let Some(joined) = self.background.tasks.try_join_next() { + Self::on_background_task_end(joined)?; + } + while let Ok(output) = self.background.results.try_recv() { + self.on_background_output(output)?; + } + self.spawn_background_merge(); + Ok(()) + } + + /// Waits for everything handed to the background. When `merge` is set, merges that + /// become due meanwhile are started and waited for as well. + async fn wait_for_background_tasks(&mut self, merge: bool) -> Result<()> { + enum Event { + Output(Option>), + TaskEnd(std::result::Result<(), tokio::task::JoinError>), + } + + let wait_time = self.metrics.background.wait_time.clone(); + let _timer = wait_time.timer(); + while self.background.pending > 0 { + let event = { + let background = &mut self.background; + let tasks_running = !background.tasks.is_empty(); + tokio::select! { + output = background.results.recv() => Event::Output(output), + Some(joined) = background.tasks.join_next(), if tasks_running => { + Event::TaskEnd(joined) + } + } + }; + match event { + Event::Output(Some(output)) => { + self.on_background_output(output)?; + if merge { + self.spawn_background_merge(); + } + } + Event::Output(None) => { + return internal_err!("Background sort results ended early"); + } + Event::TaskEnd(joined) => Self::on_background_task_end(joined)?, + } + } + Ok(()) + } + + /// Surfaces a background task's panic. Tasks that end normally sent what they produced. + fn on_background_task_end( + joined: std::result::Result<(), tokio::task::JoinError>, + ) -> Result<()> { + match joined { + Ok(()) => Ok(()), + Err(e) if e.is_panic() => std::panic::resume_unwind(e.into_panic()), + Err(e) => Err(DataFusionError::ExecutionJoin(Box::new(e))), + } + } + + /// Takes in what a background task produced. + fn on_background_output(&mut self, output: Result) -> Result<()> { + self.background.pending -= 1; + match output? { + BackgroundOutput::Sorted(mut run) => { + self.background.pending_sorts -= 1; + run.fit_reservation(&self.reservation)?; + self.background.sorted.push(run); + if self.background.pending_sorts == 0 { + self.release_background_inputs(); + } + } + BackgroundOutput::Merged(mut run) => { + self.background.merging = false; + self.metrics.background.merges.add(1); + run.fit_reservation(&self.reservation)?; + self.background.runs.push(run); + } + } + Ok(()) + } + + /// Once every buffered batch is sorted, no batch is left to keep the buffers of the + /// batches they were sorted from alive, so the input part of their reservation (see + /// [`Self::send_to_background_sort`]) goes first to what their sorted runs need to be + /// merged, and the rest is freed. + /// + /// Each sorted run gets this once: runs that have what they need to be merged then wait + /// to be merged in the background, the others to be merged at the end. + fn release_background_inputs(&mut self) { + for mut run in std::mem::take(&mut self.background.sorted) { + let top_up = run.missing_reservation().min(self.background.inputs); + run.reserved += top_up; + self.background.inputs -= top_up; + if run.missing_reservation() == 0 { + self.background.ready.push_back(run); + } else { + self.background.runs.push(run); + } + } + self.reservation.shrink(self.background.inputs); + self.background.inputs = 0; + self.background.counted_inputs.clear(); + // Nothing buffered shares buffers with what comes next anymore. + self.in_mem_batches_counter = self.new_buffered_batches_counter(); + } + + /// Merges `background_merge_batches` sorted batches, reserved for being merged, in a + /// background task, unless one is already merging. + /// + /// Like any other merge of the sort, it reserves what it needs from the merge + /// workspace first, and from the pool beyond it, failing the sort if the pool can't + /// give it. A merge that may need more than the workspace, which is all that is set + /// aside for merging while more input is taken in, is left for the final merge. + fn spawn_background_merge(&mut self) { + let merge_batches = self.background_merge_batches; + if self.background.merging || merge_batches <= 1 { + return; + } + if self.background.ready.len() < merge_batches { + return; + } + // Up to two batches per run, a batch stays in the merge until its last row is + // copied out, which can be after the next one is read, and a cursor over the keys + // of each. + let merge_size = self + .background + .ready + .iter() + .take(merge_batches) + .map(|run| 2 * run.max_batch_size + run.max_batch_cursor_size) + .sum::(); + let workspace = if self.runtime.disk_manager.tmp_files_enabled() { + self.sort_spill_reservation_bytes + } else { + 0 + }; + if merge_size > workspace { + return; + } + // The merged copy is kept, so it is reserved like the sort's other memory, and the + // merge waits while it can't be. + let merged_size = self + .background + .ready + .iter() + .take(merge_batches) + .map(SortedRun::memory_size) + .sum::(); + if self.reservation.try_grow(merged_size).is_err() { + return; + } + let runs = self + .background + .ready + .drain(..merge_batches) + .collect::>(); + let merge = merge_sorted_runs( + runs, + merged_size, + self.merge_reservation.new_empty(), + Arc::clone(&self.schema), + self.expr.clone(), + self.batch_size, + self.metrics.background.clone(), + ); + let results = self.background.results_sender.clone(); + self.background.tasks.spawn(async move { + // Only fails when the sort is being dropped. + let _ = results.send(merge.await); + }); + self.background.merging = true; + self.background.pending += 1; + } + + /// Waits for the background tasks and returns all the runs they sorted, for the final + /// merge, `merge` set, which the merges that are due run first to shrink, or for a + /// spill, which they would only delay. + async fn finish_background_tasks(&mut self, merge: bool) -> Result> { + if merge { + self.spawn_background_merge(); + } + self.wait_for_background_tasks(merge).await?; + let mut runs = std::mem::take(&mut self.background.runs); + runs.extend(self.background.ready.drain(..)); + runs.append(&mut self.background.sorted); + Ok(runs) + } + + /// Merges the runs sorted in the background into a single sorted stream. + fn merge_sorted_runs_stream( + &mut self, + runs: Vec, + is_output_stream: bool, + ) -> Result { + let mut streams = runs + .into_iter() + .map(|run| run.into_stream(Arc::clone(&self.schema), &self.reservation)) + .collect::>>()?; + + match streams.len() { + 0 => { + let empty_stream = + Box::pin(EmptyRecordBatchStream::new(Arc::clone(&self.schema))); + Ok(self.observe_if_output(empty_stream, is_output_stream)) + } + 1 => { + let stream = streams.remove(0); + Ok(self.observe_if_output(stream, is_output_stream)) + } + _ => StreamingMergeBuilder::new() + .with_streams(streams) + .with_schema(Arc::clone(&self.schema)) + .with_expressions(&self.expr.clone()) + .with_metrics(if is_output_stream { + self.metrics.baseline.clone() + } else { + self.metrics.baseline.intermediate() + }) + .with_batch_size(self.batch_size) + .with_fetch(None) + .with_reservation(self.merge_reservation.new_empty()) + .build(), + } + } + /// Returns the final sorted output of all batches inserted via /// [`Self::insert_batch`] as a stream of [`RecordBatch`]es. /// @@ -376,7 +774,7 @@ impl ExternalSorter { // Sort `in_mem_batches` and spill it first. If there are many // `in_mem_batches` and the memory limit is almost reached, merging // them with the spilled files at the same time might cause OOM. - if !self.in_mem_batches.is_empty() { + if self.has_buffered_batches() { self.sort_and_spill_in_mem_batches().await?; } @@ -399,11 +797,19 @@ impl ExternalSorter { .with_merge_pool(Arc::clone(&self.merge_pool)) .build() } else { + let background_runs = if self.background_merge_batches > 0 { + Some(self.finish_background_tasks(true).await?) + } else { + None + }; // Final output needs no reserve for future spills. Return unused // workspace so another sorter can start while this stream is alive. self.merge_pool.release_unused(); self.merge_reservation.free(); - self.in_mem_sort_stream(true, true) + match background_runs { + Some(runs) => self.merge_sorted_runs_stream(runs, true), + None => self.in_mem_sort_stream(true, true), + } } } @@ -514,20 +920,29 @@ impl ExternalSorter { async fn try_sort_and_spill_in_mem_batches(&mut self) -> Result<()> { assert_or_internal_err!( - !self.in_mem_batches.is_empty(), + self.has_buffered_batches(), "in_mem_batches must not be empty when attempting to sort and spill" ); + let background_runs = if self.background_merge_batches > 0 { + Some(self.finish_background_tasks(false).await?) + } else { + None + }; + // Reuse the pre-reserved workspace across cursor, encoded-row, and // batch reservations. Returning it to the execution pool here can // make spilling fail if another task consumes it or our share shrinks. self.merge_reservation.free(); - let mut sorted_stream = self.in_mem_sort_stream( - false, - // No coalescing on the spill path: it raises per-run peak memory. - false, - )?; + let mut sorted_stream = match background_runs { + Some(runs) => self.merge_sorted_runs_stream(runs, false)?, + None => self.in_mem_sort_stream( + false, + // No coalescing on the spill path: it raises per-run peak memory. + false, + )?, + }; // After `in_mem_sort_stream()` is constructed, all `in_mem_batches` is taken // to construct a globally sorted stream. assert_or_internal_err!( @@ -924,6 +1339,19 @@ impl ExternalSorter { // Reserve headroom for next merge sort if self.runtime.disk_manager.tmp_files_enabled() { let size = self.sort_spill_reservation_bytes; + if self.background_merge_batches > 0 { + // Kept as idle merge workspace, which the background merges use too, see + // `Self::spawn_background_merge`. Those are all done before any other + // merge, and leave it as it was. + if self.merge_pool.reserved() < size { + self.merge_reservation + .try_resize(size) + .map_err(Self::err_with_oom_context)?; + self.merge_pool.retain(size); + self.merge_reservation.free(); + } + return Ok(()); + } if self.merge_reservation.size() != size { self.merge_reservation .try_resize(size) @@ -937,22 +1365,32 @@ impl ExternalSorter { /// Reserves memory to be able to accommodate the given batch. /// If memory is scarce, tries to spill current in-memory batches to disk first. + /// Returns how much was reserved. async fn reserve_memory_for_batch_and_maybe_spill( &mut self, input: &RecordBatch, - ) -> Result<()> { + ) -> Result { let size = get_reserved_bytes_for_next_record_batch( input, self.in_mem_batches_counter.as_mut(), )?; match self.reservation.try_grow(size) { - Ok(_) => Ok(()), + Ok(_) => Ok(size), Err(e) => { - if self.in_mem_batches.is_empty() { + if !self.has_buffered_batches() { return Err(Self::err_with_oom_context(e)); } + // The background tasks free what their batches no longer need once + // they are done, which may be enough to avoid spilling. + if self.background.pending > 0 { + self.wait_for_background_tasks(true).await?; + if self.reservation.try_grow(size).is_ok() { + return Ok(size); + } + } + // Spill and try again. Nothing is buffered after the spill, so // `input` shares no buffers with a buffered batch. self.sort_and_spill_in_mem_batches().await?; @@ -962,7 +1400,8 @@ impl ExternalSorter { )?; self.reservation .try_grow(size) - .map_err(Self::err_with_oom_context) + .map_err(Self::err_with_oom_context)?; + Ok(size) } } } @@ -998,6 +1437,341 @@ impl ExternalSorter { } } +/// Batches sorted into one order, with the memory they hold. +struct SortedRun { + batches: Vec, + /// How much of the sort's reservation is for this run: its batches, and once reserved + /// for it (see [`Self::missing_reservation`]), what merging them needs. Only the sort + /// holds a reservation, so the pool judges its growth by all it holds, and a run gets + /// one of its own only to be merged at the end, see [`Self::into_stream`]. + reserved: usize, + /// The memory of the batches. + memory_size: usize, + /// The memory of the largest batch. + max_batch_size: usize, + /// Upper bound of what a merge's cursor holds for one batch of this run: the sort keys + /// of the largest batch, plus their row format overhead. + max_batch_cursor_size: usize, +} + +impl SortedRun { + fn new( + batches: Vec, + reserved: usize, + expr: &LexOrdering, + ) -> Result { + let mut memory_size = 0; + let mut max_batch_size = 0; + let mut max_batch_cursor_size = 0; + for batch in &batches { + let batch_size = get_record_batch_memory_size(batch); + memory_size += batch_size; + max_batch_size = max_batch_size.max(batch_size); + max_batch_cursor_size = max_batch_cursor_size.max(cursor_size(batch, expr)?); + } + Ok(Self { + batches, + reserved, + memory_size, + max_batch_size, + max_batch_cursor_size, + }) + } + + fn memory_size(&self) -> usize { + self.memory_size + } + + /// What the run needs reserved to be merged by the final merge, or a spill: its + /// batches, and a cursor over the keys of the batch the merge reads. The merge + /// accounts for the batches it holds with what their runs release, and its output is + /// reserved by whoever takes it. A background merge reserves the merged copy it keeps + /// when it starts, see [`ExternalSorter::spawn_background_merge`]. + fn reservation_to_merge(&self) -> usize { + self.memory_size() + self.max_batch_cursor_size + } + + /// How much more [`Self::reservation_to_merge`] needs than is reserved. + fn missing_reservation(&self) -> usize { + self.reservation_to_merge().saturating_sub(self.reserved) + } + + /// Fits the reservation to the run, now that it was sorted or merged and only its + /// batches are left, freeing what the batches it came from needed beyond what merging + /// it needs (a slice of larger buffers, buffers with spare capacity). + /// + /// The batches themselves are already in memory, so if they take more than is + /// reserved, they are reserved in full, failing like the sort would without background + /// sorting when the pool can't take them. + fn fit_reservation(&mut self, reservation: &MemoryReservation) -> Result<()> { + let memory_size = self.memory_size(); + let to_merge = self.reservation_to_merge(); + if self.reserved > to_merge { + reservation.shrink(self.reserved - to_merge); + self.reserved = to_merge; + } else if self.reserved < memory_size { + reservation + .try_grow(memory_size - self.reserved) + .map_err(ExternalSorter::err_with_oom_context)?; + self.reserved = memory_size; + } + Ok(()) + } + + /// Streams the run's batches into the final merge, or a spill, taking its share of + /// `reservation`, the sort's, releasing what was kept to merge them, and then the + /// memory of each batch as it is emitted, so the merge can account for it instead. + fn into_stream( + self, + schema: SchemaRef, + reservation: &MemoryReservation, + ) -> Result { + let memory_size = self.memory_size(); + let run_reservation = reservation.split(self.reserved); + if self.reserved > memory_size { + run_reservation.shrink(self.reserved - memory_size); + } else { + run_reservation + .try_grow(memory_size - self.reserved) + .map_err(ExternalSorter::err_with_oom_context)?; + } + Ok(Box::pin(ReservationStream::new( + Arc::clone(&schema), + Box::pin(RecordBatchStreamAdapter::new( + schema, + futures::stream::iter(self.batches.into_iter().map(Ok)), + )), + run_reservation, + ))) + } +} + +/// Sorting and merging that an [`ExternalSorter`] runs in the background, see +/// `datafusion.execution.sort_background_merge_batches`. +/// +/// The admitted batches are sorted one after another as they arrive, by a task that is +/// started when a batch is queued and none is sorting, and ends once nothing is queued, +/// and every `background_merge_batches` sorted batches are merged into one sorted run in a +/// task of its own, one merge at a time. So a sort and a merge can run alongside reading +/// the input, but no more: each partition of the plan already has a thread of its own, and +/// more work in the background only takes CPU time from them. A merge only reads those few +/// batches, so their keys and the rows copied out of them stay in cache, unlike one merge +/// of every batch at the end. Merged runs are kept apart until the final merge, they are +/// never merged again into later runs. +/// +/// ```text +/// input batch +/// │ +/// ▼ +/// ExternalSorter::insert_batch ── reserve the batch's memory ──┐ not enough +/// │ ▼ +/// │ wait for the background tasks, which free what +/// │ their batches no longer need, and try again; +/// │ spill (merge everything into a file) if still +/// │ not enough +/// │ queue to be sorted, with the batch's reservation (but the part for +/// │ buffers it may share with other buffered batches) +/// ┌ ─ ┼ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ┐ +/// background: at most one sort task, and one merge task alongside it +/// │ ▼ │ +/// sort task: sort each queued batch on its own, drop the unsorted input +/// │ │ │ +/// ▼ +/// │ sorted: [s1] [s2] .. [sN] (one sorted run per batch) │ +/// │ every `background_merge_batches` of them reserved for merging +/// │ ▼ │ +/// merge task (one at a time): merge sort them, with the merge workspace +/// │ │ (if it may not hold the merge, they wait for the final one)│ +/// ▼ +/// │ runs: [run 1] [run 2] .. [run K] (never merged again here) │ +/// └ ─ ┼ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ┘ +/// │ the input ended, or memory ran out: wait for every task +/// ▼ +/// ┌──────────────────────────────────────────┐ +/// │ merge of [run 1] .. [run K] [s1] .. [sM] │──▶ sorted output, or a spill +/// │ (with the spill files, if any) │ file when memory ran out +/// └──────────────────────────────────────────┘ +/// ``` +/// +/// # Memory +/// +/// Every batch is admitted with the same reservation as without sorting in the +/// background, buffers it shares with other buffered batches counted once. The part for +/// its sorted copy goes with it into its task. The part for its input buffers stays with +/// the sort until every buffered batch is sorted, as another one may still hold buffers +/// it shares, unless batches are not counted together, when it goes along too. +/// +/// Once a batch is sorted, or batches are merged, the run's reservation is fitted to its +/// batches and what merging them needs, see [`SortedRun::fit_reservation`], freeing what +/// the batches it came from took beyond that. Released input reservations first make up +/// what sorted runs miss to be merged, and only runs that have it are merged in the +/// background, and a merge reserves the merged copy it keeps before it starts, so it is +/// always covered. What a merge holds besides (the +/// current batches of its runs, a cursor over their keys), it reserves like any other +/// merge of the sort: from the sort's merge workspace, which the sort keeps reserved +/// anyway, and which every background merge is done with before any other merge uses it, +/// and beyond it from the pool, failing the sort if the pool can't give it. +/// +/// The batches queued to be sorted are only the ones admitted, each with its reservation, +/// so the queue is bounded by memory, not by length. +/// +/// All of it is held in the sort's own reservation, not in reservations of the batches +/// and runs, as the pool may judge whether a reservation can grow by its own size, which +/// would then leave out what the sort holds in the others. It is adjusted by the sort +/// itself as it collects what the background produced, never by the background, so it +/// always knows what that holds. When the sort is dropped, the running tasks are aborted, +/// and its reservation, with what they hold, is queued, or was produced and not collected, +/// is freed. +struct BackgroundSorts { + /// Batches waiting to be sorted. + sort_queue: Arc>, + /// The running sort and merge tasks, to abort them when the sort is dropped, and to + /// notice if one panics. + tasks: JoinSet<()>, + /// What the sort and merge tasks produce. + results: tokio::sync::mpsc::UnboundedReceiver>, + results_sender: tokio::sync::mpsc::UnboundedSender>, + /// How many batches and merges were handed to the background and did not come back. + pending: usize, + /// How many batches were handed to the background and are not sorted. + pending_sorts: usize, + /// How much of the sort's reservation is for the input buffers of the batches handed + /// to the background since every batch was last sorted, see + /// [`ExternalSorter::release_background_inputs`]. + inputs: usize, + /// Those batches, kept until then as the buffer counter knows their buffers, see + /// [`ExternalSorter::send_to_background_sort`]. + counted_inputs: Vec, + /// Whether a merge is running, only one runs at a time. + merging: bool, + /// Sorted batches, waiting for the input reservation of the batches they came from to + /// be released, see [`ExternalSorter::release_background_inputs`]. + sorted: Vec, + /// Sorted batches reserved for being merged, waiting for enough of them to be merged. + ready: std::collections::VecDeque, + /// Runs that are no longer merged in the background: merged runs, and sorted batches + /// that are merged at the end. + runs: Vec, +} + +impl Default for BackgroundSorts { + fn default() -> Self { + let (results_sender, results) = tokio::sync::mpsc::unbounded_channel(); + Self { + sort_queue: Default::default(), + tasks: JoinSet::new(), + results, + results_sender, + pending: 0, + pending_sorts: 0, + inputs: 0, + counted_inputs: vec![], + merging: false, + sorted: vec![], + ready: Default::default(), + runs: vec![], + } + } +} + +/// Batches waiting to be sorted, see [`BackgroundSorts`]. +#[derive(Default)] +struct SortQueue { + /// Each with how much of the sort's reservation goes to sorting it. + batches: std::collections::VecDeque<(RecordBatch, usize)>, + /// Whether a task is sorting them. + sorting: bool, +} + +impl BackgroundSorts { + fn is_empty(&self) -> bool { + self.pending == 0 + && self.sorted.is_empty() + && self.ready.is_empty() + && self.runs.is_empty() + } +} + +/// What a background task produces. +enum BackgroundOutput { + /// A sorted batch. + Sorted(SortedRun), + /// The merged run, with the reservations of the runs merged into it. + Merged(SortedRun), +} + +/// Merges `runs` into a single run, see [`BackgroundSorts`], reserving what it holds +/// besides them and their merged copy, for which `merged_size_reserved` of the sort's +/// reservation was reserved, in `merge_reservation`, a reservation of the sort's merge +/// workspace pool. +async fn merge_sorted_runs( + runs: Vec, + merged_size_reserved: usize, + merge_reservation: MemoryReservation, + schema: SchemaRef, + expr: LexOrdering, + batch_size: usize, + metrics: BackgroundSortMetrics, +) -> Result { + let _timer = metrics.timer(); + // The runs' reservations, and what was reserved for the merged copy, are kept whole + // while the runs' batches are dropped as the merge is done with them. + let reserved = + runs.iter().map(|run| run.reserved).sum::() + merged_size_reserved; + let streams = runs + .into_iter() + .map(|run| { + Box::pin(RecordBatchStreamAdapter::new( + Arc::clone(&schema), + futures::stream::iter(run.batches.into_iter().map(Ok)), + )) as SendableRecordBatchStream + }) + .collect(); + + let mut merged_stream = StreamingMergeBuilder::new() + .with_streams(streams) + .with_schema(Arc::clone(&schema)) + .with_expressions(&expr) + // Not reported, the merge is timed by `metrics`, and its output is not the sort's + // output. + .with_metrics(BaselineMetrics::new(&ExecutionPlanMetricsSet::new(), 0)) + .with_batch_size(batch_size) + .with_fetch(None) + .with_reservation(merge_reservation) + .build()?; + let mut batches = vec![]; + while let Some(batch) = merged_stream.next().await { + batches.push(batch?); + } + // Frees what the merge reserved. + drop(merged_stream); + + Ok(BackgroundOutput::Merged(SortedRun::new( + batches, reserved, &expr, + )?)) +} + +/// Upper bound of what a merge's cursor holds for `batch`: its sort key, or its sort keys +/// encoded in the row format when there are several (one null byte per key and an offset +/// per row on top). +fn cursor_size(batch: &RecordBatch, expr: &LexOrdering) -> Result { + let keys_size = expr + .iter() + .map(|sort_expr| { + Ok(sort_expr + .expr + .evaluate(batch)? + .into_array(batch.num_rows())? + .get_array_memory_size()) + }) + .sum::>()?; + let row_overhead = match expr.len() { + 1 => 0, + keys => batch.num_rows() * (keys + size_of::()), + }; + Ok(keys_size + row_overhead) +} + /// Estimate how much memory is needed to sort a `RecordBatch`. /// /// This is used to pre-reserve memory for the sort/merge. The sort/merge process involves @@ -1654,7 +2428,10 @@ impl ExecutionPlan for SortExec { context.session_config().spill_compression(), &self.metrics_set, context.runtime_env(), - )?; + )? + .with_background_merge_batches( + execution_options.sort_background_merge_batches, + ); Ok(Box::pin(RecordBatchStreamAdapter::new( self.schema(), futures::stream::once(async move { @@ -4425,4 +5202,340 @@ mod tests { assert_eq!(num_rows(&topk(false), None)?, Precision::Exact(10)); Ok(()) } + + /// Batches of an unsorted two column key (with nulls and repeated values) and a payload. + fn background_sort_input() -> Result<(SchemaRef, Vec)> { + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int32, true), + Field::new("b", DataType::Utf8, true), + Field::new("payload", DataType::Int64, false), + ])); + let batches = (0..37) + .map(|batch_idx| { + let rows = 0..100 + batch_idx * 7; + let a: Int32Array = rows + .clone() + .map(|row| { + let v = (batch_idx * 1000 + row) * 7919 % 211; + (v % 13 != 0).then_some(v) + }) + .collect(); + let b: StringArray = rows + .clone() + .map(|row| Some(format!("key-{}", (batch_idx + row) % 17))) + .collect(); + let payload: Int64Array = + rows.map(|row| (batch_idx * 1000 + row) as i64).collect(); + Ok(RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(a), Arc::new(b), Arc::new(payload)], + )?) + }) + .collect::>>()?; + Ok((schema, batches)) + } + + fn background_sort_ordering(schema: &Schema) -> Result { + Ok([ + PhysicalSortExpr::new(col("a", schema)?, SortOptions::default()), + PhysicalSortExpr::new( + col("b", schema)?, + SortOptions { + descending: true, + nulls_first: false, + }, + ), + PhysicalSortExpr::new_default(col("payload", schema)?), + ] + .into()) + } + + /// Sorts [`background_sort_input`], checking all memory is returned once it is done. + async fn run_background_sort( + background_merge_batches: usize, + memory_limit: Option, + sort_spill_reservation_bytes: Option, + ) -> Result<(Vec, MetricsSet)> { + let (schema, batches) = background_sort_input()?; + let mut runtime = RuntimeEnvBuilder::new(); + if let Some(limit) = memory_limit { + runtime = runtime.with_memory_limit(limit, 1.0); + } + let mut session_config = SessionConfig::new() + .with_batch_size(64) + .with_sort_background_merge_batches(background_merge_batches); + if let Some(bytes) = sort_spill_reservation_bytes { + session_config = session_config.with_sort_spill_reservation_bytes(bytes); + } + let task_ctx = Arc::new( + TaskContext::default() + .with_session_config(session_config) + .with_runtime(runtime.build_arc()?), + ); + let sort_exec = Arc::new(SortExec::new( + background_sort_ordering(&schema)?, + TestMemoryExec::try_new_exec(&[batches], schema, None)?, + )); + let result = collect(Arc::clone(&sort_exec) as _, Arc::clone(&task_ctx)).await?; + assert_eq!( + task_ctx.runtime_env().memory_pool.reserved(), + 0, + "The sort should have returned all memory used back to the memory manager" + ); + Ok((result, sort_exec.metrics().unwrap())) + } + + fn metric(metrics: &MetricsSet, name: &str) -> usize { + metrics + .sum_by_name(name) + .map_or(0, |value| value.as_usize()) + } + + /// Sorting in the background, merging every few sorted batches, gives the same output + /// as sorting everything at the end, in memory and when spilling. + async fn assert_background_sort_matches() -> Result<()> { + let (expected, _) = run_background_sort(0, None, None).await?; + let expected = batches_to_string(&expected); + for background_merge_batches in [1, 3, 8, 100] { + let (result, metrics) = + run_background_sort(background_merge_batches, None, None).await?; + assert_eq!(batches_to_string(&result), expected); + assert_eq!(metrics.spill_count(), Some(0)); + let background_sort_time = metric(&metrics, "background_sort_time"); + assert!(background_sort_time > 0); + // The background work is part of the operator's elapsed compute. + assert!(metrics.elapsed_compute().unwrap() >= background_sort_time); + let merges = metric(&metrics, "background_merges"); + assert_eq!( + merges > 0, + background_merge_batches > 1 && background_merge_batches <= 37, + "{merges} background merges of {background_merge_batches} batches" + ); + + let (result, metrics) = run_background_sort( + background_merge_batches, + Some(10 * 1024 * 1024 + 64 * 1024), + None, + ) + .await?; + assert_eq!(batches_to_string(&result), expected); + assert!( + metrics.spill_count().unwrap() > 0, + "expected the memory limit to cause spills" + ); + } + Ok(()) + } + + #[tokio::test] + async fn test_background_sort_matches() -> Result<()> { + assert_background_sort_matches().await + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] + async fn test_background_sort_matches_multi_thread() -> Result<()> { + assert_background_sort_matches().await + } + + /// Sorting in the background never needs more memory than sorting at the end: under + /// every memory limit the plain sort completes with, the background sort completes too, + /// with the same output and with all of its memory returned, across merge workspace + /// sizes. + async fn assert_background_sort_needs_no_more_memory() -> Result<()> { + let (expected, _) = run_background_sort(0, None, None).await?; + let expected = batches_to_string(&expected); + for workspace in [16 * 1024, 64 * 1024, 1024 * 1024] { + let mut baseline_completed = 0; + for limit in (1..=48).map(|step| workspace + step * 8 * 1024) { + if run_background_sort(0, Some(limit), Some(workspace)) + .await + .is_err() + { + continue; + } + baseline_completed += 1; + for background_merge_batches in [2, 8] { + let (result, _) = run_background_sort( + background_merge_batches, + Some(limit), + Some(workspace), + ) + .await + .unwrap_or_else(|e| { + panic!( + "background sort of {background_merge_batches} batches failed \ + under a {limit} bytes limit the plain sort completes with: {e}" + ) + }); + assert_eq!(batches_to_string(&result), expected); + } + } + assert!(baseline_completed > 0, "no memory limit was large enough"); + } + Ok(()) + } + + #[tokio::test] + async fn test_background_sort_needs_no_more_memory() -> Result<()> { + assert_background_sort_needs_no_more_memory().await + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] + async fn test_background_sort_needs_no_more_memory_multi_thread() -> Result<()> { + assert_background_sort_needs_no_more_memory().await + } + + fn background_sorter( + schema: &SchemaRef, + runtime: &Arc, + metrics: &ExecutionPlanMetricsSet, + background_merge_batches: usize, + ) -> Result { + Ok(ExternalSorter::new( + 0, + Arc::clone(schema), + background_sort_ordering(schema)?, + 64, + 1024 * 1024, + 1024 * 1024, + SpillCompression::Uncompressed, + metrics, + Arc::clone(runtime), + )? + .with_background_merge_batches(background_merge_batches)) + } + + /// At most one sort task and one merge task run in the background, however many batches + /// are queued. + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] + async fn test_background_sort_runs_at_most_one_sort_and_one_merge() -> Result<()> { + let (schema, batches) = background_sort_input()?; + let runtime = RuntimeEnvBuilder::new().build_arc()?; + let metrics = ExecutionPlanMetricsSet::new(); + let mut sorter = background_sorter(&schema, &runtime, &metrics, 2)?; + for batch in &batches { + sorter.insert_batch(batch.clone()).await?; + assert!( + sorter.background.tasks.len() <= 2, + "{} background tasks", + sorter.background.tasks.len() + ); + } + sorter.wait_for_background_tasks(true).await?; + assert!(sorter.background.tasks.len() <= 2); + assert!(metric(&metrics.clone_inner(), "background_merges") > 0); + Ok(()) + } + + /// Dropping a sorter while batches are being sorted or merged in the background returns + /// all of their memory. + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] + async fn test_background_sort_drop_frees_memory() -> Result<()> { + let (schema, batches) = background_sort_input()?; + let runtime = RuntimeEnvBuilder::new().build_arc()?; + for inserted in [1, 2, 5, 17, batches.len()] { + let metrics = ExecutionPlanMetricsSet::new(); + let mut sorter = background_sorter(&schema, &runtime, &metrics, 4)?; + for batch in &batches[..inserted] { + sorter.insert_batch(batch.clone()).await?; + } + assert!(runtime.memory_pool.reserved() > 0); + drop(sorter); + + // The background tasks are aborted on drop, and free what they hold once + // they stop. + for _ in 0..1000 { + if runtime.memory_pool.reserved() == 0 { + break; + } + tokio::time::sleep(std::time::Duration::from_millis(1)).await; + } + assert_eq!( + runtime.memory_pool.reserved(), + 0, + "dropping the sorter after {inserted} batches should return all its memory" + ); + } + Ok(()) + } + + /// Slices of one batch share its buffers, which are reserved once, with the first + /// slice, and stay reserved until every slice is sorted. Then what the sorted runs need + /// to be merged is made up from them, and the rest is freed. + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] + async fn test_background_sort_keeps_shared_buffers_until_all_sorted() -> Result<()> { + let (schema, batches) = background_sort_input()?; + let parent = concat_batches(&schema, &batches)?; + let runtime = RuntimeEnvBuilder::new().build_arc()?; + let metrics = ExecutionPlanMetricsSet::new(); + // Never merged, so the sorted batches are what is left. + let mut sorter = background_sorter(&schema, &runtime, &metrics, 1000)?; + let slices = (0..10) + .map(|i| parent.slice(i * 100, 100)) + .collect::>(); + let mut counter = RecordBatchMemoryCounter::new(); + let mut reserved_for_slices = 0; + for slice in &slices { + reserved_for_slices += + get_reserved_bytes_for_next_record_batch(slice, Some(&mut counter))?; + sorter.insert_batch(slice.clone()).await?; + } + // Shared buffers are counted once. + assert!(reserved_for_slices < 2 * get_record_batch_memory_size(&parent)); + + sorter.wait_for_background_tasks(false).await?; + + // Released once every slice was sorted, each is either ready to be merged or left + // for the final merge. + assert!(sorter.background.sorted.is_empty()); + let sorted = sorter + .background + .ready + .iter() + .chain(&sorter.background.runs) + .collect::>(); + assert_eq!(sorted.len(), slices.len()); + assert_eq!(sorter.background.pending_sorts, 0); + // The input reservation was released into the runs or freed, so the sort's + // reservation is just the runs'. + assert_eq!(sorter.background.inputs, 0); + let runs_reserved: usize = sorted.iter().map(|run| run.reserved).sum(); + assert_eq!(sorter.reservation.size(), runs_reserved); + assert_eq!( + runtime.memory_pool.reserved(), + runs_reserved + sorter.merge_pool.reserved() + ); + assert!(runs_reserved <= reserved_for_slices); + for run in sorted { + assert!(run.reserved >= run.memory_size()); + } + Ok(()) + } + + /// Background merges run in the merge workspace the sort keeps reserved, so they run + /// even when the pool has nothing to spare beyond it and the batches' reservations. + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] + async fn test_background_merges_run_in_merge_workspace() -> Result<()> { + let (_, batches) = background_sort_input()?; + let workspace = 1024 * 1024; + let mut counter = RecordBatchMemoryCounter::new(); + let batches_reserved: usize = batches + .iter() + .map(|batch| { + get_reserved_bytes_for_next_record_batch(batch, Some(&mut counter)) + }) + .sum::>()?; + let (_, metrics) = run_background_sort( + 8, + Some(workspace + batches_reserved + 4 * 1024), + Some(workspace), + ) + .await?; + assert_eq!(metrics.spill_count(), Some(0)); + assert!( + metric(&metrics, "background_merges") > 0, + "no merge ran in the merge workspace" + ); + Ok(()) + } } diff --git a/datafusion/sqllogictest/test_files/copy.slt b/datafusion/sqllogictest/test_files/copy.slt index 6f62827412738..da7093d16030c 100644 --- a/datafusion/sqllogictest/test_files/copy.slt +++ b/datafusion/sqllogictest/test_files/copy.slt @@ -229,7 +229,7 @@ EXPLAIN ANALYZE COPY (SELECT col1, upper(col2) AS col2_upper FROM source_table O ---- Plan with Metrics 01)DataSinkExec: sink=ParquetSink(file_groups=[]), metrics=[elapsed_compute=, bytes_written=, rows_written=2] -02)--SortExec: expr=[col1@0 ASC NULLS LAST], preserve_partitioning=[false], metrics=[output_rows=2, elapsed_compute=, output_bytes=, output_batches=, spill_count=0, spilled_bytes=0.0 B, spilled_rows=0] +02)--SortExec: expr=[col1@0 ASC NULLS LAST], preserve_partitioning=[false], metrics=[output_rows=2, elapsed_compute=, output_bytes=, output_batches=, spill_count=0, spilled_bytes=0.0 B, spilled_rows=0, background_merges=0, background_sort_time=, background_sort_wait_time=] 03)----ProjectionExec: expr=[col1@0 as col1, upper(col2@1) as col2_upper], metrics=[output_rows=2, elapsed_compute=, output_bytes=, output_batches=1, expr_0_eval_time=, expr_1_eval_time=] 04)------DataSourceExec: partitions=1, partition_sizes=[1], metrics=[] diff --git a/datafusion/sqllogictest/test_files/copy_avro.slt b/datafusion/sqllogictest/test_files/copy_avro.slt index a95786a968031..8b0c85cf00b67 100644 --- a/datafusion/sqllogictest/test_files/copy_avro.slt +++ b/datafusion/sqllogictest/test_files/copy_avro.slt @@ -138,7 +138,7 @@ EXPLAIN ANALYZE COPY (SELECT col1, upper(col2) AS col2_upper FROM source_table O ---- Plan with Metrics 01)DataSinkExec: sink=AvroFileSink(file_groups=[]), metrics=[] -02)--SortExec: expr=[col1@0 ASC NULLS LAST], preserve_partitioning=[false], metrics=[output_rows=2, elapsed_compute=, output_bytes=, output_batches=, spill_count=0, spilled_bytes=0.0 B, spilled_rows=0] +02)--SortExec: expr=[col1@0 ASC NULLS LAST], preserve_partitioning=[false], metrics=[output_rows=2, elapsed_compute=, output_bytes=, output_batches=, spill_count=0, spilled_bytes=0.0 B, spilled_rows=0, background_merges=0, background_sort_time=, background_sort_wait_time=] 03)----ProjectionExec: expr=[col1@0 as col1, upper(col2@1) as col2_upper], metrics=[output_rows=2, elapsed_compute=, output_bytes=, output_batches=1, expr_0_eval_time=, expr_1_eval_time=] 04)------DataSourceExec: partitions=1, partition_sizes=[1], metrics=[] diff --git a/datafusion/sqllogictest/test_files/information_schema.slt b/datafusion/sqllogictest/test_files/information_schema.slt index fb6d2d9dc78ec..a5bbf59f4591f 100644 --- a/datafusion/sqllogictest/test_files/information_schema.slt +++ b/datafusion/sqllogictest/test_files/information_schema.slt @@ -281,6 +281,7 @@ datafusion.execution.skip_partial_aggregation_probe_rows_threshold 100000 datafusion.execution.skip_physical_aggregate_schema_check false datafusion.execution.soft_max_bytes_per_output_file 4294967295 datafusion.execution.soft_max_rows_per_output_file 50000000 +datafusion.execution.sort_background_merge_batches 8 datafusion.execution.sort_in_place_threshold_bytes 1048576 datafusion.execution.sort_pushdown_buffer_capacity 1073741824 datafusion.execution.sort_spill_reservation_bytes 10485760 @@ -446,6 +447,7 @@ datafusion.execution.skip_partial_aggregation_probe_rows_threshold 100000 Number datafusion.execution.skip_physical_aggregate_schema_check false When set to true, skips verifying that the schema produced by planning the input of `LogicalPlan::Aggregate` exactly matches the schema of the input plan. When set to false, if the schema does not match exactly (including nullability and metadata), a planning error will be raised. This is used to workaround bugs in the planner that are now caught by the new schema verification step. datafusion.execution.soft_max_bytes_per_output_file 4294967295 Target encoded size in bytes of output files when writing multiple. Writers asynchronously report the cumulative encoded size as they process RecordBatches. The final file size may exceed this limit due to batches buffered before the limit is observed, the size of a batch, and file metadata written when the file is finalized. datafusion.execution.soft_max_rows_per_output_file 50000000 Target number of rows in output files when writing multiple. This is a soft max, so it can be exceeded slightly. There also will be one file smaller than the limit if the total number of rows written is not roughly divisible by the soft max +datafusion.execution.sort_background_merge_batches 8 When greater than 0, a sort sorts every input batch in a background task as soon as it is buffered, and merges every this many sorted batches into one sorted run in a background task, so sorting overlaps reading the input and little is left to merge once it ends. Each merge only reads that many batches, so keep it small enough for their sort keys and rows to fit in the CPU cache. 0 sorts the buffered batches once the input ends or memory runs out. datafusion.execution.sort_in_place_threshold_bytes 1048576 When sorting, below what size should data be concatenated and sorted in a single RecordBatch rather than sorted in batches and merged. datafusion.execution.sort_pushdown_buffer_capacity 1073741824 Maximum buffer capacity (in bytes) per partition for BufferExec inserted during sort pushdown optimization. When PushdownSort eliminates a SortExec under SortPreservingMergeExec, a BufferExec is inserted to replace SortExec's buffering role. This prevents I/O stalls by allowing the scan to run ahead of the merge. This uses strictly less memory than the SortExec it replaces (which buffers the entire partition). The buffer respects the global memory pool limit. Setting this to a large value is safe — actual memory usage is bounded by partition size and global memory limits. datafusion.execution.sort_spill_reservation_bytes 10485760 Specifies the reserved memory for each spillable sort operation to facilitate an in-memory merge. When a sort operation spills to disk, the in-memory data must be sorted and merged before being written to a file. This setting reserves a specific amount of memory for that in-memory sort/merge process. Note: This setting is irrelevant if the sort operation cannot spill (i.e., if there's no `DiskManager` configured). diff --git a/docs/source/user-guide/configs.md b/docs/source/user-guide/configs.md index 27356ab8c1b77..072acafd80da8 100644 --- a/docs/source/user-guide/configs.md +++ b/docs/source/user-guide/configs.md @@ -128,6 +128,7 @@ The following configuration settings are available: | datafusion.execution.spill_compression | uncompressed | Sets the compression codec used when spilling data to disk. Since datafusion writes spill files using the Arrow IPC Stream format, only codecs supported by the Arrow IPC Stream Writer are allowed. Valid values are: uncompressed, lz4_frame, zstd. Note: lz4_frame offers faster (de)compression, but typically results in larger spill files. In contrast, zstd achieves higher compression ratios at the cost of slower (de)compression speed. | | datafusion.execution.sort_spill_reservation_bytes | 10485760 | Specifies the reserved memory for each spillable sort operation to facilitate an in-memory merge. When a sort operation spills to disk, the in-memory data must be sorted and merged before being written to a file. This setting reserves a specific amount of memory for that in-memory sort/merge process. Note: This setting is irrelevant if the sort operation cannot spill (i.e., if there's no `DiskManager` configured). | | datafusion.execution.sort_in_place_threshold_bytes | 1048576 | When sorting, below what size should data be concatenated and sorted in a single RecordBatch rather than sorted in batches and merged. | +| datafusion.execution.sort_background_merge_batches | 8 | When greater than 0, a sort sorts every input batch in a background task as soon as it is buffered, and merges every this many sorted batches into one sorted run in a background task, so sorting overlaps reading the input and little is left to merge once it ends. Each merge only reads that many batches, so keep it small enough for their sort keys and rows to fit in the CPU cache. 0 sorts the buffered batches once the input ends or memory runs out. | | datafusion.execution.sort_pushdown_buffer_capacity | 1073741824 | Maximum buffer capacity (in bytes) per partition for BufferExec inserted during sort pushdown optimization. When PushdownSort eliminates a SortExec under SortPreservingMergeExec, a BufferExec is inserted to replace SortExec's buffering role. This prevents I/O stalls by allowing the scan to run ahead of the merge. This uses strictly less memory than the SortExec it replaces (which buffers the entire partition). The buffer respects the global memory pool limit. Setting this to a large value is safe — actual memory usage is bounded by partition size and global memory limits. | | datafusion.execution.max_spill_file_size_bytes | 134217728 | Maximum size in bytes for individual spill files before rotating to a new file. When operators spill data to disk (e.g., RepartitionExec), they write multiple batches to the same file until this size limit is reached, then rotate to a new file. This reduces syscall overhead compared to one-file-per-batch while preventing files from growing too large. A larger value reduces file creation overhead but may hold more disk space. A smaller value creates more files but allows finer-grained space reclamation as files can be deleted once fully consumed. Now only `RepartitionExec` supports this spill file rotation feature, other spilling operators may create spill files larger than the limit. Default: 128 MB | | datafusion.execution.enable_nlj_coordinated_fallback | true | Enables the memory-limited fallback for `NestedLoopJoinExec` join types that emit unmatched left rows in the final output (LEFT, LEFT SEMI, LEFT ANTI, LEFT MARK, FULL) when the right side has multiple partitions. This fallback shares left-side state (the current left chunk, the visited bitmap and the probe-thread counter) across all right-side partitions, which assumes every partition runs in the same process. Distributed engines that execute each output partition as an independent task (e.g. Ballista, datafusion-distributed) give each task its own copy of this state and poll only one partition, so the cross-partition counter never reaches zero and the fallback would stall. Such engines should set this to `false`: the coordinated fallback is then disabled for left-emitting multi-partition joins, which instead fail with a resource-exhaustion error under memory pressure rather than deadlocking. Single-partition and non-left-emitting joins are unaffected and always keep the fallback. |