Optimizer: avoid every rule must recursive children in optimizer

datafusion/optimizer/src/eliminate_limit.rs

@@ -20,7 +20,9 @@
 //! on a plan with an empty relation.
 //! This rule also removes OFFSET 0 from the [LogicalPlan]
 //! This saves time in planning and executing the query.
-use crate::{utils, OptimizerConfig, OptimizerRule};
+use crate::optimizer::ApplyOrder;
+use crate::optimizer::ApplyOrder::BottomUp;
+use crate::{OptimizerConfig, OptimizerRule};
 use datafusion_common::Result;
 use datafusion_expr::logical_plan::{EmptyRelation, LogicalPlan};
 
@@ -40,45 +42,49 @@ impl OptimizerRule for EliminateLimit {
     fn try_optimize(
         &self,
         plan: &LogicalPlan,
-        optimizer_config: &mut OptimizerConfig,
+        _optimizer_config: &mut OptimizerConfig,
     ) -> Result<Option<LogicalPlan>> {
-        if let LogicalPlan::Limit(limit) = plan {
-            match limit.fetch {
-                Some(fetch) => {
-                    if fetch == 0 {
-                        return Ok(Some(LogicalPlan::EmptyRelation(EmptyRelation {
-                            produce_one_row: false,
-                            schema: limit.input.schema().clone(),
-                        })));
-                    }
+        let limit = match plan {
+            LogicalPlan::Limit(limit) => limit,
+            _ => return Ok(None),
+        };
+
+        match limit.fetch {
+            Some(fetch) => {
+                if fetch == 0 {
+                    return Ok(Some(LogicalPlan::EmptyRelation(EmptyRelation {
+                        produce_one_row: false,
+                        schema: limit.input.schema().clone(),
+                    })));
                 }
-                None => {
-                    if limit.skip == 0 {
-                        let input = &*limit.input;
-                        return Ok(Some(utils::optimize_children(
-                            self,
-                            input,
-                            optimizer_config,
-                        )?));
-                    }
+            }
+            None => {
+                if limit.skip == 0 {
+                    let input = limit.input.as_ref();
+                    // input also can be Limit, so we should apply again.
+                    return Ok(Some(
+                        self.try_optimize(input, _optimizer_config)?
+                            .unwrap_or_else(|| input.clone()),
+                    ));
                 }
             }
         }
-        Ok(Some(utils::optimize_children(
-            self,
-            plan,
-            optimizer_config,
-        )?))
+        Ok(None)
     }
 
     fn name(&self) -> &str {
         "eliminate_limit"
     }
+
+    fn apply_order(&self) -> Option<ApplyOrder> {
+        Some(BottomUp)
+    }
 }
 
 #[cfg(test)]
 mod tests {
     use super::*;
+    use crate::optimizer::Optimizer;
     use crate::push_down_limit::PushDownLimit;
     use crate::test::*;
     use datafusion_common::Column;
@@ -87,12 +93,18 @@ mod tests {
         logical_plan::{builder::LogicalPlanBuilder, JoinType},
         sum,
     };
+    use std::sync::Arc;
 
     fn assert_optimized_plan_eq(plan: &LogicalPlan, expected: &str) -> Result<()> {
-        let optimized_plan = EliminateLimit::new()
-            .try_optimize(plan, &mut OptimizerConfig::new())
-            .unwrap()
-            .expect("failed to optimize plan");
+        let optimizer = Optimizer::with_rules(vec![Arc::new(EliminateLimit::new())]);
+        let optimized_plan = optimizer
+            .optimize_recursively(
+                optimizer.rules.get(0).unwrap(),
+                plan,
+                &mut OptimizerConfig::new(),
+            )?
+            .unwrap_or_else(|| plan.clone());
+
         let formatted_plan = format!("{:?}", optimized_plan);
         assert_eq!(formatted_plan, expected);
         assert_eq!(plan.schema(), optimized_plan.schema());
@@ -103,13 +115,14 @@ mod tests {
         plan: &LogicalPlan,
         expected: &str,
     ) -> Result<()> {
-        let optimized_plan = PushDownLimit::new()
-            .try_optimize(plan, &mut OptimizerConfig::new())
-            .unwrap()
-            .expect("failed to optimize plan");
-        let optimized_plan = EliminateLimit::new()
-            .try_optimize(&optimized_plan, &mut OptimizerConfig::new())
-            .unwrap()
+        fn observe(_plan: &LogicalPlan, _rule: &dyn OptimizerRule) {}
+        let mut config = OptimizerConfig::new().with_max_passes(1);
+        let optimizer = Optimizer::with_rules(vec![
+            Arc::new(PushDownLimit::new()),
+            Arc::new(EliminateLimit::new()),
+        ]);
+        let optimized_plan = optimizer
+            .optimize(plan, &mut config, observe)
             .expect("failed to optimize plan");
         let formatted_plan = format!("{:?}", optimized_plan);
         assert_eq!(formatted_plan, expected);

datafusion/optimizer/src/optimizer.rs

@@ -59,6 +59,13 @@ pub trait OptimizerRule {
 
     /// A human readable name for this optimizer rule
     fn name(&self) -> &str;
+
+    /// How should the rule be applied by the optimizer? See comments on [`ApplyOrder`] for details.
+    ///
+    /// If a rule use default None, its should traverse recursively plan inside itself
+    fn apply_order(&self) -> Option<ApplyOrder> {
+        None
+    }
 }
 
 /// Options to control the DataFusion Optimizer.
@@ -147,6 +154,44 @@ 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.
+///
+/// Notice: **sometime** result after optimize still can be optimized, we need apply again.
+///
+/// Usage Example: Merge Limit (subtree pattern is: Limit-Limit)
+/// ```rust
+/// use datafusion_expr::{Limit, LogicalPlan, LogicalPlanBuilder};
+/// use datafusion_common::Result;
+/// fn merge_limit(parent: &Limit, child: &Limit) -> LogicalPlan {
+///     // just for run
+///     return parent.input.as_ref().clone();
+/// }
+/// fn try_optimize(plan: &LogicalPlan) -> Result<Option<LogicalPlan>> {
+///     match plan {
+///         LogicalPlan::Limit(limit) => match limit.input.as_ref() {
+///             LogicalPlan::Limit(child_limit) => {
+///                 // merge limit ...
+///                 let optimized_plan = merge_limit(limit, child_limit);
+///                 // due to optimized_plan may be optimized again,
+///                 // for example: plan is Limit-Limit-Limit
+///                 Ok(Some(
+///                     try_optimize(&optimized_plan)?
+///                         .unwrap_or_else(|| optimized_plan.clone()),
+///                 ))
+///             }
+///             _ => Ok(None),
+///         },
+///         _ => Ok(None),
+///     }
+/// }
+/// ```
+pub enum ApplyOrder {
+    TopDown,
+    BottomUp,
+}
+
 impl Optimizer {
     /// Create a new optimizer using the recommended list of rules
     pub fn new(config: &OptimizerConfig) -> Self {
@@ -213,7 +258,7 @@ impl Optimizer {
             log_plan(&format!("Optimizer input (pass {})", i), &new_plan);
 
             for rule in &self.rules {
-                let result = rule.try_optimize(&new_plan, optimizer_config);
+                let result = self.optimize_recursively(rule, &new_plan, optimizer_config);
                 match result {
                     Ok(Some(plan)) => {
                         if !plan.schema().equivalent_names_and_types(new_plan.schema()) {
@@ -274,6 +319,80 @@ impl Optimizer {
         debug!("Optimizer took {} ms", start_time.elapsed().as_millis());
         Ok(new_plan)
     }
+
+    fn optimize_node(
+        &self,
+        rule: &Arc<dyn OptimizerRule + Send + Sync>,
+        plan: &LogicalPlan,
+        optimizer_config: &mut OptimizerConfig,
+    ) -> Result<Option<LogicalPlan>> {
+        // TODO: future feature: We can do Batch optimize
+        rule.try_optimize(plan, optimizer_config)
+    }
+
+    fn optimize_inputs(
+        &self,
+        rule: &Arc<dyn OptimizerRule + Send + Sync>,
+        plan: &LogicalPlan,
+        optimizer_config: &mut OptimizerConfig,
+    ) -> Result<Option<LogicalPlan>> {
+        let inputs = plan.inputs();
+        let result = inputs
+            .iter()
+            .map(|sub_plan| self.optimize_recursively(rule, sub_plan, optimizer_config))
+            .collect::<Result<Vec<_>>>()?;
+        if result.is_empty() || result.iter().all(|o| o.is_none()) {
+            return Ok(None);
+        }
+
+        let new_inputs = result
+            .into_iter()
+            .enumerate()
+            .map(|(i, o)| match o {
+                Some(plan) => plan,
+                None => (*(inputs.get(i).unwrap())).clone(),
+            })
+            .collect::<Vec<_>>();
+
+        Ok(Some(plan.with_new_inputs(new_inputs.as_slice())?))
+    }
+
+    /// Use a rule to optimize the whole plan.
+    /// If the rule with `ApplyOrder`, we don't need to recursively handle children in rule.
+    pub fn optimize_recursively(
+        &self,
+        rule: &Arc<dyn OptimizerRule + Send + Sync>,
+        plan: &LogicalPlan,
+        optimizer_config: &mut OptimizerConfig,
+    ) -> Result<Option<LogicalPlan>> {
+        match rule.apply_order() {
+            Some(order) => match order {
+                ApplyOrder::TopDown => {
+                    let optimize_self_opt =
+                        self.optimize_node(rule, plan, optimizer_config)?;
+                    let optimize_inputs_opt = match &optimize_self_opt {
+                        Some(optimized_plan) => {
+                            self.optimize_inputs(rule, optimized_plan, optimizer_config)?
+                        }
+                        _ => self.optimize_inputs(rule, plan, optimizer_config)?,
+                    };
+                    Ok(optimize_inputs_opt.or(optimize_self_opt))
+                }
+                ApplyOrder::BottomUp => {
+                    let optimize_inputs_opt =
+                        self.optimize_inputs(rule, plan, optimizer_config)?;
+                    let optimize_self_opt = match &optimize_inputs_opt {
+                        Some(optimized_plan) => {
+                            self.optimize_node(rule, optimized_plan, optimizer_config)?
+                        }
+                        _ => self.optimize_node(rule, plan, optimizer_config)?,
+                    };
+                    Ok(optimize_self_opt.or(optimize_inputs_opt))
+                }
+            },
+            _ => rule.try_optimize(plan, optimizer_config),
+        }
+    }
 }
 
 /// Log the plan in debug/tracing mode after some part of the optimizer runs