Simplify streaming_merge function parameters

datafusion/core/tests/fuzz_cases/sort_preserving_repartition_fuzz.rs

@@ -29,7 +29,7 @@ mod sp_repartition_fuzz_tests {
         metrics::{BaselineMetrics, ExecutionPlanMetricsSet},
         repartition::RepartitionExec,
         sorts::sort_preserving_merge::SortPreservingMergeExec,
-        sorts::streaming_merge::streaming_merge,
+        sorts::streaming_merge::{streaming_merge, StreamingMergeConfig},
         stream::RecordBatchStreamAdapter,
         ExecutionPlan, Partitioning,
     };
@@ -246,15 +246,15 @@ mod sp_repartition_fuzz_tests {
                 MemoryConsumer::new("test".to_string()).register(context.memory_pool());
 
             // Internally SortPreservingMergeExec uses this function for merging.
-            let res = streaming_merge(
+            let res = streaming_merge(StreamingMergeConfig {
                 streams,
                 schema,
-                &exprs,
-                BaselineMetrics::new(&ExecutionPlanMetricsSet::new(), 0),
-                1,
-                None,
-                mem_reservation,
-            )?;
+                expressions: &exprs,
+                metrics: BaselineMetrics::new(&ExecutionPlanMetricsSet::new(), 0),
+                batch_size: 1,
+                fetch: None,
+                reservation: mem_reservation,
+            })?;
             let res = collect(res).await?;
             // Contains the merged result.
             let res = concat_batches(&res[0].schema(), &res)?;

datafusion/physical-plan/src/aggregates/row_hash.rs

@@ -29,7 +29,7 @@ use crate::aggregates::{
 };
 use crate::metrics::{BaselineMetrics, MetricBuilder, RecordOutput};
 use crate::sorts::sort::sort_batch;
-use crate::sorts::streaming_merge;
+use crate::sorts::streaming_merge::{streaming_merge, StreamingMergeConfig};
 use crate::spill::{read_spill_as_stream, spill_record_batch_by_size};
 use crate::stream::RecordBatchStreamAdapter;
 use crate::{aggregates, metrics, ExecutionPlan, PhysicalExpr};
@@ -1001,15 +1001,15 @@ impl GroupedHashAggregateStream {
             streams.push(stream);
         }
         self.spill_state.is_stream_merging = true;
-        self.input = streaming_merge(
+        self.input = streaming_merge(StreamingMergeConfig {
             streams,
             schema,
-            &self.spill_state.spill_expr,
-            self.baseline_metrics.clone(),
-            self.batch_size,
-            None,
-            self.reservation.new_empty(),
-        )?;
+            expressions: &self.spill_state.spill_expr,
+            metrics: self.baseline_metrics.clone(),
+            batch_size: self.batch_size,
+            fetch: None,
+            reservation: self.reservation.new_empty(),
+        })?;
         self.input_done = false;
         self.group_ordering = GroupOrdering::Full(GroupOrderingFull::new());
         Ok(())

datafusion/physical-plan/src/repartition/mod.rs

@@ -34,7 +34,7 @@ use crate::metrics::BaselineMetrics;
 use crate::repartition::distributor_channels::{
     channels, partition_aware_channels, DistributionReceiver, DistributionSender,
 };
-use crate::sorts::streaming_merge;
+use crate::sorts::streaming_merge::{streaming_merge, StreamingMergeConfig};
 use crate::stream::RecordBatchStreamAdapter;
 use crate::{DisplayFormatType, ExecutionPlan, Partitioning, PlanProperties, Statistics};
 
@@ -637,15 +637,15 @@ impl ExecutionPlan for RepartitionExec {
                 let merge_reservation =
                     MemoryConsumer::new(format!("{}[Merge {partition}]", name))
                         .register(context.memory_pool());
-                streaming_merge(
-                    input_streams,
-                    schema_captured,
-                    &sort_exprs,
-                    BaselineMetrics::new(&metrics, partition),
-                    context.session_config().batch_size(),
+                streaming_merge(StreamingMergeConfig {
+                    streams: input_streams,
+                    schema: schema_captured,
+                    expressions: &sort_exprs,
+                    metrics: BaselineMetrics::new(&metrics, partition),
+                    batch_size: context.session_config().batch_size(),
                     fetch,
-                    merge_reservation,
-                )
+                    reservation: merge_reservation,
+                })
             } else {
                 Ok(Box::pin(RepartitionStream {
                     num_input_partitions,

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

@@ -28,4 +28,3 @@ mod stream;
 pub mod streaming_merge;
 
 pub use index::RowIndex;
-pub(crate) use streaming_merge::streaming_merge;

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

@@ -30,7 +30,7 @@ use crate::limit::LimitStream;
 use crate::metrics::{
     BaselineMetrics, Count, ExecutionPlanMetricsSet, MetricBuilder, MetricsSet,
 };
-use crate::sorts::streaming_merge::streaming_merge;
+use crate::sorts::streaming_merge::{streaming_merge, StreamingMergeConfig};
 use crate::spill::{read_spill_as_stream, spill_record_batches};
 use crate::stream::RecordBatchStreamAdapter;
 use crate::topk::TopK;
@@ -342,15 +342,15 @@ impl ExternalSorter {
                 streams.push(stream);
             }
 
-            streaming_merge(
+            streaming_merge(StreamingMergeConfig {
                 streams,
-                Arc::clone(&self.schema),
-                &self.expr,
-                self.metrics.baseline.clone(),
-                self.batch_size,
-                self.fetch,
-                self.reservation.new_empty(),
-            )
+                schema: Arc::clone(&self.schema),
+                expressions: &self.expr,
+                metrics: self.metrics.baseline.clone(),
+                batch_size: self.batch_size,
+                fetch: self.fetch,
+                reservation: self.reservation.new_empty(),
+            })
         } else {
             self.in_mem_sort_stream(self.metrics.baseline.clone())
         }
@@ -534,15 +534,15 @@ impl ExternalSorter {
             })
             .collect::<Result<_>>()?;
 
-        streaming_merge(
+        streaming_merge(StreamingMergeConfig {
             streams,
-            Arc::clone(&self.schema),
-            &self.expr,
+            schema: Arc::clone(&self.schema),
+            expressions: &self.expr,
             metrics,
-            self.batch_size,
-            self.fetch,
-            self.merge_reservation.new_empty(),
-        )
+            batch_size: self.batch_size,
+            fetch: self.fetch,
+            reservation: self.merge_reservation.new_empty(),
+        })
     }
 
     /// Sorts a single `RecordBatch` into a single stream.

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

@@ -24,7 +24,7 @@ use crate::common::spawn_buffered;
 use crate::expressions::PhysicalSortExpr;
 use crate::limit::LimitStream;
 use crate::metrics::{BaselineMetrics, ExecutionPlanMetricsSet, MetricsSet};
-use crate::sorts::streaming_merge;
+use crate::sorts::streaming_merge::{streaming_merge, StreamingMergeConfig};
 use crate::{
     DisplayAs, DisplayFormatType, Distribution, ExecutionPlan, ExecutionPlanProperties,
     Partitioning, PlanProperties, SendableRecordBatchStream, Statistics,
@@ -273,15 +273,15 @@ impl ExecutionPlan for SortPreservingMergeExec {
 
                 debug!("Done setting up sender-receiver for SortPreservingMergeExec::execute");
 
-                let result = streaming_merge(
-                    receivers,
+                let result = streaming_merge(StreamingMergeConfig {
+                    streams: receivers,
                     schema,
-                    &self.expr,
-                    BaselineMetrics::new(&self.metrics, partition),
-                    context.session_config().batch_size(),
-                    self.fetch,
+                    expressions: &self.expr,
+                    metrics: BaselineMetrics::new(&self.metrics, partition),
+                    batch_size: context.session_config().batch_size(),
+                    fetch: self.fetch,
                     reservation,
-                )?;
+                })?;
 
                 debug!("Got stream result from SortPreservingMergeStream::new_from_receivers");
 
@@ -960,15 +960,15 @@ mod tests {
             MemoryConsumer::new("test").register(&task_ctx.runtime_env().memory_pool);
 
         let fetch = None;
-        let merge_stream = streaming_merge(
+        let merge_stream = streaming_merge(StreamingMergeConfig {
             streams,
-            batches.schema(),
-            sort.as_slice(),
-            BaselineMetrics::new(&metrics, 0),
-            task_ctx.session_config().batch_size(),
+            schema: batches.schema(),
+            expressions: sort.as_slice(),
+            metrics: BaselineMetrics::new(&metrics, 0),
+            batch_size: task_ctx.session_config().batch_size(),
             fetch,
             reservation,
-        )
+        })
         .unwrap();
 
         let mut merged = common::collect(merge_stream).await.unwrap();

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

@@ -36,62 +36,67 @@ macro_rules! primitive_merge_helper {
 }
 
 macro_rules! merge_helper {
-    ($t:ty, $sort:ident, $streams:ident, $schema:ident, $tracking_metrics:ident, $batch_size:ident, $fetch:ident, $reservation:ident) => {{
-        let streams = FieldCursorStream::<$t>::new($sort, $streams);
+    ($t:ty, $sort:ident, $config:ident) => {{
+        let streams = FieldCursorStream::<$t>::new($sort, $config.streams);
         return Ok(Box::pin(SortPreservingMergeStream::new(
             Box::new(streams),
-            $schema,
-            $tracking_metrics,
-            $batch_size,
-            $fetch,
-            $reservation,
+            $config.schema,
+            $config.metrics,
+            $config.batch_size,
+            $config.fetch,
+            $config.reservation,
         )));
     }};
 }
 
+/// Configuration parameters to initialize a `SortPreservingMergeStream`
+pub struct StreamingMergeConfig<'a> {
+    pub streams: Vec<SendableRecordBatchStream>,
+    pub schema: SchemaRef,
+    pub expressions: &'a [PhysicalSortExpr],
+    pub metrics: BaselineMetrics,
+    pub batch_size: usize,
+    pub fetch: Option<usize>,
+    pub reservation: MemoryReservation,
+}
+
 /// Perform a streaming merge of [`SendableRecordBatchStream`] based on provided sort expressions
 /// while preserving order.
 pub fn streaming_merge(
-    streams: Vec<SendableRecordBatchStream>,
-    schema: SchemaRef,
-    expressions: &[PhysicalSortExpr],
-    metrics: BaselineMetrics,
-    batch_size: usize,
-    fetch: Option<usize>,
-    reservation: MemoryReservation,
+    config: StreamingMergeConfig,
 ) -> Result<SendableRecordBatchStream> {
     // If there are no sort expressions, preserving the order
     // doesn't mean anything (and result in infinite loops)
-    if expressions.is_empty() {
+    if config.expressions.is_empty() {
         return internal_err!("Sort expressions cannot be empty for streaming merge");
     }
     // Special case single column comparisons with optimized cursor implementations
-    if expressions.len() == 1 {
-        let sort = expressions[0].clone();
-        let data_type = sort.expr.data_type(schema.as_ref())?;
+    if config.expressions.len() == 1 {
+        let sort = config.expressions[0].clone();
+        let data_type = sort.expr.data_type(config.schema.as_ref())?;
         downcast_primitive! {
-            data_type => (primitive_merge_helper, sort, streams, schema, metrics, batch_size, fetch, reservation),
-            DataType::Utf8 => merge_helper!(StringArray, sort, streams, schema, metrics, batch_size, fetch, reservation)
-            DataType::LargeUtf8 => merge_helper!(LargeStringArray, sort, streams, schema, metrics, batch_size, fetch, reservation)
-            DataType::Binary => merge_helper!(BinaryArray, sort, streams, schema, metrics, batch_size, fetch, reservation)
-            DataType::LargeBinary => merge_helper!(LargeBinaryArray, sort, streams, schema, metrics, batch_size, fetch, reservation)
+            data_type => (primitive_merge_helper, sort, config),
+            DataType::Utf8 => merge_helper!(StringArray, sort, config)
+            DataType::LargeUtf8 => merge_helper!(LargeStringArray, sort, config)
+            DataType::Binary => merge_helper!(BinaryArray, sort, config)
+            DataType::LargeBinary => merge_helper!(LargeBinaryArray, sort, config)
             _ => {}
         }
     }
 
     let streams = RowCursorStream::try_new(
-        schema.as_ref(),
-        expressions,
-        streams,
-        reservation.new_empty(),
+        config.schema.as_ref(),
+        config.expressions,
+        config.streams,
+        config.reservation.new_empty(),
     )?;
 
     Ok(Box::pin(SortPreservingMergeStream::new(
         Box::new(streams),
-        schema,
-        metrics,
-        batch_size,
-        fetch,
-        reservation,
+        config.schema,
+        config.metrics,
+        config.batch_size,
+        config.fetch,
+        config.reservation,
     )))
 }