milenkovicm commented on code in PR #2793: URL: https://github.com/apache/datafusion-ballista/pull/2793#discussion_r4518381902
########## ballista/core/src/optimizer/semi_join_pushdown.rs: ########## @@ -0,0 +1,639 @@ +// 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. + +//! Pushes semi and anti joins below the joins and projections that sit between +//! them and the input their keys come from. +//! +//! DataFusion decorrelates `IN`, `EXISTS` and `NOT EXISTS` subqueries into +//! semi and anti joins placed directly above the `FROM` clause's join tree, so +//! the subquery filters the joined result. In TPC-H q18 the subquery keeps +//! 6,400 orders, but only after `orders` has been joined to `lineitem`: +//! +//! ```text +//! LeftSemi Join: orders.o_orderkey = __correlated_sq_1.l_orderkey +//! Inner Join: orders.o_orderkey = lineitem.l_orderkey +//! TableScan: orders +//! TableScan: lineitem +//! SubqueryAlias: __correlated_sq_1 +//! ``` +//! +//! The semi join only reads `orders.o_orderkey`, so this rule moves it onto +//! `orders`, before the join: +//! +//! ```text +//! Inner Join: orders.o_orderkey = lineitem.l_orderkey +//! LeftSemi Join: orders.o_orderkey = __correlated_sq_1.l_orderkey +//! TableScan: orders +//! SubqueryAlias: __correlated_sq_1 +//! TableScan: lineitem +//! ``` +//! +//! This is Spark's `PushDownLeftSemiAntiJoin`. A semi or anti join keeps or +//! drops each row of its left input based only on that row's values, so it +//! commutes with any operator that passes the referenced columns through +//! unchanged and never makes a row depend on another row: +//! +//! - a `Projection` that outputs every referenced column as a plain column +//! and has no volatile expression; +//! - an inner join, onto the side that holds every referenced column, unless +//! the other side is filtered; +//! - a left (right) outer join, onto its left (right), preserved, side only. +//! Pushing into the null-supplying side would turn a dropped row into a +//! null-extended one. +//! +//! It deliberately stops at two places where the move is sound but slower. A +//! `Filter` is the cheaper of the two operators, so the semi join should see +//! its output (TPC-H q4). An inner join whose other input is filtered can +//! remove most rows itself, so a semi join moved below it probes rows the join +//! would have dropped (TPC-H q21's `EXISTS` and `NOT EXISTS`; Spark makes this +//! move unconditionally). `PushDownFilter` would also move a filter back above +//! the semi join on the next pass. +//! +//! The rule runs after DataFusion's own rules in each optimizer pass, so it +//! sees decorrelated joins with their filters already pushed down. + +use std::collections::HashSet; +use std::sync::Arc; + +use datafusion::common::tree_node::{Transformed, TreeNode}; +use datafusion::common::{Column, DFSchema, Result}; +use datafusion::logical_expr::{Expr, Join, JoinType, LogicalPlan, Projection}; +use datafusion::optimizer::optimizer::ApplyOrder; +use datafusion::optimizer::{OptimizerConfig, OptimizerRule}; + +use crate::config::BallistaConfig; + +/// Pushes semi and anti joins below inner joins and projections, +/// onto the join input that holds every column they reference. See the +/// module documentation for the rewrite and why it is sound. +/// +/// Enabled by default and controlled by the +/// `ballista.optimizer.semi_join_pushdown` configuration key. +#[derive(Debug, Default)] +pub struct SemiJoinPushdown {} + +impl SemiJoinPushdown { + /// Creates a new instance of the rule. + pub fn new() -> Self { + Self {} + } +} + +impl OptimizerRule for SemiJoinPushdown { + fn name(&self) -> &str { + "semi_join_pushdown" + } + + fn apply_order(&self) -> Option<ApplyOrder> { + // Top-down: a semi join pushed one level lands in a child, which the + // optimizer visits next and can push further. + Some(ApplyOrder::TopDown) + } + + fn supports_rewrite(&self) -> bool { + true + } + + fn rewrite( + &self, + plan: LogicalPlan, + config: &dyn OptimizerConfig, + ) -> Result<Transformed<LogicalPlan>> { + // Cheap node checks first: this runs on every node of every pass. + let LogicalPlan::Join(join) = plan else { + return Ok(Transformed::no(plan)); + }; + if !is_pushable_semi_join(&join) { + return Ok(Transformed::no(LogicalPlan::Join(join))); + } + let enabled = config + .options() + .extensions + .get::<BallistaConfig>() + .map(|c| c.semi_join_pushdown_enabled()) + .unwrap_or_else(|| BallistaConfig::default().semi_join_pushdown_enabled()); + if !enabled { + return Ok(Transformed::no(LogicalPlan::Join(join))); + } + let referenced = left_referenced_columns(&join); + match push_one_level(&join, &referenced)? { + Some(plan) => Ok(Transformed::yes(plan)), + None => Ok(Transformed::no(LogicalPlan::Join(join))), + } + } +} + +/// A left semi or anti join whose keep/drop decision depends only on the +/// left row, over an operator this rule can move it below. Null-aware anti +/// joins (`NOT IN`) also depend only on the left row and the whole subquery, +/// but they are left alone: Ballista executes them in a single task and the +/// `NotInSubqueryRewrite` rule handles them. +fn is_pushable_semi_join(join: &Join) -> bool { + matches!(join.join_type, JoinType::LeftSemi | JoinType::LeftAnti) + && !join.null_aware + && matches!( + join.left.as_ref(), + LogicalPlan::Projection(_) | LogicalPlan::Join(_) + ) +} + +/// Columns of the semi join's left input that its keys and filter read. +fn left_referenced_columns(join: &Join) -> HashSet<Column> { + let left_schema = join.left.schema(); + let mut cols: HashSet<Column> = HashSet::new(); + for (l, _) in &join.on { + cols.extend(l.column_refs().into_iter().cloned()); + } + if let Some(filter) = &join.filter { + cols.extend( + filter + .column_refs() + .into_iter() + .filter(|c| left_schema.has_column(c)) + .cloned(), + ); + } + cols +} + +fn has_all(schema: &DFSchema, cols: &HashSet<Column>) -> bool { + cols.iter().all(|c| schema.has_column(c)) +} + +/// Whether `plan` filters any of its inputs, so a join with it may remove +/// rows: a `Filter` anywhere in it, a scan with pushed-down filters, or a +/// semi or anti join. Logical plans carry no statistics, so this syntactic test +/// stands in for "the other input is selective". +fn is_filtered(plan: &LogicalPlan) -> Result<bool> { + plan.exists(|node| { + Ok(match node { + LogicalPlan::Filter(_) => true, + LogicalPlan::TableScan(scan) => !scan.filters.is_empty(), + LogicalPlan::Join(j) => matches!( + j.join_type, + JoinType::LeftSemi + | JoinType::LeftAnti + | JoinType::RightSemi + | JoinType::RightAnti + ), + _ => false, + }) + }) +} + +/// Rebuilds `join` with new inputs. +fn with_inputs( + join: &Join, + left: Arc<LogicalPlan>, + right: Arc<LogicalPlan>, +) -> Result<LogicalPlan> { + Ok(LogicalPlan::Join(Join::try_new( + left, + right, + join.on.clone(), + join.filter.clone(), + join.join_type, + join.join_constraint, + join.null_equality, + join.null_aware, + )?)) +} + +/// Rebuilds `semi` with a new left input. +fn semi_over(semi: &Join, left: Arc<LogicalPlan>) -> Result<LogicalPlan> { + with_inputs(semi, left, Arc::clone(&semi.right)) +} + +/// Moves `semi` one operator down its left input, or returns `None` when the +/// operator below it does not commute with it. +fn push_one_level( + semi: &Join, + referenced: &HashSet<Column>, +) -> Result<Option<LogicalPlan>> { + if referenced.is_empty() { + // An uncorrelated EXISTS keeps all or none of the rows; leave it. + return Ok(None); + } + match semi.left.as_ref() { Review Comment: should semi.filter expression be checked for volatile expressions? -- 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]
