|
18 | 18 | use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; |
19 | 19 | use datafusion_common::config::ConfigOptions; |
20 | 20 | use datafusion_common::tree_node::{Transformed, TransformedResult, TreeNode}; |
21 | | -use datafusion_common::{plan_err, Result, ScalarValue}; |
| 21 | +use datafusion_common::{plan_err, DataFusionError, Result, ScalarValue}; |
22 | 22 | use datafusion_expr::{ |
23 | 23 | AggregateUDF, Between, Expr, Filter, LogicalPlan, ScalarUDF, TableSource, WindowUDF, |
24 | 24 | }; |
25 | 25 | use datafusion_optimizer::analyzer::{Analyzer, AnalyzerRule}; |
26 | | -use datafusion_optimizer::optimizer::Optimizer; |
27 | | -use datafusion_optimizer::{utils, OptimizerConfig, OptimizerContext, OptimizerRule}; |
| 26 | +use datafusion_optimizer::optimizer::{ApplyOrder, Optimizer}; |
| 27 | +use datafusion_optimizer::{OptimizerConfig, OptimizerContext, OptimizerRule}; |
28 | 28 | use datafusion_sql::planner::{ContextProvider, SqlToRel}; |
29 | 29 | use datafusion_sql::sqlparser::dialect::PostgreSqlDialect; |
30 | 30 | use datafusion_sql::sqlparser::parser::Parser; |
@@ -131,32 +131,28 @@ impl OptimizerRule for MyOptimizerRule { |
131 | 131 | "my_optimizer_rule" |
132 | 132 | } |
133 | 133 |
|
134 | | - fn try_optimize( |
| 134 | + fn apply_order(&self) -> Option<ApplyOrder> { |
| 135 | + Some(ApplyOrder::BottomUp) |
| 136 | + } |
| 137 | + |
| 138 | + fn supports_rewrite(&self) -> bool { |
| 139 | + true |
| 140 | + } |
| 141 | + |
| 142 | + fn rewrite( |
135 | 143 | &self, |
136 | | - plan: &LogicalPlan, |
137 | | - config: &dyn OptimizerConfig, |
138 | | - ) -> Result<Option<LogicalPlan>> { |
139 | | - // recurse down and optimize children first |
140 | | - let optimized_plan = utils::optimize_children(self, plan, config)?; |
141 | | - match optimized_plan { |
142 | | - Some(LogicalPlan::Filter(filter)) => { |
| 144 | + plan: LogicalPlan, |
| 145 | + _config: &dyn OptimizerConfig, |
| 146 | + ) -> Result<Transformed<LogicalPlan>, DataFusionError> { |
| 147 | + match plan { |
| 148 | + LogicalPlan::Filter(filter) => { |
143 | 149 | let predicate = my_rewrite(filter.predicate.clone())?; |
144 | | - Ok(Some(LogicalPlan::Filter(Filter::try_new( |
| 150 | + Ok(Transformed::yes(LogicalPlan::Filter(Filter::try_new( |
145 | 151 | predicate, |
146 | | - filter.input, |
| 152 | + filter.input.clone(), |
147 | 153 | )?))) |
148 | 154 | } |
149 | | - Some(optimized_plan) => Ok(Some(optimized_plan)), |
150 | | - None => match plan { |
151 | | - LogicalPlan::Filter(filter) => { |
152 | | - let predicate = my_rewrite(filter.predicate.clone())?; |
153 | | - Ok(Some(LogicalPlan::Filter(Filter::try_new( |
154 | | - predicate, |
155 | | - filter.input.clone(), |
156 | | - )?))) |
157 | | - } |
158 | | - _ => Ok(None), |
159 | | - }, |
| 155 | + _ => Ok(Transformed::no(plan)), |
160 | 156 | } |
161 | 157 | } |
162 | 158 | } |
|
0 commit comments