From 4b30eced586ee35b55cc204b353d34a9f7c64926 Mon Sep 17 00:00:00 2001 From: discord9 Date: Wed, 22 Jul 2026 19:37:23 +0800 Subject: [PATCH] fix: preserve global limit for multi-partition fetch Signed-off-by: discord9 --- .../physical_optimizer/limit_pushdown.rs | 461 +++++++++++++++++- .../physical-optimizer/src/limit_pushdown.rs | 228 +++++---- 2 files changed, 579 insertions(+), 110 deletions(-) diff --git a/datafusion/core/tests/physical_optimizer/limit_pushdown.rs b/datafusion/core/tests/physical_optimizer/limit_pushdown.rs index b8ebc80348134..d2526cb3531c4 100644 --- a/datafusion/core/tests/physical_optimizer/limit_pushdown.rs +++ b/datafusion/core/tests/physical_optimizer/limit_pushdown.rs @@ -15,30 +15,38 @@ // specific language governing permissions and limitations // under the License. +use std::fmt::Formatter; use std::sync::Arc; use crate::physical_optimizer::test_utils::{ - coalesce_partitions_exec, global_limit_exec, hash_join_exec, local_limit_exec, - sort_exec, sort_preserving_merge_exec, stream_exec, + TestScan, coalesce_partitions_exec, global_limit_exec, hash_join_exec, + local_limit_exec, sort_exec, sort_preserving_merge_exec, stream_exec, }; use arrow::compute::SortOptions; use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; +use datafusion_common::Statistics; use datafusion_common::config::ConfigOptions; use datafusion_common::error::Result; +use datafusion_execution::{SendableRecordBatchStream, TaskContext}; use datafusion_expr::{JoinType, Operator}; -use datafusion_physical_expr::Partitioning; use datafusion_physical_expr::expressions::{BinaryExpr, col, lit}; +use datafusion_physical_expr::{EquivalenceProperties, Partitioning}; use datafusion_physical_expr_common::physical_expr::PhysicalExprRef; use datafusion_physical_expr_common::sort_expr::{LexOrdering, PhysicalSortExpr}; use datafusion_physical_optimizer::PhysicalOptimizerRule; use datafusion_physical_optimizer::limit_pushdown::LimitPushdown; use datafusion_physical_plan::empty::EmptyExec; +use datafusion_physical_plan::execution_plan::{Boundedness, EmissionType}; use datafusion_physical_plan::filter::FilterExec; use datafusion_physical_plan::joins::NestedLoopJoinExec; use datafusion_physical_plan::projection::ProjectionExec; use datafusion_physical_plan::repartition::RepartitionExec; -use datafusion_physical_plan::{ExecutionPlan, get_plan_string}; +use datafusion_physical_plan::union::UnionExec; +use datafusion_physical_plan::{ + DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, StatisticsArgs, + get_plan_string, +}; fn create_schema() -> SchemaRef { Arc::new(Schema::new(vec![ @@ -103,6 +111,75 @@ fn format_plan(plan: &Arc) -> String { get_plan_string(plan).join("\n") } +#[derive(Debug)] +struct TestCombinerExec { + input: Arc, + properties: Arc, +} + +impl TestCombinerExec { + fn new(input: Arc) -> Self { + let properties = PlanProperties::new( + EquivalenceProperties::new(input.schema()), + Partitioning::UnknownPartitioning(1), + EmissionType::Incremental, + Boundedness::Bounded, + ); + Self { + input, + properties: Arc::new(properties), + } + } +} + +impl DisplayAs for TestCombinerExec { + fn fmt_as(&self, _t: DisplayFormatType, f: &mut Formatter) -> std::fmt::Result { + write!(f, "TestCombinerExec") + } +} + +impl ExecutionPlan for TestCombinerExec { + fn name(&self) -> &str { + "TestCombinerExec" + } + + fn properties(&self) -> &Arc { + &self.properties + } + + fn children(&self) -> Vec<&Arc> { + vec![&self.input] + } + + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + assert_eq!(children.len(), 1); + Ok(Arc::new(Self::new(children[0].clone()))) + } + + fn execute( + &self, + _partition: usize, + _context: Arc, + ) -> Result { + unreachable!("TestCombinerExec is only used by optimizer tests") + } + + fn statistics_from_inputs( + &self, + _input_stats: &[Arc], + _args: &StatisticsArgs, + ) -> Result> { + Ok(Arc::new(Statistics::new_unknown(self.schema().as_ref()))) + } + + fn supports_limit_pushdown(&self) -> bool { + true + } +} + #[test] fn transforms_streaming_table_exec_into_fetching_version_when_skip_is_zero() -> Result<()> { @@ -162,6 +239,343 @@ fn transforms_streaming_table_exec_into_fetching_version_and_keeps_the_global_li Ok(()) } +#[test] +fn keeps_global_limit_above_fetch_capable_multi_partition_scan() -> Result<()> { + let schema = create_schema(); + let scan = Arc::new( + TestScan::new(schema, vec![]) + .with_supports_fetch(true) + .with_partition_count(2), + ); + let global_limit = global_limit_exec(scan, 0, Some(5)); + + let optimized = LimitPushdown::new().optimize(global_limit, &ConfigOptions::new())?; + + insta::assert_snapshot!( + format_plan(&optimized), + @r" + CoalescePartitionsExec: fetch=5 + TestScan: fetch=5 + " + ); + + Ok(()) +} + +#[test] +fn keeps_global_offset_limit_above_fetch_capable_multi_partition_scan() -> Result<()> { + let schema = create_schema(); + let scan = Arc::new( + TestScan::new(schema, vec![]) + .with_supports_fetch(true) + .with_partition_count(2), + ); + let global_limit = global_limit_exec(scan, 2, Some(5)); + + let optimized = LimitPushdown::new().optimize(global_limit, &ConfigOptions::new())?; + + insta::assert_snapshot!( + format_plan(&optimized), + @r" + GlobalLimitExec: skip=2, fetch=5 + CoalescePartitionsExec: fetch=7 + TestScan: fetch=7 + " + ); + + Ok(()) +} + +#[test] +fn preserves_existing_per_partition_fetch_under_global_limit() -> Result<()> { + let schema = create_schema(); + let scan = Arc::new( + TestScan::new(schema, vec![]) + .with_supports_fetch(true) + .with_partition_count(2), + ); + let scan = scan.with_fetch(Some(3)).unwrap(); + let global_limit = global_limit_exec(scan, 0, Some(5)); + + let optimized = LimitPushdown::new().optimize(global_limit, &ConfigOptions::new())?; + + insta::assert_snapshot!( + format_plan(&optimized), + @r" + CoalescePartitionsExec: fetch=5 + TestScan: fetch=3 + " + ); + + Ok(()) +} + +#[test] +fn adds_global_boundary_above_unfetchable_multi_partition_scan() -> Result<()> { + let schema = create_schema(); + let scan = Arc::new(TestScan::new(schema, vec![]).with_partition_count(2)); + let global_limit = global_limit_exec(scan, 0, Some(5)); + + let optimized = LimitPushdown::new().optimize(global_limit, &ConfigOptions::new())?; + + insta::assert_snapshot!( + format_plan(&optimized), + @r" + CoalescePartitionsExec: fetch=5 + TestScan + " + ); + + Ok(()) +} + +#[test] +fn materializes_global_boundary_before_pushing_into_union_children() -> Result<()> { + let schema = create_schema(); + let left = + Arc::new(TestScan::new(Arc::clone(&schema), vec![]).with_supports_fetch(true)); + let right = Arc::new(TestScan::new(schema, vec![]).with_supports_fetch(true)); + let union = UnionExec::try_new(vec![left, right])?; + let global_limit = global_limit_exec(union, 0, Some(5)); + + let optimized = LimitPushdown::new().optimize(global_limit, &ConfigOptions::new())?; + + insta::assert_snapshot!( + format_plan(&optimized), + @r" + CoalescePartitionsExec: fetch=5 + UnionExec + TestScan: fetch=5 + TestScan: fetch=5 + " + ); + + Ok(()) +} + +#[test] +fn materializes_global_boundary_for_offset_only_multi_partition_scan() -> Result<()> { + let schema = create_schema(); + let scan = Arc::new(TestScan::new(schema, vec![]).with_partition_count(2)); + let global_limit = global_limit_exec(scan, 2, None); + + let optimized = LimitPushdown::new().optimize(global_limit, &ConfigOptions::new())?; + + insta::assert_snapshot!( + format_plan(&optimized), + @r" + GlobalLimitExec: skip=2, fetch=None + CoalescePartitionsExec + TestScan + " + ); + + Ok(()) +} + +#[test] +fn removes_noop_global_limit_without_materializing_boundary() -> Result<()> { + let schema = create_schema(); + let scan = Arc::new(TestScan::new(schema, vec![])); + let noop_global_limit = global_limit_exec(scan, 0, None); + + let optimized = + LimitPushdown::new().optimize(noop_global_limit, &ConfigOptions::new())?; + + insta::assert_snapshot!( + format_plan(&optimized), + @"TestScan" + ); + + Ok(()) +} + +#[test] +fn preserves_outer_global_limit_across_nested_global_limit() -> Result<()> { + let schema = create_schema(); + let scan = Arc::new( + TestScan::new(schema, vec![]) + .with_supports_fetch(true) + .with_partition_count(2), + ); + let inner = global_limit_exec(scan, 0, Some(10)); + let outer = global_limit_exec(inner, 0, Some(5)); + + let optimized = LimitPushdown::new().optimize(outer, &ConfigOptions::new())?; + + insta::assert_snapshot!( + format_plan(&optimized), + @r" + CoalescePartitionsExec: fetch=5 + TestScan: fetch=5 + " + ); + + Ok(()) +} + +#[test] +fn preserves_outer_global_limit_across_noop_global_limit() -> Result<()> { + let schema = create_schema(); + let scan = Arc::new( + TestScan::new(schema, vec![]) + .with_supports_fetch(true) + .with_partition_count(2), + ); + let noop = global_limit_exec(scan, 0, None); + let outer = global_limit_exec(noop, 0, Some(5)); + + let optimized = LimitPushdown::new().optimize(outer, &ConfigOptions::new())?; + + insta::assert_snapshot!( + format_plan(&optimized), + @r" + CoalescePartitionsExec: fetch=5 + TestScan: fetch=5 + " + ); + + Ok(()) +} + +#[test] +fn materializes_pending_global_limit_below_extension_combiner() -> Result<()> { + let schema = create_schema(); + let left = + Arc::new(TestScan::new(Arc::clone(&schema), vec![]).with_supports_fetch(true)); + let right = Arc::new(TestScan::new(schema, vec![]).with_supports_fetch(true)); + let union = UnionExec::try_new(vec![left, right])?; + let combiner = Arc::new(TestCombinerExec::new(union)); + let global_limit = global_limit_exec(combiner, 0, Some(5)); + + let optimized = LimitPushdown::new().optimize(global_limit, &ConfigOptions::new())?; + + insta::assert_snapshot!( + format_plan(&optimized), + @r" + TestCombinerExec + CoalescePartitionsExec: fetch=5 + UnionExec + TestScan: fetch=5 + TestScan: fetch=5 + " + ); + + Ok(()) +} + +#[test] +fn upgrades_pending_local_limit_before_extension_combiner() -> Result<()> { + let schema = create_schema(); + let inner_left = + Arc::new(TestScan::new(Arc::clone(&schema), vec![]).with_supports_fetch(true)); + let inner_right = + Arc::new(TestScan::new(Arc::clone(&schema), vec![]).with_supports_fetch(true)); + let inner_union = UnionExec::try_new(vec![inner_left, inner_right])?; + let combiner = Arc::new(TestCombinerExec::new(inner_union)); + let outer_child = Arc::new(TestScan::new(schema, vec![]).with_supports_fetch(true)); + let outer_union = UnionExec::try_new(vec![combiner, outer_child])?; + let local_limit = local_limit_exec(outer_union, 5); + + let optimized = LimitPushdown::new().optimize(local_limit, &ConfigOptions::new())?; + + insta::assert_snapshot!( + format_plan(&optimized), + @r" + UnionExec + TestCombinerExec + CoalescePartitionsExec: fetch=5 + UnionExec + TestScan: fetch=5 + TestScan: fetch=5 + TestScan: fetch=5 + " + ); + + Ok(()) +} + +#[test] +fn upgrades_pending_local_limit_before_noop_global_wrapper() -> Result<()> { + let schema = create_schema(); + let inner_left = + Arc::new(TestScan::new(Arc::clone(&schema), vec![]).with_supports_fetch(true)); + let inner_right = + Arc::new(TestScan::new(Arc::clone(&schema), vec![]).with_supports_fetch(true)); + let inner_union = UnionExec::try_new(vec![inner_left, inner_right])?; + let noop_global = global_limit_exec(inner_union, 0, None); + let outer_child = Arc::new(TestScan::new(schema, vec![]).with_supports_fetch(true)); + let outer_union = UnionExec::try_new(vec![noop_global, outer_child])?; + let local_limit = local_limit_exec(outer_union, 5); + + let optimized = LimitPushdown::new().optimize(local_limit, &ConfigOptions::new())?; + + insta::assert_snapshot!( + format_plan(&optimized), + @r" + UnionExec + CoalescePartitionsExec: fetch=5 + UnionExec + TestScan: fetch=5 + TestScan: fetch=5 + TestScan: fetch=5 + " + ); + + Ok(()) +} + +#[test] +fn keeps_global_limit_above_local_limit_on_multi_partition_union() -> Result<()> { + let schema = create_schema(); + let left = + Arc::new(TestScan::new(Arc::clone(&schema), vec![]).with_supports_fetch(true)); + let right = Arc::new(TestScan::new(schema, vec![]).with_supports_fetch(true)); + let union = UnionExec::try_new(vec![left, right])?; + let local_limit = local_limit_exec(union, 3); + let global_limit = global_limit_exec(local_limit, 0, Some(5)); + + let optimized = LimitPushdown::new().optimize(global_limit, &ConfigOptions::new())?; + + insta::assert_snapshot!( + format_plan(&optimized), + @r" + CoalescePartitionsExec: fetch=5 + UnionExec + TestScan: fetch=3 + TestScan: fetch=3 + " + ); + + Ok(()) +} + +#[test] +fn keeps_global_offset_limit_above_local_limit_on_multi_partition_union() -> Result<()> { + let schema = create_schema(); + let left = + Arc::new(TestScan::new(Arc::clone(&schema), vec![]).with_supports_fetch(true)); + let right = Arc::new(TestScan::new(schema, vec![]).with_supports_fetch(true)); + let union = UnionExec::try_new(vec![left, right])?; + let local_limit = local_limit_exec(union, 3); + let global_limit = global_limit_exec(local_limit, 2, Some(5)); + + let optimized = LimitPushdown::new().optimize(global_limit, &ConfigOptions::new())?; + + insta::assert_snapshot!( + format_plan(&optimized), + @r" + GlobalLimitExec: skip=2, fetch=5 + CoalescePartitionsExec: fetch=7 + UnionExec + TestScan: fetch=3 + TestScan: fetch=3 + " + ); + + Ok(()) +} + fn join_on_columns( left_col: &str, right_col: &str, @@ -180,8 +594,9 @@ fn join_on_columns( fn absorbs_limit_into_hash_join_inner() -> Result<()> { // HashJoinExec with Inner join should absorb limit via with_fetch let schema = create_schema(); - let left = empty_exec(Arc::clone(&schema)); - let right = empty_exec(Arc::clone(&schema)); + let left = + Arc::new(TestScan::new(Arc::clone(&schema), vec![]).with_supports_fetch(true)); + let right = Arc::new(TestScan::new(schema, vec![]).with_supports_fetch(true)); let on = join_on_columns("c1", "c1"); let hash_join = hash_join_exec(left, right, on, None, &JoinType::Inner)?; let global_limit = global_limit_exec(hash_join, 0, Some(5)); @@ -192,8 +607,8 @@ fn absorbs_limit_into_hash_join_inner() -> Result<()> { @r" GlobalLimitExec: skip=0, fetch=5 HashJoinExec: mode=Partitioned, join_type=Inner, on=[(c1@0, c1@0)] - EmptyExec - EmptyExec + TestScan + TestScan " ); @@ -205,8 +620,8 @@ fn absorbs_limit_into_hash_join_inner() -> Result<()> { optimized, @r" HashJoinExec: mode=Partitioned, join_type=Inner, on=[(c1@0, c1@0)], fetch=5 - EmptyExec - EmptyExec + TestScan + TestScan " ); @@ -217,8 +632,9 @@ fn absorbs_limit_into_hash_join_inner() -> Result<()> { fn absorbs_limit_into_hash_join_right() -> Result<()> { // HashJoinExec with Right join should absorb limit via with_fetch let schema = create_schema(); - let left = empty_exec(Arc::clone(&schema)); - let right = empty_exec(Arc::clone(&schema)); + let left = + Arc::new(TestScan::new(Arc::clone(&schema), vec![]).with_supports_fetch(true)); + let right = Arc::new(TestScan::new(schema, vec![]).with_supports_fetch(true)); let on = join_on_columns("c1", "c1"); let hash_join = hash_join_exec(left, right, on, None, &JoinType::Right)?; let global_limit = global_limit_exec(hash_join, 0, Some(10)); @@ -229,8 +645,8 @@ fn absorbs_limit_into_hash_join_right() -> Result<()> { @r" GlobalLimitExec: skip=0, fetch=10 HashJoinExec: mode=Partitioned, join_type=Right, on=[(c1@0, c1@0)] - EmptyExec - EmptyExec + TestScan + TestScan " ); @@ -242,8 +658,8 @@ fn absorbs_limit_into_hash_join_right() -> Result<()> { optimized, @r" HashJoinExec: mode=Partitioned, join_type=Right, on=[(c1@0, c1@0)], fetch=10 - EmptyExec - EmptyExec + TestScan + TestScan " ); @@ -254,8 +670,9 @@ fn absorbs_limit_into_hash_join_right() -> Result<()> { fn absorbs_limit_into_hash_join_left() -> Result<()> { // during probing, then unmatched rows at the end, stopping when limit is reached let schema = create_schema(); - let left = empty_exec(Arc::clone(&schema)); - let right = empty_exec(Arc::clone(&schema)); + let left = + Arc::new(TestScan::new(Arc::clone(&schema), vec![]).with_supports_fetch(true)); + let right = Arc::new(TestScan::new(schema, vec![]).with_supports_fetch(true)); let on = join_on_columns("c1", "c1"); let hash_join = hash_join_exec(left, right, on, None, &JoinType::Left)?; let global_limit = global_limit_exec(hash_join, 0, Some(5)); @@ -266,8 +683,8 @@ fn absorbs_limit_into_hash_join_left() -> Result<()> { @r" GlobalLimitExec: skip=0, fetch=5 HashJoinExec: mode=Partitioned, join_type=Left, on=[(c1@0, c1@0)] - EmptyExec - EmptyExec + TestScan + TestScan " ); @@ -279,8 +696,8 @@ fn absorbs_limit_into_hash_join_left() -> Result<()> { optimized, @r" HashJoinExec: mode=Partitioned, join_type=Left, on=[(c1@0, c1@0)], fetch=5 - EmptyExec - EmptyExec + TestScan + TestScan " ); diff --git a/datafusion/physical-optimizer/src/limit_pushdown.rs b/datafusion/physical-optimizer/src/limit_pushdown.rs index 01a288f7f1632..a6780ab309bc1 100644 --- a/datafusion/physical-optimizer/src/limit_pushdown.rs +++ b/datafusion/physical-optimizer/src/limit_pushdown.rs @@ -91,12 +91,21 @@ pub struct LimitPushdown {} /// For example: If the plan is satisfied with current fetch info, we decide to not add a LocalLimit /// /// [`LimitPushdown`]: crate::limit_pushdown::LimitPushdown -#[derive(Default, Clone, Debug)] +#[derive(Clone, Debug)] pub struct GlobalRequirements { fetch: Option, skip: usize, - satisfied: bool, preserve_order: bool, + status: LimitStatus, +} + +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] +enum LimitStatus { + #[default] + None, + PendingLocal, + PendingGlobal, + Enforced, } impl LimitPushdown { @@ -115,8 +124,8 @@ impl PhysicalOptimizerRule for LimitPushdown { let global_state = GlobalRequirements { fetch: None, skip: 0, - satisfied: false, preserve_order: false, + status: LimitStatus::None, }; pushdown_limits(plan, global_state) } @@ -130,13 +139,6 @@ impl PhysicalOptimizerRule for LimitPushdown { } } -struct LimitInfo { - input: Arc, - fetch: Option, - skip: usize, - preserve_order: bool, -} - /// This function is the main helper function of the `LimitPushDown` rule. /// The helper takes an `ExecutionPlan` and a global (algorithm) state which is /// an instance of `GlobalRequirements` and modifies these parameters while @@ -148,46 +150,97 @@ pub fn pushdown_limit_helper( mut pushdown_plan: Arc, mut global_state: GlobalRequirements, ) -> Result<(Transformed>, GlobalRequirements)> { - // Extract limit, if exist, and return child inputs. - if let Some(limit_info) = extract_limit(&pushdown_plan) { - // If we have fetch/skip info in the global state already, we need to - // decide which one to continue with: - let (skip, fetch) = combine_limit( + if global_state.status == LimitStatus::PendingLocal + && pushdown_plan.output_partitioning().partition_count() == 1 + { + global_state.status = LimitStatus::PendingGlobal; + } + + if let Some(global_limit) = pushdown_plan.downcast_ref::() + && global_limit.skip() == 0 + && global_limit.fetch().is_none() + { + return Ok(( + Transformed { + data: Arc::clone(global_limit.input()), + transformed: true, + tnr: TreeNodeRecursion::Stop, + }, + global_state, + )); + } + + if global_state.status == LimitStatus::PendingGlobal + && pushdown_plan.output_partitioning().partition_count() > 1 + { + let hint = global_state.fetch.map(|fetch| fetch + global_state.skip); + if let Some(hint) = hint { + let hint = pushdown_plan.fetch().map_or(hint, |fetch| fetch.min(hint)); + if pushdown_plan.fetch() != Some(hint) + && let Some(plan_with_fetch) = pushdown_plan.with_fetch(Some(hint)) + { + pushdown_plan = plan_with_fetch; + } + } + + let plan = materialize_global_requirement( + pushdown_plan, global_state.skip, global_state.fetch, - limit_info.skip, - limit_info.fetch, + global_state.preserve_order, ); - global_state.skip = skip; - global_state.fetch = fetch; - global_state.preserve_order = limit_info.preserve_order; - global_state.satisfied = false; + global_state.fetch = hint; + global_state.skip = 0; + global_state.status = LimitStatus::Enforced; + return Ok((Transformed::yes(plan), global_state)); + } - if let Some(fetch) = fetch - && limit_satisfied_by_input(&limit_info.input, skip, fetch)? + if let Some(global_limit) = pushdown_plan.downcast_ref::() { + let input = Arc::clone(global_limit.input()); + let skip = global_limit.skip(); + let fetch = global_limit.fetch(); + + (global_state.skip, global_state.fetch) = + combine_limit(global_state.skip, global_state.fetch, skip, fetch); + global_state.preserve_order |= global_limit.required_ordering().is_some(); + global_state.status = LimitStatus::PendingGlobal; + if let Some(fetch) = global_state.fetch + && limit_satisfied_by_input(&input, global_state.skip, fetch)? { - // The input already produces at most `fetch` rows, so no new limit - // node is needed. Mark satisfied so downstream won't re-add one, - // but preserve skip/fetch so any nested limit nodes (e.g. an inner - // GlobalLimitExec) can still be merged with the outer constraint. - global_state.satisfied = true; - - return Ok(( - Transformed { - data: limit_info.input, - transformed: true, - tnr: TreeNodeRecursion::Stop, - }, - global_state, - )); + global_state.status = LimitStatus::Enforced; } + return Ok(( + Transformed { + data: input, + transformed: true, + tnr: TreeNodeRecursion::Stop, + }, + global_state, + )); + } - // Now the global state has the most recent information, we can remove - // the limit node. We will decide later if we should add it again or - // not. + if let Some(local_limit) = pushdown_plan.downcast_ref::() { + let input = Arc::clone(local_limit.input()); + (global_state.skip, global_state.fetch) = combine_limit( + global_state.skip, + global_state.fetch, + 0, + Some(local_limit.fetch()), + ); + global_state.preserve_order |= local_limit.required_ordering().is_some(); + global_state.status = if input.output_partitioning().partition_count() == 1 { + LimitStatus::PendingGlobal + } else { + LimitStatus::PendingLocal + }; + if let Some(fetch) = global_state.fetch + && limit_satisfied_by_input(&input, global_state.skip, fetch)? + { + global_state.status = LimitStatus::Enforced; + } return Ok(( Transformed { - data: limit_info.input, + data: input, transformed: true, tnr: TreeNodeRecursion::Stop, }, @@ -199,7 +252,7 @@ pub fn pushdown_limit_helper( // state as necessary: if pushdown_plan.fetch().is_some() { if global_state.skip == 0 { - global_state.satisfied = true; + global_state.status = LimitStatus::Enforced; } (global_state.skip, global_state.fetch) = combine_limit( global_state.skip, @@ -211,17 +264,11 @@ pub fn pushdown_limit_helper( let Some(global_fetch) = global_state.fetch else { // There's no valid fetch information, exit early: - return if global_state.skip > 0 && !global_state.satisfied { + return if global_state.skip > 0 && global_state.status != LimitStatus::Enforced { // There might be a case with only offset, if so add a global limit: - global_state.satisfied = true; - Ok(( - Transformed::yes(add_global_limit( - pushdown_plan, - global_state.skip, - None, - )), - global_state, - )) + let new_plan = add_global_limit(pushdown_plan, global_state.skip, None); + global_state.status = LimitStatus::Enforced; + Ok((Transformed::yes(new_plan), global_state)) } else { // There's no info on offset or fetch, nothing to do: Ok((Transformed::no(pushdown_plan), global_state)) @@ -241,28 +288,22 @@ pub fn pushdown_limit_helper( // with the information from the global state. let mut new_plan = plan_with_fetch; // Execution plans can't (yet) handle skip, so if we have one, - // we still need to add a global limit + // we still need to add a global limit. if global_state.skip > 0 { new_plan = add_global_limit(new_plan, global_state.skip, global_state.fetch); } global_state.fetch = skip_and_fetch; global_state.skip = 0; - global_state.satisfied = true; + global_state.status = LimitStatus::Enforced; Ok((Transformed::yes(new_plan), global_state)) - } else if global_state.satisfied { + } else if global_state.status == LimitStatus::Enforced { // If the plan is already satisfied, do not add a limit: Ok((Transformed::no(pushdown_plan), global_state)) } else { - global_state.satisfied = true; - Ok(( - Transformed::yes(add_limit( - pushdown_plan, - global_state.skip, - global_fetch, - )), - global_state, - )) + let new_plan = add_limit(pushdown_plan, global_state.skip, global_fetch); + global_state.status = LimitStatus::Enforced; + Ok((Transformed::yes(new_plan), global_state)) } } else { // The plan does not support push down and it is not a limit. We will need @@ -275,7 +316,7 @@ pub fn pushdown_limit_helper( global_state.skip = 0; let maybe_fetchable = pushdown_plan.with_fetch(skip_and_fetch); - if global_state.satisfied { + if global_state.status == LimitStatus::Enforced { if let Some(plan_with_fetch) = maybe_fetchable { let plan_with_preserve_order = plan_with_fetch .with_preserve_order(global_state.preserve_order) @@ -285,7 +326,6 @@ pub fn pushdown_limit_helper( Ok((Transformed::no(pushdown_plan), global_state)) } } else { - global_state.satisfied = true; pushdown_plan = if let Some(plan_with_fetch) = maybe_fetchable { let plan_with_preserve_order = plan_with_fetch .with_preserve_order(global_state.preserve_order) @@ -303,6 +343,7 @@ pub fn pushdown_limit_helper( } else { add_limit(pushdown_plan, global_skip, global_fetch) }; + global_state.status = LimitStatus::Enforced; Ok((Transformed::yes(pushdown_plan), global_state)) } } @@ -382,7 +423,7 @@ pub(crate) fn pushdown_limits( // subtrees should not inherit its `skip`. Keep `fetch`, but clear // `skip` before recursing so child-local limits are not merged with // an `OFFSET` that has already been applied. - if global_state.satisfied { + if global_state.status == LimitStatus::Enforced { global_state.skip = 0; } @@ -409,27 +450,6 @@ pub(crate) fn pushdown_limits( } } -/// Extracts limit information from the [`ExecutionPlan`] if it is a -/// [`GlobalLimitExec`] or a [`LocalLimitExec`]. -fn extract_limit(plan: &Arc) -> Option { - if let Some(global_limit) = plan.downcast_ref::() { - Some(LimitInfo { - input: Arc::clone(global_limit.input()), - fetch: global_limit.fetch(), - skip: global_limit.skip(), - preserve_order: global_limit.required_ordering().is_some(), - }) - } else { - plan.downcast_ref::() - .map(|local_limit| LimitInfo { - input: Arc::clone(local_limit.input()), - fetch: Some(local_limit.fetch()), - skip: 0, - preserve_order: local_limit.required_ordering().is_some(), - }) - } -} - /// Checks if the given plan combines input partitions. fn combines_input_partitions(plan: &Arc) -> bool { plan.is::() || plan.is::() @@ -449,6 +469,38 @@ fn add_limit( } } +/// Materializes a global requirement at a single-partition boundary. A fetch +/// on a multi-partition plan is only a per-partition hint, so it must be +/// followed by a partition combiner before the requirement is satisfied. +fn materialize_global_requirement( + pushdown_plan: Arc, + skip: usize, + fetch: Option, + preserve_order: bool, +) -> Arc { + if pushdown_plan.output_partitioning().partition_count() == 1 { + return add_global_limit(pushdown_plan, skip, fetch); + } + + let skip_and_fetch = fetch.map(|fetch| fetch + skip); + let limited: Arc = if preserve_order + && let Some(ordering) = pushdown_plan.output_ordering().cloned() + { + Arc::new( + SortPreservingMergeExec::new(ordering, pushdown_plan) + .with_fetch(skip_and_fetch), + ) + } else { + Arc::new(CoalescePartitionsExec::new(pushdown_plan).with_fetch(skip_and_fetch)) + }; + + if skip > 0 { + add_global_limit(limited, skip, fetch) + } else { + limited + } +} + /// Adds a global limit to the plan. fn add_global_limit( pushdown_plan: Arc,