This is an automated email from the ASF dual-hosted git repository.
alamb pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion.git
The following commit(s) were added to refs/heads/main by this push:
new 6dffc53e76 Deprecate `OptimizerRule::try_optimize` (#11022)
6dffc53e76 is described below
commit 6dffc53e76fb9a785dafd55106dd4033e744bf8e
Author: 张林伟 <[email protected]>
AuthorDate: Sat Jun 22 00:20:21 2024 +0800
Deprecate `OptimizerRule::try_optimize` (#11022)
* Deprecate OptimizerRule::try_optimize
* optimize_children
* Apply review suggestions
* Fix clippy lint
---
datafusion-examples/examples/rewrite_expr.rs | 48 ++++++++++++----------
.../core/tests/user_defined/user_defined_plan.rs | 44 +++++++++++++-------
datafusion/optimizer/src/lib.rs | 1 +
datafusion/optimizer/src/optimizer.rs | 21 +++++++++-
datafusion/optimizer/src/utils.rs | 17 ++++++--
5 files changed, 89 insertions(+), 42 deletions(-)
diff --git a/datafusion-examples/examples/rewrite_expr.rs
b/datafusion-examples/examples/rewrite_expr.rs
index d8965888ea..556687a46a 100644
--- a/datafusion-examples/examples/rewrite_expr.rs
+++ b/datafusion-examples/examples/rewrite_expr.rs
@@ -18,13 +18,13 @@
use arrow::datatypes::{DataType, Field, Schema, SchemaRef};
use datafusion_common::config::ConfigOptions;
use datafusion_common::tree_node::{Transformed, TransformedResult, TreeNode};
-use datafusion_common::{plan_err, Result, ScalarValue};
+use datafusion_common::{plan_err, DataFusionError, Result, ScalarValue};
use datafusion_expr::{
AggregateUDF, Between, Expr, Filter, LogicalPlan, ScalarUDF, TableSource,
WindowUDF,
};
use datafusion_optimizer::analyzer::{Analyzer, AnalyzerRule};
-use datafusion_optimizer::optimizer::Optimizer;
-use datafusion_optimizer::{utils, OptimizerConfig, OptimizerContext,
OptimizerRule};
+use datafusion_optimizer::optimizer::{ApplyOrder, Optimizer};
+use datafusion_optimizer::{OptimizerConfig, OptimizerContext, OptimizerRule};
use datafusion_sql::planner::{ContextProvider, SqlToRel};
use datafusion_sql::sqlparser::dialect::PostgreSqlDialect;
use datafusion_sql::sqlparser::parser::Parser;
@@ -133,30 +133,34 @@ impl OptimizerRule for MyOptimizerRule {
fn try_optimize(
&self,
- plan: &LogicalPlan,
- config: &dyn OptimizerConfig,
+ _plan: &LogicalPlan,
+ _config: &dyn OptimizerConfig,
) -> Result<Option<LogicalPlan>> {
- // recurse down and optimize children first
- let optimized_plan = utils::optimize_children(self, plan, config)?;
- match optimized_plan {
- Some(LogicalPlan::Filter(filter)) => {
+ unreachable!()
+ }
+
+ fn apply_order(&self) -> Option<ApplyOrder> {
+ Some(ApplyOrder::BottomUp)
+ }
+
+ fn supports_rewrite(&self) -> bool {
+ true
+ }
+
+ fn rewrite(
+ &self,
+ plan: LogicalPlan,
+ _config: &dyn OptimizerConfig,
+ ) -> Result<Transformed<LogicalPlan>, DataFusionError> {
+ match plan {
+ LogicalPlan::Filter(filter) => {
let predicate = my_rewrite(filter.predicate.clone())?;
- Ok(Some(LogicalPlan::Filter(Filter::try_new(
+ Ok(Transformed::yes(LogicalPlan::Filter(Filter::try_new(
predicate,
- filter.input,
+ filter.input.clone(),
)?)))
}
- Some(optimized_plan) => Ok(Some(optimized_plan)),
- None => match plan {
- LogicalPlan::Filter(filter) => {
- let predicate = my_rewrite(filter.predicate.clone())?;
- Ok(Some(LogicalPlan::Filter(Filter::try_new(
- predicate,
- filter.input.clone(),
- )?)))
- }
- _ => Ok(None),
- },
+ _ => Ok(Transformed::no(plan)),
}
}
}
diff --git a/datafusion/core/tests/user_defined/user_defined_plan.rs
b/datafusion/core/tests/user_defined/user_defined_plan.rs
index 07622e48af..ebf907c5e2 100644
--- a/datafusion/core/tests/user_defined/user_defined_plan.rs
+++ b/datafusion/core/tests/user_defined/user_defined_plan.rs
@@ -80,7 +80,7 @@ use datafusion::{
Expr, Extension, Limit, LogicalPlan, Sort, UserDefinedLogicalNode,
UserDefinedLogicalNodeCore,
},
- optimizer::{optimize_children, OptimizerConfig, OptimizerRule},
+ optimizer::{OptimizerConfig, OptimizerRule},
physical_expr::EquivalenceProperties,
physical_plan::{
DisplayAs, DisplayFormatType, Distribution, ExecutionMode,
ExecutionPlan,
@@ -92,6 +92,8 @@ use datafusion::{
};
use async_trait::async_trait;
+use datafusion_common::tree_node::Transformed;
+use datafusion_optimizer::optimizer::ApplyOrder;
use futures::{Stream, StreamExt};
/// Execute the specified sql and return the resulting record batches
@@ -282,9 +284,29 @@ impl OptimizerRule for TopKOptimizerRule {
// Example rewrite pass to insert a user defined LogicalPlanNode
fn try_optimize(
&self,
- plan: &LogicalPlan,
- config: &dyn OptimizerConfig,
+ _plan: &LogicalPlan,
+ _config: &dyn OptimizerConfig,
) -> Result<Option<LogicalPlan>> {
+ unreachable!()
+ }
+
+ fn name(&self) -> &str {
+ "topk"
+ }
+
+ fn apply_order(&self) -> Option<ApplyOrder> {
+ Some(ApplyOrder::TopDown)
+ }
+
+ fn supports_rewrite(&self) -> bool {
+ true
+ }
+
+ fn rewrite(
+ &self,
+ plan: LogicalPlan,
+ _config: &dyn OptimizerConfig,
+ ) -> Result<Transformed<LogicalPlan>, DataFusionError> {
// Note: this code simply looks for the pattern of a Limit followed by
a
// Sort and replaces it by a TopK node. It does not handle many
// edge cases (e.g multiple sort columns, sort ASC / DESC), etc.
@@ -292,7 +314,7 @@ impl OptimizerRule for TopKOptimizerRule {
fetch: Some(fetch),
input,
..
- }) = plan
+ }) = &plan
{
if let LogicalPlan::Sort(Sort {
ref expr,
@@ -302,12 +324,10 @@ impl OptimizerRule for TopKOptimizerRule {
{
if expr.len() == 1 {
// we found a sort with a single sort expr, replace with a
a TopK
- return Ok(Some(LogicalPlan::Extension(Extension {
+ return
Ok(Transformed::yes(LogicalPlan::Extension(Extension {
node: Arc::new(TopKPlanNode {
k: *fetch,
- input: self
- .try_optimize(input.as_ref(), config)?
- .unwrap_or_else(|| input.as_ref().clone()),
+ input: input.as_ref().clone(),
expr: expr[0].clone(),
}),
})));
@@ -315,13 +335,7 @@ impl OptimizerRule for TopKOptimizerRule {
}
}
- // If we didn't find the Limit/Sort combination, recurse as
- // normal and build the result.
- optimize_children(self, plan, config)
- }
-
- fn name(&self) -> &str {
- "topk"
+ Ok(Transformed::no(plan))
}
}
diff --git a/datafusion/optimizer/src/lib.rs b/datafusion/optimizer/src/lib.rs
index c172d59797..a6a9e5cf26 100644
--- a/datafusion/optimizer/src/lib.rs
+++ b/datafusion/optimizer/src/lib.rs
@@ -61,6 +61,7 @@ pub mod test;
pub use analyzer::{Analyzer, AnalyzerRule};
pub use optimizer::{Optimizer, OptimizerConfig, OptimizerContext,
OptimizerRule};
+#[allow(deprecated)]
pub use utils::optimize_children;
pub(crate) mod join_key_set;
diff --git a/datafusion/optimizer/src/optimizer.rs
b/datafusion/optimizer/src/optimizer.rs
index 75ebc92ece..7a3ea6ed4c 100644
--- a/datafusion/optimizer/src/optimizer.rs
+++ b/datafusion/optimizer/src/optimizer.rs
@@ -77,6 +77,10 @@ pub trait OptimizerRule {
/// Note this API will be deprecated in the future as it requires
`clone`ing
/// the input plan, which can be expensive. OptimizerRules should implement
/// [`Self::rewrite`] instead.
+ #[deprecated(
+ since = "40.0.0",
+ note = "please implement supports_rewrite and rewrite instead"
+ )]
fn try_optimize(
&self,
plan: &LogicalPlan,
@@ -332,6 +336,7 @@ fn optimize_plan_node(
return rule.rewrite(plan, config);
}
+ #[allow(deprecated)]
rule.try_optimize(&plan, config).map(|maybe_plan| {
match maybe_plan {
Some(new_plan) => {
@@ -483,7 +488,7 @@ mod tests {
use std::sync::{Arc, Mutex};
use datafusion_common::tree_node::Transformed;
- use datafusion_common::{plan_err, DFSchema, DFSchemaRef, Result};
+ use datafusion_common::{plan_err, DFSchema, DFSchemaRef, DataFusionError,
Result};
use datafusion_expr::logical_plan::EmptyRelation;
use datafusion_expr::{col, lit, LogicalPlan, LogicalPlanBuilder,
Projection};
@@ -667,12 +672,24 @@ mod tests {
_: &LogicalPlan,
_: &dyn OptimizerConfig,
) -> Result<Option<LogicalPlan>> {
- plan_err!("rule failed")
+ unreachable!()
}
fn name(&self) -> &str {
"bad rule"
}
+
+ fn supports_rewrite(&self) -> bool {
+ true
+ }
+
+ fn rewrite(
+ &self,
+ _plan: LogicalPlan,
+ _config: &dyn OptimizerConfig,
+ ) -> Result<Transformed<LogicalPlan>, DataFusionError> {
+ plan_err!("rule failed")
+ }
}
/// Replaces whatever plan with a single table scan
diff --git a/datafusion/optimizer/src/utils.rs
b/datafusion/optimizer/src/utils.rs
index fd47cb23b1..6218140409 100644
--- a/datafusion/optimizer/src/utils.rs
+++ b/datafusion/optimizer/src/utils.rs
@@ -35,6 +35,10 @@ use log::{debug, trace};
/// This also handles the case when the `plan` is a [`LogicalPlan::Explain`].
///
/// Returning `Ok(None)` indicates that the plan can't be optimized by the
`optimizer`.
+#[deprecated(
+ since = "40.0.0",
+ note = "please use OptimizerRule::apply_order with ApplyOrder::BottomUp
instead"
+)]
pub fn optimize_children(
optimizer: &impl OptimizerRule,
plan: &LogicalPlan,
@@ -43,9 +47,16 @@ pub fn optimize_children(
let mut new_inputs = Vec::with_capacity(plan.inputs().len());
let mut plan_is_changed = false;
for input in plan.inputs() {
- let new_input = optimizer.try_optimize(input, config)?;
- plan_is_changed = plan_is_changed || new_input.is_some();
- new_inputs.push(new_input.unwrap_or_else(|| input.clone()))
+ if optimizer.supports_rewrite() {
+ let new_input = optimizer.rewrite(input.clone(), config)?;
+ plan_is_changed = plan_is_changed || new_input.transformed;
+ new_inputs.push(new_input.data);
+ } else {
+ #[allow(deprecated)]
+ let new_input = optimizer.try_optimize(input, config)?;
+ plan_is_changed = plan_is_changed || new_input.is_some();
+ new_inputs.push(new_input.unwrap_or_else(|| input.clone()))
+ }
}
if plan_is_changed {
let exprs = plan.expressions();
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
