refine

Author: xudong963

data/date=2025-03-01/TCAqpuqxDrvx1JUr.parquet

@@ -949,14 +949,14 @@ fn add_spm_on_top(input: DistributionContext) -> DistributionContext {
         let should_preserve_ordering = input.plan.output_ordering().is_some();
 
         let new_plan = if should_preserve_ordering {
-            SortPreservingMergeExec::new(
+            Arc::new(SortPreservingMergeExec::new(
                 input
                     .plan
                     .output_ordering()
                     .unwrap_or(&LexOrdering::default())
                     .clone(),
                 Arc::clone(&input.plan),
-            ) as _
+            )) as _
         } else {
             Arc::new(CoalescePartitionsExec::new(Arc::clone(&input.plan))) as _
         };

data/date=2025-03-02/TCAqpuqxDrvx1JUr.parquet

@@ -394,7 +394,7 @@ pub fn parallelize_sorts(
             SortPreservingMergeExec::new(sort_exprs, Arc::clone(&requirements.plan));
         Ok(Transformed::yes(
             PlanWithCorrespondingCoalescePartitions::new(
-                spm.with_fetch(fetch).unwrap(),
+                Arc::new(spm.with_fetch(fetch)),
                 false,
                 vec![requirements],
             ),
@@ -509,9 +509,10 @@ fn analyze_immediate_sort_removal(
                 // Replace the sort with a sort-preserving merge:
                 let expr = LexOrdering::new(sort_exec.expr().to_vec());
 
-                SortPreservingMergeExec::new(expr, Arc::clone(sort_input))
-                    .with_fetch(sort_exec.fetch())
-                    .unwrap()
+                Arc::new(
+                    SortPreservingMergeExec::new(expr, Arc::clone(sort_input))
+                        .with_fetch(sort_exec.fetch()),
+                ) as _
             } else {
                 // Remove the sort:
                 node.children = node.children.swap_remove(0).children;
@@ -737,9 +738,10 @@ fn remove_corresponding_sort_from_sub_plan(
         let plan = Arc::clone(&node.plan);
         let fetch = plan.fetch();
         let plan = if let Some(ordering) = plan.output_ordering() {
-            SortPreservingMergeExec::new(LexOrdering::new(ordering.to_vec()), plan)
-                .with_fetch(fetch)
-                .unwrap()
+            Arc::new(
+                SortPreservingMergeExec::new(LexOrdering::new(ordering.to_vec()), plan)
+                    .with_fetch(fetch),
+            ) as _
         } else {
             Arc::new(CoalescePartitionsExec::new(plan)) as _
         };

data/date=2025-03-03/TCAqpuqxDrvx1JUr.parquet

@@ -140,9 +140,10 @@ fn plan_with_order_preserving_variants(
         if let Some(ordering) = child.output_ordering() {
             // When the input of a `CoalescePartitionsExec` has an ordering,
             // replace it with a `SortPreservingMergeExec` if appropriate:
-            let spm = SortPreservingMergeExec::new(ordering.clone(), Arc::clone(child))
-                .with_fetch(fetch)
-                .unwrap();
+            let spm = Arc::new(
+                SortPreservingMergeExec::new(ordering.clone(), Arc::clone(child))
+                    .with_fetch(fetch),
+            );
             sort_input.plan = spm;
             sort_input.children[0].data = true;
             return Ok(sort_input);

data/date=2025-03-04/TCAqpuqxDrvx1JUr.parquet

@@ -26,7 +26,6 @@ use crate::{
     DisplayAs, DisplayFormatType, Distribution, ExecutionPlan, ExecutionPlanProperties,
     Partitioning, PlanProperties, SendableRecordBatchStream, Statistics,
 };
-use itertools::Itertools;
 use std::any::Any;
 use std::sync::Arc;
 
@@ -38,8 +37,6 @@ use datafusion_physical_expr_common::sort_expr::{LexOrdering, LexRequirement};
 
 use crate::sorts::progressive_eval::ProgressiveEvalExec;
 use crate::statistics::MinMaxStatistics;
-use crate::stream::RecordBatchStreamAdapter;
-use futures::StreamExt;
 use log::{debug, trace};
 
 /// Sort preserving merge execution plan
@@ -102,14 +99,13 @@ pub struct SortPreservingMergeExec {
     ///
     /// See [`Self::with_round_robin_repartition`] for more information.
     enable_round_robin_repartition: bool,
+    ///
+    progressive_eval_exec: Option<Arc<ProgressiveEvalExec>>,
 }
 
 impl SortPreservingMergeExec {
     /// Create a new sort execution plan
-    pub fn new(
-        expr: LexOrdering,
-        input: Arc<dyn ExecutionPlan>,
-    ) -> Arc<dyn ExecutionPlan> {
+    pub fn new(expr: LexOrdering, input: Arc<dyn ExecutionPlan>) -> Self {
         let cache = Self::compute_properties(&input, expr.clone());
 
         // Todo: check if partition statistic is accurate
@@ -136,24 +132,23 @@ impl SortPreservingMergeExec {
             .filter(|groups| {
                 groups.len() < input.properties().partitioning.partition_count()
             });
+        let mut progressive_eval_exec = None;
         if let Some(partition_groups) = partition_groups {
-            // Return ProgressiveEvalExec when partition_groups exists
-            Arc::new(ProgressiveEvalExec::new(
-                input,
+            progressive_eval_exec = Some(Arc::new(ProgressiveEvalExec::new(
+                input.clone(),
                 None,
                 None,
                 partition_groups,
-            ))
-        } else {
-            // Return SortPreservingMergeExec otherwise
-            Arc::new(Self {
-                input,
-                expr,
-                metrics: ExecutionPlanMetricsSet::new(),
-                fetch: None,
-                cache,
-                enable_round_robin_repartition: true,
-            })
+            )));
+        }
+        Self {
+            input,
+            expr,
+            metrics: ExecutionPlanMetricsSet::new(),
+            fetch: None,
+            cache,
+            enable_round_robin_repartition: true,
+            progressive_eval_exec,
         }
     }
 
@@ -221,11 +216,14 @@ impl DisplayAs for SortPreservingMergeExec {
     ) -> std::fmt::Result {
         match t {
             DisplayFormatType::Default | DisplayFormatType::Verbose => {
-                write!(f, "SortPreservingMergeExec: [{}]", self.expr)?;
-                if let Some(fetch) = self.fetch {
-                    write!(f, ", fetch={fetch}")?;
-                };
-
+                if let Some(progressive_eval) = &self.progressive_eval_exec {
+                    progressive_eval.fmt_as(t, f)?;
+                } else {
+                    write!(f, "SortPreservingMergeExec: [{}]", self.expr)?;
+                    if let Some(fetch) = self.fetch {
+                        write!(f, ", fetch={fetch}")?;
+                    };
+                }
                 Ok(())
             }
             DisplayFormatType::TreeRender => {
@@ -272,6 +270,7 @@ impl ExecutionPlan for SortPreservingMergeExec {
             fetch: limit,
             cache: self.cache.clone(),
             enable_round_robin_repartition: true,
+            progressive_eval_exec: self.progressive_eval_exec.clone(),
         }))
     }
 
@@ -299,18 +298,20 @@ impl ExecutionPlan for SortPreservingMergeExec {
         self: Arc<Self>,
         children: Vec<Arc<dyn ExecutionPlan>>,
     ) -> Result<Arc<dyn ExecutionPlan>> {
-        Ok(
+        Ok(Arc::new(
             SortPreservingMergeExec::new(self.expr.clone(), Arc::clone(&children[0]))
-                .with_fetch(self.fetch)
-                .unwrap(),
-        )
+                .with_fetch(self.fetch),
+        ))
     }
 
     fn execute(
         &self,
         partition: usize,
         context: Arc<TaskContext>,
     ) -> Result<SendableRecordBatchStream> {
+        if let Some(progressive_eval) = &self.progressive_eval_exec {
+            return progressive_eval.execute(partition, context);
+        }
         trace!(
             "Start SortPreservingMergeExec::execute for partition: {}",
             partition
@@ -418,11 +419,13 @@ impl ExecutionPlan for SortPreservingMergeExec {
             });
         }
 
-        Ok(SortPreservingMergeExec::new(
-            updated_exprs,
-            make_with_child(projection, self.input())?,
-        )
-        .with_fetch(self.fetch()))
+        Ok(Some(Arc::new(
+            SortPreservingMergeExec::new(
+                updated_exprs,
+                make_with_child(projection, self.input())?,
+            )
+            .with_fetch(self.fetch()),
+        )))
     }
 }
 

datafusion/physical-optimizer/src/enforce_distribution.rs

@@ -966,9 +966,9 @@ impl AsExecutionPlan for protobuf::PhysicalPlanNode {
                 } else {
                     Some(sort.fetch as usize)
                 };
-                Ok(SortPreservingMergeExec::new(exprs, input)
-                    .with_fetch(fetch)
-                    .unwrap())
+                Ok(Arc::new(
+                    SortPreservingMergeExec::new(exprs, input).with_fetch(fetch),
+                ))
             }
             PhysicalPlanType::Extension(extension) => {
                 let inputs: Vec<Arc<dyn ExecutionPlan>> = extension