datafusion_distributed/protocol/worker_channel.rs
1use crate::{MaybeEncoded, ProducerHead, WorkUnit};
2use async_trait::async_trait;
3use datafusion::arrow::record_batch::RecordBatch;
4use datafusion::common::Result;
5use datafusion::execution::TaskContext;
6use datafusion::physical_plan::ExecutionPlan;
7use datafusion::physical_plan::metrics::{ExecutionPlanMetricsSet, MetricsSet};
8use futures::stream::BoxStream;
9use http::HeaderMap;
10use std::sync::Arc;
11use url::Url;
12use uuid::Uuid;
13
14/// Abstraction over the specific transport protocol implementation.
15///
16/// WARNING: The API in this trait is unstable, and it's subject to change as more things get properly
17/// decoupled from details like protobuf serialization and http headers.
18#[async_trait]
19pub trait WorkerChannel: Send + Sync {
20 /// Establishes a bidirectional message stream between a coordinator and a worker, over which messages
21 /// will be exchanged at any time during a query's lifetime. It's expected to be one coordinator channel
22 /// per task.
23 async fn coordinator_channel(
24 &mut self,
25 headers: HeaderMap,
26 set_plan_request: SetPlanRequest,
27 c2w_stream: BoxStream<'static, CoordinatorToWorkerMsg>,
28 metrics: ExecutionPlanMetricsSet,
29 task_ctx: &Arc<TaskContext>,
30 ) -> Result<BoxStream<'static, Result<WorkerToCoordinatorMsg>>>;
31
32 /// Executes the requested partition range of a subplan previously sent by the coordinator channel.
33 async fn execute_task(
34 &mut self,
35 headers: HeaderMap,
36 request: ExecuteTaskRequest,
37 metrics: ExecutionPlanMetricsSet,
38 task_ctx: &Arc<TaskContext>,
39 ) -> Result<Vec<BoxStream<'static, Result<RecordBatch>>>>;
40
41 /// Returns metadata about a worker. Currently only used for worker versioning.
42 async fn get_worker_info(
43 &mut self,
44 request: GetWorkerInfoRequest,
45 ) -> Result<GetWorkerInfoResponse>;
46}
47
48pub enum CoordinatorToWorkerMsg {
49 /// A batch of messages from a work unit feed belonging to different partitions from one node from the plan set in
50 /// set_plan_request. A work unit feed is a per-partition stream of information that tells the node what should
51 /// be executed within a partition, for example, a stream of file addresses that should be read.
52 WorkUnitBatch(WorkUnitBatch),
53 /// Signals an EOS for WorkUnits. After this message is received, no more WorkUnits will be sent.
54 WorkUnitEos,
55}
56
57#[derive(Clone, Copy, Debug, Hash, PartialEq, Eq)]
58pub struct TaskKey {
59 /// Our query id.
60 pub query_id: Uuid,
61 /// Our stage id.
62 pub stage_id: usize,
63 /// The task number within the stage.
64 pub task_number: usize,
65}
66
67pub struct WorkUnitFeedDeclaration {
68 /// Unique identifier of the node to which work unit feeds are expected to be streamed.
69 pub id: Uuid,
70 /// The amount of partitions expected to be streamed.
71 pub partitions: usize,
72}
73
74pub struct SetPlanRequest {
75 /// The unique identifier of the task to which the subplan belongs to.
76 pub task_key: TaskKey,
77 /// The amount of tasks that share the same subplan. Necessary for building the DistributedTaskContext during execution.
78 pub task_count: usize,
79 /// The subplan the worker is expected to execute.
80 pub plan: MaybeEncoded<Arc<dyn ExecutionPlan>>,
81 /// Information about all the work unit feeds that will be streamed from coordinator to worker.
82 /// This information is needed here because at the moment of setting the plan, all the appropriate
83 /// channels for the incoming work unit feeds need to be constructed.
84 ///
85 /// If no WorkUnitFeedExec nodes are present in the plan, this should be empty.
86 pub work_unit_feed_declarations: Vec<WorkUnitFeedDeclaration>,
87 /// The worker URL to which this message will go. The receiving worker will use this information
88 /// to identify itself, and avoid further calls in case it needs to call itself for executing tasks.
89 pub target_worker_url: Url,
90 /// Unix nanos when the query started as reported by the coordinator. Used for collecting temporal metrics
91 /// relative to when the query was fired in the coordinator.
92 pub query_start_time_ns: usize,
93}
94
95pub struct WorkUnitBatch {
96 /// A batch of WorkUnits.
97 pub batch: Vec<WorkUnitMsg>,
98}
99
100pub struct WorkUnitMsg {
101 /// Identifier of the node to which this work unit feed belongs to.
102 pub id: Uuid,
103 /// The partition index within the node to which the work unit feed belongs to.
104 pub partition: usize,
105 /// Arbitrary user-defined data (e.g., a file address) necessary during execution.
106 pub body: MaybeEncoded<Box<dyn WorkUnit>>,
107 /// Unix timestamp in nanoseconds at which this message was created in the coordinator.
108 pub created_timestamp_unix_nanos: usize,
109 /// Unix timestamp in nanoseconds at which this message was sent by the coordinator.
110 pub sent_timestamp_unix_nanos: usize,
111 /// Unix timestamp in nanoseconds at which this message was received by a worker.
112 pub received_timestamp_unix_nanos: usize,
113 /// Unix timestamp in nanoseconds at which this message started being processed.
114 pub processed_timestamp_unix_nanos: usize,
115}
116
117pub enum WorkerToCoordinatorMsg {
118 /// Sends the metrics collected during task execution back to the coordinator.
119 /// This is sent after all partitions of a task have finished (or been dropped),
120 /// ensuring metrics are never lost due to early stream termination.
121 /// metrics[i] is the set of metrics for plan node i in pre-order traversal order.
122 TaskMetrics(TaskMetrics),
123 /// Load information reported by a task. This information is used for dynamically
124 /// sizing the number of workers involved in a query.
125 LoadInfo(LoadInfo),
126 LoadInfoEos,
127}
128
129#[derive(Clone, Debug)]
130pub struct TaskMetrics {
131 /// Metrics for a single task's plan nodes in pre-order traversal order.
132 /// The TaskKey is implicit — it is determined by the SetPlanRequest that
133 /// opened this coordinator channel connection.
134 pub pre_order_plan_metrics: Vec<MetricsSet>,
135 /// Metrics related to the execution of a task within a stage. This metrics, instead of being
136 /// associated to a specific node, they are global to the task, like the time at which the plan
137 /// was fed by the coordinator to the worker.
138 pub task_metrics: MetricsSet,
139}
140
141#[derive(Default)]
142pub struct LoadInfo {
143 /// The partition index to which this message belongs to.
144 pub partition: usize,
145 /// The amount of rows ready to be returned.
146 pub rows_ready: usize,
147 /// The amount of bytes ready to be returned per column.
148 pub per_column_bytes_ready: Vec<usize>,
149 /// Approximate ratio of NDV for each column.
150 pub per_column_ndv_percentage: Vec<f32>,
151 /// Approximate ratio of null count for each column.
152 pub per_column_null_percentage: Vec<f32>,
153 /// The amount of rows that were pulling from leaf nodes while the partition to which this
154 /// LoadInfo belongs to was sampling data. Used for estimating how much data is left by
155 /// comparing this value to the estimated total rows pulled from leaf nodes.
156 pub rows_pulled_from_leaf: usize,
157 /// Whether the sampled partition stream reached end-of-stream (i.e. the partition finished
158 /// producing all of its output) by the time this LoadInfo was captured. When true, `rows_ready`
159 /// and `per_column_bytes_ready` are final rather than a partial snapshot.
160 pub reached_eos: bool,
161}
162
163pub struct ExecuteTaskRequest {
164 /// The unique identifier of the task that is going to get executed.
165 pub task_key: TaskKey,
166 /// The start of the partition range of the specified task that is going to be executed.
167 pub target_partition_start: usize,
168 /// The end of the partition range of the specified task that is going to be executed.
169 pub target_partition_end: usize,
170 /// The head node the requested task should have. Depending on the network boundary executing
171 /// the task, the head node should be prepared differently, for example:
172 /// - A RepartitionExecHead implies a RepartitionExec at the head of the task.
173 /// - A BroadcastExecHead implies a BroadcastExec at the head of the task.
174 /// - A NoneHead does not need any specific head.
175 pub producer_head: ProducerHead,
176}
177
178pub struct GetWorkerInfoRequest {}
179
180pub struct GetWorkerInfoResponse {
181 pub version: String,
182}