arrow_ipc/reader/
stream.rs

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

arrow_ipc/reader/
stream.rs
