Finish the distributed insert select pipeline.

RinChanNOWWW · RinChanNOWWW · commit 07fc475fd68e · 2022-09-13T14:11:41.000+08:00
diff --git a/src/query/catalog/src/table.rs b/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(),
diff --git a/src/query/pipeline/sinks/src/processors/sinks/context_sink.rs b/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);
diff --git a/src/query/service/src/interpreters/interpreter_common.rs b/src/query/service/src/interpreters/interpreter_common.rs
@@ -19,6 +19,7 @@ use chrono::TimeZone;
 use chrono::Utc;
 use common_base::base::GlobalIORuntime;
 use common_base::base::TrySpawn;
+use common_datablocks::DataBlock;
 use common_datavalues::DataSchemaRef;
 use common_exception::ErrorCode;
 use common_exception::Result;
@@ -74,7 +75,7 @@ pub fn append2table(
         &mut build_res.main_pipeline,
     )?;
 
-    table.append2(ctx.clone(), &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);
@@ -85,12 +86,12 @@ pub fn append2table(
     executor.execute()
 }
 
-pub async fn commit2table(
+pub async fn commit2table_with_append_entries(
     ctx: Arc<QueryContext>,
     table: Arc<dyn Table>,
     overwrite: bool,
+    append_entries: Vec<DataBlock>,
 ) -> Result<()> {
-    let append_entries = ctx.consume_precommit_blocks();
     // We must put the commit operation to global runtime, which will avoid the "dispatch dropped without returning error" in tower
     let catalog_name = ctx.get_current_catalog();
     let handler = GlobalIORuntime::instance().spawn(async move {
@@ -109,6 +110,15 @@ pub async fn commit2table(
     }
 }
 
+pub async fn commit2table(
+    ctx: Arc<QueryContext>,
+    table: Arc<dyn Table>,
+    overwrite: bool,
+) -> Result<()> {
+    let append_entries = ctx.consume_precommit_blocks();
+    commit2table_with_append_entries(ctx, table, overwrite, append_entries).await
+}
+
 pub async fn validate_grant_object_exists(
     ctx: &Arc<QueryContext>,
     object: &GrantObject,
diff --git a/src/query/service/src/interpreters/interpreter_copy_v2.rs b/src/query/service/src/interpreters/interpreter_copy_v2.rs
@@ -176,7 +176,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();
diff --git a/src/query/service/src/interpreters/interpreter_insert_v2.rs b/src/query/service/src/interpreters/interpreter_insert_v2.rs
@@ -26,14 +26,15 @@ use common_streams::SendableDataBlockStream;
 use parking_lot::Mutex;
 
 use super::commit2table;
+use super::commit2table_with_append_entries;
 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::executor::PipelinePullingExecutor;
 use crate::pipelines::processors::port::OutputPort;
 use crate::pipelines::processors::BlocksSource;
 use crate::pipelines::processors::TransformCastSchema;
@@ -218,7 +219,7 @@ impl InsertInterpreterV2 {
             catalog,
             table_info: table.get_table_info().clone(),
             select_schema: plan.schema(),
-            insert_schema: self.schema(),
+            insert_schema: self.plan.schema(),
             cast_needed: self.check_schema_cast(plan)?,
         });
 
@@ -234,15 +235,15 @@ impl InsertInterpreterV2 {
         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(
+
+        let executor = PipelinePullingExecutor::from_pipelines(
             query_need_abort,
-            pipelines,
+            build_res,
             executor_settings,
         )?;
-        executor.execute()?;
-        commit2table(self.ctx.clone(), table.clone(), self.plan.overwrite).await?;
+
+        commit2table_with_append_entries(self.ctx.clone(), table, self.plan.overwrite, vec![])
+            .await?;
 
         Ok(Box::pin(DataBlockStream::create(
             self.plan.schema(),
diff --git a/src/query/service/src/interpreters/mod.rs b/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::commit2table_with_append_entries;
 pub use interpreter_common::fill_missing_columns;
 pub use interpreter_common::list_files_from_dal;
 pub use interpreter_database_create::CreateDatabaseInterpreter;
diff --git a/src/query/service/src/sql/executor/physical_plan.rs b/src/query/service/src/sql/executor/physical_plan.rs
@@ -21,13 +21,13 @@ 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;
+use common_planners::SINK_SCHEMA;
 
 use super::physical_scalar::PhysicalScalar;
 use super::AggregateFunctionDesc;
@@ -317,10 +317,7 @@ pub struct DistributedInsertSelect {
 
 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()),
-        ]))
+        Ok(SINK_SCHEMA.clone())
     }
 }
 
diff --git a/src/query/service/src/sql/executor/pipeline_builder.rs b/src/query/service/src/sql/executor/pipeline_builder.rs
@@ -126,7 +126,7 @@ impl PipelineBuilder {
             PhysicalPlan::ExchangeSource(source) => self.build_exchange_source(source),
             PhysicalPlan::UnionAll(union_all) => self.build_union_all(union_all),
             PhysicalPlan::DistributedInsertSelect(insert_select) => {
-                self.build_insert_select(insert_select)
+                self.build_distributed_insert_select(insert_select)
             }
             PhysicalPlan::Exchange(_) => Err(ErrorCode::LogicalError(
                 "Invalid physical plan with PhysicalPlan::Exchange",
@@ -532,7 +532,10 @@ impl PipelineBuilder {
         Ok(())
     }
 
-    pub fn build_insert_select(&mut self, insert_select: &DistributedInsertSelect) -> Result<()> {
+    pub fn build_distributed_insert_select(
+        &mut self,
+        insert_select: &DistributedInsertSelect,
+    ) -> Result<()> {
         let select_schema = &insert_select.select_schema;
         let insert_schema = &insert_select.insert_schema;
 
@@ -576,7 +579,7 @@ impl PipelineBuilder {
             &mut self.main_pipeline,
         )?;
 
-        table.append2(self.ctx.clone(), &mut self.main_pipeline)?;
+        table.append2(self.ctx.clone(), &mut self.main_pipeline, true)?;
 
         Ok(())
     }
diff --git a/src/query/service/src/storages/stage/stage_table.rs b/src/query/service/src/storages/stage/stage_table.rs
@@ -123,7 +123,7 @@ impl Table for StageTable {
         Ok(())
     }
 
-    fn append2(&self, ctx: Arc<dyn TableContext>, pipeline: &mut Pipeline) -> Result<()> {
+    fn append2(&self, ctx: Arc<dyn TableContext>, pipeline: &mut Pipeline, _: bool) -> Result<()> {
         let mut sink_pipeline_builder = SinkPipeBuilder::create();
         for _ in 0..pipeline.output_len() {
             let input_port = InputPort::create();
diff --git a/src/query/service/tests/it/storages/fuse/operations/navigate.rs b/src/query/service/tests/it/storages/fuse/operations/navigate.rs
@@ -158,7 +158,7 @@ async fn test_fuse_historical_table_is_read_only() -> Result<()> {
         .await?;
 
     // check append2
-    let res = tbl.append2(ctx.clone(), &mut Pipeline::create());
+    let res = tbl.append2(ctx.clone(), &mut Pipeline::create(), false);
     assert_not_writable(res, "append2");
 
     // check append_data
diff --git a/src/query/storages/fuse/src/fuse_table.rs b/src/query/storages/fuse/src/fuse_table.rs
@@ -325,9 +325,14 @@ impl Table for FuseTable {
         self.do_read2(ctx, plan, pipeline)
     }
 
-    fn append2(&self, ctx: Arc<dyn TableContext>, pipeline: &mut Pipeline) -> Result<()> {
+    fn append2(
+        &self,
+        ctx: Arc<dyn TableContext>,
+        pipeline: &mut Pipeline,
+        need_output: bool,
+    ) -> Result<()> {
         self.check_mutable()?;
-        self.do_append2(ctx, pipeline)
+        self.do_append2(ctx, pipeline, need_output)
     }
 
     #[tracing::instrument(level = "debug", name = "fuse_table_commit_insertion", skip(self, ctx, operations), fields(ctx.id = ctx.get_id().as_str()))]
diff --git a/src/query/storages/fuse/src/operations/append.rs b/src/query/storages/fuse/src/operations/append.rs
@@ -36,7 +36,12 @@ use crate::DEFAULT_BLOCK_PER_SEGMENT;
 use crate::FUSE_OPT_KEY_BLOCK_PER_SEGMENT;
 
 impl FuseTable {
-    pub fn do_append2(&self, ctx: Arc<dyn TableContext>, pipeline: &mut Pipeline) -> Result<()> {
+    pub fn do_append2(
+        &self,
+        ctx: Arc<dyn TableContext>,
+        pipeline: &mut Pipeline,
+        need_output: bool,
+    ) -> Result<()> {
         let block_per_seg =
             self.get_option(FUSE_OPT_KEY_BLOCK_PER_SEGMENT, DEFAULT_BLOCK_PER_SEGMENT);
 
@@ -74,23 +79,38 @@ impl FuseTable {
         }
 
         let da = ctx.get_storage_operator()?;
-        let mut sink_pipeline_builder = SinkPipeBuilder::create();
-        for _ in 0..pipeline.output_len() {
-            let input_port = InputPort::create();
-            sink_pipeline_builder.add_sink(
-                input_port.clone(),
+        if need_output {
+            pipeline.add_transform(|transform_input_port, transform_output_port| {
                 FuseTableSink::try_create(
-                    input_port,
+                    transform_input_port,
                     ctx.clone(),
                     block_per_seg,
                     da.clone(),
                     self.meta_location_generator().clone(),
                     cluster_stats_gen.clone(),
-                )?,
-            );
+                    Some(transform_output_port),
+                )
+            })?;
+        } else {
+            let mut sink_pipeline_builder = SinkPipeBuilder::create();
+            for _ in 0..pipeline.output_len() {
+                let input_port = InputPort::create();
+                sink_pipeline_builder.add_sink(
+                    input_port.clone(),
+                    FuseTableSink::try_create(
+                        input_port,
+                        ctx.clone(),
+                        block_per_seg,
+                        da.clone(),
+                        self.meta_location_generator().clone(),
+                        cluster_stats_gen.clone(),
+                        None,
+                    )?,
+                );
+            }
+
+            pipeline.add_pipe(sink_pipeline_builder.finalize());
         }
-
-        pipeline.add_pipe(sink_pipeline_builder.finalize());
         Ok(())
     }
 
diff --git a/src/query/storages/fuse/src/operations/compact.rs b/src/query/storages/fuse/src/operations/compact.rs
@@ -111,6 +111,7 @@ impl FuseTable {
                     mutator.get_storage_operator(),
                     self.meta_location_generator().clone(),
                     ClusterStatsGenerator::default(),
+                    None,
                 )?,
             );
         }
diff --git a/src/query/storages/fuse/src/operations/fuse_sink.rs b/src/query/storages/fuse/src/operations/fuse_sink.rs
@@ -27,6 +27,7 @@ use common_exception::Result;
 use common_fuse_meta::meta::Location;
 use common_fuse_meta::meta::SegmentInfo;
 use common_fuse_meta::meta::Statistics;
+use common_pipeline_core::processors::port::OutputPort;
 use common_storages_index::*;
 use opendal::Operator;
 
@@ -63,6 +64,7 @@ enum State {
         location: String,
         segment: Arc<SegmentInfo>,
     },
+    Output(DataBlock),
     Finished,
 }
 
@@ -75,6 +77,10 @@ pub struct FuseTableSink {
     meta_locations: TableMetaLocationGenerator,
     accumulator: StatisticsAccumulator,
     cluster_stats_gen: ClusterStatsGenerator,
+
+    // For distributed insert select.
+    // We should output the append logs to the exchange sink.
+    output: Option<Arc<OutputPort>>,
 }
 
 impl FuseTableSink {
@@ -85,6 +91,7 @@ impl FuseTableSink {
         data_accessor: Operator,
         meta_locations: TableMetaLocationGenerator,
         cluster_stats_gen: ClusterStatsGenerator,
+        output: Option<Arc<OutputPort>>,
     ) -> Result<ProcessorPtr> {
         Ok(ProcessorPtr::create(Box::new(FuseTableSink {
             ctx,
@@ -95,6 +102,7 @@ impl FuseTableSink {
             accumulator: Default::default(),
             num_block_threshold: num_block_threshold as u64,
             cluster_stats_gen,
+            output,
         })))
     }
 }
@@ -124,12 +132,23 @@ impl Processor for FuseTableSink {
             return Ok(Event::Async);
         }
 
+        if let State::Output(block) = std::mem::replace(&mut self.state, State::Finished) {
+            if let Some(output) = &self.output {
+                output.push_data(Ok(block));
+                return Ok(Event::NeedConsume);
+            } else {
+                return Err(ErrorCode::LogicalError("No output port"));
+            }
+        }
+
         if self.input.is_finished() {
             if self.accumulator.summary_row_count != 0 {
                 self.state = State::GenerateSegment;
                 return Ok(Event::Sync);
             }
-
+            if let Some(output) = &self.output {
+                output.finish();
+            }
             self.state = State::Finished;
             return Ok(Event::Finished);
         }
@@ -260,8 +279,12 @@ impl Processor for FuseTableSink {
 
                 // TODO: dyn operation for table trait
                 let log_entry = AppendOperationLogEntry::new(location, segment);
-                self.ctx
-                    .push_precommit_block(DataBlock::try_from(log_entry)?);
+                let data_block = DataBlock::try_from(log_entry)?;
+                if self.output.is_some() {
+                    self.state = State::Output(data_block)
+                } else {
+                    self.ctx.push_precommit_block(data_block);
+                }
             }
             _state => {
                 return Err(ErrorCode::LogicalError(
diff --git a/src/query/storages/fuse/src/operations/recluster.rs b/src/query/storages/fuse/src/operations/recluster.rs
@@ -191,6 +191,7 @@ impl FuseTable {
                     da.clone(),
                     self.meta_location_generator().clone(),
                     cluster_stats_gen.clone(),
+                    None,
                 )?,
             );
         }
diff --git a/src/query/storages/preludes/src/memory/memory_table.rs b/src/query/storages/preludes/src/memory/memory_table.rs
@@ -215,7 +215,7 @@ impl Table for MemoryTable {
         Ok(())
     }
 
-    fn append2(&self, ctx: Arc<dyn TableContext>, pipeline: &mut Pipeline) -> Result<()> {
+    fn append2(&self, ctx: Arc<dyn TableContext>, pipeline: &mut Pipeline, _: bool) -> Result<()> {
         let mut sink_pipeline_builder = SinkPipeBuilder::create();
         for _ in 0..pipeline.output_len() {
             let input_port = InputPort::create();
diff --git a/src/query/storages/preludes/src/null/null_table.rs b/src/query/storages/preludes/src/null/null_table.rs

Original file line number	Diff line number	Diff line change
`@@ -34,7 +34,7 @@ impl ContextSink {`
`34`	`34`	`}`
`35`	`35`
`36`	`36`	`impl Sink for ContextSink {`
`37`		`- const NAME: &'static str = "ContextSink ";`
	`37`	`+ const NAME: &'static str = "ContextSink";`
`38`	`38`
`39`	`39`	`fn consume(&mut self, block: DataBlock) -> Result<()> {`
`40`	`40`	`self.ctx.push_precommit_block(block);`
Original file line number	Diff line number	Diff line change
`@@ -123,7 +123,7 @@ impl Table for StageTable {`
`123`	`123`	`Ok(())`
`124`	`124`	`}`
`125`	`125`
`126`		`- fn append2(&self, ctx: Arc<dyn TableContext>, pipeline: &mut Pipeline) -> Result<()> {`
	`126`	`+ fn append2(&self, ctx: Arc<dyn TableContext>, pipeline: &mut Pipeline, _: bool) -> Result<()> {`
`127`	`127`	`let mut sink_pipeline_builder = SinkPipeBuilder::create();`
`128`	`128`	`for _ in 0..pipeline.output_len() {`
`129`	`129`	`let input_port = InputPort::create();`