Multiple optimizer passes

datafusion/core/src/config.rs

@@ -47,6 +47,9 @@ pub const OPT_COALESCE_TARGET_BATCH_SIZE: &str =
 pub const OPT_OPTIMIZER_SKIP_FAILED_RULES: &str =
     "datafusion.optimizer.skip_failed_rules";
 
+/// Configuration option "datafusion.optimizer.max_passes"
+pub const OPT_OPTIMIZER_MAX_PASSES: &str = "datafusion.optimizer.max_passes";
+
 /// Configuration option "datafusion.execution.time_zone"
 pub const OPT_TIME_ZONE: &str = "datafusion.execution.time_zone";
 
@@ -201,6 +204,11 @@ impl BuiltInConfigs {
                 rule. When set to false, any rules that produce errors will cause the query to fail.",
                 true
             ),
+         ConfigDefinition::new_u64(
+             OPT_OPTIMIZER_MAX_PASSES,
+             "Number of times that the optimizer will attempt to optimize the plan",
+             5
+         ),
             ConfigDefinition::new_string(
                 OPT_TIME_ZONE,
                 "The session time zone which some function require \

datafusion/core/src/execution/context.rs

@@ -78,7 +78,7 @@ use crate::physical_optimizer::repartition::Repartition;
 
 use crate::config::{
     ConfigOptions, OPT_BATCH_SIZE, OPT_COALESCE_BATCHES, OPT_COALESCE_TARGET_BATCH_SIZE,
-    OPT_FILTER_NULL_JOIN_KEYS, OPT_OPTIMIZER_SKIP_FAILED_RULES,
+    OPT_FILTER_NULL_JOIN_KEYS, OPT_OPTIMIZER_MAX_PASSES, OPT_OPTIMIZER_SKIP_FAILED_RULES,
 };
 use crate::datasource::file_format::file_type::{FileCompressionType, FileType};
 use crate::execution::{runtime_env::RuntimeEnv, FunctionRegistry};
@@ -1592,6 +1592,13 @@ impl SessionState {
                     .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,
             );

datafusion/core/tests/sql/subqueries.rs

@@ -394,9 +394,9 @@ order by cntrycode;"#;
               Projection: AVG(customer.c_acctbal) AS __value, alias=__sq_1
                 Aggregate: groupBy=[[]], aggr=[[AVG(customer.c_acctbal)]]
                   Filter: customer.c_acctbal > Decimal128(Some(0),15,2) AND substr(customer.c_phone, Int64(1), Int64(2)) IN ([Utf8("13"), Utf8("31"), Utf8("23"), Utf8("29"), Utf8("30"), Utf8("18"), Utf8("17")])
-                    TableScan: customer projection=[c_phone, c_acctbal], partial_filters=[CAST(customer.c_acctbal AS Decimal128(30, 15)) > Decimal128(Some(0),30,15), substr(customer.c_phone, Int64(1), Int64(2)) IN ([Utf8("13"), Utf8("31"), Utf8("23"), Utf8("29"), Utf8("30"), Utf8("18"), Utf8("17")])]"#
+                    TableScan: customer projection=[c_phone, c_acctbal], partial_filters=[CAST(customer.c_acctbal AS Decimal128(30, 15)) > Decimal128(Some(0),30,15), substr(customer.c_phone, Int64(1), Int64(2)) IN ([Utf8("13"), Utf8("31"), Utf8("23"), Utf8("29"), Utf8("30"), Utf8("18"), Utf8("17")]), customer.c_acctbal > Decimal128(Some(0),15,2)]"#
         .to_string();
-    assert_eq!(actual, expected);
+    assert_eq!(expected, actual);
 
     // assert data
     let results = execute_to_batches(&ctx, sql).await;

datafusion/optimizer/src/optimizer.rs

@@ -71,6 +71,8 @@ pub struct OptimizerConfig {
     skip_failing_rules: bool,
     /// Specify whether to enable the filter_null_keys rule
     filter_null_keys: bool,
+    /// Maximum number of times to run optimizer against a plan
+    max_passes: u8,
 }
 
 impl OptimizerConfig {
@@ -81,6 +83,7 @@ impl OptimizerConfig {
             next_id: 0, // useful for generating things like unique subquery aliases
             skip_failing_rules: true,
             filter_null_keys: true,
+            max_passes: 5,
         }
     }
 
@@ -107,6 +110,12 @@ impl OptimizerConfig {
         self
     }
 
+    /// Specify how many times to attempt to optimize the plan
+    pub fn with_max_passes(mut self, v: u8) -> Self {
+        self.max_passes = v;
+        self
+    }
+
     /// Generate the next ID needed
     pub fn next_id(&mut self) -> usize {
         self.next_id += 1;
@@ -189,38 +198,53 @@ impl Optimizer {
     where
         F: FnMut(&LogicalPlan, &dyn OptimizerRule),
     {
+        let mut plan_str = format!("{}", plan.display_indent());
         let mut new_plan = plan.clone();
-        log_plan("Optimizer input", plan);
+        let mut i = 0;
+        while i < optimizer_config.max_passes {
+            log_plan(&format!("Optimizer input (pass {})", i), &new_plan);
 
-        for rule in &self.rules {
-            let result = rule.optimize(&new_plan, optimizer_config);
-            match result {
-                Ok(plan) => {
-                    new_plan = plan;
-                    observer(&new_plan, rule.as_ref());
-                    log_plan(rule.name(), &new_plan);
-                }
-                Err(ref e) => {
-                    if optimizer_config.skip_failing_rules {
-                        // Note to future readers: if you see this warning it signals a
-                        // bug in the DataFusion optimizer. Please consider filing a ticket
-                        // https://github.com/apache/arrow-datafusion
-                        warn!(
+            for rule in &self.rules {
+                let result = rule.optimize(&new_plan, optimizer_config);
+                match result {
+                    Ok(plan) => {
+                        new_plan = plan;
+                        observer(&new_plan, rule.as_ref());
+                        log_plan(rule.name(), &new_plan);
+                    }
+                    Err(ref e) => {
+                        if optimizer_config.skip_failing_rules {
+                            // Note to future readers: if you see this warning it signals a
+                            // bug in the DataFusion optimizer. Please consider filing a ticket
+                            // https://github.com/apache/arrow-datafusion
+                            warn!(
                             "Skipping optimizer rule '{}' due to unexpected error: {}",
                             rule.name(),
                             e
                         );
-                    } else {
-                        return Err(DataFusionError::Internal(format!(
-                            "Optimizer rule '{}' failed due to unexpected error: {}",
-                            rule.name(),
-                            e
-                        )));
+                        } else {
+                            return Err(DataFusionError::Internal(format!(
+                                "Optimizer rule '{}' failed due to unexpected error: {}",
+                                rule.name(),
+                                e
+                            )));
+                        }
                     }
                 }
             }
+            log_plan(&format!("Optimized plan (pass {})", i), &new_plan);
+            println!("Optimized plan (pass {}): {}", i, new_plan.display_indent());
+
+            let new_plan_str = format!("{}", new_plan.display_indent());
+            if plan_str == new_plan_str {
+                // plan did not change, so no need to continue trying to optimize
+                println!("optimizer pass {} did not make changes", i);
+                break;
+            }
+            plan_str = new_plan_str;
+            i += 1;
         }
-        log_plan("Optimized plan", &new_plan);
+        log_plan("Final optimized plan", &new_plan);
         Ok(new_plan)
     }
 }