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