Cleaning.

Author: JSOD11

src/distributed_ext.rs

@@ -229,14 +229,6 @@ pub trait DistributedExt: Sized {
         resolver: T,
     );
 
-    /// Allows users to select the worker that should execute each task when a distributed query
-    /// is about to run.
-    fn with_distributed_task_router<T: TaskRouter + Send + Sync + 'static>(self, router: T)
-    -> Self;
-
-    /// Same as [DistributedExt::with_distributed_task_router] but with an in-place mutation.
-    fn set_distributed_task_router<T: TaskRouter + Send + Sync + 'static>(&mut self, router: T);
-
     /// This is what tells Distributed DataFusion how to build a Worker gRPC client out of a worker URL.
     ///
     /// There's a default implementation that caches the Worker client instances so that there's
@@ -333,6 +325,14 @@ pub trait DistributedExt: Sized {
         estimator: T,
     );
 
+    /// Allows users to select the worker that should execute each task when a distributed query
+    /// is about to run.
+    fn with_distributed_task_router<T: TaskRouter + Send + Sync + 'static>(self, router: T)
+    -> Self;
+
+    /// Same as [DistributedExt::with_distributed_task_router] but with an in-place mutation.
+    fn set_distributed_task_router<T: TaskRouter + Send + Sync + 'static>(&mut self, router: T);
+
     /// Sets the maximum number of files each task in a stage with a FileScanConfig node will
     /// handle. Reducing this number will increment the amount of tasks. By default, this
     /// is close to the number of cores in the machine.
@@ -572,10 +572,6 @@ impl DistributedExt for SessionConfig {
         set_distributed_worker_resolver(self, resolver);
     }
 
-    fn set_distributed_task_router<T: TaskRouter + Send + Sync + 'static>(&mut self, router: T) {
-        set_distributed_task_router(self, router);
-    }
-
     fn set_distributed_channel_resolver<T: ChannelResolver + Send + Sync + 'static>(
         &mut self,
         resolver: T,
@@ -590,6 +586,10 @@ impl DistributedExt for SessionConfig {
         set_distributed_task_estimator(self, estimator)
     }
 
