feat: stage ORDER BY expression evaluation

datafusion/common/src/config.rs

@@ -1001,6 +1001,13 @@ 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.
+        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.
         ///

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,8 @@ 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,
 
     // ========================================================================
     // STATE BUFFERS:
@@ -325,9 +332,22 @@ impl ExternalSorter {
             batch_size,
             sort_spill_reservation_bytes,
             sort_in_place_threshold_bytes,
+            enable_staged_sort: false,
+            sort_key_group_size: 5,
         })
     }
 
+    fn with_staged_sort(mut self, enabled: bool, group_size: usize) -> Result<Self> {
+        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
@@ -648,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
@@ -657,7 +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)?;
+            let sorted_stream = self.sort_batch_stream(batch, reservation, staged)?;
             return Ok(self.observe_if_output(sorted_stream, is_output_stream));
         }
 
@@ -674,7 +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)?;
+            let sorted_stream = self.sort_batch_stream(batch, reservation, staged)?;
             return Ok(self.observe_if_output(sorted_stream, is_output_stream));
         }
 
@@ -696,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)?;
+                let input = self.sort_batch_stream(batch, reservation, staged)?;
                 Ok(spawn_buffered(input, 1))
             })
             .collect::<Result<_>>()?;
@@ -781,6 +803,7 @@ impl ExternalSorter {
         &self,
         batch: RecordBatch,
         reservation: MemoryReservation,
+        staged: bool,
     ) -> Result<SendableRecordBatchStream> {
         assert_eq!(
             get_reserved_bytes_for_record_batch(&batch)?,
@@ -790,13 +813,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 +874,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::<Column>() && !key.expr.is::<Literal>())
+    }
+
     /// 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 +1001,141 @@ impl Debug for ExternalSorter {
     }
 }
 
+fn staged_sort_batch_chunked(
+    batch: &RecordBatch,
+    expressions: &LexOrdering,
+    batch_size: usize,
+    group_size: usize,
+) -> Result<Vec<RecordBatch>> {
+    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<UInt32Array> {
+    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<u32>,
+    unresolved: &[Range<usize>],
+    first_stage: bool,
+    find_ties: bool,
+) -> Result<Vec<Range<usize>>> {
+    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::<Result<Vec<_>>>()?;
+
+    // Key arrays are already evaluated above. This comparator finds ties that
+    // need the next stage.
+    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::<Vec<_>>();
+        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`.
+fn append_tied_ranges(
+    comparator: &LexicographicalComparator,
+    order: &UInt32Array,
+    selected_offset: usize,
+    output_offset: usize,
+    ties: &mut Vec<Range<usize>>,
+) {
+    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 +1696,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(),

datafusion/physical-plan/src/sorts/stream.rs

@@ -288,8 +288,9 @@ impl<T: CursorArray> PartitionedStream for FieldCursorStream<T> {
 /// 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,11 @@ impl IncrementalSortIterator {
             indices: None,
         }
     }
+
+    pub(crate) fn with_sorted_indices(mut self, indices: UInt32Array) -> Self {
+        self.indices = Some(indices);
+        self
+    }
 }
 
 impl Iterator for IncrementalSortIterator {

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 5
 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 <https://github.com/apache/datafusion/issues/22710> 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 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.