Implement a skeleton of distributed insert select.

Author: RinChanNOWWW

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,26 +40,39 @@ 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(),
-                )
-            })?;
-    }
+    fill_missing_columns(
+        ctx.clone(),
+        &source_schema,
+        &table.schema(),
+        &mut build_res.main_pipeline,
+    )?;
 
     table.append2(ctx.clone(), &mut build_res.main_pipeline)?;
     let query_need_abort = ctx.query_need_abort();

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

@@ -15,19 +15,25 @@
 use std::collections::VecDeque;
 use std::sync::Arc;
 
+use common_catalog::table::Table;
 use common_datavalues::DataType;
 use common_exception::ErrorCode;
 use common_exception::Result;
 use common_functions::scalars::CastFunction;
+use common_planners::StageKind;
 use common_streams::DataBlockStream;
 use common_streams::SendableDataBlockStream;
 use parking_lot::Mutex;
 
 use super::commit2table;
 use super::interpreter_common::append2table;
+use super::plan_schedulers::build_schedule_pipepline;
+use crate::clusters::ClusterHelper;
 use crate::interpreters::Interpreter;
 use crate::interpreters::InterpreterPtr;
 use crate::interpreters::SelectInterpreterV2;
+use crate::pipelines::executor::ExecutorSettings;
+use crate::pipelines::executor::PipelineCompleteExecutor;
 use crate::pipelines::processors::port::OutputPort;
 use crate::pipelines::processors::BlocksSource;
 use crate::pipelines::processors::TransformCastSchema;
@@ -36,6 +42,10 @@ 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::Exchange;
+use crate::sql::executor::PhysicalPlan;
+use crate::sql::executor::PhysicalPlanBuilder;
 use crate::sql::plans::Insert;
 use crate::sql::plans::InsertInputSource;
 use crate::sql::plans::Plan;
@@ -101,6 +111,13 @@ impl InsertInterpreterV2 {
                     );
                 }
                 InsertInputSource::SelectPlan(plan) => {
+                    if !self.ctx.get_cluster().is_empty() {
+                        // distributed insert select
+                        return self
+                            .schedule_insert_select(plan, self.plan.catalog.clone(), table.clone())
+                            .await;
+                    }
+
                     let select_interpreter = match &**plan {
                         Plan::Query {
                             s_expr,
@@ -177,6 +194,62 @@ impl InsertInterpreterV2 {
         let cast_needed = select_schema != *output_schema;
         Ok(cast_needed)
     }
+
+    async fn schedule_insert_select(
+        &self,
+        plan: &Box<Plan>,
+        catalog: String,
+        table: Arc<dyn Table>,
+    ) -> Result<SendableDataBlockStream> {
+        let inner_plan = match &**plan {
+            Plan::Query {
+                s_expr, metadata, ..
+            } => {
+                let builder = PhysicalPlanBuilder::new(metadata.clone(), self.ctx.clone());
+                builder.build(s_expr).await?
+            }
+            _ => unreachable!(),
+        };
+
+        table.get_table_info();
+
+        let insert_select_plan = PhysicalPlan::DistributedInsertSelect(DistributedInsertSelect {
+            input: Box::new(inner_plan),
+            catalog,
+            table_info: table.get_table_info().clone(),
+            select_schema: plan.schema(),
+            insert_schema: self.schema(),
+            cast_needed: self.check_schema_cast(plan)?,
+        });
+
+        let final_plan = PhysicalPlan::Exchange(Exchange {
+            input: Box::new(insert_select_plan),
+            kind: StageKind::Merge,
+            keys: vec![],
+        });
+
+        let mut build_res = build_schedule_pipepline(self.ctx.clone(), &final_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 mut pipelines = build_res.sources_pipelines;
+        pipelines.push(build_res.main_pipeline);
+        let executor = PipelineCompleteExecutor::from_pipelines(
+            query_need_abort,
+            pipelines,
+            executor_settings,
+        )?;
+        executor.execute()?;
+        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"))
         }
     }

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

@@ -21,9 +21,11 @@ use common_datavalues::DataField;
 use common_datavalues::DataSchemaRef;
 use common_datavalues::DataSchemaRefExt;
 use common_datavalues::NullableType;
+use common_datavalues::StringType;
 use common_datavalues::ToDataType;
 use common_datavalues::Vu8;
 use common_exception::Result;
+use common_meta_app::schema::TableInfo;
 use common_planners::ReadDataSourcePlan;
 use common_planners::StageKind;
 
@@ -303,6 +305,25 @@ impl UnionAll {
     }
 }
 
+#[derive(Clone, Debug, serde::Serialize, serde::Deserialize)]
+pub struct DistributedInsertSelect {
+    pub input: Box<PhysicalPlan>,
+    pub catalog: String,
+    pub table_info: TableInfo,
+    pub insert_schema: DataSchemaRef,
+    pub select_schema: DataSchemaRef,
+    pub cast_needed: bool,
+}
+
+impl DistributedInsertSelect {
+    pub fn output_schema(&self) -> Result<DataSchemaRef> {
+        Ok(DataSchemaRefExt::create(vec![
+            DataField::new("seg_loc", StringType::new_impl()),
+            DataField::new("seg_info", StringType::new_impl()),
+        ]))
+    }
+}
+
 #[derive(Clone, Debug, serde::Serialize, serde::Deserialize)]
 pub enum PhysicalPlan {
     TableScan(TableScan),
@@ -317,6 +338,9 @@ pub enum PhysicalPlan {
     Exchange(Exchange),
     UnionAll(UnionAll),
 
+    /// For insert into ... select ... in cluster
+    DistributedInsertSelect(DistributedInsertSelect),
+
     /// Synthesized by fragmenter
     ExchangeSource(ExchangeSource),
     ExchangeSink(ExchangeSink),
@@ -346,6 +370,7 @@ impl PhysicalPlan {
             PhysicalPlan::ExchangeSource(plan) => plan.output_schema(),
             PhysicalPlan::ExchangeSink(plan) => plan.output_schema(),
             PhysicalPlan::UnionAll(plan) => plan.output_schema(),
+            PhysicalPlan::DistributedInsertSelect(plan) => plan.output_schema(),
         }
     }
 
@@ -368,6 +393,9 @@ impl PhysicalPlan {
             PhysicalPlan::UnionAll(plan) => Box::new(
                 std::iter::once(plan.left.as_ref()).chain(std::iter::once(plan.right.as_ref())),
             ),
+            PhysicalPlan::DistributedInsertSelect(plan) => {
+                Box::new(std::iter::once(plan.input.as_ref()))
+            }
         }
     }
 }

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

@@ -18,6 +18,7 @@ use std::fmt::Formatter;
 use common_datavalues::format_data_type_sql;
 use itertools::Itertools;
 
+use super::DistributedInsertSelect;
 use crate::sql::executor::AggregateFinal;
 use crate::sql::executor::AggregatePartial;
 use crate::sql::executor::EvalScalar;
@@ -64,6 +65,7 @@ impl<'a> Display for PhysicalPlanIndentFormatDisplay<'a> {
             PhysicalPlan::ExchangeSource(source) => write!(f, "{}", source)?,
             PhysicalPlan::ExchangeSink(sink) => write!(f, "{}", sink)?,
             PhysicalPlan::UnionAll(union_all) => write!(f, "{}", union_all)?,
+            PhysicalPlan::DistributedInsertSelect(insert_select) => write!(f, "{}", insert_select)?,
         }
 
         for node in self.node.children() {
@@ -313,3 +315,9 @@ impl Display for UnionAll {
         write!(f, "UnionAll")
     }
 }
+
+impl Display for DistributedInsertSelect {
+    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+        write!(f, "DistributedInsertSelect")
+    }
+}

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

