arrow_ipc/reader/
stream.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: See the NOTICE file distributed with this work for additional information regarding copyright ownership
3// SPDX-FileContributor: Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements
4
5use std::collections::HashMap;
6use std::fmt::Debug;
7use std::sync::Arc;
8
9use arrow_array::{ArrayRef, RecordBatch};
10use arrow_buffer::{Buffer, MutableBuffer};
11use arrow_data::UnsafeFlag;
12use arrow_schema::{ArrowError, SchemaRef};
13
14use crate::convert::MessageBuffer;
15use crate::reader::{MAX_PREALLOC_BYTES, RecordBatchDecoder, read_dictionary_impl};
16use crate::{CONTINUATION_MARKER, MessageHeader};
17
18/// A low-level interface for reading [`RecordBatch`] data from a stream of bytes
19///
20/// See [StreamReader](crate::reader::StreamReader) for a higher-level interface
21#[derive(Debug, Default)]
22pub struct StreamDecoder {
23    /// The schema of this decoder, if read
24    schema: Option<SchemaRef>,
25    /// Lookup table for dictionaries by ID
26    dictionaries: HashMap<i64, ArrayRef>,
27    /// The decoder state
28    state: DecoderState,
29    /// A scratch buffer when a read is split across multiple `Buffer`
30    buf: MutableBuffer,
31    /// Whether or not array data in input buffers are required to be aligned
32    require_alignment: bool,
33    /// Should validation be skipped when reading data? Defaults to false.
34    ///
35    /// See [`StreamDecoder::with_skip_validation`] for details.
36    ///
37    skip_validation: UnsafeFlag,
38}
39
40#[derive(Debug)]
41enum DecoderState {
42    /// Decoding the message header
43    Header {
44        /// Temporary buffer
45        buf: [u8; 4],
46        /// Number of bytes read into buf
47        read: u8,
48        /// If we have read a continuation token
49        continuation: bool,
50    },
51    /// Decoding the message flatbuffer
52    Message {
53        /// The size of the message flatbuffer
54        size: u32,
55    },
56    /// Decoding the message body
57    Body {
58        /// The message flatbuffer
59        message: MessageBuffer,
60    },
61    /// Reached the end of the stream
62    Finished,
63}
64
65impl Default for DecoderState {
66    fn default() -> Self {
67        Self::Header {
68            buf: [0; 4],
69            read: 0,
70            continuation: false,
71        }
72    }
73}
74
75impl StreamDecoder {
76    /// Create a new [`StreamDecoder`]
77    pub fn new() -> Self {
78        Self::default()
79    }
80
81    /// Specifies whether or not array data in input buffers is required to be properly aligned.
82    ///
83    /// If `require_alignment` is true, this decoder will return an error if any array data in the
84    /// input `buf` is not properly aligned.
85    /// Under the hood it will use [`arrow_data::ArrayDataBuilder::build`] to construct
86    /// [`arrow_data::ArrayData`].
87    ///
88    /// If `require_alignment` is false (the default), this decoder will automatically allocate a
89    /// new aligned buffer and copy over the data if any array data in the input `buf` is not
90    /// properly aligned. (Properly aligned array data will remain zero-copy.)
91    /// Under the hood it will use [`arrow_data::ArrayDataBuilder::align_buffers`] to construct
92    /// [`arrow_data::ArrayData`].
93    pub fn with_require_alignment(mut self, require_alignment: bool) -> Self {
94        self.require_alignment = require_alignment;
95        self
96    }
97
98    /// Return the schema if decoded, else None.
99    pub fn schema(&self) -> Option<SchemaRef> {
100        self.schema.as_ref().map(|schema| schema.clone())
101    }
102
103    /// Specifies if validation should be skipped when reading data (defaults to `false`)
104    ///
105    /// # Safety
106    ///
107    /// This flag must only be set to `true` when you trust the input data and are
108    /// sure the data you are reading is valid Arrow IPC stream data, otherwise
109    /// undefined behavior may result.
110    ///
111    /// For example, DataFusion uses this when reading spill files it wrote itself.
112    pub unsafe fn with_skip_validation(mut self, skip_validation: bool) -> Self {
113        unsafe { self.skip_validation.set(skip_validation) };
114        self
115    }
116
117    /// Try to read the next [`RecordBatch`] from the provided [`Buffer`]
118    ///
119    /// [`Buffer::advance`] will be called on `buffer` for any consumed bytes.
120    ///
121    /// The push-based interface facilitates integration with sources that yield arbitrarily
122    /// delimited bytes ranges, such as a chunked byte stream received from object storage
123    ///
124    /// ```
125    /// # use arrow_array::RecordBatch;
126    /// # use arrow_buffer::Buffer;
127    /// # use arrow_ipc::reader::StreamDecoder;
128    /// # use arrow_schema::ArrowError;
129    /// #
130    /// fn print_stream<I>(src: impl Iterator<Item = Buffer>) -> Result<(), ArrowError> {
131    ///     let mut decoder = StreamDecoder::new();
132    ///     for mut x in src {
133    ///         while !x.is_empty() {
134    ///             if let Some(x) = decoder.decode(&mut x)? {
135    ///                 println!("{x:?}");
136    ///             }
137    ///             if let Some(schema) = decoder.schema() {
138    ///                 println!("Schema: {schema:?}");
139    ///             }
140    ///         }
141    ///     }
142    ///     decoder.finish().unwrap();
143    ///     Ok(())
144    /// }
145    /// ```
146    pub fn decode(&mut self, buffer: &mut Buffer) -> Result<Option<RecordBatch>, ArrowError> {
147        while !buffer.is_empty()
148            || matches!(
149                    &self.state,
150                    DecoderState::Body { message } if message.as_ref().bodyLength() == 0
151            )
152        {
153            match &mut self.state {
154                DecoderState::Header {
155                    buf,
156                    read,
157                    continuation,
158                } => {
159                    let offset_buf = &mut buf[*read as usize..];
160                    let to_read = buffer.len().min(offset_buf.len());
161                    offset_buf[..to_read].copy_from_slice(&buffer[..to_read]);
162                    *read += to_read as u8;
163                    buffer.advance(to_read);
164                    if *read == 4 {
165                        if !*continuation && buf == &CONTINUATION_MARKER {
166                            *continuation = true;
167                            *read = 0;
168                            continue;
169                        }
170                        let size = u32::from_le_bytes(*buf);
171
172                        if size == 0 {
173                            self.state = DecoderState::Finished;
174                            continue;
175                        }
176                        self.state = DecoderState::Message { size };
177                    }
178                }
179                DecoderState::Message { size } => {
180                    let len = *size as usize;
181                    if self.buf.is_empty() && buffer.len() > len {
182                        let message = MessageBuffer::try_new(buffer.slice_with_length(0, len))?;
183                        self.state = DecoderState::Body { message };
184                        buffer.advance(len);
185                        continue;
186                    }
187
188                    reserve_declared(&mut self.buf, len)?;
189                    let to_read = buffer.len().min(len - self.buf.len());
190                    self.buf
191                        .try_extend_from_slice(&buffer[..to_read])
192                        .map_err(|e| ArrowError::MemoryError(e.to_string()))?;
193                    buffer.advance(to_read);
194                    if self.buf.len() == len {
195                        let message = MessageBuffer::try_new(std::mem::take(&mut self.buf).into())?;
196                        self.state = DecoderState::Body { message };
197                    }
198                }
199                DecoderState::Body { message } => {
200                    let message = message.as_ref();
201                    let body_length = message.bodyLength() as usize;
202
203                    let body = if self.buf.is_empty() && buffer.len() >= body_length {
204                        let body = buffer.slice_with_length(0, body_length);
205                        buffer.advance(body_length);
206                        body
207                    } else {
208                        reserve_declared(&mut self.buf, body_length)?;
209                        let to_read = buffer.len().min(body_length - self.buf.len());
210                        self.buf
211                            .try_extend_from_slice(&buffer[..to_read])
212                            .map_err(|e| ArrowError::MemoryError(e.to_string()))?;
213                        buffer.advance(to_read);
214
215                        if self.buf.len() != body_length {
216                            continue;
217                        }
218                        std::mem::take(&mut self.buf).into()
219                    };
220
221                    let version = message.version();
222                    match message.header_type() {
223                        MessageHeader::Schema => {
224                            if self.schema.is_some() {
225                                return Err(ArrowError::IpcError(
226                                    "Not expecting a schema when messages are read".to_string(),
227                                ));
228                            }
229
230                            let ipc_schema = message.header_as_schema().ok_or_else(|| {
231                                ArrowError::ParseError(
232                                    "Unable to read IPC message as schema".to_string(),
233                                )
234                            })?;
235                            let schema = crate::convert::try_fb_to_schema(ipc_schema)?;
236                            self.state = DecoderState::default();
237                            self.schema = Some(Arc::new(schema));
238                        }
239                        MessageHeader::RecordBatch => {
240                            let batch = message.header_as_record_batch().unwrap();
241                            let schema = self.schema.clone().ok_or_else(|| {
242                                ArrowError::IpcError("Missing schema".to_string())
243                            })?;
244                            let batch = RecordBatchDecoder::try_new(
245                                &body,
246                                batch,
247                                schema,
248                                &self.dictionaries,
249                                &version,
250                            )?
251                            .with_require_alignment(self.require_alignment)
252                            .with_skip_validation(self.skip_validation.clone())
253                            .read_record_batch()?;
254                            self.state = DecoderState::default();
255                            return Ok(Some(batch));
256                        }
257                        MessageHeader::DictionaryBatch => {
258                            let dictionary = message.header_as_dictionary_batch().unwrap();
259                            let schema = self.schema.as_deref().ok_or_else(|| {
260                                ArrowError::IpcError("Missing schema".to_string())
261                            })?;
262                            read_dictionary_impl(
263                                &body,
264                                dictionary,
265                                schema,
266                                &mut self.dictionaries,
267                                &version,
268                                self.require_alignment,
269                                self.skip_validation.clone(),
270                            )?;
271                            self.state = DecoderState::default();
272                        }
273                        MessageHeader::NONE => {
274                            self.state = DecoderState::default();
275                        }
276                        t => {
277                            return Err(ArrowError::IpcError(format!(
278                                "Message type unsupported by StreamDecoder: {t:?}"
279                            )));
280                        }
281                    }
282                }
283                DecoderState::Finished => {
284                    return Err(ArrowError::IpcError("Unexpected EOS".to_string()));
285                }
286            }
287        }
288        Ok(None)
289    }
290
291    /// Signal the end of stream
292    ///
293    /// Returns an error if any partial data remains in the stream
294    pub fn finish(&mut self) -> Result<(), ArrowError> {
295        match self.state {
296            DecoderState::Finished
297            | DecoderState::Header {
298                read: 0,
299                continuation: false,
300                ..
301            } => Ok(()),
302            _ => Err(ArrowError::IpcError("Unexpected End of Stream".to_string())),
303        }
304    }
305}
306
307/// Reserves a message's declared length before copying its first bytes into an empty `buf`,
308/// bounded as in `read_body_bounded`. Growing with each input buffer instead doubles the
309/// capacity, and the decoded batch, which is zero-copy over `buf`, keeps up to twice its size.
310fn reserve_declared(buf: &mut MutableBuffer, len: usize) -> Result<(), ArrowError> {
311    if buf.is_empty() {
312        buf.try_reserve(len.min(MAX_PREALLOC_BYTES))
313            .map_err(|e| ArrowError::MemoryError(e.to_string()))?;
314    }
315    Ok(())
316}
317
318#[cfg(test)]
319mod tests {
320    use super::*;
321    use crate::writer::{IpcWriteOptions, StreamWriter};
322    use arrow_array::{
323        DictionaryArray, Int32Array, Int64Array, RecordBatch, RunArray, StringArray,
324        types::Int32Type,
325    };
326    use arrow_schema::{DataType, Field, Schema};
327
328    // Further tests in arrow-integration-testing/tests/ipc_reader.rs
329
330    #[test]
331    fn test_eos() {
332        let schema = Arc::new(Schema::new(vec![
333            Field::new("int32", DataType::Int32, false),
334            Field::new("int64", DataType::Int64, false),
335        ]));
336
337        let input = RecordBatch::try_new(
338            schema.clone(),
339            vec![
340                Arc::new(Int32Array::from(vec![1, 2, 3])) as _,
341                Arc::new(Int64Array::from(vec![1, 2, 3])) as _,
342            ],
343        )
344        .unwrap();
345
346        let mut buf = Vec::with_capacity(1024);
347        let mut s = StreamWriter::try_new(&mut buf, &schema).unwrap();
348        s.write(&input).unwrap();
349        s.finish().unwrap();
350        drop(s);
351
352        let buffer = Buffer::from_vec(buf);
353
354        let mut b = buffer.slice_with_length(0, buffer.len() - 1);
355        let mut decoder = StreamDecoder::new();
356        let output = decoder.decode(&mut b).unwrap().unwrap();
357        assert_eq!(output, input);
358        assert_eq!(b.len(), 7); // 8 byte EOS truncated by 1 byte
359        assert!(decoder.decode(&mut b).unwrap().is_none());
360
361        let err = decoder.finish().unwrap_err().to_string();
362        assert_eq!(err, "Ipc error: Unexpected End of Stream");
363    }
364
365    #[test]
366    fn test_schema() {
367        let schema = Arc::new(Schema::new(vec![
368            Field::new("int32", DataType::Int32, false),
369            Field::new("int64", DataType::Int64, false),
370        ]));
371
372        let mut buf = Vec::with_capacity(1024);
373        let mut s = StreamWriter::try_new(&mut buf, &schema).unwrap();
374        s.finish().unwrap();
375        drop(s);
376
377        let buffer = Buffer::from_vec(buf);
378
379        let mut b = buffer.slice_with_length(0, buffer.len() - 1);
380        let mut decoder = StreamDecoder::new();
381        let output = decoder.decode(&mut b).unwrap();
382        assert!(output.is_none());
383        let decoded_schema = decoder.schema().unwrap();
384        assert_eq!(schema, decoded_schema);
385
386        let err = decoder.finish().unwrap_err().to_string();
387        assert_eq!(err, "Ipc error: Unexpected End of Stream");
388    }
389
390    #[test]
391    fn test_schema_without_eos() {
392        let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)]));
393
394        let mut buf = Vec::new();
395        let writer = StreamWriter::try_new(&mut buf, &schema).unwrap();
396        // The schema has a zero-length body. Drop without finish to omit EOS.
397        drop(writer);
398
399        let mut buffer = Buffer::from_vec(buf);
400        let mut decoder = StreamDecoder::new();
401        let output = decoder.decode(&mut buffer).unwrap();
402
403        assert!(output.is_none());
404        assert!(buffer.is_empty());
405        assert_eq!(decoder.schema(), Some(schema));
406        decoder.finish().unwrap();
407    }
408
409    #[test]
410    fn test_read_ree_dict_record_batches_from_buffer() {
411        let schema = Schema::new(vec![Field::new(
412            "test1",
413            DataType::RunEndEncoded(
414                Arc::new(Field::new(
415                    Field::REE_RUN_ENDS_FIELD_DEFAULT_NAME,
416                    DataType::Int32,
417                    false,
418                )),
419                #[expect(deprecated)]
420                Arc::new(Field::new_dict(
421                    Field::REE_VALUES_FIELD_DEFAULT_NAME,
422                    DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
423                    true,
424                    0,
425                    false,
426                )),
427            ),
428            true,
429        )]);
430        let batch = RecordBatch::try_new(
431            schema.clone().into(),
432            vec![Arc::new(
433                RunArray::try_new(
434                    &Int32Array::from(vec![1, 2, 3]),
435                    &vec![Some("a"), None, Some("a")]
436                        .into_iter()
437                        .collect::<DictionaryArray<Int32Type>>(),
438                )
439                .expect("Failed to create RunArray"),
440            )],
441        )
442        .expect("Failed to create RecordBatch");
443
444        let mut buffer = vec![];
445        {
446            let mut writer = StreamWriter::try_new_with_options(
447                &mut buffer,
448                &schema,
449                IpcWriteOptions::default(),
450            )
451            .expect("Failed to create StreamWriter");
452            writer.write(&batch).expect("Failed to write RecordBatch");
453            writer.finish().expect("Failed to finish StreamWriter");
454        }
455
456        let mut decoder = StreamDecoder::new();
457        let buf = &mut Buffer::from(buffer.as_slice());
458        while let Some(batch) = decoder
459            .decode(buf)
460            .map_err(|e| {
461                ArrowError::ExternalError(format!("Failed to decode record batch: {e}").into())
462            })
463            .expect("Failed to decode record batch")
464        {
465            assert_eq!(batch, batch);
466        }
467
468        decoder.finish().expect("Failed to finish decoder");
469    }
470
471    /// Fed one byte at a time, every message is split across input buffers, so its body is
472    /// assembled in the scratch buffer that the decoded batch then shares.
473    #[test]
474    fn test_split_body_is_buffered_at_its_declared_size() {
475        let schema = Arc::new(Schema::new(vec![Field::new("s", DataType::Utf8, false)]));
476        let input = RecordBatch::try_new(
477            schema.clone(),
478            vec![Arc::new(StringArray::from_iter_values(
479                (0..7000).map(|i| format!("value{i:05}")),
480            )) as _],
481        )
482        .unwrap();
483
484        let mut stream = Vec::new();
485        let mut writer = StreamWriter::try_new(&mut stream, &schema).unwrap();
486        writer.write(&input).unwrap();
487        writer.finish().unwrap();
488        drop(writer);
489
490        let mut decoder = StreamDecoder::new();
491        let mut output = None;
492        for byte in stream.chunks(1) {
493            if let Some(batch) = decoder.decode(&mut Buffer::from(byte)).unwrap() {
494                output = Some(batch);
495            }
496        }
497        decoder.finish().unwrap();
498        let output = output.unwrap();
499        assert_eq!(output, input);
500
501        let held = output.column(0).to_data().buffers()[0].capacity();
502        assert!(
503            held <= stream.len(),
504            "decoded batch holds {held} bytes, more than the {} byte stream it came from",
505            stream.len()
506        );
507    }
508
509    /// A declared length is untrusted, so reserving it up front is bounded: an implausible
510    /// one must leave the decoder waiting for more bytes rather than abort on allocation.
511    #[test]
512    fn test_implausible_body_length_reserves_a_bounded_buffer() {
513        for body_length in [i64::MAX, 1 << 50] {
514            let mut fbb = flatbuffers::FlatBufferBuilder::new();
515            let mut message = crate::MessageBuilder::new(&mut fbb);
516            message.add_version(crate::MetadataVersion::V5);
517            message.add_header_type(MessageHeader::NONE);
518            message.add_bodyLength(body_length);
519            let root = message.finish();
520            fbb.finish(root, None);
521            let metadata = fbb.finished_data();
522
523            let mut stream = Vec::new();
524            stream.extend_from_slice(&CONTINUATION_MARKER);
525            stream.extend_from_slice(&(metadata.len() as i32).to_le_bytes());
526            stream.extend_from_slice(metadata);
527            stream.extend_from_slice(&[0xAB; 16]);
528
529            let mut decoder = StreamDecoder::new();
530            let decoded = decoder
531                .decode(&mut Buffer::from(stream.as_slice()))
532                .unwrap();
533            assert!(decoded.is_none(), "body_length {body_length}");
534            assert!(decoder.finish().is_err(), "body_length {body_length}");
535        }
536    }
537}
arrow_ipc/reader/stream.rs

arrow_ipc/reader/
stream.rs
