Define routing in terms of events

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());

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<Url> {
-        self.task_ctx
-            .session_config()
-            .get_extension::<LocalWorkerContext>()
-            .map(|v| v.self_url.clone())
-    }
-
-    pub(super) fn find_input_stage_with_single_url(&self) -> Option<Url> {
-        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)
     }
 }
 

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, 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,
+                //    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 }))
     }

src/events/common.rs

@@ -45,6 +45,17 @@ impl<H: ?Sized + Send + Sync + 'static> EventHandlerChain<H> {
         cfg.set_extension(Arc::new(handlers));
     }
 
+    pub(crate) fn extend_builtin(cfg: &mut SessionConfig, handler_list: Vec<Arc<H>>) {
+        let mut handlers = cfg
+            .get_extension::<Self>()
+            .map(|v| v.as_ref().clone())
+            .unwrap_or_default();
+        for handler in handler_list {
+            handlers.builtin.push(handler);
+        }
+        cfg.set_extension(Arc::new(handlers));
+    }
+
     pub(crate) fn push_custom(cfg: &mut SessionConfig, handler: Arc<H>) {
         let mut handlers = cfg
             .get_extension::<Self>()

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,
+};

src/events/defaults/routing.rs

@@ -0,0 +1,74 @@
+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<Result<RouteTasksEventResponse>> {
+    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::<Vec<_>>();
+
+    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<Result<RouteTasksEventResponse>> {
+    if ev.task_count != 1 {
+        return None;
+    }
+    ev.task_ctx
+        .session_config()
+        .get_extension::<LocalWorkerContext>()
+        .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<Result<RouteTasksEventResponse>> {
+    if ev.task_count != 1 {
+        return None;
+    }
+    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])))
+}

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::{