gene-bordegaray commented on code in PR #21621: URL: https://github.com/apache/datafusion/pull/21621#discussion_r3092738430
########## datafusion/optimizer/src/push_down_topk_through_join.rs: ########## @@ -0,0 +1,405 @@ +// 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::Result; +use datafusion_common::tree_node::Transformed; +use datafusion_expr::logical_plan::{JoinType, LogicalPlan, Sort as SortPlan}; + +/// 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)); + }; + + // Check if the child is a Join (look through Projection) + let (has_projection, join) = match sort.input.as_ref() { + LogicalPlan::Join(join) => (false, join), + LogicalPlan::Projection(proj) => match proj.input.as_ref() { + LogicalPlan::Join(join) => (true, join), + _ => return Ok(Transformed::no(plan)), + }, + _ => return Ok(Transformed::no(plan)), + }; + + // Only LEFT or RIGHT, no non-equijoin filter + let preserved_is_left = match join.join_type { + JoinType::Left => true, + JoinType::Right => false, + _ => return Ok(Transformed::no(plan)), + }; + if join.filter.is_some() { + return Ok(Transformed::no(plan)); + } + + // Check all sort expression columns come from the preserved side + 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 = sort + .expr + .iter() + .all(|sort_expr| has_all_column_refs(&sort_expr.expr, &preserved_cols)); + if !all_from_preserved { + return Ok(Transformed::no(plan)); + } + + // Don't push if preserved child is already a Sort (redundant) + let preserved_child = if preserved_is_left { + &join.left + } else { + &join.right + }; + if matches!(preserved_child.as_ref(), LogicalPlan::Sort(_)) { + return Ok(Transformed::no(plan)); + } + + // Create the new Sort(fetch) on the preserved child Review Comment: Also about this `clone`. I dont think we can just do this in all cases. Say that the `sort.expr` is non deterministic function like `random()`. The cloned `random()` that is pushed down is not necessarily then same as the one on top. I don't know if this will matter in practice because the top one may be eliminated but something to think about. ########## datafusion/optimizer/src/push_down_topk_through_join.rs: ########## @@ -0,0 +1,405 @@ +// 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::Result; +use datafusion_common::tree_node::Transformed; +use datafusion_expr::logical_plan::{JoinType, LogicalPlan, Sort as SortPlan}; + +/// 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)); + }; + + // Check if the child is a Join (look through Projection) + let (has_projection, join) = match sort.input.as_ref() { + LogicalPlan::Join(join) => (false, join), + LogicalPlan::Projection(proj) => match proj.input.as_ref() { + LogicalPlan::Join(join) => (true, join), + _ => return Ok(Transformed::no(plan)), + }, + _ => return Ok(Transformed::no(plan)), + }; + + // Only LEFT or RIGHT, no non-equijoin filter + let preserved_is_left = match join.join_type { + JoinType::Left => true, + JoinType::Right => false, + _ => return Ok(Transformed::no(plan)), + }; + if join.filter.is_some() { + return Ok(Transformed::no(plan)); + } + + // Check all sort expression columns come from the preserved side + 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 = sort + .expr + .iter() + .all(|sort_expr| has_all_column_refs(&sort_expr.expr, &preserved_cols)); + if !all_from_preserved { + return Ok(Transformed::no(plan)); + } + + // Don't push if preserved child is already a Sort (redundant) + let preserved_child = if preserved_is_left { + &join.left + } else { + &join.right + }; + if matches!(preserved_child.as_ref(), LogicalPlan::Sort(_)) { + return Ok(Transformed::no(plan)); + } + + // Create the new Sort(fetch) on the preserved child + let new_child_sort = Arc::new(LogicalPlan::Sort(SortPlan { + expr: sort.expr.clone(), + input: Arc::clone(preserved_child), + fetch: Some(fetch), + })); + Review Comment: I don't think we shoudl just clone `sort.expr` here. Since the `Sort` can sit on top of a `Projection`, in this case the `ORDER BY` clause of the query is interpreted against the projection output columns, not directly on the join's child. When we push down the `Sort(fetch)` rather than cloning the `Sort` columns we need to push down the columns that were projected. I believe behavior for this right now would work like this: ```text Sort: b, fetch=1 Projection: -t1.b AS b Join The optimizer rewrites it into: Sort: b, fetch=1 Projection: -t1.b AS b Join Sort: b, fetch=1 -> This is using the post-projected value! ``` -- 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]
