fix: fetch is missed in the EnsureSorting (apache#14192)

* fix: fetch is missed in the EnsureSorting

* fix conflict

* resolve comments from alamb

* update

Author: xudong963

datafusion/core/src/physical_optimizer/enforce_sorting.rs

@@ -186,13 +186,11 @@ impl PhysicalOptimizerRule for EnforceSorting {
                 )
             })
             .data()?;
-
         // Execute a top-down traversal to exploit sort push-down opportunities
         // missed by the bottom-up traversal:
         let mut sort_pushdown = SortPushDown::new_default(updated_plan.plan);
         assign_initial_requirements(&mut sort_pushdown);
         let adjusted = pushdown_sorts(sort_pushdown)?;
-
         adjusted
             .plan
             .transform_up(|plan| Ok(Transformed::yes(replace_with_partial_sort(plan)?)))

datafusion/core/src/physical_optimizer/sort_pushdown.rs

@@ -64,7 +64,9 @@ pub fn assign_initial_requirements(node: &mut SortPushDown) {
     for (child, requirement) in node.children.iter_mut().zip(reqs) {
         child.data = ParentRequirements {
             ordering_requirement: requirement,
-            fetch: None,
+            // If the parent has a fetch value, assign it to the children
+            // Or use the fetch value of the child.
+            fetch: child.plan.fetch(),
         };
     }
 }
@@ -95,14 +97,16 @@ fn pushdown_sorts_helper(
         .equivalence_properties()
         .ordering_satisfy_requirement(parent_reqs);
     if is_sort(plan) {
+        let sort_fetch = plan.fetch();
         let required_ordering = plan
             .output_ordering()
             .map(PhysicalSortRequirement::from_sort_exprs)
             .unwrap_or_default();
         if !satisfy_parent {
             // Make sure this `SortExec` satisfies parent requirements:
             let sort_reqs = requirements.data.ordering_requirement.unwrap_or_default();
-            let fetch = requirements.data.fetch;
+            // It's possible current plan (`SortExec`) has a fetch value.
+            let fetch = requirements.data.fetch.or(sort_fetch);
             requirements = requirements.children.swap_remove(0);
             requirements = add_sort_above(requirements, sort_reqs, fetch);
         };
@@ -112,7 +116,7 @@ fn pushdown_sorts_helper(
         if let Some(adjusted) =
             pushdown_requirement_to_children(&child.plan, &required_ordering)?
         {
-            let fetch = child.plan.fetch();
+            let fetch = sort_fetch.or_else(|| child.plan.fetch());
             for (grand_child, order) in child.children.iter_mut().zip(adjusted) {
                 grand_child.data = ParentRequirements {
                     ordering_requirement: order,

datafusion/sqllogictest/test_files/topk.slt

@@ -49,7 +49,12 @@ select * from topk order by x desc limit 3;
 8
 5
 
-
+query I
+select * from (select * from topk limit 8) order by x limit 3;
+----
+0
+1
+2
 
 
 statement ok