feat: support using fragment forest to execute additional broadcast operations. (#17872)

* update

* avoid committing multiple times

* fix merge

* fix fragment type

* rm unused modify

* fix

* improve err msg

* fix build pipeline

* fix exhchange

* refactor

* fix root pipeline

* fix pipeline

* improve error msg

* fix serialize block

* fix sink

* add log

* fix send

* merge rf

* fix

* fix

* fix

* fix conflict

* update

* fix

* rm unused modify

* make lint

* fix

* rm log

* rm empty impl

* rm log

* fix

* add comment

Author: SkyFan2002

src/query/catalog/src/table_context.rs

@@ -418,6 +418,15 @@ pub trait TableContext: Send + Sync {
     fn set_pruned_partitions_stats(&self, _partitions: PartStatistics) {
         unimplemented!()
     }
+
+    /// Calling this function will automatically create a pipeline for broadcast data in `build_distributed_pipeline()`
+    ///
+    /// The returned id can be used to get sender and receiver for broadcasting data.
+    fn get_next_broadcast_id(&self) -> u32;
+
+    fn reset_broadcast_id(&self) {
+        unimplemented!()
+    }
 }
 
 pub type AbortChecker = Arc<dyn CheckAbort + Send + Sync>;

src/query/service/src/pipelines/builders/builder_broadcast.rs

@@ -0,0 +1,39 @@
+// Copyright 2021 Datafuse Labs
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use databend_common_exception::Result;
+use databend_common_sql::executor::physical_plans::BroadcastSink;
+use databend_common_sql::executor::physical_plans::BroadcastSource;
+
+use crate::pipelines::processors::transforms::BroadcastSinkProcessor;
+use crate::pipelines::processors::transforms::BroadcastSourceProcessor;
+use crate::pipelines::PipelineBuilder;
+
+impl PipelineBuilder {
+    pub(crate) fn build_broadcast_source(&mut self, source: &BroadcastSource) -> Result<()> {
+        let receiver = self.ctx.broadcast_source_receiver(source.broadcast_id);
+        self.main_pipeline.add_source(
+            |output| BroadcastSourceProcessor::create(self.ctx.clone(), receiver.clone(), output),
+            1,
+        )
+    }
+
+    pub(crate) fn build_broadcast_sink(&mut self, sink: &BroadcastSink) -> Result<()> {
+        self.build_pipeline(&sink.input)?;
+        self.main_pipeline.resize(1, true)?;
+        self.main_pipeline.add_sink(|input| {
+            BroadcastSinkProcessor::create(input, self.ctx.broadcast_sink_sender(sink.broadcast_id))
+        })
+    }
+}

src/query/service/src/pipelines/builders/mod.rs

@@ -16,6 +16,7 @@ mod builder_add_stream_column;
 mod builder_aggregate;
 mod builder_append_table;
 mod builder_async_function;
+mod builder_broadcast;
 mod builder_column_mutation;
 mod builder_commit;
 mod builder_compact;

src/query/service/src/pipelines/pipeline_builder.rs

@@ -220,6 +220,8 @@ impl PipelineBuilder {
             }
             PhysicalPlan::Shuffle(shuffle) => self.build_shuffle(shuffle),
             PhysicalPlan::Duplicate(duplicate) => self.build_duplicate(duplicate),
+            PhysicalPlan::BroadcastSource(source) => self.build_broadcast_source(source),
+            PhysicalPlan::BroadcastSink(sink) => self.build_broadcast_sink(sink),
 
             // ==============================
             // 4. Data Modification Operations

src/query/service/src/pipelines/processors/transforms/broadcast.rs

@@ -0,0 +1,97 @@
+// Copyright 2021 Datafuse Labs
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use std::sync::Arc;
+
+use async_channel::Receiver;
+use async_channel::Sender;
+use databend_common_catalog::table_context::TableContext;
+use databend_common_exception::ErrorCode;
+use databend_common_exception::Result;
+use databend_common_expression::BlockMetaInfoPtr;
+use databend_common_expression::DataBlock;
+use databend_common_pipeline_core::processors::InputPort;
+use databend_common_pipeline_core::processors::OutputPort;
+use databend_common_pipeline_core::processors::ProcessorPtr;
+use databend_common_pipeline_sinks::AsyncSink;
+use databend_common_pipeline_sinks::AsyncSinker;
+use databend_common_pipeline_sources::AsyncSource;
+use databend_common_pipeline_sources::AsyncSourcer;
+
+pub struct BroadcastSourceProcessor {
+    pub receiver: Receiver<BlockMetaInfoPtr>,
+}
+
+impl BroadcastSourceProcessor {
+    pub fn create(
+        ctx: Arc<dyn TableContext>,
+        receiver: Receiver<BlockMetaInfoPtr>,
+        output_port: Arc<OutputPort>,
+    ) -> Result<ProcessorPtr> {
+        AsyncSourcer::create(ctx, output_port, Self { receiver })
+    }
+}
+
+#[async_trait::async_trait]
+impl AsyncSource for BroadcastSourceProcessor {
+    const NAME: &'static str = "BroadcastSource";
+    const SKIP_EMPTY_DATA_BLOCK: bool = false;
+
+    #[async_backtrace::framed]
+    async fn generate(&mut self) -> Result<Option<DataBlock>> {
+        let received = self.receiver.recv().await;
+        match received {
+            Ok(meta) => Ok(Some(DataBlock::empty_with_meta(meta))),
+            Err(_) => {
+                // The channel is closed, we should return None to stop generating
+                Ok(None)
+            }
+        }
+    }
+}
+
+pub struct BroadcastSinkProcessor {
+    received: Vec<BlockMetaInfoPtr>,
+    sender: Sender<BlockMetaInfoPtr>,
+}
+
+impl BroadcastSinkProcessor {
+    pub fn create(input: Arc<InputPort>, sender: Sender<BlockMetaInfoPtr>) -> Result<ProcessorPtr> {
+        Ok(ProcessorPtr::create(AsyncSinker::create(input, Self {
+            received: vec![],
+            sender,
+        })))
+    }
+}
+
+#[async_trait::async_trait]
+impl AsyncSink for BroadcastSinkProcessor {
+    const NAME: &'static str = "BroadcastSink";
+
+    async fn on_finish(&mut self) -> Result<()> {
+        self.sender.close();
+        Ok(())
+    }
+
+    async fn consume(&mut self, mut data_block: DataBlock) -> Result<bool> {
+        let meta = data_block
+            .take_meta()
+            .ok_or_else(|| ErrorCode::Internal("Cannot downcast meta to BroadcastMeta"))?;
+        self.sender
+            .send(meta)
+            .await
+            .map_err(|_| ErrorCode::Internal("BroadcastSinkProcessor send error"))?;
+        Ok(false)
+    }
+}

src/query/service/src/pipelines/processors/transforms/mod.rs

@@ -13,6 +13,8 @@
 // limitations under the License.
 
 pub mod aggregator;
+#[allow(dead_code)]
+mod broadcast;
 mod hash_join;
 pub(crate) mod range_join;
 mod runtime_pool;
@@ -40,6 +42,8 @@ mod transform_udf_script;
 mod transform_udf_server;
 mod window;
 
+pub use broadcast::BroadcastSinkProcessor;
+pub use broadcast::BroadcastSourceProcessor;
 pub use hash_join::*;
 pub use transform_add_computed_columns::TransformAddComputedColumns;
 pub use transform_add_const_columns::TransformAddConstColumns;

src/query/service/src/pipelines/processors/transforms/transform_recursive_cte_source.rs

@@ -351,7 +351,9 @@ async fn create_memory_table_for_cte_scan(
         | PhysicalPlan::ChunkFillAndReorder(_)
         | PhysicalPlan::ChunkAppendData(_)
         | PhysicalPlan::ChunkMerge(_)
-        | PhysicalPlan::ChunkCommitInsert(_) => {}
+        | PhysicalPlan::ChunkCommitInsert(_)
+        | PhysicalPlan::BroadcastSource(_)
+        | PhysicalPlan::BroadcastSink(_) => {}
     }
     Ok(())
 }

