fix(exchange): precreate inbound channel sets

Register and lazily create inbound channel sets for do_exchange
receivers so remote-only exchange sources can build pipelines before
network senders connect.

- pre-register incoming exchange fragments during query env init
- add get-or-create helpers on DataExchangeManager and QueryCoordinator
- switch broadcast/global-shuffle receive paths to use get-or-create

Validation:
- cargo check -p databend-query --lib
- db-slt --cluster --run-dir cluster --run-file subquery.test

Author: SkyFan2002

src/query/service/src/servers/flight/v1/exchange/exchange_manager.rs

@@ -249,6 +249,7 @@ impl DataExchangeManager {
 
         let mut request_exchanges = HashMap::new();
         let mut targets_exchanges = HashMap::<String, Vec<FlightExchange>>::new();
+        let mut inbound_channel_sizes = HashMap::<String, usize>::new();
 
         for index in env.dataflow_diagram.node_indices() {
             if env.dataflow_diagram[index].id == config.query.node_id {
@@ -297,9 +298,25 @@ impl DataExchangeManager {
                                 })
                             }));
                         }
-                        Edge::ExchangeFragment { .. } => {
-                            // Skip: remote sender will call do_exchange on us,
-                            // handled by handle_do_exchange → NetworkInboundSender
+                        Edge::ExchangeFragment {
+                            exchange_id,
+                            channels,
+                        } => {
+                            // Pre-register inbound channel sets so source-only exchange
+                            // receivers can build their pipelines before the first
+                            // do_exchange connection arrives.
+                            match inbound_channel_sizes.entry(exchange_id) {
+                                Entry::Occupied(v) => {
+                                    if *v.get() != channels.len() {
+                                        return Err(ErrorCode::Internal(
+                                            "Mismatched inbound exchange parallelism",
+                                        ));
+                                    }
+                                }
+                                Entry::Vacant(v) => {
+                                    v.insert(channels.len());
+                                }
+                            }
                         }
                     }
                 }
@@ -419,6 +436,7 @@ impl DataExchangeManager {
                         query_coordinator.info = query_info;
                         query_coordinator.is_request_server =
                             GlobalConfig::instance().query.node_id == env.request_server_id;
+                        query_coordinator.register_inbound_channel_sets(&inbound_channel_sizes)?;
                         query_coordinator.register_flight_channel_receiver(targets_exchanges)?;
                         query_coordinator.register_ping_pong_exchanges(ping_pong_exchanges);
                         query_coordinator.add_statistics_exchanges(request_exchanges)?;
@@ -428,6 +446,7 @@ impl DataExchangeManager {
                         query_coordinator.info = query_info;
                         query_coordinator.is_request_server =
                             GlobalConfig::instance().query.node_id == env.request_server_id;
+                        query_coordinator.register_inbound_channel_sets(&inbound_channel_sizes)?;
                         query_coordinator.register_flight_channel_receiver(targets_exchanges)?;
                         query_coordinator.register_ping_pong_exchanges(ping_pong_exchanges);
                         query_coordinator.add_statistics_exchanges(request_exchanges)?;
@@ -672,6 +691,25 @@ impl DataExchangeManager {
         }
     }
 
+    pub fn get_or_create_exchange_channel_set(
+        &self,
+        query_id: &str,
+        channel_id: &str,
+        num_threads: usize,
+    ) -> Result<Arc<NetworkInboundChannelSet>> {
+        let queries_coordinator_guard = self.queries_coordinator.lock();
+        let queries_coordinator = unsafe { &mut *queries_coordinator_guard.deref().get() };
+
+        match queries_coordinator.entry(query_id.to_string()) {
+            Entry::Occupied(mut v) => v
+                .get_mut()
+                .get_or_create_inbound_channel_set(channel_id, num_threads),
+            Entry::Vacant(v) => v
+                .insert(QueryCoordinator::create())
+                .get_or_create_inbound_channel_set(channel_id, num_threads),
+        }
+    }
+
     /// Take the PingPongExchanges for a given query and channel.
     ///
     /// Returns the exchanges that were created during init_query_env.
@@ -1034,6 +1072,41 @@ impl QueryCoordinator {
         }
     }
 
+    pub fn register_inbound_channel_sets(
+        &mut self,
+        channel_sizes: &HashMap<String, usize>,
+    ) -> Result<()> {
+        for (channel_id, num_threads) in channel_sizes {
+            self.get_or_create_inbound_channel_set(channel_id, *num_threads)?;
+        }
+
+        Ok(())
+    }
+
+    pub fn get_or_create_inbound_channel_set(
+        &mut self,
+        channel_id: &str,
+        num_threads: usize,
+    ) -> Result<Arc<NetworkInboundChannelSet>> {
+        match self.inbound_channel_sets.entry(channel_id.to_string()) {
+            Entry::Occupied(v) => {
+                if v.get().channels.len() != num_threads {
+                    return Err(ErrorCode::Internal(format!(
+                        "Mismatched inbound channel set parallelism, channel_id: {}, expected: {}, actual: {}",
+                        channel_id,
+                        num_threads,
+                        v.get().channels.len()
+                    )));
+                }
+
+                Ok(v.get().clone())
+            }
+            Entry::Vacant(v) => Ok(v
+                .insert(Arc::new(NetworkInboundChannelSet::new(num_threads)))
+                .clone()),
+        }
+    }
+
     /// Create a NetworkInboundSender for a new do_exchange connection.
     ///
     /// The `num_threads` value is provided by the coordinator via DoExchangeParams.

src/query/service/src/servers/flight/v1/exchange/exchange_sink.rs

