oxigeo_distributed/flight/
wire.rs

1//! Wire protocol for the `execute_task` Flight action.
2//!
3//! This module is the glue that connects the three previously-disconnected
4//! distributed components — the [`Coordinator`](crate::coordinator::Coordinator)
5//! (task bookkeeping), the [`Worker`](crate::worker::Worker) (task execution),
6//! and Arrow Flight (data transport) — into an actual multi-process pipeline.
7//!
8//! A task-execution request/response is a single self-describing frame:
9//!
10//! ```text
11//! ┌──────────────┬───────────────────────┬──────────────────────────────┐
12//! │ u32 LE       │ JSON header           │ Arrow IPC stream (optional)   │
13//! │ header_len   │ (Task / result meta)  │ input or output RecordBatch   │
14//! └──────────────┴───────────────────────┴──────────────────────────────┘
15//! ```
16//!
17//! Carrying the input batch inline keeps the RPC hermetic: the coordinator does
18//! not have to pre-stage data under a separately-negotiated ticket, and the
19//! worker returns its result batch in the same envelope.
20
21use crate::error::{DistributedError, Result};
22use crate::task::Task;
23use arrow::record_batch::RecordBatch;
24use arrow_ipc::reader::StreamReader;
25use arrow_ipc::writer::StreamWriter;
26use bytes::Bytes;
27use serde::{Deserialize, Serialize};
28
29/// Flight action type used to submit a task to a worker's Flight server.
30pub const EXECUTE_TASK_ACTION: &str = "execute_task";
31
32/// JSON header of an [`EXECUTE_TASK_ACTION`] request.
33#[derive(Debug, Clone, Serialize, Deserialize)]
34pub struct ExecuteTaskRequest {
35    /// The task to execute.
36    pub task: Task,
37    /// Whether an input `RecordBatch` (Arrow IPC) follows the header.
38    pub has_input: bool,
39}
40
41/// JSON header of an [`EXECUTE_TASK_ACTION`] response.
42#[derive(Debug, Clone, Serialize, Deserialize)]
43pub struct ExecuteTaskResponse {
44    /// Whether the task completed successfully.
45    pub success: bool,
46    /// Error message when the task failed.
47    pub error: Option<String>,
48    /// Wall-clock execution time reported by the worker, in milliseconds.
49    pub execution_time_ms: u64,
50    /// Number of rows in the output batch (0 when there is none).
51    pub num_rows: usize,
52    /// Whether an output `RecordBatch` (Arrow IPC) follows the header.
53    pub has_output: bool,
54}
55
56/// Serialize a single [`RecordBatch`] to an Arrow IPC stream byte buffer.
57pub fn encode_batch_ipc(batch: &RecordBatch) -> Result<Vec<u8>> {
58    let mut buffer: Vec<u8> = Vec::new();
59    {
60        let mut writer = StreamWriter::try_new(&mut buffer, batch.schema().as_ref())
61            .map_err(|e| DistributedError::arrow(format!("IPC writer init failed: {e}")))?;
62        writer
63            .write(batch)
64            .map_err(|e| DistributedError::arrow(format!("IPC write failed: {e}")))?;
65        writer
66            .finish()
67            .map_err(|e| DistributedError::arrow(format!("IPC finish failed: {e}")))?;
68    }
69    Ok(buffer)
70}
71
72/// Decode a single [`RecordBatch`] from an Arrow IPC stream byte buffer.
73pub fn decode_batch_ipc(bytes: &[u8]) -> Result<RecordBatch> {
74    let reader = StreamReader::try_new(std::io::Cursor::new(bytes), None)
75        .map_err(|e| DistributedError::arrow(format!("IPC reader init failed: {e}")))?;
76    let mut batch = None;
77    for item in reader {
78        let b = item.map_err(|e| DistributedError::arrow(format!("IPC read failed: {e}")))?;
79        // The envelope always carries exactly one batch; keep the first.
80        if batch.is_none() {
81            batch = Some(b);
82        }
83    }
84    batch.ok_or_else(|| DistributedError::arrow("IPC stream carried no record batch"))
85}
86
87/// Frame a JSON header and an optional batch into a single wire buffer.
88fn frame(header_json: Vec<u8>, batch: Option<&RecordBatch>) -> Result<Bytes> {
89    let header_len = u32::try_from(header_json.len())
90        .map_err(|_| DistributedError::flight_rpc("execute_task header too large"))?;
91
92    let mut out = Vec::with_capacity(4 + header_json.len());
93    out.extend_from_slice(&header_len.to_le_bytes());
94    out.extend_from_slice(&header_json);
95    if let Some(batch) = batch {
96        out.extend_from_slice(&encode_batch_ipc(batch)?);
97    }
98    Ok(Bytes::from(out))
99}
100
101/// Split a wire buffer back into its JSON header bytes and trailing IPC bytes.
102fn unframe(bytes: &[u8]) -> Result<(&[u8], &[u8])> {
103    if bytes.len() < 4 {
104        return Err(DistributedError::flight_rpc(
105            "execute_task frame shorter than its length prefix",
106        ));
107    }
108    let header_len = u32::from_le_bytes([bytes[0], bytes[1], bytes[2], bytes[3]]) as usize;
109    let header_end = 4usize
110        .checked_add(header_len)
111        .filter(|&end| end <= bytes.len())
112        .ok_or_else(|| DistributedError::flight_rpc("execute_task header length out of range"))?;
113    Ok((&bytes[4..header_end], &bytes[header_end..]))
114}
115
116/// Encode an execute-task request (header + optional input batch) into a Flight
117/// action body.
118pub fn encode_execute_request(task: &Task, input: Option<&RecordBatch>) -> Result<Bytes> {
119    let header = ExecuteTaskRequest {
120        task: task.clone(),
121        has_input: input.is_some(),
122    };
123    let header_json = serde_json::to_vec(&header)
124        .map_err(|e| DistributedError::task_serialization(format!("encode request: {e}")))?;
125    frame(header_json, input)
126}
127
128/// Decode an execute-task request body into its task and optional input batch.
129pub fn decode_execute_request(bytes: &[u8]) -> Result<(Task, Option<RecordBatch>)> {
130    let (header_bytes, tail) = unframe(bytes)?;
131    let header: ExecuteTaskRequest = serde_json::from_slice(header_bytes)
132        .map_err(|e| DistributedError::task_serialization(format!("decode request: {e}")))?;
133    let input = if header.has_input {
134        Some(decode_batch_ipc(tail)?)
135    } else {
136        None
137    };
138    Ok((header.task, input))
139}
140
141/// Encode an execute-task response (header + optional output batch) into a
142/// Flight action result body.
143pub fn encode_execute_response(
144    response: &ExecuteTaskResponse,
145    output: Option<&RecordBatch>,
146) -> Result<Bytes> {
147    let header_json = serde_json::to_vec(response)
148        .map_err(|e| DistributedError::task_serialization(format!("encode response: {e}")))?;
149    frame(header_json, output)
150}
151
152/// Decode an execute-task response body into its metadata and optional output
153/// batch.
154pub fn decode_execute_response(bytes: &[u8]) -> Result<(ExecuteTaskResponse, Option<RecordBatch>)> {
155    let (header_bytes, tail) = unframe(bytes)?;
156    let header: ExecuteTaskResponse = serde_json::from_slice(header_bytes)
157        .map_err(|e| DistributedError::task_serialization(format!("decode response: {e}")))?;
158    let output = if header.has_output {
159        Some(decode_batch_ipc(tail)?)
160    } else {
161        None
162    };
163    Ok((header, output))
164}
165
166#[cfg(test)]
167mod tests {
168    use super::*;
169    use crate::task::{PartitionId, TaskId, TaskOperation};
170    use arrow::array::Int32Array;
171    use arrow::datatypes::{DataType, Field, Schema};
172    use std::sync::Arc;
173
174    fn sample_batch() -> RecordBatch {
175        let schema = Arc::new(Schema::new(vec![Field::new(
176            "value",
177            DataType::Int32,
178            false,
179        )]));
180        RecordBatch::try_new(
181            schema,
182            vec![Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5]))],
183        )
184        .expect("batch")
185    }
186
187    fn sample_task() -> Task {
188        Task::new(
189            TaskId(7),
190            PartitionId(2),
191            TaskOperation::Filter {
192                expression: "value > 2".to_string(),
193            },
194        )
195    }
196
197    #[test]
198    fn batch_ipc_round_trips() {
199        let batch = sample_batch();
200        let bytes = encode_batch_ipc(&batch).expect("encode");
201        let restored = decode_batch_ipc(&bytes).expect("decode");
202        assert_eq!(restored.num_rows(), 5);
203        assert_eq!(restored.num_columns(), 1);
204    }
205
206    #[test]
207    fn request_round_trips_with_input() {
208        let task = sample_task();
209        let batch = sample_batch();
210        let body = encode_execute_request(&task, Some(&batch)).expect("encode");
211        let (decoded_task, input) = decode_execute_request(&body).expect("decode");
212        assert_eq!(decoded_task.id, TaskId(7));
213        let input = input.expect("input present");
214        assert_eq!(input.num_rows(), 5);
215    }
216
217    #[test]
218    fn request_round_trips_without_input() {
219        let task = sample_task();
220        let body = encode_execute_request(&task, None).expect("encode");
221        let (decoded_task, input) = decode_execute_request(&body).expect("decode");
222        assert_eq!(decoded_task.partition_id, PartitionId(2));
223        assert!(input.is_none());
224    }
225
226    #[test]
227    fn response_round_trips_with_output() {
228        let batch = sample_batch();
229        let resp = ExecuteTaskResponse {
230            success: true,
231            error: None,
232            execution_time_ms: 42,
233            num_rows: 5,
234            has_output: true,
235        };
236        let body = encode_execute_response(&resp, Some(&batch)).expect("encode");
237        let (decoded, output) = decode_execute_response(&body).expect("decode");
238        assert!(decoded.success);
239        assert_eq!(decoded.execution_time_ms, 42);
240        assert_eq!(output.expect("output").num_rows(), 5);
241    }
242
243    #[test]
244    fn unframe_rejects_truncated_frame() {
245        assert!(unframe(&[0, 1]).is_err());
246        // header_len claims more bytes than present.
247        assert!(unframe(&[255, 255, 255, 255, 0]).is_err());
248    }
249}
oxigeo_distributed/flight/wire.rs

oxigeo_distributed/flight/
wire.rs
