diff --git a/datafusion/core/tests/physical_optimizer/limit_pushdown.rs b/datafusion/core/tests/physical_optimizer/limit_pushdown.rs index b8ebc80348134..ecce63dab0e28 100644 --- a/datafusion/core/tests/physical_optimizer/limit_pushdown.rs +++ b/datafusion/core/tests/physical_optimizer/limit_pushdown.rs @@ -15,30 +15,39 @@ // 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_common::tree_node::TreeNodeRecursion; +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 +112,292 @@ 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 apply_expressions( + &self, + _f: &mut dyn FnMut(&PhysicalExprRef) -> Result, + ) -> Result { + // `TestCombinerExec` owns no `PhysicalExpr`s. + Ok(TreeNodeRecursion::Continue) + } + + 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 plan that reports a fixed `fetch` but cannot change it through +/// `with_fetch`. It can optionally allow limits to be pushed to its child. +#[derive(Debug)] +struct TestFetchOnlyExec { + input: Arc, + fetch: Option, + supports_limit_pushdown: bool, + properties: Arc, +} + +impl TestFetchOnlyExec { + fn new(input: Arc, fetch: Option) -> Self { + let properties = PlanProperties::new( + EquivalenceProperties::new(input.schema()), + Partitioning::UnknownPartitioning(1), + EmissionType::Incremental, + Boundedness::Bounded, + ); + Self { + input, + fetch, + supports_limit_pushdown: false, + properties: Arc::new(properties), + } + } + + /// Set whether limits may be pushed through this operator to its child. + fn with_supports_limit_pushdown(mut self, supports: bool) -> Self { + self.supports_limit_pushdown = supports; + self + } +} + +impl DisplayAs for TestFetchOnlyExec { + fn fmt_as(&self, _t: DisplayFormatType, f: &mut Formatter) -> std::fmt::Result { + write!(f, "TestFetchOnlyExec")?; + if let Some(fetch) = self.fetch { + write!(f, ": fetch={fetch}")?; + } + Ok(()) + } +} + +impl ExecutionPlan for TestFetchOnlyExec { + fn name(&self) -> &str { + "TestFetchOnlyExec" + } + + fn properties(&self) -> &Arc { + &self.properties + } + + fn children(&self) -> Vec<&Arc> { + vec![&self.input] + } + + fn apply_expressions( + &self, + _f: &mut dyn FnMut(&PhysicalExprRef) -> Result, + ) -> Result { + // `TestFetchOnlyExec` owns no `PhysicalExpr`s. + Ok(TreeNodeRecursion::Continue) + } + + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + assert_eq!(children.len(), 1); + Ok(Arc::new( + Self::new(children[0].clone(), self.fetch) + .with_supports_limit_pushdown(self.supports_limit_pushdown), + )) + } + + fn execute( + &self, + _partition: usize, + _context: Arc, + ) -> Result { + unreachable!("TestFetchOnlyExec 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 fetch(&self) -> Option { + self.fetch + } + + fn supports_limit_pushdown(&self) -> bool { + self.supports_limit_pushdown + } +} + +/// Test multi-child plan with a single output partition that allows limit +/// pushdown. Optionally absorbs a fetch via `with_fetch`. +#[derive(Debug, Clone)] +struct TestMultiChildExec { + inputs: Vec>, + properties: Arc, + supports_fetch: bool, + fetch: Option, +} + +impl TestMultiChildExec { + fn new(inputs: Vec>) -> Self { + let properties = PlanProperties::new( + EquivalenceProperties::new(inputs[0].schema()), + Partitioning::UnknownPartitioning(1), + EmissionType::Incremental, + Boundedness::Bounded, + ); + Self { + inputs, + properties: Arc::new(properties), + supports_fetch: false, + fetch: None, + } + } + + /// Set whether `with_fetch()` returns `Some` (true) or `None` (false). + fn with_supports_fetch(mut self, supports: bool) -> Self { + self.supports_fetch = supports; + self + } +} + +impl DisplayAs for TestMultiChildExec { + fn fmt_as(&self, _t: DisplayFormatType, f: &mut Formatter) -> std::fmt::Result { + write!(f, "TestMultiChildExec")?; + if let Some(fetch) = self.fetch { + write!(f, ": fetch={fetch}")?; + } + Ok(()) + } +} + +impl ExecutionPlan for TestMultiChildExec { + fn name(&self) -> &str { + "TestMultiChildExec" + } + + fn properties(&self) -> &Arc { + &self.properties + } + + fn children(&self) -> Vec<&Arc> { + self.inputs.iter().collect() + } + + fn apply_expressions( + &self, + _f: &mut dyn FnMut(&PhysicalExprRef) -> Result, + ) -> Result { + // `TestMultiChildExec` owns no `PhysicalExpr`s. + Ok(TreeNodeRecursion::Continue) + } + + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + assert_eq!(children.len(), self.inputs.len()); + let mut new_plan = Self::new(children).with_supports_fetch(self.supports_fetch); + new_plan.fetch = self.fetch; + Ok(Arc::new(new_plan)) + } + + fn execute( + &self, + _partition: usize, + _context: Arc, + ) -> Result { + unreachable!("TestMultiChildExec 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 + } + + fn with_fetch(&self, fetch: Option) -> Option> { + if self.supports_fetch { + let mut new_plan = self.clone(); + new_plan.fetch = fetch; + Some(Arc::new(new_plan)) + } else { + None + } + } + + fn fetch(&self) -> Option { + self.fetch + } +} + #[test] fn transforms_streaming_table_exec_into_fetching_version_when_skip_is_zero() -> Result<()> { @@ -162,6 +457,448 @@ 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 materializes_global_limit_before_multi_child_extension() -> Result<()> { + // Regression test: a pending global limit used to be cloned to every + // child of a multi-child node, so each child applied the full LIMIT and + // the merged output exceeded it. The limit must stay above the node. + 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 custom = Arc::new(TestMultiChildExec::new(vec![left, right])); + let global_limit = global_limit_exec(custom, 0, Some(5)); + + let optimized = LimitPushdown::new().optimize(global_limit, &ConfigOptions::new())?; + + insta::assert_snapshot!( + format_plan(&optimized), + @r" + GlobalLimitExec: skip=0, fetch=5 + TestMultiChildExec + TestScan: fetch=5 + TestScan: fetch=5 + " + ); + + Ok(()) +} + +#[test] +fn materializes_global_offset_limit_before_multi_child_extension() -> Result<()> { + // The offset stays in the GlobalLimitExec; children only get a fetch hint + // of skip + fetch for early stopping. + 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 custom = Arc::new(TestMultiChildExec::new(vec![left, right])); + let global_limit = global_limit_exec(custom, 2, Some(5)); + + let optimized = LimitPushdown::new().optimize(global_limit, &ConfigOptions::new())?; + + insta::assert_snapshot!( + format_plan(&optimized), + @r" + GlobalLimitExec: skip=2, fetch=5 + TestMultiChildExec + TestScan: fetch=7 + TestScan: fetch=7 + " + ); + + Ok(()) +} + +#[test] +fn multi_child_extension_absorbs_global_limit_and_hints_children() -> Result<()> { + // When the multi-child node absorbs the fetch itself, no extra limit is + // needed; children still receive the same fetch for early stopping. + 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 custom = + Arc::new(TestMultiChildExec::new(vec![left, right]).with_supports_fetch(true)); + let global_limit = global_limit_exec(custom, 0, Some(5)); + + let optimized = LimitPushdown::new().optimize(global_limit, &ConfigOptions::new())?; + + insta::assert_snapshot!( + format_plan(&optimized), + @r" + TestMultiChildExec: fetch=5 + TestScan: fetch=5 + TestScan: fetch=5 + " + ); + + Ok(()) +} + +#[test] +fn materializes_local_limit_before_multi_child_extension() -> Result<()> { + // A local limit also cannot be replicated to every child of a multi-child + // node with a single output partition. + 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 custom = Arc::new(TestMultiChildExec::new(vec![left, right])); + let local_limit = local_limit_exec(custom, 5); + + let optimized = LimitPushdown::new().optimize(local_limit, &ConfigOptions::new())?; + + insta::assert_snapshot!( + format_plan(&optimized), + @r" + GlobalLimitExec: skip=0, fetch=5 + TestMultiChildExec + 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 +917,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 +930,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 +943,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 +955,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 +968,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 +981,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 +993,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 +1006,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 +1019,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 " ); @@ -834,3 +1574,125 @@ fn outer_offset_with_same_sort_key_still_pushes_limit() -> Result<()> { Ok(()) } + +#[test] +fn keeps_global_limit_when_existing_fetch_is_looser_than_owed() -> Result<()> { + // This operator's `fetch=10` is weaker than `LIMIT 5` and cannot be lowered. + // Keep `GlobalLimitExec` so the query still returns at most five rows. + let schema = create_schema(); + let scan = Arc::new(TestScan::new(schema, vec![])); + let fetch_only = Arc::new(TestFetchOnlyExec::new(scan, Some(10))); + let global_limit = global_limit_exec(fetch_only, 0, Some(5)); + + let optimized = LimitPushdown::new().optimize(global_limit, &ConfigOptions::new())?; + + insta::assert_snapshot!( + format_plan(&optimized), + @r" + GlobalLimitExec: skip=0, fetch=5 + TestFetchOnlyExec: fetch=10 + TestScan + " + ); + + Ok(()) +} + +#[test] +fn pushes_owed_limit_below_fetch_only_unary_when_limit_pushdown_supported() -> Result<()> +{ + // This operator allows limit pushdown but cannot lower its own `fetch` from + // 10 to 5. Its child cannot accept a fetch either, so keep + // `GlobalLimitExec(fetch=5)` between them. + let schema = create_schema(); + let scan = Arc::new(TestScan::new(schema, vec![])); + let fetch_only = Arc::new( + TestFetchOnlyExec::new(scan, Some(10)).with_supports_limit_pushdown(true), + ); + let global_limit = global_limit_exec(fetch_only, 0, Some(5)); + + let optimized = LimitPushdown::new().optimize(global_limit, &ConfigOptions::new())?; + + insta::assert_snapshot!( + format_plan(&optimized), + @r" + TestFetchOnlyExec: fetch=10 + GlobalLimitExec: skip=0, fetch=5 + TestScan + " + ); + + Ok(()) +} + +#[test] +fn does_not_add_redundant_wrapper_when_existing_fetch_is_tighter_than_owed() -> Result<()> +{ + // `fetch=3` is stricter than `LIMIT 5`, so no additional limit is needed. + let schema = create_schema(); + let scan = Arc::new(TestScan::new(schema, vec![])); + let fetch_only = Arc::new(TestFetchOnlyExec::new(scan, Some(3))); + let global_limit = global_limit_exec(fetch_only, 0, Some(5)); + + let optimized = LimitPushdown::new().optimize(global_limit, &ConfigOptions::new())?; + + insta::assert_snapshot!( + format_plan(&optimized), + @r" + TestFetchOnlyExec: fetch=3 + TestScan + " + ); + + Ok(()) +} + +#[test] +fn tightens_existing_sort_fetch_to_owed_limit() -> Result<()> { + // `SortExec` can lower its `fetch` from 10 to 5, so the separate + // `GlobalLimitExec` is unnecessary. + let schema = create_schema(); + let scan = Arc::new(TestScan::new(schema.clone(), vec![])); + let ordering: LexOrdering = [PhysicalSortExpr { + expr: col("c1", &schema)?, + options: SortOptions::default(), + }] + .into(); + let sort = sort_exec(ordering, scan).with_fetch(Some(10)).unwrap(); + let global_limit = global_limit_exec(sort, 0, Some(5)); + + let optimized = LimitPushdown::new().optimize(global_limit, &ConfigOptions::new())?; + + insta::assert_snapshot!( + format_plan(&optimized), + @r" + SortExec: TopK(fetch=5), expr=[c1@0 ASC], preserve_partitioning=[false] + TestScan + " + ); + + Ok(()) +} + +#[test] +fn keeps_global_offset_limit_when_existing_fetch_is_looser() -> Result<()> { + // An operator `fetch` cannot apply `OFFSET 2`; keep `GlobalLimitExec` to + // enforce both the offset and `LIMIT 5`. + let schema = create_schema(); + let scan = Arc::new(TestScan::new(schema, vec![])); + let fetch_only = Arc::new(TestFetchOnlyExec::new(scan, Some(10))); + let global_limit = global_limit_exec(fetch_only, 2, Some(5)); + + let optimized = LimitPushdown::new().optimize(global_limit, &ConfigOptions::new())?; + + insta::assert_snapshot!( + format_plan(&optimized), + @r" + GlobalLimitExec: skip=2, fetch=5 + TestFetchOnlyExec: fetch=10 + TestScan + " + ); + + Ok(()) +} diff --git a/datafusion/physical-optimizer/src/limit_pushdown.rs b/datafusion/physical-optimizer/src/limit_pushdown.rs index f88a2be14e984..b34bb34631a51 100644 --- a/datafusion/physical-optimizer/src/limit_pushdown.rs +++ b/datafusion/physical-optimizer/src/limit_pushdown.rs @@ -78,26 +78,40 @@ use datafusion_physical_plan::placeholder_row::PlaceholderRowExec; use datafusion_physical_plan::projection::ProjectionExec; use datafusion_physical_plan::sorts::sort_preserving_merge::SortPreservingMergeExec; use datafusion_physical_plan::statistics::{StatisticsArgs, StatisticsContext}; +use datafusion_physical_plan::union::UnionExec; use datafusion_physical_plan::{ExecutionPlan, ExecutionPlanProperties}; /// This rule inspects [`ExecutionPlan`]'s and pushes down the fetch limit from /// the parent to the child if applicable. #[derive(Default, Debug)] pub struct LimitPushdown {} -/// This is a "data class" we use within the [`LimitPushdown`] rule to push -/// down limits in the plan. GlobalRequirements are hold as a rule-wide state -/// and holds the fetch and skip information. The struct also has a field named -/// satisfied which means if the "current" plan is valid in terms of limits or not. +/// State carried through [`LimitPushdown`] while it pushes limits down the plan. /// -/// For example: If the plan is satisfied with current fetch info, we decide to not add a LocalLimit +/// `fetch` and `skip` hold the limit currently being pushed down. `pending` +/// says whether a local or global limit still has to be enforced. After it +/// becomes `None`, `fetch` may still be copied to children to help them stop +/// early, but it is no longer responsible for correctness. /// /// [`LimitPushdown`]: crate::limit_pushdown::LimitPushdown #[derive(Default, Clone, Debug)] pub struct GlobalRequirements { fetch: Option, skip: usize, - satisfied: bool, preserve_order: bool, + pending: Option, +} + +/// The scope of a row limit: what the limited row count applies to. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +enum LimitScope { + /// The limit caps each output partition independently, as enforced by + /// [`LocalLimitExec`]. + Local, + /// The limit caps the combined output of all partitions, as enforced by + /// [`GlobalLimitExec`] over a single-partition input. A fetch on a + /// multi-partition operator cannot satisfy this scope; the partitions + /// must first be merged into one stream. + Global, } impl LimitPushdown { @@ -116,8 +130,8 @@ impl PhysicalOptimizerRule for LimitPushdown { let global_state = GlobalRequirements { fetch: None, skip: 0, - satisfied: false, preserve_order: false, + pending: None, }; pushdown_limits(plan, global_state) } @@ -131,13 +145,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 @@ -149,46 +156,109 @@ 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.pending == Some(LimitScope::Local) + && pushdown_plan.output_partitioning().partition_count() == 1 + { + // A local limit on one output partition also limits the total output. + // Treat it as global before descending because this node may have + // multi-partition children. + global_state.pending = Some(LimitScope::Global); + } + + if let Some(global_limit) = pushdown_plan.downcast_ref::() + && global_limit.skip() == 0 + && global_limit.fetch().is_none() + { + // A `GlobalLimitExec` with no skip and no fetch enforces nothing, so + // remove it and keep the carried state unchanged. General limit handling + // would mark a global limit as pending even when none is needed. Keep the + // local-to-global promotion above: it remains valid because this node has + // one output partition. + return Ok(( + Transformed { + data: Arc::clone(global_limit.input()), + transformed: true, + tnr: TreeNodeRecursion::Stop, + }, + global_state, + )); + } + + if global_state.pending == Some(LimitScope::Global) + && pushdown_plan.output_partitioning().partition_count() > 1 + { + // Handle this before the generic `fetch` case: on a multi-partition plan, + // `fetch` limits partitions separately and cannot limit their combined output. + // Keep it for early stopping, then combine the partitions to enforce the + // global limit. + 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.pending = None; + 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.pending = Some(LimitScope::Global); + 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.pending = None; } + 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.pending = if input.output_partitioning().partition_count() == 1 { + Some(LimitScope::Global) + } else { + Some(LimitScope::Local) + }; + if let Some(fetch) = global_state.fetch + && limit_satisfied_by_input(&input, global_state.skip, fetch)? + { + global_state.pending = None; + } return Ok(( Transformed { - data: limit_info.input, + data: input, transformed: true, tnr: TreeNodeRecursion::Stop, }, @@ -196,33 +266,33 @@ pub fn pushdown_limit_helper( )); } - // If we have a non-limit operator with fetch capability, update global - // state as necessary: - if pushdown_plan.fetch().is_some() { - if global_state.skip == 0 { - global_state.satisfied = true; + // An existing `fetch` satisfies the carried limit only if there is no + // offset and it is no larger than the requested fetch. For example, + // `fetch=10` does not enforce `LIMIT 5`; if the operator cannot be changed to + // `fetch=5`, the explicit limit must remain. Still combine the fetch values + // below so descendants can use the tighter value to stop early. + if let Some(existing_fetch) = pushdown_plan.fetch() { + if global_state.skip == 0 + && let Some(required_fetch) = global_state.fetch + && existing_fetch <= required_fetch + { + global_state.pending = None; } (global_state.skip, global_state.fetch) = combine_limit( global_state.skip, global_state.fetch, 0, - pushdown_plan.fetch(), + Some(existing_fetch), ); } 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.pending.is_some() { // 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.pending = None; + 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)) @@ -232,43 +302,59 @@ pub fn pushdown_limit_helper( let skip_and_fetch = Some(global_fetch + global_state.skip); if pushdown_plan.supports_limit_pushdown() { + // A pending limit cannot be replicated to every child of a multi-child + // node: each child would enforce it and their merged output could + // exceed it. Only a unary node can pass it through transparently, plus + // `UnionExec` with `Local` (each output partition comes from one child). + let can_delegate_pending = global_state.pending.is_none() + || pushdown_plan.children().len() <= 1 + || (global_state.pending == Some(LimitScope::Local) + && pushdown_plan.is::()); + if !can_delegate_pending { + // Enforce the limit at this node's output, then let children stop + // early with the same fetch hint. + let new_plan = if global_state.skip > 0 { + add_limit(pushdown_plan, global_state.skip, global_fetch) + } else if let Some(plan_with_fetch) = pushdown_plan.with_fetch(skip_and_fetch) + { + plan_with_fetch + } else { + add_limit(pushdown_plan, 0, global_fetch) + }; + global_state.fetch = skip_and_fetch; + global_state.skip = 0; + global_state.pending = None; + return Ok((Transformed::yes(new_plan), global_state)); + } if !combines_input_partitions(&pushdown_plan) { // We have information in the global state and the plan pushes down, // continue: Ok((Transformed::no(pushdown_plan), global_state)) } else if let Some(plan_with_fetch) = pushdown_plan.with_fetch(skip_and_fetch) { - // This plan is combining input partitions, so we need to add the - // fetch info to plan if possible. If not, we must add a limit node - // with the information from the global state. + // This partition-combining operator accepted `fetch`, so the fetch now + // bounds its combined output. 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.pending = None; Ok((Transformed::yes(new_plan), global_state)) - } else if global_state.satisfied { - // If the plan is already satisfied, do not add a limit: + } else if global_state.pending.is_none() { + // The required limit is already enforced, so do not add another one. 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.pending = None; + Ok((Transformed::yes(new_plan), global_state)) } } else { - // The plan does not support push down and it is not a limit. We will need - // to add a limit or a fetch. If the plan is already satisfied, we will try - // to add the fetch info and return the plan. + // This operator stops limit pushdown. If the limit is already enforced, try + // only to add a fetch for early stopping; otherwise enforce the limit here. // There's no push down, change fetch & skip to default values: let global_skip = global_state.skip; @@ -276,7 +362,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.pending.is_none() { if let Some(plan_with_fetch) = maybe_fetchable { let plan_with_preserve_order = plan_with_fetch .with_preserve_order(global_state.preserve_order) @@ -286,7 +372,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) @@ -304,6 +389,7 @@ pub fn pushdown_limit_helper( } else { add_limit(pushdown_plan, global_skip, global_fetch) }; + global_state.pending = None; Ok((Transformed::yes(pushdown_plan), global_state)) } } @@ -379,11 +465,11 @@ pub(crate) fn pushdown_limits( (new_node, global_state) = pushdown_limit_helper(new_node.data, global_state)?; } - // Once a limit has been materialized above the current node, child - // 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 { + // Once the limit is enforced, clear `skip` before visiting children so it + // cannot combine with a nested limit and skip rows twice. Each child may + // receive the same `fetch` for early stopping; those values are not added + // together because the limit has already been enforced above. + if global_state.pending.is_none() { global_state.skip = 0; } @@ -410,27 +496,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::() @@ -450,6 +515,38 @@ fn add_limit( } } +/// Enforces one limit across all output partitions. Multiple partitions are +/// merged first, preserving order when required, because a fetch on each +/// partition cannot limit their combined output. +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,