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.split_first() {
71        Some((&JSON_TAG, json)) => serde_json::from_slice(json)
72            .map(Frame::Json)
73            .map_err(|e| FrameError::Malformed(format!("invalid JSON: {e}"))),
74        Some((&ARROW_TAG, _)) => {
75            body.remove(0);
76            Ok(Frame::Arrow(body))
77        }
78        Some((t, _)) => Err(FrameError::Malformed(format!(
79            "unknown frame type byte 0x{t:02x}"
80        ))),
81        None => Err(FrameError::Malformed("empty frame".into())),
82    }
83}
84
85pub fn write_json(w: &mut impl Write, v: &impl serde::Serialize) -> io::Result<()> {
86    let mut body = vec![JSON_TAG];
87    serde_json::to_writer(&mut body, v)?;
88    write_body(w, &body)
89}
90
91pub fn write_arrow(w: &mut impl Write, ipc: &[u8]) -> io::Result<()> {
92    let mut body = Vec::with_capacity(ipc.len() + 1);
93    body.push(ARROW_TAG);
94    body.extend_from_slice(ipc);
95    write_body(w, &body)
96}
97
98fn write_body(w: &mut impl Write, body: &[u8]) -> io::Result<()> {
99    let len = u32::try_from(body.len())
100        .ok()
101        .filter(|l| *l <= MAX_FRAME)
102        .ok_or_else(|| io::Error::other("frame too large"))?;
103    w.write_all(&len.to_be_bytes())?;
104    w.write_all(body)?;
105    w.flush()
106}
107
108/// Encode one batch as a self-contained Arrow IPC stream (schema + batch).
109pub fn encode_batch(batch: &RecordBatch) -> Result<Vec<u8>, arrow::error::ArrowError> {
110    let mut buf = Vec::new();
111    {
112        let mut w = StreamWriter::try_new(&mut buf, &batch.schema())?;
113        w.write(batch)?;
114        w.finish()?;
115    }
116    Ok(buf)
117}
118
119/// Decode an Arrow frame into its schema and batches.
120pub fn decode_batches(ipc: &[u8]) -> Result<(SchemaRef, Vec<RecordBatch>), arrow::error::ArrowError> {
121    let reader = StreamReader::try_new(io::Cursor::new(ipc), None)?;
122    let schema = reader.schema();
123    let batches = reader.collect::<Result<Vec<_>, _>>()?;
124    Ok((schema, batches))
125}
126
127#[cfg(test)]
128mod tests {
129    use super::*;
130    use arrow::array::Int64Array;
131    use std::sync::Arc;
132
133    #[test]
134    fn json_and_arrow_frames_round_trip() {
135        let batch = RecordBatch::try_from_iter([("n", Arc::new(Int64Array::from(vec![1, 2])) as _)]).unwrap();
136        let mut buf = Vec::new();
137        write_json(&mut buf, &serde_json::json!({"type": "ok"})).unwrap();
138        write_arrow(&mut buf, &encode_batch(&batch).unwrap()).unwrap();
139        let mut r = io::Cursor::new(buf);
140        assert!(matches!(read_frame(&mut r).unwrap(), Frame::Json(v) if v["type"] == "ok"));
141        let Frame::Arrow(ipc) = read_frame(&mut r).unwrap() else {
142            panic!()
143        };
144        let (_, batches) = decode_batches(&ipc).unwrap();
145        assert_eq!(batches, vec![batch]);
146        assert!(matches!(read_frame(&mut r), Err(FrameError::Eof)));
147    }
148
149    #[test]
150    fn rejects_malformed_frames() {
151        for bad in [
152            vec![0, 0, 0, 0],
153            vec![0, 0, 0, 2, b'X', 1],
154            vec![0, 0, 0, 9, b'J'],
155            vec![0, 0],
156        ] {
157            assert!(matches!(
158                read_frame(&mut io::Cursor::new(bad)),
159                Err(FrameError::Malformed(_))
160            ));
161        }
162    }
163}
dre_protocol/frame.rs

dre_protocol/
frame.rs
