fix(query): avoid waiting idle fragment nodes

Author: SkyFan2002

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

@@ -161,7 +161,7 @@ impl PlanFragment {
 
         let data_sources = self.collect_data_sources()?;
 
-        let executors = Fragmenter::get_executors_nodes(ctx);
+        let executors = Fragmenter::get_executors_nodes(ctx.clone());
 
         let mut executor_partitions: HashMap<String, HashMap<u32, DataSource>> = HashMap::new();
 
@@ -213,12 +213,38 @@ impl PlanFragment {
             }
         }
 
+        let local_executor = Fragmenter::get_local_executor(ctx);
+        let mut empty_action_sources = None;
+        let mut non_empty_actions = Vec::new();
+
         for (executor, sources) in executor_partitions {
+            let is_empty = sources.values().all(DataSource::is_empty);
+
+            if is_empty {
+                if executor == local_executor {
+                    empty_action_sources = Some((executor, sources));
+                }
+                continue;
+            }
+
+            non_empty_actions.push((executor, sources));
+        }
+
+        if non_empty_actions.is_empty() {
+            if let Some((executor, sources)) = empty_action_sources {
+                let mut handle = ReadSourceDeriveHandle::create(sources);
+                let plan = self.plan.derive_with(&mut handle);
+                fragment_actions.add_action(QueryFragmentAction::create(executor, plan));
+            }
+            return Ok(());
+        }
+
+        for (executor, sources) in non_empty_actions {
             // Replace `ReadDataSourcePlan` with rewritten one and generate new fragment for it.
             let mut handle = ReadSourceDeriveHandle::create(sources);
             let plan = self.plan.derive_with(&mut handle);
 
-            fragment_actions.add_action(QueryFragmentAction::create(executor.clone(), plan));
+            fragment_actions.add_action(QueryFragmentAction::create(executor, plan));
         }
 
         Ok(())
@@ -564,6 +590,15 @@ enum DataSource {
     ConstTable(ConstTableColumn),
 }
 
