update

Author: SkyFan2002

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

@@ -82,6 +82,53 @@ pub struct PlanFragment {
 }
 
 impl PlanFragment {
+    fn requires_coordinator_source(&self) -> bool {
+        self.source_fragments
+            .iter()
+            .any(|fragment| !fragment.can_subscribe_on_all_executors())
+    }
+
+    fn can_subscribe_on_all_executors(&self) -> bool {
+        match &self.exchange {
+            Some(DataExchange::Merge(_)) => false,
+            Some(DataExchange::Broadcast(_))
+            | Some(DataExchange::NodeToNodeExchange(_))
+            | Some(DataExchange::GlobalShuffleExchange(_)) => true,
+            None => !self.requires_coordinator_source(),
+        }
+    }
+
+    fn add_source_only_actions(
+        &self,
+        ctx: Arc<QueryContext>,
+        fragment_actions: &mut QueryFragmentActions,
+    ) {
+        let Some(exchange) = &self.exchange else {
+            return;
+        };
+
+        if !matches!(
+            exchange,
+            DataExchange::Broadcast(_)
+                | DataExchange::NodeToNodeExchange(_)
+                | DataExchange::GlobalShuffleExchange(_)
+        ) {
+            return;
+        }
+
+        let local_executor = Fragmenter::get_local_executor(ctx);
+        for destination in exchange.get_destinations() {
+            if destination == local_executor {
+                continue;
+            }
+
+            fragment_actions.add_action(QueryFragmentAction::create_source_only(
+                destination,
+                self.plan.clone(),
+            ));
+        }
+    }
+
     pub fn get_actions(
         &self,
         ctx: Arc<QueryContext>,
@@ -92,6 +139,7 @@ impl PlanFragment {
         // }
 
         let mut fragment_actions = QueryFragmentActions::create(self.fragment_id);
+        let requires_coordinator_source = self.requires_coordinator_source();
 
         match &self.fragment_type {
             FragmentType::Root => {
@@ -101,25 +149,19 @@ impl PlanFragment {
                 );
                 fragment_actions.add_action(action);
             }
+            _ if requires_coordinator_source => {
+                let action = QueryFragmentAction::create(
+                    Fragmenter::get_local_executor(ctx.clone()),
+                    self.plan.clone(),
+                );
+                fragment_actions.add_action(action);
+                self.add_source_only_actions(ctx, &mut fragment_actions);
+            }
             FragmentType::Intermediate => {
-                if self
-                    .source_fragments
-                    .iter()
-                    .any(|fragment| matches!(&fragment.exchange, Some(DataExchange::Merge(_))))
-                {
-                    // If this is a intermediate fragment with merge input,
-                    // we will only send it to coordinator node.
-                    let action = QueryFragmentAction::create(
-                        Fragmenter::get_local_executor(ctx),
-                        self.plan.clone(),
-                    );
+                // Otherwise distribute the fragment to all the executors.
+                for executor in Fragmenter::get_executors(ctx) {
+                    let action = QueryFragmentAction::create(executor, self.plan.clone());
                     fragment_actions.add_action(action);
-                } else {
-                    // Otherwise distribute the fragment to all the executors.
-                    for executor in Fragmenter::get_executors(ctx) {
-                        let action = QueryFragmentAction::create(executor, self.plan.clone());
-                        fragment_actions.add_action(action);
-                    }
                 }
             }
             FragmentType::Source => {
@@ -553,6 +595,68 @@ impl PlanFragment {
     }
 }
 
+#[cfg(test)]
+mod tests {
+    use databend_common_expression::DataSchemaRefExt;
+
+    use super::*;
+    use crate::physical_plans::ConstantTableScan;
+    use crate::physical_plans::PhysicalPlanMeta;
+    use crate::servers::flight::v1::exchange::BroadcastExchange;
+    use crate::servers::flight::v1::exchange::MergeExchange;
+
+    fn mock_plan() -> PhysicalPlan {
+        PhysicalPlan::new(ConstantTableScan {
+            meta: PhysicalPlanMeta::new("ConstantTableScan"),
+            values: vec![],
+            num_rows: 0,
+            output_schema: DataSchemaRefExt::create(vec![]),
+        })
+    }
+
+    #[test]
+    fn test_merge_dependent_fragment_can_use_source_only_exchange() {
+        let merge_source = PlanFragment {
+            plan: mock_plan(),
+            fragment_type: FragmentType::Source,
+            fragment_id: 0,
+            exchange: Some(MergeExchange::create(
+                "coordinator".to_string(),
+                false,
+                true,
+            )),
+            query_id: "q".to_string(),
+            source_fragments: vec![],
+        };
+        assert!(!merge_source.requires_coordinator_source());
+        assert!(!merge_source.can_subscribe_on_all_executors());
+
+        let broadcast_fragment = PlanFragment {
+            plan: mock_plan(),
+            fragment_type: FragmentType::Intermediate,
+            fragment_id: 1,
+            exchange: Some(BroadcastExchange::create(
+                vec!["coordinator".to_string(), "worker".to_string()],
+                1,
+            )),
+            query_id: "q".to_string(),
+            source_fragments: vec![merge_source],
+        };
+        assert!(broadcast_fragment.requires_coordinator_source());
+        assert!(broadcast_fragment.can_subscribe_on_all_executors());
+
+        let downstream_fragment = PlanFragment {
+            plan: mock_plan(),
+            fragment_type: FragmentType::Intermediate,
+            fragment_id: 2,
+            exchange: None,
+            query_id: "q".to_string(),
+            source_fragments: vec![broadcast_fragment],
+        };
+        assert!(!downstream_fragment.requires_coordinator_source());
+    }
+}
+
 struct ConstTableColumn {
     columns: Vec<Column>,
     num_rows: usize,

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

@@ -41,13 +41,26 @@ use crate::sessions::TableContext;
 pub struct QueryFragmentAction {
     pub executor: String,
     pub physical_plan: PhysicalPlan,
+    pub source_only: bool,
 }
 
 impl QueryFragmentAction {
     pub fn create(executor: String, physical_plan: PhysicalPlan) -> QueryFragmentAction {
         QueryFragmentAction {
             executor,
             physical_plan,
+            source_only: false,
+        }
+    }
+
+    pub fn create_source_only(
+        executor: String,
+        physical_plan: PhysicalPlan,
+    ) -> QueryFragmentAction {
+        QueryFragmentAction {
+            executor,
+            physical_plan,
+            source_only: true,
         }
     }
 }
@@ -231,6 +244,10 @@ impl QueryFragmentsActions {
                 let use_do_exchange = exchange.use_do_exchange();
 
                 for fragment_action in &fragment_actions.fragment_actions {
+                    if fragment_action.source_only {
+                        continue;
+                    }
+
                     let source = fragment_action.executor.to_string();
 
                     for destination in &destinations {
@@ -284,6 +301,7 @@ impl QueryFragmentsActions {
                     fragment_actions.fragment_id,
                     fragment_actions.data_exchange.clone(),
                     fragment_action.physical_plan.clone(),
+                    fragment_action.source_only,
                 );
 
                 match fragments_packets.entry(fragment_action.executor.clone()) {

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

@@ -53,6 +53,9 @@ use super::exchange_params::GlobalExchangeParams;
 use super::exchange_params::MergeExchangeParams;
 use super::exchange_params::ShuffleExchangeParams;
 use super::exchange_sink::ExchangeSink;
+use super::exchange_source::via_broadcast_exchange_source;
+use super::exchange_source::via_hash_exchange_source;
+use super::exchange_source::via_node_shuffle_exchange_source;
 use super::exchange_transform::ExchangeTransform;
 use super::statistics_receiver::StatisticsReceiver;
 use super::statistics_sender::StatisticsSender;
@@ -1087,6 +1090,11 @@ impl QueryCoordinator {
         // Merge pipelines if exist locally pipeline
         if let Some(mut fragment_coordinator) = self.fragments_coordinator.remove(&fragment_id) {
             let info = self.info.as_ref().expect("QueryInfo is none");
+
+            if fragment_coordinator.source_only {
+                return fragment_coordinator.build_source_only_pipeline(ctx, info, injector);
+            }
+
             fragment_coordinator.prepare_pipeline(ctx.clone())?;
 
             if fragment_coordinator.pipeline_build_res.is_none() {
@@ -1163,16 +1171,28 @@ impl QueryCoordinator {
             }
         }
 
-        if self.fragments_coordinator.is_empty() {
+        let active_fragment_ids = self
+            .fragments_coordinator
+            .iter()
+            .filter_map(|(fragment_id, coordinator)| {
+                (!coordinator.source_only).then_some(*fragment_id)
+            })
+            .collect::<Vec<_>>();
+
+        if active_fragment_ids.is_empty() {
             // Empty fragments if it is a request server, because the pipelines may have been linked.
             return Ok(());
         }
 
         let max_threads = info.query_ctx.get_settings().get_max_threads()?;
-        let mut pipelines = Vec::with_capacity(self.fragments_coordinator.len());
-
-        let mut params = Vec::with_capacity(self.fragments_coordinator.len());
-        for coordinator in self.fragments_coordinator.values() {
+        let mut pipelines = Vec::with_capacity(active_fragment_ids.len());
+
+        let mut params = Vec::with_capacity(active_fragment_ids.len());
+        for fragment_id in &active_fragment_ids {
+            let coordinator = self
+                .fragments_coordinator
+                .get(fragment_id)
+                .expect("active fragment must exist");
             params.push(
                 coordinator.create_exchange_params(
                     info,
@@ -1187,7 +1207,11 @@ impl QueryCoordinator {
             );
         }
 
-        for ((_, coordinator), params) in self.fragments_coordinator.iter_mut().zip(params) {
+        for (fragment_id, params) in active_fragment_ids.into_iter().zip(params) {
+            let coordinator = self
+                .fragments_coordinator
+                .get_mut(&fragment_id)
+                .expect("active fragment must exist");
             if let Some(mut build_res) = coordinator.pipeline_build_res.take() {
                 build_res.set_max_threads(max_threads as usize);
 
@@ -1272,6 +1296,7 @@ impl QueryCoordinator {
 
 struct FragmentCoordinator {
     initialized: bool,
+    source_only: bool,
     fragment_id: usize,
     physical_plan: PhysicalPlan,
     data_exchange: Option<DataExchange>,
@@ -1282,6 +1307,7 @@ impl FragmentCoordinator {
     pub fn create(packet: &QueryFragment) -> Box<FragmentCoordinator> {
         Box::new(FragmentCoordinator {
             initialized: false,
+            source_only: packet.source_only,
             physical_plan: packet.physical_plan.clone(),
             fragment_id: packet.fragment_id,
             data_exchange: packet.data_exchange.clone(),
@@ -1346,7 +1372,47 @@ impl FragmentCoordinator {
         }
     }
 
+    fn build_source_only_pipeline(
+        &self,
+        ctx: &Arc<QueryContext>,
+        info: &QueryInfo,
+        injector: Arc<dyn ExchangeInjector>,
+    ) -> Result<PipelineBuildResult> {
+        let Some(exchange_params) = self.create_exchange_params(info, injector.clone())? else {
+            return Err(ErrorCode::Internal(
+                "source-only fragment must have data exchange",
+            ));
+        };
+
+        let mut build_res = PipelineBuildResult::create();
+        match &exchange_params {
+            ExchangeParams::BroadcastExchange(params) => {
+                via_broadcast_exchange_source(ctx, params, &mut build_res.main_pipeline)?
+            }
+            ExchangeParams::NodeShuffleExchange(params) => via_node_shuffle_exchange_source(
+                ctx,
+                params,
+                injector,
+                &mut build_res.main_pipeline,
+            )?,
+            ExchangeParams::GlobalShuffleExchange(params) => {
+                via_hash_exchange_source(ctx, params, &mut build_res.main_pipeline)?
+            }
+            ExchangeParams::MergeExchange(_) => {
+                return Err(ErrorCode::Internal(
+                    "merge exchange does not support source-only fragments",
+                ));
+            }
+        }
+
+        Ok(build_res)
+    }
+
     pub fn prepare_pipeline(&mut self, ctx: Arc<QueryContext>) -> Result<()> {
+        if self.source_only {
+            return Ok(());
+        }
+
         if !self.initialized {
             self.initialized = true;