From 9c179b869a1fd3a3fde509048ba1ff3de79a5aaf Mon Sep 17 00:00:00 2001 From: yashrb24 Date: Wed, 16 Sep 2026 08:58:00 +0530 Subject: [PATCH 1/9] 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/9] 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/9] 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/9] 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/9] 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/9] 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/9] 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)])); From cfc9e495be43539870e81f88bc3fcaaa505367d7 Mon Sep 17 00:00:00 2001 From: yashrb24 Date: Tue, 3 Nov 2026 19:55:45 +0530 Subject: [PATCH 8/9] fix: refine automatic staged sort evaluation --- datafusion/common/src/config.rs | 8 +- .../physical-expr/src/expressions/binary.rs | 2 +- .../physical-expr/src/expressions/mod.rs | 2 +- datafusion/physical-plan/src/sorts/sort.rs | 1201 +++++++++++++---- datafusion/physical-plan/src/sorts/stream.rs | 5 +- .../test_files/information_schema.slt | 4 +- docs/source/user-guide/configs.md | 3 +- 7 files changed, 915 insertions(+), 310 deletions(-) diff --git a/datafusion/common/src/config.rs b/datafusion/common/src/config.rs index 7ebb4b5eeffb0..8b840a6e2b397 100644 --- a/datafusion/common/src/config.rs +++ b/datafusion/common/src/config.rs @@ -998,13 +998,11 @@ config_namespace! { /// batches and merged. pub sort_in_place_threshold_bytes: usize, default = 1024 * 1024 - /// Enables staged in-memory sorting, evaluating later sort keys only for rows - /// tied on earlier keys. + /// Enables staged sorting for eligible in-memory sorts. + /// Later keys may be skipped for rows already in order + /// TopK and spill sorting do not use staged sorting. 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 = 5 - /// Maximum buffer capacity (in bytes) per partition for BufferExec /// inserted during sort pushdown optimization. /// diff --git a/datafusion/physical-expr/src/expressions/binary.rs b/datafusion/physical-expr/src/expressions/binary.rs index d3cb36e3f6ab9..9f7efd0c5408c 100644 --- a/datafusion/physical-expr/src/expressions/binary.rs +++ b/datafusion/physical-expr/src/expressions/binary.rs @@ -1443,7 +1443,7 @@ fn rows_to_filter_before( /// Whether evaluating `expr` on all rows is cheap and cannot fail. /// Inspired by `reorder_predicates`' cheap forms, with fallible forms excluded. -fn is_cheap_and_infallible(expr: &Arc, schema: &Schema) -> bool { +pub fn is_cheap_and_infallible(expr: &Arc, schema: &Schema) -> bool { use crate::expressions::{ CastExpr, Column, InListExpr, IsNotNullExpr, IsNullExpr, Literal, NotExpr, }; diff --git a/datafusion/physical-expr/src/expressions/mod.rs b/datafusion/physical-expr/src/expressions/mod.rs index 767c0b4ce40f3..d4f933c91de69 100644 --- a/datafusion/physical-expr/src/expressions/mod.rs +++ b/datafusion/physical-expr/src/expressions/mod.rs @@ -28,7 +28,7 @@ pub use crate::PhysicalSortExpr; /// Module with some convenient methods used in expression building pub use crate::aggregate::stats::StatsType; -pub use binary::{BinaryExpr, binary, similar_to}; +pub use binary::{BinaryExpr, binary, is_cheap_and_infallible, similar_to}; pub use case::{CaseExpr, case}; pub use cast::{CastExpr, cast}; pub use column::{Column, col, with_new_schema}; diff --git a/datafusion/physical-plan/src/sorts/sort.rs b/datafusion/physical-plan/src/sorts/sort.rs index 6cbe174608b83..c444e7cda5994 100644 --- a/datafusion/physical-plan/src/sorts/sort.rs +++ b/datafusion/physical-plan/src/sorts/sort.rs @@ -69,10 +69,11 @@ use datafusion_execution::runtime_env::RuntimeEnv; use datafusion_physical_expr::LexOrdering; use datafusion_physical_expr::PhysicalExpr; use datafusion_physical_expr::expressions::{ - Column, DynamicFilterPhysicalExpr, Literal, lit, + DynamicFilterPhysicalExpr, is_cheap_and_infallible, lit, }; use datafusion_physical_expr::projection::ProjectionExprs; use datafusion_physical_expr::utils::collect_columns; +use datafusion_physical_expr_common::physical_expr::is_volatile; use futures::{Stream, StreamExt, TryStreamExt}; use log::{debug, trace}; @@ -380,63 +381,15 @@ impl ExternalSorter { batch_size, sort_spill_reservation_bytes, sort_in_place_threshold_bytes, - sort_stages: None, sort_concurrency: Arc::new(OnceLock::new()), + sort_stages: None, }) } - 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(), - )); + fn with_staged_sort(mut self, enabled: bool) -> Result { + if enabled { + self.sort_stages = create_sort_stages(&self.expr, &self.schema)?; } - 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) } @@ -1240,105 +1193,85 @@ impl ExternalSorter { } } -/// 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 -/// creating sorted copies of sorted columns in record batches for speeding up comparison -/// in sorting and merging. The sorted copies are in either row format or array format. -/// Please refer to cursor.rs and stream.rs for more details. No matter what format the -/// sorted copies are, they will use more memory than the original record batch. -/// -/// This can basically be calculated as the sum of the actual space it takes in -/// memory (which would be larger for a sliced batch), and the size of the actual data. -pub(crate) fn get_reserved_bytes_for_record_batch_size( - record_batch_size: usize, - sliced_size: usize, -) -> usize { - // Even 2x may not be enough for some cases, but it's a good enough estimation as a baseline. - // If 2x is not enough, user can set a larger value for `sort_spill_reservation_bytes` - // to compensate for the extra memory needed. - record_batch_size + sliced_size -} - -/// Estimate how much memory is needed to sort a `RecordBatch`. -/// This will just call `get_reserved_bytes_for_record_batch_size` with the -/// memory size of the record batch and its sliced size. -pub(crate) fn get_reserved_bytes_for_record_batch(batch: &RecordBatch) -> Result { - batch.get_sliced_size().map(|sliced_size| { - get_reserved_bytes_for_record_batch_size( - get_record_batch_memory_size(batch), - sliced_size, - ) - }) +struct SortStage { + input_columns: Vec, + expressions: Vec, } -/// Estimate how much memory is needed to sort `batch` when it is buffered with -/// the batches `counter` has already counted. +/// Preserve key order and start a group before each expensive expression. /// -/// Like [`get_reserved_bytes_for_record_batch`], but only counts the buffers of -/// `batch` that `counter` has not seen. Buffered batches can share buffers, for -/// example zero-copy slices of one larger batch, and a shared buffer only needs -/// to be reserved once. Without a `counter`, every buffer of `batch` is counted. -fn get_reserved_bytes_for_next_record_batch( - batch: &RecordBatch, - counter: Option<&mut RecordBatchMemoryCounter>, -) -> Result { - let Some(counter) = counter else { - return get_reserved_bytes_for_record_batch(batch); - }; - let sliced_size = batch.get_sliced_size()?; - Ok(get_reserved_bytes_for_record_batch_size( - counter.count_batch(batch), - sliced_size, - )) -} - -/// Returns whether the sorted copy of an array of `data_type` keeps -/// referencing some of the array's buffers. +/// Examples: /// -/// Sorting takes the selected rows into new buffers, except that `take` keeps -/// the data buffers of view arrays and the values of dictionaries and list -/// views. Every sorted run is charged for those buffers in full, so buffered -/// batches that share them must be charged for them each, too. Otherwise the -/// sorter buffers more batches than the merge of their sorted runs can hold. -fn sorted_copy_keeps_buffers(data_type: &DataType) -> bool { - match data_type { - DataType::Utf8View - | DataType::BinaryView - | DataType::Dictionary(_, _) - | DataType::ListView(_) - | DataType::LargeListView(_) => true, - DataType::List(field) - | DataType::LargeList(field) - | DataType::FixedSizeList(field, _) - | DataType::Map(field, _) => sorted_copy_keeps_buffers(field.data_type()), - DataType::Struct(fields) => fields - .iter() - .any(|field| sorted_copy_keeps_buffers(field.data_type())), - DataType::Union(fields, _) => fields - .iter() - .any(|(_, field)| sorted_copy_keeps_buffers(field.data_type())), - DataType::RunEndEncoded(_, values) => { - sorted_copy_keeps_buffers(values.data_type()) +/// ```text +/// a, lower(name), id => [a] -> [lower(name), id] +/// a, b, lower(name), upper(city), id => [a, b] -> [lower(name)] -> [upper(city), id] +/// lower(name), upper(city), id => [lower(name)] -> [upper(city), id] +/// a, CAST(text AS INT), id => [a] -> [CAST(text AS INT), id] +/// ``` +fn create_sort_stages( + expressions: &LexOrdering, + schema: &SchemaRef, +) -> Result>> { + if expressions.iter().any(|key| is_volatile(&key.expr)) { + return Ok(None); + } + + let mut starts = vec![0]; + for (index, key) in expressions.iter().enumerate().skip(1) { + if !is_cheap_and_infallible(&key.expr, schema) { + starts.push(index); } - _ => false, } -} -impl Debug for ExternalSorter { - fn fmt(&self, f: &mut Formatter) -> fmt::Result { - f.debug_struct("ExternalSorter") - .field("memory_used", &self.used()) - .field("spilled_bytes", &self.spilled_bytes()) - .field("spilled_rows", &self.spilled_rows()) - .field("spill_count", &self.spill_count()) - .finish() + if starts.len() <= 1 { + return Ok(None); + } + starts.push(expressions.len()); + + let mut stages = vec![SortStage { + input_columns: vec![], + expressions: expressions[starts[0]..starts[1]].to_vec(), + }]; + for bounds in starts[1..].windows(2) { + stages.push(project_sort_stage( + &expressions[bounds[0]..bounds[1]], + schema, + )?); } + + Ok(Some(stages.into())) } -struct SortStage { - columns: Vec, - expressions: Vec, +/// Prepare the columns and expressions needed for a later sort stage +fn project_sort_stage( + keys: &[PhysicalSortExpr], + schema: &SchemaRef, +) -> Result { + let mut columns = keys + .iter() + .flat_map(|key| collect_columns(&key.expr)) + .map(|column| column.index()) + .collect::>(); + + // Collect each input column only once + columns.sort_unstable(); + columns.dedup(); + + let projection = ProjectionExprs::from_indices(&columns, schema); + let expressions = keys + .iter() + .map(|key| { + Ok(PhysicalSortExpr::new( + projection.project_expr(&key.expr)?, + key.options, + )) + }) + .collect::>>()?; + + Ok(SortStage { + input_columns: columns, + expressions, + }) } fn staged_sort_batch_chunked( @@ -1347,55 +1280,159 @@ fn staged_sort_batch_chunked( batch_size: usize, stages: &[SortStage], ) -> Result> { - let indices = staged_sort_indices(batch, stages)?; + let indices = staged_sort_indices(batch, stages, expressions)?; 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, 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(); - for (index, stage) in stages.iter().enumerate() { - if unresolved.is_empty() { +/// Return row indices in sort order, evaluating later stages only for ties. +/// +/// For `ORDER BY department, lower(name)`: +/// - Sort all rows by `department`. +/// - Evaluate `lower(name)` only for rows with the same department, then +/// sort those rows within each department. +/// - Return the original row indices in the resulting order. +/// +/// If there are many tied groups, we evaluate the remaining keys together +/// for the rows still tied. +fn staged_sort_indices( + batch: &RecordBatch, + stages: &[SortStage], + expressions: &LexOrdering, +) -> Result { + let (mut indices, mut tied_groups) = sort_first_stage(batch, &stages[0])?; + let mut stage_rows = batch.num_rows(); + let mut key_offset = stages[0].expressions.len(); + + for (index, stage) in stages.iter().enumerate().skip(1) { + if tied_groups.is_empty() { + break; + } + + let tied_rows: usize = tied_groups.iter().map(Range::len).sum(); + let resolved_rows = stage_rows - tied_rows; + + // Compare the number of groups with the reduction from the last stage. + if tied_groups.len() > resolved_rows { + let remaining = + project_sort_stage(&expressions[key_offset..], &batch.schema())?; + sort_remaining_keys(batch, &remaining, &mut indices, &tied_groups)?; break; } - unresolved = refine_tied_groups( + + stage_rows = tied_rows; + tied_groups = refine_tied_groups( batch, stage, &mut indices, - &unresolved, - index == 0, + &tied_groups, index + 1 < stages.len(), )?; + key_offset += stage.expressions.len(); } Ok(UInt32Array::from(indices)) } -/// Evaluate keys once for all unresolved rows, then sort each tied group independently. +/// Sort the first stage and find the groups needing later keys. +fn sort_first_stage( + batch: &RecordBatch, + stage: &SortStage, +) -> Result<(Vec, Vec>)> { + let columns = stage + .expressions + .iter() + .map(|expr| expr.evaluate_to_sort_column(batch)) + .collect::>>()?; + let order = lexsort_to_indices(&columns, None)?; + let comparator = LexicographicalComparator::try_new(&columns)?; + + let mut tied_groups = Vec::new(); + append_tied_ranges(&comparator, &order, 0, 0, &mut tied_groups); + let (_, values, _) = order.into_parts(); + Ok((values.into(), tied_groups)) +} + +/// Gather the columns and rows needed to evaluate the next keys. +fn gather_tied_rows( + batch: &RecordBatch, + stage: &SortStage, + indices: &[u32], + tied_groups: &[Range], +) -> Result<(RecordBatch, UInt32Array)> { + let selected_rows = tied_groups.iter().map(Range::len).sum(); + let mut selection = Vec::with_capacity(selected_rows); + for range in tied_groups { + selection.extend_from_slice(&indices[range.clone()]); + } + let selection = UInt32Array::from(selection); + let projected = batch.project(&stage.input_columns)?; + + let selected = if stage.input_columns.is_empty() { + RecordBatch::try_new_with_options( + projected.schema(), + vec![], + &RecordBatchOptions::new().with_row_count(Some(selected_rows)), + )? + } else { + arrow::compute::take_record_batch(&projected, &selection)? + }; + Ok((selected, selection)) +} + +/// Evaluate all remaining keys for tied rows and sort within each group. +fn sort_remaining_keys( + batch: &RecordBatch, + stage: &SortStage, + indices: &mut [u32], + tied_groups: &[Range], +) -> Result<()> { + let (selected, original_rows) = gather_tied_rows(batch, stage, indices, tied_groups)?; + let columns = stage + .expressions + .iter() + .map(|expr| expr.evaluate_to_sort_column(&selected)) + .collect::>>()?; + let comparator = LexicographicalComparator::try_new(&columns)?; + + // Comparator indices address the compact batch; map them back afterward. + let mut compact_order: Vec = (0..original_rows.len() as u32).collect(); + let mut offset = 0; + for range in tied_groups { + let end = offset + range.len(); + compact_order[offset..end] + .sort_unstable_by(|a, b| comparator.compare(*a as usize, *b as usize)); + for (output, compact_row) in indices[range.clone()] + .iter_mut() + .zip(&compact_order[offset..end]) + { + *output = original_rows.value(*compact_row as usize); + } + offset = end; + } + Ok(()) +} + +/// Sort each tied group using the next stage's keys. +/// +/// For example, given ties on a given key: +/// +/// ```text +/// Before: A: [Zoe, Ada] B: [Mia, Bob] C: [Lea] +/// After: A: [Ada, Zoe] B: [Bob, Mia] C: [Lea] +/// ``` +/// +/// This updates `indices` in place. fn refine_tied_groups( batch: &RecordBatch, stage: &SortStage, - indices: &mut Vec, + indices: &mut [u32], unresolved: &[Range], - first_stage: bool, find_ties: bool, ) -> Result>> { - let selected = if first_stage { - batch.clone() - } else { - 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.project(&stage.columns)?, - &UInt32Array::from(selection), - )? - }; + let (selected, _) = gather_tied_rows(batch, stage, indices, unresolved)?; + let sort_columns = stage .expressions .iter() @@ -1409,16 +1446,6 @@ fn refine_tied_groups( .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 @@ -1428,12 +1455,14 @@ fn refine_tied_groups( 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, @@ -1445,11 +1474,14 @@ fn refine_tied_groups( } 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`. +/// Find groups of equal keys that need to be sorted by the next stage. +/// +/// For sorted keys `[A, A, B, C, C]`, append the ranges `0..2` and `3..5`. +/// A single row, like `B`, is already in order and needs no further sorting. fn append_tied_ranges( comparator: &LexicographicalComparator, order: &UInt32Array, @@ -1474,6 +1506,102 @@ fn append_tied_ranges( } } +/// 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 +/// creating sorted copies of sorted columns in record batches for speeding up comparison +/// in sorting and merging. The sorted copies are in either row format or array format. +/// Please refer to cursor.rs and stream.rs for more details. No matter what format the +/// sorted copies are, they will use more memory than the original record batch. +/// +/// This can basically be calculated as the sum of the actual space it takes in +/// memory (which would be larger for a sliced batch), and the size of the actual data. +pub(crate) fn get_reserved_bytes_for_record_batch_size( + record_batch_size: usize, + sliced_size: usize, +) -> usize { + // Even 2x may not be enough for some cases, but it's a good enough estimation as a baseline. + // If 2x is not enough, user can set a larger value for `sort_spill_reservation_bytes` + // to compensate for the extra memory needed. + record_batch_size + sliced_size +} + +/// Estimate how much memory is needed to sort a `RecordBatch`. +/// This will just call `get_reserved_bytes_for_record_batch_size` with the +/// memory size of the record batch and its sliced size. +pub(crate) fn get_reserved_bytes_for_record_batch(batch: &RecordBatch) -> Result { + batch.get_sliced_size().map(|sliced_size| { + get_reserved_bytes_for_record_batch_size( + get_record_batch_memory_size(batch), + sliced_size, + ) + }) +} + +/// Estimate how much memory is needed to sort `batch` when it is buffered with +/// the batches `counter` has already counted. +/// +/// Like [`get_reserved_bytes_for_record_batch`], but only counts the buffers of +/// `batch` that `counter` has not seen. Buffered batches can share buffers, for +/// example zero-copy slices of one larger batch, and a shared buffer only needs +/// to be reserved once. Without a `counter`, every buffer of `batch` is counted. +fn get_reserved_bytes_for_next_record_batch( + batch: &RecordBatch, + counter: Option<&mut RecordBatchMemoryCounter>, +) -> Result { + let Some(counter) = counter else { + return get_reserved_bytes_for_record_batch(batch); + }; + let sliced_size = batch.get_sliced_size()?; + Ok(get_reserved_bytes_for_record_batch_size( + counter.count_batch(batch), + sliced_size, + )) +} + +/// Returns whether the sorted copy of an array of `data_type` keeps +/// referencing some of the array's buffers. +/// +/// Sorting takes the selected rows into new buffers, except that `take` keeps +/// the data buffers of view arrays and the values of dictionaries and list +/// views. Every sorted run is charged for those buffers in full, so buffered +/// batches that share them must be charged for them each, too. Otherwise the +/// sorter buffers more batches than the merge of their sorted runs can hold. +fn sorted_copy_keeps_buffers(data_type: &DataType) -> bool { + match data_type { + DataType::Utf8View + | DataType::BinaryView + | DataType::Dictionary(_, _) + | DataType::ListView(_) + | DataType::LargeListView(_) => true, + DataType::List(field) + | DataType::LargeList(field) + | DataType::FixedSizeList(field, _) + | DataType::Map(field, _) => sorted_copy_keeps_buffers(field.data_type()), + DataType::Struct(fields) => fields + .iter() + .any(|field| sorted_copy_keeps_buffers(field.data_type())), + DataType::Union(fields, _) => fields + .iter() + .any(|(_, field)| sorted_copy_keeps_buffers(field.data_type())), + DataType::RunEndEncoded(_, values) => { + sorted_copy_keeps_buffers(values.data_type()) + } + _ => false, + } +} + +impl Debug for ExternalSorter { + fn fmt(&self, f: &mut Formatter) -> fmt::Result { + f.debug_struct("ExternalSorter") + .field("memory_used", &self.used()) + .field("spilled_bytes", &self.spilled_bytes()) + .field("spilled_rows", &self.spilled_rows()) + .field("spill_count", &self.spill_count()) + .finish() + } +} + pub fn sort_batch( batch: &RecordBatch, expressions: &LexOrdering, @@ -2054,10 +2182,7 @@ impl ExecutionPlan for SortExec { &self.metrics_set, context.runtime_env(), )? - .with_staged_sort( - execution_options.enable_staged_sort, - execution_options.sort_key_group_size, - )?; + .with_staged_sort(execution_options.enable_staged_sort)?; sorter.sort_concurrency = Arc::clone(&self.sort_concurrency); Ok(Box::pin(RecordBatchStreamAdapter::new( self.schema(), @@ -2707,7 +2832,6 @@ mod tests { use std::collections::HashMap; use std::path::Path; use std::pin::Pin; - use std::sync::atomic::{AtomicUsize, Ordering as AtomicOrdering}; use std::task::{Context, Poll}; use super::*; @@ -2745,12 +2869,15 @@ mod tests { use datafusion_execution::{ AsyncSpillWriter, SpillFile, SpillWriter, TempFileFactory, }; + use datafusion_expr::Operator; + use datafusion_physical_expr::expressions::{BinaryExpr, CastExpr}; use datafusion_physical_expr::expressions::{Column, Literal}; use datafusion_physical_expr::{DynamicFilterTracking, EquivalenceProperties}; use datafusion_physical_expr_common::metrics::MetricValue; use futures::{FutureExt, Stream, TryStreamExt}; use insta::assert_snapshot; + use std::sync::atomic::{AtomicUsize, Ordering as AtomicOrdering}; use tokio::sync::Notify; #[tokio::test] @@ -2908,6 +3035,52 @@ mod tests { abort_count: Arc, } + #[test] + fn test_staged_sort_automatic_groups() -> Result<()> { + let schema = Arc::new(Schema::new(vec![ + Field::new("i", DataType::Int32, false), + Field::new("s", DataType::Utf8, false), + ])); + let cheap = col("i", &schema)?; + let expensive: Arc = + Arc::new(CastExpr::new(col("s", &schema)?, DataType::Int32, None)); + let group_sizes = |classes: &[bool]| -> Result> { + let keys = + LexOrdering::new(classes.iter().enumerate().map(|(index, cheap_key)| { + let expr = if *cheap_key { + Arc::new(BinaryExpr::new( + Arc::clone(&cheap), + Operator::Plus, + lit(index as i32), + )) as Arc + } else { + Arc::new(BinaryExpr::new( + Arc::clone(&expensive), + Operator::Plus, + lit(index as i32), + )) + }; + PhysicalSortExpr::new_default(expr) + })) + .unwrap(); + Ok(create_sort_stages(&keys, &schema)? + .map(|plan| plan.iter().map(|stage| stage.expressions.len()).collect()) + .unwrap_or_default()) + }; + assert_eq!( + group_sizes(&[true, true, false, false, true, false])?, + vec![2, 1, 2, 1] + ); + assert_eq!(group_sizes(&[false, false, true, false])?, vec![1, 2, 1]); + assert_eq!(group_sizes(&[true, false])?, vec![1, 1]); + assert_eq!(group_sizes(&[false, false])?, vec![1, 1]); + assert_eq!(group_sizes(&[false; 5])?, vec![1, 1, 1, 1, 1]); + assert!(group_sizes(&[true; 5])?.is_empty()); + assert!(group_sizes(&[false])?.is_empty()); + + Ok(()) + } + impl TempFileFactory for PendingWriteTempFileFactory { fn create_temp_file(&self, _description: &str) -> Result> { Ok(Arc::new(PendingWriteSpillFile { @@ -4304,8 +4477,41 @@ mod tests { Ok((sorted_batches, metrics)) } + fn counting_identity( + name: &str, + input: Arc, + schema: &SchemaRef, + volatility: datafusion_expr::Volatility, + ) -> Result<(Arc, Arc)> { + use datafusion_expr::create_udf; + use datafusion_physical_expr::ScalarFunctionExpr; + + let rows = Arc::new(AtomicUsize::new(0)); + let counter = Arc::clone(&rows); + let udf = create_udf( + name, + vec![DataType::Int32], + DataType::Int32, + volatility, + Arc::new(move |args| { + counter.fetch_add( + args[0].clone().into_array(1)?.len(), + AtomicOrdering::Relaxed, + ); + Ok(args[0].clone()) + }), + ); + let expression = Arc::new(ScalarFunctionExpr::try_new( + Arc::new(udf), + vec![input], + schema, + Arc::new(ConfigOptions::default()), + )?); + Ok((expression, rows)) + } + // Exercise SortExec directly so the planner cannot project the sort expressions - // into columns before they reach the sorter. Each key is a separate stage. + // into columns before they reach the sorter. async fn staged_sort_test_output( batch: &RecordBatch, ordering: &LexOrdering, @@ -4322,8 +4528,7 @@ mod tests { 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); + .set_bool("datafusion.execution.enable_staged_sort", enabled); let result = collect( Arc::clone(&sort) as Arc, Arc::new(TaskContext::default().with_session_config(config)), @@ -4340,21 +4545,340 @@ mod tests { result } + async fn check_staged_multi_run_path(concurrent: bool) -> Result<()> { + use datafusion_expr::Volatility; + + let batch = record_batch!( + ("p", Int32, [8, 6, 4, 2, 7, 5, 3, 1]), + ("v", Int32, [8, 6, 4, 2, 7, 5, 3, 1]) + )?; + let schema = batch.schema(); + for enabled in [false, true] { + let (suffix, rows) = counting_identity( + "count_suffix", + col("v", &schema)?, + &schema, + Volatility::Immutable, + )?; + + let ordering = [ + PhysicalSortExpr::new_default(col("p", &schema)?), + PhysicalSortExpr::new_default(suffix), + ] + .into(); + let input = TestMemoryExec::try_new_exec( + &[vec![batch.slice(0, 4), batch.slice(4, 4)]], + Arc::clone(&schema), + None, + )?; + let sort = Arc::new(SortExec::new(ordering, input)); + let config = SessionConfig::new() + .with_batch_size(2) + .with_sort_in_place_threshold_bytes(1) + .set_bool("datafusion.execution.enable_staged_sort", enabled); + let output = collect( + Arc::clone(&sort) as Arc, + Arc::new(TaskContext::default().with_session_config(config)), + ) + .await?; + assert_eq!(sort.sort_concurrency.get().is_some(), concurrent); + assert_eq!( + rows.load(AtomicOrdering::Relaxed), + if enabled { 8 } else { 16 } + ); + let expected = arrow::compute::take_record_batch( + &batch, + &UInt32Array::from(vec![7, 3, 6, 2, 5, 1, 4, 0]), + )?; + assert_eq!(concat_batches(&schema, &output)?, expected); + } + Ok(()) + } + + #[tokio::test] + async fn test_staged_sort_multi_run_inline() -> Result<()> { + check_staged_multi_run_path(false).await + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn test_staged_sort_multi_run_concurrent() -> Result<()> { + check_staged_multi_run_path(true).await + } + + #[tokio::test] + async fn test_staged_sort_zero_arg_udf() -> Result<()> { + use datafusion_expr::{ + ColumnarValue, ScalarFunctionArgs, ScalarUDF, ScalarUDFImpl, Signature, + Volatility, + }; + use datafusion_physical_expr::ScalarFunctionExpr; + static ROWS: AtomicUsize = AtomicUsize::new(0); + #[derive(Debug, PartialEq, Eq, Hash)] + struct ZeroArg { + signature: Signature, + } + impl ScalarUDFImpl for ZeroArg { + fn name(&self) -> &str { + "staged_zero_arg" + } + fn signature(&self) -> &Signature { + &self.signature + } + fn return_type(&self, _: &[DataType]) -> Result { + Ok(DataType::Int32) + } + fn invoke_with_args( + &self, + args: ScalarFunctionArgs, + ) -> Result { + assert!(args.args.is_empty()); + ROWS.fetch_add(args.number_rows, AtomicOrdering::Relaxed); + Ok(ColumnarValue::Array(Arc::new(Int32Array::from( + vec![7; args.number_rows], + )))) + } + } + let schema = Arc::new(Schema::new(vec![Field::new("p", DataType::Int32, false)])); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(Int32Array::from(vec![5, 1, 3, 2, 4, 1, 6, 7]))], + )?; + let udf = Arc::new(ScalarUDF::from(ZeroArg { + signature: Signature::exact(vec![], Volatility::Immutable), + })); + let expression = Arc::new(ScalarFunctionExpr::try_new( + udf, + vec![], + &schema, + Arc::new(ConfigOptions::default()), + )?); + let ordering = [ + PhysicalSortExpr::new_default(col("p", &schema)?), + PhysicalSortExpr::new_default(expression), + ] + .into(); + for (input_batches, threshold) in [(1, 1), (2, usize::MAX)] { + ROWS.store(0, AtomicOrdering::Relaxed); + let result = staged_sort_test_output( + &batch, + &ordering, + true, + input_batches, + threshold, + ) + .await?; + assert_eq!(ROWS.load(AtomicOrdering::Relaxed), 2); + let actual = concat_batches(&schema, &result)?; + assert_eq!( + actual.column(0).as_ref(), + &Int32Array::from(vec![1, 1, 2, 3, 4, 5, 6, 7]) + ); + } + Ok(()) + } + + #[tokio::test] + async fn test_staged_sort_fallback_reuses_prefix() -> Result<()> { + use datafusion_expr::Volatility; + + let schema = Arc::new(Schema::new(vec![ + Field::new("p", DataType::Int32, false), + Field::new("s", DataType::Utf8, false), + ])); + let (prefix, rows) = counting_identity( + "count_prefix", + col("p", &schema)?, + &schema, + Volatility::Immutable, + )?; + let suffix = Arc::new(CastExpr::new(col("s", &schema)?, DataType::Int32, None)); + let ordering = [ + PhysicalSortExpr::new_default(prefix), + PhysicalSortExpr::new_default(suffix), + ] + .into(); + for (prefix, suffix, expected_indices) in [ + ( + vec![1; 8], + vec!["8", "7", "6", "5", "4", "3", "2", "1"], + vec![7, 6, 5, 4, 3, 2, 1, 0], + ), + ( + vec![0, 0, 1, 1, 2, 2, 3, 3], + vec!["8", "7", "6", "5", "4", "3", "2", "1"], + vec![1, 0, 3, 2, 5, 4, 7, 6], + ), + ] { + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int32Array::from(prefix)), + Arc::new(StringArray::from(suffix)), + ], + )?; + let expected = arrow::compute::take_record_batch( + &batch, + &UInt32Array::from(expected_indices), + )?; + rows.store(0, AtomicOrdering::Relaxed); + let result = staged_sort_test_output(&batch, &ordering, true, 1, 1).await?; + assert_eq!(rows.load(AtomicOrdering::Relaxed), 8); + assert_eq!(concat_batches(&schema, &result)?, expected); + } + Ok(()) + } + + #[tokio::test] + async fn test_staged_sort_fallback_ignores_resolved_rows() -> Result<()> { + let batch = record_batch!( + ("prefix", Int32, [0, 0, 1, 1, 2, 3, 3, 4, 4]), + ( + "suffix", + Utf8, + ["8", "7", "6", "5", "invalid", "4", "3", "2", "1"] + ) + )?; + let schema = batch.schema(); + let suffix = Arc::new(CastExpr::new( + col("suffix", &schema)?, + DataType::Int32, + None, + )); + let ordering = [ + PhysicalSortExpr::new_default(col("prefix", &schema)?), + PhysicalSortExpr::new_default(suffix), + ] + .into(); + let expected = arrow::compute::take_record_batch( + &batch, + &UInt32Array::from(vec![1, 0, 3, 2, 4, 6, 5, 8, 7]), + )?; + + let output = staged_sort_test_output(&batch, &ordering, true, 1, 1).await?; + assert_eq!(concat_batches(&schema, &output)?, expected); + Ok(()) + } + + #[tokio::test] + async fn test_staged_sort_fallback_after_second_stage() -> Result<()> { + use datafusion_expr::Volatility; + + let batch = record_batch!( + ("prefix", Int32, [0, 0, 0, 0, 1, 1, 1, 1, 2, 3, 4, 5]), + ("second", Int32, [0, 0, 1, 1, 0, 0, 1, 1, 0, 0, 0, 0]), + ("third", Int32, [2, 1, 2, 1, 2, 1, 2, 1, 0, 0, 0, 0]), + ("fourth", Int32, [9, 8, 7, 6, 5, 4, 3, 2, 0, 0, 0, 0]) + )?; + let schema = batch.schema(); + let (second, second_rows) = counting_identity( + "fallback_second", + col("second", &schema)?, + &schema, + Volatility::Immutable, + )?; + let (third, third_rows) = counting_identity( + "fallback_third", + col("third", &schema)?, + &schema, + Volatility::Immutable, + )?; + let (fourth, fourth_rows) = counting_identity( + "fallback_fourth", + col("fourth", &schema)?, + &schema, + Volatility::Immutable, + )?; + let ordering = [ + PhysicalSortExpr::new_default(col("prefix", &schema)?), + PhysicalSortExpr::new_default(second), + PhysicalSortExpr::new_default(third), + PhysicalSortExpr::new_default(fourth), + ] + .into(); + let expected = arrow::compute::take_record_batch( + &batch, + &UInt32Array::from(vec![1, 0, 3, 2, 5, 4, 7, 6, 8, 9, 10, 11]), + )?; + + let output = staged_sort_test_output(&batch, &ordering, true, 1, 1).await?; + assert_eq!(concat_batches(&schema, &output)?, expected); + assert_eq!(second_rows.load(AtomicOrdering::Relaxed), 8); + assert_eq!(third_rows.load(AtomicOrdering::Relaxed), 8); + assert_eq!(fourth_rows.load(AtomicOrdering::Relaxed), 8); + Ok(()) + } + + #[tokio::test] + async fn test_staged_sort_successive_deferred_groups() -> Result<()> { + use datafusion_expr::Volatility; + + let batch = record_batch!( + ("prefix", Int32, [0, 0, 0, 1, 1, 1, 2, 3, 4, 5, 6, 7]), + ("second", Int32, [1, 1, 2, 1, 1, 2, 0, 0, 0, 0, 0, 0]), + ("third", Int32, [2, 1, 0, 2, 1, 0, 0, 0, 0, 0, 0, 0]) + )?; + let schema = batch.schema(); + let (second, second_rows) = counting_identity( + "second_stage", + col("second", &schema)?, + &schema, + Volatility::Immutable, + )?; + let (third, third_rows) = counting_identity( + "third_stage", + col("third", &schema)?, + &schema, + Volatility::Immutable, + )?; + let ordering = [ + PhysicalSortExpr::new_default(col("prefix", &schema)?), + PhysicalSortExpr::new_default(second), + PhysicalSortExpr::new_default(third), + ] + .into(); + let expected = arrow::compute::take_record_batch( + &batch, + &UInt32Array::from(vec![1, 0, 2, 4, 3, 5, 6, 7, 8, 9, 10, 11]), + )?; + for (input_batches, threshold) in [(1, 1), (2, usize::MAX)] { + second_rows.store(0, AtomicOrdering::Relaxed); + third_rows.store(0, AtomicOrdering::Relaxed); + let output = staged_sort_test_output( + &batch, + &ordering, + true, + input_batches, + threshold, + ) + .await?; + assert_eq!(concat_batches(&schema, &output)?, expected); + assert_eq!(second_rows.load(AtomicOrdering::Relaxed), 6); + assert_eq!(third_rows.load(AtomicOrdering::Relaxed), 4); + } + Ok(()) + } + #[tokio::test] async fn test_staged_sort_matches_eager() -> Result<()> { - use datafusion_physical_expr::expressions::CastExpr; + use datafusion_expr::Volatility; - // Keep payload columns between the keys to exercise projected column remapping. + // Deferred keys reference source columns 1 and 3, which are remapped + // to columns 0 and 1 after projecting away the payload columns. let schema = Arc::new(Schema::new(vec![ + Field::new("payload", DataType::Utf8, false), Field::new("id", DataType::Int32, false), + Field::new("other_payload", 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_iter_values( + (0..16).map(|i| format!("payload-{i}")), + )), + Arc::new(Int32Array::from_iter_values(0..16)), + Arc::new(Int32Array::from_iter_values(100..116)), Arc::new(StringArray::from(vec![ Some("3"), Some("2"), @@ -4366,10 +4890,13 @@ mod tests { Some("2"), Some("0"), Some("10"), + Some("1"), + Some("1"), + Some("1"), + Some("1"), + Some("1"), + Some("1"), ])), - Arc::new(StringArray::from_iter_values( - (0..10).map(|i| format!("payload-{i}")), - )), Arc::new(Int32Array::from(vec![ Some(2), Some(1), @@ -4381,9 +4908,22 @@ mod tests { Some(1), Some(3), Some(1), + Some(4), + Some(5), + Some(6), + Some(7), + Some(8), + Some(9), ])), ], )?; + let cast = Arc::new(CastExpr::new( + col("suffix", &schema)?, + DataType::Int32, + None, + )); + let (suffix, suffix_rows) = + counting_identity("count_cast_rows", cast, &schema, Volatility::Immutable)?; let ordering = [ PhysicalSortExpr::new( col("prefix", &schema)?, @@ -4393,11 +4933,7 @@ mod tests { }, ), PhysicalSortExpr::new( - Arc::new(CastExpr::new( - col("suffix", &schema)?, - DataType::Int32, - None, - )), + suffix, SortOptions { descending: true, nulls_first: true, @@ -4406,15 +4942,20 @@ mod tests { PhysicalSortExpr::new_default(col("id", &schema)?), ] .into(); - // Prefix ties survive the numeric suffix stage and must be ordered by id. + + // Nine rows remain tied in three groups; seven rows are already resolved. + // The prefix and numeric suffix still tie for some rows, so id breaks ties. let expected = arrow::compute::take_record_batch( &batch, - &UInt32Array::from(vec![3, 5, 9, 1, 7, 0, 4, 8, 6, 2]), + &UInt32Array::from(vec![ + 3, 5, 9, 1, 7, 0, 4, 8, 10, 11, 12, 13, 14, 15, 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] { + suffix_rows.store(0, AtomicOrdering::Relaxed); let output = staged_sort_test_output( &batch, &ordering, @@ -4423,11 +4964,15 @@ mod tests { threshold, ) .await?; + assert_eq!( concat_batches(&schema, &output)?, expected, "enabled={enabled}, input_batches={input_batches}, threshold={threshold}" ); + if enabled && (input_batches == 1 || threshold == usize::MAX) { + assert_eq!(suffix_rows.load(AtomicOrdering::Relaxed), 9); + } } } Ok(()) @@ -4450,45 +4995,98 @@ mod tests { ))), ] .into(); - for all_tied in [false, true] { - let prefix = if all_tied { - vec![1; 6] - } else { - vec![3, 1, 6, 2, 5, 4] - }; + struct Case<'a> { + name: &'a str, + prefix: &'a [i32; 6], + enabled: bool, + input_batches: usize, + threshold: usize, + expect_error: bool, + } + + let distinct = [3, 1, 6, 2, 5, 4]; + let all_tied = [1; 6]; + let cases = [ + Case { + name: "disabled evaluates eagerly", + prefix: &distinct, + enabled: false, + input_batches: 1, + threshold: 1, + expect_error: true, + }, + Case { + name: "distinct rows skip suffix", + prefix: &distinct, + enabled: true, + input_batches: 1, + threshold: 1, + expect_error: false, + }, + Case { + name: "ties evaluate suffix", + prefix: &all_tied, + enabled: true, + input_batches: 1, + threshold: 1, + expect_error: true, + }, + Case { + name: "concatenated rows skip suffix", + prefix: &distinct, + enabled: true, + input_batches: 2, + threshold: usize::MAX, + expect_error: false, + }, + Case { + name: "merge reevaluates suffix", + prefix: &distinct, + enabled: true, + input_batches: 2, + threshold: 1, + expect_error: true, + }, + ]; + + for Case { + name, + prefix, + enabled, + input_batches, + threshold, + expect_error, + } in cases + { let batch = RecordBatch::try_new( Arc::clone(&schema), vec![ - Arc::new(Int32Array::from(prefix)), + Arc::new(Int32Array::from(prefix.to_vec())), 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), - ); - } - } + let result = staged_sort_test_output( + &batch, + &ordering, + enabled, + input_batches, + threshold, + ) + .await; + + if expect_error { + let error = result.unwrap_err().strip_backtrace(); + assert!( + error.contains("Cannot cast string 'invalid'"), + "{name}: {error}" + ); + } else { + let output = concat_batches(&schema, &result?)?; + assert_eq!( + output.column(0).as_ref(), + &Int32Array::from_iter_values(1..=6), + "{name}" + ); } } Ok(()) @@ -4496,77 +5094,90 @@ mod tests { #[tokio::test] async fn test_staged_sort_volatile_evaluation() -> Result<()> { - use datafusion_expr::{Volatility, create_udf}; - use datafusion_physical_expr::ScalarFunctionExpr; + use datafusion_expr::Volatility; 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), - ] { + struct Case<'a> { + name: &'a str, + prefix: &'a [i32; 6], + enabled: bool, + input_batches: usize, + threshold: usize, + expected_rows: usize, + } + + let distinct = [3, 1, 6, 2, 5, 4]; + let partly_tied = [2, 1, 1, 3, 4, 4]; + let cases = [ + Case { + name: "distinct enabled", + prefix: &distinct, + enabled: true, + input_batches: 1, + threshold: 1, + expected_rows: 6, + }, + Case { + name: "partly tied disabled", + prefix: &partly_tied, + enabled: false, + input_batches: 1, + threshold: 1, + expected_rows: 6, + }, + Case { + name: "partly tied enabled", + prefix: &partly_tied, + enabled: true, + input_batches: 1, + threshold: 1, + expected_rows: 6, + }, + Case { + name: "merge reevaluates", + prefix: &partly_tied, + enabled: true, + input_batches: 2, + threshold: 1, + expected_rows: 12, + }, + ]; + + for Case { + name, + prefix, + enabled, + input_batches, + threshold, + expected_rows, + } in cases + { let batch = RecordBatch::try_new( Arc::clone(&schema), vec![ - Arc::new(Int32Array::from(prefix)), + Arc::new(Int32Array::from(prefix.to_vec())), 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}" - ); - } - } + let (volatile, rows) = counting_identity( + "count_evaluations", + col("id", &schema)?, + &schema, + Volatility::Volatile, + )?; + let nested = Arc::new(BinaryExpr::new(volatile, Operator::Plus, lit(0i32))); + let ordering = [ + PhysicalSortExpr::new_default(col("prefix", &schema)?), + PhysicalSortExpr::new_default(nested), + ] + .into(); + + staged_sort_test_output(&batch, &ordering, enabled, input_batches, threshold) + .await?; + assert_eq!(rows.load(AtomicOrdering::Relaxed), expected_rows, "{name}"); } Ok(()) } diff --git a/datafusion/physical-plan/src/sorts/stream.rs b/datafusion/physical-plan/src/sorts/stream.rs index fc19b3307f00a..c0f92d785357d 100644 --- a/datafusion/physical-plan/src/sorts/stream.rs +++ b/datafusion/physical-plan/src/sorts/stream.rs @@ -327,9 +327,8 @@ 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). /// -/// 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 +/// On the first call to `next()`, a sorted index array (`UInt32Array`) is +/// computed via `lexsort_to_indices`. Subsequent calls yield 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 diff --git a/datafusion/sqllogictest/test_files/information_schema.slt b/datafusion/sqllogictest/test_files/information_schema.slt index 8776fdbd9f9d7..682be97aa2e68 100644 --- a/datafusion/sqllogictest/test_files/information_schema.slt +++ b/datafusion/sqllogictest/test_files/information_schema.slt @@ -275,7 +275,6 @@ 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_in_place_threshold_bytes 1048576 -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 @@ -381,7 +380,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 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. 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.enable_staged_sort false Enables staged sorting for eligible in-memory sorts. Later keys may be skipped for rows already in order TopK and spill sorting do not use staged sorting. 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 @@ -443,7 +442,6 @@ datafusion.execution.skip_physical_aggregate_schema_check false When set to true 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_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 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 079bf7e24e018..ea67d1db869cc 100644 --- a/docs/source/user-guide/configs.md +++ b/docs/source/user-guide/configs.md @@ -75,12 +75,10 @@ The following configuration settings are available: | datafusion.catalog.has_header | true | Default value for `format.has_header` for `CREATE EXTERNAL TABLE` if not specified explicitly in the statement. | | datafusion.catalog.newlines_in_values | false | Specifies whether newlines in (quoted) CSV values are supported. This is the default value for `format.newlines_in_values` for `CREATE EXTERNAL TABLE` if not specified explicitly in the statement. Parsing newlines in quoted values may be affected by execution behaviour such as parallel file scanning. Setting this to `true` ensures that newlines in values are parsed successfully, which may reduce performance. | | datafusion.execution.batch_size | 8192 | Default batch size while creating new batches, it's especially useful for buffer-in-memory batches since creating tiny batches would result in too much metadata memory consumption | -| datafusion.execution.enable_staged_sort | false | Enables staged in-memory sorting, evaluating later sort keys only for rows tied on earlier keys. | | datafusion.execution.perfect_hash_join_small_build_threshold | 1024 | A perfect hash join (see `HashJoinExec` for more details) will be considered if the range of keys (max - min) on the build side is < this threshold. This provides a fast path for joins with very small key ranges, bypassing the density check. Currently only supports cases where build_side.num_rows() < u32::MAX. Support for build_side.num_rows() >= u32::MAX will be added in the future. | | datafusion.execution.perfect_hash_join_min_key_density | 0.15 | The minimum required density of join keys on the build side to consider a perfect hash join (see `HashJoinExec` for more details). Density is calculated as: `(number of rows) / (max_key - min_key + 1)`. A perfect hash join may be used if the actual key density > this value. Currently only supports cases where build_side.num_rows() < u32::MAX. Support for build_side.num_rows() >= u32::MAX will be added in the future. | | datafusion.execution.coalesce_batches | true | When set to true, record batches will be examined between each operator and small batches will be coalesced into larger batches. This is helpful when there are highly selective filters or joins that could produce tiny output batches. The target batch size is determined by the configuration setting | | datafusion.execution.collect_statistics | true | Should DataFusion collect statistics when first creating a table. Has no effect after the table is created. Defaults to true. | -| datafusion.execution.sort_key_group_size | 5 | Number of sort keys evaluated together by staged in-memory sorting. | | datafusion.execution.target_partitions | 0 | Number of partitions for query execution. Increasing partitions can increase concurrency. Defaults to the number of CPU cores on the system | | datafusion.execution.time_zone | NULL | The default time zone Some functions, e.g. `now`, return timestamps in this time zone. This is also used to interpret timezone-naive timestamps in comparisons and subtraction with timezone-aware timestamps. | | datafusion.execution.parquet.enable_page_index | true | (reading) If true, reads the Parquet data page level metadata (the Page Index), if present, to reduce the I/O and number of rows decoded. | @@ -131,6 +129,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.enable_staged_sort | false | Enables staged sorting for eligible in-memory sorts. Later keys may be skipped for rows already in order TopK and spill sorting do not use staged 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 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. | From 22eaf89f7ea3908efc23341fbba53a6a1eb6c340 Mon Sep 17 00:00:00 2001 From: yashrb24 Date: Tue, 3 Nov 2026 22:30:15 +0530 Subject: [PATCH 9/9] test: trim staged sort coverage --- datafusion/physical-plan/src/sorts/sort.rs | 495 ++++----------------- 1 file changed, 77 insertions(+), 418 deletions(-) diff --git a/datafusion/physical-plan/src/sorts/sort.rs b/datafusion/physical-plan/src/sorts/sort.rs index c444e7cda5994..fd411e51594e6 100644 --- a/datafusion/physical-plan/src/sorts/sort.rs +++ b/datafusion/physical-plan/src/sorts/sort.rs @@ -3035,52 +3035,6 @@ mod tests { abort_count: Arc, } - #[test] - fn test_staged_sort_automatic_groups() -> Result<()> { - let schema = Arc::new(Schema::new(vec![ - Field::new("i", DataType::Int32, false), - Field::new("s", DataType::Utf8, false), - ])); - let cheap = col("i", &schema)?; - let expensive: Arc = - Arc::new(CastExpr::new(col("s", &schema)?, DataType::Int32, None)); - let group_sizes = |classes: &[bool]| -> Result> { - let keys = - LexOrdering::new(classes.iter().enumerate().map(|(index, cheap_key)| { - let expr = if *cheap_key { - Arc::new(BinaryExpr::new( - Arc::clone(&cheap), - Operator::Plus, - lit(index as i32), - )) as Arc - } else { - Arc::new(BinaryExpr::new( - Arc::clone(&expensive), - Operator::Plus, - lit(index as i32), - )) - }; - PhysicalSortExpr::new_default(expr) - })) - .unwrap(); - Ok(create_sort_stages(&keys, &schema)? - .map(|plan| plan.iter().map(|stage| stage.expressions.len()).collect()) - .unwrap_or_default()) - }; - assert_eq!( - group_sizes(&[true, true, false, false, true, false])?, - vec![2, 1, 2, 1] - ); - assert_eq!(group_sizes(&[false, false, true, false])?, vec![1, 2, 1]); - assert_eq!(group_sizes(&[true, false])?, vec![1, 1]); - assert_eq!(group_sizes(&[false, false])?, vec![1, 1]); - assert_eq!(group_sizes(&[false; 5])?, vec![1, 1, 1, 1, 1]); - assert!(group_sizes(&[true; 5])?.is_empty()); - assert!(group_sizes(&[false])?.is_empty()); - - Ok(()) - } - impl TempFileFactory for PendingWriteTempFileFactory { fn create_temp_file(&self, _description: &str) -> Result> { Ok(Arc::new(PendingWriteSpillFile { @@ -4545,66 +4499,6 @@ mod tests { result } - async fn check_staged_multi_run_path(concurrent: bool) -> Result<()> { - use datafusion_expr::Volatility; - - let batch = record_batch!( - ("p", Int32, [8, 6, 4, 2, 7, 5, 3, 1]), - ("v", Int32, [8, 6, 4, 2, 7, 5, 3, 1]) - )?; - let schema = batch.schema(); - for enabled in [false, true] { - let (suffix, rows) = counting_identity( - "count_suffix", - col("v", &schema)?, - &schema, - Volatility::Immutable, - )?; - - let ordering = [ - PhysicalSortExpr::new_default(col("p", &schema)?), - PhysicalSortExpr::new_default(suffix), - ] - .into(); - let input = TestMemoryExec::try_new_exec( - &[vec![batch.slice(0, 4), batch.slice(4, 4)]], - Arc::clone(&schema), - None, - )?; - let sort = Arc::new(SortExec::new(ordering, input)); - let config = SessionConfig::new() - .with_batch_size(2) - .with_sort_in_place_threshold_bytes(1) - .set_bool("datafusion.execution.enable_staged_sort", enabled); - let output = collect( - Arc::clone(&sort) as Arc, - Arc::new(TaskContext::default().with_session_config(config)), - ) - .await?; - assert_eq!(sort.sort_concurrency.get().is_some(), concurrent); - assert_eq!( - rows.load(AtomicOrdering::Relaxed), - if enabled { 8 } else { 16 } - ); - let expected = arrow::compute::take_record_batch( - &batch, - &UInt32Array::from(vec![7, 3, 6, 2, 5, 1, 4, 0]), - )?; - assert_eq!(concat_batches(&schema, &output)?, expected); - } - Ok(()) - } - - #[tokio::test] - async fn test_staged_sort_multi_run_inline() -> Result<()> { - check_staged_multi_run_path(false).await - } - - #[tokio::test(flavor = "multi_thread", worker_threads = 2)] - async fn test_staged_sort_multi_run_concurrent() -> Result<()> { - check_staged_multi_run_path(true).await - } - #[tokio::test] async fn test_staged_sort_zero_arg_udf() -> Result<()> { use datafusion_expr::{ @@ -4704,9 +4598,9 @@ mod tests { vec![7, 6, 5, 4, 3, 2, 1, 0], ), ( - vec![0, 0, 1, 1, 2, 2, 3, 3], - vec!["8", "7", "6", "5", "4", "3", "2", "1"], - vec![1, 0, 3, 2, 5, 4, 7, 6], + vec![0, 0, 1, 1, 2, 3, 3, 4, 4], + vec!["8", "7", "6", "5", "invalid", "4", "3", "2", "1"], + vec![1, 0, 3, 2, 4, 6, 5, 8, 7], ), ] { let batch = RecordBatch::try_new( @@ -4722,43 +4616,12 @@ mod tests { )?; rows.store(0, AtomicOrdering::Relaxed); let result = staged_sort_test_output(&batch, &ordering, true, 1, 1).await?; - assert_eq!(rows.load(AtomicOrdering::Relaxed), 8); + assert_eq!(rows.load(AtomicOrdering::Relaxed), batch.num_rows()); assert_eq!(concat_batches(&schema, &result)?, expected); } Ok(()) } - #[tokio::test] - async fn test_staged_sort_fallback_ignores_resolved_rows() -> Result<()> { - let batch = record_batch!( - ("prefix", Int32, [0, 0, 1, 1, 2, 3, 3, 4, 4]), - ( - "suffix", - Utf8, - ["8", "7", "6", "5", "invalid", "4", "3", "2", "1"] - ) - )?; - let schema = batch.schema(); - let suffix = Arc::new(CastExpr::new( - col("suffix", &schema)?, - DataType::Int32, - None, - )); - let ordering = [ - PhysicalSortExpr::new_default(col("prefix", &schema)?), - PhysicalSortExpr::new_default(suffix), - ] - .into(); - let expected = arrow::compute::take_record_batch( - &batch, - &UInt32Array::from(vec![1, 0, 3, 2, 4, 6, 5, 8, 7]), - )?; - - let output = staged_sort_test_output(&batch, &ordering, true, 1, 1).await?; - assert_eq!(concat_batches(&schema, &output)?, expected); - Ok(()) - } - #[tokio::test] async fn test_staged_sort_fallback_after_second_stage() -> Result<()> { use datafusion_expr::Volatility; @@ -4808,98 +4671,30 @@ mod tests { Ok(()) } - #[tokio::test] - async fn test_staged_sort_successive_deferred_groups() -> Result<()> { - use datafusion_expr::Volatility; - - let batch = record_batch!( - ("prefix", Int32, [0, 0, 0, 1, 1, 1, 2, 3, 4, 5, 6, 7]), - ("second", Int32, [1, 1, 2, 1, 1, 2, 0, 0, 0, 0, 0, 0]), - ("third", Int32, [2, 1, 0, 2, 1, 0, 0, 0, 0, 0, 0, 0]) - )?; - let schema = batch.schema(); - let (second, second_rows) = counting_identity( - "second_stage", - col("second", &schema)?, - &schema, - Volatility::Immutable, - )?; - let (third, third_rows) = counting_identity( - "third_stage", - col("third", &schema)?, - &schema, - Volatility::Immutable, - )?; - let ordering = [ - PhysicalSortExpr::new_default(col("prefix", &schema)?), - PhysicalSortExpr::new_default(second), - PhysicalSortExpr::new_default(third), - ] - .into(); - let expected = arrow::compute::take_record_batch( - &batch, - &UInt32Array::from(vec![1, 0, 2, 4, 3, 5, 6, 7, 8, 9, 10, 11]), - )?; - for (input_batches, threshold) in [(1, 1), (2, usize::MAX)] { - second_rows.store(0, AtomicOrdering::Relaxed); - third_rows.store(0, AtomicOrdering::Relaxed); - let output = staged_sort_test_output( - &batch, - &ordering, - true, - input_batches, - threshold, - ) - .await?; - assert_eq!(concat_batches(&schema, &output)?, expected); - assert_eq!(second_rows.load(AtomicOrdering::Relaxed), 6); - assert_eq!(third_rows.load(AtomicOrdering::Relaxed), 4); - } - Ok(()) - } - #[tokio::test] async fn test_staged_sort_matches_eager() -> Result<()> { use datafusion_expr::Volatility; - // Deferred keys reference source columns 1 and 3, which are remapped - // to columns 0 and 1 after projecting away the payload columns. - let schema = Arc::new(Schema::new(vec![ - Field::new("payload", DataType::Utf8, false), - Field::new("id", DataType::Int32, false), - Field::new("other_payload", DataType::Int32, false), - Field::new("suffix", DataType::Utf8, true), - Field::new("prefix", DataType::Int32, true), - ])); - let batch = RecordBatch::try_new( - Arc::clone(&schema), - vec![ - Arc::new(StringArray::from_iter_values( - (0..16).map(|i| format!("payload-{i}")), - )), - Arc::new(Int32Array::from_iter_values(0..16)), - Arc::new(Int32Array::from_iter_values(100..116)), - Arc::new(StringArray::from(vec![ + let batch = record_batch!( + ("id", Int32, [7, 6, 5, 4, 3, 2, 1, 0]), + ( + "suffix", + Utf8, + [ Some("3"), Some("2"), - Some("5"), + Some("0"), None, Some("3"), Some("10"), - None, - Some("2"), - Some("0"), - Some("10"), - Some("1"), Some("1"), - Some("1"), - Some("1"), - Some("1"), - Some("1"), - ])), - Arc::new(Int32Array::from(vec![ - Some(2), - Some(1), + Some("2") + ] + ), + ( + "prefix", + Int32, + [ None, Some(1), Some(2), @@ -4907,16 +4702,11 @@ mod tests { None, Some(1), Some(3), - Some(1), - Some(4), - Some(5), - Some(6), - Some(7), - Some(8), - Some(9), - ])), - ], + Some(1) + ] + ) )?; + let schema = batch.schema(); let cast = Arc::new(CastExpr::new( col("suffix", &schema)?, DataType::Int32, @@ -4924,6 +4714,12 @@ mod tests { )); let (suffix, suffix_rows) = counting_identity("count_cast_rows", cast, &schema, Volatility::Immutable)?; + let (id, id_rows) = counting_identity( + "count_id_rows", + col("id", &schema)?, + &schema, + Volatility::Immutable, + )?; let ordering = [ PhysicalSortExpr::new( col("prefix", &schema)?, @@ -4939,39 +4735,29 @@ mod tests { nulls_first: true, }, ), - PhysicalSortExpr::new_default(col("id", &schema)?), + PhysicalSortExpr::new_default(id), ] .into(); - - // Nine rows remain tied in three groups; seven rows are already resolved. - // The prefix and numeric suffix still tie for some rows, so id breaks ties. let expected = arrow::compute::take_record_batch( &batch, - &UInt32Array::from(vec![ - 3, 5, 9, 1, 7, 0, 4, 8, 10, 11, 12, 13, 14, 15, 6, 2, - ]), + &UInt32Array::from(vec![3, 5, 7, 1, 2, 6, 4, 0]), )?; - // Single batch, concatenated batches, and separate runs followed by a merge. - for (input_batches, threshold) in [(1, 1), (2, usize::MAX), (2, 1)] { + for input_batches in [1, 2] { for enabled in [false, true] { suffix_rows.store(0, AtomicOrdering::Relaxed); - let output = staged_sort_test_output( - &batch, - &ordering, - enabled, - input_batches, - threshold, - ) - .await?; - + id_rows.store(0, AtomicOrdering::Relaxed); + let output = + staged_sort_test_output(&batch, &ordering, enabled, input_batches, 1) + .await?; assert_eq!( concat_batches(&schema, &output)?, expected, - "enabled={enabled}, input_batches={input_batches}, threshold={threshold}" + "enabled={enabled}, input_batches={input_batches}" ); - if enabled && (input_batches == 1 || threshold == usize::MAX) { - assert_eq!(suffix_rows.load(AtomicOrdering::Relaxed), 9); + if enabled && input_batches == 1 { + assert_eq!(suffix_rows.load(AtomicOrdering::Relaxed), 6); + assert_eq!(id_rows.load(AtomicOrdering::Relaxed), 4); } } } @@ -4980,12 +4766,11 @@ mod tests { #[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 batch = record_batch!( + ("prefix", Int32, [3, 1, 6, 2, 5, 4]), + ("invalid", Utf8, vec!["invalid"; 6]) + )?; + let schema = batch.schema(); let ordering = [ PhysicalSortExpr::new_default(col("prefix", &schema)?), PhysicalSortExpr::new_default(Arc::new(CastExpr::new( @@ -4995,99 +4780,34 @@ mod tests { ))), ] .into(); - struct Case<'a> { - name: &'a str, - prefix: &'a [i32; 6], - enabled: bool, - input_batches: usize, - threshold: usize, - expect_error: bool, - } - let distinct = [3, 1, 6, 2, 5, 4]; - let all_tied = [1; 6]; - let cases = [ - Case { - name: "disabled evaluates eagerly", - prefix: &distinct, - enabled: false, - input_batches: 1, - threshold: 1, - expect_error: true, - }, - Case { - name: "distinct rows skip suffix", - prefix: &distinct, - enabled: true, - input_batches: 1, - threshold: 1, - expect_error: false, - }, - Case { - name: "ties evaluate suffix", - prefix: &all_tied, - enabled: true, - input_batches: 1, - threshold: 1, - expect_error: true, - }, - Case { - name: "concatenated rows skip suffix", - prefix: &distinct, - enabled: true, - input_batches: 2, - threshold: usize::MAX, - expect_error: false, - }, - Case { - name: "merge reevaluates suffix", - prefix: &distinct, - enabled: true, - input_batches: 2, - threshold: 1, - expect_error: true, - }, - ]; + let eager = staged_sort_test_output(&batch, &ordering, false, 1, 1).await; + let skipped = staged_sort_test_output(&batch, &ordering, true, 1, 1).await?; + assert_eq!( + concat_batches(&schema, &skipped)?.column(0).as_ref(), + &Int32Array::from_iter_values(1..=6) + ); - for Case { - name, - prefix, - enabled, - input_batches, - threshold, - expect_error, - } in cases - { - let batch = RecordBatch::try_new( - Arc::clone(&schema), - vec![ - Arc::new(Int32Array::from(prefix.to_vec())), - Arc::new(StringArray::from(vec!["invalid"; 6])), - ], - )?; - let result = staged_sort_test_output( - &batch, - &ordering, - enabled, - input_batches, - threshold, - ) - .await; + let all_tied = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int32Array::from(vec![1; 6])), + Arc::clone(batch.column(1)), + ], + )?; + let tied = staged_sort_test_output(&all_tied, &ordering, true, 1, 1).await; + let merged = staged_sort_test_output(&batch, &ordering, true, 2, 1).await; - if expect_error { - let error = result.unwrap_err().strip_backtrace(); - assert!( - error.contains("Cannot cast string 'invalid'"), - "{name}: {error}" - ); - } else { - let output = concat_batches(&schema, &result?)?; - assert_eq!( - output.column(0).as_ref(), - &Int32Array::from_iter_values(1..=6), - "{name}" - ); - } + for (name, result) in [ + ("disabled evaluates eagerly", eager), + ("ties evaluate suffix", tied), + ("merge reevaluates suffix", merged), + ] { + let error = result.unwrap_err().strip_backtrace(); + assert!( + error.contains("Cannot cast string 'invalid'"), + "{name}: {error}" + ); } Ok(()) } @@ -5096,72 +4816,12 @@ mod tests { async fn test_staged_sort_volatile_evaluation() -> Result<()> { use datafusion_expr::Volatility; - let schema = Arc::new(Schema::new(vec![ - Field::new("prefix", DataType::Int32, false), - Field::new("id", DataType::Int32, false), - ])); - struct Case<'a> { - name: &'a str, - prefix: &'a [i32; 6], - enabled: bool, - input_batches: usize, - threshold: usize, - expected_rows: usize, - } - - let distinct = [3, 1, 6, 2, 5, 4]; - let partly_tied = [2, 1, 1, 3, 4, 4]; - let cases = [ - Case { - name: "distinct enabled", - prefix: &distinct, - enabled: true, - input_batches: 1, - threshold: 1, - expected_rows: 6, - }, - Case { - name: "partly tied disabled", - prefix: &partly_tied, - enabled: false, - input_batches: 1, - threshold: 1, - expected_rows: 6, - }, - Case { - name: "partly tied enabled", - prefix: &partly_tied, - enabled: true, - input_batches: 1, - threshold: 1, - expected_rows: 6, - }, - Case { - name: "merge reevaluates", - prefix: &partly_tied, - enabled: true, - input_batches: 2, - threshold: 1, - expected_rows: 12, - }, - ]; - - for Case { - name, - prefix, - enabled, - input_batches, - threshold, - expected_rows, - } in cases - { - let batch = RecordBatch::try_new( - Arc::clone(&schema), - vec![ - Arc::new(Int32Array::from(prefix.to_vec())), - Arc::new(Int32Array::from_iter_values(0..6)), - ], - )?; + let batch = record_batch!( + ("prefix", Int32, [2, 1, 1, 3, 4, 4]), + ("id", Int32, [0, 1, 2, 3, 4, 5]) + )?; + let schema = batch.schema(); + for enabled in [false, true] { let (volatile, rows) = counting_identity( "count_evaluations", col("id", &schema)?, @@ -5175,9 +4835,8 @@ mod tests { ] .into(); - staged_sort_test_output(&batch, &ordering, enabled, input_batches, threshold) - .await?; - assert_eq!(rows.load(AtomicOrdering::Relaxed), expected_rows, "{name}"); + staged_sort_test_output(&batch, &ordering, enabled, 1, 1).await?; + assert_eq!(rows.load(AtomicOrdering::Relaxed), 6, "enabled={enabled}"); } Ok(()) }