+impl DataSource {
+    fn is_empty(&self) -> bool {
+        match self {
+            DataSource::Table(plan) => plan.parts.is_empty(),
+            DataSource::ConstTable(columns) => columns.num_rows == 0,
+        }
+    }
+}
+
 impl TryFrom<DataSource> for DataSourcePlan {
     type Error = ErrorCode;
 

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

@@ -13,6 +13,7 @@
 // limitations under the License.
 
 use std::collections::HashMap;
+use std::collections::HashSet;
 use std::collections::hash_map::Entry;
 use std::fmt::Debug;
 use std::fmt::Formatter;
@@ -189,7 +190,8 @@ impl QueryFragmentsActions {
 
     pub fn get_query_env(&self) -> Result<QueryEnv> {
         let workload_group = self.ctx.get_current_session().get_current_workload_group();
-        let mut builder = DataflowDiagramBuilder::create(self.ctx.get_cluster().nodes.clone());
+        let mut builder =
+            DataflowDiagramBuilder::create(self.query_env_nodes().into_values().collect());
 
         self.fragments_connections(&mut builder)?;
         self.statistics_connections(&mut builder)?;
@@ -213,7 +215,7 @@ impl QueryFragmentsActions {
     }
 
     pub fn prepared_query(&self) -> Result<HashMap<String, String>> {
-        let nodes_info = Self::nodes_info(&self.ctx);
+        let nodes_info = self.execution_nodes();
         let mut execute_partial_query_packets = HashMap::with_capacity(nodes_info.len());
 
         for node_id in nodes_info.keys() {
@@ -258,13 +260,49 @@ impl QueryFragmentsActions {
     fn statistics_connections(&self, builder: &mut DataflowDiagramBuilder) -> Result<()> {
         let local_id = self.ctx.get_cluster().local_id.clone();
 
-        for (_id, node_info) in Self::nodes_info(&self.ctx) {
+        for (_id, node_info) in self.execution_nodes() {
             builder.add_statistics_edge(&node_info.id, &local_id)?;
         }
 
         Ok(())
     }
 
+    fn query_env_nodes(&self) -> HashMap<String, Arc<NodeInfo>> {
+        let mut node_ids = self.execution_node_ids();
+        node_ids.insert(self.get_local_executor());
+
+        for fragment_actions in &self.fragments_actions {
+            if let Some(exchange) = &fragment_actions.data_exchange {
+                node_ids.extend(exchange.get_destinations());
+            }
+        }
+
+        Self::nodes_info(&self.ctx)
+            .into_iter()
+            .filter(|(node_id, _)| node_ids.contains(node_id))
+            .collect()
+    }
+
+    fn execution_nodes(&self) -> HashMap<String, Arc<NodeInfo>> {
+        let node_ids = self.execution_node_ids();
+        Self::nodes_info(&self.ctx)
+            .into_iter()
+            .filter(|(node_id, _)| node_ids.contains(node_id))
+            .collect()
+    }
+
+    fn execution_node_ids(&self) -> HashSet<String> {
+        let mut node_ids = HashSet::new();
+
+        for fragment_actions in &self.fragments_actions {
+            for fragment_action in &fragment_actions.fragment_actions {
+                node_ids.insert(fragment_action.executor.clone());
+            }
+        }
+
+        node_ids
+    }
+
     fn nodes_info(ctx: &Arc<QueryContext>) -> HashMap<String, Arc<NodeInfo>> {
         let nodes = ctx.get_cluster().get_nodes();
         let mut nodes_info = HashMap::with_capacity(nodes.len());

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

@@ -899,6 +899,36 @@ impl DataExchangeManager {
         }
     }
 
+    pub fn get_merge_flight_receiver_if_exists(
+        &self,
+        params: &MergeExchangeParams,
+    ) -> Result<Vec<FlightReceiver>> {
+        let queries_coordinator_guard = self.queries_coordinator.lock();
+        let queries_coordinator = unsafe { &mut *queries_coordinator_guard.deref().get() };
+
+        match queries_coordinator.get_mut(&params.query_id) {
+            None => Err(ErrorCode::Internal("Query not exists.")),
+            Some(coordinator) => {
+                Ok(params.take_flight_receiver_if_exists(&mut coordinator.flight_data_receivers))
+            }
+        }
+    }
+
+    pub fn get_flight_receiver_if_exists(
+        &self,
+        params: &ShuffleExchangeParams,
+    ) -> Result<Vec<FlightReceiver>> {
+        let queries_coordinator_guard = self.queries_coordinator.lock();
+        let queries_coordinator = unsafe { &mut *queries_coordinator_guard.deref().get() };
+
+        match queries_coordinator.get_mut(&params.query_id) {
+            None => Err(ErrorCode::Internal("Query not exists.")),
+            Some(coordinator) => {
+                Ok(params.take_flight_receiver_if_exists(&mut coordinator.flight_data_receivers))
+            }
+        }
+    }
+
     pub fn get_fragment_source(
         &self,
         query_id: &str,

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

@@ -152,6 +152,13 @@ impl MergeExchangeParams {
 
         Ok(receivers)
     }
+
+    pub fn take_flight_receiver_if_exists(
+        &self,
+        receivers: &mut HashMap<String, Vec<FlightReceiver>>,
+    ) -> Vec<FlightReceiver> {
+        receivers.remove(&self.channel_id).unwrap_or_default()
+    }
 }
 
 impl BroadcastExchangeParams {
@@ -251,15 +258,36 @@ impl ShuffleExchangeParams {
         for (destination, channels) in &self.destination_channels {
             if destination == &self.executor_id {
                 for channel in channels {
-                    // Local shuffle data stays in-process. Missing receivers here only
-                    // means no remote executor targets this local destination.
+                    let Some(receivers) = receivers.remove(channel) else {
+                        return Err(ErrorCode::UnknownFragmentExchange(format!(
+                            "Unknown fragment flight receiver, {}, {}",
+                            self.executor_id, self.fragment_id
+                        )));
+                    };
+                    exchanges.extend(receivers.into_iter());
+                }
+            }
+        }
+
+        Ok(exchanges)
+    }
+
+    pub fn take_flight_receiver_if_exists(
+        &self,
+        receivers: &mut HashMap<String, Vec<FlightReceiver>>,
+    ) -> Vec<FlightReceiver> {
+        let mut exchanges = Vec::with_capacity(self.destination_channels.len());
+
+        for (destination, channels) in &self.destination_channels {
+            if destination == &self.executor_id {
+                for channel in channels {
                     if let Some(receivers) = receivers.remove(channel) {
                         exchanges.extend(receivers.into_iter());
                     }
                 }
             }
         }
 
-        Ok(exchanges)
+        exchanges
     }
 }

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

@@ -17,23 +17,16 @@ use std::sync::Arc;
 use databend_common_base::runtime::GlobalIORuntime;
 use databend_common_exception::ErrorCode;
 use databend_common_exception::Result;
-use databend_common_expression::BlockMetaInfoDowncast;
-use databend_common_expression::DataBlock;
 use databend_common_pipeline::core::Pipe;
-use databend_common_pipeline::core::PipeItem;
 use databend_common_pipeline::core::Pipeline;
-use databend_common_pipeline::core::ProcessorPtr;
 
 use super::broadcast_send_sink::BroadcastSendSink;
 use super::exchange_params::BroadcastExchangeParams;
 use super::exchange_params::ExchangeParams;
 use super::exchange_params::GlobalExchangeParams;
 use super::exchange_sink_writer::create_writer_item;
-use super::exchange_sorting::ExchangeSorting;
-use super::exchange_sorting::TransformExchangeSorting;
 use super::exchange_transform_shuffle::exchange_shuffle;
 use super::hash_send_sink::HashSendSink;
-use super::serde::ExchangeSerializeMeta;
 use crate::clusters::ClusterHelper;
 use crate::servers::flight::v1::exchange::DataExchangeManager;
 use crate::servers::flight::v1::network::OutboundChannel;
@@ -65,28 +58,6 @@ impl ExchangeSink {
                     )));
                 }
 
-                let exchange_injector = &params.exchange_injector;
-
-                if !params.ignore_exchange {
-                    let settings = ctx.get_settings();
-                    let compression = settings.get_query_flight_compression()?;
-                    exchange_injector.apply_merge_serializer(params, compression, pipeline)?;
-                }
-
-                if !params.ignore_exchange && exchange_injector.exchange_sorting().is_some() {
-                    let output_len = pipeline.output_len();
-                    let sorting = SinkExchangeSorting::create();
-                    let transform = TransformExchangeSorting::create(output_len, sorting);
-
-                    let output = transform.get_output();
-                    let inputs = transform.get_inputs();
-                    pipeline.add_pipe(Pipe::create(output_len, 1, vec![PipeItem::create(
-                        ProcessorPtr::create(Box::new(transform)),
-                        inputs,
-                        vec![output],
-                    )]));
-                }
-
                 let exchange_manager = ctx.get_exchange_manager();
                 let senders = exchange_manager
                     .get_flight_sender(&ExchangeParams::MergeExchange(params.clone()))?;
@@ -230,30 +201,6 @@ impl ExchangeSink {
     }
 }
 
-struct SinkExchangeSorting;
-
-impl SinkExchangeSorting {
-    pub fn create() -> Arc<dyn ExchangeSorting> {
-        Arc::new(SinkExchangeSorting {})
-    }
-}
-
-impl ExchangeSorting for SinkExchangeSorting {
-    fn block_number(&self, data_block: &DataBlock) -> Result<isize> {
-        let block_meta = data_block.get_meta();
-        let shuffle_meta = block_meta
-            .and_then(ExchangeSerializeMeta::downcast_ref_from)
-            .ok_or_else(|| {
-                ErrorCode::Internal(format!(
-                    "Failed to downcast ExchangeSerializeMeta from BlockMeta: {:?}",
-                    block_meta
-                ))
-            })?;
-
-        Ok(shuffle_meta.block_number)
-    }
-}
-
 /// Build OutboundChannels for broadcast exchange using PingPongExchange.
 pub(super) fn build_broadcast_outbound_channels(
     params: &BroadcastExchangeParams,

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

@@ -62,9 +62,17 @@ pub fn via_exchange_source(
 
     let exchange_params = ExchangeParams::MergeExchange(params.clone());
     let exchange_manager = ctx.get_exchange_manager();
-    let flight_receivers = exchange_manager.get_flight_receiver(&exchange_params)?;
+    let had_local_outputs = pipeline.output_len() > 0;
 
     let last_output_len = pipeline.output_len();
+    let flight_receivers = if had_local_outputs {
+        // A locally subscribed merge fragment can legitimately have no remote
+        // senders; in that case the local in-process output is the entire
+        // fragment input for the root.
+        exchange_manager.get_merge_flight_receiver_if_exists(params)?
+    } else {
+        exchange_manager.get_flight_receiver(&exchange_params)?
+    };
     let mut items = Vec::with_capacity(last_output_len + flight_receivers.len());
 
     for _index in 0..last_output_len {

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

@@ -103,7 +103,10 @@ impl ExchangeTransform {
         }
 
         let mut nodes_source = 0;
-        let receivers = exchange_manager.get_flight_receiver(&exchange_params)?;
+        // When the shuffle source fragment also runs locally, the local bucket
+        // stays in-process via the dummy/resize branch above. In that case it is
+        // valid for this executor to have no remote receiver registered.
+        let receivers = exchange_manager.get_flight_receiver_if_exists(params)?;
         for receiver in receivers {
             nodes_source += 1;
             items.push(create_reader_item(receiver));