Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Introduce TreeNode::exists() API, avoid copying expressions #10008

Merged
merged 5 commits into from
Apr 9, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
17 changes: 17 additions & 0 deletions datafusion/common/src/tree_node.rs
Original file line number Diff line number Diff line change
Expand Up @@ -292,6 +292,23 @@ pub trait TreeNode: Sized {
)
}

/// Returns true if `f` returns true for node in the tree.
///
/// Stops recursion as soon as a matching node is found
fn exists<F: FnMut(&Self) -> bool>(&self, mut f: F) -> bool {
alamb marked this conversation as resolved.
Show resolved Hide resolved
let mut found = false;
self.apply(&mut |n| {
Ok(if f(n) {
found = true;
TreeNodeRecursion::Stop
} else {
TreeNodeRecursion::Continue
})
})
.unwrap();
found
}

/// Apply the closure `F` to the node's children.
fn apply_children<F: FnMut(&Self) -> Result<TreeNodeRecursion>>(
&self,
Expand Down
4 changes: 2 additions & 2 deletions datafusion/expr/src/expr.rs
Original file line number Diff line number Diff line change
Expand Up @@ -25,7 +25,7 @@ use std::sync::Arc;

use crate::expr_fn::binary_expr;
use crate::logical_plan::Subquery;
use crate::utils::{expr_to_columns, find_out_reference_exprs};
use crate::utils::expr_to_columns;
use crate::window_frame;
use crate::{
aggregate_function, built_in_function, built_in_window_function, udaf,
Expand Down Expand Up @@ -1232,7 +1232,7 @@ impl Expr {

/// Return true when the expression contains out reference(correlated) expressions.
pub fn contains_outer(&self) -> bool {
!find_out_reference_exprs(self).is_empty()
self.exists(|expr| matches!(expr, Expr::OuterReferenceColumn { .. }))
}

/// Recursively find all [`Expr::Placeholder`] expressions, and
Expand Down
15 changes: 15 additions & 0 deletions datafusion/expr/src/logical_plan/plan.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1321,6 +1321,21 @@ impl LogicalPlan {
| LogicalPlan::Extension(_) => None,
}
}

/// If this node's expressions contains any references to an outer subquery
pub fn contains_outer_reference(&self) -> bool {
alamb marked this conversation as resolved.
Show resolved Hide resolved
let mut contains = false;
self.apply_expressions(|expr| {
Ok(if expr.contains_outer() {
contains = true;
TreeNodeRecursion::Stop
} else {
TreeNodeRecursion::Continue
})
})
.unwrap();
contains
}
}

/// This macro is used to determine continuation during combined transforming
Expand Down
9 changes: 1 addition & 8 deletions datafusion/optimizer/src/analyzer/subquery.rs
Original file line number Diff line number Diff line change
Expand Up @@ -140,7 +140,7 @@ fn check_inner_plan(
is_aggregate: bool,
can_contain_outer_ref: bool,
) -> Result<()> {
if !can_contain_outer_ref && contains_outer_reference(inner_plan) {
if !can_contain_outer_ref && inner_plan.contains_outer_reference() {
return plan_err!("Accessing outer reference columns is not allowed in the plan");
}
// We want to support as many operators as possible inside the correlated subquery
Expand Down Expand Up @@ -233,13 +233,6 @@ fn check_inner_plan(
}
}

fn contains_outer_reference(inner_plan: &LogicalPlan) -> bool {
inner_plan
.expressions()
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

❤️ for removing this copy 🚀

.iter()
.any(|expr| expr.contains_outer())
}

fn check_aggregation_in_scalar_subquery(
inner_plan: &LogicalPlan,
agg: &Aggregate,
Expand Down
2 changes: 1 addition & 1 deletion datafusion/optimizer/src/decorrelate.rs
Original file line number Diff line number Diff line change
Expand Up @@ -91,7 +91,7 @@ impl TreeNodeRewriter for PullUpCorrelatedExpr {
_ => Ok(Transformed::no(plan)),
}
}
_ if plan.expressions().iter().any(|expr| expr.contains_outer()) => {
_ if plan.contains_outer_reference() => {
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

And another copy gone!

// the unsupported cases, the plan expressions contain out reference columns(like window expressions)
self.can_pull_up = false;
Ok(Transformed::new(plan, false, TreeNodeRecursion::Jump))
Expand Down
Loading