refactor: one vocabulary for moving LIMIT through the physical plan

datafusion/core/tests/physical_optimizer/ensure_requirements.rs

@@ -510,6 +510,35 @@ fn test_skip_and_fetch_multi_partition() {
     ");
 }
 
+/// A merge with its own fetch already selects rows, so the fetch of a
+/// stricter TopK above it must stay above it.
+#[test]
+fn test_fetched_merge_keeps_stricter_topk_above() {
+    let source = Arc::new(MockMultiPartitionExec::new(4));
+    let a = sort_expr_on("a", 0, false, false);
+    let a_b = LexOrdering::new(
+        a.iter()
+            .cloned()
+            .chain(sort_expr_on("b", 1, false, false).iter().cloned()),
+    )
+    .unwrap();
+    let inner = Arc::new(
+        SortExec::new(a.clone(), source)
+            .with_fetch(Some(5))
+            .with_preserve_partitioning(true),
+    );
+    let merge = Arc::new(SortPreservingMergeExec::new(a, inner).with_fetch(Some(5)));
+    let topk: Arc<dyn ExecutionPlan> =
+        Arc::new(SortExec::new(a_b, merge).with_fetch(Some(2)));
+
+    assert_ensure_requirements_plan!(topk, @r"
+    SortExec: TopK(fetch=2), expr=[a@0 ASC NULLS LAST, b@1 ASC NULLS LAST], preserve_partitioning=[false], sort_prefix=[a@0 ASC NULLS LAST]
+      SortPreservingMergeExec: [a@0 ASC NULLS LAST], fetch=5
+        LocalLimitExec: fetch=5
+          MockMultiPartitionExec
+    ");
+}
+
 // ========================================================================
 // Multiple sort columns test
 // ========================================================================

datafusion/physical-optimizer/src/ensure_requirements/enforce_sorting/mod.rs

@@ -49,6 +49,7 @@ pub mod sort_pushdown;
 use std::sync::Arc;
 
 use crate::output_requirements::OutputRequirementExec;
+use crate::row_bound::{RowBound, limit_node};
 use crate::utils::{
     add_sort_above, add_sort_above_with_check, is_coalesce_partitions, is_limit,
     is_repartition, is_sort, is_sort_preserving_merge, is_window,
@@ -60,7 +61,7 @@ use datafusion_common::tree_node::Transformed;
 use datafusion_physical_expr::{Distribution, Partitioning};
 use datafusion_physical_expr_common::sort_expr::{LexOrdering, LexRequirement};
 use datafusion_physical_plan::coalesce_partitions::CoalescePartitionsExec;
-use datafusion_physical_plan::limit::{GlobalLimitExec, LocalLimitExec};
+use datafusion_physical_plan::limit::LocalLimitExec;
 use datafusion_physical_plan::repartition::RepartitionExec;
 use datafusion_physical_plan::sorts::partial_sort::PartialSortExec;
 use datafusion_physical_plan::sorts::sort::SortExec;
@@ -512,17 +513,13 @@ fn analyze_immediate_sort_removal(
             .with_fetch(sort_exec.fetch()),
         ) as _
     } else {
-        // Remove the sort:
+        // Remove the sort, keeping its fetch as a limit:
         node.children = node.children.swap_remove(0).children;
-        if let Some(fetch) = sort_exec.fetch() {
-            // If the sort has a fetch, we need to add a limit:
-            if properties.output_partitioning().partition_count() == 1 {
-                Arc::new(GlobalLimitExec::new(Arc::clone(sort_input), 0, Some(fetch)))
-            } else {
-                Arc::new(LocalLimitExec::new(Arc::clone(sort_input), fetch))
+        match sort_exec.fetch() {
+            Some(fetch) => {
+                limit_node(Arc::clone(sort_input), RowBound::fetch(Some(fetch)))
             }
-        } else {
-            Arc::clone(sort_input)
+            None => Arc::clone(sort_input),
         }
     };
     for child in node.children.iter_mut() {

datafusion/physical-optimizer/src/ensure_requirements/enforce_sorting/sort_pushdown.rs

@@ -37,7 +37,6 @@ use datafusion_physical_expr_common::sort_expr::{
     PhysicalSortRequirement,
 };
 use datafusion_physical_plan::aggregates::AggregateExec;
-use datafusion_physical_plan::execution_plan::CardinalityEffect;
 use datafusion_physical_plan::filter::FilterExec;
 use datafusion_physical_plan::joins::utils::{
     ColumnIndex, calculate_join_output_ordering,
@@ -49,7 +48,9 @@ use datafusion_physical_plan::repartition::RepartitionExec;
 use datafusion_physical_plan::sorts::sort::SortExec;
 use datafusion_physical_plan::sorts::sort_preserving_merge::SortPreservingMergeExec;
 use datafusion_physical_plan::tree_node::PlanContext;
-use datafusion_physical_plan::{ExecutionPlan, ExecutionPlanProperties};
+use datafusion_physical_plan::{ExecutionPlan, ExecutionPlanProperties, LimitBehavior};
+
+use crate::row_bound::{RowBound, limit_node};
 
 /// "Data class" used by sort pushdown (now driven from `EnsureRequirements`)
 /// to push down [`SortExec`] in the plan. In some cases the total
@@ -94,21 +95,19 @@ pub type SortPushDown = PlanContext<ParentRequirements>;
 /// fetch is the tighter bound. Using the bare `fetch` would be wrong too, as
 /// it would turn `LIMIT 10 OFFSET 5` into `TopK(10)` below the limit, i.e. 5
 /// result rows.
-///
-/// `skip` and `fetch` are independent `usize`s, so their sum can overflow. Like
-/// [`combine_limit`], we saturate: `usize::MAX` input rows is never a smaller
-/// bound than the real one, so the pushed-down fetch stays correct.
-///
-/// [`combine_limit`]: datafusion_common::utils::combine_limit
 fn input_fetch(
     plan: &Arc<dyn ExecutionPlan>,
     parent_fetch: Option<usize>,
 ) -> Option<usize> {
-    let fetch = min_fetch(plan.fetch(), parent_fetch)?;
     let skip = plan
         .downcast_ref::<GlobalLimitExec>()
         .map_or(0, |limit| limit.skip());
-    Some(fetch.saturating_add(skip))
+    let output = RowBound::fetch(plan.fetch()).min(RowBound::fetch(parent_fetch));
+    RowBound {
+        skip,
+        fetch: output.fetch,
+    }
+    .input_rows()
 }
 
 /// Assigns the ordering requirement of the root node to the its children.
@@ -140,18 +139,15 @@ pub fn pushdown_sorts(sort_push_down: SortPushDown) -> Result<SortPushDown> {
 }
 
 fn min_fetch(f1: Option<usize>, f2: Option<usize>) -> Option<usize> {
-    match (f1, f2) {
-        (Some(f1), Some(f2)) => Some(f1.min(f2)),
-        (Some(_), _) => f1,
-        (_, Some(_)) => f2,
-        _ => None,
-    }
+    RowBound::fetch(f1).min(RowBound::fetch(f2)).fetch
 }
 
-/// Returns whether a fetch on `plan` can also be applied to each child.
+/// Returns whether a fetch on `plan` can move to its input.
 fn can_push_fetch_through(plan: &Arc<dyn ExecutionPlan>) -> bool {
-    plan.supports_limit_pushdown()
-        && matches!(plan.cardinality_effect(), CardinalityEffect::Equal)
+    matches!(
+        plan.limit_behavior(),
+        LimitBehavior::MoveToInput | LimitBehavior::Merge { fetch: None }
+    )
 }
 
 /// Returns a plan when the fast path for an already-satisfied ordering can also
@@ -188,11 +184,7 @@ fn preserve_fetch_above_ordered_plan(
         node.plan.output_partitioning().partition_count() > 1;
     let input = Arc::clone(&node.plan);
 
-    let limit: Arc<dyn ExecutionPlan> = if input_has_multiple_partitions {
-        Arc::new(LocalLimitExec::new(input, fetch))
-    } else {
-        Arc::new(GlobalLimitExec::new(input, 0, Some(fetch)))
-    };
+    let limit = limit_node(input, RowBound::fetch(Some(fetch)));
     let limit_node = PlanContext::new(limit, ParentRequirements::default(), vec![node]);
 
     if input_has_multiple_partitions
@@ -452,29 +444,10 @@ fn pushdown_requirement_to_children(
     parent_required: OrderingRequirements,
     parent_fetch: Option<usize>,
 ) -> Result<Option<Vec<Option<OrderingRequirements>>>> {
-    // A parent fetch can be pushed through a plan only if the plan explicitly
-    // supports limit pushdown and does not change cardinality. For example, if
-    // LIMIT 2 is applied below a FilterExec that removes half the rows, only one
-    // row may remain. Applied after the filter, it correctly returns two rows
-    // when at least two are available.
-    //
-    // Checking `supports_limit_pushdown()` alone is not enough. For example:
-    //
-    // UnionExec advertises `supports_limit_pushdown() == true` because it can
-    // forward a LIMIT k to each of its children—i.e. apply “LIMIT k” separately
-    // on each branch before merging them together.
-    //
-    // However, UnionExec’s `cardinality_effect() == GreaterEqual` (it sums up
-    // all child row counts), so pushing a global TopK/LIMIT through it would
-    // break the semantics of “take the first k rows of the combined result.”
-    //
-    // For example, with two branches A and B and k = 3:
-    //   — Global LIMIT: take the first 3 rows from (A ∪ B) after merging.
-    //   — Pushed down: take 3 from A, 3 from B, then merge → up to 6 rows!
-    //
-    // That’s why we also require equal cardinality: even though UnionExec can
-    // push a LIMIT to its children, its GreaterEqual effect means it cannot
-    // preserve the global TopK semantics.
+    // Pushing the sort below `plan` moves the parent fetch with it. For
+    // example, if LIMIT 2 is applied below a FilterExec that removes half the
+    // rows, only one row may remain; and a LIMIT 3 pushed into both inputs of
+    // a UnionExec keeps up to 6 rows. See `LimitBehavior`.
     if parent_fetch.is_some() && !can_push_fetch_through(plan) {
         return Ok(None);
     }
@@ -549,7 +522,10 @@ fn pushdown_requirement_to_children(
             Ok(None)
         }
     } else if plan.fetch().is_some()
-        && plan.supports_limit_pushdown()
+        && !matches!(
+            plan.limit_behavior(),
+            LimitBehavior::Stops | LimitBehavior::Window(_)
+        )
         && plan
             .maintains_input_order()
             .into_iter()

datafusion/physical-optimizer/src/lib.rs

@@ -45,6 +45,7 @@ pub mod projection_pushdown;
 pub use datafusion_pruning as pruning;
 pub mod hash_join_buffering;
 pub mod pushdown_sort;
+mod row_bound;
 pub mod sanity_checker;
 pub mod topk_aggregation;
 pub mod topk_repartition;