From 5faed6700699b297f7b052f9b3d0427220f8937e Mon Sep 17 00:00:00 2001 From: Gabriel Date: Fri, 24 Jul 2026 13:22:40 +0200 Subject: [PATCH 1/3] Define routing in terms of events --- src/coordinator/prepare_dynamic_plan.rs | 18 +---- src/coordinator/query_coordinator.rs | 60 +++------------- .../session_state_builder_ext.rs | 22 +++++- src/events/common.rs | 11 +++ src/events/defaults/mod.rs | 4 ++ src/events/defaults/routing.rs | 71 +++++++++++++++++++ src/events/mod.rs | 3 +- 7 files changed, 117 insertions(+), 72 deletions(-) create mode 100644 src/events/defaults/routing.rs diff --git a/src/coordinator/prepare_dynamic_plan.rs b/src/coordinator/prepare_dynamic_plan.rs index b109cc2c..101a6c96 100644 --- a/src/coordinator/prepare_dynamic_plan.rs +++ b/src/coordinator/prepare_dynamic_plan.rs @@ -84,23 +84,7 @@ pub(super) async fn prepare_dynamic_plan( let mut workers = Vec::with_capacity(input_stage.tasks); let mut load_info_rxs = Vec::with_capacity(input_stage.tasks); - let routed_urls = if input_stage.tasks == 1 { - match stage_coordinator - // If the current coordinating context is running within the scope of a local - // worker (same coordinating machine happens to also be a worker), we prefer to - // co-locate single-tasked stages on it. - .find_self_url() - // If there's an input stage with a single worker, and the current stage is also - // going to run in a single worker, we want to co-locate them so that unnecessary - // network transfers are avoided. - .or_else(|| stage_coordinator.find_input_stage_with_single_url()) - { - Some(single_url) => vec![single_url], - None => stage_coordinator.routed_urls()?, - } - } else { - stage_coordinator.routed_urls()? - }; + let routed_urls = stage_coordinator.routed_urls()?; for (i, routed_url) in routed_urls.into_iter().enumerate() { workers.push(routed_url.clone()); diff --git a/src/coordinator/query_coordinator.rs b/src/coordinator/query_coordinator.rs index 8c862862..3acced18 100644 --- a/src/coordinator/query_coordinator.rs +++ b/src/coordinator/query_coordinator.rs @@ -8,18 +8,16 @@ use crate::passthrough_headers::get_passthrough_headers; use crate::stage::LocalStage; use crate::work_unit_feed::WorkUnitFeedRegistry; use crate::work_unit_feed::{build_work_unit_batch_msg, set_work_unit_send_time}; -use crate::worker::LocalWorkerContext; use crate::{ BytesCounterMetric, BytesMetricExt, CoordinatorToWorkerMsg, DISTRIBUTED_DATAFUSION_TASK_ID_LABEL, DistributedCodec, DistributedTaskContext, - DistributedWorkUnitFeedContext, LoadInfo, NetworkBoundaryExt, SetPlanRequest, Stage, TaskKey, - WorkUnitFeedDeclaration, WorkerToCoordinatorMsg, get_distributed_channel_resolver, - get_distributed_worker_resolver, + DistributedWorkUnitFeedContext, LoadInfo, SetPlanRequest, TaskKey, WorkUnitFeedDeclaration, + WorkerToCoordinatorMsg, get_distributed_channel_resolver, }; use datafusion::common::DataFusionError; use datafusion::common::instant::Instant; use datafusion::common::runtime::JoinSet; -use datafusion::common::tree_node::{Transformed, TreeNode, TreeNodeRecursion}; +use datafusion::common::tree_node::{Transformed, TreeNodeRecursion}; use datafusion::common::{Result, exec_err}; use datafusion::execution::TaskContext; use datafusion::physical_expr_common::metrics::{ExecutionPlanMetricsSet, Label, MetricBuilder}; @@ -29,7 +27,6 @@ use datafusion_proto::physical_plan::AsExecutionPlan; use datafusion_proto::protobuf::PhysicalPlanNode; use futures::{Stream, StreamExt, TryStreamExt}; use prost::Message; -use rand::Rng; use std::ops::DerefMut; use std::sync::{Arc, Mutex}; use tokio::sync::Notify; @@ -382,59 +379,18 @@ impl<'a> StageCoordinator<'a> { plan: self.plan, task_count: self.task_count, }; - let routed_urls = match RouteTasksHandlers::handle(ev) { - Some(Ok(response)) => response.urls, - // If the user has not defined custom routing with a `route_tasks` implementation, we - // default to round-robin task assignation from a randomized starting point. - None => { - let session_config = self.task_ctx.session_config(); - let worker_resolver = get_distributed_worker_resolver(session_config)?; - let available_urls = worker_resolver.get_urls()?; - let start_idx = rand::rng().random_range(0..available_urls.len()); - (0..self.task_count) - .map(|i| available_urls[(start_idx + i) % available_urls.len()].clone()) - .collect() - } - Some(Err(e)) => return exec_err!("error routing tasks to workers: {e}"), + let Some(routed) = RouteTasksHandlers::handle(ev).transpose()? else { + return exec_err!("No task routing handler was able to resolve URLs for stage"); }; - if routed_urls.len() != self.task_count { + if routed.urls.len() != self.task_count { return exec_err!( "number of tasks ({}) was not equal to number of urls ({}) at execution time", self.task_count, - routed_urls.len() + routed.urls.len() ); } - Ok(routed_urls) - } - - pub(super) fn find_self_url(&self) -> Option { - self.task_ctx - .session_config() - .get_extension::() - .map(|v| v.self_url.clone()) - } - - pub(super) fn find_input_stage_with_single_url(&self) -> Option { - let mut single_stage_url = None; - self.plan - .apply(|plan| { - let Some(nb) = plan.as_network_boundary() else { - return Ok(TreeNodeRecursion::Continue); - }; - - if let Stage::Remote(remote) = nb.input_stage() - && remote.workers.len() == 1 - { - single_stage_url = Some(remote.workers[0].clone()); - return Ok(TreeNodeRecursion::Stop); - } - - Ok(TreeNodeRecursion::Jump) - }) - .expect("Cannot fail"); - - single_stage_url + Ok(routed.urls) } } diff --git a/src/distributed_planner/session_state_builder_ext.rs b/src/distributed_planner/session_state_builder_ext.rs index 45340582..0d364502 100644 --- a/src/distributed_planner/session_state_builder_ext.rs +++ b/src/distributed_planner/session_state_builder_ext.rs @@ -1,8 +1,9 @@ use crate::distributed_planner::DistributedConfig; use crate::distributed_planner::distributed_query_planner::DistributedQueryPlanner; use crate::events::{ - DesiredTaskCountHandlers, ScaleUpLeafNodeHandlers, file_scan_config_desired_task_count, - file_scan_config_scale_up_leaf_node, + DesiredTaskCountHandlers, RouteTasksHandlers, ScaleUpLeafNodeHandlers, + file_scan_config_desired_task_count, file_scan_config_scale_up_leaf_node, random_routing, + single_task_child_url_routing, single_task_coordinator_routing, }; use datafusion::execution::SessionStateBuilder; use std::sync::Arc; @@ -27,8 +28,25 @@ impl SessionStateBuilderExt for SessionStateBuilder { .optimizer .enable_physical_uncorrelated_scalar_subquery = false; + // Add default event handlers for FileScanConfig nodes. DesiredTaskCountHandlers::push_builtin(cfg, Arc::new(file_scan_config_desired_task_count)); ScaleUpLeafNodeHandlers::push_builtin(cfg, Arc::new(file_scan_config_scale_up_leaf_node)); + + // Add default routing event handlers: + RouteTasksHandlers::extend_builtin( + cfg, + vec![ + // 1. If there's a single task to route, plate it in the coordinator if it can also act as + // a worker. + Arc::new(single_task_coordinator_routing), + // 2. If there's a single task to route, and it cannot be placed in the coordinator, + // co-locate it in one of the workers from the stage below to avoid network transfers. + Arc::new(single_task_child_url_routing), + // 3. If everything above fails, just place randomly. + Arc::new(random_routing), + ], + ); + let prev = std::mem::take(self.query_planner()); self.with_query_planner(Arc::new(DistributedQueryPlanner { prev })) } diff --git a/src/events/common.rs b/src/events/common.rs index c9b884dc..d68b4fd7 100644 --- a/src/events/common.rs +++ b/src/events/common.rs @@ -45,6 +45,17 @@ impl EventHandlerChain { cfg.set_extension(Arc::new(handlers)); } + pub(crate) fn extend_builtin(cfg: &mut SessionConfig, handler_list: Vec>) { + let mut handlers = cfg + .get_extension::() + .map(|v| v.as_ref().clone()) + .unwrap_or_default(); + for handler in handler_list.into_iter().rev() { + handlers.builtin.push(handler); + } + cfg.set_extension(Arc::new(handlers)); + } + pub(crate) fn push_custom(cfg: &mut SessionConfig, handler: Arc) { let mut handlers = cfg .get_extension::() diff --git a/src/events/defaults/mod.rs b/src/events/defaults/mod.rs index c0d6f911..751adcc5 100644 --- a/src/events/defaults/mod.rs +++ b/src/events/defaults/mod.rs @@ -1,5 +1,9 @@ mod file_scan_config; +mod routing; pub(crate) use file_scan_config::{ file_scan_config_desired_task_count, file_scan_config_scale_up_leaf_node, }; +pub(crate) use routing::{ + random_routing, single_task_child_url_routing, single_task_coordinator_routing, +}; diff --git a/src/events/defaults/routing.rs b/src/events/defaults/routing.rs new file mode 100644 index 00000000..e2649c8d --- /dev/null +++ b/src/events/defaults/routing.rs @@ -0,0 +1,71 @@ +use crate::{ + DistributedGetterExt, LocalWorkerContext, NetworkBoundaryExt, RouteTasksEvent, + RouteTasksEventResponse, Stage, WorkerResolver, +}; +use datafusion::common::tree_node::{TreeNode, TreeNodeRecursion}; +use datafusion::common::{Result, exec_err}; +use rand::Rng; + +/// Randomly chooses `ev.task_count` urls from the registered URLs. +pub(crate) fn random_routing(ev: RouteTasksEvent) -> Option> { + let worker_resolver = match ev + .task_ctx + .session_config() + .get_distributed_worker_resolver() + { + Ok(r) => r, + Err(err) => return Some(Err(err)), + }; + + let available_urls = match worker_resolver.get_urls() { + Ok(urls) if !urls.is_empty() => urls, + Ok(_) => return Some(exec_err!("0 URLs available during routing")), + Err(err) => return Some(Err(err)), + }; + + let start_idx = rand::rng().random_range(0..available_urls.len()); + let urls = (0..ev.task_count) + .map(|i| available_urls[(start_idx + i) % available_urls.len()].clone()) + .collect::>(); + + Some(Ok(RouteTasksEventResponse::new(urls))) +} + +/// If there's a single task, it co-locates it in the coordinator if it can also act as a worker. +pub(crate) fn single_task_coordinator_routing( + ev: RouteTasksEvent, +) -> Option> { + if ev.task_count != 1 { + return None; + } + ev.task_ctx + .session_config() + .get_extension::() + .map(|v| Ok(RouteTasksEventResponse::new(vec![v.self_url.clone()]))) +} + +/// If there's a single task, it co-locates it one of the remote workers that is already handling +/// a child task, avoiding network transfers. +pub(crate) fn single_task_child_url_routing( + ev: RouteTasksEvent, +) -> Option> { + let mut single_stage_url = None; + ev.plan + .apply(|plan| { + let Some(nb) = plan.as_network_boundary() else { + return Ok(TreeNodeRecursion::Continue); + }; + + if let Stage::Remote(remote) = nb.input_stage() + && remote.workers.len() == 1 + { + single_stage_url = Some(remote.workers[0].clone()); + return Ok(TreeNodeRecursion::Stop); + } + + Ok(TreeNodeRecursion::Jump) + }) + .expect("Cannot fail"); + + single_stage_url.map(|url| Ok(RouteTasksEventResponse::new(vec![url]))) +} diff --git a/src/events/mod.rs b/src/events/mod.rs index 150c9918..144904e1 100644 --- a/src/events/mod.rs +++ b/src/events/mod.rs @@ -5,7 +5,8 @@ mod route_tasks; mod scale_up_leaf_node; pub(crate) use defaults::{ - file_scan_config_desired_task_count, file_scan_config_scale_up_leaf_node, + file_scan_config_desired_task_count, file_scan_config_scale_up_leaf_node, random_routing, + single_task_child_url_routing, single_task_coordinator_routing, }; pub(crate) use desired_task_count::DesiredTaskCountHandlers; pub use desired_task_count::{ From 814f2861869fb2df25a8e9f823c9450b4311256e Mon Sep 17 00:00:00 2001 From: Gabriel Date: Mon, 27 Jul 2026 10:57:46 +0200 Subject: [PATCH 2/3] Preserve built-in event handler order --- src/events/common.rs | 2 +- src/events/defaults/routing.rs | 3 +++ 2 files changed, 4 insertions(+), 1 deletion(-) diff --git a/src/events/common.rs b/src/events/common.rs index d68b4fd7..7a056b3d 100644 --- a/src/events/common.rs +++ b/src/events/common.rs @@ -50,7 +50,7 @@ impl EventHandlerChain { .get_extension::() .map(|v| v.as_ref().clone()) .unwrap_or_default(); - for handler in handler_list.into_iter().rev() { + for handler in handler_list { handlers.builtin.push(handler); } cfg.set_extension(Arc::new(handlers)); diff --git a/src/events/defaults/routing.rs b/src/events/defaults/routing.rs index e2649c8d..a24d2480 100644 --- a/src/events/defaults/routing.rs +++ b/src/events/defaults/routing.rs @@ -49,6 +49,9 @@ pub(crate) fn single_task_coordinator_routing( pub(crate) fn single_task_child_url_routing( ev: RouteTasksEvent, ) -> Option> { + if ev.task_count != 1 { + return None; + } let mut single_stage_url = None; ev.plan .apply(|plan| { From 2dadcb20d0e0215d46c700598dab25ee96496cb6 Mon Sep 17 00:00:00 2001 From: Gabriel <45515538+gabotechs@users.noreply.github.com> Date: Fri, 31 Jul 2026 15:36:39 +0200 Subject: [PATCH 3/3] Update src/distributed_planner/session_state_builder_ext.rs Co-authored-by: Jayant Shrivastava --- src/distributed_planner/session_state_builder_ext.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/distributed_planner/session_state_builder_ext.rs b/src/distributed_planner/session_state_builder_ext.rs index 0d364502..d17e1c36 100644 --- a/src/distributed_planner/session_state_builder_ext.rs +++ b/src/distributed_planner/session_state_builder_ext.rs @@ -36,7 +36,7 @@ impl SessionStateBuilderExt for SessionStateBuilder { RouteTasksHandlers::extend_builtin( cfg, vec![ - // 1. If there's a single task to route, plate it in the coordinator if it can also act as + // 1. If there's a single task to route, place it in the coordinator if it can also act as // a worker. Arc::new(single_task_coordinator_routing), // 2. If there's a single task to route, and it cannot be placed in the coordinator,