Push down topk through join by SubhamSinghal · Pull Request #21621 · apache/datafusion

SubhamSinghal · 2026-04-14T14:29:24Z

Which issue does this PR close?

#11900

Rationale for this change

When a query has ORDER BY <cols> LIMIT N on top of an outer join and all sort columns come from the preserved side,
DataFusion currently runs the full join first, then sorts and limits. We can push a copy of the Sort(fetch=N) to the preserved input, reducing the number of rows entering the join.

Before:

Sort: t1.b ASC, fetch=3
   Left Join: t1.a = t2.a
     Scan: t1     ← scans ALL rows
     Scan: t2

After:

  Sort: t1.b ASC, fetch=3
    Left Join: t1.a = t2.a
      Sort: t1.b ASC, fetch=3  ← pushed down
        Scan: t1               ← only top-3 rows enter join
      Scan: t2

What changes are included in this PR?

A new logical optimizer rule PushDownTopKThroughJoin that:

Matches Sort with fetch = Some(N) (TopK)
Looks through an optional Projection to find a Join
Checks join type is LEFT or RIGHT with no non-equijoin filter
Verifies all sort expression columns come from the preserved side
Inserts a copy of the Sort(fetch=N) on the preserved child
Keeps the top-level sort for correctness

Are these changes tested?

Yes through UT

Are there any user-facing changes?

No API changes.

martin-g · 2026-04-15T12:52:50Z

+        } else {
+            &join.right
+        };
+        if matches!(preserved_child.as_ref(), LogicalPlan::Sort(_)) {


This condition looks a bit broad.
If the child has no fetch limit or a larger fetch limit than the current one then pushing down the current Sort with its fetch limit would be beneficial, no ?
The optimization should be skipped only if the Sort expr is different or its fetch limit is non-zero but smaller than the current one.

gene-bordegaray

The idea behind the optimizer rule itself makes sense. I think there is some room to add more regression testing and ensure we uphold correctness standards for all cases. Let me know what you think :)

gene-bordegaray · 2026-04-16T11:26:17Z

+
+        // Create the new Sort(fetch) on the preserved child
+        let new_child_sort = Arc::new(LogicalPlan::Sort(SortPlan {
+            expr: sort.expr.clone(),


I don't think we shoudl just clone sort.expr here. Since the Sort can sit on top of a Projection, in this case the ORDER BY clause of the query is interpreted against the projection output columns, not directly on the join's child.

When we push down the Sort(fetch) rather than cloning the Sort columns we need to push down the columns that were projected.

I believe behavior for this right now would work like this:

Sort: b, fetch=1 Projection: -t1.b AS b Join The optimizer rewrites it into: Sort: b, fetch=1 Projection: -t1.b AS b Join Sort: b, fetch=1 -> This is using the post-projected value!

neilconway · 2026-04-17T14:22:02Z

+# Child has larger fetch: push our tighter limit
+# The inner Sort(fetch=5) has a larger limit than our outer Sort(fetch=2),
+# so pushing fetch=2 to the preserved child reduces data further.
+query TT
+EXPLAIN SELECT * FROM (
+    SELECT t1.a, t1.b, t2.x
+    FROM (SELECT * FROM t1 ORDER BY b ASC LIMIT 5) t1
+    LEFT JOIN t2 ON t1.a = t2.x
+) sub
+ORDER BY b ASC LIMIT 2;
+----
+logical_plan
+01)Sort: sub.b ASC NULLS LAST, fetch=2
+02)--SubqueryAlias: sub
+03)----Left Join: t1.a = t2.x
+04)------SubqueryAlias: t1
+05)--------Sort: t1.b ASC NULLS LAST, fetch=5
+06)----------TableScan: t1 projection=[a, b]
+07)------TableScan: t2 projection=[x]


Seems like we don't actually push down the fetch=2 tighter limit into the nested Sort here.

