You are viewing a plain text version of this content. The canonical link for it is here.
Posted to github@arrow.apache.org by "alamb (via GitHub)" <gi...@apache.org> on 2024/04/04 13:06:48 UTC
[PR] Rewrite Optimizer to use TreeNode API [arrow-datafusion]
alamb opened a new pull request, #9948:
URL: https://github.com/apache/arrow-datafusion/pull/9948
## Which issue does this PR close?
<!--
We generally require a GitHub issue to be filed for all bug fixes and enhancements and this helps us generate change logs for our releases. You can link an issue to this PR using the GitHub syntax. For example `Closes #123` indicates that this PR will close issue #123.
-->
Closes #.
## Rationale for this change
The current structure of the optimizer copies Logical plans a large number of times
We now have a TreeNode API (TODO link) that does mostly the same thing and we are in the process of optimizing it to avoid copies (see https://github.com/apache/arrow-datafusion/pull/9780)
Thus it makes sense to rewrite the optimizer to use the TreeNode API to both be simpler as well as take advantage of the performance improvements we are adding.
## What changes are included in this PR?
1. Rewrite Optimizer to use TreeNode API
2. Change the Optimizer::optimize api to take an owned LogicalPlan rather than force a copy
## Are these changes tested?
<!--
We typically require tests for all PRs in order to:
1. Prevent the code from being accidentally broken by subsequent changes
2. Serve as another way to document the expected behavior of the code
If tests are not included in your PR, please explain why (for example, are they covered by existing tests)?
-->
## Are there any user-facing changes?
`Optimizer::optimize` now takes an owned `LogicalPlan` rather a reference (which forces a copy)
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Refactor `Optimizer` to use owned plans and `TreeNode` API (10% faster planning) [arrow-datafusion]
Posted by "jayzhan211 (via GitHub)" <gi...@apache.org>.
jayzhan211 commented on code in PR #9948:
URL: https://github.com/apache/arrow-datafusion/pull/9948#discussion_r1558790795
##########
datafusion/optimizer/src/optimizer.rs:
##########
@@ -299,44 +330,71 @@ impl Optimizer {
log_plan(&format!("Optimizer input (pass {i})"), &new_plan);
for rule in &self.rules {
- let result =
- self.optimize_recursively(rule, &new_plan, config)
- .and_then(|plan| {
- if let Some(plan) = &plan {
- assert_schema_is_the_same(rule.name(), plan, &new_plan)?;
- }
- Ok(plan)
- });
- match result {
- Ok(Some(plan)) => {
- new_plan = plan;
- observer(&new_plan, rule.as_ref());
- log_plan(rule.name(), &new_plan);
- }
- Ok(None) => {
+ // If skipping failed rules, copy plan before attempting to rewrite
+ // as rewriting is destructive
+ let prev_plan = options
+ .optimizer
+ .skip_failed_rules
+ .then(|| new_plan.clone());
+
+ let starting_schema = new_plan.schema().clone();
+
+ let result = match rule.apply_order() {
+ // optimizer handles recursion
+ Some(apply_order) => new_plan.rewrite(&mut Rewriter::new(
Review Comment:
Oh, I see the `map_children`
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Refactor `Optimizer` to use owned plans and `TreeNode` API (10% faster planning) [arrow-datafusion]
Posted by "jackwener (via GitHub)" <gi...@apache.org>.
jackwener commented on code in PR #9948:
URL: https://github.com/apache/arrow-datafusion/pull/9948#discussion_r1558673852
##########
datafusion/optimizer/src/optimizer.rs:
##########
@@ -299,44 +330,71 @@ impl Optimizer {
log_plan(&format!("Optimizer input (pass {i})"), &new_plan);
for rule in &self.rules {
- let result =
- self.optimize_recursively(rule, &new_plan, config)
- .and_then(|plan| {
- if let Some(plan) = &plan {
- assert_schema_is_the_same(rule.name(), plan, &new_plan)?;
- }
- Ok(plan)
- });
- match result {
- Ok(Some(plan)) => {
- new_plan = plan;
- observer(&new_plan, rule.as_ref());
- log_plan(rule.name(), &new_plan);
- }
- Ok(None) => {
+ // If skipping failed rules, copy plan before attempting to rewrite
+ // as rewriting is destructive
+ let prev_plan = options
+ .optimizer
+ .skip_failed_rules
+ .then(|| new_plan.clone());
+
+ let starting_schema = new_plan.schema().clone();
+
+ let result = match rule.apply_order() {
+ // optimizer handles recursion
+ Some(apply_order) => new_plan.rewrite(&mut Rewriter::new(
Review Comment:
It's a generic api (so using rewrite makes sense to me), it's not introduced by this PR
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Refactor `Optimizer` to use owned plans and `TreeNode` API (10% faster planning) [arrow-datafusion]
Posted by "jayzhan211 (via GitHub)" <gi...@apache.org>.
jayzhan211 commented on code in PR #9948:
URL: https://github.com/apache/arrow-datafusion/pull/9948#discussion_r1558786969
##########
datafusion/optimizer/src/optimizer.rs:
##########
@@ -299,44 +330,71 @@ impl Optimizer {
log_plan(&format!("Optimizer input (pass {i})"), &new_plan);
for rule in &self.rules {
- let result =
- self.optimize_recursively(rule, &new_plan, config)
- .and_then(|plan| {
- if let Some(plan) = &plan {
- assert_schema_is_the_same(rule.name(), plan, &new_plan)?;
- }
- Ok(plan)
- });
- match result {
- Ok(Some(plan)) => {
- new_plan = plan;
- observer(&new_plan, rule.as_ref());
- log_plan(rule.name(), &new_plan);
- }
- Ok(None) => {
+ // If skipping failed rules, copy plan before attempting to rewrite
+ // as rewriting is destructive
+ let prev_plan = options
+ .optimizer
+ .skip_failed_rules
+ .then(|| new_plan.clone());
+
+ let starting_schema = new_plan.schema().clone();
+
+ let result = match rule.apply_order() {
+ // optimizer handles recursion
+ Some(apply_order) => new_plan.rewrite(&mut Rewriter::new(
Review Comment:
I got a question here.
In previous code, `optimize_inputs` get the `plan.inputs()`. how does `rewrite` get the `plan.inputs()` here? How does `childnode` in `rewrite` equals to `plan.inputs()`
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Refactor `Optimizer` to use owned plans and `TreeNode` API (10% faster planning) [arrow-datafusion]
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb commented on code in PR #9948:
URL: https://github.com/apache/arrow-datafusion/pull/9948#discussion_r1559203377
##########
datafusion/optimizer/src/optimizer.rs:
##########
@@ -299,44 +330,71 @@ impl Optimizer {
log_plan(&format!("Optimizer input (pass {i})"), &new_plan);
for rule in &self.rules {
- let result =
- self.optimize_recursively(rule, &new_plan, config)
- .and_then(|plan| {
- if let Some(plan) = &plan {
- assert_schema_is_the_same(rule.name(), plan, &new_plan)?;
- }
- Ok(plan)
- });
- match result {
- Ok(Some(plan)) => {
- new_plan = plan;
- observer(&new_plan, rule.as_ref());
- log_plan(rule.name(), &new_plan);
- }
- Ok(None) => {
+ // If skipping failed rules, copy plan before attempting to rewrite
+ // as rewriting is destructive
+ let prev_plan = options
+ .optimizer
+ .skip_failed_rules
+ .then(|| new_plan.clone());
+
+ let starting_schema = new_plan.schema().clone();
+
+ let result = match rule.apply_order() {
+ // optimizer handles recursion
+ Some(apply_order) => new_plan.rewrite(&mut Rewriter::new(
Review Comment:
Yes, exactly
`rewrite` : https://github.com/apache/arrow-datafusion/blob/cb21404bd3736ff9a6d8d443a67c64ece4c551a9/datafusion/common/src/tree_node.rs#L117-L124
(eventually) calls `map_children`: https://github.com/apache/arrow-datafusion/blob/cb21404bd3736ff9a6d8d443a67c64ece4c551a9/datafusion/common/src/tree_node.rs#L27-L31
This is the beauty of these `TreeNode` APIs and @peter-toth's plan to make them all consistent.
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Refactor `Optimizer` to use owned plans and `TreeNode` API (10% faster planning) [arrow-datafusion]
Posted by "jayzhan211 (via GitHub)" <gi...@apache.org>.
jayzhan211 commented on code in PR #9948:
URL: https://github.com/apache/arrow-datafusion/pull/9948#discussion_r1558556613
##########
datafusion/optimizer/src/test/mod.rs:
##########
@@ -150,22 +150,19 @@ pub fn assert_analyzer_check_err(
}
}
}
+
+fn observe(_plan: &LogicalPlan, _rule: &dyn OptimizerRule) {}
+
pub fn assert_optimized_plan_eq(
rule: Arc<dyn OptimizerRule + Send + Sync>,
- plan: &LogicalPlan,
+ plan: LogicalPlan,
expected: &str,
) -> Result<()> {
- let optimizer = Optimizer::with_rules(vec![rule.clone()]);
- let optimized_plan = optimizer
- .optimize_recursively(
- optimizer.rules.first().unwrap(),
- plan,
- &OptimizerContext::new(),
- )?
- .unwrap_or_else(|| plan.clone());
+ // Apply the rule once
+ let opt_context = OptimizerContext::new().with_max_passes(1);
Review Comment:
I'm not sure about this. It seems that the previous code does not limit the pass to 1. Why do we need to limit it now to have the same behavior as the previous one? 😕
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Refactor `Optimizer` to use owned plans and `TreeNode` API (10% faster planning) [arrow-datafusion]
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb commented on PR #9948:
URL: https://github.com/apache/arrow-datafusion/pull/9948#issuecomment-2045280477
@jackwener since you implemented some of the original optimizer recursion I wonder if you would have some time to review this PR
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Refactor `Optimizer` to use owned plans and `TreeNode` API (10% faster planning) [arrow-datafusion]
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb merged PR #9948:
URL: https://github.com/apache/arrow-datafusion/pull/9948
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Refactor `Optimizer` to use owned plans and `TreeNode` API (10% faster planning) [arrow-datafusion]
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb commented on code in PR #9948:
URL: https://github.com/apache/arrow-datafusion/pull/9948#discussion_r1557570334
##########
datafusion/optimizer/src/push_down_projection.rs:
##########
@@ -638,25 +638,23 @@ mod tests {
\n WindowAggr: windowExpr=[[MAX(test.a) PARTITION BY [test.b] ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING]]\
\n TableScan: test projection=[a, b]";
- assert_optimized_plan_eq(&plan, expected)
+ assert_optimized_plan_eq(plan, expected)
}
- fn assert_optimized_plan_eq(plan: &LogicalPlan, expected: &str) -> Result<()> {
+ fn assert_optimized_plan_eq(plan: LogicalPlan, expected: &str) -> Result<()> {
let optimized_plan = optimize(plan).expect("failed to optimize plan");
let formatted_plan = format!("{optimized_plan:?}");
assert_eq!(formatted_plan, expected);
Ok(())
}
- fn optimize(plan: &LogicalPlan) -> Result<LogicalPlan> {
+ fn optimize(plan: LogicalPlan) -> Result<LogicalPlan> {
let optimizer = Optimizer::with_rules(vec![Arc::new(OptimizeProjections::new())]);
- let optimized_plan = optimizer
- .optimize_recursively(
- optimizer.rules.first().unwrap(),
- plan,
- &OptimizerContext::new(),
- )?
- .unwrap_or_else(|| plan.clone());
+ let optimized_plan =
Review Comment:
Since the optimizer correctly applies rules recursively now, there is no need to explicitly call optimize recursively
##########
datafusion/optimizer/src/test/mod.rs:
##########
@@ -150,22 +150,19 @@ pub fn assert_analyzer_check_err(
}
}
}
+
+fn observe(_plan: &LogicalPlan, _rule: &dyn OptimizerRule) {}
+
pub fn assert_optimized_plan_eq(
rule: Arc<dyn OptimizerRule + Send + Sync>,
- plan: &LogicalPlan,
+ plan: LogicalPlan,
expected: &str,
) -> Result<()> {
- let optimizer = Optimizer::with_rules(vec![rule.clone()]);
- let optimized_plan = optimizer
- .optimize_recursively(
- optimizer.rules.first().unwrap(),
- plan,
- &OptimizerContext::new(),
- )?
- .unwrap_or_else(|| plan.clone());
+ // Apply the rule once
+ let opt_context = OptimizerContext::new().with_max_passes(1);
Review Comment:
Without this some of the tests don't pass. By default the optimizer runs a few times until no changes are detected. Limiting to 1 pass mimics the previous test behavior
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Refactor `Optimizer` to use owned plans and `TreeNode` API (10% faster planning) [arrow-datafusion]
Posted by "jackwener (via GitHub)" <gi...@apache.org>.
jackwener commented on code in PR #9948:
URL: https://github.com/apache/arrow-datafusion/pull/9948#discussion_r1558671039
##########
datafusion-examples/examples/rewrite_expr.rs:
##########
@@ -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)?;
Review Comment:
A great progress!
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Refactor `Optimizer` to use owned plans and `TreeNode` API (10% faster planning) [arrow-datafusion]
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb commented on code in PR #9948:
URL: https://github.com/apache/arrow-datafusion/pull/9948#discussion_r1551678414
##########
datafusion/optimizer/src/eliminate_limit.rs:
##########
@@ -93,24 +93,19 @@ mod tests {
use crate::push_down_limit::PushDownLimit;
- fn assert_optimized_plan_eq(plan: &LogicalPlan, expected: &str) -> Result<()> {
+ fn observe(_plan: &LogicalPlan, _rule: &dyn OptimizerRule) {}
+ fn assert_optimized_plan_eq(plan: LogicalPlan, expected: &str) -> Result<()> {
let optimizer = Optimizer::with_rules(vec![Arc::new(EliminateLimit::new())]);
- let optimized_plan = optimizer
- .optimize_recursively(
- optimizer.rules.first().unwrap(),
- plan,
- &OptimizerContext::new(),
- )?
- .unwrap_or_else(|| plan.clone());
+ let optimized_plan =
+ optimizer.optimize(plan, &OptimizerContext::new(), observe)?;
let formatted_plan = format!("{optimized_plan:?}");
assert_eq!(formatted_plan, expected);
- assert_eq!(plan.schema(), optimized_plan.schema());
Review Comment:
I changed the tests to call `Optimizer::optimize` directly, which already checks the schema doesn't change, so this test is redundant
This applies to several other changes in this PR
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Refactor `Optimizer` to use owned plans and `TreeNode` API (10% faster planning) [arrow-datafusion]
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb commented on code in PR #9948:
URL: https://github.com/apache/arrow-datafusion/pull/9948#discussion_r1551705192
##########
datafusion/optimizer/src/optimizer.rs:
##########
@@ -356,97 +423,22 @@ impl Optimizer {
debug!("Optimizer took {} ms", start_time.elapsed().as_millis());
Ok(new_plan)
}
-
- fn optimize_node(
Review Comment:
This code implemented plan recursion within the optimizer and is (now) redundant with the `TreeNode` API
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Refactor `Optimizer` to use owned plans and `TreeNode` API (10% faster planning) [arrow-datafusion]
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb commented on PR #9948:
URL: https://github.com/apache/arrow-datafusion/pull/9948#issuecomment-2047152077
Thank you very much @jackwener and @jayzhan211 for the reviews 🙏
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Refactor `Optimizer` to use owned plans and `TreeNode` API (10% faster planning) [arrow-datafusion]
Posted by "mustafasrepo (via GitHub)" <gi...@apache.org>.
mustafasrepo commented on PR #9948:
URL: https://github.com/apache/arrow-datafusion/pull/9948#issuecomment-2055474638
Thanks @alamb for this work.
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Refactor `Optimizer` to use `TreeNode` API [arrow-datafusion]
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb commented on code in PR #9948:
URL: https://github.com/apache/arrow-datafusion/pull/9948#discussion_r1551676725
##########
datafusion-examples/examples/rewrite_expr.rs:
##########
@@ -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)?;
Review Comment:
This illustrates the API change -- the optimizer now takes an owned plan rather than a reference
##########
datafusion/core/tests/optimizer_integration.rs:
##########
@@ -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, |_, _| {})
Review Comment:
A large amount of this PR is changes to test to pass in an owned plan
##########
datafusion/optimizer/src/optimizer.rs:
##########
@@ -506,21 +498,27 @@ mod tests {
produce_one_row: false,
schema: Arc::new(DFSchema::empty()),
});
- let err = opt.optimize(&plan, &config, &observe).unwrap_err();
+ let err = opt.optimize(plan, &config, &observe).unwrap_err();
assert_eq!(
- "Optimizer rule 'get table_scan rule' failed\ncaused by\nget table_scan rule\ncaused by\n\
- Internal error: Failed due to a difference in schemas, original schema: \
- DFSchema { inner: Schema { fields: \
- [Field { name: \"a\", data_type: UInt32, nullable: false, dict_id: 0, dict_is_ordered: false, metadata: {} }, \
- Field { name: \"b\", data_type: UInt32, nullable: false, dict_id: 0, dict_is_ordered: false, metadata: {} }, \
- Field { name: \"c\", data_type: UInt32, nullable: false, dict_id: 0, dict_is_ordered: false, metadata: {} }], metadata: {} }, \
- field_qualifiers: [Some(Bare { table: \"test\" }), Some(Bare { table: \"test\" }), Some(Bare { table: \"test\" })], \
- functional_dependencies: FunctionalDependencies { deps: [] } }, \
+ "Optimizer rule 'get table_scan rule' failed\n\
Review Comment:
The original error actually is incorrect that it reports the reversed schemas (the "new schema" was actually the original schema)
##########
datafusion/optimizer/src/optimizer.rs:
##########
@@ -274,22 +248,82 @@ impl Optimizer {
pub fn with_rules(rules: Vec<Arc<dyn OptimizerRule + Send + Sync>>) -> Self {
Self { rules }
}
+}
+
+/// Rewrites LogicalPlan nodes
+struct Rewriter<'a> {
+ apply_order: ApplyOrder,
+ rule: &'a dyn OptimizerRule,
+ config: &'a dyn OptimizerConfig,
+}
+
+impl<'a> Rewriter<'a> {
+ fn new(
+ apply_order: ApplyOrder,
+ rule: &'a dyn OptimizerRule,
+ config: &'a dyn OptimizerConfig,
+ ) -> Self {
+ Self {
+ apply_order,
+ rule,
+ config,
+ }
+ }
+}
+
+impl<'a> TreeNodeRewriter for Rewriter<'a> {
+ type Node = LogicalPlan;
+
+ fn f_down(&mut self, node: LogicalPlan) -> Result<Transformed<LogicalPlan>> {
+ if self.apply_order == ApplyOrder::TopDown {
+ optimize_plan_node(node, self.rule, self.config)
+ } else {
+ Ok(Transformed::no(node))
+ }
+ }
+ fn f_up(&mut self, node: LogicalPlan) -> Result<Transformed<LogicalPlan>> {
+ if self.apply_order == ApplyOrder::BottomUp {
+ optimize_plan_node(node, self.rule, self.config)
+ } else {
+ Ok(Transformed::no(node))
+ }
+ }
+}
+
+/// Invokes the Optimizer rule to rewrite the LogicalPlan in place.
+fn optimize_plan_node(
+ plan: LogicalPlan,
+ rule: &dyn OptimizerRule,
+ config: &dyn OptimizerConfig,
+) -> Result<Transformed<LogicalPlan>> {
+ // TODO: introduce a better API to OptimizerRule to allow rewriting by ownership
Review Comment:
Each rule still requires the plan to be copied, which is not ideal. I plan to add a better API to avoid this
##########
datafusion/optimizer/src/optimizer.rs:
##########
@@ -299,44 +333,77 @@ impl Optimizer {
log_plan(&format!("Optimizer input (pass {i})"), &new_plan);
for rule in &self.rules {
- let result =
- self.optimize_recursively(rule, &new_plan, config)
- .and_then(|plan| {
- if let Some(plan) = &plan {
- assert_schema_is_the_same(rule.name(), plan, &new_plan)?;
- }
- Ok(plan)
- });
- match result {
- Ok(Some(plan)) => {
- new_plan = plan;
- observer(&new_plan, rule.as_ref());
- log_plan(rule.name(), &new_plan);
+ // If we need to skip failed rules, must copy plan before attempting to rewrite
Review Comment:
this change ensures we only copy the plan if we need it to remain unchanged
##########
datafusion/optimizer/src/optimizer.rs:
##########
@@ -356,97 +423,22 @@ impl Optimizer {
debug!("Optimizer took {} ms", start_time.elapsed().as_millis());
Ok(new_plan)
}
-
- fn optimize_node(
Review Comment:
this code handles recursion and is redundant with the tree node api
##########
datafusion/optimizer/src/eliminate_limit.rs:
##########
@@ -93,24 +93,19 @@ mod tests {
use crate::push_down_limit::PushDownLimit;
- fn assert_optimized_plan_eq(plan: &LogicalPlan, expected: &str) -> Result<()> {
+ fn observe(_plan: &LogicalPlan, _rule: &dyn OptimizerRule) {}
+ fn assert_optimized_plan_eq(plan: LogicalPlan, expected: &str) -> Result<()> {
let optimizer = Optimizer::with_rules(vec![Arc::new(EliminateLimit::new())]);
- let optimized_plan = optimizer
- .optimize_recursively(
- optimizer.rules.first().unwrap(),
- plan,
- &OptimizerContext::new(),
- )?
- .unwrap_or_else(|| plan.clone());
+ let optimized_plan =
+ optimizer.optimize(plan, &OptimizerContext::new(), observe)?;
let formatted_plan = format!("{optimized_plan:?}");
assert_eq!(formatted_plan, expected);
- assert_eq!(plan.schema(), optimized_plan.schema());
Review Comment:
I changed the tests to call `Optimizer::optimize` directly, which already checks the schema doesn't change, so this test is redundant
##########
datafusion/optimizer/src/optimizer.rs:
##########
@@ -184,39 +185,12 @@ pub struct Optimizer {
pub rules: Vec<Arc<dyn OptimizerRule + Send + Sync>>,
}
-/// If a rule is with `ApplyOrder`, it means the optimizer will derive to handle children instead of
-/// recursively handling in rule.
-/// We just need handle a subtree pattern itself.
+/// Specifies how recursion for an `OptimizerRule` should be handled.
///
-/// Notice: **sometime** result after optimize still can be optimized, we need apply again.
Review Comment:
I do not think this comment is applicable anymore -- the optimizer handles the recursion internally as well as applying multiple optimizer passes
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Refactor `Optimizer` to use owned plans and `TreeNode` API (10% faster planning) [arrow-datafusion]
Posted by "jayzhan211 (via GitHub)" <gi...@apache.org>.
jayzhan211 commented on code in PR #9948:
URL: https://github.com/apache/arrow-datafusion/pull/9948#discussion_r1558566105
##########
datafusion/optimizer/src/optimizer.rs:
##########
@@ -299,44 +330,71 @@ impl Optimizer {
log_plan(&format!("Optimizer input (pass {i})"), &new_plan);
for rule in &self.rules {
- let result =
- self.optimize_recursively(rule, &new_plan, config)
- .and_then(|plan| {
- if let Some(plan) = &plan {
- assert_schema_is_the_same(rule.name(), plan, &new_plan)?;
- }
- Ok(plan)
- });
- match result {
- Ok(Some(plan)) => {
- new_plan = plan;
- observer(&new_plan, rule.as_ref());
- log_plan(rule.name(), &new_plan);
- }
- Ok(None) => {
+ // If skipping failed rules, copy plan before attempting to rewrite
+ // as rewriting is destructive
+ let prev_plan = options
+ .optimizer
+ .skip_failed_rules
+ .then(|| new_plan.clone());
+
+ let starting_schema = new_plan.schema().clone();
+
+ let result = match rule.apply_order() {
+ // optimizer handles recursion
+ Some(apply_order) => new_plan.rewrite(&mut Rewriter::new(
Review Comment:
Does rename it to `rewrite_recurisvely` more straightforward?
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Refactor `Optimizer` to use owned plans and `TreeNode` API (10% faster planning) [arrow-datafusion]
Posted by "jayzhan211 (via GitHub)" <gi...@apache.org>.
jayzhan211 commented on code in PR #9948:
URL: https://github.com/apache/arrow-datafusion/pull/9948#discussion_r1559296615
##########
datafusion/optimizer/src/test/mod.rs:
##########
@@ -150,22 +150,19 @@ pub fn assert_analyzer_check_err(
}
}
}
+
+fn observe(_plan: &LogicalPlan, _rule: &dyn OptimizerRule) {}
+
pub fn assert_optimized_plan_eq(
rule: Arc<dyn OptimizerRule + Send + Sync>,
- plan: &LogicalPlan,
+ plan: LogicalPlan,
expected: &str,
) -> Result<()> {
- let optimizer = Optimizer::with_rules(vec![rule.clone()]);
- let optimized_plan = optimizer
- .optimize_recursively(
- optimizer.rules.first().unwrap(),
- plan,
- &OptimizerContext::new(),
- )?
- .unwrap_or_else(|| plan.clone());
+ // Apply the rule once
+ let opt_context = OptimizerContext::new().with_max_passes(1);
Review Comment:
I see!
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Refactor `Optimizer` to use owned plans and `TreeNode` API (10% faster planning) [arrow-datafusion]
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb commented on code in PR #9948:
URL: https://github.com/apache/arrow-datafusion/pull/9948#discussion_r1557568328
##########
datafusion/optimizer/src/optimizer.rs:
##########
@@ -274,22 +249,78 @@ impl Optimizer {
pub fn with_rules(rules: Vec<Arc<dyn OptimizerRule + Send + Sync>>) -> Self {
Self { rules }
}
+}
+
+/// Recursively rewrites LogicalPlans
+struct Rewriter<'a> {
Review Comment:
`datafusion/optimizer/src/optimizer.rs` has all the important changes (to use the TreeNode API and stop copying)
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Refactor `Optimizer` to use owned plans and `TreeNode` API (10% faster planning) [arrow-datafusion]
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb commented on code in PR #9948:
URL: https://github.com/apache/arrow-datafusion/pull/9948#discussion_r1559194542
##########
datafusion/optimizer/src/test/mod.rs:
##########
@@ -150,22 +150,19 @@ pub fn assert_analyzer_check_err(
}
}
}
+
+fn observe(_plan: &LogicalPlan, _rule: &dyn OptimizerRule) {}
+
pub fn assert_optimized_plan_eq(
rule: Arc<dyn OptimizerRule + Send + Sync>,
- plan: &LogicalPlan,
+ plan: LogicalPlan,
expected: &str,
) -> Result<()> {
- let optimizer = Optimizer::with_rules(vec![rule.clone()]);
- let optimized_plan = optimizer
- .optimize_recursively(
- optimizer.rules.first().unwrap(),
- plan,
- &OptimizerContext::new(),
- )?
- .unwrap_or_else(|| plan.clone());
+ // Apply the rule once
+ let opt_context = OptimizerContext::new().with_max_passes(1);
Review Comment:
My explanation goes like: The loop that applies the rule more than once calls `optimize_recursively` each time
https://github.com/apache/arrow-datafusion/blob/75c399ce7d4d5360140c64089dd7b05ffd7c49ef/datafusion/optimizer/src/optimizer.rs#L298-L303
This test only called `optimze_recursively` once (directly) and thus the `OptimizeRule` is only applied once
When I rewrote the test to use `Optimizer::optimize` the loop will now kick in and so the `OptimizeRule` will be run several times unless we set `with_max_passes`
This same reasoning applies to the other tests, but apparently they get the same answer when applied more than once
--
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: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org