@@ -156,7 +156,11 @@ impl ExchangeSink {
         let exchange_id = &params.exchange_id;
         let exchange_manager = DataExchangeManager::instance();
 
-        let channel_set = exchange_manager.get_exchange_channel_set(query_id, exchange_id)?;
+        let channel_set = exchange_manager.get_or_create_exchange_channel_set(
+            query_id,
+            exchange_id,
+            local_threads,
+        )?;
         assert_eq!(channel_set.channels.len(), local_threads);
 
         let local_outbound = create_local_channels(&channel_set);

src/query/service/src/servers/flight/v1/exchange/exchange_source.rs

@@ -62,6 +62,22 @@ fn add_source_pipe(pipeline: &mut Pipeline, source_items: Vec<PipeItem>) {
     pipeline.add_pipe(Pipe::create(last_output_len, items.len(), items));
 }
 
+fn local_exchange_threads(
+    executor_id: &str,
+    destination_channels: &[(String, Vec<String>)],
+) -> Result<usize> {
+    destination_channels
+        .iter()
+        .find_map(|(destination, channels)| (destination == executor_id).then_some(channels.len()))
+        .filter(|threads| *threads > 0)
+        .ok_or_else(|| {
+            ErrorCode::Internal(format!(
+                "Executor {} is not a destination of exchange",
+                executor_id
+            ))
+        })
+}
+
 /// Add Exchange Source to the pipeline.
 pub fn via_exchange_source(
     ctx: Arc<QueryContext>,
@@ -117,8 +133,13 @@ pub fn via_hash_exchange_source(
     let query_id = &params.query_id;
     let exchange_id = &params.exchange_id;
     let exchange_manager = DataExchangeManager::instance();
+    let local_threads = local_exchange_threads(&params.executor_id, &params.destination_channels)?;
 
-    let channel_set = exchange_manager.get_exchange_channel_set(query_id, exchange_id)?;
+    let channel_set = exchange_manager.get_or_create_exchange_channel_set(
+        query_id,
+        exchange_id,
+        local_threads,
+    )?;
     let waker = pipeline.get_waker();
 
     let num_receivers = channel_set.channels.len();
@@ -137,12 +158,17 @@ pub fn via_hash_exchange_source(
 }
 
 pub fn via_broadcast_exchange_source(
+    _ctx: &Arc<QueryContext>,
     params: &BroadcastExchangeParams,
     pipeline: &mut Pipeline,
 ) -> Result<()> {
     let exchange_manager = DataExchangeManager::instance();
-    let channel_set =
-        exchange_manager.get_exchange_channel_set(&params.query_id, &params.exchange_id)?;
+    let local_threads = local_exchange_threads(&params.executor_id, &params.destination_channels)?;
+    let channel_set = exchange_manager.get_or_create_exchange_channel_set(
+        &params.query_id,
+        &params.exchange_id,
+        local_threads,
+    )?;
     let waker = pipeline.get_waker();
     let mut source_items = Vec::with_capacity(channel_set.channels.len());
 
@@ -159,27 +185,15 @@ pub fn via_broadcast_exchange_source(
 }
 
 pub fn via_shuffle_exchange_source(
-    ctx: &Arc<QueryContext>,
+    _ctx: &Arc<QueryContext>,
     params: &ShuffleExchangeParams,
     injector: Arc<dyn ExchangeInjector>,
     pipeline: &mut Pipeline,
 ) -> Result<()> {
-    let local_channels = params
-        .destination_channels
-        .iter()
-        .find_map(|(destination, channels)| {
-            (destination == &params.executor_id).then_some(channels.len())
-        })
-        .unwrap_or_default();
-
-    if local_channels == 0 {
-        return Err(ErrorCode::Internal(format!(
-            "Executor {} is not a destination of fragment {} exchange",
-            params.executor_id, params.fragment_id
-        )));
-    }
+    let _local_channels =
+        local_exchange_threads(&params.executor_id, &params.destination_channels)?;
 
-    let exchange_manager = ctx.get_exchange_manager();
+    let exchange_manager = DataExchangeManager::instance();
     let exchange_params = ExchangeParams::NodeShuffleExchange(params.clone());
     let flight_receivers = exchange_manager.get_flight_receiver(&exchange_params)?;
     let mut source_items = Vec::with_capacity(flight_receivers.len());
@@ -255,6 +269,7 @@ pub fn via_remote_exchange_source(
             pipeline,
         ),
         DataExchange::Broadcast(exchange) => via_broadcast_exchange_source(
+            &ctx,
             &create_broadcast_source_params(&ctx, query_id, schema, exchange),
             pipeline,
         ),

src/query/service/src/servers/flight/v1/exchange/exchange_transform.rs

@@ -144,7 +144,11 @@ impl ExchangeTransform {
         let exchange_id = &params.exchange_id;
         let exchange_manager = DataExchangeManager::instance();
 
-        let channel_set = exchange_manager.get_exchange_channel_set(query_id, exchange_id)?;
+        let channel_set = exchange_manager.get_or_create_exchange_channel_set(
+            query_id,
+            exchange_id,
+            local_threads,
+        )?;
 
         assert_eq!(channel_set.channels.len(), local_threads);
 
@@ -205,7 +209,11 @@ impl ExchangeTransform {
         let exchange_id = &params.exchange_id;
         let exchange_manager = DataExchangeManager::instance();
 
-        let channel_set = exchange_manager.get_exchange_channel_set(query_id, exchange_id)?;
+        let channel_set = exchange_manager.get_or_create_exchange_channel_set(
+            query_id,
+            exchange_id,
+            local_threads,
+        )?;
         assert_eq!(channel_set.channels.len(), local_threads);
 
         let local_outbound = create_local_channels(&channel_set);