It is being blocked by subqueryAlias between sort and join. I think I need to update the comment.

Right; couldn't we push the topk down despite the alias? This seems like a fairly common query structure that it would be nice to support.

refactored code to handle SubqueryAlias

neilconway · 2026-04-17T14:24:53Z

+        if join.filter.is_some() {
+            return Ok(Transformed::no(plan));
+        }


Might not be necessary for this PR, but would be pretty easy to check if the filter only references non-preserved-side columns, in which case I think we can still do the pushdown?

removed filter. added UT to verify results are correct.

martin-g · 2026-04-21T05:25:43Z

+/// (`Option<TableReference>`) structurally.  A `Bare("t1")` and
+/// `Full { catalog, schema, table: "t1" }` are NOT equal even though they
+/// refer to the same column.  After resolving through SubqueryAlias the
+/// variant may differ, so we compare by display string instead.


How Expr::to_string() helps with the missing TableReference on one of the sides ?
I don't understand how this is better than Column::eq().

updated comment and removed to_string()

gene-bordegaray

this is looking good 👍

gene-bordegaray · 2026-04-21T12:03:47Z

+/// Input sort exprs:  [neg_b ASC]
+/// Output sort exprs: [(- t1.b) ASC]
+/// ```
+fn resolve_sort_exprs_through_projection(


may be worth adding unit tests for this guy and resolve_sort_exprs_through_subquery_alias rather than just in slt files to make behavior expectations very clear

SubhamSinghal · 2026-04-25T17:47:49Z

@martin-g @neilconway I have addressed review comments. PTAL

gene-bordegaray · 2026-05-06T18:03:35Z

is this still being worked on @SubhamSinghal ?

SubhamSinghal · 2026-05-06T18:17:30Z

@gene-bordegaray yes, just fixed UT. Please let me know the next step.

gene-bordegaray

re-review due to staleness (Github is amazing and I can't find where to remove approval)

gene-bordegaray

my approval is stale, will re-review soon. request changes is only way to remove my approval it looks :(

kumarUjjawal · 2026-05-08T12:12:34Z

            Arc::new(EliminateOuterJoin::new()),
            // Filters can't be pushed down past Limits, we should do PushDownFilter after PushDownLimit
            Arc::new(PushDownLimit::new()),
+            Arc::new(PushDownTopKThroughJoin::new()),


this runs before PushDownFilter. For Sort(fetch) -> Filter -> Join, it only fires on a later optimizer pass after the filter has moved. That works with the default max_passes = 3, but with max_passes = 1 the pushdown is missed. Should we run this after PushDownFilter?

changed rule order to run after PushDownFilter

kumarUjjawal · 2026-05-08T12:14:29Z

Could we add a few more coverage cases? Like DESC ordering, NULLS FIRST, a multi-level outer join, a CROSS JOIN negative case, and maybe a tied-sort-key case. A custom volatile UDF case would also be useful if it is easy to add.

Added more test cases. Would require help for custom volatile UDF case.

kumarUjjawal · 2026-05-08T12:19:15Z

+        let (inner_child, child_resolved_exprs) = match preserved_child.as_ref() {
+            LogicalPlan::SubqueryAlias(sq) => {
+                let exprs =
+                    resolve_sort_exprs_through_subquery_alias(&resolved_sort_exprs, sq)?;
+                (sq.input.as_ref(), exprs)
+            }
+            _ => (preserved_child.as_ref(), resolved_sort_exprs.clone()),
+        };


This only looks through one SubqueryAlias on the preserved side. If there are nested aliases, we may miss an existing inner Sort(fetch) and add another sort instead of tightening it. We should handle nested aliases.

@kumarUjjawal Thanks for highlighting this, fixed issue with iterative traversal of nested aliases

Subham Singhal added 2 commits April 14, 2026 19:57

Push TopK (Sort with fetch) through outer joins

4675c43

lint fix

9aede67

github-actions Bot added optimizer Optimizer rules sqllogictest SQL Logic Tests (.slt) labels Apr 14, 2026

Subham Singhal added 2 commits April 14, 2026 20:46

Merge branch 'main' into push-down-topk-through-join

3899d6a

fix build failure

19b0edc

martin-g reviewed Apr 15, 2026

View reviewed changes

This comment was marked as duplicate.

Sign in to view

gene-bordegaray reviewed Apr 16, 2026

View reviewed changes

Subham Singhal and others added 6 commits April 17, 2026 14:17

Merge branch 'main' into push-down-topk-through-join

367d8a5

Handle edge cases

baf25ef

Handle volatile expr early

67f9265

Merge branch 'main' into push-down-topk-through-join

0ca184e

Merge branch 'main' into push-down-topk-through-join

16f61cf

Fix build failure

d12aefa

neilconway reviewed Apr 17, 2026

View reviewed changes

Subham Singhal and others added 6 commits April 18, 2026 19:12

Handle subquery alias

902ef77

Merge branch 'main' into push-down-topk-through-join

6b232e1

Update comment

254b224

Doc fix

c648e71

Merge branch 'main' into push-down-topk-through-join

34832c8

Merge branch 'main' into push-down-topk-through-join

e02b82a

SubhamSinghal requested review from gene-bordegaray, martin-g and neilconway April 20, 2026 15:35

martin-g reviewed Apr 20, 2026

View reviewed changes

Comment thread datafusion/optimizer/src/push_down_topk_through_join.rs Outdated

Comment thread datafusion/optimizer/src/push_down_topk_through_join.rs

Subham Singhal added 2 commits April 21, 2026 10:29

Handle volatile expr in projection

03f6499

Merge branch 'main' into push-down-topk-through-join

e051707

martin-g reviewed Apr 21, 2026

View reviewed changes

use structural equality

051868a

gene-bordegaray approved these changes Apr 21, 2026

View reviewed changes

Adds UT

1cfeb76

SubhamSinghal requested a review from martin-g April 22, 2026 04:31

SubhamSinghal added 2 commits May 2, 2026 23:10

Merge branch 'main' into push-down-topk-through-join

c6edd07

Merge branch 'main' into push-down-topk-through-join

f72d845

Fix UT

0371004

github-actions Bot added the core Core DataFusion crate label May 6, 2026

gene-bordegaray reviewed May 8, 2026

View reviewed changes

gene-bordegaray suggested changes May 8, 2026

View reviewed changes

SubhamSinghal requested a review from gene-bordegaray May 8, 2026 11:21

kumarUjjawal reviewed May 8, 2026

View reviewed changes

Subham Singhal added 3 commits May 9, 2026 12:40

Merge branch 'main' into push-down-topk-through-join

5d62f22

Resolve comment

38bbf87

Merge branch 'main' into push-down-topk-through-join

036fb9f

SubhamSinghal force-pushed the push-down-topk-through-join branch from 2711783 to 036fb9f Compare May 9, 2026 08:21

Subham Singhal and others added 3 commits May 9, 2026 14:45

Adds back missing UT

8d4fc2c

Merge branch 'main' into push-down-topk-through-join

53662c1

Fix explain slt test

80f84bc

Conversation

SubhamSinghal commented Apr 14, 2026

Which issue does this PR close?

Rationale for this change

What changes are included in this PR?

Are these changes tested?

Are there any user-facing changes?

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

This comment was marked as duplicate.

Uh oh!

gene-bordegaray left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

neilconway Apr 17, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

gene-bordegaray left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

SubhamSinghal commented Apr 25, 2026

Uh oh!

gene-bordegaray commented May 6, 2026

Uh oh!

SubhamSinghal commented May 6, 2026

Uh oh!

gene-bordegaray left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

gene-bordegaray left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

neilconway Apr 17, 2026 •

edited

Loading

gene-bordegaray left a comment •

edited

Loading