@@ -16,6 +16,7 @@ use common_exception::Result;
 
 use super::AggregateFinal;
 use super::AggregatePartial;
+use super::DistributedInsertSelect;
 use super::EvalScalar;
 use super::Exchange;
 use super::ExchangeSink;
@@ -45,6 +46,7 @@ pub trait PhysicalPlanReplacer {
             PhysicalPlan::ExchangeSource(plan) => self.replace_exchange_source(plan),
             PhysicalPlan::ExchangeSink(plan) => self.replace_exchange_sink(plan),
             PhysicalPlan::UnionAll(plan) => self.replace_union(plan),
+            PhysicalPlan::DistributedInsertSelect(plan) => self.replace_insert_select(plan),
         }
     }
 
@@ -174,6 +176,21 @@ pub trait PhysicalPlanReplacer {
             schema: plan.schema.clone(),
         }))
     }
+
+    fn replace_insert_select(&mut self, plan: &DistributedInsertSelect) -> Result<PhysicalPlan> {
+        let input = self.replace(&plan.input)?;
+
+        Ok(PhysicalPlan::DistributedInsertSelect(
+            DistributedInsertSelect {
+                input: Box::new(input),
+                catalog: plan.catalog.clone(),
+                table_info: plan.table_info.clone(),
+                select_schema: plan.select_schema.clone(),
+                insert_schema: plan.insert_schema.clone(),
+                cast_needed: plan.cast_needed,
+            },
+        ))
+    }
 }
 
 impl PhysicalPlan {
@@ -223,6 +240,9 @@ impl PhysicalPlan {
                     Self::traverse(&plan.left, pre_visit, visit, post_visit);
                     Self::traverse(&plan.right, pre_visit, visit, post_visit);
                 }
+                PhysicalPlan::DistributedInsertSelect(plan) => {
+                    Self::traverse(&plan.input, pre_visit, visit, post_visit);
+                }
             }
             post_visit(plan);
         }