http_streams_core/
arrow_ipc_codec.rs

1//! Decoding an Arrow IPC stream.
2
3use crate::error::{StreamError, StreamErrorKind};
4use arrow::array::RecordBatch;
5use arrow::ipc::reader::StreamDecoder;
6use bytes::{Buf, BytesMut};
7
8/// A [`Decoder`](tokio_util::codec::Decoder) that yields one [`RecordBatch`] per IPC message.
9///
10/// The schema is read from the stream, so nothing needs to be configured to decode one.
11#[derive(Debug)]
12pub struct ArrowIpcCodec {
13    max_length: usize,
14    decoder: StreamDecoder,
15    current_obj_len: usize,
16}
17
18impl ArrowIpcCodec {
19    /// A codec that rejects any single batch longer than `max_length` bytes.
20    pub fn new_with_max_length(max_length: usize) -> Self {
21        ArrowIpcCodec {
22            max_length,
23            decoder: StreamDecoder::new(),
24            current_obj_len: 0,
25        }
26    }
27}
28
29impl tokio_util::codec::Decoder for ArrowIpcCodec {
30    /// Always `Ok(_)`: the decoder carries dictionary state across the whole stream, so a
31    /// message it could not read leaves it unable to interpret the ones after it.
32    type Item = Result<RecordBatch, StreamError>;
33    type Error = StreamError;
34
35    fn decode(&mut self, buf: &mut BytesMut) -> Result<Option<Self::Item>, StreamError> {
36        let buf_len = buf.len();
37        if buf_len == 0 {
38            return Ok(None);
39        }
40
41        let obj_bytes = buf.as_ref();
42        let obj_bytes_len = obj_bytes.len();
43        let mut buffer = arrow::buffer::Buffer::from(obj_bytes);
44        let maybe_record = self.decoder.decode(&mut buffer).map_err(|e| {
45            StreamError::new(
46                StreamErrorKind::CodecError,
47                Some(Box::new(e)),
48                Some("Decode arrow IPC record error".into()),
49            )
50        })?;
51
52        if maybe_record.is_none() {
53            self.current_obj_len += obj_bytes_len;
54        } else {
55            self.current_obj_len = 0;
56        }
57
58        if self.current_obj_len > self.max_length {
59            return Err(StreamError::new(
60                StreamErrorKind::MaxLenReachedError,
61                None,
62                Some("Object length exceeds the maximum length".into()),
63            ));
64        }
65
66        buf.advance(obj_bytes_len - buffer.len());
67        Ok(maybe_record.map(Ok))
68    }
69
70    fn decode_eof(&mut self, buf: &mut BytesMut) -> Result<Option<Self::Item>, StreamError> {
71        self.decode(buf)
72    }
73}
http_streams_core/arrow_ipc_codec.rs

http_streams_core/
arrow_ipc_codec.rs
