Add Fetch Property to OutputRequirementExec (apache#16892)

* ore fetch

* Update output_requirements.rs

Author: berkaysynnada

datafusion/core/tests/physical_optimizer/enforce_sorting.rs

@@ -989,6 +989,7 @@ async fn test_soft_hard_requirements_with_multiple_soft_requirements_and_output_
         bounded_window2,
         Some(OrderingRequirements::new(requirement)),
         Distribution::SinglePartition,
+        None,
     ));
 
     let expected_input = [

datafusion/core/tests/physical_optimizer/projection_pushdown.rs

@@ -712,6 +712,7 @@ fn test_output_req_after_projection() -> Result<()> {
             Arc::new(Column::new("a", 0)),
             Arc::new(Column::new("b", 1)),
         ]),
+        None,
     ));
     let projection: Arc<dyn ExecutionPlan> = Arc::new(ProjectionExec::try_new(
         vec![

datafusion/physical-optimizer/src/output_requirements.rs

@@ -32,6 +32,7 @@ use datafusion_common::{Result, Statistics};
 use datafusion_execution::TaskContext;
 use datafusion_physical_expr::Distribution;
 use datafusion_physical_expr_common::sort_expr::OrderingRequirements;
+use datafusion_physical_plan::execution_plan::Boundedness;
 use datafusion_physical_plan::projection::{
     make_with_child, update_expr, update_ordering_requirement, ProjectionExec,
 };
@@ -98,20 +99,23 @@ pub struct OutputRequirementExec {
     order_requirement: Option<OrderingRequirements>,
     dist_requirement: Distribution,
     cache: PlanProperties,
+    fetch: Option<usize>,
 }
 
 impl OutputRequirementExec {
     pub fn new(
         input: Arc<dyn ExecutionPlan>,
         requirements: Option<OrderingRequirements>,
         dist_requirement: Distribution,
+        fetch: Option<usize>,
     ) -> Self {
-        let cache = Self::compute_properties(&input);
+        let cache = Self::compute_properties(&input, &fetch);
         Self {
             input,
             order_requirement: requirements,
             dist_requirement,
             cache,
+            fetch,
         }
     }
 
@@ -120,14 +124,28 @@ impl OutputRequirementExec {
     }
 
     /// This function creates the cache object that stores the plan properties such as schema, equivalence properties, ordering, partitioning, etc.
-    fn compute_properties(input: &Arc<dyn ExecutionPlan>) -> PlanProperties {
+    fn compute_properties(
+        input: &Arc<dyn ExecutionPlan>,
+        fetch: &Option<usize>,
+    ) -> PlanProperties {
+        let boundedness = if fetch.is_some() {
+            Boundedness::Bounded
+        } else {
+            input.boundedness()
+        };
+
         PlanProperties::new(
             input.equivalence_properties().clone(), // Equivalence Properties
             input.output_partitioning().clone(),    // Output Partitioning
             input.pipeline_behavior(),              // Pipeline Behavior
-            input.boundedness(),                    // Boundedness
+            boundedness,                            // Boundedness
         )
     }
+
+    /// Get fetch
+    pub fn fetch(&self) -> Option<usize> {
+        self.fetch
+    }
 }
 
 impl DisplayAs for OutputRequirementExec {
@@ -214,6 +232,7 @@ impl ExecutionPlan for OutputRequirementExec {
             children.remove(0), // has a single child
             self.order_requirement.clone(),
             self.dist_requirement.clone(),
+            self.fetch,
         )))
     }
 
@@ -273,10 +292,14 @@ impl ExecutionPlan for OutputRequirementExec {
         };
 
         make_with_child(projection, &self.input()).map(|input| {
-            let e = OutputRequirementExec::new(input, requirements, dist_req);
+            let e = OutputRequirementExec::new(input, requirements, dist_req, self.fetch);
             Some(Arc::new(e) as _)
         })
     }
+
+    fn fetch(&self) -> Option<usize> {
+        self.fetch
+    }
 }
 
 impl PhysicalOptimizerRule for OutputRequirements {
@@ -323,6 +346,7 @@ fn require_top_ordering(plan: Arc<dyn ExecutionPlan>) -> Result<Arc<dyn Executio
             // there is no ordering requirement
             None,
             Distribution::UnspecifiedDistribution,
+            None,
         )) as _)
     }
 }
@@ -341,20 +365,29 @@ fn require_top_ordering_helper(
         // In case of constant columns, output ordering of the `SortExec` would
         // be an empty set. Therefore; we check the sort expression field to
         // assign the requirements.
-        let req_ordering = sort_exec.expr();
         let req_dist = sort_exec.required_input_distribution().swap_remove(0);
+        let req_ordering = sort_exec.expr();
         let reqs = OrderingRequirements::from(req_ordering.clone());
+        let fetch = sort_exec.fetch();
+
         Ok((
-            Arc::new(OutputRequirementExec::new(plan, Some(reqs), req_dist)) as _,
+            Arc::new(OutputRequirementExec::new(
+                plan,
+                Some(reqs),
+                req_dist,
+                fetch,
+            )) as _,
             true,
         ))
     } else if let Some(spm) = plan.as_any().downcast_ref::<SortPreservingMergeExec>() {
         let reqs = OrderingRequirements::from(spm.expr().clone());
+        let fetch = spm.fetch();
         Ok((
             Arc::new(OutputRequirementExec::new(
                 plan,
                 Some(reqs),
                 Distribution::SinglePartition,
+                fetch,
             )) as _,
             true,
         ))