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

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

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!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

neilconway Apr 17, 2026 •

edited

Loading