Test cloud distributed merge into

src/binaries/query/entry.rs

@@ -83,9 +83,9 @@ pub async fn init_services(conf: &InnerConfig) -> Result<()> {
 
 async fn precheck_services(conf: &InnerConfig) -> Result<()> {
     if conf.query.max_memory_limit_enabled {
-        let size = conf.query.max_server_memory_usage as i64;
-        info!("Set memory limit: {}", size);
-        GLOBAL_MEM_STAT.set_limit(size);
+        // let size = conf.query.max_server_memory_usage as i64;
+        // info!("Set memory limit: {}", size);
+        // GLOBAL_MEM_STAT.set_limit(size);
     }
 
     let tenant = conf.query.tenant_id.clone();

src/query/ast/src/ast/statements/merge_into.rs

@@ -31,18 +31,13 @@ use crate::ast::TableReference;
 
 #[derive(Debug, Clone, PartialEq)]
 pub struct MergeUpdateExpr {
-    pub catalog: Option<Identifier>,
     pub table: Option<Identifier>,
     pub name: Identifier,
     pub expr: Expr,
 }
 
 impl Display for MergeUpdateExpr {
     fn fmt(&self, f: &mut Formatter) -> std::fmt::Result {
-        if self.catalog.is_some() {
-            write!(f, "{}.", self.catalog.clone().unwrap())?;
-        }
-
         if self.table.is_some() {
             write!(f, "{}.", self.table.clone().unwrap())?;
         }
@@ -103,7 +98,7 @@ impl Display for MergeIntoStmt {
         write!(f, "MERGE INTO ")?;
         write_dot_separated_list(
             f,
-            self.catalog
+            self.database
                 .iter()
                 .chain(&self.database)
                 .chain(Some(&self.table_ident)),

src/query/ast/src/parser/statement.rs

@@ -2850,12 +2850,7 @@ pub fn udf_definition(i: Input) -> IResult<UDFDefinition> {
 
 pub fn merge_update_expr(i: Input) -> IResult<MergeUpdateExpr> {
     map(
-        rule! { ( #dot_separated_idents_1_to_3 ~ "=" ~ ^#expr ) },
-        |((catalog, table, name), _, expr)| MergeUpdateExpr {
-            catalog,
-            table,
-            name,
-            expr,
-        },
+        rule! { ( #dot_separated_idents_1_to_2 ~ "=" ~ ^#expr ) },
+        |((table, name), _, expr)| MergeUpdateExpr { table, name, expr },
     )(i)
 }

src/query/expression/src/schema.rs

@@ -51,6 +51,7 @@ pub const SEGMENT_NAME_COLUMN_ID: u32 = u32::MAX - 2;
 pub const SNAPSHOT_NAME_COLUMN_ID: u32 = u32::MAX - 3;
 
 pub const ROW_ID_COL_NAME: &str = "_row_id";
+pub const ROW_NUMBER_COL_NAME: &str = "_row_number";
 pub const SNAPSHOT_NAME_COL_NAME: &str = "_snapshot_name";
 pub const SEGMENT_NAME_COL_NAME: &str = "_segment_name";
 pub const BLOCK_NAME_COL_NAME: &str = "_block_name";

src/query/pipeline/core/src/pipe.rs

@@ -144,4 +144,10 @@ impl TransformPipeBuilder {
         }
         self.items = items
     }
+
+    pub fn add_items(&mut self, items: Vec<PipeItem>) {
+        for item in items {
+            self.items.push(item)
+        }
+    }
 }

src/query/service/src/api/rpc/exchange/exchange_manager.rs

@@ -58,6 +58,7 @@ use crate::pipelines::executor::ExecutorSettings;
 use crate::pipelines::executor::PipelineCompleteExecutor;
 use crate::pipelines::PipelineBuildResult;
 use crate::pipelines::PipelineBuilder;
+use crate::pipelines::PipelineBuilderData;
 use crate::schedulers::QueryFragmentActions;
 use crate::schedulers::QueryFragmentsActions;
 use crate::sessions::QueryContext;
@@ -210,6 +211,28 @@ impl DataExchangeManager {
         }
     }
 
+    // Create a pipeline based on query plan
+    #[minitrace::trace]
+    pub fn init_query_fragments_plan_local(
+        &self,
+        ctx: &Arc<QueryContext>,
+        packet: &QueryFragmentsPlanPacket,
+    ) -> Result<()> {
+        let queries_coordinator_guard = self.queries_coordinator.lock();
+        let queries_coordinator = unsafe { &mut *queries_coordinator_guard.deref().get() };
+
+        // TODO: When the query is not executed for a long time after submission, we need to remove it
+        match queries_coordinator.get_mut(&packet.query_id) {
+            None => Err(ErrorCode::Internal(format!(
+                "Query {} not found in cluster.",
+                packet.query_id
+            ))),
+            Some(query_coordinator) => {
+                query_coordinator.prepare_pipeline_local(ctx, packet.enable_profiling, packet)
+            }
+        }
+    }
+
     #[minitrace::trace]
     pub fn handle_statistics_exchange(
         &self,
@@ -297,7 +320,7 @@ impl DataExchangeManager {
             .await?;
 
         // Submit tasks to localhost
-        self.init_query_fragments_plan(&ctx, &local_query_fragments_plan_packet)?;
+        self.init_query_fragments_plan_local(&ctx, &local_query_fragments_plan_packet)?;
 
         // Get local pipeline of local task
         let build_res = self.get_root_pipeline(ctx, enable_profiling, root_actions)?;
@@ -590,6 +613,45 @@ impl QueryCoordinator {
         Ok(())
     }
 
+    pub fn prepare_pipeline_local(
+        &mut self,
+        ctx: &Arc<QueryContext>,
+        enable_profiling: bool,
+        packet: &QueryFragmentsPlanPacket,
+    ) -> Result<()> {
+        self.info = Some(QueryInfo {
+            query_ctx: ctx.clone(),
+            query_id: packet.query_id.clone(),
+            current_executor: packet.executor.clone(),
+            query_executor: None,
+        });
+
+        for fragment in &packet.fragments {
+            self.fragments_coordinator.insert(
+                fragment.fragment_id.to_owned(),
+                FragmentCoordinator::create(fragment),
+            );
+        }
+
+        let mut input_builder_data = PipelineBuilderData {
+            input_join_state: None,
+            input_probe_schema: None,
+        };
+
+        for fragment in &packet.fragments {
+            let fragment_id = fragment.fragment_id;
+            if let Some(coordinator) = self.fragments_coordinator.get_mut(&fragment_id) {
+                input_builder_data = coordinator.prepare_pipeline_local(
+                    ctx.clone(),
+                    enable_profiling,
+                    input_builder_data,
+                )?;
+            }
+        }
+
+        Ok(())
+    }
+
     pub fn subscribe_fragment(
         &mut self,
         ctx: &Arc<QueryContext>,
@@ -821,16 +883,53 @@ impl FragmentCoordinator {
             self.initialized = true;
 
             let pipeline_ctx = QueryContext::create_from(ctx);
+
             let pipeline_builder = PipelineBuilder::create(
                 pipeline_ctx.get_function_context()?,
                 pipeline_ctx.get_settings(),
                 pipeline_ctx,
                 enable_profiling,
                 SharedProcessorProfiles::default(),
             );
-            self.pipeline_build_res = Some(pipeline_builder.finalize(&self.physical_plan)?);
+            let res = pipeline_builder.finalize(&self.physical_plan)?;
+
+            self.pipeline_build_res = Some(res);
         }
 
         Ok(())
     }
+
+    pub fn prepare_pipeline_local(
+        &mut self,
+        ctx: Arc<QueryContext>,
+        enable_profiling: bool,
+        input_builder_data: PipelineBuilderData,
+    ) -> Result<PipelineBuilderData> {
+        if !self.initialized {
+            self.initialized = true;
+
+            let pipeline_ctx = QueryContext::create_from(ctx);
+
+            let mut pipeline_builder = PipelineBuilder::create(
+                pipeline_ctx.get_function_context()?,
+                pipeline_ctx.get_settings(),
+                pipeline_ctx,
+                enable_profiling,
+                SharedProcessorProfiles::default(),
+            );
+
+            pipeline_builder.join_state = input_builder_data.input_join_state;
+            pipeline_builder.probe_data_fields = input_builder_data.input_probe_schema;
+
+            let (res, builder_data) = pipeline_builder.finalize_local(&self.physical_plan)?;
+
+            self.pipeline_build_res = Some(res);
+            return Ok(builder_data);
+        }
+
+        Ok(PipelineBuilderData {
+            input_join_state: None,
+            input_probe_schema: None,
+        })
+    }
 }