fix

Author: SkyFan2002

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

@@ -15,7 +15,6 @@
 use databend_common_exception::Result;
 use databend_common_sql::executor::physical_plans::BroadcastSink;
 use databend_common_sql::executor::physical_plans::BroadcastSource;
-use databend_common_storages_fuse::TableContext;
 
 use crate::pipelines::processors::transforms::BroadcastSinkProcessor;
 use crate::pipelines::processors::transforms::BroadcastSourceProcessor;
@@ -33,13 +32,8 @@ impl PipelineBuilder {
     pub(crate) fn build_broadcast_sink(&mut self, sink: &BroadcastSink) -> Result<()> {
         self.build_pipeline(&sink.input)?;
         self.main_pipeline.resize(1, true)?;
-        let node_num = self.ctx.get_cluster().nodes.len();
         self.main_pipeline.add_sink(|input| {
-            BroadcastSinkProcessor::create(
-                input,
-                node_num,
-                self.ctx.broadcast_sink_sender(sink.broadcast_id),
-            )
+            BroadcastSinkProcessor::create(input, self.ctx.broadcast_sink_sender(sink.broadcast_id))
         })
     }
 }

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

@@ -69,19 +69,13 @@ impl AsyncSource for BroadcastSourceProcessor {
 }
 
 pub struct BroadcastSinkProcessor {
-    expect_num: usize,
     received: Vec<BlockMetaInfoPtr>,
-    sender: Sender<Vec<BlockMetaInfoPtr>>,
+    sender: Sender<BlockMetaInfoPtr>,
 }
 
 impl BroadcastSinkProcessor {
-    pub fn create(
-        input: Arc<InputPort>,
-        node_num: usize,
-        sender: Sender<Vec<BlockMetaInfoPtr>>,
-    ) -> Result<ProcessorPtr> {
+    pub fn create(input: Arc<InputPort>, sender: Sender<BlockMetaInfoPtr>) -> Result<ProcessorPtr> {
         Ok(ProcessorPtr::create(AsyncSinker::create(input, Self {
-            expect_num: node_num,
             received: vec![],
             sender,
         })))
@@ -95,6 +89,7 @@ impl AsyncSink for BroadcastSinkProcessor {
     const NAME: &'static str = "BroadcastSink";
 
     async fn on_finish(&mut self) -> Result<()> {
+        self.sender.close();
         Ok(())
     }
 
@@ -103,8 +98,10 @@ impl AsyncSink for BroadcastSinkProcessor {
             .take_meta()
             .ok_or_else(|| ErrorCode::Internal("Cannot downcast meta to BroadcastMeta"))?;
         log::info!("BroadcastSinkProcessor recv meta: {:?}", meta);
-        self.received.push(meta);
-        let all_recv = self.expect_num == self.received.len();
-        Ok(all_recv)
+        self.sender
+            .send(meta)
+            .await
+            .map_err(|_| ErrorCode::Internal("BroadcastSinkProcessor send error"))?;
+        Ok(false)
     }
 }

src/query/service/src/sessions/query_ctx.rs

@@ -295,10 +295,10 @@ impl QueryContext {
         self.shared.broadcast_source_sender(broadcast_id)
     }
 
-    pub fn broadcast_sink_receiver(&self, broadcast_id: u32) -> Receiver<Vec<BlockMetaInfoPtr>> {
+    pub fn broadcast_sink_receiver(&self, broadcast_id: u32) -> Receiver<BlockMetaInfoPtr> {
         self.shared.broadcast_sink_receiver(broadcast_id)
     }
-    pub fn broadcast_sink_sender(&self, broadcast_id: u32) -> Sender<Vec<BlockMetaInfoPtr>> {
+    pub fn broadcast_sink_sender(&self, broadcast_id: u32) -> Sender<BlockMetaInfoPtr> {
         self.shared.broadcast_sink_sender(broadcast_id)
     }
 

src/query/service/src/sessions/query_ctx_shared.rs

@@ -183,8 +183,8 @@ pub struct QueryContextShared {
 pub struct BroadcastChannel {
     pub source_sender: Option<Sender<BlockMetaInfoPtr>>,
     pub source_receiver: Option<Receiver<BlockMetaInfoPtr>>,
-    pub sink_sender: Option<Sender<Vec<BlockMetaInfoPtr>>>,
-    pub sink_receiver: Option<Receiver<Vec<BlockMetaInfoPtr>>>,
+    pub sink_sender: Option<Sender<BlockMetaInfoPtr>>,
+    pub sink_receiver: Option<Receiver<BlockMetaInfoPtr>>,
 }
 
 impl QueryContextShared {
@@ -285,7 +285,7 @@ impl QueryContextShared {
         }
     }
 
-    pub fn broadcast_sink_receiver(&self, broadcast_id: u32) -> Receiver<Vec<BlockMetaInfoPtr>> {
+    pub fn broadcast_sink_receiver(&self, broadcast_id: u32) -> Receiver<BlockMetaInfoPtr> {
         let mut broadcast_channels = self.broadcast_channels.lock();
         let entry = broadcast_channels.entry(broadcast_id).or_default();
         match entry.sink_receiver.take() {
@@ -297,7 +297,7 @@ impl QueryContextShared {
             }
         }
     }
-    pub fn broadcast_sink_sender(&self, broadcast_id: u32) -> Sender<Vec<BlockMetaInfoPtr>> {
+    pub fn broadcast_sink_sender(&self, broadcast_id: u32) -> Sender<BlockMetaInfoPtr> {
         let mut broadcast_channels = self.broadcast_channels.lock();
         let entry = broadcast_channels.entry(broadcast_id).or_default();
         match entry.sink_sender.take() {