Task Routing with the Task Estimator: Allow users to map tasks to URLs

examples/custom_execution_plan.rs

@@ -341,6 +341,13 @@ impl TaskEstimator for NumbersTaskEstimator {
 
         Some(Arc::new(NumbersExec::new(ranges_per_task, plan.schema())))
     }
+
+    fn route_tasks(
+        &self,
+        _routing_ctx: &datafusion_distributed::TaskRoutingContext<'_>,
+    ) -> Result<Option<Vec<url::Url>>> {
+        Ok(None)
+    }
 }
 
 #[derive(StructOpt)]

src/distributed_planner/mod.rs

@@ -10,5 +10,5 @@ mod task_estimator;
 pub use distributed_config::DistributedConfig;
 pub use network_boundary::{NetworkBoundary, NetworkBoundaryExt};
 pub use session_state_builder_ext::SessionStateBuilderExt;
-pub(crate) use task_estimator::set_distributed_task_estimator;
-pub use task_estimator::{TaskCountAnnotation, TaskEstimation, TaskEstimator};
+pub use task_estimator::{TaskCountAnnotation, TaskEstimation, TaskEstimator, TaskRoutingContext};
+pub(crate) use task_estimator::{get_distributed_task_estimator, set_distributed_task_estimator};

src/distributed_planner/plan_annotator.rs

