@@ -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+ }
0 commit comments