refactor(optimizer): optimize trait and pipeline (#17712)

* refactor(optimizer): optimize trait

* add optimizer trait

* prepare for pipeline

* memo&optimize_query to use pipeline

* impl cascade memo trait

Author: BohuTANG

src/query/service/tests/it/sql/planner/optimizer/agg_index_query_rewrite.rs

@@ -394,7 +394,7 @@ async fn test_query_rewrite_impl(format: &str) -> Result<()> {
 
         let opt_ctx = OptimizerContext::new(ctx.clone(), metadata.clone());
         let result = RecursiveOptimizer::new(opt_ctx.clone(), &[RuleID::TryApplyAggIndex])
-            .optimize(&query)?;
+            .optimize_sync(&query)?;
         let agg_index = find_push_down_index_info(&result)?;
         assert_eq!(
             suite.is_matched,
@@ -449,7 +449,8 @@ async fn plan_sql(
     {
         let s_expr = if optimize {
             let opt_ctx = OptimizerContext::new(ctx.clone(), metadata.clone());
-            RecursiveOptimizer::new(opt_ctx.clone(), &DEFAULT_REWRITE_RULES).optimize(&s_expr)?
+            RecursiveOptimizer::new(opt_ctx.clone(), &DEFAULT_REWRITE_RULES)
+                .optimize_sync(&s_expr)?
         } else {
             *s_expr
         };

src/query/sql/src/planner/binder/bind_mutation/mutation_expression.rs

@@ -313,7 +313,7 @@ impl MutationExpression {
                     let opt_ctx =
                         OptimizerContext::new(binder.ctx.clone(), binder.metadata.clone());
                     let mut rewriter = SubqueryRewriter::new(opt_ctx, None);
-                    let s_expr = rewriter.optimize(&s_expr)?;
+                    let s_expr = rewriter.optimize_sync(&s_expr)?;
 
                     Ok(MutationExpressionBindResult {
                         input: s_expr,

src/query/sql/src/planner/optimizer/cost/cost.rs

@@ -55,7 +55,7 @@ impl AddAssign for Cost {
     }
 }
 
-pub trait CostModel: Send {
+pub trait CostModel: Send + Sync {
     /// Compute cost of given `MExpr`(children are not encapsulated).
     fn compute_cost(&self, memo: &Memo, m_expr: &MExpr) -> Result<Cost>;
 }

src/query/sql/src/planner/optimizer/mod.rs

@@ -16,13 +16,14 @@ mod cost;
 pub mod ir;
 #[allow(clippy::module_inception)]
 mod optimizer;
+mod optimizer_api;
 mod optimizer_context;
 pub mod optimizers;
+pub mod pipeline;
 mod statistics;
-mod util;
 
 pub use optimizer::optimize;
 pub use optimizer::optimize_query;
+pub use optimizer_api::Optimizer;
 pub use optimizer_context::OptimizerContext;
 pub use optimizers::rule::agg_index;
-pub use util::contains_local_table_scan;

src/query/sql/src/planner/optimizer/optimizer.rs

@@ -37,9 +37,8 @@ use crate::optimizer::optimizers::rule::RuleID;
 use crate::optimizer::optimizers::rule::DEFAULT_REWRITE_RULES;
 use crate::optimizer::optimizers::CascadesOptimizer;
 use crate::optimizer::optimizers::DPhpy;
+use crate::optimizer::pipeline::OptimizerPipeline;
 use crate::optimizer::statistics::CollectStatisticsOptimizer;
-use crate::optimizer::util::contains_local_table_scan;
-use crate::optimizer::util::contains_warehouse_table_scan;
 use crate::optimizer::OptimizerContext;
 use crate::plans::ConstantTableScan;
 use crate::plans::CopyIntoLocationPlan;
@@ -92,11 +91,8 @@ pub async fn optimize(opt_ctx: Arc<OptimizerContext>, plan: Plan) -> Result<Plan
                     ));
                 };
 
-                let mut s_expr = s_expr;
-                if s_expr.contain_subquery() {
-                    s_expr =
-                        Box::new(SubqueryRewriter::new(opt_ctx.clone(), None).optimize(&s_expr)?);
-                }
+                let s_expr =
+                    Box::new(SubqueryRewriter::new(opt_ctx.clone(), None).optimize_sync(&s_expr)?);
                 Ok(Plan::Explain {
                     kind,
                     config,
@@ -241,129 +237,86 @@ pub async fn optimize(opt_ctx: Arc<OptimizerContext>, plan: Plan) -> Result<Plan
     }
 }
 
-pub async fn optimize_query(opt_ctx: Arc<OptimizerContext>, mut s_expr: SExpr) -> Result<SExpr> {
-    let metadata = opt_ctx.get_metadata();
-
-    // 1. Configure distributed optimization based on table types
-    if contains_local_table_scan(&s_expr, &metadata) {
-        opt_ctx.set_enable_distributed_optimization(false);
-        info!("Disable distributed optimization due to local table scan.");
-    } else if contains_warehouse_table_scan(&s_expr, &metadata) {
-        let warehouse = opt_ctx.get_table_ctx().get_warehouse_cluster().await?;
-
-        if !warehouse.is_empty() {
-            opt_ctx.set_enable_distributed_optimization(true);
-            info!("Enable distributed optimization due to warehouse table scan.");
-        }
-    }
-
-    // 2. Eliminate subqueries by rewriting them into more efficient forms
-    s_expr = SubqueryRewriter::new(opt_ctx.clone(), None).optimize(&s_expr)?;
-
-    // 3. Apply statistics aggregation to gather and propagate statistics
-    s_expr = RuleStatsAggregateOptimizer::new(opt_ctx.clone())
-        .optimize(&s_expr)
-        .await?;
-
-    // 4. Collect statistics for SExpr nodes to support cost estimation
-    s_expr = CollectStatisticsOptimizer::new(opt_ctx.clone())
-        .optimize(&s_expr)
-        .await?;
-
-    // 5. Normalize aggregate, it should be executed before RuleSplitAggregate.
-    s_expr = RuleNormalizeAggregateOptimizer::new().optimize(&s_expr)?;
-
-    // 6. Pull up and infer filter.
-    s_expr = PullUpFilterOptimizer::new(opt_ctx.clone()).optimize(&s_expr)?;
-
-    // 7. Run default rewrite rules
-    s_expr = RecursiveOptimizer::new(opt_ctx.clone(), &DEFAULT_REWRITE_RULES).optimize(&s_expr)?;
-
-    // 8. Run post rewrite rules
-    s_expr =
-        RecursiveOptimizer::new(opt_ctx.clone(), &[RuleID::SplitAggregate]).optimize(&s_expr)?;
-
-    // 9. Apply DPhyp algorithm for cost-based join reordering
-    s_expr = DPhpy::new(opt_ctx.clone()).optimize(&s_expr).await?;
-
-    // 10. After join reorder, Convert some single join to inner join.
-    s_expr = SingleToInnerOptimizer::new().optimize(&s_expr)?;
-
-    // 11. Deduplicate join conditions.
-    s_expr = DeduplicateJoinConditionOptimizer::new().optimize(&s_expr)?;
-
-    // 12. Apply join commutativity to further optimize join ordering
-    if opt_ctx.get_enable_join_reorder() {
-        s_expr = RecursiveOptimizer::new(opt_ctx.clone(), [RuleID::CommuteJoin].as_slice())
-            .optimize(&s_expr)?;
-    }
-
-    // 13. Cascades optimizer may fail due to timeout, fallback to heuristic optimizer in this case.
-    s_expr = CascadesOptimizer::new(opt_ctx.clone())?.optimize(s_expr)?;
-
-    // 14. Eliminate unnecessary scalar calculations to clean up the final plan
-    if !opt_ctx.get_planning_agg_index() {
-        s_expr = RecursiveOptimizer::new(opt_ctx.clone(), [RuleID::EliminateEvalScalar].as_slice())
-            .optimize(&s_expr)?;
-    }
+pub async fn optimize_query(opt_ctx: Arc<OptimizerContext>, s_expr: SExpr) -> Result<SExpr> {
+    let mut pipeline = OptimizerPipeline::new(opt_ctx.clone(), s_expr.clone())
+        .await?
+        // 2. Eliminate subqueries by rewriting them into more efficient form
+        .add(SubqueryRewriter::new(opt_ctx.clone(), None))
+        // 3. Apply statistics aggregation to gather and propagate statistics
+        .add(RuleStatsAggregateOptimizer::new(opt_ctx.clone()))
+        // 4. Collect statistics for SExpr nodes to support cost estimation
+        .add(CollectStatisticsOptimizer::new(opt_ctx.clone()))
+        // 5. Normalize aggregate, it should be executed before RuleSplitAggregate.
+        .add(RuleNormalizeAggregateOptimizer::new())
+        // 6. Pull up and infer filter.
+        .add(PullUpFilterOptimizer::new(opt_ctx.clone()))
+        // 7. Run default rewrite rules
+        .add(RecursiveOptimizer::new(
+            opt_ctx.clone(),
+            &DEFAULT_REWRITE_RULES,
+        ))
+        // 8. Run post rewrite rules
+        .add(RecursiveOptimizer::new(opt_ctx.clone(), &[
+            RuleID::SplitAggregate,
+        ]))
+        // 9. Apply DPhyp algorithm for cost-based join reordering
+        .add(DPhpy::new(opt_ctx.clone()))
+        // 10. After join reorder, Convert some single join to inner join.
+        .add(SingleToInnerOptimizer::new())
+        // 11. Deduplicate join conditions.
+        .add(DeduplicateJoinConditionOptimizer::new())
+        // 12. Apply join commutativity to further optimize join ordering
+        .add_if(
+            opt_ctx.get_enable_join_reorder(),
+            RecursiveOptimizer::new(opt_ctx.clone(), [RuleID::CommuteJoin].as_slice()),
+        )
+        // 13. Cascades optimizer may fail due to timeout, fallback to heuristic optimizer in this case.
+        .add(CascadesOptimizer::new(opt_ctx.clone())?)
+        // 14. Eliminate unnecessary scalar calculations to clean up the final plan
+        .add_if(
+            !opt_ctx.get_planning_agg_index(),
+            RecursiveOptimizer::new(opt_ctx.clone(), [RuleID::EliminateEvalScalar].as_slice()),
+        );
+
+    // 15. Execute the pipeline
+    let s_expr = pipeline.execute().await?;
 
     Ok(s_expr)
 }
 
-// TODO(leiysky): reuse the optimization logic with `optimize_query`
-async fn get_optimized_memo(opt_ctx: Arc<OptimizerContext>, mut s_expr: SExpr) -> Result<Memo> {
-    let metadata = opt_ctx.get_metadata();
-    let _table_ctx = opt_ctx.get_table_ctx();
-
-    if contains_local_table_scan(&s_expr, &metadata) {
-        opt_ctx.set_enable_distributed_optimization(false);
-        info!("Disable distributed optimization due to local table scan.");
-    } else if contains_warehouse_table_scan(&s_expr, &metadata) {
-        let warehouse = opt_ctx.get_table_ctx().get_warehouse_cluster().await?;
-
-        if !warehouse.is_empty() {
-            opt_ctx.set_enable_distributed_optimization(true);
-            info!("Enable distributed optimization due to warehouse table scan.");
-        }
-    }
-
-    // Decorrelate subqueries, after this step, there should be no subquery in the expression.
-    if s_expr.contain_subquery() {
-        s_expr = SubqueryRewriter::new(opt_ctx.clone(), None).optimize(&s_expr)?;
-    }
-
-    s_expr = RuleStatsAggregateOptimizer::new(opt_ctx.clone())
-        .optimize(&s_expr)
-        .await?;
-
-    // Collect statistics for each leaf node in SExpr.
-    s_expr = CollectStatisticsOptimizer::new(opt_ctx.clone())
-        .optimize(&s_expr)
-        .await?;
-
-    // Pull up and infer filter.
-    s_expr = PullUpFilterOptimizer::new(opt_ctx.clone()).optimize(&s_expr)?;
-    // Run default rewrite rules
-    s_expr = RecursiveOptimizer::new(opt_ctx.clone(), &DEFAULT_REWRITE_RULES).optimize(&s_expr)?;
-    // Run post rewrite rules
-    s_expr =
-        RecursiveOptimizer::new(opt_ctx.clone(), &[RuleID::SplitAggregate]).optimize(&s_expr)?;
-
-    // Cost based optimization
-    if opt_ctx.get_enable_dphyp() && opt_ctx.get_enable_join_reorder() {
-        s_expr = DPhpy::new(opt_ctx.clone()).optimize(&s_expr).await?;
-    }
-    let mut cascades = CascadesOptimizer::new(opt_ctx.clone())?;
-    cascades.optimize(s_expr)?;
-
-    Ok(cascades.memo)
+async fn get_optimized_memo(opt_ctx: Arc<OptimizerContext>, s_expr: SExpr) -> Result<Memo> {
+    let mut pipeline = OptimizerPipeline::new(opt_ctx.clone(), s_expr.clone())
+        .await?
+        // Decorrelate subqueries, after this step, there should be no subquery in the expression.
+        .add(SubqueryRewriter::new(opt_ctx.clone(), None))
+        .add(RuleStatsAggregateOptimizer::new(opt_ctx.clone()))
+        // Collect statistics for each leaf node in SExpr.
+        .add(CollectStatisticsOptimizer::new(opt_ctx.clone()))
+        // Pull up and infer filter.
+        .add(PullUpFilterOptimizer::new(opt_ctx.clone()))
+        // Run default rewrite rules
+        .add(RecursiveOptimizer::new(
+            opt_ctx.clone(),
+            &DEFAULT_REWRITE_RULES,
+        ))
+        // Run post rewrite rules
+        .add(RecursiveOptimizer::new(opt_ctx.clone(), &[
+            RuleID::SplitAggregate,
+        ]))
+        // Cost based optimization
+        .add(DPhpy::new(opt_ctx.clone()))
+        .add(CascadesOptimizer::new(opt_ctx.clone())?);
+
+    let _s_expr = pipeline.execute().await?;
+
+    Ok(pipeline.memo())
 }
 
 async fn optimize_mutation(opt_ctx: Arc<OptimizerContext>, s_expr: SExpr) -> Result<Plan> {
     // Optimize the input plan.
     let mut input_s_expr = optimize_query(opt_ctx.clone(), s_expr.child(0)?.clone()).await?;
     input_s_expr = RecursiveOptimizer::new(opt_ctx.clone(), &[RuleID::MergeFilterIntoMutation])
-        .optimize(&input_s_expr)?;
+        .optimize_sync(&input_s_expr)?;
 
     // For distributed query optimization, we need to remove the Exchange operator at the top of the plan.
     if let &RelOperator::Exchange(_) = input_s_expr.plan() {

src/query/sql/src/planner/optimizer/optimizer_api.rs

@@ -0,0 +1,35 @@
+// Copyright 2021 Datafuse Labs
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use databend_common_exception::Result;
+
+use crate::optimizer::ir::Memo;
+use crate::optimizer::ir::SExpr;
+
+/// Trait defining the interface for query optimizers.
+
+#[async_trait::async_trait]
+pub trait Optimizer: Send + Sync {
+    /// Optimize the given expression and return the optimized version.
+    async fn optimize(&mut self, expr: &SExpr) -> Result<SExpr>;
+
+    /// Returns a unique identifier for this optimizer.
+    fn name(&self) -> &'static str;
+
+    /// Get the memo if this optimizer maintains one.
+    /// Default implementation returns None for optimizers that don't use a memo.
+    fn memo(&self) -> Option<&Memo> {
+        None
+    }
+}