From 9c179b869a1fd3a3fde509048ba1ff3de79a5aaf Mon Sep 17 00:00:00 2001 From: yashrb24 Date: Wed, 16 Sep 2026 08:58:00 +0530 Subject: [PATCH 1/7] feat: stage ORDER BY expression evaluation --- datafusion/common/src/config.rs | 12 + datafusion/physical-plan/src/sorts/sort.rs | 206 +++++++++++++++++- datafusion/physical-plan/src/sorts/stream.rs | 12 +- .../test_files/information_schema.slt | 4 + docs/source/user-guide/configs.md | 2 + 5 files changed, 227 insertions(+), 9 deletions(-) diff --git a/datafusion/common/src/config.rs b/datafusion/common/src/config.rs index 360586b0e9bae..6360648ebdf08 100644 --- a/datafusion/common/src/config.rs +++ b/datafusion/common/src/config.rs @@ -1001,6 +1001,18 @@ config_namespace! { /// batches and merged. pub sort_in_place_threshold_bytes: usize, default = 1024 * 1024 + /// Enables staged evaluation for eligible final, single-run, in-memory sorts. + /// Later sort keys are evaluated only for rows still tied on earlier keys. + /// Requires a key after the first group that is neither a column nor a literal. + /// Spill preparation, multi-run merges, and TopK sorting continue to use eager + /// key evaluation. This experimental option can suppress expression errors + /// or change volatile expression results by skipping evaluations. + pub enable_staged_sort: bool, default = false + + /// Number of sort keys evaluated together by staged in-memory sorting. + /// Must be greater than zero when `enable_staged_sort` is true. + pub sort_key_group_size: usize, default = 3 + /// Maximum buffer capacity (in bytes) per partition for BufferExec /// inserted during sort pushdown optimization. /// diff --git a/datafusion/physical-plan/src/sorts/sort.rs b/datafusion/physical-plan/src/sorts/sort.rs index 9a149cc7c38c1..a39d681a71056 100644 --- a/datafusion/physical-plan/src/sorts/sort.rs +++ b/datafusion/physical-plan/src/sorts/sort.rs @@ -19,8 +19,10 @@ //! It will do in-memory sorting if it has enough memory budget //! but spills to disk if needed. +use std::cmp::Ordering as CmpOrdering; use std::fmt; use std::fmt::{Debug, Formatter}; +use std::ops::Range; use std::sync::Arc; use parking_lot::RwLock; @@ -57,9 +59,10 @@ use crate::{ PlanProperties, ReplaceChildrenOptions, SendableRecordBatchStream, Statistics, }; -use arrow::array::{RecordBatch, RecordBatchOptions}; -use arrow::compute::{concat_batches, lexsort_to_indices, take_arrays}; +use arrow::array::{Array, RecordBatch, RecordBatchOptions, UInt32Array}; +use arrow::compute::{SortColumn, concat_batches, lexsort_to_indices, take_arrays}; use arrow::datatypes::SchemaRef; +use arrow_ord::sort::LexicographicalComparator; use datafusion_common::config::SpillCompression; use datafusion_common::tree_node::TreeNodeRecursion; use datafusion_common::{ @@ -73,7 +76,9 @@ use datafusion_execution::memory_pool::{ 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::expressions::{ + Column, DynamicFilterPhysicalExpr, Literal, lit, +}; use futures::{StreamExt, TryStreamExt}; use log::{debug, trace}; @@ -229,6 +234,9 @@ struct ExternalSorter { /// the data will be concatenated and sorted in place rather than /// sort/merged. sort_in_place_threshold_bytes: usize, + /// Staged evaluation is restricted to final, single-run in-memory sorts. + enable_staged_sort: bool, + sort_key_group_size: usize, // ======================================================================== // STATE BUFFERS: @@ -325,9 +333,22 @@ impl ExternalSorter { batch_size, sort_spill_reservation_bytes, sort_in_place_threshold_bytes, + enable_staged_sort: false, + sort_key_group_size: 3, }) } + fn with_staged_sort(mut self, enabled: bool, group_size: usize) -> Result { + if enabled && group_size == 0 { + return Err(DataFusionError::Configuration( + "datafusion.execution.sort_key_group_size must be greater than zero when datafusion.execution.enable_staged_sort is true".to_string(), + )); + } + self.enable_staged_sort = enabled; + self.sort_key_group_size = group_size; + Ok(self) + } + /// Appends an unsorted [`RecordBatch`] to `in_mem_batches` /// /// Updates memory usage metrics, and possibly triggers spilling to disk @@ -657,7 +678,11 @@ 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)?; + let sorted_stream = self.sort_batch_stream( + batch, + reservation, + is_output_stream && self.use_staged_sort(), + )?; return Ok(self.observe_if_output(sorted_stream, is_output_stream)); } @@ -674,7 +699,11 @@ 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)?; + let sorted_stream = self.sort_batch_stream( + batch, + reservation, + is_output_stream && self.use_staged_sort(), + )?; return Ok(self.observe_if_output(sorted_stream, is_output_stream)); } @@ -696,7 +725,7 @@ impl ExternalSorter { let reservation = self .reservation .split(get_reserved_bytes_for_record_batch(&batch)?); - let input = self.sort_batch_stream(batch, reservation)?; + let input = self.sort_batch_stream(batch, reservation, false)?; Ok(spawn_buffered(input, 1)) }) .collect::>()?; @@ -781,6 +810,7 @@ impl ExternalSorter { &self, batch: RecordBatch, reservation: MemoryReservation, + staged: bool, ) -> Result { assert_eq!( get_reserved_bytes_for_record_batch(&batch)?, @@ -790,13 +820,23 @@ impl ExternalSorter { let schema = batch.schema(); let expressions = self.expr.clone(); let batch_size = self.batch_size; + let sort_key_group_size = self.sort_key_group_size; let merge_pool = Arc::clone(&self.merge_pool); let stream = futures::stream::once(async move { let schema = batch.schema(); // Sort the batch immediately and get all output batches - let sorted_batches = sort_batch_chunked(&batch, &expressions, batch_size)?; + let sorted_batches = if staged { + staged_sort_batch_chunked( + &batch, + &expressions, + batch_size, + sort_key_group_size, + )? + } else { + sort_batch_chunked(&batch, &expressions, batch_size)? + }; // Chunked output can retain shared buffers in every batch and // exceed the input estimate. Borrow only already-reserved spill @@ -841,6 +881,17 @@ impl ExternalSorter { Ok(Box::pin(RecordBatchStreamAdapter::new(schema, stream))) } + fn use_staged_sort(&self) -> bool { + // The first group is always evaluated for every row. Only computation + // in later groups can be avoided. + self.enable_staged_sort + && self + .expr + .iter() + .skip(self.sort_key_group_size) + .any(|key| !key.expr.is::() && !key.expr.is::()) + } + /// If this sort may spill, pre-allocates /// `sort_spill_reservation_bytes` of memory to guarantee memory /// left for the in memory sort/merge. @@ -957,6 +1008,143 @@ impl Debug for ExternalSorter { } } +fn staged_sort_batch_chunked( + batch: &RecordBatch, + expressions: &LexOrdering, + batch_size: usize, + group_size: usize, +) -> Result> { + let indices = staged_sort_indices(batch, expressions, group_size)?; + IncrementalSortIterator::new(batch.clone(), expressions.clone(), batch_size) + .with_sorted_indices(indices) + .collect() +} + +/// Compute a complete permutation by evaluating successive key groups only for ties. +fn staged_sort_indices( + batch: &RecordBatch, + expressions: &LexOrdering, + group_size: usize, +) -> Result { + let mut indices = Vec::new(); + // Ranges address positions in `indices`, whose values address the input batch. + let mut unresolved: Vec<_> = std::iter::once(0..batch.num_rows()).collect(); + let stage_count = expressions.len().div_ceil(group_size); + for (stage, keys) in expressions.chunks(group_size).enumerate() { + if unresolved.is_empty() { + break; + } + unresolved = refine_tied_groups( + batch, + keys, + &mut indices, + &unresolved, + stage == 0, + stage + 1 < stage_count, + )?; + } + + Ok(UInt32Array::from(indices)) +} + +/// Evaluate keys once for all unresolved rows, then sort each tied group independently. +fn refine_tied_groups( + batch: &RecordBatch, + keys: &[PhysicalSortExpr], + indices: &mut Vec, + unresolved: &[Range], + first_stage: bool, + find_ties: bool, +) -> Result>> { + let selected = if first_stage { + batch.clone() + } else { + // Gather the original row indices for all tied groups into one buffer. + let mut selection = Vec::with_capacity(unresolved.iter().map(Range::len).sum()); + for range in unresolved { + selection.extend_from_slice(&indices[range.clone()]); + } + arrow::compute::take_record_batch(batch, &UInt32Array::from(selection))? + }; + let sort_columns = keys + .iter() + .map(|expr| expr.evaluate_to_sort_column(&selected)) + .collect::>>()?; + + // Key arrays are already evaluated above. This comparator finds ties that + // need the next stage; the final stage does not need it. + let comparator = find_ties + .then(|| LexicographicalComparator::try_new(&sort_columns)) + .transpose()?; + + let mut next_unresolved = Vec::new(); + if first_stage { + let order = lexsort_to_indices(&sort_columns, None)?; + if let Some(comparator) = &comparator { + append_tied_ranges(comparator, &order, 0, 0, &mut next_unresolved); + } + let (_, values, _) = order.into_parts(); + *indices = values.into(); + return Ok(next_unresolved); + } + + let mut selected_offset = 0; + for range in unresolved { + let columns = sort_columns + .iter() + .map(|column| SortColumn { + values: column.values.slice(selected_offset, range.len()), + options: column.options, + }) + .collect::>(); + let local_order = lexsort_to_indices(&columns, None)?; + let previous = indices[range.clone()].to_vec(); + for (output, local) in indices[range.clone()].iter_mut().zip(local_order.values()) + { + *output = previous[*local as usize]; + } + if let Some(comparator) = &comparator { + append_tied_ranges( + comparator, + &local_order, + selected_offset, + range.start, + &mut next_unresolved, + ); + } + selected_offset += range.len(); + } + Ok(next_unresolved) +} + +/// Find consecutive equal keys in a sorted group; only these rows need later keys. +/// For sorted keys `[a, a, b]` at output offset 5, append the range `5..7`. +/// `selected_offset` locates this group's key values in the gathered batch; +/// `output_offset` locates its positions in the complete output permutation. +fn append_tied_ranges( + comparator: &LexicographicalComparator, + order: &UInt32Array, + selected_offset: usize, + output_offset: usize, + ties: &mut Vec>, +) { + let mut start = 0; + while start < order.len() { + let first = selected_offset + order.value(start) as usize; + let mut end = start + 1; + while end < order.len() + && comparator.compare(first, selected_offset + order.value(end) as usize) + == CmpOrdering::Equal + { + end += 1; + } + if end - start > 1 { + ties.push((output_offset + start)..(output_offset + end)); + } + start = end; + } +} + pub fn sort_batch( batch: &RecordBatch, expressions: &LexOrdering, @@ -1517,6 +1705,10 @@ impl ExecutionPlan for SortExec { context.session_config().spill_compression(), &self.metrics_set, context.runtime_env(), + )? + .with_staged_sort( + execution_options.enable_staged_sort, + execution_options.sort_key_group_size, )?; Ok(Box::pin(RecordBatchStreamAdapter::new( self.schema(), diff --git a/datafusion/physical-plan/src/sorts/stream.rs b/datafusion/physical-plan/src/sorts/stream.rs index e2a840fdedfd5..5d9169a9b7c8a 100644 --- a/datafusion/physical-plan/src/sorts/stream.rs +++ b/datafusion/physical-plan/src/sorts/stream.rs @@ -288,8 +288,9 @@ impl PartitionedStream for FieldCursorStream { /// spill when there is not enough memory for an eager sort (which requires ~2x /// peak memory to hold both the unsorted and sorted copies simultaneously). /// -/// On the first call to `next()`, a sorted index array (`UInt32Array`) is -/// computed via `lexsort_to_indices`. Subsequent calls yield chunks of +/// Unless supplied via [`Self::with_sorted_indices`], a sorted index array +/// (`UInt32Array`) is computed via `lexsort_to_indices` on the first call to +/// `next()`. The iterator yields chunks of /// `batch_size` rows by `take`-ing from the original batch using slices of /// this index array. Each `take` copies data for the chunk (not zero-copy), /// but only one chunk is live at a time since the caller consumes it before @@ -320,6 +321,13 @@ impl IncrementalSortIterator { indices: None, } } + + /// Use a precomputed sort permutation instead of evaluating the expressions. + /// Must be called before iteration; indices must contain every input row once. + pub(crate) fn with_sorted_indices(mut self, indices: UInt32Array) -> Self { + self.indices = Some(indices); + self + } } impl Iterator for IncrementalSortIterator { diff --git a/datafusion/sqllogictest/test_files/information_schema.slt b/datafusion/sqllogictest/test_files/information_schema.slt index b270eba99d7b0..3d7039045ef0c 100644 --- a/datafusion/sqllogictest/test_files/information_schema.slt +++ b/datafusion/sqllogictest/test_files/information_schema.slt @@ -221,6 +221,7 @@ datafusion.execution.enable_file_stream_work_stealing true datafusion.execution.enable_migration_aggregate true datafusion.execution.enable_nlj_coordinated_fallback true datafusion.execution.enable_recursive_ctes true +datafusion.execution.enable_staged_sort false datafusion.execution.enforce_batch_size_in_joins false datafusion.execution.hash_join_buffering_capacity 0 datafusion.execution.keep_partition_by_columns false @@ -278,6 +279,7 @@ datafusion.execution.skip_partial_aggregation_probe_rows_threshold 100000 datafusion.execution.skip_physical_aggregate_schema_check false datafusion.execution.soft_max_rows_per_output_file 50000000 datafusion.execution.sort_in_place_threshold_bytes 1048576 +datafusion.execution.sort_key_group_size 3 datafusion.execution.sort_pushdown_buffer_capacity 1073741824 datafusion.execution.sort_spill_reservation_bytes 10485760 datafusion.execution.spill_compression uncompressed @@ -382,6 +384,7 @@ datafusion.execution.enable_file_stream_work_stealing true When `true` (the defa datafusion.execution.enable_migration_aggregate true Whether aggregation uses the implementation from the major refactor completed in the 56.0.0 release. When set to `false`, aggregation falls back to the implementation used before 55.0.0. The fallback exists only as a workaround for bugs in the new implementation and will be removed, together with this option, after the 56.0.0 release. See for details. 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 coordinates per-chunk left state (visited bitmap and 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) build a separate coordinator per task 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. datafusion.execution.enable_recursive_ctes true Should DataFusion support recursive CTEs +datafusion.execution.enable_staged_sort false Enables staged evaluation for eligible final, single-run, in-memory sorts. Later sort keys are evaluated only for rows still tied on earlier keys. Requires a key after the first group that is neither a column nor a literal. Spill preparation, multi-run merges, and TopK sorting continue to use eager key evaluation. This experimental option can suppress expression errors or change volatile expression results by skipping evaluations. datafusion.execution.enforce_batch_size_in_joins false Should DataFusion enforce batch size in joins or not. By default, DataFusion will not enforce batch size in joins. Enforcing batch size in joins can reduce memory usage when joining large tables with a highly-selective join filter, but is also slightly slower. Note: this option currently only applies to the symmetric hash join. datafusion.execution.hash_join_buffering_capacity 0 How many bytes to buffer in the probe side of hash joins while the build side is concurrently being built. Without this, hash joins will wait until the full materialization of the build side before polling the probe side. This is useful in scenarios where the query is not completely CPU bounded, allowing to do some early work concurrently and reducing the latency of the query. Note that when hash join buffering is enabled, the probe side will start eagerly polling data, not giving time for the producer side of dynamic filters to produce any meaningful predicate. Queries with dynamic filters might see performance degradation. Disabled by default, set to a number greater than 0 for enabling it. datafusion.execution.keep_partition_by_columns false Should DataFusion keep the columns used for partition_by in the output RecordBatches @@ -439,6 +442,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_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_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_key_group_size 3 Number of sort keys evaluated together by staged in-memory sorting. Must be greater than zero when `enable_staged_sort` is true. 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). 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. diff --git a/docs/source/user-guide/configs.md b/docs/source/user-guide/configs.md index 0085d4ac7c1fa..12bea7aab6c56 100644 --- a/docs/source/user-guide/configs.md +++ b/docs/source/user-guide/configs.md @@ -125,6 +125,8 @@ 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.enable_staged_sort | false | Enables staged evaluation for eligible final, single-run, in-memory sorts. Later sort keys are evaluated only for rows still tied on earlier keys. Requires a key after the first group that is neither a column nor a literal. Spill preparation, multi-run merges, and TopK sorting continue to use eager key evaluation. This experimental option can suppress expression errors or change volatile expression results by skipping evaluations. | +| datafusion.execution.sort_key_group_size | 3 | Number of sort keys evaluated together by staged in-memory sorting. Must be greater than zero when `enable_staged_sort` is true. | | 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 coordinates per-chunk left state (visited bitmap and 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) build a separate coordinator per task 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. | From ca766eb669920513b8d2c0f29f765e865a43be29 Mon Sep 17 00:00:00 2001 From: yashrb24 Date: Wed, 16 Sep 2026 09:06:26 +0530 Subject: [PATCH 2/7] chore: remove staged-sort config docs and snapshot changes --- datafusion/sqllogictest/test_files/information_schema.slt | 4 ---- docs/source/user-guide/configs.md | 2 -- 2 files changed, 6 deletions(-) diff --git a/datafusion/sqllogictest/test_files/information_schema.slt b/datafusion/sqllogictest/test_files/information_schema.slt index 3d7039045ef0c..b270eba99d7b0 100644 --- a/datafusion/sqllogictest/test_files/information_schema.slt +++ b/datafusion/sqllogictest/test_files/information_schema.slt @@ -221,7 +221,6 @@ datafusion.execution.enable_file_stream_work_stealing true datafusion.execution.enable_migration_aggregate true datafusion.execution.enable_nlj_coordinated_fallback true datafusion.execution.enable_recursive_ctes true -datafusion.execution.enable_staged_sort false datafusion.execution.enforce_batch_size_in_joins false datafusion.execution.hash_join_buffering_capacity 0 datafusion.execution.keep_partition_by_columns false @@ -279,7 +278,6 @@ datafusion.execution.skip_partial_aggregation_probe_rows_threshold 100000 datafusion.execution.skip_physical_aggregate_schema_check false datafusion.execution.soft_max_rows_per_output_file 50000000 datafusion.execution.sort_in_place_threshold_bytes 1048576 -datafusion.execution.sort_key_group_size 3 datafusion.execution.sort_pushdown_buffer_capacity 1073741824 datafusion.execution.sort_spill_reservation_bytes 10485760 datafusion.execution.spill_compression uncompressed @@ -384,7 +382,6 @@ datafusion.execution.enable_file_stream_work_stealing true When `true` (the defa datafusion.execution.enable_migration_aggregate true Whether aggregation uses the implementation from the major refactor completed in the 56.0.0 release. When set to `false`, aggregation falls back to the implementation used before 55.0.0. The fallback exists only as a workaround for bugs in the new implementation and will be removed, together with this option, after the 56.0.0 release. See for details. 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 coordinates per-chunk left state (visited bitmap and 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) build a separate coordinator per task 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. datafusion.execution.enable_recursive_ctes true Should DataFusion support recursive CTEs -datafusion.execution.enable_staged_sort false Enables staged evaluation for eligible final, single-run, in-memory sorts. Later sort keys are evaluated only for rows still tied on earlier keys. Requires a key after the first group that is neither a column nor a literal. Spill preparation, multi-run merges, and TopK sorting continue to use eager key evaluation. This experimental option can suppress expression errors or change volatile expression results by skipping evaluations. datafusion.execution.enforce_batch_size_in_joins false Should DataFusion enforce batch size in joins or not. By default, DataFusion will not enforce batch size in joins. Enforcing batch size in joins can reduce memory usage when joining large tables with a highly-selective join filter, but is also slightly slower. Note: this option currently only applies to the symmetric hash join. datafusion.execution.hash_join_buffering_capacity 0 How many bytes to buffer in the probe side of hash joins while the build side is concurrently being built. Without this, hash joins will wait until the full materialization of the build side before polling the probe side. This is useful in scenarios where the query is not completely CPU bounded, allowing to do some early work concurrently and reducing the latency of the query. Note that when hash join buffering is enabled, the probe side will start eagerly polling data, not giving time for the producer side of dynamic filters to produce any meaningful predicate. Queries with dynamic filters might see performance degradation. Disabled by default, set to a number greater than 0 for enabling it. datafusion.execution.keep_partition_by_columns false Should DataFusion keep the columns used for partition_by in the output RecordBatches @@ -442,7 +439,6 @@ 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_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_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_key_group_size 3 Number of sort keys evaluated together by staged in-memory sorting. Must be greater than zero when `enable_staged_sort` is true. 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). 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. diff --git a/docs/source/user-guide/configs.md b/docs/source/user-guide/configs.md index 12bea7aab6c56..0085d4ac7c1fa 100644 --- a/docs/source/user-guide/configs.md +++ b/docs/source/user-guide/configs.md @@ -125,8 +125,6 @@ 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.enable_staged_sort | false | Enables staged evaluation for eligible final, single-run, in-memory sorts. Later sort keys are evaluated only for rows still tied on earlier keys. Requires a key after the first group that is neither a column nor a literal. Spill preparation, multi-run merges, and TopK sorting continue to use eager key evaluation. This experimental option can suppress expression errors or change volatile expression results by skipping evaluations. | -| datafusion.execution.sort_key_group_size | 3 | Number of sort keys evaluated together by staged in-memory sorting. Must be greater than zero when `enable_staged_sort` is true. | | 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 coordinates per-chunk left state (visited bitmap and 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) build a separate coordinator per task 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. | From f9ca9facd201433395d24b8956626dd732c65907 Mon Sep 17 00:00:00 2001 From: yashrb24 Date: Wed, 16 Sep 2026 09:17:14 +0530 Subject: [PATCH 3/7] chore: simplify staged-sort comments --- datafusion/common/src/config.rs | 9 ++------- datafusion/physical-plan/src/sorts/sort.rs | 5 +---- datafusion/physical-plan/src/sorts/stream.rs | 2 -- 3 files changed, 3 insertions(+), 13 deletions(-) diff --git a/datafusion/common/src/config.rs b/datafusion/common/src/config.rs index 6360648ebdf08..663cdc6dba7f9 100644 --- a/datafusion/common/src/config.rs +++ b/datafusion/common/src/config.rs @@ -1001,16 +1001,11 @@ config_namespace! { /// batches and merged. pub sort_in_place_threshold_bytes: usize, default = 1024 * 1024 - /// Enables staged evaluation for eligible final, single-run, in-memory sorts. - /// Later sort keys are evaluated only for rows still tied on earlier keys. - /// Requires a key after the first group that is neither a column nor a literal. - /// Spill preparation, multi-run merges, and TopK sorting continue to use eager - /// key evaluation. This experimental option can suppress expression errors - /// or change volatile expression results by skipping evaluations. + /// Enables staged in-memory sorting, evaluating later sort keys only for rows + /// tied on earlier keys. pub enable_staged_sort: bool, default = false /// Number of sort keys evaluated together by staged in-memory sorting. - /// Must be greater than zero when `enable_staged_sort` is true. pub sort_key_group_size: usize, default = 3 /// Maximum buffer capacity (in bytes) per partition for BufferExec diff --git a/datafusion/physical-plan/src/sorts/sort.rs b/datafusion/physical-plan/src/sorts/sort.rs index a39d681a71056..8f268f55cdf69 100644 --- a/datafusion/physical-plan/src/sorts/sort.rs +++ b/datafusion/physical-plan/src/sorts/sort.rs @@ -234,7 +234,6 @@ struct ExternalSorter { /// the data will be concatenated and sorted in place rather than /// sort/merged. sort_in_place_threshold_bytes: usize, - /// Staged evaluation is restricted to final, single-run in-memory sorts. enable_staged_sort: bool, sort_key_group_size: usize, @@ -1072,7 +1071,7 @@ fn refine_tied_groups( .collect::>>()?; // Key arrays are already evaluated above. This comparator finds ties that - // need the next stage; the final stage does not need it. + // need the next stage. let comparator = find_ties .then(|| LexicographicalComparator::try_new(&sort_columns)) .transpose()?; @@ -1119,8 +1118,6 @@ fn refine_tied_groups( /// Find consecutive equal keys in a sorted group; only these rows need later keys. /// For sorted keys `[a, a, b]` at output offset 5, append the range `5..7`. -/// `selected_offset` locates this group's key values in the gathered batch; -/// `output_offset` locates its positions in the complete output permutation. fn append_tied_ranges( comparator: &LexicographicalComparator, order: &UInt32Array, diff --git a/datafusion/physical-plan/src/sorts/stream.rs b/datafusion/physical-plan/src/sorts/stream.rs index 5d9169a9b7c8a..a9533e0e801b0 100644 --- a/datafusion/physical-plan/src/sorts/stream.rs +++ b/datafusion/physical-plan/src/sorts/stream.rs @@ -322,8 +322,6 @@ impl IncrementalSortIterator { } } - /// Use a precomputed sort permutation instead of evaluating the expressions. - /// Must be called before iteration; indices must contain every input row once. pub(crate) fn with_sorted_indices(mut self, indices: UInt32Array) -> Self { self.indices = Some(indices); self From 8bf323a28a5bb91b50f6e864d2531335a2434bb4 Mon Sep 17 00:00:00 2001 From: yashrb24 Date: Wed, 16 Sep 2026 10:53:13 +0530 Subject: [PATCH 4/7] feat: stage sort keys for in-memory runs --- datafusion/physical-plan/src/sorts/sort.rs | 16 +++++----------- 1 file changed, 5 insertions(+), 11 deletions(-) diff --git a/datafusion/physical-plan/src/sorts/sort.rs b/datafusion/physical-plan/src/sorts/sort.rs index 8f268f55cdf69..87dada6ea8781 100644 --- a/datafusion/physical-plan/src/sorts/sort.rs +++ b/datafusion/physical-plan/src/sorts/sort.rs @@ -668,6 +668,8 @@ impl ExternalSorter { let elapsed_compute = self.metrics.baseline.elapsed_compute().clone(); let _timer = elapsed_compute.timer(); + let staged = is_output_stream && self.use_staged_sort(); + // Please pay attention that any operation inside of `in_mem_sort_stream` will // not perform any memory reservation. This is for avoiding the need of handling // reservation failure and spilling in the middle of the sort/merge. The memory @@ -677,11 +679,7 @@ 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, - is_output_stream && self.use_staged_sort(), - )?; + let sorted_stream = self.sort_batch_stream(batch, reservation, staged)?; return Ok(self.observe_if_output(sorted_stream, is_output_stream)); } @@ -698,11 +696,7 @@ 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, - is_output_stream && self.use_staged_sort(), - )?; + let sorted_stream = self.sort_batch_stream(batch, reservation, staged)?; return Ok(self.observe_if_output(sorted_stream, is_output_stream)); } @@ -724,7 +718,7 @@ impl ExternalSorter { let reservation = self .reservation .split(get_reserved_bytes_for_record_batch(&batch)?); - let input = self.sort_batch_stream(batch, reservation, false)?; + let input = self.sort_batch_stream(batch, reservation, staged)?; Ok(spawn_buffered(input, 1)) }) .collect::>()?; From 5b9ffdbb8bc46e6dbb6c858fa4254a43d702ce2e Mon Sep 17 00:00:00 2001 From: yashrb24 Date: Wed, 16 Sep 2026 11:05:52 +0530 Subject: [PATCH 5/7] fix: synchronize staged-sort configuration entries --- datafusion/sqllogictest/test_files/information_schema.slt | 4 ++++ docs/source/user-guide/configs.md | 2 ++ 2 files changed, 6 insertions(+) diff --git a/datafusion/sqllogictest/test_files/information_schema.slt b/datafusion/sqllogictest/test_files/information_schema.slt index b270eba99d7b0..825b6738db79f 100644 --- a/datafusion/sqllogictest/test_files/information_schema.slt +++ b/datafusion/sqllogictest/test_files/information_schema.slt @@ -221,6 +221,7 @@ datafusion.execution.enable_file_stream_work_stealing true datafusion.execution.enable_migration_aggregate true datafusion.execution.enable_nlj_coordinated_fallback true datafusion.execution.enable_recursive_ctes true +datafusion.execution.enable_staged_sort false datafusion.execution.enforce_batch_size_in_joins false datafusion.execution.hash_join_buffering_capacity 0 datafusion.execution.keep_partition_by_columns false @@ -278,6 +279,7 @@ datafusion.execution.skip_partial_aggregation_probe_rows_threshold 100000 datafusion.execution.skip_physical_aggregate_schema_check false datafusion.execution.soft_max_rows_per_output_file 50000000 datafusion.execution.sort_in_place_threshold_bytes 1048576 +datafusion.execution.sort_key_group_size 3 datafusion.execution.sort_pushdown_buffer_capacity 1073741824 datafusion.execution.sort_spill_reservation_bytes 10485760 datafusion.execution.spill_compression uncompressed @@ -382,6 +384,7 @@ datafusion.execution.enable_file_stream_work_stealing true When `true` (the defa datafusion.execution.enable_migration_aggregate true Whether aggregation uses the implementation from the major refactor completed in the 56.0.0 release. When set to `false`, aggregation falls back to the implementation used before 55.0.0. The fallback exists only as a workaround for bugs in the new implementation and will be removed, together with this option, after the 56.0.0 release. See for details. 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 coordinates per-chunk left state (visited bitmap and 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) build a separate coordinator per task 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. datafusion.execution.enable_recursive_ctes true Should DataFusion support recursive CTEs +datafusion.execution.enable_staged_sort false Enables staged in-memory sorting, evaluating later sort keys only for rows tied on earlier keys. datafusion.execution.enforce_batch_size_in_joins false Should DataFusion enforce batch size in joins or not. By default, DataFusion will not enforce batch size in joins. Enforcing batch size in joins can reduce memory usage when joining large tables with a highly-selective join filter, but is also slightly slower. Note: this option currently only applies to the symmetric hash join. datafusion.execution.hash_join_buffering_capacity 0 How many bytes to buffer in the probe side of hash joins while the build side is concurrently being built. Without this, hash joins will wait until the full materialization of the build side before polling the probe side. This is useful in scenarios where the query is not completely CPU bounded, allowing to do some early work concurrently and reducing the latency of the query. Note that when hash join buffering is enabled, the probe side will start eagerly polling data, not giving time for the producer side of dynamic filters to produce any meaningful predicate. Queries with dynamic filters might see performance degradation. Disabled by default, set to a number greater than 0 for enabling it. datafusion.execution.keep_partition_by_columns false Should DataFusion keep the columns used for partition_by in the output RecordBatches @@ -439,6 +442,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_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_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_key_group_size 3 Number of sort keys evaluated together by staged in-memory sorting. 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). 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. diff --git a/docs/source/user-guide/configs.md b/docs/source/user-guide/configs.md index 0085d4ac7c1fa..fae5c3cc2f084 100644 --- a/docs/source/user-guide/configs.md +++ b/docs/source/user-guide/configs.md @@ -125,6 +125,8 @@ 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.enable_staged_sort | false | Enables staged in-memory sorting, evaluating later sort keys only for rows tied on earlier keys. | +| datafusion.execution.sort_key_group_size | 3 | Number of sort keys evaluated together by staged in-memory sorting. | | 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 coordinates per-chunk left state (visited bitmap and 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) build a separate coordinator per task 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. | From a266039447d65de3d10cc8a8a142206cc9fe7215 Mon Sep 17 00:00:00 2001 From: yashrb24 Date: Sat, 3 Oct 2026 20:04:34 +0530 Subject: [PATCH 6/7] chore: default staged sort key groups to five --- datafusion/common/src/config.rs | 2 +- datafusion/physical-plan/src/sorts/sort.rs | 2 +- datafusion/sqllogictest/test_files/information_schema.slt | 4 ++-- docs/source/user-guide/configs.md | 2 +- 4 files changed, 5 insertions(+), 5 deletions(-) diff --git a/datafusion/common/src/config.rs b/datafusion/common/src/config.rs index 663cdc6dba7f9..37ce34b49c847 100644 --- a/datafusion/common/src/config.rs +++ b/datafusion/common/src/config.rs @@ -1006,7 +1006,7 @@ config_namespace! { pub enable_staged_sort: bool, default = false /// Number of sort keys evaluated together by staged in-memory sorting. - pub sort_key_group_size: usize, default = 3 + pub sort_key_group_size: usize, default = 5 /// Maximum buffer capacity (in bytes) per partition for BufferExec /// inserted during sort pushdown optimization. diff --git a/datafusion/physical-plan/src/sorts/sort.rs b/datafusion/physical-plan/src/sorts/sort.rs index 87dada6ea8781..748db94579d99 100644 --- a/datafusion/physical-plan/src/sorts/sort.rs +++ b/datafusion/physical-plan/src/sorts/sort.rs @@ -333,7 +333,7 @@ impl ExternalSorter { sort_spill_reservation_bytes, sort_in_place_threshold_bytes, enable_staged_sort: false, - sort_key_group_size: 3, + sort_key_group_size: 5, }) } diff --git a/datafusion/sqllogictest/test_files/information_schema.slt b/datafusion/sqllogictest/test_files/information_schema.slt index 825b6738db79f..5bcb39b06d9fa 100644 --- a/datafusion/sqllogictest/test_files/information_schema.slt +++ b/datafusion/sqllogictest/test_files/information_schema.slt @@ -279,7 +279,7 @@ datafusion.execution.skip_partial_aggregation_probe_rows_threshold 100000 datafusion.execution.skip_physical_aggregate_schema_check false datafusion.execution.soft_max_rows_per_output_file 50000000 datafusion.execution.sort_in_place_threshold_bytes 1048576 -datafusion.execution.sort_key_group_size 3 +datafusion.execution.sort_key_group_size 5 datafusion.execution.sort_pushdown_buffer_capacity 1073741824 datafusion.execution.sort_spill_reservation_bytes 10485760 datafusion.execution.spill_compression uncompressed @@ -442,7 +442,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_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_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_key_group_size 3 Number of sort keys evaluated together by staged in-memory sorting. +datafusion.execution.sort_key_group_size 5 Number of sort keys evaluated together by staged in-memory sorting. 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). 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. diff --git a/docs/source/user-guide/configs.md b/docs/source/user-guide/configs.md index fae5c3cc2f084..f068e11c7757e 100644 --- a/docs/source/user-guide/configs.md +++ b/docs/source/user-guide/configs.md @@ -126,7 +126,7 @@ The following configuration settings are available: | 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.enable_staged_sort | false | Enables staged in-memory sorting, evaluating later sort keys only for rows tied on earlier keys. | -| datafusion.execution.sort_key_group_size | 3 | Number of sort keys evaluated together by staged in-memory sorting. | +| datafusion.execution.sort_key_group_size | 5 | Number of sort keys evaluated together by staged in-memory sorting. | | 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 coordinates per-chunk left state (visited bitmap and 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) build a separate coordinator per task 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. | From 094d0a2b6459640bb4aa7631d317896b18a778cd Mon Sep 17 00:00:00 2001 From: yashrb24 Date: Tue, 13 Oct 2026 00:34:46 +0530 Subject: [PATCH 7/7] perf: cache staged sort projections and add focused tests --- datafusion/physical-plan/src/sorts/sort.rs | 386 ++++++++++++++++++--- 1 file changed, 344 insertions(+), 42 deletions(-) diff --git a/datafusion/physical-plan/src/sorts/sort.rs b/datafusion/physical-plan/src/sorts/sort.rs index 748db94579d99..acf9cafeda9ed 100644 --- a/datafusion/physical-plan/src/sorts/sort.rs +++ b/datafusion/physical-plan/src/sorts/sort.rs @@ -79,6 +79,8 @@ use datafusion_physical_expr::PhysicalExpr; use datafusion_physical_expr::expressions::{ Column, DynamicFilterPhysicalExpr, Literal, lit, }; +use datafusion_physical_expr::projection::ProjectionExprs; +use datafusion_physical_expr::utils::collect_columns; use futures::{StreamExt, TryStreamExt}; use log::{debug, trace}; @@ -234,8 +236,7 @@ struct ExternalSorter { /// the data will be concatenated and sorted in place rather than /// sort/merged. sort_in_place_threshold_bytes: usize, - enable_staged_sort: bool, - sort_key_group_size: usize, + sort_stages: Option>, // ======================================================================== // STATE BUFFERS: @@ -332,8 +333,7 @@ impl ExternalSorter { batch_size, sort_spill_reservation_bytes, sort_in_place_threshold_bytes, - enable_staged_sort: false, - sort_key_group_size: 5, + sort_stages: None, }) } @@ -343,8 +343,52 @@ impl ExternalSorter { "datafusion.execution.sort_key_group_size must be greater than zero when datafusion.execution.enable_staged_sort is true".to_string(), )); } - self.enable_staged_sort = enabled; - self.sort_key_group_size = group_size; + self.sort_stages = if enabled + && self + .expr + .iter() + .skip(group_size) + .any(|key| !key.expr.is::() && !key.expr.is::()) + { + let stages = self + .expr + .chunks(group_size) + .enumerate() + .map(|(index, keys)| { + if index == 0 { + return Ok(SortStage { + columns: vec![], + expressions: keys.to_vec(), + }); + } + let mut columns = keys + .iter() + .flat_map(|key| collect_columns(&key.expr)) + .map(|column| column.index()) + .collect::>(); + columns.sort_unstable(); + columns.dedup(); + let projection = + ProjectionExprs::from_indices(&columns, &self.schema); + let expressions = keys + .iter() + .map(|key| { + Ok(PhysicalSortExpr::new( + projection.project_expr(&key.expr)?, + key.options, + )) + }) + .collect::>>()?; + Ok(SortStage { + columns, + expressions, + }) + }) + .collect::>>()?; + Some(stages.into()) + } else { + None + }; Ok(self) } @@ -668,7 +712,7 @@ impl ExternalSorter { let elapsed_compute = self.metrics.baseline.elapsed_compute().clone(); let _timer = elapsed_compute.timer(); - let staged = is_output_stream && self.use_staged_sort(); + let staged = is_output_stream && self.sort_stages.is_some(); // Please pay attention that any operation inside of `in_mem_sort_stream` will // not perform any memory reservation. This is for avoiding the need of handling @@ -813,20 +857,19 @@ impl ExternalSorter { let schema = batch.schema(); let expressions = self.expr.clone(); let batch_size = self.batch_size; - let sort_key_group_size = self.sort_key_group_size; + let stages = if staged { + self.sort_stages.clone() + } else { + None + }; let merge_pool = Arc::clone(&self.merge_pool); let stream = futures::stream::once(async move { let schema = batch.schema(); // Sort the batch immediately and get all output batches - let sorted_batches = if staged { - staged_sort_batch_chunked( - &batch, - &expressions, - batch_size, - sort_key_group_size, - )? + let sorted_batches = if let Some(stages) = stages { + staged_sort_batch_chunked(&batch, &expressions, batch_size, &stages)? } else { sort_batch_chunked(&batch, &expressions, batch_size)? }; @@ -874,17 +917,6 @@ impl ExternalSorter { Ok(Box::pin(RecordBatchStreamAdapter::new(schema, stream))) } - fn use_staged_sort(&self) -> bool { - // The first group is always evaluated for every row. Only computation - // in later groups can be avoided. - self.enable_staged_sort - && self - .expr - .iter() - .skip(self.sort_key_group_size) - .any(|key| !key.expr.is::() && !key.expr.is::()) - } - /// If this sort may spill, pre-allocates /// `sort_spill_reservation_bytes` of memory to guarantee memory /// left for the in memory sort/merge. @@ -1001,39 +1033,39 @@ impl Debug for ExternalSorter { } } +struct SortStage { + columns: Vec, + expressions: Vec, +} + fn staged_sort_batch_chunked( batch: &RecordBatch, expressions: &LexOrdering, batch_size: usize, - group_size: usize, + stages: &[SortStage], ) -> Result> { - let indices = staged_sort_indices(batch, expressions, group_size)?; + let indices = staged_sort_indices(batch, stages)?; IncrementalSortIterator::new(batch.clone(), expressions.clone(), batch_size) .with_sorted_indices(indices) .collect() } /// Compute a complete permutation by evaluating successive key groups only for ties. -fn staged_sort_indices( - batch: &RecordBatch, - expressions: &LexOrdering, - group_size: usize, -) -> Result { +fn staged_sort_indices(batch: &RecordBatch, stages: &[SortStage]) -> Result { let mut indices = Vec::new(); // Ranges address positions in `indices`, whose values address the input batch. let mut unresolved: Vec<_> = std::iter::once(0..batch.num_rows()).collect(); - let stage_count = expressions.len().div_ceil(group_size); - for (stage, keys) in expressions.chunks(group_size).enumerate() { + for (index, stage) in stages.iter().enumerate() { if unresolved.is_empty() { break; } unresolved = refine_tied_groups( batch, - keys, + stage, &mut indices, &unresolved, - stage == 0, - stage + 1 < stage_count, + index == 0, + index + 1 < stages.len(), )?; } @@ -1043,7 +1075,7 @@ fn staged_sort_indices( /// Evaluate keys once for all unresolved rows, then sort each tied group independently. fn refine_tied_groups( batch: &RecordBatch, - keys: &[PhysicalSortExpr], + stage: &SortStage, indices: &mut Vec, unresolved: &[Range], first_stage: bool, @@ -1052,14 +1084,17 @@ fn refine_tied_groups( let selected = if first_stage { batch.clone() } else { - // Gather the original row indices for all tied groups into one buffer. let mut selection = Vec::with_capacity(unresolved.iter().map(Range::len).sum()); for range in unresolved { selection.extend_from_slice(&indices[range.clone()]); } - arrow::compute::take_record_batch(batch, &UInt32Array::from(selection))? + arrow::compute::take_record_batch( + &batch.project(&stage.columns)?, + &UInt32Array::from(selection), + )? }; - let sort_columns = keys + let sort_columns = stage + .expressions .iter() .map(|expr| expr.evaluate_to_sort_column(&selected)) .collect::>>()?; @@ -3785,6 +3820,273 @@ mod tests { Ok((sorted_batches, metrics)) } + // Exercise SortExec directly so the planner cannot project the sort expressions + // into columns before they reach the sorter. Each key is a separate stage. + async fn staged_sort_test_output( + batch: &RecordBatch, + ordering: &LexOrdering, + enabled: bool, + input_batches: usize, + in_place_threshold: usize, + ) -> Result> { + let rows_per_batch = batch.num_rows() / input_batches; + let batches = (0..input_batches) + .map(|i| batch.slice(i * rows_per_batch, rows_per_batch)) + .collect(); + let input = TestMemoryExec::try_new_exec(&[batches], batch.schema(), None)?; + let sort = Arc::new(SortExec::new(ordering.clone(), input)); + let config = SessionConfig::new() + .with_batch_size(2) + .with_sort_in_place_threshold_bytes(in_place_threshold) + .set_bool("datafusion.execution.enable_staged_sort", enabled) + .set_usize("datafusion.execution.sort_key_group_size", 1); + let result = collect( + Arc::clone(&sort) as Arc, + Arc::new(TaskContext::default().with_session_config(config)), + ) + .await; + assert_eq!(sort.metrics().unwrap().spill_count().unwrap_or(0), 0); + if let Ok(batches) = &result { + assert_eq!( + batches.iter().map(RecordBatch::num_rows).sum::(), + batch.num_rows() + ); + assert!(batches.iter().all(|b| (1..=2).contains(&b.num_rows()))); + } + result + } + + #[tokio::test] + async fn test_staged_sort_matches_eager() -> Result<()> { + use datafusion_physical_expr::expressions::CastExpr; + + // Keep payload columns between the keys to exercise projected column remapping. + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new("suffix", DataType::Utf8, true), + Field::new("payload", DataType::Utf8, false), + Field::new("prefix", DataType::Int32, true), + ])); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int32Array::from_iter_values(0..10)), + Arc::new(StringArray::from(vec![ + Some("3"), + Some("2"), + Some("5"), + None, + Some("3"), + Some("10"), + None, + Some("2"), + Some("0"), + Some("10"), + ])), + Arc::new(StringArray::from_iter_values( + (0..10).map(|i| format!("payload-{i}")), + )), + Arc::new(Int32Array::from(vec![ + Some(2), + Some(1), + None, + Some(1), + Some(2), + Some(1), + None, + Some(1), + Some(3), + Some(1), + ])), + ], + )?; + let ordering = [ + PhysicalSortExpr::new( + col("prefix", &schema)?, + SortOptions { + descending: false, + nulls_first: false, + }, + ), + PhysicalSortExpr::new( + Arc::new(CastExpr::new( + col("suffix", &schema)?, + DataType::Int32, + None, + )), + SortOptions { + descending: true, + nulls_first: true, + }, + ), + PhysicalSortExpr::new_default(col("id", &schema)?), + ] + .into(); + // Prefix ties survive the numeric suffix stage and must be ordered by id. + let expected = arrow::compute::take_record_batch( + &batch, + &UInt32Array::from(vec![3, 5, 9, 1, 7, 0, 4, 8, 6, 2]), + )?; + + // Single batch, concatenated batches, and separate runs followed by a merge. + for (input_batches, threshold) in [(1, 1), (2, usize::MAX), (2, 1)] { + for enabled in [false, true] { + let output = staged_sort_test_output( + &batch, + &ordering, + enabled, + input_batches, + threshold, + ) + .await?; + assert_eq!( + concat_batches(&schema, &output)?, + expected, + "enabled={enabled}, input_batches={input_batches}, threshold={threshold}" + ); + } + } + Ok(()) + } + + #[tokio::test] + async fn test_staged_sort_expression_errors() -> Result<()> { + use datafusion_physical_expr::expressions::CastExpr; + + let schema = Arc::new(Schema::new(vec![ + Field::new("prefix", DataType::Int32, false), + Field::new("invalid", DataType::Utf8, false), + ])); + let ordering = [ + PhysicalSortExpr::new_default(col("prefix", &schema)?), + PhysicalSortExpr::new_default(Arc::new(CastExpr::new( + col("invalid", &schema)?, + DataType::Int32, + None, + ))), + ] + .into(); + for all_tied in [false, true] { + let prefix = if all_tied { + vec![1; 6] + } else { + vec![3, 1, 6, 2, 5, 4] + }; + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int32Array::from(prefix)), + Arc::new(StringArray::from(vec!["invalid"; 6])), + ], + )?; + for (input_batches, threshold) in [(1, 1), (2, usize::MAX), (2, 1)] { + for enabled in [false, true] { + let result = staged_sort_test_output( + &batch, + &ordering, + enabled, + input_batches, + threshold, + ) + .await; + // A merge evaluates all keys again, even if each run skipped the cast. + let merge = input_batches == 2 && threshold == 1; + if !enabled || all_tied || merge { + let error = result.unwrap_err().strip_backtrace(); + assert!( + error.contains("Cannot cast string 'invalid'"), + "{error}" + ); + } else { + let output = concat_batches(&schema, &result?)?; + assert_eq!( + output.column(0).as_ref(), + &Int32Array::from_iter_values(1..=6), + ); + } + } + } + } + Ok(()) + } + + #[tokio::test] + async fn test_staged_sort_volatile_evaluation() -> Result<()> { + use datafusion_expr::{Volatility, create_udf}; + use datafusion_physical_expr::ScalarFunctionExpr; + + let schema = Arc::new(Schema::new(vec![ + Field::new("prefix", DataType::Int32, false), + Field::new("id", DataType::Int32, false), + ])); + // Count rows, not calls: expression evaluation is batched. The middle + // case has four tied rows, with two in each of the split input batches. + for (prefix, tied_rows) in [ + (vec![3, 1, 6, 2, 5, 4], 0), + (vec![2, 1, 1, 3, 4, 4], 4), + (vec![1; 6], 6), + ] { + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int32Array::from(prefix)), + Arc::new(Int32Array::from_iter_values(0..6)), + ], + )?; + for (input_batches, threshold) in [(1, 1), (2, usize::MAX), (2, 1)] { + for enabled in [false, true] { + let rows_evaluated = Arc::new(AtomicUsize::new(0)); + let counter = Arc::clone(&rows_evaluated); + let function = create_udf( + "count_evaluations", + vec![DataType::Int32], + DataType::Int32, + Volatility::Volatile, + Arc::new(move |args| { + counter.fetch_add( + args[0].clone().into_array(1)?.len(), + AtomicOrdering::Relaxed, + ); + Ok(args[0].clone()) + }), + ); + let ordering = [ + PhysicalSortExpr::new_default(col("prefix", &schema)?), + PhysicalSortExpr::new_default(Arc::new( + ScalarFunctionExpr::try_new( + Arc::new(function), + vec![col("id", &schema)?], + &schema, + Arc::new(ConfigOptions::default()), + )?, + )), + ] + .into(); + staged_sort_test_output( + &batch, + &ordering, + enabled, + input_batches, + threshold, + ) + .await?; + let merge_rows = if input_batches == 2 && threshold == 1 { + 6 + } else { + 0 + }; + let sort_rows = if enabled { tied_rows } else { 6 }; + assert_eq!( + rows_evaluated.load(AtomicOrdering::Relaxed), + sort_rows + merge_rows, + "enabled={enabled}, tied_rows={tied_rows}, input_batches={input_batches}, threshold={threshold}" + ); + } + } + } + Ok(()) + } + #[tokio::test] async fn test_sort_batch_chunked_basic() -> Result<()> { let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)]));