chore(planner): make optimize function async (#15121)

chore: make optimize async

Author: Dousir9

src/query/service/src/interpreters/interpreter_delete.rs

@@ -386,7 +386,7 @@ pub async fn subquery_filter(
         .with_enable_join_reorder(unsafe { !ctx.get_settings().get_disable_join_reorder()? })
         .with_enable_dphyp(ctx.get_settings().get_enable_dphyp()?);
 
-    s_expr = optimize_query(opt_ctx, s_expr.clone())?;
+    s_expr = optimize_query(opt_ctx, s_expr.clone()).await?;
 
     // Create `input_expr` pipeline and execute it to get `_row_id` data block.
     let select_interpreter = SelectInterpreter::try_create(

src/query/sql/src/planner/binder/ddl/index.rs

@@ -319,7 +319,7 @@ impl Binder {
         let plan = if let Statement::Query(_) = &stmt {
             let select_plan = self.bind_statement(bind_context, &stmt).await?;
             let opt_ctx = OptimizerContext::new(self.ctx.clone(), self.metadata.clone());
-            Ok(optimize(opt_ctx, select_plan)?)
+            Ok(optimize(opt_ctx, select_plan).await?)
         } else {
             Err(ErrorCode::UnsupportedIndex("statement is not query"))
         };

src/query/sql/src/planner/binder/ddl/table.rs

@@ -665,7 +665,7 @@ impl Binder {
                 let select_plan = self.bind_statement(&mut bind_context, &stmt).await?;
                 // Don't enable distributed optimization for `CREATE TABLE ... AS SELECT ...` for now
                 let opt_ctx = OptimizerContext::new(self.ctx.clone(), self.metadata.clone());
-                let optimized_plan = optimize(opt_ctx, select_plan)?;
+                let optimized_plan = optimize(opt_ctx, select_plan).await?;
                 Some(Box::new(optimized_plan))
             } else {
                 None

src/query/sql/src/planner/binder/insert.rs

@@ -181,7 +181,7 @@ impl Binder {
                     }
                 }
 
-                let optimized_plan = optimize(opt_ctx, select_plan)?;
+                let optimized_plan = optimize(opt_ctx, select_plan).await?;
                 Ok(InsertInputSource::SelectPlan(Box::new(optimized_plan)))
             }
         };

src/query/sql/src/planner/binder/replace.rs

@@ -166,7 +166,7 @@ impl Binder {
                 }
                 let opt_ctx = OptimizerContext::new(self.ctx.clone(), self.metadata.clone())
                     .with_enable_distributed_optimization(false);
-                let optimized_plan = optimize(opt_ctx, select_plan)?;
+                let optimized_plan = optimize(opt_ctx, select_plan).await?;
                 Ok(InsertInputSource::SelectPlan(Box::new(optimized_plan)))
             }
         };

src/query/sql/src/planner/dataframe.rs

@@ -35,11 +35,8 @@ use databend_common_exception::Result;
 use databend_common_expression::DataSchemaRef;
 use parking_lot::RwLock;
 
-use crate::optimizer::optimize;
-use crate::optimizer::OptimizerContext;
 use crate::planner::optimizer::s_expr::SExpr;
 use crate::plans::Limit;
-use crate::plans::Plan;
 use crate::BindContext;
 use crate::Binder;
 use crate::Metadata;
