Skip to content

Commit b308014

Browse files
committed
Initial work
1 parent 74e4974 commit b308014

3 files changed

Lines changed: 191 additions & 7 deletions

File tree

‎datafusion/core/tests/physical_optimizer/ensure_requirements.rs‎

Lines changed: 68 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1533,7 +1533,7 @@ fn test_collect_left_join_keeps_hash_partitioned_build_side_coalesce() -> Result
15331533
}
15341534

15351535
// ========================================================================
1536-
// Limits with a `skip`
1536+
// Sort pushdown through limits
15371537
// ========================================================================
15381538

15391539
/// Keep sorting above an unordered limit so its input can stop early and only
@@ -1556,10 +1556,32 @@ fn test_keep_sort_above_unordered_limit_with_skip() -> Result<()> {
15561556
Ok(())
15571557
}
15581558

1559-
/// Refining an existing ordering can still push a TopK below the limit. It must
1560-
/// retain `skip + fetch` rows so the offset does not reduce the result count.
1559+
/// `OFFSET` without `LIMIT` is a `GlobalLimitExec` with a `skip` and no
1560+
/// `fetch`. It selects rows by input position like any limit, so the sort
1561+
/// stays above it.
15611562
#[test]
1562-
fn test_sort_pushed_below_ordered_limit_with_skip_keeps_skip_rows() -> Result<()> {
1563+
fn test_keep_sort_above_limit_with_skip_only() -> Result<()> {
1564+
let source = Arc::new(MockMultiPartitionExec::new(4));
1565+
let coalesce = Arc::new(CoalescePartitionsExec::new(source));
1566+
let limit = Arc::new(GlobalLimitExec::new(coalesce, 5, None));
1567+
let sort: Arc<dyn ExecutionPlan> =
1568+
Arc::new(SortExec::new(sort_expr_on("a", 0, true, true), limit));
1569+
1570+
let optimized = optimize_and_sanity_check(sort)?;
1571+
assert_snapshot!(plan_string(&optimized), @r"
1572+
SortExec: expr=[a@0 DESC], preserve_partitioning=[false]
1573+
GlobalLimitExec: skip=5, fetch=None
1574+
CoalescePartitionsExec
1575+
MockMultiPartitionExec
1576+
");
1577+
Ok(())
1578+
}
1579+
1580+
/// A limit whose input ordering comes from the source, not a sort, keeps the
1581+
/// finer sort above it: a TopK on `[a, b]` below the limit would keep the
1582+
/// smallest rows by `[a, b]` instead of the first input rows.
1583+
#[test]
1584+
fn test_keep_sort_above_ordered_limit_with_skip() -> Result<()> {
15631585
let source = Arc::new(MockMultiPartitionExec::new(4));
15641586
let mut ordering = sort_expr_on("a", 0, false, false);
15651587
let ordered_input = Arc::new(MockReqExec::new(
@@ -1573,11 +1595,51 @@ fn test_sort_pushed_below_ordered_limit_with_skip_keeps_skip_rows() -> Result<()
15731595

15741596
let optimized = optimize_and_sanity_check(sort)?;
15751597
assert_snapshot!(plan_string(&optimized), @r"
1576-
GlobalLimitExec: skip=5, fetch=10
1577-
SortExec: TopK(fetch=15), expr=[a@0 ASC NULLS LAST, b@1 ASC NULLS LAST], preserve_partitioning=[false], sort_prefix=[a@0 ASC NULLS LAST]
1598+
SortExec: expr=[a@0 ASC NULLS LAST, b@1 ASC NULLS LAST], preserve_partitioning=[false]
1599+
GlobalLimitExec: skip=5, fetch=10
15781600
MockReqExec
15791601
SortPreservingMergeExec: [a@0 ASC NULLS LAST]
15801602
MockMultiPartitionExec
15811603
");
15821604
Ok(())
15831605
}
1606+
1607+
/// A limit directly above a `SortExec` is the exception: the finer requirement
1608+
/// replaces that sort's key rather than adding a sort below the limit. The TopK
1609+
/// retains `skip + fetch` rows so the offset does not shrink the result.
1610+
#[test]
1611+
fn test_sort_pushed_into_sort_below_limit_with_skip() -> Result<()> {
1612+
let source = Arc::new(MockMultiPartitionExec::new(1));
1613+
let mut ordering = sort_expr_on("b", 1, false, false);
1614+
let inner_sort = Arc::new(SortExec::new(ordering.clone(), source));
1615+
let limit = Arc::new(GlobalLimitExec::new(inner_sort, 5, Some(10)));
1616+
ordering.extend(sort_expr_on("a", 0, false, false));
1617+
let sort: Arc<dyn ExecutionPlan> = Arc::new(SortExec::new(ordering, limit));
1618+
1619+
let optimized = optimize_and_sanity_check(sort)?;
1620+
assert_snapshot!(plan_string(&optimized), @r"
1621+
GlobalLimitExec: skip=5, fetch=10
1622+
SortExec: TopK(fetch=15), expr=[b@1 ASC NULLS LAST, a@0 ASC NULLS LAST], preserve_partitioning=[false]
1623+
MockMultiPartitionExec
1624+
");
1625+
Ok(())
1626+
}
1627+
1628+
/// `LocalLimitExec` keeps the first `fetch` rows of each partition, so a finer
1629+
/// sort is not pushed below it either.
1630+
#[test]
1631+
fn test_keep_sort_above_ordered_local_limit() -> Result<()> {
1632+
let source = Arc::new(MockMultiPartitionExec::new(1));
1633+
let limit = Arc::new(LocalLimitExec::new(source, 10));
1634+
let mut ordering = sort_expr_on("a", 0, false, false);
1635+
ordering.extend(sort_expr_on("b", 1, false, false));
1636+
let sort: Arc<dyn ExecutionPlan> = Arc::new(SortExec::new(ordering, limit));
1637+
1638+
let optimized = optimize_and_sanity_check(sort)?;
1639+
assert_snapshot!(plan_string(&optimized), @r"
1640+
SortExec: expr=[a@0 ASC NULLS LAST, b@1 ASC NULLS LAST], preserve_partitioning=[false]
1641+
LocalLimitExec: fetch=10
1642+
MockMultiPartitionExec
1643+
");
1644+
Ok(())
1645+
}

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

Lines changed: 23 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -536,6 +536,27 @@ fn pushdown_requirement_to_children(
536536
)
537537
.then(|| Ok(vec![Some(parent_required)]))
538538
.transpose()
539+
} else if plan.is::<GlobalLimitExec>() || plan.is::<LocalLimitExec>() {
540+
// A limit selects rows by input position, so a sort pushed below it
541+
// would change which rows it skips and keeps. The exception is a limit
542+
// directly above a `SortExec` whose ordering the requirement refines:
543+
// the pushdown replaces that sort's key rather than adding a sort
544+
// below the limit, and only changes the tie-break of an explicit sort,
545+
// as pushing into a `SortExec` with a fetch already does. Otherwise
546+
// keep the sort above the limit, so the input can stop after
547+
// `skip + fetch` rows and only the retained rows are sorted.
548+
let refines_sort_below = plan.children()[0].is::<SortExec>()
549+
&& plan.output_ordering().is_some_and(|ordering| {
550+
plan.equivalence_properties().requirements_compatible(
551+
parent_required.first().clone(),
552+
ordering.clone().into(),
553+
)
554+
});
555+
if refines_sort_below {
556+
Ok(Some(vec![Some(parent_required)]))
557+
} else {
558+
Ok(None)
559+
}
539560
} else if plan.fetch().is_some()
540561
&& plan.supports_limit_pushdown()
541562
&& plan
@@ -567,7 +588,8 @@ fn pushdown_requirement_to_children(
567588
parent_required.clone()
568589
};
569590

570-
// Keep sorting above an unordered limit so the input can stop early.
591+
// With no output ordering to refine, a sort below this fetch would
592+
// change which rows it keeps, so keep the sort above it.
571593
let Some(ordering) = plan.properties().output_ordering() else {
572594
return Ok(None);
573595
};

‎datafusion/sqllogictest/test_files/limit.slt‎

Lines changed: 100 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -361,6 +361,106 @@ SELECT * FROM (SELECT a FROM t1 LIMIT 4 OFFSET 3) ORDER BY a;
361361
6
362362
7
363363

364+
# Which rows a limit keeps depends on its input order, so repeat the check on
365+
# unsorted data: `LIMIT 4 OFFSET 3` returns input rows 4 to 7 (1, 7, 2, 10),
366+
# not the 4th to 7th smallest values.
367+
statement ok
368+
CREATE TABLE unsorted (a INT) AS VALUES (5), (3), (9), (1), (7), (2), (10), (4), (8), (6);
369+
370+
query I
371+
SELECT * FROM (SELECT a FROM unsorted LIMIT 4 OFFSET 3) ORDER BY a;
372+
----
373+
1
374+
2
375+
7
376+
10
377+
378+
# `OFFSET` without `LIMIT` also selects rows by position, so the sort stays
379+
# above it: it skips the first 3 input rows (5, 3, 9), not the 3 smallest values.
380+
query TT
381+
EXPLAIN SELECT * FROM (SELECT a FROM unsorted OFFSET 3) ORDER BY a;
382+
----
383+
logical_plan
384+
01)Sort: unsorted.a ASC NULLS LAST
385+
02)--Limit: skip=3, fetch=None
386+
03)----TableScan: unsorted projection=[a]
387+
physical_plan
388+
01)SortExec: expr=[a@0 ASC NULLS LAST], preserve_partitioning=[false]
389+
02)--GlobalLimitExec: skip=3, fetch=None
390+
03)----DataSourceExec: partitions=1, partition_sizes=[1]
391+
392+
query I
393+
SELECT * FROM (SELECT a FROM unsorted OFFSET 3) ORDER BY a;
394+
----
395+
1
396+
2
397+
4
398+
6
399+
7
400+
8
401+
10
402+
403+
# The same holds when the input is ordered on a prefix of the requested
404+
# ordering: a `TopK` on `(a, b)` below the limit would keep the 7 smallest rows
405+
# by `(a, b)` instead of the first 7 input rows.
406+
query I
407+
COPY (SELECT column1 AS a, column2 AS b FROM (VALUES
408+
(1, 9), (1, 8), (1, 7), (1, 6), (1, 5), (1, 4), (1, 3), (1, 2), (2, 1), (3, 0)))
409+
TO 'test_files/scratch/limit/ordered_prefix.parquet' STORED AS PARQUET;
410+
----
411+
10
412+
413+
statement ok
414+
CREATE EXTERNAL TABLE ordered_prefix (a BIGINT, b BIGINT)
415+
STORED AS PARQUET
416+
WITH ORDER (a ASC)
417+
LOCATION 'test_files/scratch/limit/ordered_prefix.parquet';
418+
419+
query TT
420+
EXPLAIN SELECT * FROM (SELECT a, b FROM ordered_prefix LIMIT 4 OFFSET 3) ORDER BY a, b;
421+
----
422+
logical_plan
423+
01)Sort: ordered_prefix.a ASC NULLS LAST, ordered_prefix.b ASC NULLS LAST
424+
02)--Limit: skip=3, fetch=4
425+
03)----TableScan: ordered_prefix projection=[a, b], fetch=7
426+
physical_plan
427+
01)SortExec: expr=[a@0 ASC NULLS LAST, b@1 ASC NULLS LAST], preserve_partitioning=[false]
428+
02)--GlobalLimitExec: skip=3, fetch=4
429+
03)----DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/limit/ordered_prefix.parquet]]}, projection=[a, b], limit=7, output_ordering=[a@0 ASC NULLS LAST], file_type=parquet
430+
431+
query II
432+
SELECT * FROM (SELECT a, b FROM ordered_prefix LIMIT 4 OFFSET 3) ORDER BY a, b;
433+
----
434+
1 3
435+
1 4
436+
1 5
437+
1 6
438+
439+
# When the input ordering comes from an explicit sort, the finer sort is pushed
440+
# into that sort instead, as already happens when a limit without OFFSET is
441+
# fused into its sort. The TopK keeps `skip + fetch` (3 + 4 = 7) rows so the
442+
# limit still returns 4.
443+
query TT
444+
EXPLAIN SELECT * FROM (SELECT a, b FROM ordered_prefix ORDER BY b LIMIT 4 OFFSET 3) ORDER BY b, a;
445+
----
446+
logical_plan
447+
01)Sort: ordered_prefix.b ASC NULLS LAST, ordered_prefix.a ASC NULLS LAST
448+
02)--Limit: skip=3, fetch=4
449+
03)----Sort: ordered_prefix.b ASC NULLS LAST, fetch=7
450+
04)------TableScan: ordered_prefix projection=[a, b]
451+
physical_plan
452+
01)GlobalLimitExec: skip=3, fetch=4
453+
02)--SortExec: TopK(fetch=7), expr=[b@1 ASC NULLS LAST, a@0 ASC NULLS LAST], preserve_partitioning=[false]
454+
03)----DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/limit/ordered_prefix.parquet]]}, projection=[a, b], file_type=parquet, predicate=DynamicFilter [ empty ], sort_order_for_reorder=[b@1 ASC NULLS LAST, a@0 ASC NULLS LAST], dynamic_rg_pruning=eligible
455+
456+
query II
457+
SELECT * FROM (SELECT a, b FROM ordered_prefix ORDER BY b LIMIT 4 OFFSET 3) ORDER BY b, a;
458+
----
459+
1 3
460+
1 4
461+
1 5
462+
1 6
463+
364464
# The aggregate does not need to be computed because the input statistics are exact and
365465
# an OFFSET, but no LIMIT, is specified.
366466
query TT

0 commit comments

Comments
 (0)