http_streams_core/
arrow_format.rs

1//! Apache Arrow IPC, both directions.
2
3use crate::content_type::ContentType;
4use crate::error::{StreamError, StreamErrorKind};
5use crate::format::{
6    DecodeOptions, IdentityParser, ItemEncoder, StreamFormat, StreamFormatDecode,
7    StreamFormatEncode,
8};
9use crate::arrow_ipc_codec::ArrowIpcCodec;
10use arrow::array::RecordBatch;
11use arrow::datatypes::{Schema, SchemaRef};
12use arrow::error::ArrowError;
13use arrow::ipc::writer::{
14    write_message, DictionaryTracker, IpcDataGenerator, IpcWriteContext, IpcWriteOptions,
15};
16use bytes::{BufMut, BytesMut};
17use std::io::Write;
18use std::sync::Arc;
19
20const ARROW_CONTENT_TYPE: &str = "application/vnd.apache.arrow.stream";
21
22/// The four-byte continuation marker plus a zero length, which ends an Arrow IPC stream.
23const CONTINUATION_MARKER: [u8; 4] = [0xff; 4];
24const TOTAL_LEN: [u8; 4] = [0; 4];
25
26/// Arrow record batches in IPC stream framing.
27///
28///
29/// Encoding needs the schema up front, because it is written once ahead of the first batch.
30/// Decoding does not: the schema arrives in the stream.
31#[derive(Debug, Clone)]
32pub struct ArrowRecordBatchIpcStreamFormat {
33    schema: SchemaRef,
34    options: IpcWriteOptions,
35}
36
37impl ArrowRecordBatchIpcStreamFormat {
38    /// A format writing batches of `schema` with default IPC options.
39    pub fn new(schema: Arc<Schema>) -> Self {
40        Self::with_options(schema, IpcWriteOptions::default())
41    }
42
43    /// A format writing batches of `schema` with the given IPC options.
44    pub fn with_options(schema: Arc<Schema>, options: IpcWriteOptions) -> Self {
45        Self { schema, options }
46    }
47
48    /// A format for **decoding** only.
49    ///
50    /// An Arrow IPC stream carries its own schema, so a decoder needs none. Named for what it
51    /// is rather than offered as a `Default`, because encoding with it would write an empty
52    /// schema and produce a useless stream.
53    pub fn for_decoding() -> Self {
54        Self::new(Arc::new(Schema::empty()))
55    }
56}
57
58impl crate::format::DefaultFormat for ArrowRecordBatchIpcStreamFormat {
59    /// Decoding needs no configuration; see [`for_decoding`](Self::for_decoding).
60    fn default_format() -> Self {
61        Self::for_decoding()
62    }
63}
64
65impl StreamFormat for ArrowRecordBatchIpcStreamFormat {
66    fn format_name(&self) -> &'static str {
67        "arrow"
68    }
69
70    fn default_content_type(&self) -> &'static str {
71        ARROW_CONTENT_TYPE
72    }
73
74    fn accepts_content_type(&self, ct: &ContentType<'_>) -> bool {
75        ct.matches(ARROW_CONTENT_TYPE)
76    }
77}
78
79fn arrow_error(err: ArrowError) -> StreamError {
80    StreamError::new(StreamErrorKind::CodecError, Some(Box::new(err)), None)
81}
82
83/// Per-stream state for [`ArrowRecordBatchIpcStreamFormat`].
84///
85/// Genuinely stateful, unlike every other format here: the dictionary tracker spans the whole
86/// stream, so batches cannot be encoded independently of one another.
87pub struct ArrowIpcEncoder {
88    schema: SchemaRef,
89    options: IpcWriteOptions,
90    data_gen: IpcDataGenerator,
91    dictionary_tracker: DictionaryTracker,
92    write_context: IpcWriteContext,
93}
94
95impl ItemEncoder<RecordBatch> for ArrowIpcEncoder {
96    fn encode(
97        &mut self,
98        item: &RecordBatch,
99        index: u64,
100        buf: &mut BytesMut,
101    ) -> Result<(), StreamError> {
102        let mut writer = buf.writer();
103
104        // The schema message goes ahead of the first batch and nowhere else.
105        if index == 0 {
106            let encoded = self.data_gen.schema_to_bytes_with_dictionary_tracker(
107                &self.schema,
108                &mut self.dictionary_tracker,
109                &self.options,
110            );
111            write_message(&mut writer, encoded, &self.options).map_err(arrow_error)?;
112        }
113
114        let (encoded_dictionaries, encoded_message) = self
115            .data_gen
116            .encode(
117                item,
118                &mut self.dictionary_tracker,
119                &self.options,
120                &mut self.write_context,
121            )
122            .map_err(arrow_error)?;
123
124        for encoded_dictionary in encoded_dictionaries {
125            write_message(&mut writer, encoded_dictionary, &self.options).map_err(arrow_error)?;
126        }
127
128        write_message(&mut writer, encoded_message, &self.options).map_err(arrow_error)?;
129        writer.flush().map_err(|err| {
130            StreamError::new(StreamErrorKind::CodecError, Some(Box::new(err)), None)
131        })
132    }
133
134    /// The end-of-stream marker. Emitted only on a clean end, so a truncated body stays
135    /// visibly truncated rather than reading as a complete, empty-tailed stream.
136    fn epilogue(&mut self, buf: &mut BytesMut) -> Result<(), StreamError> {
137        buf.extend_from_slice(&CONTINUATION_MARKER);
138        buf.extend_from_slice(&TOTAL_LEN);
139        Ok(())
140    }
141}
142
143impl StreamFormatEncode<RecordBatch> for ArrowRecordBatchIpcStreamFormat {
144    type Encoder = ArrowIpcEncoder;
145
146    fn encoder(&self) -> Self::Encoder {
147        ArrowIpcEncoder {
148            schema: self.schema.clone(),
149            options: self.options.clone(),
150            data_gen: IpcDataGenerator::default(),
151            dictionary_tracker: DictionaryTracker::new(false),
152            write_context: IpcWriteContext::default(),
153        }
154    }
155}
156
157impl StreamFormatDecode<RecordBatch> for ArrowRecordBatchIpcStreamFormat {
158    type Frame = Result<RecordBatch, StreamError>;
159    type Framer = ArrowIpcCodec;
160    type Parser = IdentityParser;
161
162    fn framer(&self, options: &DecodeOptions) -> Self::Framer {
163        ArrowIpcCodec::new_with_max_length(options.max_obj_len)
164    }
165
166    fn parser(&self) -> Self::Parser {
167        IdentityParser
168    }
169}
http_streams_core/arrow_format.rs

http_streams_core/
arrow_format.rs
