Add affinity key to interface.

Author: JSOD11

src/execution_plans/benchmarks/shuffle_bench.rs

@@ -131,8 +131,11 @@ impl ShuffleBench {
             num: 0,
             plan: None,
             tasks: (0..self.producer_tasks)
-                .map(|i| ExecutionTask {
-                    url: Some(Url::parse(&format!("http://localhost:{i}")).unwrap()),
+                .map(|i| {
+                    ExecutionTask::new(
+                        Some(Url::parse(&format!("http://localhost:{i}")).unwrap()),
+                        None,
+                    )
                 })
                 .collect(),
         };

src/execution_plans/distributed.rs

@@ -130,15 +130,15 @@ impl DistributedExec {
                 .tasks
                 .iter()
                 .enumerate()
-                .map(|(i, _)| {
+                .map(|(i, task)| {
                     let url = task_router.route(RouterInfo {
                         task_number: i,
                         urls: urls.clone(),
                         stage_seed: start_idx,
+                        affinity_key: task.affinity_key.clone(),
                     });
-                    let execution_task = ExecutionTask {
-                        url: Some(url.clone()),
-                    };
+                    let execution_task =
+                        ExecutionTask::new(Some(url.clone()), task.affinity_key.clone());
                     let request = SetPlanRequest {
                         plan_proto: bytes.clone(),
                         task_count: stage.tasks.len() as _,

src/execution_plans/network_coalesce.rs

@@ -120,7 +120,7 @@ impl NetworkCoalesceExec {
                 query_id,
                 num,
                 plan: Some(input),
-                tasks: vec![ExecutionTask { url: None }; input_task_count],
+                tasks: vec![ExecutionTask::default(); input_task_count],
             },
             worker_connections: WorkerConnectionPool::new(input_task_count),
             metrics_collection: Default::default(),

src/execution_plans/network_shuffle.rs

@@ -162,7 +162,7 @@ impl NetworkShuffleExec {
                 query_id,
                 num,
                 plan: Some(transformed.data),
-                tasks: vec![ExecutionTask { url: None }; input_task_count],
+                tasks: vec![ExecutionTask::default(); input_task_count],
             },
             worker_connections: WorkerConnectionPool::new(input_task_count),
             properties: input.properties().clone(),

src/networking/task_router.rs

@@ -5,11 +5,17 @@ use url::Url;
 
 pub struct RouterInfo {
     pub task_number: usize,
-    pub urls: Vec<Url>,
     pub stage_seed: usize,
+    pub urls: Vec<Url>,
+    pub affinity_key: Option<Vec<u8>>,
 }
 
-/// Allows users to route tasks to worker nodes.
+/// Allows users to consistently route tasks to URLs. Defining a custom routing function may be
+/// useful for users that want to ensure that e.g. repeated file reads are directed to the same 
+/// physical machines to allow for proper caching rather than repeatedly reading from object storage.
+///
+/// To implement this, include routing information as bytes in `affinity_key`. These bytes can
+/// then be hashed to consistently point to a specific URL.
 pub trait TaskRouter {
     /// Returns the url of the worker that should execute the task. The default implementation
     /// assigns tasks to workers round-robin starting from stage_seed.

src/protobuf/distributed_codec.rs

@@ -344,6 +344,9 @@ pub struct ExecutionTaskProto {
     /// unassigned.
     #[prost(string, optional, tag = "1")]
     pub url_str: Option<String>,
+    /// An optional affinity key. Used in the task router for consistent routing.
+    #[prost(bytes, optional, tag = "2")]
+    pub affinity_key: Option<Vec<u8>>,
 }
 
 #[derive(Clone, PartialEq, ::prost::Message)]
@@ -499,6 +502,7 @@ fn encode_tasks(tasks: &[ExecutionTask]) -> Vec<ExecutionTaskProto> {
         .iter()
         .map(|task| ExecutionTaskProto {
             url_str: task.url.as_ref().map(|v| v.to_string()),
+            affinity_key: task.affinity_key.clone(),
         })
         .collect()
 }
@@ -514,6 +518,7 @@ fn decode_tasks(tasks: Vec<ExecutionTaskProto>) -> Result<Vec<ExecutionTask>, Da
                         Url::parse(&u).map_err(|_| internal_datafusion_err!("Invalid URL: {u}"))
                     })
                     .transpose()?,
+                affinity_key: task.affinity_key,
             })
         })
         .collect()

src/stage.rs

@@ -81,11 +81,19 @@ pub struct Stage {
     pub tasks: Vec<ExecutionTask>,
 }
 
-#[derive(Debug, Clone)]
+#[derive(Debug, Clone, Default)]
 pub struct ExecutionTask {
     /// The url of the worker that will execute this task.  A None value is interpreted as
     /// unassigned.
     pub(crate) url: Option<Url>,
+    // Affinity key that can be used for consistent routing with the task router.
+    pub(crate) affinity_key: Option<Vec<u8>>,
+}
+
+impl ExecutionTask {
+    pub(crate) fn new(url: Option<Url>, affinity_key: Option<Vec<u8>>) -> Self {
+        Self { url, affinity_key }
+    }
 }
 
 #[derive(Debug, Clone, PartialEq)]
@@ -118,7 +126,7 @@ impl Stage {
             query_id,
             num,
             plan: Some(plan),
-            tasks: vec![ExecutionTask { url: None }; n_tasks],
+            tasks: vec![ExecutionTask::default(); n_tasks],
         }
     }
 }
@@ -439,7 +447,7 @@ pub fn display_plan_graphviz(plan: Arc<dyn ExecutionPlan>) -> Result<String> {
             query_id: Default::default(),
             num: max_num + 1,
             plan: Some(plan.clone()),
-            tasks: vec![ExecutionTask { url: None }],
+            tasks: vec![ExecutionTask::default()],
         };
         all_stages.insert(0, &head_stage);