dre_protocol/
frame.rs

1//! Frames: a 4-byte big-endian length, then that many bytes. The first body byte is the frame
2//! type: `J` (0x4A) for a UTF-8 JSON control message, `A` (0x41) for an Arrow IPC stream holding
3//! a schema and zero or more record batches.
4
5use std::io::{self, Read, Write};
6
7use arrow::array::RecordBatch;
8use arrow::datatypes::SchemaRef;
9use arrow::ipc::reader::StreamReader;
10use arrow::ipc::writer::StreamWriter;
11
12pub const JSON_TAG: u8 = b'J';
13pub const ARROW_TAG: u8 = b'A';
14/// Largest frame body accepted: 1 GiB. Larger result sets are split across batches.
15pub const MAX_FRAME: u32 = 1 << 30;
16
17pub enum Frame {
18    Json(serde_json::Value),
19    Arrow(Vec<u8>),
20}
21
22/// Never the raw bytes of an Arrow frame: they end up in error messages.
23impl std::fmt::Debug for Frame {
24    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
25        match self {
26            Frame::Json(v) => write!(f, "Json({v})"),
27            Frame::Arrow(b) => write!(f, "Arrow frame ({} bytes)", b.len()),
28        }
29    }
30}
31
32#[derive(Debug)]
33pub enum FrameError {
34    /// Clean end of stream before a frame started.
35    Eof,
36    Io(io::Error),
37    Malformed(String),
38}
39
40impl std::fmt::Display for FrameError {
41    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
42        match self {
43            FrameError::Eof => f.write_str("end of stream"),
44            FrameError::Io(e) => write!(f, "I/O error: {e}"),
45            FrameError::Malformed(m) => write!(f, "malformed frame: {m}"),
46        }
47    }
48}
49
50impl std::error::Error for FrameError {}
51
52pub fn read_frame(r: &mut impl Read) -> Result<Frame, FrameError> {
53    let mut len = [0u8; 4];
54    match r.read(&mut len[..1]) {
55        Ok(0) => return Err(FrameError::Eof),
56        Ok(_) => {}
57        Err(e) => return Err(FrameError::Io(e)),
58    }
59    r.read_exact(&mut len[1..])
60        .map_err(|_| FrameError::Malformed("stream ended inside a frame header".into()))?;
61    let len = u32::from_be_bytes(len);
62    if len == 0 || len > MAX_FRAME {
63        return Err(FrameError::Malformed(format!(
64            "frame length {len} is out of range"
65        )));
66    }
67    let mut body = vec![0u8; len as usize];
68    r.read_exact(&mut body)
69        .map_err(|_| FrameError::Malformed(format!("stream ended inside a {len}-byte frame")))?;
70    match body[0] {
71        JSON_TAG => serde_json::from_slice(&body[1..])
72            .map(Frame::Json)
73            .map_err(|e| FrameError::Malformed(format!("invalid JSON: {e}"))),
74        ARROW_TAG => {
75            body.remove(0);
76            Ok(Frame::Arrow(body))
77        }
78        t => Err(FrameError::Malformed(format!(
79            "unknown frame type byte 0x{t:02x}"
80        ))),
81    }
82}
83
84pub fn write_json(w: &mut impl Write, v: &impl serde::Serialize) -> io::Result<()> {
85    let mut body = vec![JSON_TAG];
86    serde_json::to_writer(&mut body, v)?;
87    write_body(w, &body)
88}
89
90pub fn write_arrow(w: &mut impl Write, ipc: &[u8]) -> io::Result<()> {
91    let mut body = Vec::with_capacity(ipc.len() + 1);
92    body.push(ARROW_TAG);
93    body.extend_from_slice(ipc);
94    write_body(w, &body)
95}
96
97fn write_body(w: &mut impl Write, body: &[u8]) -> io::Result<()> {
98    let len = u32::try_from(body.len())
99        .ok()
100        .filter(|l| *l <= MAX_FRAME)
101        .ok_or_else(|| io::Error::other("frame too large"))?;
102    w.write_all(&len.to_be_bytes())?;
103    w.write_all(body)?;
104    w.flush()
105}
106
107/// Encode one batch as a self-contained Arrow IPC stream (schema + batch).
108pub fn encode_batch(batch: &RecordBatch) -> Result<Vec<u8>, arrow::error::ArrowError> {
109    let mut buf = Vec::new();
110    {
111        let mut w = StreamWriter::try_new(&mut buf, &batch.schema())?;
112        w.write(batch)?;
113        w.finish()?;
114    }
115    Ok(buf)
116}
117
118/// Decode an Arrow frame into its schema and batches.
119pub fn decode_batches(ipc: &[u8]) -> Result<(SchemaRef, Vec<RecordBatch>), arrow::error::ArrowError> {
120    let reader = StreamReader::try_new(io::Cursor::new(ipc), None)?;
121    let schema = reader.schema();
122    let batches = reader.collect::<Result<Vec<_>, _>>()?;
123    Ok((schema, batches))
124}
125
126#[cfg(test)]
127mod tests {
128    use super::*;
129    use arrow::array::Int64Array;
130    use std::sync::Arc;
131
132    #[test]
133    fn json_and_arrow_frames_round_trip() {
134        let batch = RecordBatch::try_from_iter([("n", Arc::new(Int64Array::from(vec![1, 2])) as _)]).unwrap();
135        let mut buf = Vec::new();
136        write_json(&mut buf, &serde_json::json!({"type": "ok"})).unwrap();
137        write_arrow(&mut buf, &encode_batch(&batch).unwrap()).unwrap();
138        let mut r = io::Cursor::new(buf);
139        assert!(matches!(read_frame(&mut r).unwrap(), Frame::Json(v) if v["type"] == "ok"));
140        let Frame::Arrow(ipc) = read_frame(&mut r).unwrap() else {
141            panic!()
142        };
143        let (_, batches) = decode_batches(&ipc).unwrap();
144        assert_eq!(batches, vec![batch]);
145        assert!(matches!(read_frame(&mut r), Err(FrameError::Eof)));
146    }
147
148    #[test]
149    fn rejects_malformed_frames() {
150        for bad in [
151            vec![0, 0, 0, 0],
152            vec![0, 0, 0, 2, b'X', 1],
153            vec![0, 0, 0, 9, b'J'],
154            vec![0, 0],
155        ] {
156            assert!(matches!(
157                read_frame(&mut io::Cursor::new(bad)),
158                Err(FrameError::Malformed(_))
159            ));
160        }
161    }
162}
dre_protocol/frame.rs

dre_protocol/
frame.rs
