Skip to content

Commit d951289

Browse files
committed
Rewrite Optimizer to use TreeNode API
1 parent b4a9ffd commit d951289

26 files changed

+532
-561
lines changed

datafusion-examples/examples/rewrite_expr.rs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -59,7 +59,7 @@ pub fn main() -> Result<()> {
5959

6060
// then run the optimizer with our custom rule
6161
let optimizer = Optimizer::with_rules(vec![Arc::new(MyOptimizerRule {})]);
62-
let optimized_plan = optimizer.optimize(&analyzed_plan, &config, observe)?;
62+
let optimized_plan = optimizer.optimize(analyzed_plan, &config, observe)?;
6363
println!(
6464
"Optimized Logical Plan:\n\n{}\n",
6565
optimized_plan.display_indent()

datafusion/common/src/tree_node.rs

Lines changed: 15 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -100,6 +100,9 @@ pub trait TreeNode: Sized {
100100
/// Visit the tree node using the given [`TreeNodeVisitor`], performing a
101101
/// depth-first walk of the node and its children.
102102
///
103+
/// See also:
104+
/// * [`Self::rewrite`] to rewrite owned `TreeNode`s
105+
///
103106
/// Consider the following tree structure:
104107
/// ```text
105108
/// ParentNode
@@ -144,6 +147,9 @@ pub trait TreeNode: Sized {
144147
/// Implements the [visitor pattern](https://en.wikipedia.org/wiki/Visitor_pattern) for
145148
/// recursively transforming [`TreeNode`]s.
146149
///
150+
/// See also:
151+
/// * [`Self::visit`] for inspecting (without modification) `TreeNode`s
152+
///
147153
/// Consider the following tree structure:
148154
/// ```text
149155
/// ParentNode
@@ -353,13 +359,15 @@ pub trait TreeNode: Sized {
353359
}
354360

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

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

500+
/// Invokes f(), depending on the value of self.tnr.
501+
///
502+
/// This is used to conditionally apply a function during a f_up tree
503+
/// traversal, if the result of children traversal was `[`TreeNodeRecursion::Continue`].
504+
///
492505
/// Handling [`TreeNodeRecursion::Continue`] and [`TreeNodeRecursion::Stop`]
493506
/// is straightforward, but [`TreeNodeRecursion::Jump`] can behave differently
494507
/// when we are traversing down or up on a tree. If [`TreeNodeRecursion`] of

datafusion/core/src/execution/context/mod.rs

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1892,7 +1892,7 @@ impl SessionState {
18921892

18931893
// optimize the child plan, capturing the output of each optimizer
18941894
let optimized_plan = self.optimizer.optimize(
1895-
&analyzed_plan,
1895+
analyzed_plan,
18961896
self,
18971897
|optimized_plan, optimizer| {
18981898
let optimizer_name = optimizer.name().to_string();
@@ -1922,7 +1922,7 @@ impl SessionState {
19221922
let analyzed_plan =
19231923
self.analyzer
19241924
.execute_and_check(plan, self.options(), |_, _| {})?;
1925-
self.optimizer.optimize(&analyzed_plan, self, |_, _| {})
1925+
self.optimizer.optimize(analyzed_plan, self, |_, _| {})
19261926
}
19271927
}
19281928

datafusion/core/tests/optimizer_integration.rs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -110,7 +110,7 @@ fn test_sql(sql: &str) -> Result<LogicalPlan> {
110110
let optimizer = Optimizer::new();
111111
// analyze and optimize the logical plan
112112
let plan = analyzer.execute_and_check(&plan, config.options(), |_, _| {})?;
113-
optimizer.optimize(&plan, &config, |_, _| {})
113+
optimizer.optimize(plan, &config, |_, _| {})
114114
}
115115

116116
#[derive(Default)]

0 commit comments

Comments
 (0)