Merge pull request #7527 from RinChanNOWWW/distributed_insert

feat(query): Support distributed insert select in new planner.

Author: mergify[bot]

src/query/catalog/src/table.rs

@@ -140,7 +140,12 @@ pub trait Table: Sync + Send {
         )))
     }
 
-    fn append2(&self, _: Arc<dyn TableContext>, _: &mut Pipeline) -> Result<()> {
+    fn append2(
+        &self,
+        _: Arc<dyn TableContext>,
+        _: &mut Pipeline,
+        _need_output: bool,
+    ) -> Result<()> {
         Err(ErrorCode::UnImplement(format!(
             "append2 operation for table {} is not implemented, table engine is {}",
             self.name(),

src/query/pipeline/sinks/src/processors/sinks/context_sink.rs

@@ -34,7 +34,7 @@ impl ContextSink {
 }
 
 impl Sink for ContextSink {
-    const NAME: &'static str = "ContextSink ";
+    const NAME: &'static str = "ContextSink";
 
     fn consume(&mut self, block: DataBlock) -> Result<()> {
         self.ctx.push_precommit_block(block);

src/query/service/src/interpreters/interpreter_common.rs

@@ -25,6 +25,7 @@ use common_exception::Result;
 use common_meta_types::GrantObject;
 use common_meta_types::StageFile;
 use common_meta_types::UserStageInfo;
+use common_pipeline_core::Pipeline;
 use futures::TryStreamExt;
 use regex::Regex;
 use tracing::debug;
@@ -39,28 +40,41 @@ use crate::sessions::TableContext;
 use crate::storages::stage::StageSourceHelper;
 use crate::storages::Table;
 
+pub fn fill_missing_columns(
+    ctx: Arc<QueryContext>,
+    source_schema: &DataSchemaRef,
+    target_schema: &DataSchemaRef,
+    pipeline: &mut Pipeline,
+) -> Result<()> {
+    let need_fill_missing_columns = target_schema != source_schema;
+    if need_fill_missing_columns {
+        pipeline.add_transform(|transform_input_port, transform_output_port| {
+            TransformAddOn::try_create(
+                transform_input_port,
+                transform_output_port,
+                source_schema.clone(),
+                target_schema.clone(),
+                ctx.clone(),
+            )
+        })?;
+    }
+    Ok(())
+}
+
 pub fn append2table(
     ctx: Arc<QueryContext>,
     table: Arc<dyn Table>,
     source_schema: DataSchemaRef,
     mut build_res: PipelineBuildResult,
 ) -> Result<()> {
-    let need_fill_missing_columns = table.schema() != source_schema;
-    if need_fill_missing_columns {
-        build_res
-            .main_pipeline
-            .add_transform(|transform_input_port, transform_output_port| {
-                TransformAddOn::try_create(
-                    transform_input_port,
-                    transform_output_port,
-                    source_schema.clone(),
-                    table.schema(),
-                    ctx.clone(),
-                )
-            })?;
-    }
-
-    table.append2(ctx.clone(), &mut build_res.main_pipeline)?;
+    fill_missing_columns(
+        ctx.clone(),
+        &source_schema,
+        &table.schema(),
+        &mut build_res.main_pipeline,
+    )?;
+
+    table.append2(ctx.clone(), &mut build_res.main_pipeline, false)?;
     let query_need_abort = ctx.query_need_abort();
     let executor_settings = ExecutorSettings::try_create(&ctx.get_settings())?;
     build_res.set_max_threads(ctx.get_settings().get_max_threads()? as usize);

src/query/service/src/interpreters/interpreter_copy_v2.rs

@@ -180,7 +180,7 @@ impl CopyInterpreterV2 {
 
         let table = ctx.get_table(catalog_name, db_name, tbl_name).await?;
 
-        table.append2(ctx.clone(), &mut pipeline)?;
+        table.append2(ctx.clone(), &mut pipeline, false)?;
         pipeline.set_max_threads(settings.get_max_threads()? as usize);
 
         let query_need_abort = ctx.query_need_abort();

src/query/service/src/interpreters/interpreter_insert_v2.rs

@@ -15,27 +15,33 @@
 use std::collections::VecDeque;
 use std::sync::Arc;
 
-use common_datavalues::DataType;
+use common_catalog::table::Table;
 use common_exception::ErrorCode;
 use common_exception::Result;
-use common_functions::scalars::CastFunction;
 use common_streams::DataBlockStream;
 use common_streams::SendableDataBlockStream;
+use futures_util::StreamExt;
 use parking_lot::Mutex;
 
 use super::commit2table;
 use super::interpreter_common::append2table;
+use super::plan_schedulers::build_schedule_pipepline;
+use super::ProcessorExecutorStream;
 use crate::interpreters::Interpreter;
 use crate::interpreters::InterpreterPtr;
-use crate::interpreters::SelectInterpreterV2;
+use crate::pipelines::executor::ExecutorSettings;
+use crate::pipelines::executor::PipelinePullingExecutor;
 use crate::pipelines::processors::port::OutputPort;
 use crate::pipelines::processors::BlocksSource;
-use crate::pipelines::processors::TransformCastSchema;
 use crate::pipelines::Pipeline;
 use crate::pipelines::PipelineBuildResult;
 use crate::pipelines::SourcePipeBuilder;
 use crate::sessions::QueryContext;
 use crate::sessions::TableContext;
+use crate::sql::executor::DistributedInsertSelect;
+use crate::sql::executor::PhysicalPlan;
+use crate::sql::executor::PhysicalPlanBuilder;
+use crate::sql::executor::PipelineBuilder;
 use crate::sql::plans::Insert;
 use crate::sql::plans::InsertInputSource;
 use crate::sql::plans::Plan;
@@ -101,53 +107,7 @@ impl InsertInterpreterV2 {
                     );
                 }
                 InsertInputSource::SelectPlan(plan) => {
-                    let select_interpreter = match &**plan {
-                        Plan::Query {
-                            s_expr,
-                            metadata,
-                            bind_context,
-                            ..
-                        } => SelectInterpreterV2::try_create(
-                            self.ctx.clone(),
-                            *bind_context.clone(),
-                            *s_expr.clone(),
-                            metadata.clone(),
-                        ),
-                        _ => unreachable!(),
-                    };
-
-                    build_res = select_interpreter?.create_new_pipeline().await?;
-
-                    if self.check_schema_cast(plan)? {
-                        let mut functions = Vec::with_capacity(self.plan.schema().fields().len());
-
-                        for (target_field, original_field) in self
-                            .plan
-                            .schema()
-                            .fields()
-                            .iter()
-                            .zip(plan.schema().fields().iter())
-                        {
-                            let target_type_name = target_field.data_type().name();
-                            let from_type = original_field.data_type().clone();
-                            let cast_function =
-                                CastFunction::create("cast", &target_type_name, from_type)?;
-                            functions.push(cast_function);
-                        }
-
-                        let func_ctx = self.ctx.try_get_function_context()?;
-                        build_res.main_pipeline.add_transform(
-                            |transform_input_port, transform_output_port| {
-                                TransformCastSchema::try_create(
-                                    transform_input_port,
-                                    transform_output_port,
-                                    self.plan.schema(),
-                                    functions.clone(),
-                                    func_ctx.clone(),
-                                )
-                            },
-                        )?;
-                    }
+                    return self.schedule_insert_select(plan, table.clone()).await;
                 }
             };
         }
@@ -177,6 +137,93 @@ impl InsertInterpreterV2 {
         let cast_needed = select_schema != *output_schema;
         Ok(cast_needed)
     }
+
+    async fn schedule_insert_select(
+        &self,
+        plan: &Plan,
+        table: Arc<dyn Table>,
+    ) -> Result<SendableDataBlockStream> {
+        // select_plan is already distributed optimized
+        let (mut select_plan, select_column_bindings) = match plan {
+            Plan::Query {
+                s_expr,
+                metadata,
+                bind_context,
+                ..
+            } => {
+                let builder = PhysicalPlanBuilder::new(metadata.clone(), self.ctx.clone());
+                (builder.build(s_expr).await?, bind_context.columns.clone())
+            }
+            _ => unreachable!(),
+        };
+
+        table.get_table_info();
+        let catalog = self.plan.catalog.clone();
+        let is_distributed_plan = select_plan.is_distributed_plan();
+
+        let insert_select_plan = match select_plan {
+            PhysicalPlan::Exchange(ref mut exchange) => {
+                // insert can be dispatched to different nodes
+                let input = exchange.input.clone();
+                exchange.input = Box::new(PhysicalPlan::DistributedInsertSelect(Box::new(
+                    DistributedInsertSelect {
+                        input,
+                        catalog,
+                        table_info: table.get_table_info().clone(),
+                        select_schema: plan.schema(),
+                        select_column_bindings,
+                        insert_schema: self.plan.schema(),
+                        cast_needed: self.check_schema_cast(plan)?,
+                    },
+                )));
+                select_plan
+            }
+            other_plan => {
+                // insert should wait until all nodes finished
+                PhysicalPlan::DistributedInsertSelect(Box::new(DistributedInsertSelect {
+                    input: Box::new(other_plan),
+                    catalog,
+                    table_info: table.get_table_info().clone(),
+                    select_schema: plan.schema(),
+                    select_column_bindings,
+                    insert_schema: self.plan.schema(),
+                    cast_needed: self.check_schema_cast(plan)?,
+                }))
+            }
+        };
+
+        let mut build_res = if !is_distributed_plan {
+            let builder = PipelineBuilder::create(self.ctx.clone());
+            builder.finalize(&insert_select_plan)?
+        } else {
+            build_schedule_pipepline(self.ctx.clone(), &insert_select_plan).await?
+        };
+
+        let settings = self.ctx.get_settings();
+        let query_need_abort = self.ctx.query_need_abort();
+        let executor_settings = ExecutorSettings::try_create(&settings)?;
+        build_res.set_max_threads(settings.get_max_threads()? as usize);
+
+        let executor = PipelinePullingExecutor::from_pipelines(
+            query_need_abort,
+            build_res,
+            executor_settings,
+        )?;
+
+        let mut stream: SendableDataBlockStream =
+            Box::pin(ProcessorExecutorStream::create(executor)?);
+        while let Some(block) = stream.next().await {
+            block?;
+        }
+
+        commit2table(self.ctx.clone(), table.clone(), self.plan.overwrite).await?;
+
+        Ok(Box::pin(DataBlockStream::create(
+            self.plan.schema(),
+            None,
+            vec![],
+        )))
+    }
 }
 
 #[async_trait::async_trait]

src/query/service/src/interpreters/mod.rs

@@ -99,6 +99,7 @@ pub use interpreter_cluster_key_drop::DropTableClusterKeyInterpreter;
 pub use interpreter_clustering_history::InterpreterClusteringHistory;
 pub use interpreter_common::append2table;
 pub use interpreter_common::commit2table;
+pub use interpreter_common::fill_missing_columns;
 pub use interpreter_common::list_files_from_dal;
 pub use interpreter_database_create::CreateDatabaseInterpreter;
 pub use interpreter_database_drop::DropDatabaseInterpreter;

src/query/service/src/interpreters/plan_schedulers/mod.rs

@@ -15,6 +15,7 @@
 mod plan_scheduler_query;
 mod plan_scheduler_rewriter;
 
+pub use plan_scheduler_query::build_schedule_pipepline;
 pub use plan_scheduler_query::schedule_query_new;
 pub use plan_scheduler_query::schedule_query_v2;
 pub use plan_scheduler_rewriter::apply_plan_rewrite;

src/query/service/src/interpreters/plan_schedulers/plan_scheduler_query.rs

@@ -81,6 +81,22 @@ pub async fn schedule_query_v2(
         return Ok(build_res);
     }
 
+    let mut build_res = build_schedule_pipepline(ctx.clone(), plan).await?;
+
+    let input_schema = plan.output_schema()?;
+    PipelineBuilderV2::render_result_set(
+        input_schema,
+        result_columns,
+        &mut build_res.main_pipeline,
+    )?;
+
+    Ok(build_res)
+}
+
+pub async fn build_schedule_pipepline(
+    ctx: Arc<QueryContext>,
+    plan: &PhysicalPlan,
+) -> Result<PipelineBuildResult> {
     let fragmenter = Fragmenter::try_create(ctx.clone())?;
     let root_fragment = fragmenter.build_fragment(plan)?;
 
@@ -95,13 +111,5 @@ pub async fn schedule_query_v2(
 
     let settings = ctx.get_settings();
     build_res.set_max_threads(settings.get_max_threads()? as usize);
-
-    let input_schema = plan.output_schema()?;
-    PipelineBuilderV2::render_result_set(
-        input_schema,
-        result_columns,
-        &mut build_res.main_pipeline,
-    )?;
-
     Ok(build_res)
 }

src/query/service/src/sql/executor/format.rs

@@ -53,7 +53,9 @@ fn to_format_tree(plan: &PhysicalPlan, metadata: &MetadataRef) -> Result<FormatT
         PhysicalPlan::HashJoin(plan) => hash_join_to_format_tree(plan, metadata),
         PhysicalPlan::Exchange(plan) => exchange_to_format_tree(plan, metadata),
         PhysicalPlan::UnionAll(plan) => union_all_to_format_tree(plan, metadata),
-        PhysicalPlan::ExchangeSource(_) | PhysicalPlan::ExchangeSink(_) => {
+        PhysicalPlan::ExchangeSource(_)
+        | PhysicalPlan::ExchangeSink(_)
+        | PhysicalPlan::DistributedInsertSelect(_) => {
             Err(ErrorCode::LogicalError("Invalid physical plan"))
         }
     }