fix

Author: SkyFan2002

src/query/catalog/src/table_context.rs

@@ -420,6 +420,10 @@ pub trait TableContext: Send + Sync {
     }
 
     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/schedulers/scheduler.rs

@@ -101,7 +101,7 @@ pub async fn build_distributed_pipeline(
     plan: &PhysicalPlan,
 ) -> Result<PipelineBuildResult> {
     let mut fragments_actions = QueryFragmentsActions::create(ctx.clone());
-    for plan in build_broadcast_plans(ctx.get_next_broadcast_id())?
+    for plan in build_broadcast_plans(ctx.as_ref())?
         .iter()
         .chain(std::iter::once(plan))
     {

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

@@ -1861,7 +1861,13 @@ impl TableContext for QueryContext {
     }
 
     fn get_next_broadcast_id(&self) -> u32 {
-        self.shared.get_next_broadcast_id()
+        self.shared
+            .next_broadcast_id
+            .fetch_add(1, Ordering::Acquire)
+    }
+
+    fn reset_broadcast_id(&self) {
+        self.shared.next_broadcast_id.store(0, Ordering::Release);
     }
 }
 

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

@@ -256,10 +256,6 @@ impl QueryContextShared {
         }))
     }
 
-    pub fn get_next_broadcast_id(&self) -> u32 {
-        self.next_broadcast_id.fetch_add(1, Ordering::AcqRel)
-    }
-
     pub fn broadcast_source_receiver(&self, broadcast_id: u32) -> Receiver<BlockMetaInfoPtr> {
         let mut broadcast_channels = self.broadcast_channels.lock();
         let entry = broadcast_channels.entry(broadcast_id).or_default();

src/query/sql/src/executor/physical_plans/physical_broadcast.rs

@@ -12,6 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+use databend_common_catalog::table_context::TableContext;
 use databend_common_exception::Result;
 
 use super::Exchange;
@@ -52,8 +53,10 @@ pub fn build_broadcast_plan(broadcast_id: u32) -> Result<PhysicalPlan> {
     Ok(broadcast_sink)
 }
 
-pub fn build_broadcast_plans(next_broadcast_id: u32) -> Result<Vec<PhysicalPlan>> {
+pub fn build_broadcast_plans(ctx: &dyn TableContext) -> Result<Vec<PhysicalPlan>> {
     let mut plans = vec![];
+    let next_broadcast_id = ctx.get_next_broadcast_id();
+    ctx.reset_broadcast_id();
     for broadcast_id in 0..next_broadcast_id {
         plans.push(build_broadcast_plan(broadcast_id)?);
     }