diff --git a/datafusion/core/tests/physical_optimizer/ensure_requirements.rs b/datafusion/core/tests/physical_optimizer/ensure_requirements.rs index 06b57a8630553..fe73518bf17d2 100644 --- a/datafusion/core/tests/physical_optimizer/ensure_requirements.rs +++ b/datafusion/core/tests/physical_optimizer/ensure_requirements.rs @@ -1531,3 +1531,32 @@ fn test_collect_left_join_keeps_hash_partitioned_build_side_coalesce() -> Result Ok(()) } + +// ======================================================================== +// Limits with a `skip` +// ======================================================================== + +/// A sort pushed below `GlobalLimitExec` with a non-zero `skip` must ask its +/// input for `skip + fetch` rows; asking for only `fetch` rows used to leave +/// `LIMIT 10 OFFSET 5` with 5 result rows. +/// +/// This checks a single pass only: the sort is inserted by `pushdown_sorts`, +/// which runs after `parallelize_sorts`, so a second pass would additionally +/// parallelize it into `SortPreservingMergeExec` + partitioned `SortExec`. +#[test] +fn test_sort_pushed_below_limit_with_skip_keeps_skip_rows() -> Result<()> { + let source = Arc::new(MockMultiPartitionExec::new(4)); + let coalesce = Arc::new(CoalescePartitionsExec::new(source)); + let limit = Arc::new(GlobalLimitExec::new(coalesce, 5, Some(10))); + let sort: Arc = + Arc::new(SortExec::new(sort_expr_on("a", 0, true, true), limit)); + + let optimized = optimize_and_sanity_check(sort)?; + assert_snapshot!(plan_string(&optimized), @r" + GlobalLimitExec: skip=5, fetch=10 + SortExec: TopK(fetch=15), expr=[a@0 DESC], preserve_partitioning=[false] + CoalescePartitionsExec + MockMultiPartitionExec + "); + Ok(()) +} diff --git a/datafusion/physical-optimizer/src/ensure_requirements/enforce_sorting/sort_pushdown.rs b/datafusion/physical-optimizer/src/ensure_requirements/enforce_sorting/sort_pushdown.rs index 19faf8ea9fa43..b80fd16bc8b83 100644 --- a/datafusion/physical-optimizer/src/ensure_requirements/enforce_sorting/sort_pushdown.rs +++ b/datafusion/physical-optimizer/src/ensure_requirements/enforce_sorting/sort_pushdown.rs @@ -79,6 +79,22 @@ impl Default for ParentRequirements { pub type SortPushDown = PlanContext; +/// Number of input rows `plan` needs from its children in order to produce +/// its own `fetch` rows. This is `plan.fetch()` for every operator except +/// [`GlobalLimitExec`], which discards `skip` rows first and therefore needs +/// `skip + fetch` input rows. Using the bare `fetch` there would turn +/// `LIMIT 10 OFFSET 5` into `TopK(10)` below the limit, i.e. 5 result rows. +/// +/// Note this is distinct from the fetch a parent imposes on `plan`'s *output* +/// (`ParentRequirements::fetch`), for which `plan.fetch()` is the right bound. +fn input_fetch(plan: &Arc) -> Option { + let fetch = plan.fetch()?; + let skip = plan + .downcast_ref::() + .map_or(0, |limit| limit.skip()); + Some(fetch + skip) +} + /// Assigns the ordering requirement of the root node to the its children. pub fn assign_initial_requirements(sort_push_down: &mut SortPushDown) { let reqs = sort_push_down.plan.required_input_ordering(); @@ -364,7 +380,7 @@ fn pushdown_sorts_helper( // For operators that can take a sort pushdown, continue with updated // requirements. If this node already outputs single partition (e.g. SPM), // don't push SinglePartition to children. - let current_fetch = sort_push_down.plan.fetch(); + let current_fetch = input_fetch(&sort_push_down.plan); let dists = sort_push_down .plan .input_distribution_requirements() diff --git a/datafusion/sqllogictest/test_files/limit.slt b/datafusion/sqllogictest/test_files/limit.slt index e17f633bc64ad..a7f6782b49348 100644 --- a/datafusion/sqllogictest/test_files/limit.slt +++ b/datafusion/sqllogictest/test_files/limit.slt @@ -339,6 +339,29 @@ SELECT COUNT(*) FROM (SELECT a FROM t1 LIMIT 3 OFFSET 8); ---- 2 +# A sort above LIMIT ... OFFSET is pushed below the limit as a TopK. The TopK +# has to keep `skip + fetch` rows (3 + 4 = 7) so that the limit can still skip +# 3 rows and return 4; keeping only `fetch` rows returned a single row. +query TT +EXPLAIN SELECT * FROM (SELECT a FROM t1 LIMIT 4 OFFSET 3) ORDER BY a; +---- +logical_plan +01)Sort: t1.a ASC NULLS LAST +02)--Limit: skip=3, fetch=4 +03)----TableScan: t1 projection=[a], fetch=7 +physical_plan +01)GlobalLimitExec: skip=3, fetch=4 +02)--SortExec: TopK(fetch=7), expr=[a@0 ASC NULLS LAST], preserve_partitioning=[false] +03)----DataSourceExec: partitions=1, partition_sizes=[1] + +query I +SELECT * FROM (SELECT a FROM t1 LIMIT 4 OFFSET 3) ORDER BY a; +---- +4 +5 +6 +7 + # The aggregate does not need to be computed because the input statistics are exact and # an OFFSET, but no LIMIT, is specified. query TT