kumarUjjawal commented on code in PR #21621:
URL: https://github.com/apache/datafusion/pull/21621#discussion_r3208553405
##########
datafusion/optimizer/src/optimizer.rs:
##########
@@ -296,6 +297,7 @@ impl Optimizer {
Arc::new(EliminateOuterJoin::new()),
// Filters can't be pushed down past Limits, we should do
PushDownFilter after PushDownLimit
Arc::new(PushDownLimit::new()),
+ Arc::new(PushDownTopKThroughJoin::new()),
Review Comment:
this runs before PushDownFilter. For Sort(fetch) -> Filter -> Join, it only
fires on a later optimizer pass after the filter has moved. That works with the
default `max_passes = 3`, but with `max_passes = 1` the pushdown is missed.
Should we run this after PushDownFilter?
##########
datafusion/optimizer/src/push_down_topk_through_join.rs:
##########
@@ -0,0 +1,1070 @@
+// 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.
+
+//! [`PushDownTopKThroughJoin`] pushes TopK (Sort with fetch) through outer
joins
+//!
+//! When a `Sort` with a fetch limit sits above an outer join and all sort
+//! expressions come from the **preserved** side, this rule inserts a copy
+//! of the `Sort(fetch)` on that input to reduce the number of rows
+//! entering the join.
+//!
+//! This is correct because:
+//! - A LEFT JOIN preserves every left row (each appears at least once in the
+//! output). The final top-N by left-side columns must come from the top-N
+//! left rows.
+//! - The same reasoning applies symmetrically for RIGHT JOIN and right-side
+//! columns.
+//!
+//! The top-level sort is kept for correctness since a 1-to-many join can
+//! produce more than N output rows from N input rows.
+//!
+//! ## Example
+//!
+//! Before:
+//! ```text
+//! Sort: t1.b ASC, fetch=3
+//! Left Join: t1.a = t2.c
+//! Scan: t1 ← scans ALL rows
+//! Scan: t2
+//! ```
+//!
+//! After:
+//! ```text
+//! Sort: t1.b ASC, fetch=3
+//! Left Join: t1.a = t2.c
+//! Sort: t1.b ASC, fetch=3 ← pushed down
+//! Scan: t1
+//! Scan: t2
+//! ```
+
+use std::sync::Arc;
+
+use crate::optimizer::ApplyOrder;
+use crate::{OptimizerConfig, OptimizerRule};
+
+use crate::utils::{has_all_column_refs, schema_columns};
+use datafusion_common::tree_node::{Transformed, TreeNode};
+use datafusion_common::{Column, Result};
+use datafusion_expr::logical_plan::{
+ JoinType, LogicalPlan, Projection, Sort as SortPlan, SubqueryAlias,
+};
+use datafusion_expr::{Expr, SortExpr};
+
+/// Optimization rule that pushes TopK (Sort with fetch) through
+/// LEFT / RIGHT outer joins when all sort expressions come from
+/// the preserved side.
+///
+/// See module-level documentation for details.
+#[derive(Default, Debug)]
+pub struct PushDownTopKThroughJoin;
+
+impl PushDownTopKThroughJoin {
+ #[expect(missing_docs)]
+ pub fn new() -> Self {
+ Self {}
+ }
+}
+
+impl OptimizerRule for PushDownTopKThroughJoin {
+ fn supports_rewrite(&self) -> bool {
+ true
+ }
+
+ fn rewrite(
+ &self,
+ plan: LogicalPlan,
+ _config: &dyn OptimizerConfig,
+ ) -> Result<Transformed<LogicalPlan>> {
+ // Match Sort with fetch (TopK)
+ let LogicalPlan::Sort(sort) = &plan else {
+ return Ok(Transformed::no(plan));
+ };
+ let Some(fetch) = sort.fetch else {
+ return Ok(Transformed::no(plan));
+ };
+
+ // 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 the result.
+ if sort.expr.iter().any(|se| se.expr.is_volatile()) {
+ return Ok(Transformed::no(plan));
+ }
+
+ // Peel through transparent nodes (SubqueryAlias, Projection) to find
+ // the Join. Track intermediate nodes so we can reconstruct the tree
+ // and resolve sort expressions through them.
+ let mut current = sort.input.as_ref();
+ let mut intermediates: Vec<&LogicalPlan> = Vec::new();
+ let join = loop {
+ match current {
+ LogicalPlan::Join(join) => break join,
+ LogicalPlan::Projection(proj) => {
+ intermediates.push(current);
+ current = proj.input.as_ref();
+ }
+ LogicalPlan::SubqueryAlias(sq) => {
+ intermediates.push(current);
+ current = sq.input.as_ref();
+ }
+ _ => return Ok(Transformed::no(plan)),
+ }
+ };
+
+ // Only outer joins where the preserved side is known.
+ // Semi/Anti joins are excluded: not all preserved-side rows appear in
+ // the output (only matched/unmatched rows do), so pushing fetch=N to
+ // the preserved child can drop rows that would have survived the
filter.
+ //
+ // Non-equijoin filters in the ON clause are safe: outer joins
guarantee
+ // all preserved-side rows appear in the output regardless of the
filter.
+ // The filter only controls matching (which non-preserved rows pair
up),
+ // not which preserved rows survive.
+ let preserved_is_left = match join.join_type {
+ JoinType::Left => true,
+ JoinType::Right => false,
+ _ => return Ok(Transformed::no(plan)),
+ };
+
+ // Resolve sort expressions through all intermediate nodes (Projection,
+ // SubqueryAlias) so that column references match the join's schema.
+ let mut resolved_sort_exprs = sort.expr.clone();
+ for node in &intermediates {
+ match node {
+ LogicalPlan::Projection(proj) => {
+ resolved_sort_exprs =
resolve_sort_exprs_through_projection(
+ &resolved_sort_exprs,
+ proj,
+ )?;
+ }
+ LogicalPlan::SubqueryAlias(sq) => {
+ resolved_sort_exprs =
resolve_sort_exprs_through_subquery_alias(
+ &resolved_sort_exprs,
+ sq,
+ )?;
+ }
+ _ => unreachable!(),
+ }
+ }
+
+ // After resolving through projections, the sort expressions may now
+ // contain volatile functions (e.g. `random() AS col`). Duplicating
+ // volatile expressions in the pushed Sort would produce different
+ // values, changing results.
+ if resolved_sort_exprs.iter().any(|se| se.expr.is_volatile()) {
+ return Ok(Transformed::no(plan));
+ }
+
+ let preserved_schema = if preserved_is_left {
+ join.left.schema()
+ } else {
+ join.right.schema()
+ };
+ let preserved_cols = schema_columns(preserved_schema);
+
+ let all_from_preserved = resolved_sort_exprs
+ .iter()
+ .all(|sort_expr| has_all_column_refs(&sort_expr.expr,
&preserved_cols));
+ if !all_from_preserved {
+ return Ok(Transformed::no(plan));
+ }
+
+ let preserved_child = if preserved_is_left {
+ &join.left
+ } else {
+ &join.right
+ };
+
+ // Resolve sort exprs further through any SubqueryAlias wrapping the
+ // preserved child, so we can compare with the inner Sort's
expressions.
+ //
+ // After intermediate resolution, resolved_sort_exprs = [t1.b ASC].
+ // The inner Sort uses [orders.b ASC]. This step maps t1.b → orders.b.
+ //
+ // ```text
+ // Sort(sub.b ASC, fetch=2)
+ // SubqueryAlias(sub) ← intermediate, already resolved
+ // Left Join
+ // SubqueryAlias(t1) ← preserved child, resolve here
+ // Sort(orders.b ASC, fetch=5)
+ // TableScan: orders
+ // ```
+ let (inner_child, child_resolved_exprs) = match
preserved_child.as_ref() {
+ LogicalPlan::SubqueryAlias(sq) => {
+ let exprs =
+
resolve_sort_exprs_through_subquery_alias(&resolved_sort_exprs, sq)?;
+ (sq.input.as_ref(), exprs)
+ }
+ _ => (preserved_child.as_ref(), resolved_sort_exprs.clone()),
+ };
Review Comment:
This only looks through one SubqueryAlias on the preserved side. If there
are nested aliases, we may miss an existing inner `Sort(fetch)` and add another
sort instead of tightening it. We should handle nested aliases.
##########
datafusion/sqllogictest/test_files/push_down_topk_through_join.slt:
##########
Review Comment:
Could we add a few more coverage cases? Like DESC ordering, NULLS FIRST, a
multi-level outer join, a CROSS JOIN negative case, and maybe a tied-sort-key
case. A custom volatile UDF case would also be useful if it is easy to add.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