src/query/service/src/schedulers/fragments/fragmenter.rs

@@ -123,9 +123,13 @@ impl Fragmenter {
 
     pub fn build_fragment(mut self, plan: &PhysicalPlan) -> Result<PlanFragment> {
         let root = self.replace(plan)?;
+        let fragment_type = match plan {
+            PhysicalPlan::BroadcastSink(_) => FragmentType::Intermediate,
+            _ => FragmentType::Root,
+        };
         let mut root_fragment = PlanFragment {
             plan: root,
-            fragment_type: FragmentType::Root,
+            fragment_type,
             fragment_id: self.ctx.get_fragment_id(),
             exchange: None,
             query_id: self.query_id.clone(),
@@ -218,6 +222,7 @@ impl PhysicalPlanReplacer for Fragmenter {
         fragments.append(&mut self.fragments);
         let probe_input = self.replace(plan.probe.as_ref())?;
         fragments.append(&mut self.fragments);
+
         self.fragments = fragments;
 
         Ok(PhysicalPlan::HashJoin(HashJoin {

src/query/service/src/schedulers/fragments/query_fragment_actions.rs

@@ -126,12 +126,16 @@ impl QueryFragmentsActions {
         self.ctx.get_cluster().local_id()
     }
 
-    pub fn get_root_actions(&self) -> Result<&QueryFragmentActions> {
-        self.fragments_actions.last().ok_or_else(|| {
-            ErrorCode::Internal(
-                "Logical error, call get_root_actions in empty QueryFragmentsActions",
-            )
-        })
+    pub fn get_root_fragment_ids(&self) -> Result<Vec<usize>> {
+        let mut fragment_ids = Vec::new();
+        for fragment_actions in &self.fragments_actions {
+            let plan = &fragment_actions.fragment_actions[0].physical_plan;
+            if !matches!(plan, PhysicalPlan::ExchangeSink(_)) {
+                fragment_ids.push(fragment_actions.fragment_id);
+            }
+        }
+
+        Ok(fragment_ids)
     }
 
     pub fn pop_root_actions(&mut self) -> Option<QueryFragmentActions> {

src/query/service/src/schedulers/scheduler.rs

@@ -17,6 +17,7 @@ use std::sync::Arc;
 use async_trait::async_trait;
 use databend_common_exception::Result;
 use databend_common_expression::DataBlock;
+use databend_common_sql::executor::build_broadcast_plans;
 use databend_common_sql::planner::QueryExecutor;
 use databend_common_sql::Planner;
 use futures_util::TryStreamExt;
@@ -99,11 +100,15 @@ pub async fn build_distributed_pipeline(
     ctx: &Arc<QueryContext>,
     plan: &PhysicalPlan,
 ) -> Result<PipelineBuildResult> {
-    let fragmenter = Fragmenter::try_create(ctx.clone())?;
-
-    let root_fragment = fragmenter.build_fragment(plan)?;
     let mut fragments_actions = QueryFragmentsActions::create(ctx.clone());
-    root_fragment.get_actions(ctx.clone(), &mut fragments_actions)?;
+    for plan in build_broadcast_plans(ctx.as_ref())?
+        .iter()
+        .chain(std::iter::once(plan))
+    {
+        let fragmenter = Fragmenter::try_create(ctx.clone())?;
+        let root_fragment = fragmenter.build_fragment(plan)?;
+        root_fragment.get_actions(ctx.clone(), &mut fragments_actions)?;
+    }
 
     let exchange_manager = ctx.get_exchange_manager();