andygrove commented on code in PR #2793: URL: https://github.com/apache/datafusion-ballista/pull/2793#discussion_r4519627320
########## 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: Yes, good call. It now leaves the semi join in place when any key or the filter is volatile. Otherwise something like `EXISTS (... WHERE s.x = a.v AND random() < 0.5)` would draw one number per `a` row instead of one per joined row. Added `keeps_semi_join_with_a_volatile_filter`. ########## ballista/core/src/optimizer.rs: ########## @@ -331,6 +350,48 @@ mod tests { .unwrap() } + /// TPC-H q18's shape through the rules every Ballista session installs: + /// the `IN` subquery filters `orders` before it is joined to `lineitem`. + #[tokio::test] + async fn session_rules_push_an_in_subquery_below_the_join() { + let state = SessionStateBuilder::new() + .with_default_features() + .with_config( + SessionConfig::new().with_option_extension(BallistaConfig::default()), + ) + .with_optimizer_rules(super::ballista_default_optimizer_rules()) + .build(); + let ctx = SessionContext::new_with_state(state); + for ddl in [ + "CREATE TABLE orders (o_orderkey INT, o_custkey INT)", + "CREATE TABLE lineitem (l_orderkey INT, l_quantity INT)", + ] { + ctx.sql(ddl).await.unwrap(); + } + let plan = ctx + .sql( + "SELECT o_custkey, l_quantity FROM orders \ + JOIN lineitem ON o_orderkey = l_orderkey \ + WHERE o_orderkey IN (SELECT l_orderkey FROM lineitem \ + GROUP BY l_orderkey HAVING sum(l_quantity) > 300)", + ) + .await + .unwrap() + .into_optimized_plan() + .unwrap(); + let display = format!("{}", plan.display_indent()); + let line_of = |needle: &str| { + display + .lines() + .position(|l| l.contains(needle)) + .unwrap_or_else(|| panic!("{needle:?} not found in:\n{display}")) + }; + assert!( Review Comment: Done, the session test now compares the whole optimized plan. ########## 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() { + LogicalPlan::Projection(projection) => { + // A volatile expression must be evaluated on the rows the query + // asked for, not on the fewer that survive the semi join. + if projection.expr.iter().any(|e| e.is_volatile()) { + return Ok(None); + } + // Every referenced column must be passed through unchanged, so + // the name the semi join uses means the same value below. + let passthrough: HashSet<&Column> = projection + .expr + .iter() + .filter_map(|e| match e { + Expr::Column(c) => Some(c), + _ => None, + }) + .collect(); + if !referenced.iter().all(|c| passthrough.contains(c)) + || !has_all(projection.input.schema(), referenced) + { + return Ok(None); + } + let pushed = semi_over(semi, Arc::clone(&projection.input))?; + Ok(Some(LogicalPlan::Projection(Projection::try_new( + projection.expr.clone(), + Arc::new(pushed), + )?))) + } + LogicalPlan::Join(inner) => { + if inner.null_aware { + return Ok(None); + } + let (to_left, to_right) = match inner.join_type { + JoinType::Inner => (true, true), + JoinType::Left => (true, false), + JoinType::Right => (false, true), + _ => (false, false), + }; + // Moving below an inner join whose other input is filtered would + // make the semi join see rows that join could have removed first. + // An outer join keeps every row of its preserved side, so there + // is nothing to wait for. + let inner_join = inner.join_type == JoinType::Inner; + let (new_left, new_right) = if to_left + && has_all(inner.left.schema(), referenced) + && !(inner_join && is_filtered(&inner.right)?) + { + ( + Arc::new(semi_over(semi, Arc::clone(&inner.left))?), + Arc::clone(&inner.right), + ) + } else if to_right + && has_all(inner.right.schema(), referenced) + && !(inner_join && is_filtered(&inner.left)?) + { + ( + Arc::clone(&inner.left), + Arc::new(semi_over(semi, Arc::clone(&inner.right))?), + ) + } else { + return Ok(None); + }; + with_inputs(inner, new_left, new_right).map(Some) + } + _ => Ok(None), + } +} + +#[cfg(test)] +mod tests { Review Comment: Added `semi_join_pushdown_preserves_results`. It runs `IN`, `EXISTS` and `NOT EXISTS` over inner and left joins, with NULL keys on both the outer and subquery sides, with the rule on and off. It compares the rows and checks where the semi join ended up. I tested it by letting the rule push into a left join's null-supplying side, and the test catches the extra null-extended row. -- 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]