+    fn set_distributed_task_router<T: TaskRouter + Send + Sync + 'static>(&mut self, router: T) {
+        set_distributed_task_router(self, router);
+    }
+
     fn set_distributed_files_per_task(
         &mut self,
         files_per_task: usize,
@@ -689,10 +689,6 @@ impl DistributedExt for SessionConfig {
             #[expr($;self)]
             fn with_distributed_worker_resolver<T: WorkerResolver + Send + Sync + 'static>(mut self, resolver: T) -> Self;
 
-            #[call(set_distributed_task_router)]
-            #[expr($;self)]
-            fn with_distributed_task_router<T: TaskRouter + Send + Sync + 'static>(mut self, router: T) -> Self;
-
             #[call(set_distributed_channel_resolver)]
             #[expr($;self)]
             fn with_distributed_channel_resolver<T: ChannelResolver + Send + Sync + 'static>(mut self, resolver: T) -> Self;
@@ -701,6 +697,10 @@ impl DistributedExt for SessionConfig {
             #[expr($;self)]
             fn with_distributed_task_estimator<T: TaskEstimator + Send + Sync + 'static>(mut self, estimator: T) -> Self;
 
+            #[call(set_distributed_task_router)]
+            #[expr($;self)]
+            fn with_distributed_task_router<T: TaskRouter + Send + Sync + 'static>(mut self, router: T) -> Self;
+
             #[call(set_distributed_files_per_task)]
             #[expr($?;Ok(self))]
             fn with_distributed_files_per_task(mut self, files_per_task: usize) -> Result<Self, DataFusionError>;
@@ -768,11 +768,6 @@ impl DistributedExt for SessionStateBuilder {
             #[expr($;self)]
             fn with_distributed_worker_resolver<T: WorkerResolver + Send + Sync + 'static>(mut self, resolver: T) -> Self;
 
-            fn set_distributed_task_router<T: TaskRouter + Send + Sync + 'static>(&mut self, router: T);
-            #[call(set_distributed_task_router)]
-            #[expr($;self)]
-            fn with_distributed_task_router<T: TaskRouter + Send + Sync + 'static>(mut self, router: T) -> Self;
-
             fn set_distributed_channel_resolver<T: ChannelResolver + Send + Sync + 'static>(&mut self, resolver: T);
             #[call(set_distributed_channel_resolver)]
             #[expr($;self)]
@@ -783,6 +778,11 @@ impl DistributedExt for SessionStateBuilder {
             #[expr($;self)]
             fn with_distributed_task_estimator<T: TaskEstimator + Send + Sync + 'static>(mut self, estimator: T) -> Self;
 
+            fn set_distributed_task_router<T: TaskRouter + Send + Sync + 'static>(&mut self, router: T);
+            #[call(set_distributed_task_router)]
+            #[expr($;self)]
+            fn with_distributed_task_router<T: TaskRouter + Send + Sync + 'static>(mut self, router: T) -> Self;
+
             fn set_distributed_files_per_task(&mut self, files_per_task: usize) -> Result<(), DataFusionError>;
             #[call(set_distributed_files_per_task)]
             #[expr($?;Ok(self))]
@@ -859,11 +859,6 @@ impl DistributedExt for SessionState {
             #[expr($;self)]
             fn with_distributed_worker_resolver<T: WorkerResolver + Send + Sync + 'static>(mut self, resolver: T) -> Self;
 
-            fn set_distributed_task_router<T: TaskRouter + Send + Sync + 'static>(&mut self, router: T);
-            #[call(set_distributed_task_router)]
-            #[expr($;self)]
-            fn with_distributed_task_router<T: TaskRouter + Send + Sync + 'static>(mut self, router: T) -> Self;
-
             fn set_distributed_channel_resolver<T: ChannelResolver + Send + Sync + 'static>(&mut self, resolver: T);
             #[call(set_distributed_channel_resolver)]
             #[expr($;self)]
@@ -874,6 +869,11 @@ impl DistributedExt for SessionState {
             #[expr($;self)]
             fn with_distributed_task_estimator<T: TaskEstimator + Send + Sync + 'static>(mut self, estimator: T) -> Self;
 
+            fn set_distributed_task_router<T: TaskRouter + Send + Sync + 'static>(&mut self, router: T);
+            #[call(set_distributed_task_router)]
+            #[expr($;self)]
+            fn with_distributed_task_router<T: TaskRouter + Send + Sync + 'static>(mut self, router: T) -> Self;
+
             fn set_distributed_files_per_task(&mut self, files_per_task: usize) -> Result<(), DataFusionError>;
             #[call(set_distributed_files_per_task)]
             #[expr($?;Ok(self))]
@@ -950,11 +950,6 @@ impl DistributedExt for SessionContext {
             #[expr($;self)]
             fn with_distributed_worker_resolver<T: WorkerResolver + Send + Sync + 'static>(self, resolver: T) -> Self;
 
-            fn set_distributed_task_router<T: TaskRouter + Send + Sync + 'static>(&mut self, router: T);
-            #[call(set_distributed_task_router)]
-            #[expr($;self)]
-            fn with_distributed_task_router<T: TaskRouter + Send + Sync + 'static>(self, router: T) -> Self;
-
             fn set_distributed_channel_resolver<T: ChannelResolver + Send + Sync + 'static>(&mut self, resolver: T);
             #[call(set_distributed_channel_resolver)]
             #[expr($;self)]
@@ -965,6 +960,11 @@ impl DistributedExt for SessionContext {
             #[expr($;self)]
             fn with_distributed_task_estimator<T: TaskEstimator + Send + Sync + 'static>(self, estimator: T) -> Self;
 
+            fn set_distributed_task_router<T: TaskRouter + Send + Sync + 'static>(&mut self, router: T);
+            #[call(set_distributed_task_router)]
+            #[expr($;self)]
+            fn with_distributed_task_router<T: TaskRouter + Send + Sync + 'static>(self, router: T) -> Self;
+
             fn set_distributed_files_per_task(&mut self, files_per_task: usize) -> Result<(), DataFusionError>;
             #[call(set_distributed_files_per_task)]
             #[expr($?;Ok(self))]

src/execution_plans/distributed.rs

@@ -76,13 +76,12 @@ 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 by applying these steps to each network boundary:
+    /// 1. Perform worker assignation according to the task router.
+    /// 2. Send the sliced subplans to the assigned URLs. For each URL assigned to a task,
+    ///    perform a network call feeding the subplan to the specified URL.
+    /// 3. Set the node's input plan to `None`. That way, network boundaries appear as leaf
+    ///    nodes, halting further traversal during execution.
     fn prepare_plan(&self, ctx: &Arc<TaskContext>) -> Result<PreparedPlan> {
         let worker_resolver = get_distributed_worker_resolver(ctx.session_config())?;
         let task_router = get_distributed_task_router(ctx.session_config());
@@ -118,7 +117,7 @@ impl DistributedExec {
                 return internal_err!("Plan is not set for stage {}", stage.num);
             };
 
-            // Right now, we assign random workers to tasks. This might change in the future.
+            // The default task router assigns workers to tasks using a random seed.
             let start_idx = rand::rng().random_range(0..urls.len());
 
             // This assumes the plan is the same for all the tasks within a stage. This is fine for
@@ -149,15 +148,15 @@ impl DistributedExec {
                     let plan_send_latency = Arc::clone(&plan_send_latency);
                     let ctx = Arc::clone(ctx);
                     // Spawns the task that feeds this subplan to this worker. There will be as
-                    // many as this spawned tasks as workers.
+                    // many spawned tasks as workers.
                     join_set.spawn(async move {
                         send_plan_task(ctx, url, request).await?;
                         plan_send_latency.record(&start);
                         Ok(())
                     });
-                    Ok(execution_task)
+                    execution_task
                 })
-                .collect::<Result<Vec<_>>>()?;
+                .collect::<Vec<_>>();
 
             Ok(Transformed::yes(plan.with_input_stage(Stage {
                 query_id: stage.query_id,

src/networking/task_router.rs

@@ -5,7 +5,8 @@ use url::Url;
 
 /// Allows users to choose which worker should execute a given distributed task.
 pub trait TaskRouter {
-    /// Returns the url of the worker that should execute the task.
+    /// Returns the index of the worker that should execute the task. The default implementation
+    /// assigns tasks to workers round-robin starting from start_index.
     fn route_fn(&self, start_index: usize, task_number: usize, num_urls: usize) -> usize {
         (start_index + task_number) % num_urls
     }