Skip to content

Commit

Permalink
Rewrite Optimizer to use TreeNode API
Browse files Browse the repository at this point in the history
  • Loading branch information
alamb committed Apr 4, 2024
1 parent 10448fb commit a755338
Show file tree
Hide file tree
Showing 26 changed files with 532 additions and 561 deletions.
2 changes: 1 addition & 1 deletion datafusion-examples/examples/rewrite_expr.rs
Original file line number Diff line number Diff line change
Expand Up @@ -59,7 +59,7 @@ pub fn main() -> Result<()> {

// then run the optimizer with our custom rule
let optimizer = Optimizer::with_rules(vec![Arc::new(MyOptimizerRule {})]);
let optimized_plan = optimizer.optimize(&analyzed_plan, &config, observe)?;
let optimized_plan = optimizer.optimize(analyzed_plan, &config, observe)?;
println!(
"Optimized Logical Plan:\n\n{}\n",
optimized_plan.display_indent()
Expand Down
17 changes: 15 additions & 2 deletions datafusion/common/src/tree_node.rs
Original file line number Diff line number Diff line change
Expand Up @@ -100,6 +100,9 @@ pub trait TreeNode: Sized {
/// Visit the tree node using the given [`TreeNodeVisitor`], performing a
/// depth-first walk of the node and its children.
///
/// See also:
/// * [`Self::rewrite`] to rewrite owned `TreeNode`s
///
/// Consider the following tree structure:
/// ```text
/// ParentNode
Expand Down Expand Up @@ -144,6 +147,9 @@ pub trait TreeNode: Sized {
/// Implements the [visitor pattern](https://en.wikipedia.org/wiki/Visitor_pattern) for
/// recursively transforming [`TreeNode`]s.
///
/// See also:
/// * [`Self::visit`] for inspecting (without modification) `TreeNode`s
///
/// Consider the following tree structure:
/// ```text
/// ParentNode
Expand Down Expand Up @@ -353,13 +359,15 @@ pub trait TreeNode: Sized {
}

/// Apply the closure `F` to the node's children.
///
/// See `mutate_children` for rewriting in place
fn apply_children<F: FnMut(&Self) -> Result<TreeNodeRecursion>>(
&self,
f: &mut F,
) -> Result<TreeNodeRecursion>;

/// Apply transform `F` to the node's children. Note that the transform `F`
/// might have a direction (pre-order or post-order).
/// Apply transform `F` to potentially rewrite the node's children. Note
/// that the transform `F` might have a direction (pre-order or post-order).
fn map_children<F: FnMut(Self) -> Result<Transformed<Self>>>(
self,
f: F,
Expand Down Expand Up @@ -489,6 +497,11 @@ impl<T> Transformed<T> {
f(self.data).map(|data| Transformed::new(data, self.transformed, self.tnr))
}

/// Invokes f(), depending on the value of self.tnr.
///
/// This is used to conditionally apply a function during a f_up tree
/// traversal, if the result of children traversal was `[`TreeNodeRecursion::Continue`].
///
/// Handling [`TreeNodeRecursion::Continue`] and [`TreeNodeRecursion::Stop`]
/// is straightforward, but [`TreeNodeRecursion::Jump`] can behave differently
/// when we are traversing down or up on a tree. If [`TreeNodeRecursion`] of
Expand Down
4 changes: 2 additions & 2 deletions datafusion/core/src/execution/context/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1892,7 +1892,7 @@ impl SessionState {

// optimize the child plan, capturing the output of each optimizer
let optimized_plan = self.optimizer.optimize(
&analyzed_plan,
analyzed_plan,
self,
|optimized_plan, optimizer| {
let optimizer_name = optimizer.name().to_string();
Expand Down Expand Up @@ -1922,7 +1922,7 @@ impl SessionState {
let analyzed_plan =
self.analyzer
.execute_and_check(plan, self.options(), |_, _| {})?;
self.optimizer.optimize(&analyzed_plan, self, |_, _| {})
self.optimizer.optimize(analyzed_plan, self, |_, _| {})
}
}

Expand Down
2 changes: 1 addition & 1 deletion datafusion/core/tests/optimizer_integration.rs
Original file line number Diff line number Diff line change
Expand Up @@ -110,7 +110,7 @@ fn test_sql(sql: &str) -> Result<LogicalPlan> {
let optimizer = Optimizer::new();
// analyze and optimize the logical plan
let plan = analyzer.execute_and_check(&plan, config.options(), |_, _| {})?;
optimizer.optimize(&plan, &config, |_, _| {})
optimizer.optimize(plan, &config, |_, _| {})
}

#[derive(Default)]
Expand Down
Loading

0 comments on commit a755338

Please sign in to comment.