@@ -415,6 +415,7 @@ mod tests {
     use datafusion::execution::SessionStateBuilder;
     use datafusion::physical_plan::coalesce_partitions::CoalescePartitionsExec;
     use datafusion::physical_plan::filter::FilterExec;
+    use url::Url;
     /* schema for the "weather" table
 
      MinTemp [type=DOUBLE] [repetitiontype=OPTIONAL]
@@ -893,6 +894,10 @@ mod tests {
         ) -> Option<Arc<dyn ExecutionPlan>> {
             None
         }
+
+        fn route_tasks(&self, _: &crate::TaskRoutingContext<'_>) -> Result<Option<Vec<Url>>> {
+            Ok(None)
+        }
     }
 
     #[derive(Debug)]
@@ -920,6 +925,10 @@ mod tests {
         ) -> Option<Arc<dyn ExecutionPlan>> {
             None
         }
+
+        fn route_tasks(&self, _: &crate::TaskRoutingContext<'_>) -> Result<Option<Vec<url::Url>>> {
+            Ok(None)
+        }
     }
 
     async fn sql_to_annotated(query: &str) -> String {

src/distributed_planner/task_estimator.rs

@@ -1,13 +1,18 @@
 use crate::config_extension_ext::set_distributed_option_extension;
+use crate::stage::ExecutionTask;
 use crate::{DistributedConfig, PartitionIsolatorExec};
 use datafusion::catalog::memory::DataSourceExec;
+use datafusion::common::exec_err;
 use datafusion::config::ConfigOptions;
 use datafusion::datasource::physical_plan::FileScanConfig;
+use datafusion::error::{DataFusionError, Result};
+use datafusion::execution::TaskContext;
 use datafusion::physical_plan::ExecutionPlan;
 use datafusion::prelude::SessionConfig;
 use delegate::delegate;
 use std::fmt::Debug;
 use std::sync::Arc;
+use url::Url;
 
 /// Annotation attached to a single [ExecutionPlan] that determines how many distributed tasks
 /// it should run on.
@@ -121,6 +126,17 @@ pub trait TaskEstimator {
         task_count: usize,
         cfg: &ConfigOptions,
     ) -> Option<Arc<dyn ExecutionPlan>>;
+
+    /// Defines a custom protocol for routing tasks to specific URLs / physical machines.
+    fn route_tasks(&self, routing_ctx: &TaskRoutingContext<'_>) -> Result<Option<Vec<Url>>>;
+}
+
+/// Context usable for routing tasks to worker URLs.
+pub struct TaskRoutingContext<'a> {
+    pub task_ctx: Arc<TaskContext>,
+    pub plan: &'a Arc<dyn ExecutionPlan>,
+    pub tasks: &'a [ExecutionTask],
+    pub urls: &'a [Url],
 }
 
 impl TaskEstimator for usize {
@@ -146,13 +162,18 @@ impl TaskEstimator for usize {
     ) -> Option<Arc<dyn ExecutionPlan>> {
         None
     }
+
+    fn route_tasks(&self, _: &TaskRoutingContext<'_>) -> Result<Option<Vec<Url>>> {
+        Ok(None)
+    }
 }
 
 impl TaskEstimator for Arc<dyn TaskEstimator> {
     delegate! {
         to self.as_ref() {
             fn task_estimation(&self, plan: &Arc<dyn ExecutionPlan>, cfg: &ConfigOptions) -> Option<TaskEstimation>;
             fn scale_up_leaf_node(&self, plan: &Arc<dyn ExecutionPlan>, task_count: usize, cfg: &ConfigOptions) -> Option<Arc<dyn ExecutionPlan>>;
+            fn route_tasks(&self, routing_ctx: &TaskRoutingContext<'_>) -> Result<Option<Vec<Url>>>;
         }
     }
 }
@@ -162,6 +183,7 @@ impl TaskEstimator for Arc<dyn TaskEstimator + Send + Sync> {
         to self.as_ref() {
             fn task_estimation(&self, plan: &Arc<dyn ExecutionPlan>, cfg: &ConfigOptions) -> Option<TaskEstimation>;
             fn scale_up_leaf_node(&self, plan: &Arc<dyn ExecutionPlan>, task_count: usize, cfg: &ConfigOptions) -> Option<Arc<dyn ExecutionPlan>>;
+            fn route_tasks(&self, routing_ctx: &TaskRoutingContext<'_>) -> Result<Option<Vec<Url>>>;
         }
     }
 }
@@ -189,6 +211,18 @@ pub(crate) fn set_distributed_task_estimator(
     }
 }
 
+pub fn get_distributed_task_estimator(
+    cfg: &SessionConfig,
+) -> Result<Arc<dyn TaskEstimator>, DataFusionError> {
+    let opts = cfg.options();
+    let Some(distributed_cfg) = opts.extensions.get::<DistributedConfig>() else {
+        return exec_err!("WorkerResolver not present in the session config");
+    };
+    let task_estimator: Arc<dyn TaskEstimator> =
+        Arc::new(distributed_cfg.__private_task_estimator.clone());
+    Ok(Arc::clone(&task_estimator))
+}
+
 /// [TaskEstimator] implementation that acts on [DataSourceExec] nodes that contain
 /// [FileScanConfig]s data sources (e.g., Parquet or CSV files). it will read the
 /// [DistributedConfig].`files_per_task` field and assigns as many tasks as needed so that
@@ -247,6 +281,10 @@ impl TaskEstimator for FileScanConfigTaskEstimator {
         let plan = DataSourceExec::from_data_source(new_file_scan);
         Some(Arc::new(PartitionIsolatorExec::new(plan, task_count)))
     }
+
+    fn route_tasks(&self, _routing_ctx: &TaskRoutingContext<'_>) -> Result<Option<Vec<Url>>> {
+        Ok(None)
+    }
 }
 
 /// Tries multiple user-provided [TaskEstimator]s until one returns an estimation. If none
@@ -300,6 +338,15 @@ impl TaskEstimator for CombinedTaskEstimator {
         }
         None
     }
+
+    fn route_tasks(&self, routing_ctx: &TaskRoutingContext<'_>) -> Result<Option<Vec<Url>>> {
+        for estimator in &self.user_provided {
+            if let Some(result) = estimator.route_tasks(routing_ctx)? {
+                return Ok(Some(result));
+            }
+        }
+        Ok(None)
+    }
 }
 
 #[cfg(test)]
@@ -400,5 +447,9 @@ mod tests {
         ) -> Option<Arc<dyn ExecutionPlan>> {
             None
         }
+
+        fn route_tasks(&self, _routing_ctx: &TaskRoutingContext<'_>) -> Result<Option<Vec<Url>>> {
+            Ok(None)
+        }
     }
 }

src/execution_plans/distributed.rs

@@ -1,6 +1,6 @@
 use crate::common::{require_one_child, serialize_uuid, task_ctx_with_extension};
 use crate::config_extension_ext::get_config_extension_propagation_headers;
-use crate::distributed_planner::NetworkBoundaryExt;
+use crate::distributed_planner::{NetworkBoundaryExt, get_distributed_task_estimator};
 use crate::execution_plans::ChildrenIsolatorUnionExec;
 use crate::networking::get_distributed_worker_resolver;
 use crate::passthrough_headers::get_passthrough_headers;
@@ -14,7 +14,8 @@ use crate::worker::generated::worker::{
 };
 use crate::{
     DISTRIBUTED_DATAFUSION_TASK_ID_LABEL, DistributedConfig, DistributedTaskContext,
-    DistributedWorkUnitFeedContext, WorkerResolver, get_distributed_channel_resolver,
+    DistributedWorkUnitFeedContext, TaskRoutingContext, WorkerResolver,
+    get_distributed_channel_resolver,
 };
 use datafusion::common::HashMap;
 use datafusion::common::instant::Instant;
@@ -156,13 +157,14 @@ impl DistributedExec {
             })
     }
 
-    /// Prepares the distributed plan for execution, which implies:
-    /// 1. Perform some worker assignation, choosing randomly from the given URLs and assigning one
-    ///    URL per task.
-    /// 2. Sending the sliced subplans to the assigned URLs. For each URL assigned to a task, a
-    ///    network call feeding the subplan is necessary.
-    /// 3. In each network boundary, set the input plan to `None`. That way, network boundaries
-    ///    become nodes without children and traversing them will not go further down in.
+    /// Prepares the distributed plan for execution.
+    /// In particular, this means we must take the following steps at each network boundary node:
+    /// 1. Assign tasks to URLs. Follow the user-specified routing defined in the TaskEstimator
+    ///    implementation, or default to random round-robin assignment.
+    /// 2. Send the sliced subplans to the assigned URLs. For each task assigned to a URL, it is here
+    ///    that we now must actually send that subplan to the URL over the wire.
+    /// 3. Set network boundary input plans to `None`. This way, network boundaries become nodes
+    ///    without children, so we stop further traversal from happening in the future.
     /// 4. Spawn a background task per worker that waits for the worker to finish and collects
     ///    its metrics into [DistributedExec::task_metrics] via the coordinator channel.
     fn prepare_plan(&self, ctx: &Arc<TaskContext>) -> Result<PreparedPlan> {
@@ -186,28 +188,59 @@ impl DistributedExec {
 
         let mut join_set = JoinSet::new();
         let prepared = Arc::clone(&self.plan).transform_up(|plan| {
-            // The following logic is just applied on network boundaries.
+            // The following logic is only relevant to network boundaries.
             let Some(plan) = plan.as_network_boundary() else {
                 return Ok(Transformed::no(plan));
             };
 
+            let task_estimator = get_distributed_task_estimator(ctx.session_config())?;
+
             let stage = plan.input_stage();
 
             let mut spawner =
                 CoordinatorToWorkerTaskSpawner::new(stage, &metrics, &codec, &mut join_set)?;
 
-            // Right now, we assign random workers to tasks. This might change in the future.
+            // Default routing assigns tasks to URLs round-robin from a random starting point.
             let start_idx = rand::rng().random_range(0..urls.len());
 
+            // If the user has not defined custom routing through the route_task API, we
+            // default to round-robin task assignation with a randomized starting point.
+            //
+            // If the user has defined custom routing through the route_task interface,
+            // we attempt to route the task to the specified URL.
+            let routed_urls = match task_estimator.route_tasks(&TaskRoutingContext {
+                task_ctx: Arc::clone(ctx),
+                plan: stage.plan.as_ref().ok_or_else(|| {
+                    internal_datafusion_err!(
+                        "network boundary stage missing input plan during routing"
+                    )
+                })?,
+                tasks: &stage.tasks,
+                urls: &urls,
+            }) {
+                Ok(Some(routed_urls)) => routed_urls,
+                Ok(None) => (0..stage.tasks.len())
+                    .map(|i| urls[(start_idx + i) % urls.len()].clone())
+                    .collect(),
+                Err(e) => return Err(exec_datafusion_err!("error routing tasks to workers: {e}")),
+            };
+
+            if routed_urls.len() != stage.tasks.len() {
+                return Err(exec_datafusion_err!(
+                    "number of tasks ({}) was not equal to number of urls ({}) at execution time",
+                    routed_urls.len(),
+                    stage.tasks.len()
+                ));
+            }
+
             let mut tasks = Vec::with_capacity(stage.tasks.len());
-            for i in 0..stage.tasks.len() {
-                let url = urls[(start_idx + i) % urls.len()].clone();
+            for (i, routed_url) in routed_urls.into_iter().enumerate() {
                 tasks.push(ExecutionTask {
-                    url: Some(url.clone()),
+                    url: Some(routed_url.clone()),
                 });
-                // Spawns the task that feeds this subplan to this worker. There will be as
-                // many as this spawned tasks as workers.
-                let (tx, worker_rx) = spawner.send_plan_task(Arc::clone(ctx), i, url)?;
+                // Spawn a task that sends the subplan to the chosen URL.
+                // There will be as many spawned tasks as workers.
+                let (tx, worker_rx) = spawner.send_plan_task(Arc::clone(ctx), i, routed_url)?;
                 spawner.metrics_collection_task(i, worker_rx, Arc::clone(&self.task_metrics));
                 spawner.work_unit_feed_task(Arc::clone(ctx), i, tx)?;
             }

src/lib.rs

@@ -22,7 +22,7 @@ pub use arrow_ipc::CompressionType;
 pub use distributed_ext::DistributedExt;
 pub use distributed_planner::{
     DistributedConfig, NetworkBoundary, NetworkBoundaryExt, SessionStateBuilderExt,
-    TaskCountAnnotation, TaskEstimation, TaskEstimator,
+    TaskCountAnnotation, TaskEstimation, TaskEstimator, TaskRoutingContext,
 };
 pub use execution_plans::{
     BroadcastExec, DistributedExec, NetworkBroadcastExec, NetworkCoalesceExec, NetworkShuffleExec,

src/test_utils/in_memory_channel_resolver.rs

@@ -1,15 +1,19 @@
 use crate::worker::generated::worker::worker_service_client::WorkerServiceClient;
 use crate::{
     BoxCloneSyncChannel, ChannelResolver, DefaultSessionBuilder, DistributedExt,
-    MappedWorkerSessionBuilderExt, Worker, WorkerResolver, WorkerSessionBuilder,
-    create_worker_client,
+    MappedWorkerSessionBuilderExt, SessionStateBuilderExt, Worker, WorkerResolver,
+    WorkerSessionBuilder, create_worker_client,
 };
 use async_trait::async_trait;
 use datafusion::common::DataFusionError;
+use datafusion::execution::SessionStateBuilder;
+use datafusion::prelude::SessionContext;
 use hyper_util::rt::TokioIo;
 use tonic::transport::{Endpoint, Server};
+use url::Url;
 
 const DUMMY_URL: &str = "http://localhost:50051";
+const DUMMY_URL_PREFIX: &str = "http://url-";
 
 /// [ChannelResolver] implementation that returns gRPC clients backed by an in-memory
 /// tokio duplex rather than a TCP connection.
@@ -85,9 +89,27 @@ impl InMemoryWorkerResolver {
 }
 
 impl WorkerResolver for InMemoryWorkerResolver {
-    fn get_urls(&self) -> Result<Vec<url::Url>, DataFusionError> {
-        // Set to a high number so that the distributed planner does not limit the maximum
-        // spawned tasks to just 1.
-        Ok(vec![url::Url::parse(DUMMY_URL).unwrap(); self.n_workers])
+    fn get_urls(&self) -> Result<Vec<Url>, DataFusionError> {
+        (0..self.n_workers)
+            .map(|i| Url::parse(&format!("{}{}", DUMMY_URL_PREFIX, i)))
+            .collect::<Result<Vec<_>, _>>()
+            .map_err(|err| DataFusionError::External(Box::new(err)))
     }
 }
+
+/// Creates a distributed session context backed by a single in-memory worker service.
+/// The set of produced worker URLs is deterministic, taking the form http://worker-<i>.
+pub async fn start_in_memory_context(
+    num_workers: usize,
+    session_builder: impl WorkerSessionBuilder + Send + Sync + 'static,
+) -> SessionContext {
+    let channel_resolver = InMemoryChannelResolver::from_session_builder(session_builder);
+    let mut state = SessionStateBuilder::new()
+        .with_default_features()
+        .with_distributed_planner()
+        .with_distributed_worker_resolver(InMemoryWorkerResolver::new(num_workers))
+        .with_distributed_channel_resolver(channel_resolver)
+        .build();
+    state.config_mut().options_mut().execution.target_partitions = 3;
+    SessionContext::from(state)
+}

src/test_utils/mod.rs

@@ -6,5 +6,6 @@ pub mod mock_exec;
 pub mod parquet;
 pub mod plans;
 pub mod property_based;
+pub mod routing;
 pub mod session_context;
 pub mod test_work_unit_feed;

src/test_utils/plans.rs

@@ -204,4 +204,11 @@ impl TaskEstimator for BuildSideOneTaskEstimator {
     ) -> Option<Arc<dyn ExecutionPlan>> {
         None
     }
+
+    fn route_tasks(
+        &self,
+        _: &crate::TaskRoutingContext<'_>,
+    ) -> datafusion::error::Result<Option<Vec<url::Url>>> {
+        Ok(None)
+    }
 }