@@ -548,20 +545,6 @@ impl Dataframe {
     pub fn get_expr(&self) -> &SExpr {
         &self.s_expr
     }
-
-    pub fn into_plan(self, enable_distributed_optimization: bool) -> Result<Plan> {
-        let plan = Plan::Query {
-            s_expr: Box::new(self.s_expr),
-            metadata: self.binder.metadata.clone(),
-            bind_context: Box::new(self.bind_context),
-            rewrite_kind: None,
-            ignore_result: false,
-            formatted_ast: None,
-        };
-        let opt_ctx = OptimizerContext::new(self.query_ctx.clone(), self.binder.metadata.clone())
-            .with_enable_distributed_optimization(enable_distributed_optimization);
-        optimize(opt_ctx, plan)
-    }
 }
 
 fn parse_cols(schema: DataSchemaRef, columns: &[&str]) -> Result<Vec<SelectTarget>> {

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

@@ -148,7 +148,8 @@ impl<'a> RecursiveOptimizer<'a> {
 }
 
 #[minitrace::trace]
-pub fn optimize(opt_ctx: OptimizerContext, plan: Plan) -> Result<Plan> {
+#[async_backtrace::framed]
+pub async fn optimize(opt_ctx: OptimizerContext, plan: Plan) -> Result<Plan> {
     match plan {
         Plan::Query {
             s_expr,
@@ -158,7 +159,7 @@ pub fn optimize(opt_ctx: OptimizerContext, plan: Plan) -> Result<Plan> {
             formatted_ast,
             ignore_result,
         } => Ok(Plan::Query {
-            s_expr: Box::new(optimize_query(opt_ctx, *s_expr)?),
+            s_expr: Box::new(optimize_query(opt_ctx, *s_expr).await?),
             bind_context,
             metadata,
             rewrite_kind,
@@ -185,7 +186,7 @@ pub fn optimize(opt_ctx: OptimizerContext, plan: Plan) -> Result<Plan> {
             }
             _ => {
                 if config.optimized || !config.logical {
-                    let optimized_plan = optimize(opt_ctx.clone(), *plan)?;
+                    let optimized_plan = Box::pin(optimize(opt_ctx.clone(), *plan)).await?;
                     Ok(Plan::Explain {
                         kind,
                         config,
@@ -197,13 +198,13 @@ pub fn optimize(opt_ctx: OptimizerContext, plan: Plan) -> Result<Plan> {
             }
         },
         Plan::ExplainAnalyze { plan } => Ok(Plan::ExplainAnalyze {
-            plan: Box::new(optimize(opt_ctx, *plan)?),
+            plan: Box::new(Box::pin(optimize(opt_ctx, *plan)).await?),
         }),
         Plan::CopyIntoLocation(CopyIntoLocationPlan { stage, path, from }) => {
             Ok(Plan::CopyIntoLocation(CopyIntoLocationPlan {
                 stage,
                 path,
-                from: Box::new(optimize(opt_ctx, *from)?),
+                from: Box::new(Box::pin(optimize(opt_ctx, *from)).await?),
             }))
         }
         Plan::CopyIntoTable(mut plan) if !plan.no_file_to_copy => {
@@ -218,14 +219,15 @@ pub fn optimize(opt_ctx: OptimizerContext, plan: Plan) -> Result<Plan> {
             );
             Ok(Plan::CopyIntoTable(plan))
         }
-        Plan::MergeInto(plan) => optimize_merge_into(opt_ctx.clone(), plan),
+        Plan::MergeInto(plan) => optimize_merge_into(opt_ctx.clone(), plan).await,
 
         // Pass through statements.
         _ => Ok(plan),
     }
 }
 
-pub fn optimize_query(opt_ctx: OptimizerContext, mut s_expr: SExpr) -> Result<SExpr> {
+#[async_backtrace::framed]
+pub async fn optimize_query(opt_ctx: OptimizerContext, mut s_expr: SExpr) -> Result<SExpr> {
     let enable_distributed_query = opt_ctx.enable_distributed_optimization
         && !contains_local_table_scan(&s_expr, &opt_ctx.metadata);
 
@@ -345,12 +347,13 @@ fn get_optimized_memo(opt_ctx: OptimizerContext, mut s_expr: SExpr) -> Result<Me
     Ok(cascades.memo)
 }
 
-fn optimize_merge_into(opt_ctx: OptimizerContext, plan: Box<MergeInto>) -> Result<Plan> {
+#[async_backtrace::framed]
+async fn optimize_merge_into(opt_ctx: OptimizerContext, plan: Box<MergeInto>) -> Result<Plan> {
     // optimize source :fix issue #13733
     // reason: if there is subquery,windowfunc exprs etc. see
     // src/planner/semantic/lowering.rs `as_raw_expr()`, we will
     // get dummy index. So we need to use optimizer to solve this.
-    let mut right_source = optimize_query(opt_ctx.clone(), plan.input.child(1)?.clone())?;
+    let mut right_source = optimize_query(opt_ctx.clone(), plan.input.child(1)?.clone()).await?;
 
     // if it's not distributed execution, we should reserve
     // exchange to merge source data.

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

@@ -164,7 +164,7 @@ impl Planner {
                     })
                     .with_enable_dphyp(self.ctx.get_settings().get_enable_dphyp()?);
 
-                let optimized_plan = optimize(opt_ctx, plan)?;
+                let optimized_plan = optimize(opt_ctx, plan).await?;
                 Ok((optimized_plan, PlanExtras {
                     metadata,
                     format,