Add fetch method to LogicalPlan

Author: shivbhatia10

datafusion/expr/src/logical_plan/plan.rs

@@ -1392,6 +1392,40 @@ impl LogicalPlan {
         }
     }
 
+    /// Returns the fetch (limit) of this plan node, if it has one.
+    ///
+    /// Only [`LogicalPlan::Sort`] and [`LogicalPlan::TableScan`] carry a fetch
+    /// value directly; all other variants return `None`.
+    pub fn fetch(&self) -> Option<usize> {
+        match self {
+            LogicalPlan::Sort(Sort { fetch, .. }) => *fetch,
+            LogicalPlan::TableScan(TableScan { fetch, .. }) => *fetch,
+            LogicalPlan::Projection(_) => None,
+            LogicalPlan::Filter(_) => None,
+            LogicalPlan::Window(_) => None,
+            LogicalPlan::Aggregate(_) => None,
+            LogicalPlan::Join(_) => None,
+            LogicalPlan::Repartition(_) => None,
+            LogicalPlan::Union(_) => None,
+            LogicalPlan::EmptyRelation(_) => None,
+            LogicalPlan::Subquery(_) => None,
+            LogicalPlan::SubqueryAlias(_) => None,
+            LogicalPlan::Limit(_) => None,
+            LogicalPlan::Statement(_) => None,
+            LogicalPlan::Values(_) => None,
+            LogicalPlan::Explain(_) => None,
+            LogicalPlan::Analyze(_) => None,
+            LogicalPlan::Extension(_) => None,
+            LogicalPlan::Distinct(_) => None,
+            LogicalPlan::Dml(_) => None,
+            LogicalPlan::Ddl(_) => None,
+            LogicalPlan::Copy(_) => None,
+            LogicalPlan::DescribeTable(_) => None,
+            LogicalPlan::Unnest(_) => None,
+            LogicalPlan::RecursiveQuery(_) => None,
+        }
+    }
+
     /// If this node's expressions contains any references to an outer subquery
     pub fn contains_outer_reference(&self) -> bool {
         let mut contains = false;

datafusion/optimizer/src/push_down_filter.rs

@@ -796,6 +796,12 @@ impl OptimizerRule for PushDownFilter {
             filter.predicate = new_predicate;
         }
 
+        // If the child has a fetch (limit), pushing a filter below it would
+        // change semantics: the limit should apply before the filter, not after.
+        if filter.input.fetch().is_some() {
+            return Ok(Transformed::no(LogicalPlan::Filter(filter)));
+        }
+
         match Arc::unwrap_or_clone(filter.input) {
             LogicalPlan::Filter(child_filter) => {
                 let parents_predicates = split_conjunction_owned(filter.predicate);
@@ -832,13 +838,6 @@ impl OptimizerRule for PushDownFilter {
                 insert_below(LogicalPlan::Distinct(distinct), new_filter)
             }
             LogicalPlan::Sort(sort) => {
-                // If the sort has a fetch (limit), pushing a filter below
-                // it would change semantics: the limit should apply before
-                // the filter, not after.
-                if sort.fetch.is_some() {
-                    filter.input = Arc::new(LogicalPlan::Sort(sort));
-                    return Ok(Transformed::no(LogicalPlan::Filter(filter)));
-                }
                 let new_filter =
                     Filter::try_new(filter.predicate, Arc::clone(&sort.input))
                         .map(LogicalPlan::Filter)?;
@@ -1137,13 +1136,6 @@ impl OptimizerRule for PushDownFilter {
             }
             LogicalPlan::Join(join) => push_down_join(join, Some(&filter.predicate)),
             LogicalPlan::TableScan(scan) => {
-                // If the scan has a fetch (limit), pushing filters into it
-                // would change semantics: the limit should apply before the
-                // filter, not after.
-                if scan.fetch.is_some() {
-                    filter.input = Arc::new(LogicalPlan::TableScan(scan));
-                    return Ok(Transformed::no(LogicalPlan::Filter(filter)));
-                }
                 let filter_predicates = split_conjunction(&filter.predicate);
 
                 let (volatile_filters, non_volatile_filters): (Vec<&Expr>, Vec<&Expr>) =