Fix spm + limit (apache#14569) v46

Author: xudong963

datafusion/core/src/physical_optimizer/enforce_sorting.rs

@@ -374,7 +374,13 @@ fn ensure_sorting(
     {
         // This `SortPreservingMergeExec` is unnecessary, input already has a
         // single partition.
-        let child_node = requirements.children.swap_remove(0);
+        // single partition and no fetch is required.
+        let mut child_node = requirements.children.swap_remove(0);
+        if let Some(fetch) = plan.fetch() {
+            // Add the limit exec if the spm has a fetch
+            child_node.plan =
+                Arc::new(LocalLimitExec::new(Arc::clone(&child_node.plan), fetch));
+        }
         return Ok(Transformed::yes(child_node));
     }
 
@@ -660,6 +666,7 @@ mod tests {
         sort_preserving_merge_exec, spr_repartition_exec, union_exec,
         RequirementsTestExec,
     };
+    use crate::physical_optimizer::utils::sort_preserving_merge_exec_with_fetch;
     use crate::physical_plan::{displayable, get_plan_string, Partitioning};
     use crate::prelude::{SessionConfig, SessionContext};
     use crate::test::{csv_exec_ordered, csv_exec_sorted, stream_exec_ordered};
@@ -1399,6 +1406,30 @@ mod tests {
         Ok(())
     }
 
+    #[tokio::test]
+    async fn test_remove_unnecessary_spm2() -> Result<()> {
+        let schema = create_test_schema()?;
+        let source = memory_exec(&schema);
+        let input = sort_preserving_merge_exec_with_fetch(
+            vec![sort_expr("non_nullable_col", &schema)],
+            source,
+            100,
+        );
+
+        let expected_input = [
+            "SortPreservingMergeExec: [non_nullable_col@1 ASC], fetch=100",
+            "  MemoryExec: partitions=1, partition_sizes=[0]",
+        ];
+        let expected_optimized = [
+            "LocalLimitExec: fetch=100",
+            "  SortExec: expr=[non_nullable_col@1 ASC], preserve_partitioning=[false]",
+            "    MemoryExec: partitions=1, partition_sizes=[0]",
+        ];
+        assert_optimized!(expected_input, expected_optimized, input, true);
+
+        Ok(())
+    }
+
     #[tokio::test]
     async fn test_union_inputs_different_sorted2() -> Result<()> {
         let schema = create_test_schema()?;

datafusion/core/src/physical_optimizer/utils.rs

@@ -27,7 +27,9 @@ use crate::physical_plan::union::UnionExec;
 use crate::physical_plan::windows::{BoundedWindowAggExec, WindowAggExec};
 use crate::physical_plan::{ExecutionPlan, ExecutionPlanProperties};
 
-use datafusion_physical_expr::{LexRequirement, PhysicalSortRequirement};
+use datafusion_physical_expr::{
+    LexRequirement, PhysicalSortExpr, PhysicalSortRequirement,
+};
 use datafusion_physical_plan::limit::{GlobalLimitExec, LocalLimitExec};
 use datafusion_physical_plan::tree_node::PlanContext;
 
@@ -52,6 +54,15 @@ pub fn add_sort_above<T: Clone + Default>(
     PlanContext::new(Arc::new(new_sort), T::default(), vec![node])
 }
 
+pub fn sort_preserving_merge_exec_with_fetch(
+    sort_exprs: impl IntoIterator<Item = PhysicalSortExpr>,
+    input: Arc<dyn ExecutionPlan>,
+    fetch: usize,
+) -> Arc<dyn ExecutionPlan> {
+    let sort_exprs = sort_exprs.into_iter().collect();
+    Arc::new(SortPreservingMergeExec::new(sort_exprs, input).with_fetch(Some(fetch)))
+}
+
 /// This utility function adds a `SortExec` above an operator according to the
 /// given ordering requirements while preserving the original partitioning. If
 /// requirement is already satisfied no `SortExec` is added.