Make OptimizerConfig a trait (#4631) (#4638)

datafusion-examples/examples/rewrite_expr.rs

@@ -22,7 +22,7 @@ use datafusion_expr::{
     AggregateUDF, Between, Expr, Filter, LogicalPlan, ScalarUDF, TableSource,
 };
 use datafusion_optimizer::optimizer::Optimizer;
-use datafusion_optimizer::{utils, OptimizerConfig, OptimizerRule};
+use datafusion_optimizer::{utils, OptimizerConfig, OptimizerContext, OptimizerRule};
 use datafusion_sql::planner::{ContextProvider, SqlToRel};
 use datafusion_sql::sqlparser::dialect::PostgreSqlDialect;
 use datafusion_sql::sqlparser::parser::Parser;
@@ -47,9 +47,8 @@ pub fn main() -> Result<()> {
 
     // now run the optimizer with our custom rule
     let optimizer = Optimizer::with_rules(vec![Arc::new(MyRule {})]);
-    let mut optimizer_config = OptimizerConfig::default().with_skip_failing_rules(false);
-    let optimized_plan =
-        optimizer.optimize(&logical_plan, &mut optimizer_config, observe)?;
+    let config = OptimizerContext::default().with_skip_failing_rules(false);
+    let optimized_plan = optimizer.optimize(&logical_plan, &config, observe)?;
     println!(
         "Optimized Logical Plan:\n\n{}\n",
         optimized_plan.display_indent()
@@ -76,10 +75,10 @@ impl OptimizerRule for MyRule {
     fn try_optimize(
         &self,
         plan: &LogicalPlan,
-        _config: &mut OptimizerConfig,
+        config: &dyn OptimizerConfig,
     ) -> Result<Option<LogicalPlan>> {
         // recurse down and optimize children first
-        let plan = utils::optimize_children(self, plan, _config)?;
+        let plan = utils::optimize_children(self, plan, config)?;
 
         match plan {
             LogicalPlan::Filter(filter) => {

datafusion/core/src/execution/context.rs

@@ -68,7 +68,7 @@ use crate::logical_expr::{
     CreateView, DropTable, DropView, Explain, LogicalPlan, LogicalPlanBuilder,
     SetVariable, TableSource, TableType, UNNAMED_TABLE,
 };
-use crate::optimizer::optimizer::{OptimizerConfig, OptimizerRule};
+use crate::optimizer::{OptimizerContext, OptimizerRule};
 use datafusion_sql::{ResolvedTableReference, TableReference};
 
 use crate::physical_optimizer::coalesce_batches::CoalesceBatches;
@@ -1557,14 +1557,6 @@ impl SessionState {
                 .register_catalog(config.default_catalog.clone(), default_catalog);
         }
 
-        let optimizer_config = OptimizerConfig::new().filter_null_keys(
-            config
-                .config_options
-                .read()
-                .get_bool(OPT_FILTER_NULL_JOIN_KEYS)
-                .unwrap_or_default(),
-        );
-
         let mut physical_optimizers: Vec<Arc<dyn PhysicalOptimizerRule + Sync + Send>> = vec![
             Arc::new(AggregateStatistics::new()),
             Arc::new(JoinSelection::new()),
@@ -1593,7 +1585,7 @@ impl SessionState {
 
         SessionState {
             session_id,
-            optimizer: Optimizer::new(&optimizer_config),
+            optimizer: Optimizer::new(),
             physical_optimizers,
             query_planner: Arc::new(DefaultQueryPlanner {}),
             catalog_list,
@@ -1741,32 +1733,37 @@ impl SessionState {
 
     /// Optimizes the logical plan by applying optimizer rules.
     pub fn optimize(&self, plan: &LogicalPlan) -> Result<LogicalPlan> {
-        let mut optimizer_config = OptimizerConfig::new()
-            .with_skip_failing_rules(
-                self.config
-                    .config_options
-                    .read()
-                    .get_bool(OPT_OPTIMIZER_SKIP_FAILED_RULES)
-                    .unwrap_or_default(),
-            )
-            .with_max_passes(
-                self.config
-                    .config_options
-                    .read()
-                    .get_u64(OPT_OPTIMIZER_MAX_PASSES)
-                    .unwrap_or_default() as u8,
-            )
-            .with_query_execution_start_time(
-                self.execution_props.query_execution_start_time,
-            );
+        // TODO: Implement OptimizerContext directly on DataFrame (#4631) (#4626)
+        let config = {
+            let config_options = self.config.config_options.read();
+            OptimizerContext::new()
+                .with_skip_failing_rules(
+                    config_options
+                        .get_bool(OPT_OPTIMIZER_SKIP_FAILED_RULES)
+                        .unwrap_or_default(),
+                )
+                .with_max_passes(
+                    config_options
+                        .get_u64(OPT_OPTIMIZER_MAX_PASSES)
+                        .unwrap_or_default() as u8,
+                )
+                .with_query_execution_start_time(
+                    self.execution_props.query_execution_start_time,
+                )
+                .filter_null_keys(
+                    config_options
+                        .get_bool(OPT_FILTER_NULL_JOIN_KEYS)
+                        .unwrap_or_default(),
+                )
+        };
 
         if let LogicalPlan::Explain(e) = plan {
             let mut stringified_plans = e.stringified_plans.clone();
 
             // optimize the child plan, capturing the output of each optimizer
             let plan = self.optimizer.optimize(
                 e.plan.as_ref(),
-                &mut optimizer_config,
+                &config,
                 |optimized_plan, optimizer| {
                     let optimizer_name = optimizer.name().to_string();
                     let plan_type = PlanType::OptimizedLogicalPlan { optimizer_name };
@@ -1781,8 +1778,7 @@ impl SessionState {
                 schema: e.schema.clone(),
             }))
         } else {
-            self.optimizer
-                .optimize(plan, &mut optimizer_config, |_, _| {})
+            self.optimizer.optimize(plan, &config, |_, _| {})
         }
     }
 

datafusion/core/tests/user_defined_plan.rs

@@ -285,7 +285,7 @@ impl OptimizerRule for TopKOptimizerRule {
     fn try_optimize(
         &self,
         plan: &LogicalPlan,
-        optimizer_config: &mut OptimizerConfig,
+        config: &dyn OptimizerConfig,
     ) -> Result<Option<LogicalPlan>> {
         // 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
@@ -308,7 +308,7 @@ impl OptimizerRule for TopKOptimizerRule {
                         node: Arc::new(TopKPlanNode {
                             k: *fetch,
                             input: self
-                                .try_optimize(input.as_ref(), optimizer_config)?
+                                .try_optimize(input.as_ref(), config)?
                                 .unwrap_or_else(|| input.as_ref().clone()),
                             expr: expr[0].clone(),
                         }),
@@ -319,7 +319,7 @@ impl OptimizerRule for TopKOptimizerRule {
 
         // If we didn't find the Limit/Sort combination, recurse as
         // normal and build the result.
-        Ok(Some(optimize_children(self, plan, optimizer_config)?))
+        Ok(Some(optimize_children(self, plan, config)?))
     }
 
     fn name(&self) -> &str {

datafusion/optimizer/README.md

@@ -42,9 +42,9 @@ and applying it to a logical plan to produce an optimized logical plan.
 // The `datafusion` crate provides a DataFrame API that can create a LogicalPlan
 let logical_plan = ...
 
-let mut config = OptimizerConfig::default();
+let mut config = OptimizerContext::default();
 let optimizer = Optimizer::new(&config);
-let optimized_plan = optimizer.optimize(&logical_plan, &mut config, observe)?;
+let optimized_plan = optimizer.optimize(&logical_plan, &config, observe)?;
 
 fn observe(plan: &LogicalPlan, rule: &dyn OptimizerRule) {
     println!(
@@ -82,7 +82,7 @@ pub trait OptimizerRule {
     fn optimize(
         &self,
         plan: &LogicalPlan,
-        optimizer_config: &mut OptimizerConfig,
+        config: &dyn OptimizerConfig,
     ) -> Result<LogicalPlan>;
 
     /// A human readable name for this optimizer rule

datafusion/optimizer/src/common_subexpr_eliminate.rs

@@ -17,8 +17,11 @@
 
 //! Eliminate common sub-expression.
 
-use crate::{utils, OptimizerConfig, OptimizerRule};
+use std::collections::{BTreeSet, HashMap};
+use std::sync::Arc;
+
 use arrow::datatypes::DataType;
+
 use datafusion_common::{DFField, DFSchema, DFSchemaRef, DataFusionError, Result};
 use datafusion_expr::{
     col,
@@ -27,8 +30,8 @@ use datafusion_expr::{
     logical_plan::{Aggregate, Filter, LogicalPlan, Projection, Sort, Window},
     Expr, ExprSchemable,
 };
-use std::collections::{BTreeSet, HashMap};
-use std::sync::Arc;
+
+use crate::{utils, OptimizerConfig, OptimizerRule};
 
 /// A map from expression's identifier to tuple including
 /// - the expression itself (cloned)
@@ -60,7 +63,7 @@ impl CommonSubexprEliminate {
         arrays_list: &[&[Vec<(usize, String)>]],
         input: &LogicalPlan,
         expr_set: &mut ExprSet,
-        optimizer_config: &mut OptimizerConfig,
+        config: &dyn OptimizerConfig,
     ) -> Result<(Vec<Vec<Expr>>, LogicalPlan)> {
         let mut affected_id = BTreeSet::<Identifier>::new();
 
@@ -80,7 +83,7 @@ impl CommonSubexprEliminate {
             .collect::<Result<Vec<_>>>()?;
 
         let mut new_input = self
-            .try_optimize(input, optimizer_config)?
+            .try_optimize(input, config)?
             .unwrap_or_else(|| input.clone());
         if !affected_id.is_empty() {
             new_input = build_project_plan(new_input, affected_id, expr_set)?;
@@ -94,7 +97,7 @@ impl OptimizerRule for CommonSubexprEliminate {
     fn try_optimize(
         &self,
         plan: &LogicalPlan,
-        optimizer_config: &mut OptimizerConfig,
+        config: &dyn OptimizerConfig,
     ) -> Result<Option<LogicalPlan>> {
         let mut expr_set = ExprSet::new();
 
@@ -107,13 +110,8 @@ impl OptimizerRule for CommonSubexprEliminate {
                 let input_schema = Arc::clone(input.schema());
                 let arrays = to_arrays(expr, input_schema, &mut expr_set)?;
 
-                let (mut new_expr, new_input) = self.rewrite_expr(
-                    &[expr],
-                    &[&arrays],
-                    input,
-                    &mut expr_set,
-                    optimizer_config,
-                )?;
+                let (mut new_expr, new_input) =
+                    self.rewrite_expr(&[expr], &[&arrays], input, &mut expr_set, config)?;
 
                 Ok(Some(LogicalPlan::Projection(
                     Projection::try_new_with_schema(
@@ -140,7 +138,7 @@ impl OptimizerRule for CommonSubexprEliminate {
                     &[&[id_array]],
                     filter.input(),
                     &mut expr_set,
-                    optimizer_config,
+                    config,
                 )?;
 
                 if let Some(predicate) = pop_expr(&mut new_expr)?.pop() {
@@ -167,7 +165,7 @@ impl OptimizerRule for CommonSubexprEliminate {
                     &[&arrays],
                     input,
                     &mut expr_set,
-                    optimizer_config,
+                    config,
                 )?;
 
                 Ok(Some(LogicalPlan::Window(Window {
@@ -192,7 +190,7 @@ impl OptimizerRule for CommonSubexprEliminate {
                     &[&group_arrays, &aggr_arrays],
                     input,
                     &mut expr_set,
-                    optimizer_config,
+                    config,
                 )?;
                 // note the reversed pop order.
                 let new_aggr_expr = pop_expr(&mut new_expr)?;
@@ -211,13 +209,8 @@ impl OptimizerRule for CommonSubexprEliminate {
                 let input_schema = Arc::clone(input.schema());
                 let arrays = to_arrays(expr, input_schema, &mut expr_set)?;
 
-                let (mut new_expr, new_input) = self.rewrite_expr(
-                    &[expr],
-                    &[&arrays],
-                    input,
-                    &mut expr_set,
-                    optimizer_config,
-                )?;
+                let (mut new_expr, new_input) =
+                    self.rewrite_expr(&[expr], &[&arrays], input, &mut expr_set, config)?;
 
                 Ok(Some(LogicalPlan::Sort(Sort {
                     expr: pop_expr(&mut new_expr)?,
@@ -249,11 +242,7 @@ impl OptimizerRule for CommonSubexprEliminate {
             | LogicalPlan::Extension(_)
             | LogicalPlan::Prepare(_) => {
                 // apply the optimization to all inputs of the plan
-                Ok(Some(utils::optimize_children(
-                    self,
-                    plan,
-                    optimizer_config,
-                )?))
+                Ok(Some(utils::optimize_children(self, plan, config)?))
             }
         }
     }
@@ -572,20 +561,25 @@ fn replace_common_expr(
 
 #[cfg(test)]
 mod test {
-    use super::*;
-    use crate::test::*;
+    use std::iter;
+
     use arrow::datatypes::{Field, Schema};
+
     use datafusion_expr::logical_plan::{table_scan, JoinType};
     use datafusion_expr::{
         avg, binary_expr, col, lit, logical_plan::builder::LogicalPlanBuilder, sum,
         Operator,
     };
-    use std::iter;
+
+    use crate::optimizer::OptimizerContext;
+    use crate::test::*;
+
+    use super::*;
 
     fn assert_optimized_plan_eq(expected: &str, plan: &LogicalPlan) {
         let optimizer = CommonSubexprEliminate {};
         let optimized_plan = optimizer
-            .try_optimize(plan, &mut OptimizerConfig::new())
+            .try_optimize(plan, &OptimizerContext::new())
             .unwrap()
             .expect("failed to optimize plan");
         let formatted_plan = format!("{:?}", optimized_plan);
@@ -831,7 +825,7 @@ mod test {
             .unwrap();
         let rule = CommonSubexprEliminate {};
         let optimized_plan = rule
-            .try_optimize(&plan, &mut OptimizerConfig::new())
+            .try_optimize(&plan, &OptimizerContext::new())
             .unwrap()
             .unwrap();
 

datafusion/optimizer/src/decorrelate_where_exists.rs

@@ -48,7 +48,7 @@ impl DecorrelateWhereExists {
     fn extract_subquery_exprs(
         &self,
         predicate: &Expr,
-        optimizer_config: &mut OptimizerConfig,
+        config: &dyn OptimizerConfig,
     ) -> Result<(Vec<SubqueryInfo>, Vec<Expr>)> {
         let filters = split_conjunction(predicate);
 
@@ -58,7 +58,7 @@ impl DecorrelateWhereExists {
             match it {
                 Expr::Exists { subquery, negated } => {
                     let subquery = self
-                        .try_optimize(&subquery.subquery, optimizer_config)?
+                        .try_optimize(&subquery.subquery, config)?
                         .map(Arc::new)
                         .unwrap_or_else(|| subquery.subquery.clone());
                     let subquery = Subquery { subquery };
@@ -77,7 +77,7 @@ impl OptimizerRule for DecorrelateWhereExists {
     fn try_optimize(
         &self,
         plan: &LogicalPlan,
-        optimizer_config: &mut OptimizerConfig,
+        config: &dyn OptimizerConfig,
     ) -> Result<Option<LogicalPlan>> {
         match plan {
             LogicalPlan::Filter(filter) => {
@@ -86,11 +86,11 @@ impl OptimizerRule for DecorrelateWhereExists {
 
                 // Apply optimizer rule to current input
                 let optimized_input = self
-                    .try_optimize(filter_input, optimizer_config)?
+                    .try_optimize(filter_input, config)?
                     .unwrap_or_else(|| filter_input.clone());
 
                 let (subqueries, other_exprs) =
-                    self.extract_subquery_exprs(predicate, optimizer_config)?;
+                    self.extract_subquery_exprs(predicate, config)?;
                 let optimized_plan = LogicalPlan::Filter(Filter::try_new(
                     predicate.clone(),
                     Arc::new(optimized_input),
@@ -114,11 +114,7 @@ impl OptimizerRule for DecorrelateWhereExists {
             }
             _ => {
                 // Apply the optimization to all inputs of the plan
-                Ok(Some(utils::optimize_children(
-                    self,
-                    plan,
-                    optimizer_config,
-                )?))
+                Ok(Some(utils::optimize_children(self, plan, config)?))
             }
         }
     }