Make limit pushdown work for SortPreservingMergeExec (apache#17893) (#361)

Author: Dandandan

datafusion/physical-optimizer/src/limit_pushdown_past_window.rs

@@ -23,6 +23,7 @@ use datafusion_expr::{WindowFrameBound, WindowFrameUnits};
 use datafusion_physical_plan::execution_plan::CardinalityEffect;
 use datafusion_physical_plan::limit::GlobalLimitExec;
 use datafusion_physical_plan::sorts::sort::SortExec;
+use datafusion_physical_plan::sorts::sort_preserving_merge::SortPreservingMergeExec;
 use datafusion_physical_plan::windows::BoundedWindowAggExec;
 use datafusion_physical_plan::ExecutionPlan;
 use std::cmp;
@@ -91,6 +92,22 @@ impl PhysicalOptimizerRule for LimitPushPastWindows {
                 return Ok(Transformed::no(node));
             }
 
+            // Apply the limit if we hit a sortpreservingmerge node
+            if let Some(spm) = node.as_any().downcast_ref::<SortPreservingMergeExec>() {
+                let latest = latest_limit.take();
+                let Some(fetch) = latest else {
+                    latest_max = 0;
+                    return Ok(Transformed::no(node));
+                };
+                let fetch = match spm.fetch() {
+                    None => fetch + latest_max,
+                    Some(existing) => cmp::min(existing, fetch + latest_max),
+                };
+                let spm: Arc<dyn ExecutionPlan> = spm.with_fetch(Some(fetch)).unwrap();
+                latest_max = 0;
+                return Ok(Transformed::complete(spm));
+            }
+
             // Apply the limit if we hit a sort node
             if let Some(sort) = node.as_any().downcast_ref::<SortExec>() {
                 let latest = latest_limit.take();