From 0e8776e4895857d565af6fce218b31105622a58b Mon Sep 17 00:00:00 2001 From: Raz Luvaton <16746759+rluvaton@users.noreply.github.com> Date: Wed, 15 Jul 2026 22:20:40 +0300 Subject: [PATCH 1/3] fix(sort): multi level merge sort should respect record batch size fetch and ordering --- .../src/sorts/streaming_merge.rs | 274 ++++++++++++++++++ 1 file changed, 274 insertions(+) diff --git a/datafusion/physical-plan/src/sorts/streaming_merge.rs b/datafusion/physical-plan/src/sorts/streaming_merge.rs index e96138ef1306c..4bacb0c552dbf 100644 --- a/datafusion/physical-plan/src/sorts/streaming_merge.rs +++ b/datafusion/physical-plan/src/sorts/streaming_merge.rs @@ -267,3 +267,277 @@ impl<'a> StreamingMergeBuilder<'a> { .into_stream()) } } + +#[cfg(test)] +mod tests { + use std::{ops::Deref, sync::Arc}; + + use arrow::{ + array::{ArrayRef, AsArray, Int32Array, RecordBatch}, + datatypes::Int32Type, + }; + use arrow_schema::{Field, Fields, Schema, SchemaRef, SortOptions}; + use datafusion_common::Result; + use datafusion_execution::{ + memory_pool::{ + GreedyMemoryPool, MemoryConsumer, MemoryPool, UnboundedMemoryPool, + }, + runtime_env::RuntimeEnv, + }; + use datafusion_physical_expr::{LexOrdering, PhysicalSortExpr, expressions::col}; + use datafusion_physical_expr_common::metrics::{ + BaselineMetrics, ExecutionPlanMetricsSet, SpillMetrics, + }; + use futures::TryStreamExt; + use itertools::Itertools; + + use crate::{ + SpillManager, + sorts::streaming_merge::{SortedSpillFile, StreamingMergeBuilder}, + }; + + fn create_batches( + arrays: Vec<(&str, ArrayRef)>, + batch_size: usize, + ) -> Vec { + let fields = arrays + .iter() + .map(|(name, arr)| { + Arc::new(Field::new( + *name, + arr.data_type().clone(), + arr.is_nullable(), + )) + }) + .collect::(); + let schema = Arc::new(Schema::new(fields)); + let batch = RecordBatch::try_new( + schema, + arrays.into_iter().map(|(_name, arr)| arr).collect(), + ) + .unwrap(); + + (0..batch.num_rows().div_ceil(batch_size)) + .map(|index| { + let offset = index * batch_size; + let left = batch.num_rows() - offset; + batch.slice(offset, left.min(batch_size)) + }) + .collect::>() + } + + fn create_builder<'a>( + schema: &SchemaRef, + merge_fan_in: Option, + ) -> StreamingMergeBuilder<'a> { + let runtime = Arc::new(RuntimeEnv::default()); + + if let Some(merge_fan_in) = merge_fan_in { + runtime + .disk_manager + .set_max_spill_merge_fan_in(merge_fan_in); + } + + let mem_pool: Arc = Arc::new(UnboundedMemoryPool::default()); + let spill_manager = SpillManager::new( + runtime, + SpillMetrics::new(&ExecutionPlanMetricsSet::new(), 0), + Arc::clone(schema), + ); + + StreamingMergeBuilder::new() + .with_spill_manager(spill_manager) + .with_reservation( + MemoryConsumer::new("merge stream mock memory").register(&mem_pool), + ) + .with_metrics(BaselineMetrics::new(&ExecutionPlanMetricsSet::default(), 0)) + .with_schema(Arc::clone(schema)) + } + + fn spill_batches( + spill_manager: &SpillManager, + batches_for_files: &[&[RecordBatch]], + ) -> Vec { + batches_for_files + .iter() + .map(|batches| { + let (file, max_record_batch_memory) = spill_manager + .spill_record_batch_iter_and_return_max_batch_memory( + batches.iter().map(Ok), + "spill", + ) + .unwrap() + .unwrap(); + + SortedSpillFile { + file, + max_record_batch_memory, + } + }) + .collect() + } + + #[tokio::test] + async fn multi_level_merge_sort_should_respect_limit_for_single_spill_file() { + let batch_size = 10; + let batches = create_batches( + vec![( + "a", + Arc::new((0..(batch_size * 2) as i32).collect::()), + )], + batch_size, + ); + + let schema = batches[0].schema(); + + let ordering = LexOrdering::new(vec![PhysicalSortExpr::new( + col("a", &schema).unwrap(), + // Match the existing sort order of the data + SortOptions::default().asc(), + )]) + .unwrap(); + + let fetch = (batch_size as f64 * 1.5) as usize; + assert_ne!(fetch % batch_size, 0); + + let merge_builder = create_builder(&schema, None) + .with_batch_size(batch_size) + .with_fetch(Some(fetch)) + .with_expressions(&ordering); + + let spilled_files = + spill_batches(merge_builder.spill_manager.as_ref().unwrap(), &[&batches]); + + let stream = merge_builder + .with_sorted_spill_files(spilled_files) + .build() + .unwrap(); + + let sorted: Result> = stream.try_collect().await; + let sorted = sorted.unwrap(); + + let output = arrow::compute::concat_batches(&schema, sorted.iter()).unwrap(); + + assert_eq!(output.num_rows(), fetch); + + let rows = output.column(0).as_primitive::(); + assert_eq!( + rows.values().deref(), + (0..(fetch as i32)).collect::>() + ); + } + + #[tokio::test] + async fn multi_level_merge_sort_should_be_stable_when_round_robin_tie_breaker_is_disabled() + { + let input_batch_size = 10; + let num_rows = input_batch_size * 6; + let batches = create_batches( + vec![ + // All the same value so it should match + ( + "a", + Arc::new((0..num_rows).map(|_| 0).collect::()), + ), + ( + "other", + Arc::new((0..num_rows as i32).collect::()), + ), + ], + input_batch_size, + ); + + let schema = batches[0].schema(); + + let files_batches = batches.chunks(2).collect::>(); + assert_eq!(files_batches.len(), 3); + + let ordering = LexOrdering::new(vec![PhysicalSortExpr::new_default( + col("a", &schema).unwrap(), + )]) + .unwrap(); + let output_batch_size = (input_batch_size as f64 * 1.5) as usize; + + let merge_builder = create_builder(&schema, Some(2)) + .with_batch_size(output_batch_size) + .with_fetch(None) + .with_expressions(&ordering) + .with_round_robin_tie_breaker(false); + + let spilled_files = spill_batches( + merge_builder.spill_manager.as_ref().unwrap(), + &files_batches, + ); + + let stream = merge_builder + .with_sorted_spill_files(spilled_files) + .build() + .unwrap(); + + let sorted: Result> = stream.try_collect().await; + let sorted = sorted.unwrap(); + + let output = arrow::compute::concat_batches(&schema, sorted.iter()).unwrap(); + + let other_rows = output.column(1).as_primitive::(); + assert_eq!( + other_rows.values().deref(), + (0..num_rows as i32).collect::>() + ); + } + + #[tokio::test] + async fn multi_level_merge_sort_should_respect_batch_size_for_single_spill_file() { + let input_batch_size = 10; + let num_rows = input_batch_size * 4; + let batches = create_batches( + vec![("a", Arc::new((0..num_rows as i32).collect::()))], + input_batch_size, + ); + + let schema = batches[0].schema(); + + let ordering = LexOrdering::new(vec![PhysicalSortExpr::new( + col("a", &schema).unwrap(), + // Match the existing sort order of the data + SortOptions::default().asc(), + )]) + .unwrap(); + let output_batch_size = (input_batch_size as f64 * 1.5) as usize; + + let merge_builder = create_builder(&schema, None) + .with_batch_size(output_batch_size) + .with_fetch(None) + .with_expressions(&ordering); + + let spilled_files = + spill_batches(merge_builder.spill_manager.as_ref().unwrap(), &[&batches]); + + let stream = merge_builder + .with_sorted_spill_files(spilled_files) + .build() + .unwrap(); + + let sorted: Result> = stream.try_collect().await; + let sorted = sorted.unwrap(); + + assert_eq!(sorted.iter().map(|a| a.num_rows()).sum::(), num_rows); + + for sorted_batch in sorted.iter().dropping_back(1) { + assert_eq!( + sorted_batch.num_rows(), + output_batch_size, + "batch size mismatch" + ); + } + + let last_batch_num_rows = sorted.last().unwrap().num_rows(); + assert!( + last_batch_num_rows <= output_batch_size, + "last batch size mismatch: {last_batch_num_rows} <= {output_batch_size}" + ); + } + + // TODO - single spill file does not respect batch size + // +} From a1344cccc36b19e86c77b662abe26548a6c73dd2 Mon Sep 17 00:00:00 2001 From: Raz Luvaton <16746759+rluvaton@users.noreply.github.com> Date: Thu, 16 Jul 2026 12:16:34 +0300 Subject: [PATCH 2/3] fix fetch --- .../src/sorts/multi_level_merge.rs | 65 +++++++++++++++---- 1 file changed, 52 insertions(+), 13 deletions(-) diff --git a/datafusion/physical-plan/src/sorts/multi_level_merge.rs b/datafusion/physical-plan/src/sorts/multi_level_merge.rs index 3ec52cc70c0a9..c6283af35d610 100644 --- a/datafusion/physical-plan/src/sorts/multi_level_merge.rs +++ b/datafusion/physical-plan/src/sorts/multi_level_merge.rs @@ -20,6 +20,7 @@ use crate::metrics::BaselineMetrics; use crate::{EmptyRecordBatchStream, SpillManager}; use arrow::array::RecordBatch; +use std::collections::VecDeque; use std::fmt::{Debug, Formatter}; use std::mem; use std::pin::Pin; @@ -56,6 +57,7 @@ use futures::{Stream, StreamExt}; /// /// - No: return that sorted stream as the final output stream /// +/// (the diagram below is for when round robin tie breaker is enabled - when disabled the order of the streams are kept) /// ```text /// Initial State: Multiple sorted streams + spill files /// ┌───────────┐ @@ -147,7 +149,7 @@ pub(crate) struct MultiLevelMergeBuilder { /// carry `batch_size`. A run re-spilled smaller to resolve skew carries its halved /// limit (see [`Self::split_spill_file_in_half`]). Tracking it here keeps this limit /// out of the public [`SortedSpillFile`], so no external caller has to set it. - sorted_spill_files: Vec<(SortedSpillFile, usize)>, + sorted_spill_files: VecDeque<(SortedSpillFile, usize)>, sorted_streams: Vec, expr: LexOrdering, metrics: BaselineMetrics, @@ -185,7 +187,7 @@ impl MultiLevelMergeBuilder { sorted_spill_files: sorted_spill_files .into_iter() .map(|file| (file, batch_size)) - .collect(), + .collect::>(), sorted_streams, expr, metrics, @@ -247,18 +249,47 @@ impl MultiLevelMergeBuilder { else { continue; }; - + // Add the spill file paired with the batch-size limit of the merge that // produced it: if that merge consumed a shrunk (skew-resolved) run, its // output was capped and this intermediate run is likewise capped, so a // later pass that re-merges it won't rebuild an oversized batch. - self.sorted_spill_files.push(( + let spill_file_to_add = ( SortedSpillFile { file: spill_file, max_record_batch_memory, }, batch_size_limit, - )); + ); + + // In round robin tie breaker the sort is not expected to be stable so we push to the back for performance reasons: + // having even spill files + // + // For example lets say we have 6 streams: a, b, c, d, e, f and the merge fan-in is 3: + // + // so we merge a, b and c and write to spill file abc' + // and put it in the end, so we now have: d, e, f, abc' + // now, we merge the next 3: d, e and f into spill file def' + // + // next we merge abc' and def' into one and return that. + // + // now, if we kept the file at the orignal order (= the start) we would merge like this: + // a, b, c and write to spill file abc' (like before) + // and put it in the start, so we now have: abc', d, e, f + // now, we merge abc', d and e and write into spill file: abcde' + // + // notice that the spill file we just wrote have data from 5 streams as opposed to before which was only 3, + // this mean that the spill file can be larger if each initial spill file have the same number of rows + if self.enable_round_robin_tie_breaker { + // Push back so we won't sort that file again right away + self.sorted_spill_files.push_back(spill_file_to_add); + } else { + // Add to front to keep the original order so the sort will be kept stable + // + // TODO - push back and keep track of original position to merge sort in the correct order of + // stream to benefit from the performance improvement described above + self.sorted_spill_files.push_front(spill_file_to_add) + } } } @@ -287,7 +318,8 @@ impl MultiLevelMergeBuilder { // Only single sorted spill file so return it (1, 0) => { - let (spill_file, batch_size) = self.sorted_spill_files.remove(0); + // TODO - check next level as well + let (spill_file, batch_size) = self.sorted_spill_files.remove(0).unwrap(); // Not reserving any memory for this disk as we are not holding it in memory let output_stream = self @@ -586,7 +618,7 @@ impl MultiLevelMergeBuilder { self.sorted_spill_files.swap(index, last); let (target, old_batch_size) = self .sorted_spill_files - .pop() + .pop_back() .expect("index is in bounds, so the vec is non-empty"); let old_max = target.max_record_batch_memory; @@ -644,17 +676,24 @@ impl MultiLevelMergeBuilder { // can't rebuild a full-size batch and reintroduce the skew. let new_batch_size_limit = (old_batch_size / 2).max(1); - // Push the re-spilled (smaller) file and swap it back into `index`, undoing - // the swap-to-back above so the order is preserved. - self.sorted_spill_files.push(( + self.sorted_spill_files.push_back(( SortedSpillFile { file, max_record_batch_memory: new_max, }, new_batch_size_limit, )); - let last = self.sorted_spill_files.len() - 1; - self.sorted_spill_files.swap(index, last); + + // Only for stable sort keep it in the current position, otherwise move it to the end + // to avoid as much as possible creating intermidate spill file with small batch size + // and only defer to last + if !self.enable_round_robin_tie_breaker { + // Push the re-spilled (smaller) file and swap it back into `index`, undoing + // the swap-to-back above so the order is preserved. + let last = self.sorted_spill_files.len() - 1; + self.sorted_spill_files.swap(index, last); + } + Ok(()) } @@ -663,7 +702,7 @@ impl MultiLevelMergeBuilder { &self, stream: SendableRecordBatchStream, ) -> SendableRecordBatchStream { - Box::pin(ObservedStream::new(stream, self.metrics.clone(), None)) + Box::pin(ObservedStream::new(stream, self.metrics.clone(), self.fetch)) } } From a5db2b3efa7882d15dad00c9df5b30fbb8c5538c Mon Sep 17 00:00:00 2001 From: Raz Luvaton <16746759+rluvaton@users.noreply.github.com> Date: Thu, 16 Jul 2026 13:16:13 +0300 Subject: [PATCH 3/3] rename --- datafusion/physical-plan/src/sorts/streaming_merge.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/datafusion/physical-plan/src/sorts/streaming_merge.rs b/datafusion/physical-plan/src/sorts/streaming_merge.rs index 4bacb0c552dbf..312f38e181fe8 100644 --- a/datafusion/physical-plan/src/sorts/streaming_merge.rs +++ b/datafusion/physical-plan/src/sorts/streaming_merge.rs @@ -428,7 +428,7 @@ mod tests { } #[tokio::test] - async fn multi_level_merge_sort_should_be_stable_when_round_robin_tie_breaker_is_disabled() + async fn multi_level_merge_sort_should_be_stable_and_have_multiple_rounds_when_round_robin_tie_breaker_is_disabled() { let input_batch_size = 10; let num_rows = input_batch_size * 6;