diff --git a/datafusion/core/tests/physical_optimizer/limit_pushdown.rs b/datafusion/core/tests/physical_optimizer/limit_pushdown.rs index b8ebc80348134..06372f87ebd35 100644 --- a/datafusion/core/tests/physical_optimizer/limit_pushdown.rs +++ b/datafusion/core/tests/physical_optimizer/limit_pushdown.rs @@ -19,13 +19,16 @@ 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, + parquet_exec, parquet_exec_with_sort, sort_exec, sort_preserving_merge_exec, + sort_preserving_merge_exec_with_fetch, stream_exec, }; use arrow::compute::SortOptions; use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; use datafusion_common::config::ConfigOptions; use datafusion_common::error::Result; +use datafusion_datasource::file_scan_config::FileScanConfig; +use datafusion_datasource::source::DataSourceExec; use datafusion_expr::{JoinType, Operator}; use datafusion_physical_expr::Partitioning; use datafusion_physical_expr::expressions::{BinaryExpr, col, lit}; @@ -405,6 +408,64 @@ fn pushes_global_limit_into_multiple_fetch_plans() -> Result<()> { Ok(()) } +#[test] +fn preserves_order_when_pushing_fetch_from_sort_preserving_merge() -> Result<()> { + let schema = create_schema(); + let ordering: LexOrdering = [PhysicalSortExpr { + expr: col("c1", &schema)?, + options: SortOptions::default(), + }] + .into(); + let scan = parquet_exec_with_sort(schema, vec![ordering.clone()]); + let local_limit = local_limit_exec(scan, 10); + let plan = sort_preserving_merge_exec_with_fetch(ordering, local_limit, 5); + + let optimized = LimitPushdown::new().optimize(plan, &ConfigOptions::new())?; + let scan = optimized.children().swap_remove(0); + let scan = scan + .downcast_ref::() + .expect("fetch should be pushed into the scan"); + let config = scan + .data_source() + .downcast_ref::() + .expect("parquet scan should use FileScanConfig"); + + assert_eq!(config.limit, Some(5)); + assert!(config.preserve_order); + + Ok(()) +} + +#[test] +fn does_not_preserve_order_for_independent_limit_below_sort() -> Result<()> { + let schema = create_schema(); + let ordering: LexOrdering = [PhysicalSortExpr { + expr: col("c1", &schema)?, + options: SortOptions::default(), + }] + .into(); + let scan = parquet_exec(schema); + let local_limit = local_limit_exec(scan, 10); + let sort = sort_exec(ordering.clone(), local_limit); + let plan = sort_preserving_merge_exec_with_fetch(ordering, sort, 5); + + let optimized = LimitPushdown::new().optimize(plan, &ConfigOptions::new())?; + let sort = optimized.children().swap_remove(0); + let scan = sort.children().swap_remove(0); + let scan = scan + .downcast_ref::() + .expect("inner fetch should be pushed into the scan"); + let config = scan + .data_source() + .downcast_ref::() + .expect("parquet scan should use FileScanConfig"); + + assert_eq!(config.limit, Some(10)); + assert!(!config.preserve_order); + + Ok(()) +} + #[test] fn keeps_pushed_local_limit_exec_when_there_are_multiple_input_partitions() -> Result<()> { diff --git a/datafusion/physical-optimizer/src/limit_pushdown.rs b/datafusion/physical-optimizer/src/limit_pushdown.rs index 01a288f7f1632..a361801dc2614 100644 --- a/datafusion/physical-optimizer/src/limit_pushdown.rs +++ b/datafusion/physical-optimizer/src/limit_pushdown.rs @@ -160,7 +160,7 @@ pub fn pushdown_limit_helper( ); global_state.skip = skip; global_state.fetch = fetch; - global_state.preserve_order = limit_info.preserve_order; + global_state.preserve_order |= limit_info.preserve_order; global_state.satisfied = false; if let Some(fetch) = fetch @@ -196,8 +196,10 @@ pub fn pushdown_limit_helper( } // If we have a non-limit operator with fetch capability, update global - // state as necessary: + // state as necessary. A fetched ordered merge selects the leading rows in + // its ordering, so preserve that ordering when pushing the fetch down. if pushdown_plan.fetch().is_some() { + global_state.preserve_order |= pushdown_plan.is::(); if global_state.skip == 0 { global_state.satisfied = true; } @@ -269,36 +271,33 @@ pub fn pushdown_limit_helper( // 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. - // There's no push down, change fetch & skip to default values: + // This operator materializes the current fetch. Clear its requirements before + // descending so they are not combined with an independent limit below. let global_skip = global_state.skip; + let global_preserve_order = std::mem::take(&mut global_state.preserve_order); global_state.fetch = None; global_state.skip = 0; - let maybe_fetchable = pushdown_plan.with_fetch(skip_and_fetch); + let maybe_fetchable = pushdown_plan.with_fetch(skip_and_fetch).map(|plan| { + if global_preserve_order { + plan.with_preserve_order(true).unwrap_or(plan) + } else { + plan + } + }); if global_state.satisfied { if let Some(plan_with_fetch) = maybe_fetchable { - let plan_with_preserve_order = plan_with_fetch - .with_preserve_order(global_state.preserve_order) - .unwrap_or(plan_with_fetch); - Ok((Transformed::yes(plan_with_preserve_order), global_state)) + Ok((Transformed::yes(plan_with_fetch), global_state)) } else { 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) - .unwrap_or(plan_with_fetch); - if global_skip > 0 { - add_global_limit( - plan_with_preserve_order, - global_skip, - Some(global_fetch), - ) + add_global_limit(plan_with_fetch, global_skip, Some(global_fetch)) } else { - plan_with_preserve_order + plan_with_fetch } } else { add_limit(pushdown_plan, global_skip, global_fetch) diff --git a/datafusion/sqllogictest/test_files/limit_pruning.slt b/datafusion/sqllogictest/test_files/limit_pruning.slt index 4ef0b5c74f3e7..0bbdb489b8f36 100644 --- a/datafusion/sqllogictest/test_files/limit_pruning.slt +++ b/datafusion/sqllogictest/test_files/limit_pruning.slt @@ -113,6 +113,79 @@ select a from fully_matched_limit where a >= 3 order by a limit 4; statement ok drop table fully_matched_limit; +# A fetched SortPreservingMergeExec must preserve scan order when pushing down +# its limit. Otherwise, limit pruning can skip the first partially matched row +# group in favor of the following fully matched row group. +statement ok +CREATE TABLE ordered_limit_source(id INT, passes INT) AS VALUES + (1, 0), + (2, 1), + (3, 1), + (4, 1), + (5, 1), + (6, 1), + (101, 1), + (102, 1), + (103, 1); + +query I +COPY (SELECT * FROM ordered_limit_source WHERE id < 100 ORDER BY id) +TO 'test_files/scratch/limit_pruning/ordered_limit/a.parquet' +STORED AS PARQUET +OPTIONS ( + 'format.max_row_group_size' '3' +); +---- +6 + +query I +COPY (SELECT * FROM ordered_limit_source WHERE id >= 100 ORDER BY id) +TO 'test_files/scratch/limit_pruning/ordered_limit/b.parquet' +STORED AS PARQUET +OPTIONS ( + 'format.max_row_group_size' '3' +); +---- +3 + +statement ok +drop table ordered_limit_source; + +statement ok +set datafusion.execution.target_partitions = 2; + +statement ok +CREATE EXTERNAL TABLE ordered_limit(id INT, passes INT) +STORED AS PARQUET +LOCATION 'test_files/scratch/limit_pruning/ordered_limit/' +WITH ORDER (id ASC); + +statement ok +set datafusion.explain.physical_plan_only = true; + +query TT +explain select * from ordered_limit where passes = 1 order by id limit 3; +---- +physical_plan +01)SortPreservingMergeExec: [id@0 ASC NULLS LAST], fetch=3 +02)--DataSourceExec: limit=3, output_ordering=[id@0 ASC NULLS LAST], file_type=parquet + +query II +select * from ordered_limit where passes = 1 order by id limit 3; +---- +2 1 +3 1 +4 1 + +statement ok +drop table ordered_limit; + +statement ok +reset datafusion.explain.physical_plan_only; + +statement ok +set datafusion.execution.target_partitions = 4; + # limit_pruned_row_groups=0 total → 0 matched # because of order by, scan needs to preserve sort, so limit pruning is disabled query TT