You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by gi...@apache.org on 2026/09/15 11:14:57 UTC
(datafusion) branch main updated: Push down topk through join (#21621)
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion.git
The following commit(s) were added to refs/heads/main by this push:
new ec97d1c902 Push down topk through join (#21621)
ec97d1c902 is described below
commit ec97d1c902bf6740a03d4995aeb276dc59a953db
Author: Subham Singhal <Su...@users.noreply.github.com>
AuthorDate: Tue Sep 15 10:57:14 2026 +0000
Push down topk through join (#21621)
## Which issue does this PR close?
- closes https://github.com/apache/datafusion/issues/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:
1. Matches `Sort` with `fetch = Some(N)` (TopK)
2. Looks through an optional `Projection` to find a `Join`
3. Checks join type is LEFT or RIGHT with no non-equijoin filter
4. Verifies all sort expression columns come from the preserved side
5. Inserts a copy of the `Sort(fetch=N)` on the preserved child
6. Keeps the top-level sort for correctness
## Are these changes tested?
Yes through UT
## Are there any user-facing changes?
No API changes.
---------
Co-authored-by: Subham Singhal <su...@Subhams-MacBook-Air.local>
Co-authored-by: Dmitrii Blaginin <dm...@blaginin.me>
Co-authored-by: Kumar Ujjawal <uj...@gmail.com>
---
datafusion/core/src/optimizer_rule_reference.md | 54 +-
datafusion/optimizer/src/push_down_limit.rs | 260 ++--
.../src/push_down_limit/topk_through_join.rs | 1635 ++++++++++++++++++++
.../test_files/push_down_topk_through_join.slt | 430 ++++-
4 files changed, 2160 insertions(+), 219 deletions(-)
diff --git a/datafusion/core/src/optimizer_rule_reference.md b/datafusion/core/src/optimizer_rule_reference.md
index 1367ed0843..6f65ee92e7 100644
--- a/datafusion/core/src/optimizer_rule_reference.md
+++ b/datafusion/core/src/optimizer_rule_reference.md
@@ -35,33 +35,33 @@ Rule order matters. The default pipeline may change between releases.
### Logical Optimizer Rules
-| order | rule | summary |
-| ----- | ----------------------------------------- | --------------------------------------------------------------------------------------------------------------------------- |
-| 1 | `rewrite_set_comparison` | Rewrites `ANY` and `ALL` set-comparison subqueries into `EXISTS`-based boolean expressions with correct SQL NULL semantics. |
-| 2 | `optimize_unions` | Flattens nested unions and removes unions with a single input. |
-| 3 | `unions_to_filter` | Merges `UNION DISTINCT` branches that share the same source into a single filtered branch with a disjunctive predicate. |
-| 4 | `simplify_expressions` | Constant-folds and simplifies expressions while preserving output names. |
-| 5 | `replace_distinct_aggregate` | Rewrites `DISTINCT` and `DISTINCT ON` operators into aggregate-based plans that later rules can optimize further. |
-| 6 | `eliminate_join` | Replaces keyless inner joins with a literal `false` filter by an empty relation. |
-| 7 | `decorrelate_predicate_subquery` | Converts eligible `IN` and `EXISTS` predicate subqueries into semi or anti joins. |
-| 8 | `scalar_subquery_to_join` | Rewrites eligible scalar subqueries into joins and adds schema-preserving projections. |
-| 9 | `decorrelate_lateral_join` | Rewrites eligible lateral joins into regular joins. |
-| 10 | `extract_equijoin_predicate` | Splits join filters into equijoin keys and residual predicates. |
-| 11 | `eliminate_duplicated_expr` | Removes duplicate expressions from projections, aggregates, and similar operators. |
-| 12 | `eliminate_filter` | Drops always-true filters and replaces always-false or NULL filters with empty relations. |
-| 13 | `eliminate_cross_join` | Uses filter predicates to replace cross joins with inner joins when join keys can be found. |
-| 14 | `eliminate_limit` | Removes no-op limits and simplifies trivial limit shapes. |
-| 15 | `propagate_empty_relation` | Pushes empty-relation knowledge upward so operators fed by no rows collapse early. |
-| 16 | `filter_null_join_keys` | Adds `IS NOT NULL` filters to nullable equijoin keys that can never match. |
-| 17 | `eliminate_outer_join` | Rewrites outer joins to inner joins when later filters reject the NULL-extended rows. |
-| 18 | `push_down_limit` | Moves literal limits closer to scans and unions and merges adjacent limits. |
-| 19 | `push_down_filter` | Moves filters as early as possible through filter-commutative operators. |
-| 20 | `single_distinct_aggregation_to_group_by` | Rewrites single-column `DISTINCT` aggregations into two-stage `GROUP BY` plans. |
-| 21 | `eliminate_group_by_constant` | Removes constant or functionally redundant expressions from `GROUP BY`. |
-| 22 | `common_sub_expression_eliminate` | Computes repeated subexpressions once and reuses the result. |
-| 23 | `extract_leaf_expressions` | Pulls cheap leaf expressions closer to data sources so later pruning and filter rules can act earlier. |
-| 24 | `push_down_leaf_projections` | Pushes the helper projections created by leaf extraction toward leaf inputs. |
-| 25 | `optimize_projections` | Prunes unused columns and removes unnecessary logical projections. |
+| order | rule | summary |
+| ----- | ----------------------------------------- | --------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
+| 1 | `rewrite_set_comparison` | Rewrites `ANY` and `ALL` set-comparison subqueries into `EXISTS`-based boolean expressions with correct SQL NULL semantics. |
+| 2 | `optimize_unions` | Flattens nested unions and removes unions with a single input. |
+| 3 | `unions_to_filter` | Merges `UNION DISTINCT` branches that share the same source into a single filtered branch with a disjunctive predicate. |
+| 4 | `simplify_expressions` | Constant-folds and simplifies expressions while preserving output names. |
+| 5 | `replace_distinct_aggregate` | Rewrites `DISTINCT` and `DISTINCT ON` operators into aggregate-based plans that later rules can optimize further. |
+| 6 | `eliminate_join` | Replaces keyless inner joins with a literal `false` filter by an empty relation. |
+| 7 | `decorrelate_predicate_subquery` | Converts eligible `IN` and `EXISTS` predicate subqueries into semi or anti joins. |
+| 8 | `scalar_subquery_to_join` | Rewrites eligible scalar subqueries into joins and adds schema-preserving projections. |
+| 9 | `decorrelate_lateral_join` | Rewrites eligible lateral joins into regular joins. |
+| 10 | `extract_equijoin_predicate` | Splits join filters into equijoin keys and residual predicates. |
+| 11 | `eliminate_duplicated_expr` | Removes duplicate expressions from projections, aggregates, and similar operators. |
+| 12 | `eliminate_filter` | Drops always-true filters and replaces always-false or NULL filters with empty relations. |
+| 13 | `eliminate_cross_join` | Uses filter predicates to replace cross joins with inner joins when join keys can be found. |
+| 14 | `eliminate_limit` | Removes no-op limits and simplifies trivial limit shapes. |
+| 15 | `propagate_empty_relation` | Pushes empty-relation knowledge upward so operators fed by no rows collapse early. |
+| 16 | `filter_null_join_keys` | Adds `IS NOT NULL` filters to nullable equijoin keys that can never match. |
+| 17 | `eliminate_outer_join` | Rewrites outer joins to inner joins when later filters reject the NULL-extended rows. |
+| 18 | `push_down_limit` | Moves literal limits to scans/unions, merges limits, and pushes `Sort(fetch=N)` onto a preserved join side only when every sort key is from it; the outer Sort stays. |
+| 19 | `push_down_filter` | Moves filters as early as possible through filter-commutative operators. |
+| 20 | `single_distinct_aggregation_to_group_by` | Rewrites single-column `DISTINCT` aggregations into two-stage `GROUP BY` plans. |
+| 21 | `eliminate_group_by_constant` | Removes constant or functionally redundant expressions from `GROUP BY`. |
+| 22 | `common_sub_expression_eliminate` | Computes repeated subexpressions once and reuses the result. |
+| 23 | `extract_leaf_expressions` | Pulls cheap leaf expressions closer to data sources so later pruning and filter rules can act earlier. |
+| 24 | `push_down_leaf_projections` | Pushes the helper projections created by leaf extraction toward leaf inputs. |
+| 25 | `optimize_projections` | Prunes unused columns and removes unnecessary logical projections. |
### Physical Optimizer Rules
diff --git a/datafusion/optimizer/src/push_down_limit.rs b/datafusion/optimizer/src/push_down_limit.rs
index 4a26cd5884..79c18fbdeb 100644
--- a/datafusion/optimizer/src/push_down_limit.rs
+++ b/datafusion/optimizer/src/push_down_limit.rs
@@ -15,7 +15,9 @@
// specific language governing permissions and limitations
// under the License.
-//! [`PushDownLimit`] pushes `LIMIT` earlier in the query plan
+//! [`PushDownLimit`] pushes `LIMIT` earlier in the query plan, and pushes a
+//! `Sort` with a fetch limit (TopK) onto a join's preserved-side child (see
+//! the `topk_through_join` submodule).
use std::cmp::min;
use std::sync::Arc;
@@ -29,8 +31,25 @@ use datafusion_common::utils::combine_limit;
use datafusion_expr::logical_plan::{Join, JoinType, Limit, LogicalPlan};
use datafusion_expr::{FetchType, SkipType, lit};
+mod topk_through_join;
+use topk_through_join::push_topk_through_join;
+
/// Optimization rule that tries to push down `LIMIT`.
-//. It will push down through projection, limits (taking the smaller limit)
+///
+/// It pushes `LIMIT` through projections, unions, joins and into scans
+/// (taking the smaller limit when one is already present), merges adjacent
+/// limits, and pushes a `Sort` carrying a fetch limit (a TopK) onto a join's
+/// preserved-side child so fewer rows enter the join. See
+/// `topk_through_join` for the latter.
+///
+/// The TopK-through-join rewrite has an implicit ordering dependency on
+/// `push_down_filter` and `eliminate_outer_join`: a `Filter` sitting between
+/// the `Sort` and the `Join` blocks this rule (a `Filter` can't sink below a
+/// fetch-bearing node without changing which rows survive), and an outer
+/// join only has a "preserved side" once `eliminate_outer_join` has already
+/// dropped joins it can prove are equivalent to an inner join. Both of those
+/// rules need a prior or concurrent pass to unblock this one; with
+/// `max_passes >= 2` this resolves on a later pass.
#[derive(Default, Debug)]
pub struct PushDownLimit {}
@@ -47,146 +66,153 @@ impl OptimizerRule for PushDownLimit {
true
}
- #[expect(clippy::only_used_in_recursion)]
fn rewrite(
&self,
plan: LogicalPlan,
- config: &dyn OptimizerConfig,
+ _config: &dyn OptimizerConfig,
) -> Result<Transformed<LogicalPlan>> {
- let LogicalPlan::Limit(mut limit) = plan else {
- return Ok(Transformed::no(plan));
- };
+ match plan {
+ LogicalPlan::Limit(limit) => rewrite_limit(limit),
+ LogicalPlan::Sort(s) if s.fetch.is_some() => push_topk_through_join(s),
+ other => Ok(Transformed::no(other)),
+ }
+ }
- // Currently only rewrite if skip and fetch are both literals
- let SkipType::Literal(skip) = limit.get_skip_type()? else {
+ fn name(&self) -> &str {
+ "push_down_limit"
+ }
+
+ fn apply_order(&self) -> Option<ApplyOrder> {
+ Some(ApplyOrder::TopDown)
+ }
+}
+
+/// Limit-side dispatch (split out from `rewrite` so that the top-level
+/// match in `OptimizerRule::rewrite` reads as a parallel branch alongside
+/// the Sort-with-fetch handler).
+fn rewrite_limit(mut limit: Limit) -> Result<Transformed<LogicalPlan>> {
+ // Currently only rewrite if skip and fetch are both literals
+ let SkipType::Literal(skip) = limit.get_skip_type()? else {
+ return Ok(Transformed::no(LogicalPlan::Limit(limit)));
+ };
+ let FetchType::Literal(fetch) = limit.get_fetch_type()? else {
+ return Ok(Transformed::no(LogicalPlan::Limit(limit)));
+ };
+
+ // Merge the Parent Limit and the Child Limit.
+ if let LogicalPlan::Limit(child) = limit.input.as_ref() {
+ let SkipType::Literal(child_skip) = child.get_skip_type()? else {
return Ok(Transformed::no(LogicalPlan::Limit(limit)));
};
- let FetchType::Literal(fetch) = limit.get_fetch_type()? else {
+ let FetchType::Literal(child_fetch) = child.get_fetch_type()? else {
return Ok(Transformed::no(LogicalPlan::Limit(limit)));
};
- // Merge the Parent Limit and the Child Limit.
- if let LogicalPlan::Limit(child) = limit.input.as_ref() {
- let SkipType::Literal(child_skip) = child.get_skip_type()? else {
- return Ok(Transformed::no(LogicalPlan::Limit(limit)));
- };
- let FetchType::Literal(child_fetch) = child.get_fetch_type()? else {
- return Ok(Transformed::no(LogicalPlan::Limit(limit)));
- };
-
- let (skip, fetch) = combine_limit(skip, fetch, child_skip, child_fetch);
- let plan = LogicalPlan::Limit(Limit {
- skip: Some(Box::new(lit(skip as i64))),
- fetch: fetch.map(|f| Box::new(lit(f as i64))),
- input: Arc::clone(&child.input),
- });
+ let (skip, fetch) = combine_limit(skip, fetch, child_skip, child_fetch);
+ let new_limit = Limit {
+ skip: Some(Box::new(lit(skip as i64))),
+ fetch: fetch.map(|f| Box::new(lit(f as i64))),
+ input: Arc::clone(&child.input),
+ };
- // recursively reapply the rule on the new plan
- return self.rewrite(plan, config);
- }
+ // recursively reapply the rule on the new limit
+ return rewrite_limit(new_limit);
+ }
- // no fetch to push, so return the original plan
- let Some(fetch) = fetch else {
- return Ok(Transformed::no(LogicalPlan::Limit(limit)));
- };
+ // no fetch to push, so return the original plan
+ let Some(fetch) = fetch else {
+ return Ok(Transformed::no(LogicalPlan::Limit(limit)));
+ };
- match Arc::unwrap_or_clone(limit.input) {
- LogicalPlan::TableScan(mut scan) => {
- let rows_needed = if fetch != 0 { fetch + skip } else { 0 };
- let new_fetch = scan
+ match Arc::unwrap_or_clone(limit.input) {
+ LogicalPlan::TableScan(mut scan) => {
+ let rows_needed = if fetch != 0 { fetch + skip } else { 0 };
+ let new_fetch = scan
+ .fetch
+ .map(|x| min(x, rows_needed))
+ .or(Some(rows_needed));
+ if new_fetch == scan.fetch {
+ original_limit(skip, fetch, LogicalPlan::TableScan(scan))
+ } else {
+ // push limit into the table scan itself
+ scan.fetch = scan
.fetch
.map(|x| min(x, rows_needed))
.or(Some(rows_needed));
- if new_fetch == scan.fetch {
- original_limit(skip, fetch, LogicalPlan::TableScan(scan))
- } else {
- // push limit into the table scan itself
- scan.fetch = scan
- .fetch
- .map(|x| min(x, rows_needed))
- .or(Some(rows_needed));
- transformed_limit(skip, fetch, LogicalPlan::TableScan(scan))
- }
- }
- LogicalPlan::Union(mut union) => {
- // push limits to each input of the union
- union.inputs = union
- .inputs
- .into_iter()
- .map(|input| make_arc_limit(0, fetch + skip, input))
- .collect();
- transformed_limit(skip, fetch, LogicalPlan::Union(union))
+ transformed_limit(skip, fetch, LogicalPlan::TableScan(scan))
}
+ }
+ LogicalPlan::Union(mut union) => {
+ // push limits to each input of the union
+ union.inputs = union
+ .inputs
+ .into_iter()
+ .map(|input| make_arc_limit(0, fetch + skip, input))
+ .collect();
+ transformed_limit(skip, fetch, LogicalPlan::Union(union))
+ }
+
+ LogicalPlan::Join(join) => {
+ Ok(push_down_join(join, fetch + skip).update_data(|join| {
+ make_limit(skip, fetch, Arc::new(LogicalPlan::Join(join)))
+ }))
+ }
- LogicalPlan::Join(join) => Ok(push_down_join(join, fetch + skip)
- .update_data(|join| {
- make_limit(skip, fetch, Arc::new(LogicalPlan::Join(join)))
- })),
-
- LogicalPlan::Sort(mut sort) => {
- let new_fetch = {
- let sort_fetch = skip + fetch;
- Some(sort.fetch.map(|f| f.min(sort_fetch)).unwrap_or(sort_fetch))
- };
- if new_fetch == sort.fetch {
- if skip > 0 {
- original_limit(skip, fetch, LogicalPlan::Sort(sort))
- } else {
- Ok(Transformed::yes(LogicalPlan::Sort(sort)))
- }
+ LogicalPlan::Sort(mut sort) => {
+ let new_fetch = {
+ let sort_fetch = skip + fetch;
+ Some(sort.fetch.map(|f| f.min(sort_fetch)).unwrap_or(sort_fetch))
+ };
+ if new_fetch == sort.fetch {
+ if skip > 0 {
+ original_limit(skip, fetch, LogicalPlan::Sort(sort))
} else {
- sort.fetch = new_fetch;
- limit.input = Arc::new(LogicalPlan::Sort(sort));
- Ok(Transformed::yes(LogicalPlan::Limit(limit)))
+ Ok(Transformed::yes(LogicalPlan::Sort(sort)))
}
+ } else {
+ sort.fetch = new_fetch;
+ limit.input = Arc::new(LogicalPlan::Sort(sort));
+ Ok(Transformed::yes(LogicalPlan::Limit(limit)))
}
- LogicalPlan::Projection(mut proj) => {
- // commute
- limit.input = Arc::clone(&proj.input);
- let new_limit = LogicalPlan::Limit(limit);
- proj.input = Arc::new(new_limit);
- Ok(Transformed::yes(LogicalPlan::Projection(proj)))
- }
- LogicalPlan::SubqueryAlias(mut subquery_alias) => {
- // commute
- limit.input = Arc::clone(&subquery_alias.input);
- let new_limit = LogicalPlan::Limit(limit);
- subquery_alias.input = Arc::new(new_limit);
- Ok(Transformed::yes(LogicalPlan::SubqueryAlias(subquery_alias)))
- }
- LogicalPlan::Extension(extension_plan)
- if extension_plan.node.supports_limit_pushdown() =>
- {
- let new_children = extension_plan
- .node
- .inputs()
- .into_iter()
- .map(|child| {
- LogicalPlan::Limit(Limit {
- skip: None,
- fetch: Some(Box::new(lit((fetch + skip) as i64))),
- input: Arc::new(child.clone()),
- })
+ }
+ LogicalPlan::Projection(mut proj) => {
+ // commute
+ limit.input = Arc::clone(&proj.input);
+ let new_limit = LogicalPlan::Limit(limit);
+ proj.input = Arc::new(new_limit);
+ Ok(Transformed::yes(LogicalPlan::Projection(proj)))
+ }
+ LogicalPlan::SubqueryAlias(mut subquery_alias) => {
+ // commute
+ limit.input = Arc::clone(&subquery_alias.input);
+ let new_limit = LogicalPlan::Limit(limit);
+ subquery_alias.input = Arc::new(new_limit);
+ Ok(Transformed::yes(LogicalPlan::SubqueryAlias(subquery_alias)))
+ }
+ LogicalPlan::Extension(extension_plan)
+ if extension_plan.node.supports_limit_pushdown() =>
+ {
+ let new_children = extension_plan
+ .node
+ .inputs()
+ .into_iter()
+ .map(|child| {
+ LogicalPlan::Limit(Limit {
+ skip: None,
+ fetch: Some(Box::new(lit((fetch + skip) as i64))),
+ input: Arc::new(child.clone()),
})
- .collect::<Vec<_>>();
+ })
+ .collect::<Vec<_>>();
- // Create a new extension node with updated inputs
- let child_plan = LogicalPlan::Extension(extension_plan);
- let new_extension =
- child_plan.with_new_exprs(child_plan.expressions(), new_children)?;
+ // Create a new extension node with updated inputs
+ let child_plan = LogicalPlan::Extension(extension_plan);
+ let new_extension =
+ child_plan.with_new_exprs(child_plan.expressions(), new_children)?;
- transformed_limit(skip, fetch, new_extension)
- }
- input => original_limit(skip, fetch, input),
+ transformed_limit(skip, fetch, new_extension)
}
- }
-
- fn name(&self) -> &str {
- "push_down_limit"
- }
-
- fn apply_order(&self) -> Option<ApplyOrder> {
- Some(ApplyOrder::TopDown)
+ input => original_limit(skip, fetch, input),
}
}
diff --git a/datafusion/optimizer/src/push_down_limit/topk_through_join.rs b/datafusion/optimizer/src/push_down_limit/topk_through_join.rs
new file mode 100644
index 0000000000..b4141a5f3b
--- /dev/null
+++ b/datafusion/optimizer/src/push_down_limit/topk_through_join.rs
@@ -0,0 +1,1635 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Sort(fetch) → Join pushdown — a sub-module of `push_down_limit`.
+//!
+//! When a `Sort` with a fetch limit (TopK) sits above a join whose
+//! preserved side is known (LEFT / RIGHT / LeftMark / RightMark) and all
+//! sort expressions come from the preserved side, we insert a copy of the
+//! `Sort(fetch)` onto that input to reduce rows entering the join. The
+//! outer `Sort` is kept because a 1-to-many join can produce more than N
+//! output rows from N preserved-side rows.
+//!
+//! CROSS JOIN is deliberately excluded, even though every row from both
+//! sides appears in its output when *both* sides are non-empty: if the
+//! *other* side turns out to be empty, the join output is empty and the
+//! original (unpushed) Sort never evaluates its expression, whereas a
+//! pushed Sort would evaluate it eagerly regardless — which can surface an
+//! error from a fallible sort expression that a correct, unoptimized
+//! execution would never have hit. Proving the other side non-empty at
+//! logical-plan time isn't something this optimizer (or DataFusion's
+//! optimizer generally) has machinery for, so this rule stays conservative
+//! and skips CROSS JOIN entirely rather than risk changing whether a query
+//! errors.
+//!
+//! Dispatched from `PushDownLimit::rewrite` when the plan node is
+//! `LogicalPlan::Sort` with `fetch.is_some()`.
+
+use std::collections::HashMap;
+use std::sync::Arc;
+
+use crate::utils::{has_all_column_refs, schema_columns};
+
+use datafusion_common::tree_node::{Transformed, TreeNode};
+use datafusion_common::{Column, Result, internal_err};
+use datafusion_expr::logical_plan::{
+ JoinType, LogicalPlan, Projection, Sort as SortPlan, SubqueryAlias,
+};
+use datafusion_expr::{Expr, SortExpr};
+
+/// Which child of a join is being treated as the preserved side.
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+enum Side {
+ Left,
+ Right,
+}
+
+/// If `node` is a transparent node (`Projection` or `SubqueryAlias`),
+/// returns its child together with `exprs` resolved through it (column
+/// references rewritten to match the child's schema). Returns `None` for
+/// any other node, signaling the walk should stop.
+///
+/// Shared by both tree walks in [`push_topk_through_join`]: peeling down
+/// from the `Sort` to find the `Join`, and scanning inside the preserved
+/// child for an existing inner `Sort`.
+fn peel_transparent_layer<'a>(
+ node: &'a LogicalPlan,
+ exprs: &[SortExpr],
+) -> Result<Option<(&'a LogicalPlan, Vec<SortExpr>)>> {
+ match node {
+ LogicalPlan::Projection(proj) => Ok(Some((
+ proj.input.as_ref(),
+ resolve_sort_exprs_through_projection(exprs, proj)?,
+ ))),
+ LogicalPlan::SubqueryAlias(sq) => Ok(Some((
+ sq.input.as_ref(),
+ resolve_sort_exprs_through_subquery_alias(exprs, sq)?,
+ ))),
+ _ => Ok(None),
+ }
+}
+
+/// Top-level pushdown for `Sort(fetch) → ... → Join` patterns. The caller
+/// (`PushDownLimit::rewrite`) has already matched `LogicalPlan::Sort(sort)`
+/// with `sort.fetch.is_some()`; taking `Sort` by value here (rather than
+/// re-wrapping it into a `LogicalPlan` just to immediately unwrap it) means
+/// early exits can hand `sort` straight back via `LogicalPlan::Sort(sort)`
+/// without cloning.
+pub(super) fn push_topk_through_join(sort: SortPlan) -> Result<Transformed<LogicalPlan>> {
+ let Some(fetch) = sort.fetch else {
+ return Ok(Transformed::no(LogicalPlan::Sort(sort)));
+ };
+
+ // Don't push if any sort expression is non-deterministic (e.g.
+ // `random()`). Duplicating such expressions would produce different
+ // values at each evaluation point, potentially changing results.
+ if sort.expr.iter().any(|se| se.expr.is_volatile()) {
+ return Ok(Transformed::no(LogicalPlan::Sort(sort)));
+ }
+
+ // Peel through transparent nodes (SubqueryAlias, Projection) to find
+ // the Join, resolving sort expressions through each layer along the
+ // way so column references end up matching the join's schema. Track
+ // intermediates so we can reconstruct the tree afterward.
+ let mut current = sort.input.as_ref();
+ let mut intermediates: Vec<&LogicalPlan> = Vec::new();
+ let mut resolved_sort_exprs = sort.expr.clone();
+ let join = loop {
+ if let LogicalPlan::Join(join) = current {
+ break join;
+ }
+ match peel_transparent_layer(current, &resolved_sort_exprs)? {
+ Some((next, new_exprs)) => {
+ intermediates.push(current);
+ current = next;
+ resolved_sort_exprs = new_exprs;
+ }
+ None => return Ok(Transformed::no(LogicalPlan::Sort(sort))),
+ }
+ };
+
+ // Determine which side of the join is preserved.
+ //
+ // - LEFT / LeftMark: only left preserved.
+ // - RIGHT / RightMark: symmetric.
+ //
+ // Non-equijoin filters in the ON clause are safe: outer joins guarantee
+ // all preserved-side rows appear in the output regardless of the
+ // filter.
+ //
+ // CROSS JOIN and other Inner joins are deliberately excluded — see the
+ // module doc comment for why CROSS JOIN specifically isn't safe to
+ // push through here despite superficially preserving every row.
+ let preserved_side = match join.join_type {
+ JoinType::Left | JoinType::LeftMark => Side::Left,
+ JoinType::Right | JoinType::RightMark => Side::Right,
+ _ => return Ok(Transformed::no(LogicalPlan::Sort(sort))),
+ };
+
+ // After resolving through projections, sort expressions may now
+ // contain volatile functions (e.g. `random() AS col`). Duplicating
+ // them would change results.
+ if resolved_sort_exprs.iter().any(|se| se.expr.is_volatile()) {
+ return Ok(Transformed::no(LogicalPlan::Sort(sort)));
+ }
+
+ // A sort key with no column references at all (e.g. `ORDER BY 'x'
+ // LIMIT 3`) can't distinguish any row from any other on either side.
+ // `has_all_column_refs` is vacuously true for such a key (there are no
+ // column refs to fail to find), so without this check we'd still push
+ // a Sort onto a preserved child even though its "top N" doesn't depend
+ // on that child's data at all — a legal result (any N rows satisfy an
+ // all-ties ORDER BY), but wasted work duplicating a Sort node for no
+ // row-reduction benefit.
+ if resolved_sort_exprs
+ .iter()
+ .all(|se| se.expr.column_refs().is_empty())
+ {
+ return Ok(Transformed::no(LogicalPlan::Sort(sort)));
+ }
+
+ // Confirm the preserved side's schema contains all referenced sort
+ // columns.
+ //
+ // Caveat: `schema_columns` adds an unqualified entry for every field
+ // alongside its qualified one, so an unqualified sort column matches as
+ // long as the preserved side has a same-named field — even if an
+ // identically-named field also exists on the *other* side, in which
+ // case an unqualified reference is genuinely ambiguous and this could
+ // wrongly treat a reference to the other side's column as if it were
+ // the preserved side's. SQL can never produce this: the planner always
+ // qualifies column references (or rejects the query as ambiguous)
+ // before this rule runs. Only a hand-built `LogicalPlan` with a
+ // deliberately unqualified `Column` can hit this. `push_down_filter`
+ // has the exact same caveat, for the same reason (it also calls
+ // `schema_columns` to accept unqualified filter columns).
+ let schema = match preserved_side {
+ Side::Left => join.left.schema(),
+ Side::Right => join.right.schema(),
+ };
+ let cols = schema_columns(schema);
+ if !resolved_sort_exprs
+ .iter()
+ .all(|se| has_all_column_refs(&se.expr, &cols))
+ {
+ return Ok(Transformed::no(LogicalPlan::Sort(sort)));
+ }
+
+ let preserved_child = match preserved_side {
+ Side::Left => &join.left,
+ Side::Right => &join.right,
+ };
+
+ // Scan deep inside the preserved child (through SubqueryAlias and
+ // Projection layers) to find an existing Sort. If found with same
+ // exprs, tighten its fetch in-place. Otherwise, insert a new Sort
+ // directly below the join as the preserved child's wrapper.
+ let mut inner_child = preserved_child.as_ref();
+ let mut deep_resolved_exprs = resolved_sort_exprs.clone();
+ while let Some((next, new_exprs)) =
+ peel_transparent_layer(inner_child, &deep_resolved_exprs)?
+ {
+ inner_child = next;
+ deep_resolved_exprs = new_exprs;
+ }
+
+ // If the inner child is a Limit, skip pushing down for now.
+ //
+ // This is only correct-but-conservative, not a correctness requirement:
+ // inserting `Sort(fetch)` *above* the Limit (`Sort(fetch=N) -> Limit ->
+ // ...`) would be sound regardless of what the Limit's rows are — it
+ // only ever reduces which of the Limit's already-capped rows survive,
+ // the same "insert a new Sort above the preserved child" move used
+ // below when the exprs differ from an existing inner Sort. We decline
+ // instead, which is a missed optimization, not a blocked one:
+ // - If a `Sort` sits below the Limit (`Limit -> Sort`), `PushDownLimit`
+ // merges the two on a later pass into `Sort(fetch)`, and this rule
+ // revisits the node then and pushes down as normal.
+ // - If there's no `Sort` at all (e.g. `LEFT JOIN (SELECT * FROM t LIMIT
+ // 100)`), there is nothing for the Limit to ever merge with, so this
+ // case never gets pushed down at all, on any pass — a real gap, not
+ // just a later-pass one, but non-blocking: the join still runs
+ // correctly against the already-limited input, just without the
+ // extra row reduction a pushed Sort could have added.
+ if matches!(inner_child, LogicalPlan::Limit(_)) {
+ return Ok(Transformed::no(LogicalPlan::Sort(sort)));
+ }
+
+ // Determine action based on existing inner Sort:
+ // - Same exprs, tighter fetch → skip (already optimal)
+ // - Same exprs, larger/no fetch → tighten in-place
+ // - Different exprs or no Sort → insert new Sort below the join
+ //
+ // If `deep_resolved_exprs` became volatile while resolving through
+ // projections inside the preserved child (e.g. `random() AS col`),
+ // structural equality with an existing inner Sort is unsound: two
+ // identical `random()` exprs evaluate to different values. Fall
+ // back to inserting a new Sort with `resolved_sort_exprs`.
+ let deep_exprs_volatile = deep_resolved_exprs.iter().any(|se| se.expr.is_volatile());
+ let inner_sort = match inner_child {
+ LogicalPlan::Sort(s) if !deep_exprs_volatile => Some(s),
+ _ => None,
+ };
+ let same_exprs_sort = inner_sort
+ .filter(|child_sort| sort_exprs_equal(&child_sort.expr, &deep_resolved_exprs));
+
+ let new_preserved_child = if let Some(child_sort) = same_exprs_sort {
+ let child_fetch_tighter = match child_sort.fetch {
+ Some(child_fetch) => child_fetch <= fetch,
+ None => false,
+ };
+ if child_fetch_tighter {
+ return Ok(Transformed::no(LogicalPlan::Sort(sort)));
+ }
+ rebuild_with_tightened_sort(
+ preserved_child.as_ref(),
+ &deep_resolved_exprs,
+ fetch,
+ )?
+ } else {
+ // Different exprs, or no existing inner Sort — insert a new Sort
+ // above the preserved child. If an inner Sort exists with no
+ // fetch, our pushed Sort is the only row reduction. If it has a
+ // fetch, re-sorting a small set is cheap and still reduces join
+ // input.
+ Arc::new(LogicalPlan::Sort(SortPlan {
+ expr: resolved_sort_exprs,
+ input: Arc::clone(preserved_child),
+ fetch: Some(fetch),
+ }))
+ };
+
+ let mut new_join = join.clone();
+ match preserved_side {
+ Side::Left => new_join.left = new_preserved_child,
+ Side::Right => new_join.right = new_preserved_child,
+ }
+
+ // Rebuild the tree: join → intermediate nodes → top-level sort.
+ let mut new_sort_input = Arc::new(LogicalPlan::Join(new_join));
+ for node in intermediates.into_iter().rev() {
+ new_sort_input = Arc::new(match node {
+ LogicalPlan::Projection(proj) => {
+ let mut new_proj = proj.clone();
+ new_proj.input = new_sort_input;
+ LogicalPlan::Projection(new_proj)
+ }
+ LogicalPlan::SubqueryAlias(sq) => LogicalPlan::SubqueryAlias(
+ SubqueryAlias::try_new(new_sort_input, sq.alias.clone())?,
+ ),
+ _ => {
+ return internal_err!(
+ "push_topk_through_join: unexpected intermediate node: {}",
+ node.display()
+ );
+ }
+ });
+ }
+
+ Ok(Transformed::yes(LogicalPlan::Sort(SortPlan {
+ expr: sort.expr,
+ input: new_sort_input,
+ fetch: sort.fetch,
+ })))
+}
+
+/// Replace column references in sort expressions using a structural
+/// `Column`→expr map. Keying by `Column` (rather than `Column::flat_name()`)
+/// avoids collisions between an unqualified column whose name happens to
+/// contain a dot (e.g. a quoted alias `"t1.b"`) and a qualified column that
+/// stringifies the same way (`t1.b`).
+fn replace_columns_in_sort_exprs(
+ sort_exprs: &[SortExpr],
+ replace_map: &HashMap<Column, Expr>,
+) -> Result<Vec<SortExpr>> {
+ sort_exprs
+ .iter()
+ .map(|sort_expr| {
+ let new_expr = sort_expr.expr.clone().transform(|expr| {
+ let replacement = match &expr {
+ Expr::Column(col) => replace_map.get(col).cloned(),
+ _ => None,
+ };
+ Ok(replacement.map_or_else(|| Transformed::no(expr), Transformed::yes))
+ })?;
+ Ok(SortExpr {
+ expr: new_expr.data,
+ ..*sort_expr
+ })
+ })
+ .collect()
+}
+
+/// Resolve sort expressions through a projection by replacing column
+/// references with the underlying projection expressions.
+fn resolve_sort_exprs_through_projection(
+ sort_exprs: &[SortExpr],
+ projection: &Projection,
+) -> Result<Vec<SortExpr>> {
+ let replace_map: HashMap<Column, Expr> = projection
+ .schema
+ .iter()
+ .zip(projection.expr.iter())
+ .map(|((qualifier, field), expr)| {
+ let key = Column::from((qualifier, field));
+ (key, expr.clone().unalias())
+ })
+ .collect();
+
+ replace_columns_in_sort_exprs(sort_exprs, &replace_map)
+}
+
+/// Compare two slices of `SortExpr` for structural equality.
+fn sort_exprs_equal(a: &[SortExpr], b: &[SortExpr]) -> bool {
+ a.len() == b.len()
+ && a.iter().zip(b.iter()).all(|(left, right)| {
+ left.asc == right.asc
+ && left.nulls_first == right.nulls_first
+ && left.expr == right.expr
+ })
+}
+
+/// Resolve sort expressions through a `SubqueryAlias` by replacing the
+/// alias qualifier with the input schema's qualifier.
+fn resolve_sort_exprs_through_subquery_alias(
+ sort_exprs: &[SortExpr],
+ subquery_alias: &SubqueryAlias,
+) -> Result<Vec<SortExpr>> {
+ let replace_map: HashMap<Column, Expr> = subquery_alias
+ .schema
+ .iter()
+ .zip(subquery_alias.input.schema().iter())
+ .map(|((alias_qual, alias_field), (input_qual, input_field))| {
+ let alias_col = Column::from((alias_qual, alias_field));
+ let input_col = Column::from((input_qual, input_field));
+ (alias_col, Expr::Column(input_col))
+ })
+ .collect();
+
+ replace_columns_in_sort_exprs(sort_exprs, &replace_map)
+}
+
+/// Rebuild the tree from `root` down to an existing Sort whose expressions
+/// match `target_exprs`, tightening its fetch to `new_fetch`.
+fn rebuild_with_tightened_sort(
+ root: &LogicalPlan,
+ target_exprs: &[SortExpr],
+ new_fetch: usize,
+) -> Result<Arc<LogicalPlan>> {
+ match root {
+ LogicalPlan::Sort(s) if sort_exprs_equal(&s.expr, target_exprs) => {
+ Ok(Arc::new(LogicalPlan::Sort(SortPlan {
+ expr: s.expr.clone(),
+ input: Arc::clone(&s.input),
+ fetch: Some(new_fetch),
+ })))
+ }
+ LogicalPlan::Projection(proj) => {
+ let new_input = rebuild_with_tightened_sort(
+ proj.input.as_ref(),
+ target_exprs,
+ new_fetch,
+ )?;
+ let mut new_proj = proj.clone();
+ new_proj.input = new_input;
+ Ok(Arc::new(LogicalPlan::Projection(new_proj)))
+ }
+ LogicalPlan::SubqueryAlias(sq) => {
+ let new_input =
+ rebuild_with_tightened_sort(sq.input.as_ref(), target_exprs, new_fetch)?;
+ Ok(Arc::new(LogicalPlan::SubqueryAlias(
+ SubqueryAlias::try_new(new_input, sq.alias.clone())?,
+ )))
+ }
+ _ => internal_err!(
+ "rebuild_with_tightened_sort: unexpected node: {}",
+ root.display()
+ ),
+ }
+}
+
+#[cfg(test)]
+mod test {
+ use super::*;
+ use crate::OptimizerContext;
+ use crate::assert_optimized_plan_eq_snapshot;
+ use crate::push_down_limit::PushDownLimit;
+ use crate::test::*;
+
+ use datafusion_expr::col;
+ use datafusion_expr::lit;
+ use datafusion_expr::logical_plan::builder::LogicalPlanBuilder;
+
+ macro_rules! assert_optimized_plan_equal {
+ (
+ $plan:expr,
+ @ $expected:literal $(,)?
+ ) => {{
+ let optimizer_ctx = OptimizerContext::new().with_max_passes(1);
+ let rules: Vec<Arc<dyn crate::OptimizerRule + Send + Sync>> =
+ vec![Arc::new(PushDownLimit::new())];
+ assert_optimized_plan_eq_snapshot!(
+ optimizer_ctx,
+ rules,
+ $plan,
+ @ $expected,
+ )
+ }};
+ }
+
+ /// TopK on left-side columns above a LEFT JOIN → pushed to left child.
+ #[test]
+ fn topk_pushed_to_left_of_left_join() -> Result<()> {
+ let t1 = test_table_scan_with_name("t1")?;
+ let t2 = test_table_scan_with_name("t2")?;
+
+ let plan = LogicalPlanBuilder::from(t1)
+ .join(
+ LogicalPlanBuilder::from(t2).build()?,
+ JoinType::Left,
+ (vec!["a"], vec!["a"]),
+ None,
+ )?
+ .sort_with_limit(vec![col("t1.b").sort(true, false)], Some(3))?
+ .build()?;
+
+ assert_optimized_plan_equal!(
+ plan,
+ @r"
+ Sort: t1.b ASC NULLS LAST, fetch=3
+ Left Join: t1.a = t2.a
+ Sort: t1.b ASC NULLS LAST, fetch=3
+ TableScan: t1
+ TableScan: t2
+ "
+ )
+ }
+
+ /// TopK on right-side columns above a RIGHT JOIN → pushed to right child.
+ #[test]
+ fn topk_pushed_to_right_of_right_join() -> Result<()> {
+ let t1 = test_table_scan_with_name("t1")?;
+ let t2 = test_table_scan_with_name("t2")?;
+
+ let plan = LogicalPlanBuilder::from(t1)
+ .join(
+ LogicalPlanBuilder::from(t2).build()?,
+ JoinType::Right,
+ (vec!["a"], vec!["a"]),
+ None,
+ )?
+ .sort_with_limit(vec![col("t2.b").sort(true, false)], Some(5))?
+ .build()?;
+
+ assert_optimized_plan_equal!(
+ plan,
+ @r"
+ Sort: t2.b ASC NULLS LAST, fetch=5
+ Right Join: t1.a = t2.a
+ TableScan: t1
+ Sort: t2.b ASC NULLS LAST, fetch=5
+ TableScan: t2
+ "
+ )
+ }
+
+ /// TopK pushed through a Projection between Sort and Join.
+ #[test]
+ fn topk_pushed_through_projection() -> Result<()> {
+ let t1 = test_table_scan_with_name("t1")?;
+ let t2 = test_table_scan_with_name("t2")?;
+
+ let plan = LogicalPlanBuilder::from(t1)
+ .join(
+ LogicalPlanBuilder::from(t2).build()?,
+ JoinType::Left,
+ (vec!["a"], vec!["a"]),
+ None,
+ )?
+ .project(vec![col("t1.a"), col("t1.b"), col("t2.c")])?
+ .sort_with_limit(vec![col("t1.b").sort(true, false)], Some(3))?
+ .build()?;
+
+ assert_optimized_plan_equal!(
+ plan,
+ @r"
+ Sort: t1.b ASC NULLS LAST, fetch=3
+ Projection: t1.a, t1.b, t2.c
+ Left Join: t1.a = t2.a
+ Sort: t1.b ASC NULLS LAST, fetch=3
+ TableScan: t1
+ TableScan: t2
+ "
+ )
+ }
+
+ /// INNER JOIN → no pushdown.
+ #[test]
+ fn topk_not_pushed_for_inner_join() -> Result<()> {
+ let t1 = test_table_scan_with_name("t1")?;
+ let t2 = test_table_scan_with_name("t2")?;
+
+ let plan = LogicalPlanBuilder::from(t1)
+ .join(
+ LogicalPlanBuilder::from(t2).build()?,
+ JoinType::Inner,
+ (vec!["a"], vec!["a"]),
+ None,
+ )?
+ .sort_with_limit(vec![col("t1.b").sort(true, false)], Some(3))?
+ .build()?;
+
+ assert_optimized_plan_equal!(
+ plan,
+ @r"
+ Sort: t1.b ASC NULLS LAST, fetch=3
+ Inner Join: t1.a = t2.a
+ TableScan: t1
+ TableScan: t2
+ "
+ )
+ }
+
+ /// CROSS JOIN is never pushed through, even when the sort key comes
+ /// entirely from one side: see the module doc comment for why (a
+ /// pushed Sort would evaluate eagerly on a side whose relevance
+ /// depends on the *other* side being non-empty, which this rule has no
+ /// way to prove at logical-plan time).
+ #[test]
+ fn topk_not_pushed_for_cross_join_left_side_sort() -> Result<()> {
+ let t1 = test_table_scan_with_name("t1")?;
+ let t2 = test_table_scan_with_name("t2")?;
+
+ let plan = LogicalPlanBuilder::from(t1)
+ .cross_join(LogicalPlanBuilder::from(t2).build()?)?
+ .sort_with_limit(vec![col("t1.b").sort(true, false)], Some(3))?
+ .build()?;
+
+ assert_optimized_plan_equal!(
+ plan,
+ @r"
+ Sort: t1.b ASC NULLS LAST, fetch=3
+ Cross Join:
+ TableScan: t1
+ TableScan: t2
+ "
+ )
+ }
+
+ /// Symmetric to the left-side case above.
+ #[test]
+ fn topk_not_pushed_for_cross_join_right_side_sort() -> Result<()> {
+ let t1 = test_table_scan_with_name("t1")?;
+ let t2 = test_table_scan_with_name("t2")?;
+
+ let plan = LogicalPlanBuilder::from(t1)
+ .cross_join(LogicalPlanBuilder::from(t2).build()?)?
+ .sort_with_limit(vec![col("t2.b").sort(true, false)], Some(3))?
+ .build()?;
+
+ assert_optimized_plan_equal!(
+ plan,
+ @r"
+ Sort: t2.b ASC NULLS LAST, fetch=3
+ Cross Join:
+ TableScan: t1
+ TableScan: t2
+ "
+ )
+ }
+
+ /// CROSS JOIN sorted by columns from both sides → no pushdown (would
+ /// still be excluded even if CROSS JOIN were otherwise supported).
+ #[test]
+ fn topk_not_pushed_for_cross_join_mixed_side_sort() -> Result<()> {
+ let t1 = test_table_scan_with_name("t1")?;
+ let t2 = test_table_scan_with_name("t2")?;
+
+ let plan = LogicalPlanBuilder::from(t1)
+ .cross_join(LogicalPlanBuilder::from(t2).build()?)?
+ .sort_with_limit(
+ vec![(col("t1.b") + col("t2.b")).sort(true, false)],
+ Some(3),
+ )?
+ .build()?;
+
+ assert_optimized_plan_equal!(
+ plan,
+ @r"
+ Sort: t1.b + t2.b ASC NULLS LAST, fetch=3
+ Cross Join:
+ TableScan: t1
+ TableScan: t2
+ "
+ )
+ }
+
+ /// Inner join with no equi-keys but a non-empty filter: filter can drop
+ /// rows from either side, so pushing fetch=N is unsafe.
+ #[test]
+ fn topk_not_pushed_for_inner_with_filter_no_on() -> Result<()> {
+ let t1 = test_table_scan_with_name("t1")?;
+ let t2 = test_table_scan_with_name("t2")?;
+
+ let plan = LogicalPlanBuilder::from(t1)
+ .join_on(
+ LogicalPlanBuilder::from(t2).build()?,
+ JoinType::Inner,
+ vec![col("t1.b").gt(col("t2.b"))],
+ )?
+ .sort_with_limit(vec![col("t1.b").sort(true, false)], Some(3))?
+ .build()?;
+
+ assert_optimized_plan_equal!(
+ plan,
+ @r"
+ Sort: t1.b ASC NULLS LAST, fetch=3
+ Inner Join: Filter: t1.b > t2.b
+ TableScan: t1
+ TableScan: t2
+ "
+ )
+ }
+
+ /// LEFT MARK join: one record per left row → pushdown to left.
+ #[test]
+ fn topk_pushed_to_left_of_left_mark_join() -> Result<()> {
+ let t1 = test_table_scan_with_name("t1")?;
+ let t2 = test_table_scan_with_name("t2")?;
+
+ let plan = LogicalPlanBuilder::from(t1)
+ .join(
+ LogicalPlanBuilder::from(t2).build()?,
+ JoinType::LeftMark,
+ (vec!["a"], vec!["a"]),
+ None,
+ )?
+ .sort_with_limit(vec![col("t1.b").sort(true, false)], Some(3))?
+ .build()?;
+
+ assert_optimized_plan_equal!(
+ plan,
+ @r"
+ Sort: t1.b ASC NULLS LAST, fetch=3
+ LeftMark Join: t1.a = t2.a
+ Sort: t1.b ASC NULLS LAST, fetch=3
+ TableScan: t1
+ TableScan: t2
+ "
+ )
+ }
+
+ /// RIGHT MARK join: symmetric to LeftMark.
+ #[test]
+ fn topk_pushed_to_right_of_right_mark_join() -> Result<()> {
+ let t1 = test_table_scan_with_name("t1")?;
+ let t2 = test_table_scan_with_name("t2")?;
+
+ let plan = LogicalPlanBuilder::from(t1)
+ .join(
+ LogicalPlanBuilder::from(t2).build()?,
+ JoinType::RightMark,
+ (vec!["a"], vec!["a"]),
+ None,
+ )?
+ .sort_with_limit(vec![col("t2.b").sort(true, false)], Some(3))?
+ .build()?;
+
+ assert_optimized_plan_equal!(
+ plan,
+ @r"
+ Sort: t2.b ASC NULLS LAST, fetch=3
+ RightMark Join: t1.a = t2.a
+ TableScan: t1
+ Sort: t2.b ASC NULLS LAST, fetch=3
+ TableScan: t2
+ "
+ )
+ }
+
+ /// LEFT JOIN but sort on right-side columns → no pushdown.
+ #[test]
+ fn topk_not_pushed_for_wrong_side() -> Result<()> {
+ let t1 = test_table_scan_with_name("t1")?;
+ let t2 = test_table_scan_with_name("t2")?;
+
+ let plan = LogicalPlanBuilder::from(t1)
+ .join(
+ LogicalPlanBuilder::from(t2).build()?,
+ JoinType::Left,
+ (vec!["a"], vec!["a"]),
+ None,
+ )?
+ .sort_with_limit(vec![col("t2.b").sort(true, false)], Some(3))?
+ .build()?;
+
+ assert_optimized_plan_equal!(
+ plan,
+ @r"
+ Sort: t2.b ASC NULLS LAST, fetch=3
+ Left Join: t1.a = t2.a
+ TableScan: t1
+ TableScan: t2
+ "
+ )
+ }
+
+ /// Join with non-equijoin filter → pushdown still happens.
+ #[test]
+ fn topk_pushed_with_join_filter() -> Result<()> {
+ let t1 = test_table_scan_with_name("t1")?;
+ let t2 = test_table_scan_with_name("t2")?;
+
+ let plan = LogicalPlanBuilder::from(t1)
+ .join_on(
+ LogicalPlanBuilder::from(t2).build()?,
+ JoinType::Left,
+ vec![col("t1.a").eq(col("t2.a"))],
+ )?
+ .sort_with_limit(vec![col("t1.b").sort(true, false)], Some(3))?
+ .build()?;
+
+ assert_optimized_plan_equal!(
+ plan,
+ @r"
+ Sort: t1.b ASC NULLS LAST, fetch=3
+ Left Join: Filter: t1.a = t2.a
+ Sort: t1.b ASC NULLS LAST, fetch=3
+ TableScan: t1
+ TableScan: t2
+ "
+ )
+ }
+
+ /// Sort without fetch → no pushdown.
+ #[test]
+ fn topk_not_pushed_without_fetch() -> Result<()> {
+ let t1 = test_table_scan_with_name("t1")?;
+ let t2 = test_table_scan_with_name("t2")?;
+
+ let plan = LogicalPlanBuilder::from(t1)
+ .join(
+ LogicalPlanBuilder::from(t2).build()?,
+ JoinType::Left,
+ (vec!["a"], vec!["a"]),
+ None,
+ )?
+ .sort(vec![col("t1.b").sort(true, false)])?
+ .build()?;
+
+ assert_optimized_plan_equal!(
+ plan,
+ @r"
+ Sort: t1.b ASC NULLS LAST
+ Left Join: t1.a = t2.a
+ TableScan: t1
+ TableScan: t2
+ "
+ )
+ }
+
+ /// LEFT SEMI JOIN: not all left rows appear in output → no pushdown.
+ #[test]
+ fn topk_not_pushed_for_left_semi_join() -> Result<()> {
+ let t1 = test_table_scan_with_name("t1")?;
+ let t2 = test_table_scan_with_name("t2")?;
+
+ let plan = LogicalPlanBuilder::from(t1)
+ .join(
+ LogicalPlanBuilder::from(t2).build()?,
+ JoinType::LeftSemi,
+ (vec!["a"], vec!["a"]),
+ None,
+ )?
+ .sort_with_limit(vec![col("t1.b").sort(true, false)], Some(3))?
+ .build()?;
+
+ assert_optimized_plan_equal!(
+ plan,
+ @r"
+ Sort: t1.b ASC NULLS LAST, fetch=3
+ LeftSemi Join: t1.a = t2.a
+ TableScan: t1
+ TableScan: t2
+ "
+ )
+ }
+
+ /// LEFT ANTI JOIN: not all left rows appear in output → no pushdown.
+ #[test]
+ fn topk_not_pushed_for_left_anti_join() -> Result<()> {
+ let t1 = test_table_scan_with_name("t1")?;
+ let t2 = test_table_scan_with_name("t2")?;
+
+ let plan = LogicalPlanBuilder::from(t1)
+ .join(
+ LogicalPlanBuilder::from(t2).build()?,
+ JoinType::LeftAnti,
+ (vec!["a"], vec!["a"]),
+ None,
+ )?
+ .sort_with_limit(vec![col("t1.b").sort(true, false)], Some(3))?
+ .build()?;
+
+ assert_optimized_plan_equal!(
+ plan,
+ @r"
+ Sort: t1.b ASC NULLS LAST, fetch=3
+ LeftAnti Join: t1.a = t2.a
+ TableScan: t1
+ TableScan: t2
+ "
+ )
+ }
+
+ /// RIGHT SEMI JOIN: not all right rows appear in output → no pushdown.
+ #[test]
+ fn topk_not_pushed_for_right_semi_join() -> Result<()> {
+ let t1 = test_table_scan_with_name("t1")?;
+ let t2 = test_table_scan_with_name("t2")?;
+
+ let plan = LogicalPlanBuilder::from(t1)
+ .join(
+ LogicalPlanBuilder::from(t2).build()?,
+ JoinType::RightSemi,
+ (vec!["a"], vec!["a"]),
+ None,
+ )?
+ .sort_with_limit(vec![col("t2.b").sort(true, false)], Some(3))?
+ .build()?;
+
+ assert_optimized_plan_equal!(
+ plan,
+ @r"
+ Sort: t2.b ASC NULLS LAST, fetch=3
+ RightSemi Join: t1.a = t2.a
+ TableScan: t1
+ TableScan: t2
+ "
+ )
+ }
+
+ /// RIGHT ANTI JOIN: not all right rows appear in output → no pushdown.
+ #[test]
+ fn topk_not_pushed_for_right_anti_join() -> Result<()> {
+ let t1 = test_table_scan_with_name("t1")?;
+ let t2 = test_table_scan_with_name("t2")?;
+
+ let plan = LogicalPlanBuilder::from(t1)
+ .join(
+ LogicalPlanBuilder::from(t2).build()?,
+ JoinType::RightAnti,
+ (vec!["a"], vec!["a"]),
+ None,
+ )?
+ .sort_with_limit(vec![col("t2.b").sort(true, false)], Some(3))?
+ .build()?;
+
+ assert_optimized_plan_equal!(
+ plan,
+ @r"
+ Sort: t2.b ASC NULLS LAST, fetch=3
+ RightAnti Join: t1.a = t2.a
+ TableScan: t1
+ TableScan: t2
+ "
+ )
+ }
+
+ /// Multi-column sort with columns from both sides → no pushdown.
+ #[test]
+ fn topk_not_pushed_for_mixed_side_sort() -> Result<()> {
+ let t1 = test_table_scan_with_name("t1")?;
+ let t2 = test_table_scan_with_name("t2")?;
+
+ let plan = LogicalPlanBuilder::from(t1)
+ .join(
+ LogicalPlanBuilder::from(t2).build()?,
+ JoinType::Left,
+ (vec!["a"], vec!["a"]),
+ None,
+ )?
+ .sort_with_limit(
+ vec![col("t1.b").sort(true, false), col("t2.b").sort(true, false)],
+ Some(3),
+ )?
+ .build()?;
+
+ assert_optimized_plan_equal!(
+ plan,
+ @r"
+ Sort: t1.b ASC NULLS LAST, t2.b ASC NULLS LAST, fetch=3
+ Left Join: t1.a = t2.a
+ TableScan: t1
+ TableScan: t2
+ "
+ )
+ }
+
+ /// Preserved child has a larger fetch → push our tighter limit.
+ #[test]
+ fn topk_pushed_when_child_has_larger_fetch() -> Result<()> {
+ let t1 = test_table_scan_with_name("t1")?;
+ let t2 = test_table_scan_with_name("t2")?;
+
+ let t1_with_sort = LogicalPlanBuilder::from(t1)
+ .sort_with_limit(vec![col("t1.b").sort(true, false)], Some(10))?
+ .build()?;
+
+ let plan = LogicalPlanBuilder::from(t1_with_sort)
+ .join(
+ LogicalPlanBuilder::from(t2).build()?,
+ JoinType::Left,
+ (vec!["a"], vec!["a"]),
+ None,
+ )?
+ .sort_with_limit(vec![col("t1.b").sort(true, false)], Some(3))?
+ .build()?;
+
+ assert_optimized_plan_equal!(
+ plan,
+ @r"
+ Sort: t1.b ASC NULLS LAST, fetch=3
+ Left Join: t1.a = t2.a
+ Sort: t1.b ASC NULLS LAST, fetch=3
+ TableScan: t1
+ TableScan: t2
+ "
+ )
+ }
+
+ /// Preserved child already has a tighter fetch → skip pushdown.
+ #[test]
+ fn topk_not_pushed_when_child_has_smaller_fetch() -> Result<()> {
+ let t1 = test_table_scan_with_name("t1")?;
+ let t2 = test_table_scan_with_name("t2")?;
+
+ let t1_with_sort = LogicalPlanBuilder::from(t1)
+ .sort_with_limit(vec![col("t1.b").sort(true, false)], Some(2))?
+ .build()?;
+
+ let plan = LogicalPlanBuilder::from(t1_with_sort)
+ .join(
+ LogicalPlanBuilder::from(t2).build()?,
+ JoinType::Left,
+ (vec!["a"], vec!["a"]),
+ None,
+ )?
+ .sort_with_limit(vec![col("t1.b").sort(true, false)], Some(5))?
+ .build()?;
+
+ assert_optimized_plan_equal!(
+ plan,
+ @r"
+ Sort: t1.b ASC NULLS LAST, fetch=5
+ Left Join: t1.a = t2.a
+ Sort: t1.b ASC NULLS LAST, fetch=2
+ TableScan: t1
+ TableScan: t2
+ "
+ )
+ }
+
+ /// Projection passthrough: sort expr matches a projected column directly.
+ #[test]
+ fn resolve_through_projection_passthrough() -> Result<()> {
+ let t1 = test_table_scan_with_name("t1")?;
+ let plan = LogicalPlanBuilder::from(t1)
+ .project(vec![col("t1.a"), col("t1.b")])?
+ .build()?;
+ let LogicalPlan::Projection(proj) = &plan else {
+ panic!("expected Projection");
+ };
+ let sort_exprs = vec![col("t1.b").sort(true, false)];
+ let resolved = resolve_sort_exprs_through_projection(&sort_exprs, proj)?;
+ assert_eq!(resolved.len(), 1);
+ assert_eq!(resolved[0].expr.to_string(), "t1.b");
+ assert!(resolved[0].asc);
+ Ok(())
+ }
+
+ /// Projection alias: sort expr references an alias mapping to a negation.
+ #[test]
+ fn resolve_through_projection_alias() -> Result<()> {
+ let t1 = test_table_scan_with_name("t1")?;
+ let plan = LogicalPlanBuilder::from(t1)
+ .project(vec![
+ col("t1.a"),
+ (Expr::Negative(Box::new(col("t1.b")))).alias("neg_b"),
+ ])?
+ .build()?;
+ let LogicalPlan::Projection(proj) = &plan else {
+ panic!("expected Projection");
+ };
+ let sort_exprs = vec![col("neg_b").sort(true, false)];
+ let resolved = resolve_sort_exprs_through_projection(&sort_exprs, proj)?;
+ assert_eq!(resolved.len(), 1);
+ assert_eq!(resolved[0].expr.to_string(), "(- t1.b)");
+ Ok(())
+ }
+
+ /// Regression test for a `flat_name()` collision: an unqualified alias
+ /// literally named "t1.b" (as produced by a quoted identifier like
+ /// `t2.y AS "t1.b"`) must not be confused with the qualified column
+ /// `t1.b` that is also present in the same projection. Both used to
+ /// stringify to the same `flat_name()` ("t1.b"), so the old
+ /// `HashMap<String, Expr>` would silently resolve the alias to the
+ /// wrong underlying expression.
+ ///
+ /// Built via `Projection::try_new` directly rather than
+ /// `LogicalPlanBuilder::project`: the builder's `project()` rejects two
+ /// expressions with the same `schema_name()` up front (so this exact
+ /// shape can't come from a `SELECT` list), but nothing stops a
+ /// `Projection` built by other means — e.g. a different optimizer rule,
+ /// or a plan constructed directly via the expr API — from having two
+ /// structurally distinct output columns whose *string* names collide.
+ /// `resolve_sort_exprs_through_projection` must handle that shape
+ /// correctly regardless of how the `Projection` was constructed.
+ #[test]
+ fn resolve_through_projection_quoted_alias_no_collision() -> Result<()> {
+ let t1 = test_table_scan_with_name("t1")?;
+ let t2 = test_table_scan_with_name("t2")?;
+ let join = LogicalPlanBuilder::from(t1)
+ .join(
+ LogicalPlanBuilder::from(t2).build()?,
+ JoinType::Left,
+ (vec!["a"], vec!["a"]),
+ None,
+ )?
+ .build()?;
+ let proj = Projection::try_new(
+ vec![col("t1.a"), col("t2.b").alias("t1.b"), col("t1.b")],
+ Arc::new(join),
+ )?;
+
+ // `"t1.b"` here is the unqualified alias output column, not the
+ // qualified `t1.b` column also present in the projection.
+ let quoted_alias = Expr::Column(Column::new_unqualified("t1.b"));
+ let sort_exprs = vec![quoted_alias.sort(true, false)];
+ let resolved = resolve_sort_exprs_through_projection(&sort_exprs, &proj)?;
+ assert_eq!(resolved.len(), 1);
+ assert_eq!(resolved[0].expr.to_string(), "t2.b");
+ Ok(())
+ }
+
+ /// Multi-column resolution preserves direction and nulls_first per column.
+ #[test]
+ fn resolve_through_projection_multi_column() -> Result<()> {
+ let t1 = test_table_scan_with_name("t1")?;
+ let plan = LogicalPlanBuilder::from(t1)
+ .project(vec![col("t1.a"), col("t1.b"), col("t1.c")])?
+ .build()?;
+ let LogicalPlan::Projection(proj) = &plan else {
+ panic!("expected Projection");
+ };
+ let sort_exprs =
+ vec![col("t1.a").sort(true, false), col("t1.b").sort(false, true)];
+ let resolved = resolve_sort_exprs_through_projection(&sort_exprs, proj)?;
+ assert_eq!(resolved.len(), 2);
+ assert_eq!(resolved[0].expr.to_string(), "t1.a");
+ assert!(resolved[0].asc);
+ assert_eq!(resolved[1].expr.to_string(), "t1.b");
+ assert!(!resolved[1].asc);
+ assert!(resolved[1].nulls_first);
+ Ok(())
+ }
+
+ /// Stacked Projection + SubqueryAlias: resolve through both layers.
+ #[test]
+ fn resolve_through_projection_and_subquery_alias() -> Result<()> {
+ let t1 = test_table_scan_with_name("t1")?;
+ let plan = LogicalPlanBuilder::from(t1)
+ .alias("sub")?
+ .project(vec![
+ col("sub.a"),
+ (Expr::Negative(Box::new(col("sub.b")))).alias("neg_b"),
+ ])?
+ .build()?;
+ let LogicalPlan::Projection(proj) = &plan else {
+ panic!("expected Projection");
+ };
+ let LogicalPlan::SubqueryAlias(sq) = proj.input.as_ref() else {
+ panic!("expected SubqueryAlias");
+ };
+ let sort_exprs = vec![col("neg_b").sort(true, false)];
+ let after_proj = resolve_sort_exprs_through_projection(&sort_exprs, proj)?;
+ assert_eq!(after_proj[0].expr.to_string(), "(- sub.b)");
+ let after_sq = resolve_sort_exprs_through_subquery_alias(&after_proj, sq)?;
+ assert_eq!(after_sq[0].expr.to_string(), "(- t1.b)");
+ assert!(after_sq[0].asc);
+ assert!(!after_sq[0].nulls_first);
+ Ok(())
+ }
+
+ /// Simple SubqueryAlias resolution: sub.b → t1.b.
+ #[test]
+ fn resolve_through_subquery_alias_simple() -> Result<()> {
+ let t1 = test_table_scan_with_name("t1")?;
+ let plan = LogicalPlanBuilder::from(t1).alias("sub")?.build()?;
+ let LogicalPlan::SubqueryAlias(sq) = &plan else {
+ panic!("expected SubqueryAlias");
+ };
+ let sort_exprs = vec![col("sub.b").sort(true, false)];
+ let resolved = resolve_sort_exprs_through_subquery_alias(&sort_exprs, sq)?;
+ assert_eq!(resolved.len(), 1);
+ assert_eq!(resolved[0].expr.to_string(), "t1.b");
+ assert!(resolved[0].asc);
+ assert!(!resolved[0].nulls_first);
+ Ok(())
+ }
+
+ /// Multi-column SubqueryAlias resolution preserves direction per column.
+ #[test]
+ fn resolve_through_subquery_alias_multi_column() -> Result<()> {
+ let t1 = test_table_scan_with_name("t1")?;
+ let plan = LogicalPlanBuilder::from(t1).alias("sub")?.build()?;
+ let LogicalPlan::SubqueryAlias(sq) = &plan else {
+ panic!("expected SubqueryAlias");
+ };
+ let sort_exprs = vec![
+ col("sub.a").sort(true, false),
+ col("sub.b").sort(false, true),
+ ];
+ let resolved = resolve_sort_exprs_through_subquery_alias(&sort_exprs, sq)?;
+ assert_eq!(resolved.len(), 2);
+ assert_eq!(resolved[0].expr.to_string(), "t1.a");
+ assert!(resolved[0].asc);
+ assert_eq!(resolved[1].expr.to_string(), "t1.b");
+ assert!(!resolved[1].asc);
+ assert!(resolved[1].nulls_first);
+ Ok(())
+ }
+
+ /// SubqueryAlias with a different alias name (foo ≠ t1).
+ #[test]
+ fn resolve_through_subquery_alias_different_name() -> Result<()> {
+ let t1 = test_table_scan_with_name("t1")?;
+ let plan = LogicalPlanBuilder::from(t1).alias("foo")?.build()?;
+ let LogicalPlan::SubqueryAlias(sq) = &plan else {
+ panic!("expected SubqueryAlias");
+ };
+ let sort_exprs = vec![col("foo.b").sort(true, false)];
+ let resolved = resolve_sort_exprs_through_subquery_alias(&sort_exprs, sq)?;
+ assert_eq!(resolved.len(), 1);
+ assert_eq!(resolved[0].expr.to_string(), "t1.b");
+ Ok(())
+ }
+
+ /// SubqueryAlias with nested expression: (- sub.b) → (- t1.b).
+ #[test]
+ fn resolve_through_subquery_alias_nested_expr() -> Result<()> {
+ let t1 = test_table_scan_with_name("t1")?;
+ let plan = LogicalPlanBuilder::from(t1).alias("sub")?.build()?;
+ let LogicalPlan::SubqueryAlias(sq) = &plan else {
+ panic!("expected SubqueryAlias");
+ };
+ let sort_exprs = vec![Expr::Negative(Box::new(col("sub.b"))).sort(true, false)];
+ let resolved = resolve_sort_exprs_through_subquery_alias(&sort_exprs, sq)?;
+ assert_eq!(resolved.len(), 1);
+ assert_eq!(resolved[0].expr.to_string(), "(- t1.b)");
+ assert!(resolved[0].asc);
+ Ok(())
+ }
+
+ /// Inner Sort has different exprs WITH fetch → stacked sorts.
+ #[test]
+ fn topk_stacked_when_child_has_different_exprs_with_fetch() -> Result<()> {
+ let t1 = test_table_scan_with_name("t1")?;
+ let t2 = test_table_scan_with_name("t2")?;
+
+ let t1_with_sort = LogicalPlanBuilder::from(t1)
+ .sort_with_limit(vec![col("t1.a").sort(true, false)], Some(5))?
+ .build()?;
+
+ let plan = LogicalPlanBuilder::from(t1_with_sort)
+ .join(
+ LogicalPlanBuilder::from(t2).build()?,
+ JoinType::Left,
+ (vec!["a"], vec!["a"]),
+ None,
+ )?
+ .sort_with_limit(vec![col("t1.b").sort(true, false)], Some(2))?
+ .build()?;
+
+ assert_optimized_plan_equal!(
+ plan,
+ @r"
+ Sort: t1.b ASC NULLS LAST, fetch=2
+ Left Join: t1.a = t2.a
+ Sort: t1.b ASC NULLS LAST, fetch=2
+ Sort: t1.a ASC NULLS LAST, fetch=5
+ TableScan: t1
+ TableScan: t2
+ "
+ )
+ }
+
+ /// Inner Sort has different exprs WITHOUT fetch → stacked sorts.
+ #[test]
+ fn topk_stacked_when_child_has_different_exprs_no_fetch() -> Result<()> {
+ let t1 = test_table_scan_with_name("t1")?;
+ let t2 = test_table_scan_with_name("t2")?;
+
+ let t1_with_sort = LogicalPlanBuilder::from(t1)
+ .sort(vec![col("t1.a").sort(true, false)])?
+ .build()?;
+
+ let plan = LogicalPlanBuilder::from(t1_with_sort)
+ .join(
+ LogicalPlanBuilder::from(t2).build()?,
+ JoinType::Left,
+ (vec!["a"], vec!["a"]),
+ None,
+ )?
+ .sort_with_limit(vec![col("t1.b").sort(true, false)], Some(2))?
+ .build()?;
+
+ assert_optimized_plan_equal!(
+ plan,
+ @r"
+ Sort: t1.b ASC NULLS LAST, fetch=2
+ Left Join: t1.a = t2.a
+ Sort: t1.b ASC NULLS LAST, fetch=2
+ Sort: t1.a ASC NULLS LAST
+ TableScan: t1
+ TableScan: t2
+ "
+ )
+ }
+
+ /// End-to-end regression for the `flat_name()` collision: `t2.b AS
+ /// "t1.b"` (an unqualified alias literally named "t1.b") sits next to
+ /// the qualified column `t1.b` in the same projection. ORDER BY the
+ /// quoted alias must push down sorted on `t2.b` (the right/preserved
+ /// side of this RIGHT JOIN), not silently resolve to `t1.b` and push
+ /// the wrong TopK onto the wrong side.
+ #[test]
+ fn topk_pushed_through_projection_quoted_alias_no_collision() -> Result<()> {
+ let t1 = test_table_scan_with_name("t1")?;
+ let t2 = test_table_scan_with_name("t2")?;
+
+ let join = LogicalPlanBuilder::from(t1)
+ .join(
+ LogicalPlanBuilder::from(t2).build()?,
+ JoinType::Right,
+ (vec!["a"], vec!["a"]),
+ None,
+ )?
+ .build()?;
+ // Built via `Projection::try_new` (see comment on
+ // `resolve_through_projection_quoted_alias_no_collision`): the
+ // builder's `project()` would reject this expression list outright.
+ let proj = LogicalPlan::Projection(Projection::try_new(
+ vec![col("t1.a"), col("t2.b").alias("t1.b"), col("t1.b")],
+ Arc::new(join),
+ )?);
+ let plan = LogicalPlanBuilder::from(proj)
+ .sort_with_limit(
+ vec![Expr::Column(Column::new_unqualified("t1.b")).sort(true, false)],
+ Some(3),
+ )?
+ .build()?;
+
+ assert_optimized_plan_equal!(
+ plan,
+ @r#"
+ Sort: t1.b ASC NULLS LAST, fetch=3
+ Projection: t1.a, t2.b AS t1.b, t1.b
+ Right Join: t1.a = t2.a
+ TableScan: t1
+ Sort: t2.b ASC NULLS LAST, fetch=3
+ TableScan: t2
+ "#
+ )
+ }
+
+ /// Plan-level regression for the CROSS JOIN empty-side hazard: a
+ /// division (fallible) sort key over a CROSS JOIN must never be pushed
+ /// down, since doing so would evaluate it eagerly on rows that a
+ /// correct, unoptimized execution would never reach if the other side
+ /// turned out to be empty (see `topk_pushed_for_cross_join_with_empty_other_side_does_not_error`
+ /// in the SLT suite for the corresponding end-to-end execution case).
+ #[test]
+ fn topk_not_pushed_for_cross_join_with_fallible_sort_expr() -> Result<()> {
+ let t1 = test_table_scan_with_name("t1")?;
+ let t2 = test_table_scan_with_name("t2")?;
+
+ let plan = LogicalPlanBuilder::from(t1)
+ .cross_join(LogicalPlanBuilder::from(t2).build()?)?
+ .sort_with_limit(
+ vec![(col("t1.a") / col("t1.b")).sort(true, false)],
+ Some(3),
+ )?
+ .build()?;
+
+ assert_optimized_plan_equal!(
+ plan,
+ @r"
+ Sort: t1.a / t1.b ASC NULLS LAST, fetch=3
+ Cross Join:
+ TableScan: t1
+ TableScan: t2
+ "
+ )
+ }
+
+ /// A sort key with no column references at all (e.g. a literal) can't
+ /// distinguish any row from any other on either side. `has_all_column_refs`
+ /// is vacuously true for it, so without an explicit guard this would
+ /// still push a Sort onto the first candidate side — legal (any N rows
+ /// satisfy an all-ties ORDER BY) but wasted work, since the pushed
+ /// Sort's "top N" doesn't depend on that side's data at all.
+ #[test]
+ fn topk_not_pushed_for_sort_key_with_no_column_refs() -> Result<()> {
+ let t1 = test_table_scan_with_name("t1")?;
+ let t2 = test_table_scan_with_name("t2")?;
+
+ let plan = LogicalPlanBuilder::from(t1)
+ .join(
+ LogicalPlanBuilder::from(t2).build()?,
+ JoinType::Left,
+ (vec!["a"], vec!["a"]),
+ None,
+ )?
+ .sort_with_limit(vec![lit(1i64).sort(true, false)], Some(3))?
+ .build()?;
+
+ assert_optimized_plan_equal!(
+ plan,
+ @r"
+ Sort: Int64(1) ASC NULLS LAST, fetch=3
+ Left Join: t1.a = t2.a
+ TableScan: t1
+ TableScan: t2
+ "
+ )
+ }
+
+ /// `Limit(skip=0) → Sort(no fetch) → Join`: `rewrite_limit` merges the
+ /// Limit's fetch into the Sort on one visit, but only strips the now-
+ /// redundant `Limit(skip=0)` wrapper on a *later* visit to that same
+ /// node (it only recognizes the wrapper as redundant once the Sort
+ /// already carries the matching fetch). Confirms the TopK-through-join
+ /// pushdown and the Limit/Sort merge both still converge to the fully
+ /// collapsed form with two passes.
+ #[test]
+ fn topk_pushed_through_limit_then_sort_with_two_passes() -> Result<()> {
+ let t1 = test_table_scan_with_name("t1")?;
+ let t2 = test_table_scan_with_name("t2")?;
+
+ let plan = LogicalPlanBuilder::from(t1)
+ .join(
+ LogicalPlanBuilder::from(t2).build()?,
+ JoinType::Left,
+ (vec!["a"], vec!["a"]),
+ None,
+ )?
+ .sort(vec![col("t1.b").sort(true, false)])?
+ .limit(0, Some(3))?
+ .build()?;
+
+ let optimizer_ctx = OptimizerContext::new().with_max_passes(2);
+ let rules: Vec<Arc<dyn crate::OptimizerRule + Send + Sync>> =
+ vec![Arc::new(PushDownLimit::new())];
+ assert_optimized_plan_eq_snapshot!(
+ optimizer_ctx,
+ rules,
+ plan,
+ @r"
+ Sort: t1.b ASC NULLS LAST, fetch=3
+ Left Join: t1.a = t2.a
+ Sort: t1.b ASC NULLS LAST, fetch=3
+ TableScan: t1
+ TableScan: t2
+ "
+ )
+ }
+
+ /// Running the optimizer a second time on its own output must be a
+ /// no-op: a rule that keeps rewriting an already-optimized plan (e.g.
+ /// oscillating between two equivalent forms) would never let the
+ /// overall optimizer converge.
+ #[test]
+ fn topk_pushed_through_join_is_idempotent() -> Result<()> {
+ let t1 = test_table_scan_with_name("t1")?;
+ let t2 = test_table_scan_with_name("t2")?;
+
+ let plan = LogicalPlanBuilder::from(t1)
+ .join(
+ LogicalPlanBuilder::from(t2).build()?,
+ JoinType::Left,
+ (vec!["a"], vec!["a"]),
+ None,
+ )?
+ .sort_with_limit(vec![col("t1.b").sort(true, false)], Some(3))?
+ .build()?;
+
+ let rules: Vec<Arc<dyn crate::OptimizerRule + Send + Sync>> =
+ vec![Arc::new(PushDownLimit::new())];
+ let optimizer = crate::Optimizer::with_rules(rules);
+ let ctx = OptimizerContext::new().with_max_passes(1);
+
+ let once = optimizer.optimize(plan, &ctx, |_, _| {})?;
+ let twice = optimizer.optimize(once.clone(), &ctx, |_, _| {})?;
+ assert_eq!(once, twice);
+ Ok(())
+ }
+
+ /// `fetch = 0` is a degenerate but valid TopK: the pushed Sort still
+ /// carries `fetch=0` rather than being special-cased away.
+ #[test]
+ fn topk_pushed_with_fetch_zero() -> Result<()> {
+ let t1 = test_table_scan_with_name("t1")?;
+ let t2 = test_table_scan_with_name("t2")?;
+
+ let plan = LogicalPlanBuilder::from(t1)
+ .join(
+ LogicalPlanBuilder::from(t2).build()?,
+ JoinType::Left,
+ (vec!["a"], vec!["a"]),
+ None,
+ )?
+ .sort_with_limit(vec![col("t1.b").sort(true, false)], Some(0))?
+ .build()?;
+
+ assert_optimized_plan_equal!(
+ plan,
+ @r"
+ Sort: t1.b ASC NULLS LAST, fetch=0
+ Left Join: t1.a = t2.a
+ Sort: t1.b ASC NULLS LAST, fetch=0
+ TableScan: t1
+ TableScan: t2
+ "
+ )
+ }
+
+ /// Preserved child is a bare `Limit` with no `Sort` beneath it at all
+ /// (`LEFT JOIN (SELECT * FROM t1 LIMIT 5)`) — there is nothing for the
+ /// Limit to ever merge with, so this is a permanent missed
+ /// optimization, not a transient one that resolves on a later pass (see
+ /// the comment at the `matches!(inner_child, LogicalPlan::Limit(_))`
+ /// check). It must still skip cleanly here rather than push past the
+ /// Limit and silently ignore the existing row cap.
+ #[test]
+ fn topk_not_pushed_when_child_is_bare_limit() -> Result<()> {
+ let t1 = test_table_scan_with_name("t1")?;
+ let t2 = test_table_scan_with_name("t2")?;
+
+ let t1_with_limit = LogicalPlanBuilder::from(t1).limit(0, Some(5))?.build()?;
+
+ let plan = LogicalPlanBuilder::from(t1_with_limit)
+ .join(
+ LogicalPlanBuilder::from(t2).build()?,
+ JoinType::Left,
+ (vec!["a"], vec!["a"]),
+ None,
+ )?
+ .sort_with_limit(vec![col("t1.b").sort(true, false)], Some(3))?
+ .build()?;
+
+ assert_optimized_plan_equal!(
+ plan,
+ @r"
+ Sort: t1.b ASC NULLS LAST, fetch=3
+ Left Join: t1.a = t2.a
+ Limit: skip=0, fetch=5
+ TableScan: t1, fetch=5
+ TableScan: t2
+ "
+ )
+ }
+
+ /// A `Filter` between the `Sort` and the `Join` blocks pushdown: `Filter`
+ /// is not one of the transparent nodes this rewrite peels through (a
+ /// `Filter` can drop preserved-side rows, which would change which rows
+ /// survive to be sorted). Running `PushDownLimit` alone can never push
+ /// through it, however many passes are allowed.
+ #[test]
+ fn topk_not_pushed_through_filter_with_push_down_limit_alone() -> Result<()> {
+ let t1 = test_table_scan_with_name("t1")?;
+ let t2 = test_table_scan_with_name("t2")?;
+
+ let plan = LogicalPlanBuilder::from(t1)
+ .join(
+ LogicalPlanBuilder::from(t2).build()?,
+ JoinType::Left,
+ (vec!["a"], vec!["a"]),
+ None,
+ )?
+ .filter(col("t1.c").eq(lit("foo")))?
+ .sort_with_limit(vec![col("t1.b").sort(true, false)], Some(3))?
+ .build()?;
+
+ assert_optimized_plan_equal!(
+ plan,
+ @r#"
+ Sort: t1.b ASC NULLS LAST, fetch=3
+ Filter: t1.c = Utf8("foo")
+ Left Join: t1.a = t2.a
+ TableScan: t1
+ TableScan: t2
+ "#
+ )
+ }
+
+ /// Same plan as above, but with `PushDownFilter` also in the rule set
+ /// and two passes allowed: pass 1 runs `push_down_limit` first (declines,
+ /// same as the single-rule case above) and then `push_down_filter` sinks
+ /// the Filter onto `t1` (the only side it references) below the Join;
+ /// pass 2 revisits the Sort, which now sits directly above the Join and
+ /// pushes through. One pass alone is not enough — see
+ /// `topk_not_pushed_through_filter_with_push_down_limit_alone` and the
+ /// module doc comment's note on this ordering dependency.
+ #[test]
+ fn topk_pushed_through_filter_after_push_down_filter_with_two_passes() -> Result<()> {
+ let t1 = test_table_scan_with_name("t1")?;
+ let t2 = test_table_scan_with_name("t2")?;
+
+ let plan = LogicalPlanBuilder::from(t1)
+ .join(
+ LogicalPlanBuilder::from(t2).build()?,
+ JoinType::Left,
+ (vec!["a"], vec!["a"]),
+ None,
+ )?
+ .filter(col("t1.c").eq(lit("foo")))?
+ .sort_with_limit(vec![col("t1.b").sort(true, false)], Some(3))?
+ .build()?;
+
+ let optimizer_ctx = OptimizerContext::new().with_max_passes(2);
+ let rules: Vec<Arc<dyn crate::OptimizerRule + Send + Sync>> = vec![
+ Arc::new(PushDownLimit::new()),
+ Arc::new(crate::push_down_filter::PushDownFilter::new()),
+ ];
+ assert_optimized_plan_eq_snapshot!(
+ optimizer_ctx,
+ rules,
+ plan,
+ @r#"
+ Sort: t1.b ASC NULLS LAST, fetch=3
+ Left Join: t1.a = t2.a
+ Sort: t1.b ASC NULLS LAST, fetch=3
+ TableScan: t1, full_filters=[t1.c = Utf8("foo")]
+ TableScan: t2
+ "#
+ )
+ }
+
+ /// Regression test using the *real* default rule set and order
+ /// (`Optimizer::new()`), not the 2-rule subset above: with
+ /// `max_passes(1)` — e.g. `datafusion.optimizer.max_passes = 1` — the
+ /// TopK is never pushed down for `Sort(fetch) -> Filter -> Left Join`,
+ /// even though `PushDownFilter` does sink the filter into the scan
+ /// within that same single pass. This is because `push_down_limit`
+ /// (which currently owns the TopK-through-join rewrite) runs *before*
+ /// `push_down_filter` in `Optimizer::new()`'s rule list, so by the time
+ /// this rewrite's tree walk reaches the Filter, the Filter hasn't been
+ /// sunk below the Join yet — and `push_down_limit` is never revisited
+ /// again within the same pass. A later fix should let the specialized
+ /// TopK-through-join transformation run after filter pushdown without
+ /// moving the ordinary Limit transformation past filters (see the
+ /// module doc comment).
+ #[test]
+ fn topk_not_pushed_with_default_optimizer_and_one_pass() -> Result<()> {
+ let t1 = test_table_scan_with_name("t1")?;
+ let t2 = test_table_scan_with_name("t2")?;
+
+ let plan = LogicalPlanBuilder::from(t1)
+ .join(
+ LogicalPlanBuilder::from(t2).build()?,
+ JoinType::Left,
+ (vec!["a"], vec!["a"]),
+ None,
+ )?
+ .filter(col("t1.c").eq(lit("foo")))?
+ .sort_with_limit(vec![col("t1.b").sort(true, false)], Some(3))?
+ .build()?;
+
+ let optimizer_ctx = OptimizerContext::new().with_max_passes(1);
+ let optimizer = crate::Optimizer::new();
+ let optimized = optimizer.optimize(plan, &optimizer_ctx, |_, _| {})?;
+
+ insta::assert_snapshot!(optimized, @r#"
+ Sort: t1.b ASC NULLS LAST, fetch=3
+ Left Join: t1.a = t2.a
+ TableScan: t1 projection=[a, b, c], full_filters=[t1.c = Utf8("foo")]
+ TableScan: t2 projection=[a, b, c]
+ "#);
+ Ok(())
+ }
+}
diff --git a/datafusion/sqllogictest/test_files/push_down_topk_through_join.slt b/datafusion/sqllogictest/test_files/push_down_topk_through_join.slt
index bdc04786f5..aafe2ec05f 100644
--- a/datafusion/sqllogictest/test_files/push_down_topk_through_join.slt
+++ b/datafusion/sqllogictest/test_files/push_down_topk_through_join.slt
@@ -17,10 +17,9 @@
# Tests for pushing a TopK (Sort with fetch) through an outer join.
#
-# These queries exercise the scenarios handled by the PushDownTopKThroughJoin
-# rule. That rule lands in a follow-up PR; the EXPLAIN plans below capture
-# current behavior, so the follow-up's diff shows exactly which plans change.
-# The query-result checks hold whether or not the rule is enabled.
+# These queries exercise the scenarios handled by the `push_down_limit`
+# rule's TopK-through-join rewrite (see
+# `datafusion/optimizer/src/push_down_limit/topk_through_join.rs`).
statement ok
set datafusion.execution.target_partitions = 1;
@@ -56,8 +55,9 @@ ORDER BY t1.b ASC LIMIT 3;
logical_plan
01)Sort: t1.b ASC NULLS LAST, fetch=3
02)--Left Join: t1.a = t2.x
-03)----TableScan: t1 projection=[a, b]
-04)----TableScan: t2 projection=[x]
+03)----Sort: t1.b ASC NULLS LAST, fetch=3
+04)------TableScan: t1 projection=[a, b]
+05)----TableScan: t2 projection=[x]
query III
SELECT t1.a, t1.b, t2.x
@@ -78,7 +78,8 @@ logical_plan
01)Sort: t2.y ASC NULLS LAST, fetch=3
02)--Right Join: t1.a = t2.x
03)----TableScan: t1 projection=[a]
-04)----TableScan: t2 projection=[x, y]
+04)----Sort: t2.y ASC NULLS LAST, fetch=3
+05)------TableScan: t2 projection=[x, y]
query III
SELECT t1.a, t2.x, t2.y
@@ -144,8 +145,9 @@ logical_plan
01)Sort: t1.b ASC NULLS LAST, fetch=3
02)--Projection: t1.a, t1.b, t2.x
03)----Left Join: t1.a = t2.x Filter: t1.b > t2.y
-04)------TableScan: t1 projection=[a, b]
-05)------TableScan: t2 projection=[x, y]
+04)------Sort: t1.b ASC NULLS LAST, fetch=3
+05)--------TableScan: t1 projection=[a, b]
+06)------TableScan: t2 projection=[x, y]
query III
SELECT t1.a, t1.b, t2.x
@@ -165,10 +167,11 @@ ORDER BY t1.b ASC LIMIT 3;
logical_plan
01)Sort: t1.b ASC NULLS LAST, fetch=3
02)--Left Join: t1.a = t2.x
-03)----TableScan: t1 projection=[a, b]
-04)----Projection: t2.x
-05)------Filter: t2.y > Int32(100)
-06)--------TableScan: t2 projection=[x, y]
+03)----Sort: t1.b ASC NULLS LAST, fetch=3
+04)------TableScan: t1 projection=[a, b]
+05)----Projection: t2.x
+06)------Filter: t2.y > Int32(100)
+07)--------TableScan: t2 projection=[x, y]
# A preserved-side filter in the ON clause suppresses matches, but the rows
# still appear NULL-filled, so it does not change which rows are preserved.
@@ -180,8 +183,9 @@ ORDER BY t1.b ASC LIMIT 3;
logical_plan
01)Sort: t1.b ASC NULLS LAST, fetch=3
02)--Left Join: t1.a = t2.x Filter: t1.b > Int32(20)
-03)----TableScan: t1 projection=[a, b]
-04)----TableScan: t2 projection=[x]
+03)----Sort: t1.b ASC NULLS LAST, fetch=3
+04)------TableScan: t1 projection=[a, b]
+05)----TableScan: t2 projection=[x]
query III
SELECT t1.a, t1.b, t2.x
@@ -231,7 +235,7 @@ logical_plan
02)--SubqueryAlias: sub
03)----Left Join: t1.a = t2.x
04)------SubqueryAlias: t1
-05)--------Sort: t1.b ASC NULLS LAST, fetch=5
+05)--------Sort: t1.b ASC NULLS LAST, fetch=2
06)----------TableScan: t1 projection=[a, b]
07)------TableScan: t2 projection=[x]
@@ -360,8 +364,9 @@ logical_plan
01)Limit: skip=1, fetch=2
02)--Sort: t1.b ASC NULLS LAST, fetch=3
03)----Left Join: t1.a = t2.x
-04)------TableScan: t1 projection=[a, b]
-05)------TableScan: t2 projection=[x]
+04)------Sort: t1.b ASC NULLS LAST, fetch=3
+05)--------TableScan: t1 projection=[a, b]
+06)------TableScan: t2 projection=[x]
query III
SELECT t1.a, t1.b, t2.x
@@ -386,8 +391,9 @@ logical_plan
01)Sort: neg_b ASC NULLS LAST, fetch=3
02)--Projection: (- t1.b) AS neg_b, t2.x
03)----Left Join: t1.a = t2.x
-04)------TableScan: t1 projection=[a, b]
-05)------TableScan: t2 projection=[x]
+04)------Sort: (- t1.b) ASC NULLS LAST, fetch=3
+05)--------TableScan: t1 projection=[a, b]
+06)------TableScan: t2 projection=[x]
# -b ascending means largest b first
query II
@@ -444,9 +450,10 @@ 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)--------TableScan: t1 projection=[a, b]
-06)------TableScan: t2 projection=[x]
+04)------Sort: t1.b ASC NULLS LAST, fetch=2
+05)--------SubqueryAlias: t1
+06)----------TableScan: t1 projection=[a, b]
+07)------TableScan: t2 projection=[x]
query III
SELECT * FROM (
@@ -475,7 +482,7 @@ logical_plan
03)----Right Join: t1.a = t2.x
04)------TableScan: t1 projection=[a]
05)------SubqueryAlias: t2
-06)--------Sort: t2.y ASC NULLS LAST, fetch=10
+06)--------Sort: t2.y ASC NULLS LAST, fetch=3
07)----------TableScan: t2 projection=[x, y]
query III
@@ -505,9 +512,10 @@ logical_plan
01)Sort: sub.b ASC NULLS LAST, fetch=3
02)--SubqueryAlias: sub
03)----Left Join: foo.a = t2.x
-04)------SubqueryAlias: foo
-05)--------TableScan: t1 projection=[a, b]
-06)------TableScan: t2 projection=[x]
+04)------Sort: foo.b ASC NULLS LAST, fetch=3
+05)--------SubqueryAlias: foo
+06)----------TableScan: t1 projection=[a, b]
+07)------TableScan: t2 projection=[x]
query III
SELECT * FROM (
@@ -568,9 +576,10 @@ logical_plan
01)Sort: sub.a ASC NULLS LAST, sub.b ASC NULLS LAST, fetch=3
02)--SubqueryAlias: sub
03)----Left Join: t1.a = t2.x
-04)------SubqueryAlias: t1
-05)--------TableScan: t1 projection=[a, b]
-06)------TableScan: t2 projection=[x]
+04)------Sort: t1.a ASC NULLS LAST, t1.b ASC NULLS LAST, fetch=3
+05)--------SubqueryAlias: t1
+06)----------TableScan: t1 projection=[a, b]
+07)------TableScan: t2 projection=[x]
query III
SELECT * FROM (
@@ -595,9 +604,10 @@ ORDER BY t1.b ASC LIMIT 3;
logical_plan
01)Sort: t1.b ASC NULLS LAST, fetch=3
02)--Left Join: t1.a = t2.x
-03)----Filter: t1.b > Int32(10)
-04)------TableScan: t1 projection=[a, b]
-05)----TableScan: t2 projection=[x]
+03)----Sort: t1.b ASC NULLS LAST, fetch=3
+04)------Filter: t1.b > Int32(10)
+05)--------TableScan: t1 projection=[a, b]
+06)----TableScan: t2 projection=[x]
query III
SELECT t1.a, t1.b, t2.x
@@ -622,8 +632,9 @@ ORDER BY t1.b DESC LIMIT 3;
logical_plan
01)Sort: t1.b DESC NULLS FIRST, fetch=3
02)--Left Join: t1.a = t2.x
-03)----TableScan: t1 projection=[a, b]
-04)----TableScan: t2 projection=[x]
+03)----Sort: t1.b DESC NULLS FIRST, fetch=3
+04)------TableScan: t1 projection=[a, b]
+05)----TableScan: t2 projection=[x]
query III
SELECT t1.a, t1.b, t2.x
@@ -643,8 +654,9 @@ ORDER BY t1.b ASC NULLS FIRST LIMIT 3;
logical_plan
01)Sort: t1.b ASC NULLS FIRST, fetch=3
02)--Left Join: t1.a = t2.x
-03)----TableScan: t1 projection=[a, b]
-04)----TableScan: t2 projection=[x]
+03)----Sort: t1.b ASC NULLS FIRST, fetch=3
+04)------TableScan: t1 projection=[a, b]
+05)----TableScan: t2 projection=[x]
query III
SELECT t1.a, t1.b, t2.x
@@ -665,7 +677,8 @@ logical_plan
01)Sort: t2.y DESC NULLS LAST, fetch=3
02)--Right Join: t1.a = t2.x
03)----TableScan: t1 projection=[a]
-04)----TableScan: t2 projection=[x, y]
+04)----Sort: t2.y DESC NULLS LAST, fetch=3
+05)------TableScan: t2 projection=[x, y]
query III
SELECT t1.a, t2.x, t2.y
@@ -680,8 +693,12 @@ NULL 6 600
### CROSS JOIN
###
-# Each left row appears |t2| times, so the top-N by left columns must come
-# from the top-N left rows.
+# CROSS JOIN is never pushed through, even when the sort key comes
+# entirely from one side: a pushed Sort would evaluate its expression
+# eagerly on that side regardless of whether the *other* side is empty, and
+# this rule has no way to prove the other side is non-empty at
+# logical-plan time. See the fallible-expression regression further below
+# for why this matters, not just for performance.
query TT
EXPLAIN SELECT t1.a, t1.b, t2.x
FROM t1 CROSS JOIN t2
@@ -735,6 +752,43 @@ logical_plan
03)----TableScan: t1 projection=[a, b]
04)----TableScan: t2 projection=[y]
+# Regression: a CROSS JOIN with a fallible sort expression (division) on
+# one side and an empty *other* side must not error. Before this rule
+# excluded CROSS JOIN, pushing `Sort(10 / l.x, fetch=1)` onto `l` would
+# evaluate the division eagerly (l.x = 0 divides by zero) even though the
+# join's actual output is empty (r has no rows), which previously returned
+# cleanly without ever evaluating the sort expression.
+statement ok
+CREATE TABLE cj_l (x INT) AS VALUES (0), (1);
+
+statement ok
+CREATE TABLE cj_r (y INT);
+
+query TT
+EXPLAIN SELECT cj_l.x
+FROM cj_l CROSS JOIN cj_r
+ORDER BY 10 / cj_l.x
+LIMIT 1;
+----
+logical_plan
+01)Sort: Int64(10) / CAST(cj_l.x AS Int64) ASC NULLS LAST, fetch=1
+02)--Cross Join:
+03)----TableScan: cj_l projection=[x]
+04)----TableScan: cj_r projection=[]
+
+query I
+SELECT cj_l.x
+FROM cj_l CROSS JOIN cj_r
+ORDER BY 10 / cj_l.x
+LIMIT 1;
+----
+
+statement ok
+DROP TABLE cj_l;
+
+statement ok
+DROP TABLE cj_r;
+
# INNER JOIN with only a non-equi filter: the filter can drop rows from either
# side, so a pushed fetch could select rows that get filtered out.
query TT
@@ -770,10 +824,12 @@ ORDER BY t1.b ASC LIMIT 2;
logical_plan
01)Sort: t1.b ASC NULLS LAST, fetch=2
02)--Left Join: t1.a = t3.p
-03)----Left Join: t1.a = t2.x
-04)------TableScan: t1 projection=[a, b]
-05)------TableScan: t2 projection=[x]
-06)----TableScan: t3 projection=[p]
+03)----Sort: t1.b ASC NULLS LAST, fetch=2
+04)------Left Join: t1.a = t2.x
+05)--------Sort: t1.b ASC NULLS LAST, fetch=2
+06)----------TableScan: t1 projection=[a, b]
+07)--------TableScan: t2 projection=[x]
+08)----TableScan: t3 projection=[p]
query IIII
SELECT t1.a, t1.b, t2.x, t3.p
@@ -812,8 +868,9 @@ ORDER BY t_tied.b ASC, t_tied.a ASC LIMIT 3;
logical_plan
01)Sort: t_tied.b ASC NULLS LAST, t_tied.a ASC NULLS LAST, fetch=3
02)--Left Join: t_tied.a = t_other.x
-03)----TableScan: t_tied projection=[a, b]
-04)----TableScan: t_other projection=[x]
+03)----Sort: t_tied.b ASC NULLS LAST, t_tied.a ASC NULLS LAST, fetch=3
+04)------TableScan: t_tied projection=[a, b]
+05)----TableScan: t_other projection=[x]
query III
SELECT t_tied.a, t_tied.b, t_other.x
@@ -849,10 +906,11 @@ logical_plan
01)Sort: outer_sub.b ASC NULLS LAST, fetch=2
02)--SubqueryAlias: outer_sub
03)----Left Join: inner_sub.a = t2.x
-04)------SubqueryAlias: inner_sub
-05)--------SubqueryAlias: inner_alias
-06)----------TableScan: t1 projection=[a, b]
-07)------TableScan: t2 projection=[x]
+04)------Sort: inner_sub.b ASC NULLS LAST, fetch=2
+05)--------SubqueryAlias: inner_sub
+06)----------SubqueryAlias: inner_alias
+07)------------TableScan: t1 projection=[a, b]
+08)------TableScan: t2 projection=[x]
query III
SELECT * FROM (
@@ -884,7 +942,7 @@ logical_plan
03)----Left Join: inner_sub.a = t2.x
04)------SubqueryAlias: inner_sub
05)--------SubqueryAlias: inner_alias
-06)----------Sort: t1.b ASC NULLS LAST, fetch=5
+06)----------Sort: t1.b ASC NULLS LAST, fetch=2
07)------------TableScan: t1 projection=[a, b]
08)------TableScan: t2 projection=[x]
@@ -951,11 +1009,12 @@ logical_plan
01)Sort: outer_sub.b ASC NULLS LAST, fetch=2
02)--SubqueryAlias: outer_sub
03)----Left Join: inner_sub.a = t2.x
-04)------SubqueryAlias: inner_sub
-05)--------SubqueryAlias: inner_alias
-06)----------Sort: t1.a ASC NULLS LAST, fetch=5
-07)------------TableScan: t1 projection=[a, b]
-08)------TableScan: t2 projection=[x]
+04)------Sort: inner_sub.b ASC NULLS LAST, fetch=2
+05)--------SubqueryAlias: inner_sub
+06)----------SubqueryAlias: inner_alias
+07)------------Sort: t1.a ASC NULLS LAST, fetch=5
+08)--------------TableScan: t1 projection=[a, b]
+09)------TableScan: t2 projection=[x]
query III
SELECT * FROM (
@@ -985,10 +1044,11 @@ logical_plan
01)Sort: outer_sub.b ASC NULLS LAST, fetch=2
02)--SubqueryAlias: outer_sub
03)----Left Join: inner_sub.a = t2.x
-04)------SubqueryAlias: inner_sub
-05)--------SubqueryAlias: inner_alias
-06)----------TableScan: t1 projection=[a, b]
-07)------TableScan: t2 projection=[x]
+04)------Sort: inner_sub.b ASC NULLS LAST, fetch=2
+05)--------SubqueryAlias: inner_sub
+06)----------SubqueryAlias: inner_alias
+07)------------TableScan: t1 projection=[a, b]
+08)------TableScan: t2 projection=[x]
query III
SELECT * FROM (
@@ -1022,7 +1082,7 @@ logical_plan
04)------SubqueryAlias: inner_sub
05)--------Projection: inner_alias.a, inner_alias.b AS renamed_b
06)----------SubqueryAlias: inner_alias
-07)------------Sort: t1.b ASC NULLS LAST, fetch=5
+07)------------Sort: t1.b ASC NULLS LAST, fetch=2
08)--------------TableScan: t1 projection=[a, b]
09)------TableScan: t2 projection=[x]
@@ -1052,10 +1112,11 @@ logical_plan
01)Sort: sub.renamed_b ASC NULLS LAST, fetch=2
02)--SubqueryAlias: sub
03)----Left Join: t1.a = t2.x
-04)------SubqueryAlias: t1
-05)--------Projection: t1.a, t1.b AS renamed_b
-06)----------TableScan: t1 projection=[a, b]
-07)------TableScan: t2 projection=[x]
+04)------Sort: t1.renamed_b ASC NULLS LAST, fetch=2
+05)--------SubqueryAlias: t1
+06)----------Projection: t1.a, t1.b AS renamed_b
+07)------------TableScan: t1 projection=[a, b]
+08)------TableScan: t2 projection=[x]
query III
SELECT * FROM (
@@ -1083,10 +1144,11 @@ logical_plan
02)--SubqueryAlias: sub
03)----Projection: t1.rand_col, t2.x
04)------Left Join: t1.a = t2.x
-05)--------SubqueryAlias: t1
-06)----------Projection: random() AS rand_col, t1.a
-07)------------TableScan: t1 projection=[a]
-08)--------TableScan: t2 projection=[x]
+05)--------Sort: t1.rand_col ASC NULLS LAST, fetch=2
+06)----------SubqueryAlias: t1
+07)------------Projection: random() AS rand_col, t1.a
+08)--------------TableScan: t1 projection=[a]
+09)--------TableScan: t2 projection=[x]
# The outer ORDER BY column resolves to random() through the Projection, and an
# existing inner Sort is also on random() -- but they are independent random()
@@ -1108,11 +1170,229 @@ logical_plan
02)--SubqueryAlias: sub
03)----Projection: t1.rand_col, t2.x
04)------Left Join: t1.a = t2.x
-05)--------SubqueryAlias: t1
-06)----------Projection: random() AS rand_col, t1.a
-07)------------Sort: random() ASC NULLS LAST, fetch=10
-08)--------------TableScan: t1 projection=[a]
-09)--------TableScan: t2 projection=[x]
+05)--------Sort: t1.rand_col ASC NULLS LAST, fetch=2
+06)----------SubqueryAlias: t1
+07)------------Projection: random() AS rand_col, t1.a
+08)--------------Sort: random() ASC NULLS LAST, fetch=10
+09)----------------TableScan: t1 projection=[a]
+10)--------TableScan: t2 projection=[x]
+
+###
+### Degenerate cases: LIMIT 0, sort keys with no column references
+###
+
+# LIMIT 0 is folded away to an EmptyRelation by an earlier optimizer pass
+# before the TopK-through-join rewrite ever sees it.
+query TT
+EXPLAIN SELECT t1.a, t1.b, t2.x
+FROM t1 LEFT JOIN t2 ON t1.a = t2.x
+ORDER BY t1.b ASC LIMIT 0;
+----
+logical_plan EmptyRelation: rows=0
+
+query III
+SELECT t1.a, t1.b, t2.x
+FROM t1 LEFT JOIN t2 ON t1.a = t2.x
+ORDER BY t1.b ASC LIMIT 0;
+----
+
+# A sort key with no column references at all (e.g. a literal) can't
+# distinguish any row from any other on either side, so pushing a Sort onto
+# a preserved child would be legal but pointless — nothing to push down.
+# `1 + 1` is constant-folded to `2` (aliased back to its original display
+# name) by an earlier optimizer pass.
+query TT
+EXPLAIN SELECT t1.a, t1.b, t2.x
+FROM t1 LEFT JOIN t2 ON t1.a = t2.x
+ORDER BY 1 + 1 LIMIT 3;
+----
+logical_plan
+01)Sort: Int64(2) AS Int64(1) + Int64(1) ASC NULLS LAST, fetch=3
+02)--Left Join: t1.a = t2.x
+03)----TableScan: t1 projection=[a, b]
+04)----TableScan: t2 projection=[x]
+
+query III
+SELECT t1.a, t1.b, t2.x
+FROM t1 LEFT JOIN t2 ON t1.a = t2.x
+ORDER BY 1 + 1 LIMIT 3;
+----
+1 10 1
+2 20 2
+3 30 3
+
+###
+### JoinConstraint::Using
+###
+
+statement ok
+CREATE TABLE t_using1 (a INT, b INT) AS VALUES (1, 10), (2, 20), (3, 30);
+
+statement ok
+CREATE TABLE t_using2 (a INT, y INT) AS VALUES (1, 100), (2, 200);
+
+# A USING join wraps the join in a Projection (to coalesce the shared join
+# column); the rewrite still peels through that Projection to reach the Join.
+query TT
+EXPLAIN SELECT t_using1.a, t_using1.b, t_using2.y
+FROM t_using1 LEFT JOIN t_using2 USING (a)
+ORDER BY t_using1.b ASC LIMIT 2;
+----
+logical_plan
+01)Sort: t_using1.b ASC NULLS LAST, fetch=2
+02)--Projection: t_using1.a, t_using1.b, t_using2.y
+03)----Left Join: Using t_using1.a = t_using2.a
+04)------Sort: t_using1.b ASC NULLS LAST, fetch=2
+05)--------TableScan: t_using1 projection=[a, b]
+06)------TableScan: t_using2 projection=[a, y]
+
+query III
+SELECT t_using1.a, t_using1.b, t_using2.y
+FROM t_using1 LEFT JOIN t_using2 USING (a)
+ORDER BY t_using1.b ASC LIMIT 2;
+----
+1 10 100
+2 20 200
+
+statement ok
+DROP TABLE t_using1;
+
+statement ok
+DROP TABLE t_using2;
+
+###
+### Mark joins produced from EXISTS combined with OR
+###
+
+# `t1.b > 25 OR EXISTS (...)` decorrelates to a LeftMark join, but — unlike
+# a bare `EXISTS` decorrelating to LeftSemi — the OR needs the mark's boolean
+# value, so a Filter always sits directly above the LeftMark join to combine
+# it with the other predicate. That Filter is not one of the transparent
+# nodes this rewrite peels through, so pushdown is (currently) blocked for
+# every mark join reachable this way from SQL, even though the LeftMark
+# preserved-side guarantee itself holds (see the unit tests
+# `topk_pushed_to_left_of_left_mark_join` / `..._right_mark_join`).
+query TT
+EXPLAIN SELECT t1.a, t1.b
+FROM t1
+WHERE t1.b > 25 OR EXISTS (SELECT 1 FROM t2 WHERE t2.x = t1.a)
+ORDER BY t1.b ASC LIMIT 3;
+----
+logical_plan
+01)Sort: t1.b ASC NULLS LAST, fetch=3
+02)--Projection: t1.a, t1.b
+03)----Filter: t1.b > Int32(25) OR __correlated_sq_1.mark
+04)------LeftMark Join: t1.a = __correlated_sq_1.x
+05)--------TableScan: t1 projection=[a, b]
+06)--------SubqueryAlias: __correlated_sq_1
+07)----------TableScan: t2 projection=[x]
+
+query II
+SELECT t1.a, t1.b
+FROM t1
+WHERE t1.b > 25 OR EXISTS (SELECT 1 FROM t2 WHERE t2.x = t1.a)
+ORDER BY t1.b ASC LIMIT 3;
+----
+1 10
+2 20
+3 30
+
+###
+### Sort key is a projected expression, not a plain column
+###
+
+# The outer ORDER BY references a projected expression built only from
+# preserved-side columns. The rule resolves neg_b through the Projection to
+# (- t1.b) and pushes a Sort on that expression below the join.
+query TT
+EXPLAIN SELECT * FROM (
+ SELECT t1.a, -t1.b AS neg_b, t2.y
+ FROM t1 LEFT JOIN t2 ON t1.a = t2.x
+) sub
+ORDER BY neg_b ASC LIMIT 3;
+----
+logical_plan
+01)Sort: sub.neg_b ASC NULLS LAST, fetch=3
+02)--SubqueryAlias: sub
+03)----Projection: t1.a, (- t1.b) AS neg_b, t2.y
+04)------Left Join: t1.a = t2.x
+05)--------Sort: (- t1.b) ASC NULLS LAST, fetch=3
+06)----------TableScan: t1 projection=[a, b]
+07)--------TableScan: t2 projection=[x, y]
+
+query III
+SELECT * FROM (
+ SELECT t1.a, -t1.b AS neg_b, t2.y
+ FROM t1 LEFT JOIN t2 ON t1.a = t2.x
+) sub
+ORDER BY neg_b ASC LIMIT 3;
+----
+5 -50 NULL
+4 -40 NULL
+3 -30 300
+
+# Negative case: the projected expression mixes both sides. After resolving
+# through the Projection the key references t2.y, so nothing is pushed.
+query TT
+EXPLAIN SELECT * FROM (
+ SELECT t1.a, t1.b + coalesce(t2.y, 0) AS mixed
+ FROM t1 LEFT JOIN t2 ON t1.a = t2.x
+) sub
+ORDER BY mixed DESC LIMIT 3;
+----
+logical_plan
+01)Sort: sub.mixed DESC NULLS FIRST, fetch=3
+02)--SubqueryAlias: sub
+03)----Projection: t1.a, CAST(t1.b AS Int64) + CASE WHEN __common_expr_1 IS NOT NULL THEN __common_expr_1 ELSE Int64(0) END AS mixed
+04)------Projection: CAST(t2.y AS Int64) AS __common_expr_1, t1.a, t1.b
+05)--------Left Join: t1.a = t2.x
+06)----------TableScan: t1 projection=[a, b]
+07)----------TableScan: t2 projection=[x, y]
+
+query II
+SELECT * FROM (
+ SELECT t1.a, t1.b + coalesce(t2.y, 0) AS mixed
+ FROM t1 LEFT JOIN t2 ON t1.a = t2.x
+) sub
+ORDER BY mixed DESC LIMIT 3;
+----
+3 330
+2 220
+1 110
+
+# Two stacked Projections: the key is computed in the inner one and re-aliased
+# in the outer one. Resolution must substitute through both layers.
+query TT
+EXPLAIN SELECT a, k2 FROM (
+ SELECT a, k1 AS k2 FROM (
+ SELECT t1.a, t1.b * 2 AS k1, t2.y
+ FROM t1 LEFT JOIN t2 ON t1.a = t2.x
+ ) s1
+) s2
+ORDER BY k2 DESC LIMIT 2;
+----
+logical_plan
+01)Sort: s2.k2 DESC NULLS FIRST, fetch=2
+02)--SubqueryAlias: s2
+03)----Projection: s1.a, s1.k1 AS k2
+04)------SubqueryAlias: s1
+05)--------Projection: t1.a, CAST(t1.b AS Int64) * Int64(2) AS k1
+06)----------Left Join: t1.a = t2.x
+07)------------Sort: CAST(t1.b AS Int64) * Int64(2) DESC NULLS FIRST, fetch=2
+08)--------------TableScan: t1 projection=[a, b]
+09)------------TableScan: t2 projection=[x]
+
+query II
+SELECT a, k2 FROM (
+ SELECT a, k1 AS k2 FROM (
+ SELECT t1.a, t1.b * 2 AS k1, t2.y
+ FROM t1 LEFT JOIN t2 ON t1.a = t2.x
+ ) s1
+) s2
+ORDER BY k2 DESC LIMIT 2;
+----
+5 100
+4 80
statement ok
set datafusion.execution.target_partitions = 4;
@@ -1124,4 +1404,4 @@ statement ok
DROP TABLE t1;
statement ok
-DROP TABLE t2;
+DROP TABLE t2;
\ No newline at end of file
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org