Skip to main content

arrow_ipc/
reader.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
18//! Arrow IPC File and Stream Readers
19//!
20//! # Notes
21//!
22//! The [`FileReader`] and [`StreamReader`] have similar interfaces,
23//! however the [`FileReader`] expects a reader that supports [`Seek`]ing
24//!
25//! [`Seek`]: std::io::Seek
26
27mod stream;
28pub use stream::*;
29
30use arrow_select::concat;
31
32use flatbuffers::{VectorIter, VerifierOptions};
33use std::collections::{HashMap, VecDeque};
34use std::fmt;
35use std::io::{BufReader, Read, Seek, SeekFrom};
36use std::sync::Arc;
37
38use arrow_array::*;
39use arrow_buffer::{
40    ArrowNativeType, BooleanBuffer, Buffer, MutableBuffer, NullBuffer, ScalarBuffer,
41};
42use arrow_data::{ArrayData, ArrayDataBuilder, UnsafeFlag};
43use arrow_schema::*;
44
45use crate::compression::{CompressionCodec, DecompressionContext};
46use crate::r#gen::Message;
47use crate::{Block, CONTINUATION_MARKER, FieldNode, MetadataVersion};
48use DataType::*;
49
50/// Read a buffer based on offset and length
51/// From <https://github.com/apache/arrow/blob/6a936c4ff5007045e86f65f1a6b6c3c955ad5103/format/Message.fbs#L58>
52/// Each constituent buffer is first compressed with the indicated
53/// compressor, and then written with the uncompressed length in the first 8
54/// bytes as a 64-bit little-endian signed integer followed by the compressed
55/// buffer bytes (and then padding as required by the protocol). The
56/// uncompressed length may be set to -1 to indicate that the data that
57/// follows is not compressed, which can be useful for cases where
58/// compression does not yield appreciable savings.
59fn read_buffer(
60    buf: &crate::Buffer,
61    a_data: &Buffer,
62    compression_codec: Option<CompressionCodec>,
63    decompression_context: &mut DecompressionContext,
64) -> Result<Buffer, ArrowError> {
65    let start_offset = buf.offset() as usize;
66    let buf_data = a_data.slice_with_length(start_offset, buf.length() as usize);
67    // corner case: empty buffer
68    match (buf_data.is_empty(), compression_codec) {
69        (true, _) | (_, None) => Ok(buf_data),
70        (false, Some(decompressor)) => {
71            decompressor.decompress_to_buffer(&buf_data, decompression_context)
72        }
73    }
74}
75impl RecordBatchDecoder<'_> {
76    /// Coordinates reading arrays based on data types.
77    ///
78    /// `variadic_counts` encodes the number of buffers to read for variadic types (e.g., Utf8View, BinaryView)
79    /// When encounter such types, we pop from the front of the queue to get the number of buffers to read.
80    ///
81    /// Notes:
82    /// * In the IPC format, null buffers are always set, but may be empty. We discard them if an array has 0 nulls
83    /// * Numeric values inside list arrays are often stored as 64-bit values regardless of their data type size.
84    ///   We thus:
85    ///     - check if the bit width of non-64-bit numbers is 64, and
86    ///     - read the buffer as 64-bit (signed integer or float), and
87    ///     - cast the 64-bit array to the appropriate data type
88    fn create_array(
89        &mut self,
90        field: &Field,
91        variadic_counts: &mut VecDeque<i64>,
92    ) -> Result<ArrayRef, ArrowError> {
93        let data_type = field.data_type();
94        match data_type {
95            Utf8 | Binary | LargeBinary | LargeUtf8 => {
96                let field_node = self.next_node(field)?;
97                let buffers = [
98                    self.next_buffer()?,
99                    self.next_buffer()?,
100                    self.next_buffer()?,
101                ];
102                self.create_primitive_array(field_node, data_type, &buffers)
103            }
104            BinaryView | Utf8View => {
105                let count = self.next_variadic_buffer_count(variadic_counts, data_type)?;
106                let buffers = (0..count)
107                    .map(|_| self.next_buffer())
108                    .collect::<Result<Vec<_>, _>>()?;
109                let field_node = self.next_node(field)?;
110                self.create_primitive_array(field_node, data_type, &buffers)
111            }
112            FixedSizeBinary(_) => {
113                let field_node = self.next_node(field)?;
114                let buffers = [self.next_buffer()?, self.next_buffer()?];
115                self.create_primitive_array(field_node, data_type, &buffers)
116            }
117            List(list_field) | LargeList(list_field) | Map(list_field, _) => {
118                let list_node = self.next_node(field)?;
119                let list_buffers = [self.next_buffer()?, self.next_buffer()?];
120                let values = self.create_array(list_field, variadic_counts)?;
121                self.create_list_array(list_node, data_type, &list_buffers, values)
122            }
123            ListView(list_field) | LargeListView(list_field) => {
124                let list_node = self.next_node(field)?;
125                let list_buffers = [
126                    self.next_buffer()?, // null buffer
127                    self.next_buffer()?, // offsets
128                    self.next_buffer()?, // sizes
129                ];
130                let values = self.create_array(list_field, variadic_counts)?;
131                self.create_list_view_array(list_node, data_type, &list_buffers, values)
132            }
133            FixedSizeList(list_field, _) => {
134                let list_node = self.next_node(field)?;
135                let list_buffers = [self.next_buffer()?];
136                let values = self.create_array(list_field, variadic_counts)?;
137                self.create_list_array(list_node, data_type, &list_buffers, values)
138            }
139            Struct(struct_fields) => {
140                let struct_node = self.next_node(field)?;
141                let null_buffer = self.next_buffer()?;
142
143                // read the arrays for each field
144                let mut struct_arrays = Vec::with_capacity(struct_fields.len());
145                // TODO investigate whether just knowing the number of buffers could
146                // still work
147                for struct_field in struct_fields {
148                    let child = self.create_array(struct_field, variadic_counts)?;
149                    struct_arrays.push(child);
150                }
151                self.create_struct_array(struct_node, null_buffer, struct_fields, struct_arrays)
152            }
153            RunEndEncoded(run_ends_field, values_field) => {
154                let run_node = self.next_node(field)?;
155                let run_ends = self.create_array(run_ends_field, variadic_counts)?;
156                let values = self.create_array(values_field, variadic_counts)?;
157
158                let run_array_length = run_node.length() as usize;
159                let builder = ArrayData::builder(data_type.clone())
160                    .len(run_array_length)
161                    .offset(0)
162                    .add_child_data(run_ends.into_data())
163                    .add_child_data(values.into_data())
164                    .null_count(run_node.null_count() as usize);
165
166                self.create_array_from_builder(builder)
167            }
168            // Create dictionary array from RecordBatch
169            Dictionary(_, _) => {
170                let index_node = self.next_node(field)?;
171                let index_buffers = [self.next_buffer()?, self.next_buffer()?];
172
173                #[expect(deprecated)]
174                let dict_id = field.dict_id().ok_or_else(|| {
175                    ArrowError::ParseError(format!("Field {field} does not have dict id"))
176                })?;
177
178                let value_array = match self.dictionaries_by_id.get(&dict_id) {
179                    Some(array) => array.clone(),
180                    None => {
181                        // Per the IPC spec, dictionary batches may be omitted when all
182                        // values in the column are null. In that case we synthesize an
183                        // empty values array so decoding can proceed.
184                        if let Dictionary(_, value_type) = data_type {
185                            arrow_array::new_empty_array(value_type.as_ref())
186                        } else {
187                            unreachable!()
188                        }
189                    }
190                };
191
192                self.create_dictionary_array(index_node, data_type, &index_buffers, value_array)
193            }
194            Union(fields, mode) => {
195                let union_node = self.next_node(field)?;
196                let len = union_node.length() as usize;
197
198                // In V4, union types has validity bitmap
199                // In V5 and later, union types have no validity bitmap
200                if self.version < MetadataVersion::V5 {
201                    self.next_buffer()?;
202                }
203
204                let type_ids: ScalarBuffer<i8> =
205                    self.next_buffer()?.slice_with_length(0, len).into();
206
207                let value_offsets = match mode {
208                    UnionMode::Dense => {
209                        let offsets: ScalarBuffer<i32> =
210                            self.next_buffer()?.slice_with_length(0, len * 4).into();
211                        Some(offsets)
212                    }
213                    UnionMode::Sparse => None,
214                };
215
216                let mut children = Vec::with_capacity(fields.len());
217
218                for (_id, field) in fields.iter() {
219                    let child = self.create_array(field, variadic_counts)?;
220                    children.push(child);
221                }
222
223                let array = if self.skip_validation.get() {
224                    // safety: flag can only be set via unsafe code
225                    unsafe {
226                        UnionArray::new_unchecked(fields.clone(), type_ids, value_offsets, children)
227                    }
228                } else {
229                    UnionArray::try_new(fields.clone(), type_ids, value_offsets, children)?
230                };
231                Ok(Arc::new(array))
232            }
233            Null => {
234                let node = self.next_node(field)?;
235                let length = node.length();
236                let null_count = node.null_count();
237
238                if length != null_count {
239                    return Err(ArrowError::SchemaError(format!(
240                        "Field {field} of NullArray has unequal null_count {null_count} and len {length}"
241                    )));
242                }
243
244                let builder = ArrayData::builder(data_type.clone())
245                    .len(length as usize)
246                    .offset(0);
247                self.create_array_from_builder(builder)
248            }
249            _ => {
250                let field_node = self.next_node(field)?;
251                let buffers = [self.next_buffer()?, self.next_buffer()?];
252                self.create_primitive_array(field_node, data_type, &buffers)
253            }
254        }
255    }
256
257    /// Reads the correct number of buffers based on data type and null_count, and creates a
258    /// primitive array ref
259    fn create_primitive_array(
260        &self,
261        field_node: &FieldNode,
262        data_type: &DataType,
263        buffers: &[Buffer],
264    ) -> Result<ArrayRef, ArrowError> {
265        let length = field_node.length() as usize;
266        let null_buffer = (field_node.null_count() > 0).then_some(buffers[0].clone());
267        let mut builder = match data_type {
268            Utf8 | Binary | LargeBinary | LargeUtf8 => {
269                // read 3 buffers: null buffer (optional), offsets buffer and data buffer
270                ArrayData::builder(data_type.clone())
271                    .len(length)
272                    .buffers(buffers[1..3].to_vec())
273                    .null_bit_buffer(null_buffer)
274            }
275            BinaryView | Utf8View => ArrayData::builder(data_type.clone())
276                .len(length)
277                .buffers(buffers[1..].to_vec())
278                .null_bit_buffer(null_buffer),
279            _ if data_type.is_primitive() || matches!(data_type, Boolean | FixedSizeBinary(_)) => {
280                // read 2 buffers: null buffer (optional) and data buffer
281                ArrayData::builder(data_type.clone())
282                    .len(length)
283                    .add_buffer(buffers[1].clone())
284                    .null_bit_buffer(null_buffer)
285            }
286            t => unreachable!("Data type {:?} either unsupported or not primitive", t),
287        };
288
289        builder = builder.null_count(field_node.null_count() as usize);
290
291        self.create_array_from_builder(builder)
292    }
293
294    /// Update the ArrayDataBuilder based on settings in this decoder
295    fn create_array_from_builder(&self, builder: ArrayDataBuilder) -> Result<ArrayRef, ArrowError> {
296        let mut builder = builder.align_buffers(!self.require_alignment);
297        if self.skip_validation.get() {
298            // SAFETY: flag can only be set via unsafe code
299            unsafe { builder = builder.skip_validation(true) }
300        }
301        Ok(make_array(builder.build()?))
302    }
303
304    /// Reads the correct number of buffers based on list type and null_count, and creates a
305    /// list array ref
306    fn create_list_array(
307        &self,
308        field_node: &FieldNode,
309        data_type: &DataType,
310        buffers: &[Buffer],
311        child_array: ArrayRef,
312    ) -> Result<ArrayRef, ArrowError> {
313        let null_buffer = (field_node.null_count() > 0).then_some(buffers[0].clone());
314        let length = field_node.length() as usize;
315        let child_data = child_array.into_data();
316        let mut builder = match data_type {
317            List(_) | LargeList(_) | Map(_, _) => ArrayData::builder(data_type.clone())
318                .len(length)
319                .add_buffer(buffers[1].clone())
320                .add_child_data(child_data)
321                .null_bit_buffer(null_buffer),
322
323            FixedSizeList(_, _) => ArrayData::builder(data_type.clone())
324                .len(length)
325                .add_child_data(child_data)
326                .null_bit_buffer(null_buffer),
327
328            _ => unreachable!("Cannot create list or map array from {:?}", data_type),
329        };
330
331        builder = builder.null_count(field_node.null_count() as usize);
332
333        self.create_array_from_builder(builder)
334    }
335
336    fn create_list_view_array(
337        &self,
338        field_node: &FieldNode,
339        data_type: &DataType,
340        buffers: &[Buffer],
341        child_array: ArrayRef,
342    ) -> Result<ArrayRef, ArrowError> {
343        assert!(matches!(data_type, ListView(_) | LargeListView(_)));
344
345        let null_buffer = (field_node.null_count() > 0).then_some(buffers[0].clone());
346        let length = field_node.length() as usize;
347        let child_data = child_array.into_data();
348
349        self.create_array_from_builder(
350            ArrayData::builder(data_type.clone())
351                .len(length)
352                .add_buffer(buffers[1].clone()) // offsets
353                .add_buffer(buffers[2].clone()) // sizes
354                .add_child_data(child_data)
355                .null_bit_buffer(null_buffer)
356                .null_count(field_node.null_count() as usize),
357        )
358    }
359
360    fn create_struct_array(
361        &self,
362        struct_node: &FieldNode,
363        null_buffer: Buffer,
364        struct_fields: &Fields,
365        struct_arrays: Vec<ArrayRef>,
366    ) -> Result<ArrayRef, ArrowError> {
367        let null_count = struct_node.null_count() as usize;
368        let len = struct_node.length() as usize;
369        let skip_validation = self.skip_validation.get();
370
371        let nulls = if null_count > 0 {
372            if !skip_validation {
373                // Check before constructing the BooleanBuffer, which would otherwise panic.
374                let needed_len = len.div_ceil(8);
375                if null_buffer.len() < needed_len {
376                    return Err(ArrowError::InvalidArgumentError(format!(
377                        "null_bit_buffer size too small. got {} needed {}",
378                        null_buffer.len(),
379                        needed_len
380                    )));
381                }
382            }
383            let validity_buffer = BooleanBuffer::new(null_buffer, 0, len);
384            let null_buffer = if skip_validation {
385                // safety: flag can only be set via unsafe code
386                unsafe { NullBuffer::new_unchecked(validity_buffer, null_count) }
387            } else {
388                let null_buffer = NullBuffer::new(validity_buffer);
389
390                if null_buffer.null_count() != null_count {
391                    return Err(ArrowError::InvalidArgumentError(format!(
392                        "null_count value ({}) doesn't match actual number of nulls in array ({})",
393                        null_count,
394                        null_buffer.null_count()
395                    )));
396                }
397
398                null_buffer
399            };
400
401            Some(null_buffer)
402        } else {
403            None
404        };
405        if struct_arrays.is_empty() {
406            // `StructArray::from` can't infer the correct row count
407            // if we have zero fields
408            return Ok(Arc::new(StructArray::new_empty_fields(len, nulls)));
409        }
410
411        let struct_array = if skip_validation {
412            // safety: flag can only be set via unsafe code
413            unsafe { StructArray::new_unchecked(struct_fields.clone(), struct_arrays, nulls) }
414        } else {
415            StructArray::try_new(struct_fields.clone(), struct_arrays, nulls)?
416        };
417
418        Ok(Arc::new(struct_array))
419    }
420
421    /// Reads the correct number of buffers based on list type and null_count, and creates a
422    /// list array ref
423    fn create_dictionary_array(
424        &self,
425        field_node: &FieldNode,
426        data_type: &DataType,
427        buffers: &[Buffer],
428        value_array: ArrayRef,
429    ) -> Result<ArrayRef, ArrowError> {
430        if let Dictionary(_, _) = *data_type {
431            let null_buffer = (field_node.null_count() > 0).then_some(buffers[0].clone());
432            let builder = ArrayData::builder(data_type.clone())
433                .len(field_node.length() as usize)
434                .add_buffer(buffers[1].clone())
435                .add_child_data(value_array.into_data())
436                .null_bit_buffer(null_buffer)
437                .null_count(field_node.null_count() as usize);
438            self.create_array_from_builder(builder)
439        } else {
440            unreachable!("Cannot create dictionary array from {:?}", data_type)
441        }
442    }
443}
444
445/// State for decoding Arrow arrays from an [IPC RecordBatch] structure to
446/// [`RecordBatch`]
447///
448/// [IPC RecordBatch]: crate::RecordBatch
449///
450pub struct RecordBatchDecoder<'a> {
451    /// The flatbuffers encoded record batch
452    batch: crate::RecordBatch<'a>,
453    /// The output schema
454    schema: SchemaRef,
455    /// Decoded dictionaries indexed by dictionary id
456    dictionaries_by_id: &'a HashMap<i64, ArrayRef>,
457    /// Optional compression codec
458    compression: Option<CompressionCodec>,
459    /// Decompression context for reusing zstd decompressor state
460    decompression_context: DecompressionContext,
461    /// The format version
462    version: MetadataVersion,
463    /// The raw data buffer
464    data: &'a Buffer,
465    /// The fields comprising this array
466    nodes: VectorIter<'a, FieldNode>,
467    /// The buffers comprising this array
468    buffers: VectorIter<'a, crate::Buffer>,
469    /// Projection (subset of columns) to read, if any
470    /// See [`RecordBatchDecoder::with_projection`] for details
471    projection: Option<&'a [usize]>,
472    /// Are buffers required to already be aligned? See
473    /// [`RecordBatchDecoder::with_require_alignment`] for details
474    require_alignment: bool,
475    /// Should validation be skipped when reading data? Defaults to false.
476    ///
477    /// See [`FileDecoder::with_skip_validation`] for details.
478    skip_validation: UnsafeFlag,
479}
480
481impl<'a> RecordBatchDecoder<'a> {
482    /// Create a reader for decoding arrays from an encoded [`RecordBatch`]
483    pub fn try_new(
484        buf: &'a Buffer,
485        batch: crate::RecordBatch<'a>,
486        schema: SchemaRef,
487        dictionaries_by_id: &'a HashMap<i64, ArrayRef>,
488        metadata: &'a MetadataVersion,
489    ) -> Result<Self, ArrowError> {
490        let buffers = batch.buffers().ok_or_else(|| {
491            ArrowError::IpcError("Unable to get buffers from IPC RecordBatch".to_string())
492        })?;
493        let field_nodes = batch.nodes().ok_or_else(|| {
494            ArrowError::IpcError("Unable to get field nodes from IPC RecordBatch".to_string())
495        })?;
496
497        let batch_compression = batch.compression();
498        let compression = batch_compression
499            .map(|batch_compression| batch_compression.codec().try_into())
500            .transpose()?;
501
502        Ok(Self {
503            batch,
504            schema,
505            dictionaries_by_id,
506            compression,
507            decompression_context: DecompressionContext::new(),
508            version: *metadata,
509            data: buf,
510            nodes: field_nodes.iter(),
511            buffers: buffers.iter(),
512            projection: None,
513            require_alignment: false,
514            skip_validation: UnsafeFlag::new(),
515        })
516    }
517
518    /// Set the projection (default: None)
519    ///
520    /// If set, the projection is the list  of column indices
521    /// that will be read
522    pub fn with_projection(mut self, projection: Option<&'a [usize]>) -> Self {
523        self.projection = projection;
524        self
525    }
526
527    /// Set require_alignment (default: false)
528    ///
529    /// If true, buffers must be aligned appropriately or error will
530    /// result. If false, buffers will be copied to aligned buffers
531    /// if necessary.
532    pub fn with_require_alignment(mut self, require_alignment: bool) -> Self {
533        self.require_alignment = require_alignment;
534        self
535    }
536
537    /// Specifies if validation should be skipped when reading data (defaults to `false`)
538    ///
539    /// When enabled, the following checks are bypassed:
540    /// - Offset bounds (e.g. list/string offsets pointing past the end of their value buffer)
541    /// - UTF-8 validity of string columns (`Utf8` / `LargeUtf8`)
542    /// - Null count consistency and buffer length checks
543    ///
544    /// # Undefined behavior
545    ///
546    /// Relies on the caller only passing a flag with `true` value if they are
547    /// certain that the data is valid. Invalid data that bypasses these checks
548    /// may cause undefined behavior when the arrays are later accessed.
549    pub fn with_skip_validation(mut self, skip_validation: UnsafeFlag) -> Self {
550        self.skip_validation = skip_validation;
551        self
552    }
553
554    /// Read the record batch, consuming the reader
555    ///
556    /// # Errors
557    ///
558    /// Returns an error if the message does not describe a batch matching the schema,
559    /// for example if it declares more variadic buffer counts than the schema uses.
560    pub fn read_record_batch(mut self) -> Result<RecordBatch, ArrowError> {
561        let mut variadic_counts: VecDeque<i64> = self
562            .batch
563            .variadicBufferCounts()
564            .into_iter()
565            .flatten()
566            .collect();
567
568        let options = RecordBatchOptions::new().with_row_count(Some(self.batch.length() as usize));
569
570        let schema = Arc::clone(&self.schema);
571        if let Some(projection) = self.projection {
572            let mut arrays = Vec::with_capacity(projection.len());
573            // project fields
574            for (idx, field) in schema.fields().iter().enumerate() {
575                // A projected field can appear more than once, so collect all matching positions.
576                let mut decoded = None;
577                for (proj_idx, projected_idx) in projection.iter().enumerate() {
578                    if *projected_idx == idx {
579                        // Reuse the decoded array for duplicate projection entries.
580                        let child = match decoded.clone() {
581                            Some(child) => child,
582                            None => {
583                                let child = self.create_array(field, &mut variadic_counts)?;
584                                decoded = Some(Arc::clone(&child));
585                                child
586                            }
587                        };
588                        arrays.push((proj_idx, child));
589                    }
590                }
591
592                if decoded.is_none() {
593                    self.skip_field(field, &mut variadic_counts)?;
594                }
595            }
596
597            arrays.sort_by_key(|t| t.0);
598
599            let schema = Arc::new(schema.project(projection)?);
600            let columns = arrays.into_iter().map(|t| t.1).collect::<Vec<_>>();
601
602            if self.skip_validation.get() {
603                // Safety: setting `skip_validation` requires `unsafe`, user assures data is valid
604                unsafe {
605                    Ok(RecordBatch::new_unchecked(
606                        schema,
607                        columns,
608                        self.batch.length() as usize,
609                    ))
610                }
611            } else {
612                check_variadic_counts_consumed(&variadic_counts)?;
613                RecordBatch::try_new_with_options(schema, columns, &options)
614            }
615        } else {
616            let mut children = Vec::with_capacity(schema.fields().len());
617            // keep track of index as lists require more than one node
618            for field in schema.fields() {
619                let child = self.create_array(field, &mut variadic_counts)?;
620                children.push(child);
621            }
622
623            if self.skip_validation.get() {
624                // Safety: setting `skip_validation` requires `unsafe`, user assures data is valid
625                unsafe {
626                    Ok(RecordBatch::new_unchecked(
627                        schema,
628                        children,
629                        self.batch.length() as usize,
630                    ))
631                }
632            } else {
633                check_variadic_counts_consumed(&variadic_counts)?;
634                RecordBatch::try_new_with_options(schema, children, &options)
635            }
636        }
637    }
638
639    fn next_buffer(&mut self) -> Result<Buffer, ArrowError> {
640        let buffer = self.buffers.next().ok_or_else(|| {
641            ArrowError::IpcError("Buffer count mismatched with metadata".to_string())
642        })?;
643        read_buffer(
644            buffer,
645            self.data,
646            self.compression,
647            &mut self.decompression_context,
648        )
649    }
650
651    fn skip_buffer(&mut self) -> Result<(), ArrowError> {
652        self.buffers.next().ok_or_else(|| {
653            ArrowError::IpcError("Buffer count mismatched with metadata".to_string())
654        })?;
655        Ok(())
656    }
657
658    fn next_node(&mut self, field: &Field) -> Result<&'a FieldNode, ArrowError> {
659        self.nodes.next().ok_or_else(|| {
660            ArrowError::SchemaError(format!(
661                "Invalid data for schema. {field} refers to node not found in schema",
662            ))
663        })
664    }
665
666    fn skip_field(
667        &mut self,
668        field: &Field,
669        variadic_count: &mut VecDeque<i64>,
670    ) -> Result<(), ArrowError> {
671        self.next_node(field)?;
672
673        match field.data_type() {
674            Utf8 | Binary | LargeBinary | LargeUtf8 => {
675                for _ in 0..3 {
676                    self.skip_buffer()?;
677                }
678            }
679            Utf8View | BinaryView => {
680                let count = self.next_variadic_buffer_count(variadic_count, field.data_type())?;
681                for _ in 0..count {
682                    self.skip_buffer()?;
683                }
684            }
685            FixedSizeBinary(_) => {
686                self.skip_buffer()?;
687                self.skip_buffer()?;
688            }
689            List(list_field) | LargeList(list_field) | Map(list_field, _) => {
690                self.skip_buffer()?;
691                self.skip_buffer()?;
692                self.skip_field(list_field, variadic_count)?;
693            }
694            ListView(list_field) | LargeListView(list_field) => {
695                self.skip_buffer()?; // Null buffer
696                self.skip_buffer()?; // Offsets
697                self.skip_buffer()?; // Sizes
698                self.skip_field(list_field, variadic_count)?;
699            }
700            FixedSizeList(list_field, _) => {
701                self.skip_buffer()?;
702                self.skip_field(list_field, variadic_count)?;
703            }
704            Struct(struct_fields) => {
705                self.skip_buffer()?;
706
707                // skip for each field
708                for struct_field in struct_fields {
709                    self.skip_field(struct_field, variadic_count)?
710                }
711            }
712            RunEndEncoded(run_ends_field, values_field) => {
713                self.skip_field(run_ends_field, variadic_count)?;
714                self.skip_field(values_field, variadic_count)?;
715            }
716            Dictionary(_, _) => {
717                self.skip_buffer()?; // Nulls
718                self.skip_buffer()?; // Indices
719            }
720            Union(fields, mode) => {
721                if self.version < MetadataVersion::V5 {
722                    self.skip_buffer()?; // Null buffer
723                }
724                self.skip_buffer()?; // Type ids
725
726                match mode {
727                    UnionMode::Dense => self.skip_buffer()?, // Offsets
728                    UnionMode::Sparse => {}
729                }
730
731                for (_, field) in fields.iter() {
732                    self.skip_field(field, variadic_count)?
733                }
734            }
735            // Null has no buffers to skip
736            Null => {}
737
738            // Fixed-width and boolean types: skip null buffer + values buffer
739            Boolean
740            | Int8
741            | Int16
742            | Int32
743            | Int64
744            | UInt8
745            | UInt16
746            | UInt32
747            | UInt64
748            | Float16
749            | Float32
750            | Float64
751            | Timestamp(_, _)
752            | Date32
753            | Date64
754            | Time32(_)
755            | Time64(_)
756            | Duration(_)
757            | Interval(_)
758            | Decimal32(_, _)
759            | Decimal64(_, _)
760            | Decimal128(_, _)
761            | Decimal256(_, _) => {
762                self.skip_buffer()?;
763                self.skip_buffer()?;
764            }
765        }
766        Ok(())
767    }
768}
769
770impl RecordBatchDecoder<'_> {
771    /// Takes the number of variadic buffers declared for one `BinaryView` or `Utf8View`
772    /// column, and returns the total number of buffers to read for it.
773    ///
774    /// The count comes from the IPC message, so it may be missing, negative, or larger
775    /// than the number of buffers the message actually has.
776    fn next_variadic_buffer_count(
777        &self,
778        variadic_counts: &mut VecDeque<i64>,
779        data_type: &DataType,
780    ) -> Result<usize, ArrowError> {
781        let count = variadic_counts.pop_front().ok_or_else(|| {
782            ArrowError::IpcError(format!("Missing variadic count for {data_type} column"))
783        })?;
784
785        let remaining = self.buffers.len();
786
787        // The view buffer and the null buffer are not counted as variadic.
788        usize::try_from(count)
789            .ok()
790            .and_then(|count| count.checked_add(2))
791            .filter(|total| *total <= remaining)
792            .ok_or_else(|| {
793                ArrowError::IpcError(format!(
794                    "Invalid variadic count {count} for {data_type} column, \
795                     with {remaining} buffer(s) left in the message"
796                ))
797            })
798    }
799}
800
801/// Creates a record batch from binary data using the `crate::RecordBatch` indexes and the `Schema`.
802///
803/// If `require_alignment` is true, this function will return an error if any array data in the
804/// input `buf` is not properly aligned.
805/// Under the hood it will use [`arrow_data::ArrayDataBuilder::build`] to construct [`arrow_data::ArrayData`].
806///
807/// If `require_alignment` is false, this function will automatically allocate a new aligned buffer
808/// and copy over the data if any array data in the input `buf` is not properly aligned.
809/// (Properly aligned array data will remain zero-copy.)
810/// Under the hood it will use [`arrow_data::ArrayDataBuilder::align_buffers`] to construct [`arrow_data::ArrayData`].
811pub fn read_record_batch(
812    buf: &Buffer,
813    batch: crate::RecordBatch,
814    schema: SchemaRef,
815    dictionaries_by_id: &HashMap<i64, ArrayRef>,
816    projection: Option<&[usize]>,
817    metadata: &MetadataVersion,
818) -> Result<RecordBatch, ArrowError> {
819    RecordBatchDecoder::try_new(buf, batch, schema, dictionaries_by_id, metadata)?
820        .with_projection(projection)
821        .with_require_alignment(false)
822        .read_record_batch()
823}
824
825/// Read the dictionary from the buffer and provided metadata,
826/// updating the `dictionaries_by_id` with the resulting dictionary
827pub fn read_dictionary(
828    buf: &Buffer,
829    batch: crate::DictionaryBatch,
830    schema: &Schema,
831    dictionaries_by_id: &mut HashMap<i64, ArrayRef>,
832    metadata: &MetadataVersion,
833) -> Result<(), ArrowError> {
834    read_dictionary_impl(
835        buf,
836        batch,
837        schema,
838        dictionaries_by_id,
839        metadata,
840        false,
841        UnsafeFlag::new(),
842    )
843}
844
845/// Low-level version of [`read_dictionary`] with alignment and validation controls
846pub fn read_dictionary_impl(
847    buf: &Buffer,
848    batch: crate::DictionaryBatch,
849    schema: &Schema,
850    dictionaries_by_id: &mut HashMap<i64, ArrayRef>,
851    metadata: &MetadataVersion,
852    require_alignment: bool,
853    skip_validation: UnsafeFlag,
854) -> Result<(), ArrowError> {
855    let id = batch.id();
856
857    let dictionary_values = get_dictionary_values(
858        buf,
859        batch,
860        schema,
861        dictionaries_by_id,
862        metadata,
863        require_alignment,
864        skip_validation,
865    )?;
866
867    update_dictionaries(dictionaries_by_id, batch.isDelta(), id, dictionary_values)?;
868
869    Ok(())
870}
871
872/// Updates the `dictionaries_by_id` with the provided dictionary values and id.
873///
874/// # Errors
875/// - If `is_delta` is true and there is no existing dictionary for the given
876///   `dict_id`
877/// - If `is_delta` is true and the concatenation of the existing and new
878///   dictionary fails. This usually signals a type mismatch between the old and
879///   new values.
880fn update_dictionaries(
881    dictionaries_by_id: &mut HashMap<i64, ArrayRef>,
882    is_delta: bool,
883    dict_id: i64,
884    dict_values: ArrayRef,
885) -> Result<(), ArrowError> {
886    if !is_delta {
887        // We don't currently record the isOrdered field. This could be general
888        // attributes of arrays.
889        // Add (possibly multiple) array refs to the dictionaries array.
890        dictionaries_by_id.insert(dict_id, dict_values.clone());
891        return Ok(());
892    }
893
894    let existing = dictionaries_by_id.get(&dict_id).ok_or_else(|| {
895        ArrowError::InvalidArgumentError(format!(
896            "No existing dictionary for delta dictionary with id '{dict_id}'"
897        ))
898    })?;
899
900    let combined = concat::concat(&[existing, &dict_values]).map_err(|e| {
901        ArrowError::InvalidArgumentError(format!("Failed to concat delta dictionary: {e}"))
902    })?;
903
904    dictionaries_by_id.insert(dict_id, combined);
905
906    Ok(())
907}
908
909/// Given a dictionary batch IPC message/body along with the full state of a
910/// stream including schema, dictionary cache, metadata, and other flags, this
911/// function will parse the buffer into an array of dictionary values.
912fn get_dictionary_values(
913    buf: &Buffer,
914    batch: crate::DictionaryBatch,
915    schema: &Schema,
916    dictionaries_by_id: &HashMap<i64, ArrayRef>,
917    metadata: &MetadataVersion,
918    require_alignment: bool,
919    skip_validation: UnsafeFlag,
920) -> Result<ArrayRef, ArrowError> {
921    let id = batch.id();
922    #[expect(deprecated)]
923    let fields_using_this_dictionary = schema.fields_with_dict_id(id);
924    let first_field = fields_using_this_dictionary.first().ok_or_else(|| {
925        ArrowError::InvalidArgumentError(format!("dictionary id {id} not found in schema"))
926    })?;
927
928    // As the dictionary batch does not contain the type of the
929    // values array, we need to retrieve this from the schema.
930    // Get an array representing this dictionary's values.
931    let dictionary_values: ArrayRef = match first_field.data_type() {
932        DataType::Dictionary(_, value_type) => {
933            // Make a fake schema for the dictionary batch.
934            let value = value_type.as_ref().clone();
935            let schema = Schema::new(vec![Field::new("", value, true)]);
936            // Read a single column
937            let Some(data) = batch.data() else {
938                return Err(ArrowError::ParseError(
939                    "Dictionary batch is missing its data".to_string(),
940                ));
941            };
942            let record_batch = RecordBatchDecoder::try_new(
943                buf,
944                data,
945                Arc::new(schema),
946                dictionaries_by_id,
947                metadata,
948            )?
949            .with_require_alignment(require_alignment)
950            .with_skip_validation(skip_validation)
951            .read_record_batch()?;
952
953            Some(record_batch.column(0).clone())
954        }
955        _ => None,
956    }
957    .ok_or_else(|| {
958        ArrowError::InvalidArgumentError(format!("dictionary id {id} not found in schema"))
959    })?;
960
961    Ok(dictionary_values)
962}
963
964/// Read the data for a given block
965fn read_block<R: Read + Seek>(mut reader: R, block: &Block) -> Result<Buffer, ArrowError> {
966    reader.seek(SeekFrom::Start(block.offset() as u64))?;
967    let body_len = block.bodyLength().to_usize().unwrap();
968    let metadata_len = block.metaDataLength().to_usize().unwrap();
969    let total_len = body_len.checked_add(metadata_len).unwrap();
970
971    let mut buf = MutableBuffer::try_from_len_zeroed(total_len)
972        .map_err(|e| ArrowError::MemoryError(e.to_string()))?;
973    reader.read_exact(&mut buf)?;
974    Ok(buf.into())
975}
976
977/// One variadic buffer count is consumed per `BinaryView` or `Utf8View` column in
978/// the schema, so any count left over means the message and the schema disagree.
979///
980/// The opposite case, too few counts, is reported by [`RecordBatchDecoder::create_array`].
981fn check_variadic_counts_consumed(variadic_counts: &VecDeque<i64>) -> Result<(), ArrowError> {
982    if variadic_counts.is_empty() {
983        Ok(())
984    } else {
985        Err(ArrowError::IpcError(format!(
986            "Mismatch between schema and data: the IPC message declares {} more variadic \
987             buffer count(s) than the schema has BinaryView or Utf8View columns",
988            variadic_counts.len()
989        )))
990    }
991}
992
993/// Parse an encapsulated message
994///
995/// <https://arrow.apache.org/docs/format/Columnar.html#encapsulated-message-format>
996fn parse_message(buf: &[u8]) -> Result<Message::Message<'_>, ArrowError> {
997    let buf = match buf[..4] == CONTINUATION_MARKER {
998        true => &buf[8..],
999        false => &buf[4..],
1000    };
1001    crate::root_as_message(buf)
1002        .map_err(|err| ArrowError::ParseError(format!("Unable to get root as message: {err:?}")))
1003}
1004
1005/// Read the footer length from the last 10 bytes of an Arrow IPC file
1006///
1007/// Expects a 4 byte footer length followed by `b"ARROW1"`
1008pub fn read_footer_length(buf: [u8; 10]) -> Result<usize, ArrowError> {
1009    if buf[4..] != super::ARROW_MAGIC {
1010        return Err(ArrowError::ParseError(
1011            "Arrow file does not contain correct footer".to_string(),
1012        ));
1013    }
1014
1015    // read footer length
1016    let footer_len = i32::from_le_bytes([buf[0], buf[1], buf[2], buf[3]]);
1017    footer_len
1018        .try_into()
1019        .map_err(|_| ArrowError::ParseError(format!("Invalid footer length: {footer_len}")))
1020}
1021
1022/// A low-level, push-based interface for reading an IPC file
1023///
1024/// For a higher-level interface see [`FileReader`]
1025///
1026/// For an example of using this API with `mmap` see the [`zero_copy_ipc`] example.
1027///
1028/// [`zero_copy_ipc`]: https://github.com/apache/arrow-rs/blob/main/arrow/examples/zero_copy_ipc.rs
1029///
1030/// ```
1031/// # use std::sync::Arc;
1032/// # use arrow_array::*;
1033/// # use arrow_array::types::Int32Type;
1034/// # use arrow_buffer::Buffer;
1035/// # use arrow_ipc::convert::try_fb_to_schema;
1036/// # use arrow_ipc::reader::{FileDecoder, read_footer_length};
1037/// # use arrow_ipc::root_as_footer;
1038/// # use arrow_ipc::writer::FileWriter;
1039/// // Write an IPC file
1040///
1041/// let batch = RecordBatch::try_from_iter([
1042///     ("a", Arc::new(Int32Array::from(vec![1, 2, 3])) as _),
1043///     ("b", Arc::new(Int32Array::from(vec![1, 2, 3])) as _),
1044///     ("c", Arc::new(DictionaryArray::<Int32Type>::from_iter(["hello", "hello", "world"])) as _),
1045/// ]).unwrap();
1046///
1047/// let schema = batch.schema();
1048///
1049/// let mut out = Vec::with_capacity(1024);
1050/// let mut writer = FileWriter::try_new(&mut out, schema.as_ref()).unwrap();
1051/// writer.write(&batch).unwrap();
1052/// writer.finish().unwrap();
1053///
1054/// drop(writer);
1055///
1056/// // Read IPC file
1057///
1058/// let buffer = Buffer::from_vec(out);
1059/// let trailer_start = buffer.len() - 10;
1060/// let footer_len = read_footer_length(buffer[trailer_start..].try_into().unwrap()).unwrap();
1061/// let footer = root_as_footer(&buffer[trailer_start - footer_len..trailer_start]).unwrap();
1062///
1063/// let back = try_fb_to_schema(footer.schema().unwrap()).unwrap();
1064/// assert_eq!(&back, schema.as_ref());
1065///
1066/// let mut decoder = FileDecoder::new(schema, footer.version());
1067///
1068/// // Read dictionaries
1069/// for block in footer.dictionaries().iter().flatten() {
1070///     let block_len = block.bodyLength() as usize + block.metaDataLength() as usize;
1071///     let data = buffer.slice_with_length(block.offset() as _, block_len);
1072///     decoder.read_dictionary(&block, &data).unwrap();
1073/// }
1074///
1075/// // Read record batch
1076/// let batches = footer.recordBatches().unwrap();
1077/// assert_eq!(batches.len(), 1); // Only wrote a single batch
1078///
1079/// let block = batches.get(0);
1080/// let block_len = block.bodyLength() as usize + block.metaDataLength() as usize;
1081/// let data = buffer.slice_with_length(block.offset() as _, block_len);
1082/// let back = decoder.read_record_batch(block, &data).unwrap().unwrap();
1083///
1084/// assert_eq!(batch, back);
1085/// ```
1086#[derive(Debug)]
1087pub struct FileDecoder {
1088    schema: SchemaRef,
1089    dictionaries: HashMap<i64, ArrayRef>,
1090    version: MetadataVersion,
1091    projection: Option<Vec<usize>>,
1092    require_alignment: bool,
1093    skip_validation: UnsafeFlag,
1094}
1095
1096impl FileDecoder {
1097    /// Create a new [`FileDecoder`] with the given schema and version
1098    pub fn new(schema: SchemaRef, version: MetadataVersion) -> Self {
1099        Self {
1100            schema,
1101            version,
1102            dictionaries: Default::default(),
1103            projection: None,
1104            require_alignment: false,
1105            skip_validation: UnsafeFlag::new(),
1106        }
1107    }
1108
1109    /// Specify a projection
1110    pub fn with_projection(mut self, projection: Vec<usize>) -> Self {
1111        self.projection = Some(projection);
1112        self
1113    }
1114
1115    /// Specifies if the array data in input buffers is required to be properly aligned.
1116    ///
1117    /// If `require_alignment` is true, this decoder will return an error if any array data in the
1118    /// input `buf` is not properly aligned.
1119    /// Under the hood it will use [`arrow_data::ArrayDataBuilder::build`] to construct
1120    /// [`arrow_data::ArrayData`].
1121    ///
1122    /// If `require_alignment` is false (the default), this decoder will automatically allocate a
1123    /// new aligned buffer and copy over the data if any array data in the input `buf` is not
1124    /// properly aligned. (Properly aligned array data will remain zero-copy.)
1125    /// Under the hood it will use [`arrow_data::ArrayDataBuilder::align_buffers`] to construct
1126    /// [`arrow_data::ArrayData`].
1127    pub fn with_require_alignment(mut self, require_alignment: bool) -> Self {
1128        self.require_alignment = require_alignment;
1129        self
1130    }
1131
1132    /// Specifies if validation should be skipped when reading data (defaults to `false`)
1133    ///
1134    /// # Safety
1135    ///
1136    /// This flag must only be set to `true` when you trust the input data and are sure the data you are
1137    /// reading is a valid Arrow IPC file, otherwise undefined behavior may
1138    /// result.
1139    ///
1140    /// For example, some programs may wish to trust reading IPC files written
1141    /// by the same process that created the files.
1142    pub unsafe fn with_skip_validation(mut self, skip_validation: bool) -> Self {
1143        unsafe { self.skip_validation.set(skip_validation) };
1144        self
1145    }
1146
1147    fn read_message<'a>(&self, buf: &'a [u8]) -> Result<Message::Message<'a>, ArrowError> {
1148        let message = parse_message(buf)?;
1149
1150        // some old test data's footer metadata is not set, so we account for that
1151        if self.version != MetadataVersion::V1 && message.version() != self.version {
1152            return Err(ArrowError::IpcError(
1153                "Could not read IPC message as metadata versions mismatch".to_string(),
1154            ));
1155        }
1156        Ok(message)
1157    }
1158
1159    /// Read the dictionary with the given block and data buffer
1160    pub fn read_dictionary(&mut self, block: &Block, buf: &Buffer) -> Result<(), ArrowError> {
1161        let message = self.read_message(buf)?;
1162        match message.header_type() {
1163            crate::MessageHeader::DictionaryBatch => {
1164                let batch = message.header_as_dictionary_batch().unwrap();
1165                read_dictionary_impl(
1166                    &buf.slice(block.metaDataLength() as _),
1167                    batch,
1168                    &self.schema,
1169                    &mut self.dictionaries,
1170                    &message.version(),
1171                    self.require_alignment,
1172                    self.skip_validation.clone(),
1173                )
1174            }
1175            t => Err(ArrowError::ParseError(format!(
1176                "Expecting DictionaryBatch in dictionary blocks, found {t:?}."
1177            ))),
1178        }
1179    }
1180
1181    /// Read the RecordBatch with the given block and data buffer
1182    pub fn read_record_batch(
1183        &self,
1184        block: &Block,
1185        buf: &Buffer,
1186    ) -> Result<Option<RecordBatch>, ArrowError> {
1187        let message = self.read_message(buf)?;
1188        match message.header_type() {
1189            crate::MessageHeader::Schema => Err(ArrowError::IpcError(
1190                "Not expecting a schema when messages are read".to_string(),
1191            )),
1192            crate::MessageHeader::RecordBatch => {
1193                let batch = message.header_as_record_batch().ok_or_else(|| {
1194                    ArrowError::IpcError("Unable to read IPC message as record batch".to_string())
1195                })?;
1196                // read the block that makes up the record batch into a buffer
1197                RecordBatchDecoder::try_new(
1198                    &buf.slice(block.metaDataLength() as _),
1199                    batch,
1200                    self.schema.clone(),
1201                    &self.dictionaries,
1202                    &message.version(),
1203                )?
1204                .with_projection(self.projection.as_deref())
1205                .with_require_alignment(self.require_alignment)
1206                .with_skip_validation(self.skip_validation.clone())
1207                .read_record_batch()
1208                .map(Some)
1209            }
1210            crate::MessageHeader::NONE => Ok(None),
1211            t => Err(ArrowError::InvalidArgumentError(format!(
1212                "Reading types other than record batches not yet supported, unable to read {t:?}"
1213            ))),
1214        }
1215    }
1216}
1217
1218/// Build an Arrow [`FileReader`] with custom options.
1219#[derive(Debug)]
1220pub struct FileReaderBuilder {
1221    /// Optional projection for which columns to load (zero-based column indices)
1222    projection: Option<Vec<usize>>,
1223    /// Passed through to construct [`VerifierOptions`]
1224    max_footer_fb_tables: usize,
1225    /// Passed through to construct [`VerifierOptions`]
1226    max_footer_fb_depth: usize,
1227}
1228
1229impl Default for FileReaderBuilder {
1230    fn default() -> Self {
1231        let verifier_options = VerifierOptions::default();
1232        Self {
1233            max_footer_fb_tables: verifier_options.max_tables,
1234            max_footer_fb_depth: verifier_options.max_depth,
1235            projection: None,
1236        }
1237    }
1238}
1239
1240impl FileReaderBuilder {
1241    /// Options for creating a new [`FileReader`].
1242    ///
1243    /// To convert a builder into a reader, call [`FileReaderBuilder::build`].
1244    pub fn new() -> Self {
1245        Self::default()
1246    }
1247
1248    /// Optional projection for which columns to load (zero-based column indices).
1249    pub fn with_projection(mut self, projection: Vec<usize>) -> Self {
1250        self.projection = Some(projection);
1251        self
1252    }
1253
1254    /// Flatbuffers option for parsing the footer. Controls the max number of fields and
1255    /// metadata key-value pairs that can be parsed from the schema of the footer.
1256    ///
1257    /// By default this is set to `1_000_000` which roughly translates to a schema with
1258    /// no metadata key-value pairs but 499,999 fields.
1259    ///
1260    /// This default limit is enforced to protect against malicious files with a massive
1261    /// amount of flatbuffer tables which could cause a denial of service attack.
1262    ///
1263    /// If you need to ingest a trusted file with a massive number of fields and/or
1264    /// metadata key-value pairs and are facing the error `"Unable to get root as
1265    /// footer: TooManyTables"` then increase this parameter as necessary.
1266    pub fn with_max_footer_fb_tables(mut self, max_footer_fb_tables: usize) -> Self {
1267        self.max_footer_fb_tables = max_footer_fb_tables;
1268        self
1269    }
1270
1271    /// Flatbuffers option for parsing the footer. Controls the max depth for schemas with
1272    /// nested fields parsed from the footer.
1273    ///
1274    /// By default this is set to `64` which roughly translates to a schema with
1275    /// a field nested 60 levels down through other struct fields.
1276    ///
1277    /// This default limit is enforced to protect against malicious files with a extremely
1278    /// deep flatbuffer structure which could cause a denial of service attack.
1279    ///
1280    /// If you need to ingest a trusted file with a deeply nested field and are facing the
1281    /// error `"Unable to get root as footer: DepthLimitReached"` then increase this
1282    /// parameter as necessary.
1283    pub fn with_max_footer_fb_depth(mut self, max_footer_fb_depth: usize) -> Self {
1284        self.max_footer_fb_depth = max_footer_fb_depth;
1285        self
1286    }
1287
1288    /// Build [`FileReader`] with given reader.
1289    pub fn build<R: Read + Seek>(self, mut reader: R) -> Result<FileReader<R>, ArrowError> {
1290        // Space for ARROW_MAGIC (6 bytes) and length (4 bytes)
1291        let mut buffer = [0; 10];
1292        reader.seek(SeekFrom::End(-10))?;
1293        reader.read_exact(&mut buffer)?;
1294
1295        let footer_len = read_footer_length(buffer)?;
1296
1297        // read footer
1298        let mut footer_data = vec![0; footer_len];
1299        reader.seek(SeekFrom::End(-10 - footer_len as i64))?;
1300        reader.read_exact(&mut footer_data)?;
1301
1302        let verifier_options = VerifierOptions {
1303            max_tables: self.max_footer_fb_tables,
1304            max_depth: self.max_footer_fb_depth,
1305            ..Default::default()
1306        };
1307        let footer = crate::root_as_footer_with_opts(&verifier_options, &footer_data[..]).map_err(
1308            |err| ArrowError::ParseError(format!("Unable to get root as footer: {err:?}")),
1309        )?;
1310
1311        let blocks = footer.recordBatches().ok_or_else(|| {
1312            ArrowError::ParseError("Unable to get record batches from IPC Footer".to_string())
1313        })?;
1314
1315        let total_blocks = blocks.len();
1316
1317        let ipc_schema = footer.schema().ok_or_else(|| {
1318            ArrowError::ParseError("Unable to get schema from IPC Footer".to_string())
1319        })?;
1320        if !ipc_schema.endianness().equals_to_target_endianness() {
1321            return Err(ArrowError::IpcError(
1322                "the endianness of the source system does not match the endianness of the target system.".to_owned()
1323            ));
1324        }
1325
1326        let schema = Arc::new(crate::convert::try_fb_to_schema(ipc_schema)?);
1327
1328        let projected_schema = match &self.projection {
1329            Some(projection) => Arc::new(schema.project(projection)?),
1330            None => schema.clone(),
1331        };
1332
1333        let mut custom_metadata = HashMap::new();
1334        if let Some(fb_custom_metadata) = footer.custom_metadata() {
1335            for kv in fb_custom_metadata {
1336                let (Some(key), Some(value)) = (kv.key(), kv.value()) else {
1337                    return Err(ArrowError::ParseError(
1338                        "Custom metadata in the IPC footer is missing a key or a value".to_string(),
1339                    ));
1340                };
1341                custom_metadata.insert(key.to_string(), value.to_string());
1342            }
1343        }
1344
1345        let mut decoder = FileDecoder::new(schema, footer.version());
1346        if let Some(projection) = self.projection {
1347            decoder = decoder.with_projection(projection)
1348        }
1349
1350        // Create an array of optional dictionary value arrays, one per field.
1351        if let Some(dictionaries) = footer.dictionaries() {
1352            for block in dictionaries {
1353                let buf = read_block(&mut reader, block)?;
1354                decoder.read_dictionary(block, &buf)?;
1355            }
1356        }
1357
1358        Ok(FileReader {
1359            reader,
1360            blocks: blocks.iter().copied().collect(),
1361            current_block: 0,
1362            total_blocks,
1363            decoder,
1364            schema: projected_schema,
1365            custom_metadata,
1366        })
1367    }
1368}
1369
1370/// Arrow File Reader
1371///
1372/// Reads Arrow [`RecordBatch`]es from bytes in the [IPC File Format],
1373/// providing random access to the record batches.
1374///
1375/// # See Also
1376///
1377/// * [`Self::set_index`] for random access
1378/// * [`StreamReader`] for reading streaming data
1379///
1380/// # Example: Reading from a `File`
1381/// ```
1382/// # use std::io::Cursor;
1383/// use arrow_array::record_batch;
1384/// # use arrow_ipc::reader::FileReader;
1385/// # use arrow_ipc::writer::FileWriter;
1386/// # let batch = record_batch!(("a", Int32, [1, 2, 3])).unwrap();
1387/// # let mut file = vec![]; // mimic a stream for the example
1388/// # {
1389/// #  let mut writer = FileWriter::try_new(&mut file, &batch.schema()).unwrap();
1390/// #  writer.write(&batch).unwrap();
1391/// #  writer.write(&batch).unwrap();
1392/// #  writer.finish().unwrap();
1393/// # }
1394/// # let mut file = Cursor::new(&file);
1395/// let projection = None; // read all columns
1396/// let mut reader = FileReader::try_new(&mut file, projection).unwrap();
1397/// // Position the reader to the second batch
1398/// reader.set_index(1).unwrap();
1399/// // read batches from the reader using the Iterator trait
1400/// let mut num_rows = 0;
1401/// for batch in reader {
1402///    let batch = batch.unwrap();
1403///    num_rows += batch.num_rows();
1404/// }
1405/// assert_eq!(num_rows, 3);
1406/// ```
1407/// # Example: Reading from `mmap`ed file
1408///
1409/// For an example creating Arrays without copying using  memory mapped (`mmap`)
1410/// files see the [`zero_copy_ipc`] example.
1411///
1412/// [IPC File Format]: https://arrow.apache.org/docs/format/Columnar.html#ipc-file-format
1413/// [`zero_copy_ipc`]: https://github.com/apache/arrow-rs/blob/main/arrow/examples/zero_copy_ipc.rs
1414pub struct FileReader<R> {
1415    /// File reader that supports reading and seeking
1416    reader: R,
1417
1418    /// The decoder
1419    decoder: FileDecoder,
1420
1421    /// Schema of the record batches produced by this reader
1422    schema: SchemaRef,
1423
1424    /// The blocks in the file
1425    ///
1426    /// A block indicates the regions in the file to read to get data
1427    blocks: Vec<Block>,
1428
1429    /// A counter to keep track of the current block that should be read
1430    current_block: usize,
1431
1432    /// The total number of blocks, which may contain record batches and other types
1433    total_blocks: usize,
1434
1435    /// User defined metadata
1436    custom_metadata: HashMap<String, String>,
1437}
1438
1439impl<R> fmt::Debug for FileReader<R> {
1440    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> Result<(), fmt::Error> {
1441        f.debug_struct("FileReader<R>")
1442            .field("decoder", &self.decoder)
1443            .field("blocks", &self.blocks)
1444            .field("current_block", &self.current_block)
1445            .field("total_blocks", &self.total_blocks)
1446            .finish_non_exhaustive()
1447    }
1448}
1449
1450impl<R: Read + Seek> FileReader<BufReader<R>> {
1451    /// Try to create a new file reader with the reader wrapped in a BufReader.
1452    ///
1453    /// See [`FileReader::try_new`] for an unbuffered version.
1454    pub fn try_new_buffered(reader: R, projection: Option<Vec<usize>>) -> Result<Self, ArrowError> {
1455        Self::try_new(BufReader::new(reader), projection)
1456    }
1457}
1458
1459impl<R: Read + Seek> FileReader<R> {
1460    /// Try to create a new file reader.
1461    ///
1462    /// There is no internal buffering. If buffered reads are needed you likely want to use
1463    /// [`FileReader::try_new_buffered`] instead.
1464    ///
1465    /// # Errors
1466    ///
1467    /// An [`Err`] may be returned if:
1468    /// - the file does not meet the Arrow Format footer requirements,
1469    /// - file endianness does not match the target endianness, or
1470    /// - the projection contains an index outside the file schema.
1471    pub fn try_new(reader: R, projection: Option<Vec<usize>>) -> Result<Self, ArrowError> {
1472        let builder = FileReaderBuilder {
1473            projection,
1474            ..Default::default()
1475        };
1476        builder.build(reader)
1477    }
1478
1479    /// Return user defined customized metadata
1480    pub fn custom_metadata(&self) -> &HashMap<String, String> {
1481        &self.custom_metadata
1482    }
1483
1484    /// Return the number of batches in the file
1485    pub fn num_batches(&self) -> usize {
1486        self.total_blocks
1487    }
1488
1489    /// Return the schema of the record batches produced by this reader
1490    pub fn schema(&self) -> SchemaRef {
1491        self.schema.clone()
1492    }
1493
1494    /// See to a specific [`RecordBatch`]
1495    ///
1496    /// Sets the current block to the index, allowing random reads
1497    pub fn set_index(&mut self, index: usize) -> Result<(), ArrowError> {
1498        if index >= self.total_blocks {
1499            Err(ArrowError::InvalidArgumentError(format!(
1500                "Cannot set batch to index {} from {} total batches",
1501                index, self.total_blocks
1502            )))
1503        } else {
1504            self.current_block = index;
1505            Ok(())
1506        }
1507    }
1508
1509    fn maybe_next(&mut self) -> Result<Option<RecordBatch>, ArrowError> {
1510        let block = &self.blocks[self.current_block];
1511        self.current_block += 1;
1512
1513        // read length
1514        let buffer = read_block(&mut self.reader, block)?;
1515        self.decoder.read_record_batch(block, &buffer)
1516    }
1517
1518    /// Gets a reference to the underlying reader.
1519    ///
1520    /// It is inadvisable to directly read from the underlying reader.
1521    pub fn get_ref(&self) -> &R {
1522        &self.reader
1523    }
1524
1525    /// Gets a mutable reference to the underlying reader.
1526    ///
1527    /// It is inadvisable to directly read from the underlying reader.
1528    pub fn get_mut(&mut self) -> &mut R {
1529        &mut self.reader
1530    }
1531
1532    /// Specifies if validation should be skipped when reading data (defaults to `false`)
1533    ///
1534    /// # Safety
1535    ///
1536    /// See [`FileDecoder::with_skip_validation`]
1537    pub unsafe fn with_skip_validation(mut self, skip_validation: bool) -> Self {
1538        self.decoder = unsafe { self.decoder.with_skip_validation(skip_validation) };
1539        self
1540    }
1541}
1542
1543impl<R: Read + Seek> Iterator for FileReader<R> {
1544    type Item = Result<RecordBatch, ArrowError>;
1545
1546    fn next(&mut self) -> Option<Self::Item> {
1547        // get current block
1548        if self.current_block < self.total_blocks {
1549            self.maybe_next().transpose()
1550        } else {
1551            None
1552        }
1553    }
1554}
1555
1556impl<R: Read + Seek> RecordBatchReader for FileReader<R> {
1557    fn schema(&self) -> SchemaRef {
1558        self.schema()
1559    }
1560}
1561
1562/// Arrow Stream Reader
1563///
1564/// Reads Arrow [`RecordBatch`]es from bytes in the [IPC Streaming Format].
1565///
1566/// # See Also
1567///
1568/// * [`FileReader`] for random access.
1569///
1570/// # Example
1571/// ```
1572/// # use arrow_array::record_batch;
1573/// # use arrow_ipc::reader::StreamReader;
1574/// # use arrow_ipc::writer::StreamWriter;
1575/// # let batch = record_batch!(("a", Int32, [1, 2, 3])).unwrap();
1576/// # let mut stream = vec![]; // mimic a stream for the example
1577/// # {
1578/// #  let mut writer = StreamWriter::try_new(&mut stream, &batch.schema()).unwrap();
1579/// #  writer.write(&batch).unwrap();
1580/// #  writer.finish().unwrap();
1581/// # }
1582/// # let stream = stream.as_slice();
1583/// let projection = None; // read all columns
1584/// let mut reader = StreamReader::try_new(stream, projection).unwrap();
1585/// // read batches from the reader using the Iterator trait
1586/// let mut num_rows = 0;
1587/// for batch in reader {
1588///    let batch = batch.unwrap();
1589///    num_rows += batch.num_rows();
1590/// }
1591/// assert_eq!(num_rows, 3);
1592/// ```
1593///
1594/// [IPC Streaming Format]: https://arrow.apache.org/docs/format/Columnar.html#ipc-streaming-format
1595pub struct StreamReader<R> {
1596    /// Stream reader
1597    reader: MessageReader<R>,
1598
1599    /// The schema that is read from the stream's first message
1600    schema: SchemaRef,
1601
1602    /// Optional dictionaries for each schema field.
1603    ///
1604    /// Dictionaries may be appended to in the streaming format.
1605    dictionaries_by_id: HashMap<i64, ArrayRef>,
1606
1607    /// An indicator of whether the stream is complete.
1608    ///
1609    /// This value is set to `true` the first time the reader's `next()` returns `None`.
1610    finished: bool,
1611
1612    /// Optional projection: column indices and the resulting projected schema
1613    projection: Option<(Vec<usize>, SchemaRef)>,
1614
1615    /// Should validation be skipped when reading data? Defaults to false.
1616    ///
1617    /// See [`FileDecoder::with_skip_validation`] for details.
1618    skip_validation: UnsafeFlag,
1619}
1620
1621impl<R> fmt::Debug for StreamReader<R> {
1622    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> std::result::Result<(), fmt::Error> {
1623        f.debug_struct("StreamReader<R>")
1624            .field("reader", &"R")
1625            .field("schema", &self.schema)
1626            .field("dictionaries_by_id", &self.dictionaries_by_id)
1627            .field("finished", &self.finished)
1628            .field("projection", &self.projection)
1629            .finish()
1630    }
1631}
1632
1633impl<R: Read> StreamReader<BufReader<R>> {
1634    /// Try to create a new stream reader with the reader wrapped in a BufReader.
1635    ///
1636    /// See [`StreamReader::try_new`] for an unbuffered version.
1637    pub fn try_new_buffered(reader: R, projection: Option<Vec<usize>>) -> Result<Self, ArrowError> {
1638        Self::try_new(BufReader::new(reader), projection)
1639    }
1640}
1641
1642impl<R: Read> StreamReader<R> {
1643    /// Try to create a new stream reader.
1644    ///
1645    /// To check if the reader is done, use [`is_finished(self)`](StreamReader::is_finished).
1646    ///
1647    /// There is no internal buffering. If buffered reads are needed you likely want to use
1648    /// [`StreamReader::try_new_buffered`] instead.
1649    ///
1650    /// # Errors
1651    ///
1652    /// An [`Err`] may be returned if the reader does not encounter a schema
1653    /// as the first message in the stream.
1654    pub fn try_new(
1655        reader: R,
1656        projection: Option<Vec<usize>>,
1657    ) -> Result<StreamReader<R>, ArrowError> {
1658        let mut msg_reader = MessageReader::new(reader);
1659        let message = msg_reader.maybe_next()?;
1660        let Some((message, _)) = message else {
1661            return Err(ArrowError::IpcError(
1662                "Expected schema message, found empty stream.".to_string(),
1663            ));
1664        };
1665
1666        if message.header_type() != Message::MessageHeader::Schema {
1667            return Err(ArrowError::IpcError(format!(
1668                "Expected a schema as the first message in the stream, got: {:?}",
1669                message.header_type()
1670            )));
1671        }
1672
1673        let schema = message.header_as_schema().ok_or_else(|| {
1674            ArrowError::ParseError("Failed to parse schema from message header".to_string())
1675        })?;
1676        let schema = crate::convert::try_fb_to_schema(schema)?;
1677
1678        // Create an array of optional dictionary value arrays, one per field.
1679        let dictionaries_by_id = HashMap::new();
1680
1681        let projection = match projection {
1682            Some(projection_indices) => {
1683                let schema = Arc::new(schema.project(&projection_indices)?);
1684                Some((projection_indices, schema))
1685            }
1686            _ => None,
1687        };
1688
1689        Ok(Self {
1690            reader: msg_reader,
1691            schema: Arc::new(schema),
1692            finished: false,
1693            dictionaries_by_id,
1694            projection,
1695            skip_validation: UnsafeFlag::new(),
1696        })
1697    }
1698
1699    /// Return the schema of the record batches produced by this reader
1700    pub fn schema(&self) -> SchemaRef {
1701        match &self.projection {
1702            Some((_, projected_schema)) => projected_schema.clone(),
1703            None => self.schema.clone(),
1704        }
1705    }
1706
1707    /// Check if the stream is finished
1708    pub fn is_finished(&self) -> bool {
1709        self.finished
1710    }
1711
1712    fn maybe_next(&mut self) -> Result<Option<RecordBatch>, ArrowError> {
1713        if self.finished {
1714            return Ok(None);
1715        }
1716
1717        // Read messages until we get a record batch or end of stream
1718        loop {
1719            let message = self.next_ipc_message()?;
1720            let Some(message) = message else {
1721                // If the message is None, we have reached the end of the stream.
1722                self.finished = true;
1723                return Ok(None);
1724            };
1725
1726            match message {
1727                IpcMessage::Schema(_) => {
1728                    return Err(ArrowError::IpcError(
1729                        "Expected a record batch, but found a schema".to_string(),
1730                    ));
1731                }
1732                IpcMessage::RecordBatch(record_batch) => {
1733                    return Ok(Some(record_batch));
1734                }
1735                IpcMessage::DictionaryBatch { .. } => {}
1736            }
1737        }
1738    }
1739
1740    /// Reads and fully parses the next IPC message from the stream. Whereas
1741    /// [`Self::maybe_next`] is a higher level method focused on reading
1742    /// `RecordBatch`es, this method returns the individual fully parsed IPC
1743    /// messages from the underlying stream.
1744    ///
1745    /// This is useful primarily for testing reader/writer behaviors as it
1746    /// allows a full view into the messages that have been written to a stream.
1747    pub(crate) fn next_ipc_message(&mut self) -> Result<Option<IpcMessage>, ArrowError> {
1748        let message = self.reader.maybe_next()?;
1749        let Some((message, body)) = message else {
1750            // If the message is None, we have reached the end of the stream.
1751            return Ok(None);
1752        };
1753
1754        let ipc_message = match message.header_type() {
1755            Message::MessageHeader::Schema => {
1756                let schema = message.header_as_schema().ok_or_else(|| {
1757                    ArrowError::ParseError("Failed to parse schema from message header".to_string())
1758                })?;
1759                let arrow_schema = crate::convert::try_fb_to_schema(schema)?;
1760                IpcMessage::Schema(arrow_schema)
1761            }
1762            Message::MessageHeader::RecordBatch => {
1763                let batch = message.header_as_record_batch().ok_or_else(|| {
1764                    ArrowError::IpcError("Unable to read IPC message as record batch".to_string())
1765                })?;
1766
1767                let version = message.version();
1768                let schema = self.schema.clone();
1769                let record_batch = RecordBatchDecoder::try_new(
1770                    &body.into(),
1771                    batch,
1772                    schema,
1773                    &self.dictionaries_by_id,
1774                    &version,
1775                )?
1776                .with_projection(self.projection.as_ref().map(|x| x.0.as_ref()))
1777                .with_require_alignment(false)
1778                .with_skip_validation(self.skip_validation.clone())
1779                .read_record_batch()?;
1780                IpcMessage::RecordBatch(record_batch)
1781            }
1782            Message::MessageHeader::DictionaryBatch => {
1783                let dict = message.header_as_dictionary_batch().ok_or_else(|| {
1784                    ArrowError::ParseError(
1785                        "Failed to parse dictionary batch from message header".to_string(),
1786                    )
1787                })?;
1788
1789                let version = message.version();
1790                let dict_values = get_dictionary_values(
1791                    &body.into(),
1792                    dict,
1793                    &self.schema,
1794                    &self.dictionaries_by_id,
1795                    &version,
1796                    false,
1797                    self.skip_validation.clone(),
1798                )?;
1799
1800                update_dictionaries(
1801                    &mut self.dictionaries_by_id,
1802                    dict.isDelta(),
1803                    dict.id(),
1804                    dict_values.clone(),
1805                )?;
1806
1807                IpcMessage::DictionaryBatch {
1808                    id: dict.id(),
1809                    is_delta: (dict.isDelta()),
1810                    values: (dict_values),
1811                }
1812            }
1813            x => {
1814                return Err(ArrowError::ParseError(format!(
1815                    "Unsupported message header type in IPC stream: '{x:?}'"
1816                )));
1817            }
1818        };
1819
1820        Ok(Some(ipc_message))
1821    }
1822
1823    /// Gets a reference to the underlying reader.
1824    ///
1825    /// It is inadvisable to directly read from the underlying reader.
1826    pub fn get_ref(&self) -> &R {
1827        self.reader.inner()
1828    }
1829
1830    /// Gets a mutable reference to the underlying reader.
1831    ///
1832    /// It is inadvisable to directly read from the underlying reader.
1833    pub fn get_mut(&mut self) -> &mut R {
1834        self.reader.inner_mut()
1835    }
1836
1837    /// Specifies if validation should be skipped when reading data (defaults to `false`)
1838    ///
1839    /// # Safety
1840    ///
1841    /// See [`FileDecoder::with_skip_validation`]
1842    pub unsafe fn with_skip_validation(mut self, skip_validation: bool) -> Self {
1843        unsafe { self.skip_validation.set(skip_validation) };
1844        self
1845    }
1846}
1847
1848impl<R: Read> Iterator for StreamReader<R> {
1849    type Item = Result<RecordBatch, ArrowError>;
1850
1851    fn next(&mut self) -> Option<Self::Item> {
1852        self.maybe_next().transpose()
1853    }
1854}
1855
1856impl<R: Read> RecordBatchReader for StreamReader<R> {
1857    fn schema(&self) -> SchemaRef {
1858        self.schema()
1859    }
1860}
1861
1862/// Representation of a fully parsed IpcMessage from the underlying stream.
1863/// Parsing this kind of message is done by higher level constructs such as
1864/// [`StreamReader`], because fully interpreting the messages into a record
1865/// batch or dictionary batch requires access to stream state such as schema
1866/// and the full dictionary cache.
1867#[derive(Debug)]
1868#[expect(dead_code)]
1869pub(crate) enum IpcMessage {
1870    Schema(arrow_schema::Schema),
1871    RecordBatch(RecordBatch),
1872    DictionaryBatch {
1873        id: i64,
1874        is_delta: bool,
1875        values: ArrayRef,
1876    },
1877}
1878
1879/// Upper bound on how much memory a single message length is allowed to reserve before any
1880/// of the bytes it promises have been read.
1881///
1882/// Message lengths come from the stream itself, so they cannot be trusted: a truncated or
1883/// corrupted stream can declare a body of arbitrary size. Reserving that up front turns a
1884/// malformed input into an allocation failure, which aborts the process rather than
1885/// returning an error the caller can handle. Beyond this size the buffer grows as the data
1886/// arrives instead, so an implausible length costs one bounded allocation and then fails as
1887/// a short read.
1888///
1889/// The value trades the size of that bounded allocation against how large a body still gets
1890/// read in a single allocation: bodies up to this size behave exactly as before, larger ones
1891/// start here and the buffer doubles as the data arrives, so a body of `n` bytes costs
1892/// `log2(n / MAX_PREALLOC_BYTES)` reallocations and the buffer never runs more than a factor
1893/// of two ahead of the bytes actually received.
1894const MAX_PREALLOC_BYTES: usize = 64 * 1024 * 1024;
1895
1896/// Reads exactly `len` bytes of message body, without reserving `len` before reading it.
1897fn read_body_bounded<R: Read>(reader: &mut R, len: usize) -> Result<MutableBuffer, ArrowError> {
1898    let mut buf = MutableBuffer::try_from_len_zeroed(len.min(MAX_PREALLOC_BYTES))
1899        .map_err(|e| ArrowError::MemoryError(e.to_string()))?;
1900    let mut filled = 0;
1901    while filled < len {
1902        let target = buf.len();
1903        reader.read_exact(&mut buf.as_slice_mut()[filled..target])?;
1904        filled = target;
1905        if filled < len {
1906            buf.try_resize(len.min(target.saturating_mul(2)), 0)
1907                .map_err(|e| ArrowError::MemoryError(e.to_string()))?;
1908        }
1909    }
1910    Ok(buf)
1911}
1912
1913/// A low-level construct that reads [`Message::Message`]s from a reader while
1914/// re-using a buffer for metadata. This is composed into [`StreamReader`].
1915struct MessageReader<R> {
1916    reader: R,
1917    buf: Vec<u8>,
1918}
1919
1920impl<R: Read> MessageReader<R> {
1921    fn new(reader: R) -> Self {
1922        Self {
1923            reader,
1924            buf: Vec::new(),
1925        }
1926    }
1927
1928    /// Reads the entire next message from the underlying reader which includes
1929    /// the metadata length, the metadata, and the body.
1930    ///
1931    /// # Returns
1932    /// - `Ok(None)` if the the reader signals the end of stream with EOF on
1933    ///   the first read
1934    /// - `Err(_)` if the reader returns an error other than on the first
1935    ///   read, or if the metadata length is invalid
1936    /// - `Ok(Some(_))` with the Message and buffer containing the
1937    ///   body bytes otherwise.
1938    fn maybe_next(&mut self) -> Result<Option<(Message::Message<'_>, MutableBuffer)>, ArrowError> {
1939        let meta_len = self.read_meta_len()?;
1940        let Some(meta_len) = meta_len else {
1941            return Ok(None);
1942        };
1943
1944        // `read_to_end` on a `Take` grows with the bytes that arrive, so an implausible
1945        // `meta_len` costs a short read rather than the allocation it asks for.
1946        self.buf.clear();
1947        let read = (&mut self.reader)
1948            .take(meta_len as u64)
1949            .read_to_end(&mut self.buf)?;
1950        if read != meta_len {
1951            return Err(ArrowError::ParseError(format!(
1952                "Unexpected end of stream: expected {meta_len} metadata bytes, got {read}"
1953            )));
1954        }
1955
1956        let message = crate::root_as_message(self.buf.as_slice()).map_err(|err| {
1957            ArrowError::ParseError(format!("Unable to get root as message: {err:?}"))
1958        })?;
1959
1960        let body_len = usize::try_from(message.bodyLength()).map_err(|_| {
1961            ArrowError::ParseError(format!(
1962                "Invalid IPC message body length: {}",
1963                message.bodyLength()
1964            ))
1965        })?;
1966        let buf = read_body_bounded(&mut self.reader, body_len)?;
1967
1968        Ok(Some((message, buf)))
1969    }
1970
1971    /// Get a mutable reference to the underlying reader.
1972    fn inner_mut(&mut self) -> &mut R {
1973        &mut self.reader
1974    }
1975
1976    /// Get an immutable reference to the underlying reader.
1977    fn inner(&self) -> &R {
1978        &self.reader
1979    }
1980
1981    /// Read the metadata length for the next message from the underlying stream.
1982    ///
1983    /// # Returns
1984    /// - `Ok(None)` if the the reader signals the end of stream with EOF on
1985    ///   the first read
1986    /// - `Err(_)` if the reader returns an error other than on the first
1987    ///   read, or if the metadata length is less than 0.
1988    /// - `Ok(Some(_))` with the length otherwise.
1989    pub fn read_meta_len(&mut self) -> Result<Option<usize>, ArrowError> {
1990        let mut meta_len: [u8; 4] = [0; 4];
1991        match self.reader.read_exact(&mut meta_len) {
1992            Ok(()) => {}
1993            Err(e) => {
1994                return if e.kind() == std::io::ErrorKind::UnexpectedEof {
1995                    // Handle EOF without the "0xFFFFFFFF 0x00000000"
1996                    // valid according to:
1997                    // https://arrow.apache.org/docs/format/Columnar.html#ipc-streaming-format
1998                    Ok(None)
1999                } else {
2000                    Err(ArrowError::from(e))
2001                };
2002            }
2003        }
2004
2005        let meta_len = {
2006            // If a continuation marker is encountered, skip over it and read
2007            // the size from the next four bytes.
2008            if meta_len == CONTINUATION_MARKER {
2009                self.reader.read_exact(&mut meta_len)?;
2010            }
2011
2012            i32::from_le_bytes(meta_len)
2013        };
2014
2015        if meta_len == 0 {
2016            return Ok(None);
2017        }
2018
2019        let meta_len = usize::try_from(meta_len)
2020            .map_err(|_| ArrowError::ParseError(format!("Invalid metadata length: {meta_len}")))?;
2021
2022        Ok(Some(meta_len))
2023    }
2024}
2025
2026#[cfg(test)]
2027mod tests {
2028    use std::io::Cursor;
2029
2030    use crate::convert::try_fb_to_schema;
2031    use crate::writer::{
2032        DictionaryTracker, IpcDataGenerator, IpcWriteOptions, unslice_run_array, write_message,
2033    };
2034
2035    use super::*;
2036
2037    use crate::{root_as_footer, root_as_message, size_prefixed_root_as_message};
2038    use arrow_array::builder::{PrimitiveRunBuilder, UnionBuilder};
2039    use arrow_array::types::*;
2040    use arrow_buffer::{NullBuffer, OffsetBuffer};
2041    use arrow_data::ArrayDataBuilder;
2042
2043    fn create_test_projection_schema() -> Schema {
2044        // define field types
2045        let list_data_type = DataType::List(Arc::new(Field::new_list_field(DataType::Int32, true)));
2046
2047        let fixed_size_list_data_type =
2048            DataType::FixedSizeList(Arc::new(Field::new_list_field(DataType::Int32, false)), 3);
2049
2050        let union_fields = UnionFields::from_fields(vec![
2051            Field::new("a", DataType::Int32, false),
2052            Field::new("b", DataType::Float64, false),
2053        ]);
2054
2055        let union_data_type = DataType::Union(union_fields, UnionMode::Dense);
2056
2057        let struct_fields = Fields::from(vec![
2058            Field::new("id", DataType::Int32, false),
2059            Field::new_list("list", Field::new_list_field(DataType::Int8, true), false),
2060        ]);
2061        let struct_data_type = DataType::Struct(struct_fields);
2062
2063        let run_encoded_data_type = DataType::RunEndEncoded(
2064            Arc::new(Field::new(
2065                Field::REE_RUN_ENDS_FIELD_DEFAULT_NAME,
2066                DataType::Int16,
2067                false,
2068            )),
2069            Arc::new(Field::new(
2070                Field::REE_VALUES_FIELD_DEFAULT_NAME,
2071                DataType::Int32,
2072                true,
2073            )),
2074        );
2075
2076        // define schema
2077        Schema::new(vec![
2078            Field::new("f0", DataType::UInt32, false),
2079            Field::new("f1", DataType::Utf8, false),
2080            Field::new("f2", DataType::Boolean, false),
2081            Field::new("f3", union_data_type, true),
2082            Field::new("f4", DataType::Null, true),
2083            Field::new("f5", DataType::Float64, true),
2084            Field::new("f6", list_data_type, false),
2085            Field::new("f7", DataType::FixedSizeBinary(3), true),
2086            Field::new("f8", fixed_size_list_data_type, false),
2087            Field::new("f9", struct_data_type, false),
2088            Field::new("f10", run_encoded_data_type, false),
2089            Field::new("f11", DataType::Boolean, false),
2090            Field::new_dictionary("f12", DataType::Int8, DataType::Utf8, false),
2091            Field::new("f13", DataType::Utf8, false),
2092        ])
2093    }
2094
2095    fn create_test_projection_batch_data(schema: &Schema) -> RecordBatch {
2096        // set test data for each column
2097        let array0 = UInt32Array::from(vec![1, 2, 3]);
2098        let array1 = StringArray::from(vec!["foo", "bar", "baz"]);
2099        let array2 = BooleanArray::from(vec![true, false, true]);
2100
2101        let mut union_builder = UnionBuilder::new_dense();
2102        union_builder.append::<Int32Type>("a", 1).unwrap();
2103        union_builder.append::<Float64Type>("b", 10.1).unwrap();
2104        union_builder.append_null::<Float64Type>("b").unwrap();
2105        let array3 = union_builder.build().unwrap();
2106
2107        let array4 = NullArray::new(3);
2108        let array5 = Float64Array::from(vec![Some(1.1), None, Some(3.3)]);
2109        let array6_values = vec![
2110            Some(vec![Some(10), Some(10), Some(10)]),
2111            Some(vec![Some(20), Some(20), Some(20)]),
2112            Some(vec![Some(30), Some(30)]),
2113        ];
2114        let array6 = ListArray::from_iter_primitive::<Int32Type, _, _>(array6_values);
2115        let array7_values = vec![vec![11, 12, 13], vec![22, 23, 24], vec![33, 34, 35]];
2116        let array7 = FixedSizeBinaryArray::try_from_iter(array7_values.into_iter()).unwrap();
2117
2118        let array8_values = ArrayData::builder(DataType::Int32)
2119            .len(9)
2120            .add_buffer(Buffer::from_slice_ref([40, 41, 42, 43, 44, 45, 46, 47, 48]))
2121            .build()
2122            .unwrap();
2123        let array8_data = ArrayData::builder(schema.field(8).data_type().clone())
2124            .len(3)
2125            .add_child_data(array8_values)
2126            .build()
2127            .unwrap();
2128        let array8 = FixedSizeListArray::from(array8_data);
2129
2130        let array9_id: ArrayRef = Arc::new(Int32Array::from(vec![1001, 1002, 1003]));
2131        let array9_list: ArrayRef =
2132            Arc::new(ListArray::from_iter_primitive::<Int8Type, _, _>(vec![
2133                Some(vec![Some(-10)]),
2134                Some(vec![Some(-20), Some(-20), Some(-20)]),
2135                Some(vec![Some(-30)]),
2136            ]));
2137        let array9 = ArrayDataBuilder::new(schema.field(9).data_type().clone())
2138            .add_child_data(array9_id.into_data())
2139            .add_child_data(array9_list.into_data())
2140            .len(3)
2141            .build()
2142            .unwrap();
2143        let array9 = StructArray::from(array9);
2144
2145        let array10_input = vec![Some(1_i32), None, None];
2146        let mut array10_builder = PrimitiveRunBuilder::<Int16Type, Int32Type>::new();
2147        array10_builder.extend(array10_input);
2148        let array10 = array10_builder.finish();
2149
2150        let array11 = BooleanArray::from(vec![false, false, true]);
2151
2152        let array12_values = StringArray::from(vec!["x", "yy", "zzz"]);
2153        let array12_keys = Int8Array::from_iter_values([1, 1, 2]);
2154        let array12 = DictionaryArray::new(array12_keys, Arc::new(array12_values));
2155
2156        let array13 = StringArray::from(vec!["a", "bb", "ccc"]);
2157
2158        // create record batch
2159        RecordBatch::try_new(
2160            Arc::new(schema.clone()),
2161            vec![
2162                Arc::new(array0),
2163                Arc::new(array1),
2164                Arc::new(array2),
2165                Arc::new(array3),
2166                Arc::new(array4),
2167                Arc::new(array5),
2168                Arc::new(array6),
2169                Arc::new(array7),
2170                Arc::new(array8),
2171                Arc::new(array9),
2172                Arc::new(array10),
2173                Arc::new(array11),
2174                Arc::new(array12),
2175                Arc::new(array13),
2176            ],
2177        )
2178        .unwrap()
2179    }
2180
2181    #[test]
2182    fn test_negative_meta_len_start_stream() {
2183        let bytes = i32::to_le_bytes(-1);
2184        let mut buf = vec![];
2185        buf.extend(CONTINUATION_MARKER);
2186        buf.extend(bytes);
2187
2188        let reader_err = StreamReader::try_new(Cursor::new(buf), None).err();
2189        assert!(reader_err.is_some());
2190        assert_eq!(
2191            reader_err.unwrap().to_string(),
2192            "Parser error: Invalid metadata length: -1"
2193        );
2194    }
2195
2196    #[test]
2197    fn test_negative_meta_len_mid_stream() {
2198        let schema = Schema::new(vec![Field::new("a", DataType::Int32, false)]);
2199        let mut buf = Vec::new();
2200        {
2201            let mut writer = crate::writer::StreamWriter::try_new(&mut buf, &schema).unwrap();
2202            let batch =
2203                RecordBatch::try_new(Arc::new(schema), vec![Arc::new(Int32Array::from(vec![1]))])
2204                    .unwrap();
2205            writer.write(&batch).unwrap();
2206        }
2207
2208        let bytes = i32::to_le_bytes(-1);
2209        buf.extend(CONTINUATION_MARKER);
2210        buf.extend(bytes);
2211
2212        let mut reader = StreamReader::try_new(Cursor::new(buf), None).unwrap();
2213        // Read the valid value
2214        assert!(reader.maybe_next().is_ok());
2215        // Read the invalid meta len
2216        let batch_err = reader.maybe_next().err();
2217        assert!(batch_err.is_some());
2218        assert_eq!(
2219            batch_err.unwrap().to_string(),
2220            "Parser error: Invalid metadata length: -1"
2221        );
2222    }
2223
2224    #[test]
2225    fn test_invalid_dictionary_batch_without_data() {
2226        use crate::r#gen::Message::*;
2227        use flatbuffers::FlatBufferBuilder;
2228
2229        let schema = Schema::new(vec![Field::new(
2230            "col",
2231            DataType::Dictionary(Box::new(DataType::Int8), Box::new(DataType::Utf8)),
2232            true,
2233        )]);
2234
2235        // DictionaryBatch.data is optional in the flatbuffer grammar and
2236        // required by the format
2237        let mut fbb = FlatBufferBuilder::new();
2238        let batch_offset = DictionaryBatch::create(
2239            &mut fbb,
2240            &DictionaryBatchArgs {
2241                id: 0,
2242                data: None,
2243                isDelta: false,
2244            },
2245        );
2246        fbb.finish_minimal(batch_offset);
2247        let batch_bytes = fbb.finished_data().to_vec();
2248        let batch = flatbuffers::root::<DictionaryBatch>(&batch_bytes).unwrap();
2249
2250        let data_buffer = Buffer::from(vec![0u8; 0]);
2251        let mut dictionaries: HashMap<i64, ArrayRef> = HashMap::new();
2252
2253        let err = read_dictionary(
2254            &data_buffer,
2255            batch,
2256            &schema,
2257            &mut dictionaries,
2258            &MetadataVersion::V5,
2259        )
2260        .unwrap_err();
2261
2262        assert_eq!(
2263            err.to_string(),
2264            "Parser error: Dictionary batch is missing its data"
2265        );
2266    }
2267
2268    #[test]
2269    fn test_missing_buffer_metadata_error() {
2270        use crate::r#gen::Message::*;
2271        use flatbuffers::FlatBufferBuilder;
2272
2273        let schema = Arc::new(Schema::new(vec![Field::new("col", DataType::Int32, true)]));
2274
2275        // create RecordBatch buffer metadata with invalid buffer count
2276        // Int32Array needs 2 buffers (validity + data) but we provide only 1
2277        let mut fbb = FlatBufferBuilder::new();
2278        let nodes = fbb.create_vector(&[FieldNode::new(2, 0)]);
2279        let buffers = fbb.create_vector(&[crate::Buffer::new(0, 8)]);
2280        let batch_offset = RecordBatch::create(
2281            &mut fbb,
2282            &RecordBatchArgs {
2283                length: 2,
2284                nodes: Some(nodes),
2285                buffers: Some(buffers),
2286                compression: None,
2287                variadicBufferCounts: None,
2288            },
2289        );
2290        fbb.finish_minimal(batch_offset);
2291        let batch_bytes = fbb.finished_data().to_vec();
2292        let batch = flatbuffers::root::<RecordBatch>(&batch_bytes).unwrap();
2293
2294        let data_buffer = Buffer::from(vec![0u8; 8]);
2295        let dictionaries: HashMap<i64, ArrayRef> = HashMap::new();
2296        let metadata = MetadataVersion::V5;
2297
2298        let decoder = RecordBatchDecoder::try_new(
2299            &data_buffer,
2300            batch,
2301            schema.clone(),
2302            &dictionaries,
2303            &metadata,
2304        )
2305        .unwrap();
2306
2307        let result = decoder.read_record_batch();
2308
2309        match result {
2310            Err(ArrowError::IpcError(msg)) => {
2311                assert_eq!(msg, "Buffer count mismatched with metadata");
2312            }
2313            other => panic!("unexpected error: {other:?}"),
2314        }
2315    }
2316
2317    /// A `Utf8View` batch whose variadic buffer count is `count`, with two buffers
2318    /// in the message. Returns the error from reading it with the given projection.
2319    fn read_batch_with_variadic_count(count: i64, projection: Option<&[usize]>) -> ArrowError {
2320        use crate::r#gen::Message::*;
2321        use flatbuffers::FlatBufferBuilder;
2322
2323        let schema = Arc::new(Schema::new(vec![Field::new(
2324            "col",
2325            DataType::Utf8View,
2326            true,
2327        )]));
2328
2329        let mut fbb = FlatBufferBuilder::new();
2330        let nodes = fbb.create_vector(&[FieldNode::new(1, 0)]);
2331        let buffers = fbb.create_vector(&[crate::Buffer::new(0, 8), crate::Buffer::new(8, 8)]);
2332        let variadic_buffer_counts = fbb.create_vector(&[count]);
2333        let batch_offset = RecordBatch::create(
2334            &mut fbb,
2335            &RecordBatchArgs {
2336                length: 1,
2337                nodes: Some(nodes),
2338                buffers: Some(buffers),
2339                compression: None,
2340                variadicBufferCounts: Some(variadic_buffer_counts),
2341            },
2342        );
2343        fbb.finish_minimal(batch_offset);
2344        let batch_bytes = fbb.finished_data().to_vec();
2345        let batch = flatbuffers::root::<RecordBatch>(&batch_bytes).unwrap();
2346
2347        let data_buffer = Buffer::from(vec![0u8; 16]);
2348        let dictionaries: HashMap<i64, ArrayRef> = HashMap::new();
2349
2350        RecordBatchDecoder::try_new(
2351            &data_buffer,
2352            batch,
2353            schema,
2354            &dictionaries,
2355            &MetadataVersion::V5,
2356        )
2357        .unwrap()
2358        .with_projection(projection)
2359        .read_record_batch()
2360        .expect_err("should get error")
2361    }
2362
2363    /// A variadic count the message cannot honour used to panic while slicing the
2364    /// buffers it did not read, both when reading the column and when skipping it.
2365    #[test]
2366    fn test_invalid_variadic_buffer_count_error() {
2367        // -2 leaves no buffers at all, -1 leaves too few, and 1 asks for more than the
2368        // message has. The projection selects nothing, so the column is skipped instead.
2369        for count in [-2, -1, 1, i64::MAX] {
2370            for projection in [None, Some([].as_slice())] {
2371                let err = read_batch_with_variadic_count(count, projection);
2372                assert_eq!(
2373                    err.to_string(),
2374                    format!(
2375                        "Ipc error: Invalid variadic count {count} for Utf8View column, \
2376                         with 2 buffer(s) left in the message"
2377                    ),
2378                    "count {count}, projection {projection:?}"
2379                );
2380            }
2381        }
2382    }
2383
2384    /// The valid count for a message with two buffers is zero.
2385    #[test]
2386    fn test_valid_variadic_buffer_count_is_accepted() {
2387        let err = read_batch_with_variadic_count(0, None);
2388        assert!(!err.to_string().contains("Invalid variadic count"), "{err}");
2389    }
2390
2391    #[test]
2392    fn test_missing_footer_schema_error() {
2393        use crate::r#gen::File::{Footer, FooterArgs};
2394        use flatbuffers::FlatBufferBuilder;
2395
2396        // a footer that verifies but has no schema table. record batches present
2397        // (so the earlier ok_or_else passes) but schema absent, which used to panic.
2398        let mut fbb = FlatBufferBuilder::new();
2399        let record_batches = fbb.create_vector::<Block>(&[]);
2400        let footer = Footer::create(
2401            &mut fbb,
2402            &FooterArgs {
2403                version: MetadataVersion::V5,
2404                schema: None,
2405                dictionaries: None,
2406                recordBatches: Some(record_batches),
2407                custom_metadata: None,
2408            },
2409        );
2410        fbb.finish(footer, None);
2411        let footer_data = fbb.finished_data();
2412
2413        // assemble a minimal IPC file: magic header, footer, footer length, magic trailer
2414        let mut buf = Vec::new();
2415        buf.extend_from_slice(&crate::ARROW_MAGIC);
2416        buf.extend_from_slice(footer_data);
2417        buf.extend_from_slice(&(footer_data.len() as i32).to_le_bytes());
2418        buf.extend_from_slice(&crate::ARROW_MAGIC);
2419
2420        let err = FileReader::try_new(Cursor::new(buf), None)
2421            .expect_err("expected an error, not a panic");
2422        assert!(
2423            matches!(err, ArrowError::ParseError(_)),
2424            "expected ParseError, got {err:?}"
2425        );
2426    }
2427
2428    /// Test that the reader can read legacy files where empty list arrays were written with a 0-byte offsets buffer.
2429    #[test]
2430    fn test_read_legacy_empty_list_without_offsets_buffer() {
2431        use crate::r#gen::Message::*;
2432        use flatbuffers::FlatBufferBuilder;
2433
2434        let schema = Arc::new(Schema::new(vec![Field::new_list(
2435            "items",
2436            Field::new_list_field(DataType::Int32, true),
2437            true,
2438        )]));
2439
2440        // Legacy arrow-rs versions wrote empty offsets buffers for empty list arrays.
2441        // Keep reader compatibility with such files by accepting a 0-byte offsets buffer.
2442        let mut fbb = FlatBufferBuilder::new();
2443        let nodes = fbb.create_vector(&[
2444            FieldNode::new(0, 0), // list node
2445            FieldNode::new(0, 0), // child int32 node
2446        ]);
2447        let buffers = fbb.create_vector(&[
2448            crate::Buffer::new(0, 0), // list validity
2449            crate::Buffer::new(0, 0), // list offsets (legacy empty buffer)
2450            crate::Buffer::new(0, 0), // child validity
2451            crate::Buffer::new(0, 0), // child values
2452        ]);
2453        let batch_offset = RecordBatch::create(
2454            &mut fbb,
2455            &RecordBatchArgs {
2456                length: 0,
2457                nodes: Some(nodes),
2458                buffers: Some(buffers),
2459                compression: None,
2460                variadicBufferCounts: None,
2461            },
2462        );
2463        fbb.finish_minimal(batch_offset);
2464        let batch_bytes = fbb.finished_data().to_vec();
2465        let batch = flatbuffers::root::<RecordBatch>(&batch_bytes).unwrap();
2466
2467        let body = Buffer::from(Vec::<u8>::new());
2468        let dictionaries: HashMap<i64, ArrayRef> = HashMap::new();
2469        let metadata = MetadataVersion::V5;
2470
2471        let decoder =
2472            RecordBatchDecoder::try_new(&body, batch, schema.clone(), &dictionaries, &metadata)
2473                .unwrap();
2474
2475        let read_batch = decoder.read_record_batch().unwrap();
2476        assert_eq!(read_batch.num_rows(), 0);
2477
2478        let list = read_batch
2479            .column(0)
2480            .as_any()
2481            .downcast_ref::<ListArray>()
2482            .unwrap();
2483        assert_eq!(list.len(), 0);
2484        assert_eq!(list.values().len(), 0);
2485    }
2486
2487    /// Test that the reader can read legacy files where empty Utf8/Binary arrays were written with a 0-byte offsets buffer.
2488    #[test]
2489    fn test_read_legacy_empty_utf8_and_binary_without_offsets_buffer() {
2490        use crate::r#gen::Message::*;
2491        use flatbuffers::FlatBufferBuilder;
2492
2493        let schema = Arc::new(Schema::new(vec![
2494            Field::new("name", DataType::Utf8, true),
2495            Field::new("payload", DataType::Binary, true),
2496        ]));
2497
2498        // Legacy arrow-rs versions wrote empty offsets buffers for empty Utf8/Binary arrays.
2499        // Keep reader compatibility with such files by accepting 0-byte offsets buffers.
2500        let mut fbb = FlatBufferBuilder::new();
2501        let nodes = fbb.create_vector(&[
2502            FieldNode::new(0, 0), // utf8 node
2503            FieldNode::new(0, 0), // binary node
2504        ]);
2505        let buffers = fbb.create_vector(&[
2506            crate::Buffer::new(0, 0), // utf8 validity
2507            crate::Buffer::new(0, 0), // utf8 offsets (legacy empty buffer)
2508            crate::Buffer::new(0, 0), // utf8 values
2509            crate::Buffer::new(0, 0), // binary validity
2510            crate::Buffer::new(0, 0), // binary offsets (legacy empty buffer)
2511            crate::Buffer::new(0, 0), // binary values
2512        ]);
2513        let batch_offset = RecordBatch::create(
2514            &mut fbb,
2515            &RecordBatchArgs {
2516                length: 0,
2517                nodes: Some(nodes),
2518                buffers: Some(buffers),
2519                compression: None,
2520                variadicBufferCounts: None,
2521            },
2522        );
2523        fbb.finish_minimal(batch_offset);
2524        let batch_bytes = fbb.finished_data().to_vec();
2525        let batch = flatbuffers::root::<RecordBatch>(&batch_bytes).unwrap();
2526
2527        let body = Buffer::from(Vec::<u8>::new());
2528        let dictionaries: HashMap<i64, ArrayRef> = HashMap::new();
2529        let metadata = MetadataVersion::V5;
2530
2531        let decoder =
2532            RecordBatchDecoder::try_new(&body, batch, schema.clone(), &dictionaries, &metadata)
2533                .unwrap();
2534
2535        let read_batch = decoder.read_record_batch().unwrap();
2536        assert_eq!(read_batch.num_rows(), 0);
2537
2538        let utf8 = read_batch
2539            .column(0)
2540            .as_any()
2541            .downcast_ref::<StringArray>()
2542            .unwrap();
2543        assert_eq!(utf8.len(), 0);
2544        assert_eq!(utf8.value_offsets(), [0]);
2545
2546        let binary = read_batch
2547            .column(1)
2548            .as_any()
2549            .downcast_ref::<BinaryArray>()
2550            .unwrap();
2551        assert_eq!(binary.len(), 0);
2552        assert_eq!(binary.value_offsets(), [0]);
2553    }
2554
2555    #[test]
2556    fn test_projection_array_values() {
2557        // define schema
2558        let schema = create_test_projection_schema();
2559
2560        // create record batch with test data
2561        let batch = create_test_projection_batch_data(&schema);
2562
2563        // write record batch in IPC format
2564        let mut buf = Vec::new();
2565        {
2566            let mut writer = crate::writer::FileWriter::try_new(&mut buf, &schema).unwrap();
2567            writer.write(&batch).unwrap();
2568            writer.finish().unwrap();
2569        }
2570
2571        // read record batch with projection
2572        for index in 0..12 {
2573            let projection = vec![index];
2574            let reader = FileReader::try_new(std::io::Cursor::new(buf.clone()), Some(projection));
2575            let read_batch = reader.unwrap().next().unwrap().unwrap();
2576            let projected_column = read_batch.column(0);
2577            let expected_column = batch.column(index);
2578
2579            // check the projected column equals the expected column
2580            assert_eq!(projected_column.as_ref(), expected_column.as_ref());
2581        }
2582
2583        {
2584            // read record batch with reversed projection
2585            let reader =
2586                FileReader::try_new(std::io::Cursor::new(buf.clone()), Some(vec![3, 2, 1]));
2587            let read_batch = reader.unwrap().next().unwrap().unwrap();
2588            let expected_batch = batch.project(&[3, 2, 1]).unwrap();
2589            assert_eq!(read_batch, expected_batch);
2590        }
2591    }
2592
2593    #[test]
2594    fn test_file_reader_projected_schema_matches_batch_schema() {
2595        let schema = create_test_projection_schema();
2596        let batch = create_test_projection_batch_data(&schema);
2597
2598        let mut buf = Vec::new();
2599        {
2600            let mut writer = crate::writer::FileWriter::try_new(&mut buf, &schema).unwrap();
2601            writer.write(&batch).unwrap();
2602            writer.finish().unwrap();
2603        }
2604
2605        let projection = vec![3, 2, 1];
2606        let mut reader = FileReader::try_new(Cursor::new(buf), Some(projection)).unwrap();
2607        let reader_schema = RecordBatchReader::schema(&reader);
2608        let read_batch = reader.next().unwrap().unwrap();
2609
2610        assert_eq!(reader_schema, read_batch.schema());
2611    }
2612
2613    #[test]
2614    fn test_stream_reader_projected_schema_matches_batch_schema() {
2615        let schema = create_test_projection_schema();
2616        let batch = create_test_projection_batch_data(&schema);
2617
2618        let mut buf = Vec::new();
2619        {
2620            let mut writer = crate::writer::StreamWriter::try_new(&mut buf, &schema).unwrap();
2621            writer.write(&batch).unwrap();
2622            writer.finish().unwrap();
2623        }
2624
2625        let projection = vec![3, 2, 1];
2626        let mut reader = StreamReader::try_new(Cursor::new(buf), Some(projection)).unwrap();
2627        let reader_schema = RecordBatchReader::schema(&reader);
2628        let read_batch = reader.next().unwrap().unwrap();
2629
2630        assert_eq!(reader_schema, read_batch.schema());
2631    }
2632
2633    #[test]
2634    fn test_file_reader_rejects_invalid_projection() {
2635        let schema = create_test_projection_schema();
2636        let batch = create_test_projection_batch_data(&schema);
2637
2638        let mut buf = Vec::new();
2639        {
2640            let mut writer = crate::writer::FileWriter::try_new(&mut buf, &schema).unwrap();
2641            writer.write(&batch).unwrap();
2642            writer.finish().unwrap();
2643        }
2644
2645        let result = FileReader::try_new(Cursor::new(buf), Some(vec![schema.fields().len()]));
2646
2647        assert!(matches!(result, Err(ArrowError::SchemaError(_))));
2648    }
2649
2650    #[test]
2651    fn test_projection_duplicate_indices() {
2652        let schema = create_test_projection_schema();
2653        let batch = create_test_projection_batch_data(&schema);
2654
2655        // Write the batch to IPC
2656        let mut buf = Vec::new();
2657        {
2658            let mut writer = crate::writer::FileWriter::try_new(&mut buf, &schema).unwrap();
2659            writer.write(&batch).unwrap();
2660            writer.finish().unwrap();
2661        }
2662
2663        // Verify duplicate([1, 1]) and reordered([2, 0, 2]) projection indices
2664        for projection in [vec![1, 1], vec![2, 0, 2]] {
2665            let reader =
2666                FileReader::try_new(std::io::Cursor::new(buf.clone()), Some(projection.clone()));
2667            let read_batch = reader.unwrap().next().unwrap().unwrap();
2668
2669            let expected_batch = batch.project(&projection).unwrap();
2670            assert_eq!(read_batch, expected_batch);
2671        }
2672    }
2673
2674    #[test]
2675    fn test_arrow_single_float_row() {
2676        let schema = Schema::new(vec![
2677            Field::new("a", DataType::Float32, false),
2678            Field::new("b", DataType::Float32, false),
2679            Field::new("c", DataType::Int32, false),
2680            Field::new("d", DataType::Int32, false),
2681        ]);
2682        let arrays = vec![
2683            Arc::new(Float32Array::from(vec![1.23])) as ArrayRef,
2684            Arc::new(Float32Array::from(vec![-6.50])) as ArrayRef,
2685            Arc::new(Int32Array::from(vec![2])) as ArrayRef,
2686            Arc::new(Int32Array::from(vec![1])) as ArrayRef,
2687        ];
2688        let batch = RecordBatch::try_new(Arc::new(schema.clone()), arrays).unwrap();
2689        // create stream writer
2690        let mut file = tempfile::tempfile().unwrap();
2691        let mut stream_writer = crate::writer::StreamWriter::try_new(&mut file, &schema).unwrap();
2692        stream_writer.write(&batch).unwrap();
2693        stream_writer.finish().unwrap();
2694
2695        drop(stream_writer);
2696
2697        file.rewind().unwrap();
2698
2699        // read stream back
2700        let reader = StreamReader::try_new(&mut file, None).unwrap();
2701
2702        reader.for_each(|batch| {
2703            let batch = batch.unwrap();
2704            assert!(
2705                batch
2706                    .column(0)
2707                    .as_any()
2708                    .downcast_ref::<Float32Array>()
2709                    .unwrap()
2710                    .value(0)
2711                    != 0.0
2712            );
2713            assert!(
2714                batch
2715                    .column(1)
2716                    .as_any()
2717                    .downcast_ref::<Float32Array>()
2718                    .unwrap()
2719                    .value(0)
2720                    != 0.0
2721            );
2722        });
2723
2724        file.rewind().unwrap();
2725
2726        // Read with projection
2727        let reader = StreamReader::try_new(file, Some(vec![0, 3])).unwrap();
2728
2729        reader.for_each(|batch| {
2730            let batch = batch.unwrap();
2731            assert_eq!(batch.schema().fields().len(), 2);
2732            assert_eq!(batch.schema().fields()[0].data_type(), &DataType::Float32);
2733            assert_eq!(batch.schema().fields()[1].data_type(), &DataType::Int32);
2734        });
2735    }
2736
2737    /// Write the record batch to an in-memory buffer in IPC File format
2738    fn write_ipc(rb: &RecordBatch) -> Vec<u8> {
2739        let mut buf = Vec::new();
2740        let mut writer = crate::writer::FileWriter::try_new(&mut buf, rb.schema_ref()).unwrap();
2741        writer.write(rb).unwrap();
2742        writer.finish().unwrap();
2743        buf
2744    }
2745
2746    /// Return the first record batch read from the IPC File buffer
2747    fn read_ipc(buf: &[u8]) -> Result<RecordBatch, ArrowError> {
2748        let mut reader = FileReader::try_new(std::io::Cursor::new(buf), None)?;
2749        reader.next().unwrap()
2750    }
2751
2752    /// Return the first record batch read from the IPC File buffer, disabling
2753    /// validation
2754    fn read_ipc_skip_validation(buf: &[u8]) -> Result<RecordBatch, ArrowError> {
2755        let mut reader = unsafe {
2756            FileReader::try_new(std::io::Cursor::new(buf), None)?.with_skip_validation(true)
2757        };
2758        reader.next().unwrap()
2759    }
2760
2761    fn roundtrip_ipc(rb: &RecordBatch) -> RecordBatch {
2762        let buf = write_ipc(rb);
2763        read_ipc(&buf).unwrap()
2764    }
2765
2766    /// Return the first record batch read from the IPC File buffer
2767    /// using the FileDecoder API
2768    fn read_ipc_with_decoder(buf: Vec<u8>) -> Result<RecordBatch, ArrowError> {
2769        read_ipc_with_decoder_inner(buf, false)
2770    }
2771
2772    /// Return the first record batch read from the IPC File buffer
2773    /// using the FileDecoder API, disabling validation
2774    fn read_ipc_with_decoder_skip_validation(buf: Vec<u8>) -> Result<RecordBatch, ArrowError> {
2775        read_ipc_with_decoder_inner(buf, true)
2776    }
2777
2778    fn read_ipc_with_decoder_inner(
2779        buf: Vec<u8>,
2780        skip_validation: bool,
2781    ) -> Result<RecordBatch, ArrowError> {
2782        let buffer = Buffer::from_vec(buf);
2783        let trailer_start = buffer.len() - 10;
2784        let footer_len = read_footer_length(buffer[trailer_start..].try_into().unwrap())?;
2785        let footer = root_as_footer(&buffer[trailer_start - footer_len..trailer_start])
2786            .map_err(|e| ArrowError::InvalidArgumentError(format!("Invalid footer: {e}")))?;
2787
2788        let schema = try_fb_to_schema(footer.schema().unwrap()).unwrap();
2789
2790        let mut decoder = unsafe {
2791            FileDecoder::new(Arc::new(schema), footer.version())
2792                .with_skip_validation(skip_validation)
2793        };
2794        // Read dictionaries
2795        for block in footer.dictionaries().iter().flatten() {
2796            let block_len = block.bodyLength() as usize + block.metaDataLength() as usize;
2797            let data = buffer.slice_with_length(block.offset() as _, block_len);
2798            decoder.read_dictionary(block, &data)?
2799        }
2800
2801        // Read record batch
2802        let batches = footer.recordBatches().unwrap();
2803        assert_eq!(batches.len(), 1); // Only wrote a single batch
2804
2805        let block = batches.get(0);
2806        let block_len = block.bodyLength() as usize + block.metaDataLength() as usize;
2807        let data = buffer.slice_with_length(block.offset() as _, block_len);
2808        Ok(decoder.read_record_batch(block, &data)?.unwrap())
2809    }
2810
2811    /// Write the record batch to an in-memory buffer in IPC Stream format
2812    fn write_stream(rb: &RecordBatch) -> Vec<u8> {
2813        let mut buf = Vec::new();
2814        let mut writer = crate::writer::StreamWriter::try_new(&mut buf, rb.schema_ref()).unwrap();
2815        writer.write(rb).unwrap();
2816        writer.finish().unwrap();
2817        buf
2818    }
2819
2820    /// Return the first record batch read from the IPC Stream buffer
2821    fn read_stream(buf: &[u8]) -> Result<RecordBatch, ArrowError> {
2822        let mut reader = StreamReader::try_new(std::io::Cursor::new(buf), None)?;
2823        reader.next().unwrap()
2824    }
2825
2826    /// Return the first record batch read from the IPC Stream buffer,
2827    /// disabling validation
2828    fn read_stream_skip_validation(buf: &[u8]) -> Result<RecordBatch, ArrowError> {
2829        let mut reader = unsafe {
2830            StreamReader::try_new(std::io::Cursor::new(buf), None)?.with_skip_validation(true)
2831        };
2832        reader.next().unwrap()
2833    }
2834
2835    fn roundtrip_ipc_stream(rb: &RecordBatch) -> RecordBatch {
2836        let buf = write_stream(rb);
2837        read_stream(&buf).unwrap()
2838    }
2839
2840    #[test]
2841    fn test_roundtrip_with_custom_metadata() {
2842        let schema = Schema::new(vec![Field::new("dummy", DataType::Float64, false)]);
2843        let mut buf = Vec::new();
2844        let mut writer = crate::writer::FileWriter::try_new(&mut buf, &schema).unwrap();
2845        let mut test_metadata = HashMap::new();
2846        test_metadata.insert("abc".to_string(), "abc".to_string());
2847        test_metadata.insert("def".to_string(), "def".to_string());
2848        for (k, v) in &test_metadata {
2849            writer.write_metadata(k, v);
2850        }
2851        writer.finish().unwrap();
2852        drop(writer);
2853
2854        let reader = crate::reader::FileReader::try_new(std::io::Cursor::new(buf), None).unwrap();
2855        assert_eq!(reader.custom_metadata(), &test_metadata);
2856    }
2857
2858    #[test]
2859    fn test_roundtrip_nested_dict() {
2860        let inner: DictionaryArray<Int32Type> = vec!["a", "b", "a"].into_iter().collect();
2861
2862        let array = Arc::new(inner) as ArrayRef;
2863
2864        let dctfield = Arc::new(Field::new("dict", array.data_type().clone(), false));
2865
2866        let s = StructArray::from(vec![(dctfield, array)]);
2867        let struct_array = Arc::new(s) as ArrayRef;
2868
2869        let schema = Arc::new(Schema::new(vec![Field::new(
2870            "struct",
2871            struct_array.data_type().clone(),
2872            false,
2873        )]));
2874
2875        let batch = RecordBatch::try_new(schema, vec![struct_array]).unwrap();
2876
2877        assert_eq!(batch, roundtrip_ipc(&batch));
2878    }
2879
2880    #[test]
2881    fn test_roundtrip_nested_dict_no_preserve_dict_id() {
2882        let inner: DictionaryArray<Int32Type> = vec!["a", "b", "a"].into_iter().collect();
2883
2884        let array = Arc::new(inner) as ArrayRef;
2885
2886        let dctfield = Arc::new(Field::new("dict", array.data_type().clone(), false));
2887
2888        let s = StructArray::from(vec![(dctfield, array)]);
2889        let struct_array = Arc::new(s) as ArrayRef;
2890
2891        let schema = Arc::new(Schema::new(vec![Field::new(
2892            "struct",
2893            struct_array.data_type().clone(),
2894            false,
2895        )]));
2896
2897        let batch = RecordBatch::try_new(schema, vec![struct_array]).unwrap();
2898
2899        let mut buf = Vec::new();
2900        let mut writer = crate::writer::FileWriter::try_new_with_options(
2901            &mut buf,
2902            batch.schema_ref(),
2903            IpcWriteOptions::default(),
2904        )
2905        .unwrap();
2906        writer.write(&batch).unwrap();
2907        writer.finish().unwrap();
2908        drop(writer);
2909
2910        let mut reader = FileReader::try_new(std::io::Cursor::new(buf), None).unwrap();
2911
2912        assert_eq!(batch, reader.next().unwrap().unwrap());
2913    }
2914
2915    fn check_union_with_builder(mut builder: UnionBuilder) {
2916        builder.append::<Int32Type>("a", 1).unwrap();
2917        builder.append_null::<Int32Type>("a").unwrap();
2918        builder.append::<Float64Type>("c", 3.0).unwrap();
2919        builder.append::<Int32Type>("a", 4).unwrap();
2920        builder.append::<Int64Type>("d", 11).unwrap();
2921        let union = builder.build().unwrap();
2922
2923        let schema = Arc::new(Schema::new(vec![Field::new(
2924            "union",
2925            union.data_type().clone(),
2926            false,
2927        )]));
2928
2929        let union_array = Arc::new(union) as ArrayRef;
2930
2931        let rb = RecordBatch::try_new(schema, vec![union_array]).unwrap();
2932        let rb2 = roundtrip_ipc(&rb);
2933        // TODO: equality not yet implemented for union, so we check that the length of the array is
2934        // the same and that all of the buffers are the same instead.
2935        assert_eq!(rb.schema(), rb2.schema());
2936        assert_eq!(rb.num_columns(), rb2.num_columns());
2937        assert_eq!(rb.num_rows(), rb2.num_rows());
2938        let union1 = rb.column(0);
2939        let union2 = rb2.column(0);
2940
2941        assert_eq!(union1, union2);
2942    }
2943
2944    #[test]
2945    fn test_roundtrip_dense_union() {
2946        check_union_with_builder(UnionBuilder::new_dense());
2947    }
2948
2949    #[test]
2950    fn test_roundtrip_sparse_union() {
2951        check_union_with_builder(UnionBuilder::new_sparse());
2952    }
2953
2954    #[test]
2955    fn test_roundtrip_struct_empty_fields() {
2956        let nulls = NullBuffer::from(&[true, true, false]);
2957        let rb = RecordBatch::try_from_iter([(
2958            "",
2959            Arc::new(StructArray::new_empty_fields(nulls.len(), Some(nulls))) as _,
2960        )])
2961        .unwrap();
2962        let rb2 = roundtrip_ipc(&rb);
2963        assert_eq!(rb, rb2);
2964    }
2965
2966    #[test]
2967    fn test_roundtrip_stream_run_array_sliced() {
2968        let run_array_1: Int32RunArray = vec!["a", "a", "a", "b", "b", "c", "c", "c"]
2969            .into_iter()
2970            .collect();
2971        let run_array_1_sliced = run_array_1.slice(2, 5);
2972
2973        let run_array_2_inupt = vec![Some(1_i32), None, None, Some(2), Some(2)];
2974        let mut run_array_2_builder = PrimitiveRunBuilder::<Int16Type, Int32Type>::new();
2975        run_array_2_builder.extend(run_array_2_inupt);
2976        let run_array_2 = run_array_2_builder.finish();
2977
2978        let schema = Arc::new(Schema::new(vec![
2979            Field::new(
2980                "run_array_1_sliced",
2981                run_array_1_sliced.data_type().clone(),
2982                false,
2983            ),
2984            Field::new("run_array_2", run_array_2.data_type().clone(), false),
2985        ]));
2986        let input_batch = RecordBatch::try_new(
2987            schema,
2988            vec![Arc::new(run_array_1_sliced.clone()), Arc::new(run_array_2)],
2989        )
2990        .unwrap();
2991        let output_batch = roundtrip_ipc_stream(&input_batch);
2992
2993        // As partial comparison not yet supported for run arrays, the sliced run array
2994        // has to be unsliced before comparing with the output. the second run array
2995        // can be compared as such.
2996        assert_eq!(input_batch.column(1), output_batch.column(1));
2997
2998        let run_array_1_unsliced = unslice_run_array(run_array_1_sliced.into_data()).unwrap();
2999        assert_eq!(run_array_1_unsliced, output_batch.column(0).into_data());
3000    }
3001
3002    #[test]
3003    fn test_roundtrip_stream_nested_dict() {
3004        let xs = vec!["AA", "BB", "AA", "CC", "BB"];
3005        let dict = Arc::new(
3006            xs.clone()
3007                .into_iter()
3008                .collect::<DictionaryArray<Int8Type>>(),
3009        );
3010        let string_array: ArrayRef = Arc::new(StringArray::from(xs.clone()));
3011        let struct_array = StructArray::from(vec![
3012            (
3013                Arc::new(Field::new("f2.1", DataType::Utf8, false)),
3014                string_array,
3015            ),
3016            (
3017                Arc::new(Field::new("f2.2_struct", dict.data_type().clone(), false)),
3018                dict.clone() as ArrayRef,
3019            ),
3020        ]);
3021        let schema = Arc::new(Schema::new(vec![
3022            Field::new("f1_string", DataType::Utf8, false),
3023            Field::new("f2_struct", struct_array.data_type().clone(), false),
3024        ]));
3025        let input_batch = RecordBatch::try_new(
3026            schema,
3027            vec![
3028                Arc::new(StringArray::from(xs.clone())),
3029                Arc::new(struct_array),
3030            ],
3031        )
3032        .unwrap();
3033        let output_batch = roundtrip_ipc_stream(&input_batch);
3034        assert_eq!(input_batch, output_batch);
3035    }
3036
3037    #[test]
3038    fn test_ipc_writers_reject_dictionary_of_dictionary_schema() {
3039        let values = Arc::new(StringArray::from(vec![Some("a"), Some("b")])) as ArrayRef;
3040        let inner = Arc::new(DictionaryArray::new(
3041            UInt32Array::from_iter_values([0, 1]),
3042            values,
3043        )) as ArrayRef;
3044        let outer = Arc::new(DictionaryArray::new(
3045            UInt32Array::from_iter_values([0, 1, 0]),
3046            inner,
3047        )) as ArrayRef;
3048
3049        let schema = Arc::new(Schema::new(vec![Field::new(
3050            "f1",
3051            outer.data_type().clone(),
3052            false,
3053        )]));
3054        let batch = RecordBatch::try_new(schema, vec![outer]).unwrap();
3055
3056        let mut stream = Vec::new();
3057        let Err(err) = crate::writer::StreamWriter::try_new(&mut stream, batch.schema_ref()) else {
3058            panic!("IPC stream writer should reject dictionary-of-dictionary schemas");
3059        };
3060        assert!(stream.is_empty());
3061
3062        assert!(
3063            err.to_string().contains("dictionary-of-dictionary values"),
3064            "unexpected error: {err}"
3065        );
3066
3067        let mut file = Vec::new();
3068        let Err(err) = crate::writer::FileWriter::try_new(&mut file, batch.schema_ref()) else {
3069            panic!("IPC file writer should reject dictionary-of-dictionary schemas");
3070        };
3071        assert!(file.is_empty());
3072
3073        assert!(
3074            err.to_string().contains("dictionary-of-dictionary values"),
3075            "unexpected error: {err}"
3076        );
3077    }
3078
3079    #[test]
3080    fn test_roundtrip_stream_nested_dict_of_map_of_dict() {
3081        let values = StringArray::from(vec![Some("a"), None, Some("b"), Some("c")]);
3082        let values = Arc::new(values) as ArrayRef;
3083        let value_dict_keys = Int8Array::from_iter_values([0, 1, 1, 2, 3, 1]);
3084        let value_dict_array = DictionaryArray::new(value_dict_keys, values.clone());
3085
3086        let key_dict_keys = Int8Array::from_iter_values([0, 0, 2, 2, 2, 3]);
3087        let key_dict_array = DictionaryArray::new(key_dict_keys, values);
3088
3089        #[expect(deprecated)]
3090        let keys_field = Arc::new(Field::new_dict(
3091            Field::MAP_KEY_FIELD_DEFAULT_NAME,
3092            DataType::Dictionary(Box::new(DataType::Int8), Box::new(DataType::Utf8)),
3093            false,
3094            1,
3095            false,
3096        ));
3097        #[expect(deprecated)]
3098        let values_field = Arc::new(Field::new_dict(
3099            Field::MAP_VALUE_FIELD_DEFAULT_NAME,
3100            DataType::Dictionary(Box::new(DataType::Int8), Box::new(DataType::Utf8)),
3101            true,
3102            2,
3103            false,
3104        ));
3105        let entry_struct = StructArray::from(vec![
3106            (keys_field, make_array(key_dict_array.into_data())),
3107            (values_field, make_array(value_dict_array.into_data())),
3108        ]);
3109        let map_data_type = DataType::Map(
3110            Arc::new(Field::new(
3111                Field::MAP_ENTRIES_FIELD_DEFAULT_NAME,
3112                entry_struct.data_type().clone(),
3113                false,
3114            )),
3115            false,
3116        );
3117
3118        let entry_offsets = Buffer::from_slice_ref([0, 2, 4, 6]);
3119        let map_data = ArrayData::builder(map_data_type)
3120            .len(3)
3121            .add_buffer(entry_offsets)
3122            .add_child_data(entry_struct.into_data())
3123            .build()
3124            .unwrap();
3125        let map_array = MapArray::from(map_data);
3126
3127        let dict_keys = Int8Array::from_iter_values([0, 1, 1, 2, 2, 1]);
3128        let dict_dict_array = DictionaryArray::new(dict_keys, Arc::new(map_array));
3129
3130        let schema = Arc::new(Schema::new(vec![Field::new(
3131            "f1",
3132            dict_dict_array.data_type().clone(),
3133            false,
3134        )]));
3135        let input_batch = RecordBatch::try_new(schema, vec![Arc::new(dict_dict_array)]).unwrap();
3136        let output_batch = roundtrip_ipc_stream(&input_batch);
3137        assert_eq!(input_batch, output_batch);
3138    }
3139
3140    fn test_roundtrip_stream_dict_of_list_of_dict_impl<
3141        OffsetSize: OffsetSizeTrait,
3142        U: ArrowNativeType,
3143    >(
3144        list_data_type: DataType,
3145        offsets: &[U; 5],
3146    ) {
3147        let values = StringArray::from(vec![Some("a"), None, Some("c"), None]);
3148        let keys = Int8Array::from_iter_values([0, 0, 1, 2, 0, 1, 3]);
3149        let dict_array = DictionaryArray::new(keys, Arc::new(values));
3150        let dict_data = dict_array.to_data();
3151
3152        let value_offsets = Buffer::from_slice_ref(offsets);
3153
3154        let list_data = ArrayData::builder(list_data_type)
3155            .len(4)
3156            .add_buffer(value_offsets)
3157            .add_child_data(dict_data)
3158            .build()
3159            .unwrap();
3160        let list_array = GenericListArray::<OffsetSize>::from(list_data);
3161
3162        let keys_for_dict = Int8Array::from_iter_values([0, 3, 0, 1, 1, 2, 0, 1, 3]);
3163        let dict_dict_array = DictionaryArray::new(keys_for_dict, Arc::new(list_array));
3164
3165        let schema = Arc::new(Schema::new(vec![Field::new(
3166            "f1",
3167            dict_dict_array.data_type().clone(),
3168            false,
3169        )]));
3170        let input_batch = RecordBatch::try_new(schema, vec![Arc::new(dict_dict_array)]).unwrap();
3171        let output_batch = roundtrip_ipc_stream(&input_batch);
3172        assert_eq!(input_batch, output_batch);
3173    }
3174
3175    #[test]
3176    fn test_roundtrip_stream_dict_of_list_of_dict() {
3177        // list
3178        #[expect(deprecated)]
3179        let list_data_type = DataType::List(Arc::new(Field::new_dict(
3180            "item",
3181            DataType::Dictionary(Box::new(DataType::Int8), Box::new(DataType::Utf8)),
3182            true,
3183            1,
3184            false,
3185        )));
3186        let offsets: &[i32; 5] = &[0, 2, 4, 4, 6];
3187        test_roundtrip_stream_dict_of_list_of_dict_impl::<i32, i32>(list_data_type, offsets);
3188
3189        // large list
3190        #[expect(deprecated)]
3191        let list_data_type = DataType::LargeList(Arc::new(Field::new_dict(
3192            "item",
3193            DataType::Dictionary(Box::new(DataType::Int8), Box::new(DataType::Utf8)),
3194            true,
3195            1,
3196            false,
3197        )));
3198        let offsets: &[i64; 5] = &[0, 2, 4, 4, 7];
3199        test_roundtrip_stream_dict_of_list_of_dict_impl::<i64, i64>(list_data_type, offsets);
3200    }
3201
3202    #[test]
3203    fn test_roundtrip_stream_dict_of_fixed_size_list_of_dict() {
3204        let values = StringArray::from(vec![Some("a"), None, Some("c"), None]);
3205        let keys = Int8Array::from_iter_values([0, 0, 1, 2, 0, 1, 3, 1, 2]);
3206        let dict_array = DictionaryArray::new(keys, Arc::new(values));
3207        let dict_data = dict_array.into_data();
3208
3209        #[expect(deprecated)]
3210        let list_data_type = DataType::FixedSizeList(
3211            Arc::new(Field::new_dict(
3212                "item",
3213                DataType::Dictionary(Box::new(DataType::Int8), Box::new(DataType::Utf8)),
3214                true,
3215                1,
3216                false,
3217            )),
3218            3,
3219        );
3220        let list_data = ArrayData::builder(list_data_type)
3221            .len(3)
3222            .add_child_data(dict_data)
3223            .build()
3224            .unwrap();
3225        let list_array = FixedSizeListArray::from(list_data);
3226
3227        let keys_for_dict = Int8Array::from_iter_values([0, 1, 0, 1, 1, 2, 0, 1, 2]);
3228        let dict_dict_array = DictionaryArray::new(keys_for_dict, Arc::new(list_array));
3229
3230        let schema = Arc::new(Schema::new(vec![Field::new(
3231            "f1",
3232            dict_dict_array.data_type().clone(),
3233            false,
3234        )]));
3235        let input_batch = RecordBatch::try_new(schema, vec![Arc::new(dict_dict_array)]).unwrap();
3236        let output_batch = roundtrip_ipc_stream(&input_batch);
3237        assert_eq!(input_batch, output_batch);
3238    }
3239
3240    const LONG_TEST_STRING: &str =
3241        "This is a long string to make sure binary view array handles it";
3242
3243    #[test]
3244    fn test_roundtrip_view_types() {
3245        let schema = Schema::new(vec![
3246            Field::new("field_1", DataType::BinaryView, true),
3247            Field::new("field_2", DataType::Utf8, true),
3248            Field::new("field_3", DataType::Utf8View, true),
3249        ]);
3250        let bin_values: Vec<Option<&[u8]>> = vec![
3251            Some(b"foo"),
3252            None,
3253            Some(b"bar"),
3254            Some(LONG_TEST_STRING.as_bytes()),
3255        ];
3256        let utf8_values: Vec<Option<&str>> =
3257            vec![Some("foo"), None, Some("bar"), Some(LONG_TEST_STRING)];
3258        let bin_view_array = BinaryViewArray::from_iter(bin_values);
3259        let utf8_array = StringArray::from_iter(utf8_values.iter());
3260        let utf8_view_array = StringViewArray::from_iter(utf8_values);
3261        let record_batch = RecordBatch::try_new(
3262            Arc::new(schema.clone()),
3263            vec![
3264                Arc::new(bin_view_array),
3265                Arc::new(utf8_array),
3266                Arc::new(utf8_view_array),
3267            ],
3268        )
3269        .unwrap();
3270
3271        assert_eq!(record_batch, roundtrip_ipc(&record_batch));
3272        assert_eq!(record_batch, roundtrip_ipc_stream(&record_batch));
3273
3274        let sliced_batch = record_batch.slice(1, 2);
3275        assert_eq!(sliced_batch, roundtrip_ipc(&sliced_batch));
3276        assert_eq!(sliced_batch, roundtrip_ipc_stream(&sliced_batch));
3277    }
3278
3279    #[test]
3280    fn test_roundtrip_view_types_nested_dict() {
3281        let bin_values: Vec<Option<&[u8]>> = vec![
3282            Some(b"foo"),
3283            None,
3284            Some(b"bar"),
3285            Some(LONG_TEST_STRING.as_bytes()),
3286            Some(b"field"),
3287        ];
3288        let utf8_values: Vec<Option<&str>> = vec![
3289            Some("foo"),
3290            None,
3291            Some("bar"),
3292            Some(LONG_TEST_STRING),
3293            Some("field"),
3294        ];
3295        let bin_view_array = Arc::new(BinaryViewArray::from_iter(bin_values));
3296        let utf8_view_array = Arc::new(StringViewArray::from_iter(utf8_values));
3297
3298        let key_dict_keys = Int8Array::from_iter_values([0, 0, 2, 2, 0, 2, 3]);
3299        let key_dict_array = DictionaryArray::new(key_dict_keys, utf8_view_array.clone());
3300        #[expect(deprecated)]
3301        let keys_field = Arc::new(Field::new_dict(
3302            Field::MAP_KEY_FIELD_DEFAULT_NAME,
3303            DataType::Dictionary(Box::new(DataType::Int8), Box::new(DataType::Utf8View)),
3304            false,
3305            1,
3306            false,
3307        ));
3308
3309        let value_dict_keys = Int8Array::from_iter_values([0, 3, 0, 1, 2, 0, 1]);
3310        let value_dict_array = DictionaryArray::new(value_dict_keys, bin_view_array);
3311        #[expect(deprecated)]
3312        let values_field = Arc::new(Field::new_dict(
3313            Field::MAP_VALUE_FIELD_DEFAULT_NAME,
3314            DataType::Dictionary(Box::new(DataType::Int8), Box::new(DataType::BinaryView)),
3315            true,
3316            2,
3317            false,
3318        ));
3319        let entry_struct = StructArray::from(vec![
3320            (keys_field, make_array(key_dict_array.into_data())),
3321            (values_field, make_array(value_dict_array.into_data())),
3322        ]);
3323
3324        let map_data_type = DataType::Map(
3325            Arc::new(Field::new(
3326                Field::MAP_ENTRIES_FIELD_DEFAULT_NAME,
3327                entry_struct.data_type().clone(),
3328                false,
3329            )),
3330            false,
3331        );
3332        let entry_offsets = Buffer::from_slice_ref([0, 2, 4, 7]);
3333        let map_data = ArrayData::builder(map_data_type)
3334            .len(3)
3335            .add_buffer(entry_offsets)
3336            .add_child_data(entry_struct.into_data())
3337            .build()
3338            .unwrap();
3339        let map_array = MapArray::from(map_data);
3340
3341        let dict_keys = Int8Array::from_iter_values([0, 1, 0, 1, 1, 2, 0, 1, 2]);
3342        let dict_dict_array = DictionaryArray::new(dict_keys, Arc::new(map_array));
3343        let schema = Arc::new(Schema::new(vec![Field::new(
3344            "f1",
3345            dict_dict_array.data_type().clone(),
3346            false,
3347        )]));
3348        let batch = RecordBatch::try_new(schema, vec![Arc::new(dict_dict_array)]).unwrap();
3349        assert_eq!(batch, roundtrip_ipc(&batch));
3350        assert_eq!(batch, roundtrip_ipc_stream(&batch));
3351
3352        let sliced_batch = batch.slice(1, 2);
3353        assert_eq!(sliced_batch, roundtrip_ipc(&sliced_batch));
3354        assert_eq!(sliced_batch, roundtrip_ipc_stream(&sliced_batch));
3355    }
3356
3357    #[test]
3358    fn test_no_columns_batch() {
3359        let schema = Arc::new(Schema::empty());
3360        let options = RecordBatchOptions::new()
3361            .with_match_field_names(true)
3362            .with_row_count(Some(10));
3363        let input_batch = RecordBatch::try_new_with_options(schema, vec![], &options).unwrap();
3364        let output_batch = roundtrip_ipc_stream(&input_batch);
3365        assert_eq!(input_batch, output_batch);
3366    }
3367
3368    #[test]
3369    fn test_unaligned() {
3370        let batch = RecordBatch::try_from_iter(vec![(
3371            "i32",
3372            Arc::new(Int32Array::from(vec![1, 2, 3, 4])) as _,
3373        )])
3374        .unwrap();
3375
3376        let r#gen = IpcDataGenerator {};
3377        let mut dict_tracker = DictionaryTracker::new(false);
3378        let (_, encoded) = r#gen
3379            .encode(
3380                &batch,
3381                &mut dict_tracker,
3382                &Default::default(),
3383                &mut Default::default(),
3384            )
3385            .unwrap();
3386
3387        let message = root_as_message(&encoded.ipc_message).unwrap();
3388
3389        // Construct an unaligned buffer
3390        let mut buffer = MutableBuffer::with_capacity(encoded.arrow_data.len() + 1);
3391        buffer.push(0_u8);
3392        buffer.extend_from_slice(&encoded.arrow_data);
3393        let b = Buffer::from(buffer).slice(1);
3394        assert_ne!(b.as_ptr().align_offset(8), 0);
3395
3396        let ipc_batch = message.header_as_record_batch().unwrap();
3397        let roundtrip = RecordBatchDecoder::try_new(
3398            &b,
3399            ipc_batch,
3400            batch.schema(),
3401            &Default::default(),
3402            &message.version(),
3403        )
3404        .unwrap()
3405        .with_require_alignment(false)
3406        .read_record_batch()
3407        .unwrap();
3408        assert_eq!(batch, roundtrip);
3409    }
3410
3411    #[test]
3412    fn test_unaligned_throws_error_with_require_alignment() {
3413        let batch = RecordBatch::try_from_iter(vec![(
3414            "i32",
3415            Arc::new(Int32Array::from(vec![1, 2, 3, 4])) as _,
3416        )])
3417        .unwrap();
3418
3419        let r#gen = IpcDataGenerator {};
3420        let mut dict_tracker = DictionaryTracker::new(false);
3421        let (_, encoded) = r#gen
3422            .encode(
3423                &batch,
3424                &mut dict_tracker,
3425                &Default::default(),
3426                &mut Default::default(),
3427            )
3428            .unwrap();
3429
3430        let message = root_as_message(&encoded.ipc_message).unwrap();
3431
3432        // Construct an unaligned buffer
3433        let mut buffer = MutableBuffer::with_capacity(encoded.arrow_data.len() + 1);
3434        buffer.push(0_u8);
3435        buffer.extend_from_slice(&encoded.arrow_data);
3436        let b = Buffer::from(buffer).slice(1);
3437        assert_ne!(b.as_ptr().align_offset(8), 0);
3438
3439        let ipc_batch = message.header_as_record_batch().unwrap();
3440        let result = RecordBatchDecoder::try_new(
3441            &b,
3442            ipc_batch,
3443            batch.schema(),
3444            &Default::default(),
3445            &message.version(),
3446        )
3447        .unwrap()
3448        .with_require_alignment(true)
3449        .read_record_batch();
3450
3451        let error = result.unwrap_err();
3452        assert_eq!(
3453            error.to_string(),
3454            "Invalid argument error: Misaligned buffers[0] in array of type Int32, \
3455             offset from expected alignment of 4 by 1"
3456        );
3457    }
3458
3459    /// Verify that misaligned IPC buffers are caught by `require_alignment = true`.
3460    ///
3461    /// For each array type we shift the IPC body by every byte offset in 0..OFFSET_RANGE and
3462    /// assert that the decoder errors exactly when the offset violates the type's
3463    /// minimum alignment requirement.  The Arrow columnar spec permits alignment to
3464    /// any multiple of 8 or 64 bytes, so multiples of 8 (which include every multiple
3465    /// of 64) must succeed; everything else must return a "Misaligned buffers" error.
3466    /// See <https://arrow.apache.org/docs/format/Columnar.html#buffer-alignment-and-padding>.
3467    #[test]
3468    fn test_misaligned_buffers_error() {
3469        const OFFSET_RANGE: usize = 128;
3470
3471        // (array, minimum required alignment in bytes)
3472        // Fixed-width: alignment == element byte width.
3473        // Variable-width (e.g. StringArray): the offsets buffer drives alignment (Int32 → 4).
3474        let cases: Vec<(ArrayRef, usize)> = vec![
3475            (Arc::new(Int32Array::from_iter(0i32..100)) as _, 4),
3476            (Arc::new(Int64Array::from_iter(0i64..100)) as _, 8),
3477            (
3478                Arc::new(StringArray::from_iter_values(
3479                    (0..100).map(|i| i.to_string()),
3480                )) as _,
3481                4,
3482            ),
3483            (
3484                Arc::new(LargeStringArray::from_iter_values(
3485                    (0..100).map(|i| i.to_string()),
3486                )) as _,
3487                8,
3488            ),
3489        ];
3490
3491        for (array, alignment) in cases {
3492            let batch = RecordBatch::try_from_iter(vec![("col", Arc::clone(&array))]).unwrap();
3493            let encoder = IpcDataGenerator {};
3494            let mut dict_tracker = DictionaryTracker::new(false);
3495            let (_, encoded) = encoder
3496                .encode(
3497                    &batch,
3498                    &mut dict_tracker,
3499                    &Default::default(),
3500                    &mut Default::default(),
3501                )
3502                .unwrap();
3503            let message = root_as_message(&encoded.ipc_message).unwrap();
3504            let ipc_batch = message.header_as_record_batch().unwrap();
3505
3506            for offset in 0..OFFSET_RANGE {
3507                // MutableBuffer always allocates at a 64-byte aligned base address.
3508                // Slicing by `offset` bytes yields a pointer at `base_ptr + offset`,
3509                // whose alignment is gcd(64, offset).  That makes `offset` the sole
3510                // determinant of the resulting alignment — without this guarantee the
3511                // test would be non-deterministic depending on what the allocator returns.
3512                let mut storage = MutableBuffer::with_capacity(encoded.arrow_data.len() + offset);
3513                for _ in 0..offset {
3514                    storage.push(0_u8);
3515                }
3516                storage.extend_from_slice(&encoded.arrow_data);
3517                let buf = Buffer::from(storage).slice(offset);
3518
3519                let result = RecordBatchDecoder::try_new(
3520                    &buf,
3521                    ipc_batch,
3522                    batch.schema(),
3523                    &Default::default(),
3524                    &message.version(),
3525                )
3526                .unwrap()
3527                .with_require_alignment(true)
3528                .read_record_batch();
3529
3530                if offset % alignment == 0 {
3531                    assert!(
3532                        result.is_ok(),
3533                        "type={} offset={offset}: expected Ok but got {:?}",
3534                        array.data_type(),
3535                        result.unwrap_err(),
3536                    );
3537                } else {
3538                    let err = result
3539                        .expect_err(&format!(
3540                            "type={} offset={offset}: expected Err for misaligned buffer",
3541                            array.data_type()
3542                        ))
3543                        .to_string();
3544                    assert!(
3545                        err.contains("Misaligned buffers"),
3546                        "type={} offset={offset}: unexpected error: {err}",
3547                        array.data_type(),
3548                    );
3549                }
3550            }
3551        }
3552    }
3553
3554    #[test]
3555    #[cfg_attr(miri, ignore)] // Takes too long
3556    fn test_file_with_massive_column_count() {
3557        // 499_999 is upper limit for default settings (1_000_000)
3558        let limit = 600_000;
3559
3560        let fields = (0..limit)
3561            .map(|i| Field::new(format!("{i}"), DataType::Boolean, false))
3562            .collect::<Vec<_>>();
3563        let schema = Arc::new(Schema::new(fields));
3564        let batch = RecordBatch::new_empty(schema);
3565
3566        let mut buf = Vec::new();
3567        let mut writer = crate::writer::FileWriter::try_new(&mut buf, batch.schema_ref()).unwrap();
3568        writer.write(&batch).unwrap();
3569        writer.finish().unwrap();
3570        drop(writer);
3571
3572        let mut reader = FileReaderBuilder::new()
3573            .with_max_footer_fb_tables(1_500_000)
3574            .build(std::io::Cursor::new(buf))
3575            .unwrap();
3576        let roundtrip_batch = reader.next().unwrap().unwrap();
3577
3578        assert_eq!(batch, roundtrip_batch);
3579    }
3580
3581    #[test]
3582    fn test_file_with_deeply_nested_columns() {
3583        // 60 is upper limit for default settings (64)
3584        let limit = 61;
3585
3586        let fields = (0..limit).fold(
3587            vec![Field::new("leaf", DataType::Boolean, false)],
3588            |field, index| vec![Field::new_struct(format!("{index}"), field, false)],
3589        );
3590        let schema = Arc::new(Schema::new(fields));
3591        let batch = RecordBatch::new_empty(schema);
3592
3593        let mut buf = Vec::new();
3594        let mut writer = crate::writer::FileWriter::try_new(&mut buf, batch.schema_ref()).unwrap();
3595        writer.write(&batch).unwrap();
3596        writer.finish().unwrap();
3597        drop(writer);
3598
3599        let mut reader = FileReaderBuilder::new()
3600            .with_max_footer_fb_depth(65)
3601            .build(std::io::Cursor::new(buf))
3602            .unwrap();
3603        let roundtrip_batch = reader.next().unwrap().unwrap();
3604
3605        assert_eq!(batch, roundtrip_batch);
3606    }
3607
3608    #[test]
3609    fn test_stream_reader_rejects_short_validity_buffer() {
3610        // Reproduce #7124: serialize an Int32Array with too few validity bits.
3611        let data = ArrayDataBuilder::new(DataType::Int32)
3612            .len(8000)
3613            .add_buffer(ScalarBuffer::<i32>::from_iter(0..8000).into())
3614            .nulls(Some(NullBuffer::from(&[true, false, true, false])));
3615        let array: ArrayRef = unsafe { Arc::new(Int32Array::from(data.build_unchecked())) };
3616        let batch = RecordBatch::try_from_iter([("a", array)]).unwrap();
3617
3618        let mut stream = Vec::new();
3619        let mut writer =
3620            crate::writer::StreamWriter::try_new(&mut stream, &batch.schema()).unwrap();
3621        writer.write(&batch).unwrap();
3622        writer.finish().unwrap();
3623
3624        let mut reader = StreamReader::try_new(Cursor::new(stream), None).unwrap();
3625        let err = reader.next().unwrap().unwrap_err();
3626        assert_eq!(
3627            err.to_string(),
3628            "Invalid argument error: null_bit_buffer size too small. got 1 needed 1000"
3629        );
3630    }
3631
3632    #[test]
3633    fn test_stream_reader_rejects_short_struct_validity_buffer() {
3634        // Serialize a StructArray with too few validity bits.
3635        let fields = Fields::from(vec![Field::new("b", DataType::Int32, false)]);
3636        let child: ArrayRef = Arc::new(Int32Array::from_iter_values(0..8000));
3637        let nulls = NullBuffer::from(&[true, false, true, false]);
3638        let array: ArrayRef =
3639            unsafe { Arc::new(StructArray::new_unchecked(fields, vec![child], Some(nulls))) };
3640        let batch = RecordBatch::try_from_iter([("a", array)]).unwrap();
3641
3642        let mut stream = Vec::new();
3643        let mut writer =
3644            crate::writer::StreamWriter::try_new(&mut stream, &batch.schema()).unwrap();
3645        writer.write(&batch).unwrap();
3646        writer.finish().unwrap();
3647
3648        let mut reader = StreamReader::try_new(Cursor::new(stream), None).unwrap();
3649        let err = reader.next().unwrap().unwrap_err();
3650        assert_eq!(
3651            err.to_string(),
3652            "Invalid argument error: null_bit_buffer size too small. got 1 needed 1000"
3653        );
3654    }
3655
3656    #[test]
3657    fn test_invalid_struct_array_ipc_read_errors() {
3658        let a_field = Field::new("a", DataType::Int32, false);
3659        let b_field = Field::new("b", DataType::Int32, false);
3660        let struct_fields = Fields::from(vec![a_field.clone(), b_field.clone()]);
3661
3662        let a_array_data = ArrayData::builder(a_field.data_type().clone())
3663            .len(4)
3664            .add_buffer(Buffer::from_slice_ref([1, 2, 3, 4]))
3665            .build()
3666            .unwrap();
3667        let b_array_data = ArrayData::builder(b_field.data_type().clone())
3668            .len(3)
3669            .add_buffer(Buffer::from_slice_ref([5, 6, 7]))
3670            .build()
3671            .unwrap();
3672
3673        let invalid_struct_arr = unsafe {
3674            StructArray::new_unchecked(
3675                struct_fields,
3676                vec![make_array(a_array_data), make_array(b_array_data)],
3677                None,
3678            )
3679        };
3680
3681        expect_ipc_validation_error(
3682            Arc::new(invalid_struct_arr),
3683            "Invalid argument error: Incorrect array length for StructArray field \"b\", expected 4 got 3",
3684        );
3685    }
3686
3687    #[test]
3688    fn test_invalid_nested_array_ipc_read_errors() {
3689        // one of the nested arrays has invalid data
3690        let a_field = Field::new("a", DataType::Int32, false);
3691        let b_field = Field::new("b", DataType::Utf8, false);
3692
3693        let schema = Arc::new(Schema::new(vec![Field::new_struct(
3694            "s",
3695            vec![a_field.clone(), b_field.clone()],
3696            false,
3697        )]));
3698
3699        let a_array_data = ArrayData::builder(a_field.data_type().clone())
3700            .len(4)
3701            .add_buffer(Buffer::from_slice_ref([1, 2, 3, 4]))
3702            .build()
3703            .unwrap();
3704        // invalid nested child array -- length is correct, but has invalid utf8 data
3705        let b_array_data = {
3706            let valid: &[u8] = b"   ";
3707            let mut invalid = vec![];
3708            invalid.extend_from_slice(b"ValidString");
3709            invalid.extend_from_slice(INVALID_UTF8_FIRST_CHAR);
3710            let binary_array =
3711                BinaryArray::from_iter(vec![None, Some(valid), None, Some(&invalid)]);
3712            let array = unsafe {
3713                StringArray::new_unchecked(
3714                    binary_array.offsets().clone(),
3715                    binary_array.values().clone(),
3716                    binary_array.nulls().cloned(),
3717                )
3718            };
3719            array.into_data()
3720        };
3721        let struct_data_type = schema.field(0).data_type();
3722
3723        let invalid_struct_arr = unsafe {
3724            make_array(
3725                ArrayData::builder(struct_data_type.clone())
3726                    .len(4)
3727                    .add_child_data(a_array_data)
3728                    .add_child_data(b_array_data)
3729                    .build_unchecked(),
3730            )
3731        };
3732        expect_ipc_validation_error(
3733            invalid_struct_arr,
3734            "Invalid argument error: Invalid UTF8 sequence at string index 3 (3..18): invalid utf-8 sequence of 1 bytes from index 11",
3735        );
3736    }
3737
3738    #[test]
3739    fn test_same_dict_id_without_preserve() {
3740        let batch = RecordBatch::try_new(
3741            Arc::new(Schema::new(
3742                ["a", "b"]
3743                    .iter()
3744                    .map(|name| {
3745                        #[expect(deprecated)]
3746                        Field::new_dict(
3747                            name.to_string(),
3748                            DataType::Dictionary(
3749                                Box::new(DataType::Int32),
3750                                Box::new(DataType::Utf8),
3751                            ),
3752                            true,
3753                            0,
3754                            false,
3755                        )
3756                    })
3757                    .collect::<Vec<Field>>(),
3758            )),
3759            vec![
3760                Arc::new(
3761                    vec![Some("c"), Some("d")]
3762                        .into_iter()
3763                        .collect::<DictionaryArray<Int32Type>>(),
3764                ) as ArrayRef,
3765                Arc::new(
3766                    vec![Some("e"), Some("f")]
3767                        .into_iter()
3768                        .collect::<DictionaryArray<Int32Type>>(),
3769                ) as ArrayRef,
3770            ],
3771        )
3772        .expect("Failed to create RecordBatch");
3773
3774        // serialize the record batch as an IPC stream
3775        let mut buf = vec![];
3776        {
3777            let mut writer = crate::writer::StreamWriter::try_new_with_options(
3778                &mut buf,
3779                batch.schema().as_ref(),
3780                crate::writer::IpcWriteOptions::default(),
3781            )
3782            .expect("Failed to create StreamWriter");
3783            writer.write(&batch).expect("Failed to write RecordBatch");
3784            writer.finish().expect("Failed to finish StreamWriter");
3785        }
3786
3787        StreamReader::try_new(std::io::Cursor::new(buf), None)
3788            .expect("Failed to create StreamReader")
3789            .for_each(|decoded_batch| {
3790                assert_eq!(decoded_batch.expect("Failed to read RecordBatch"), batch);
3791            });
3792    }
3793
3794    #[test]
3795    fn test_validation_of_invalid_list_array() {
3796        // ListArray with invalid offsets
3797        let array = unsafe {
3798            let values = Int32Array::from(vec![1, 2, 3]);
3799            let bad_offsets = ScalarBuffer::<i32>::from(vec![0, 2, 4, 2]); // offsets can't go backwards
3800            let offsets = OffsetBuffer::new_unchecked(bad_offsets); // INVALID array created
3801            let field = Field::new_list_field(DataType::Int32, true);
3802            let nulls = None;
3803            ListArray::new(Arc::new(field), offsets, Arc::new(values), nulls)
3804        };
3805
3806        expect_ipc_validation_error(
3807            Arc::new(array),
3808            "Invalid argument error: Offset invariant failure: offset at position 2 out of bounds: 4 > 2",
3809        );
3810    }
3811
3812    #[test]
3813    fn test_validation_of_invalid_string_array() {
3814        let valid: &[u8] = b"   ";
3815        let mut invalid = vec![];
3816        invalid.extend_from_slice(b"ThisStringIsCertainlyLongerThan12Bytes");
3817        invalid.extend_from_slice(INVALID_UTF8_FIRST_CHAR);
3818        let binary_array = BinaryArray::from_iter(vec![None, Some(valid), None, Some(&invalid)]);
3819        // data is not valid utf8 we can not construct a correct StringArray
3820        // safely, so purposely create an invalid StringArray
3821        let array = unsafe {
3822            StringArray::new_unchecked(
3823                binary_array.offsets().clone(),
3824                binary_array.values().clone(),
3825                binary_array.nulls().cloned(),
3826            )
3827        };
3828        expect_ipc_validation_error(
3829            Arc::new(array),
3830            "Invalid argument error: Invalid UTF8 sequence at string index 3 (3..45): invalid utf-8 sequence of 1 bytes from index 38",
3831        );
3832    }
3833
3834    #[test]
3835    fn test_validation_of_invalid_string_view_array() {
3836        let valid: &[u8] = b"   ";
3837        let mut invalid = vec![];
3838        invalid.extend_from_slice(b"ThisStringIsCertainlyLongerThan12Bytes");
3839        invalid.extend_from_slice(INVALID_UTF8_FIRST_CHAR);
3840        let binary_view_array =
3841            BinaryViewArray::from_iter(vec![None, Some(valid), None, Some(&invalid)]);
3842        // data is not valid utf8 we can not construct a correct StringArray
3843        // safely, so purposely create an invalid StringArray
3844        let array = unsafe {
3845            StringViewArray::new_unchecked(
3846                binary_view_array.views().clone(),
3847                Arc::clone(binary_view_array.data_buffers()),
3848                binary_view_array.nulls().cloned(),
3849            )
3850        };
3851        expect_ipc_validation_error(
3852            Arc::new(array),
3853            "Invalid argument error: Encountered non-UTF-8 data at index 3: invalid utf-8 sequence of 1 bytes from index 38",
3854        );
3855    }
3856
3857    /// return an invalid dictionary array (key is larger than values)
3858    /// ListArray with invalid offsets
3859    #[test]
3860    fn test_validation_of_invalid_dictionary_array() {
3861        let array = unsafe {
3862            let values = StringArray::from_iter_values(["a", "b", "c"]);
3863            let keys = Int32Array::from(vec![1, 200]); // keys are not valid for values
3864            DictionaryArray::new_unchecked(keys, Arc::new(values))
3865        };
3866
3867        expect_ipc_validation_error(
3868            Arc::new(array),
3869            "Invalid argument error: Value at position 1 out of bounds: 200 (should be in [0, 2])",
3870        );
3871    }
3872
3873    #[test]
3874    fn test_validation_of_invalid_union_array() {
3875        let array = unsafe {
3876            let fields = UnionFields::try_new(
3877                vec![1, 3], // typeids : type id 2 is not valid
3878                vec![
3879                    Field::new("a", DataType::Int32, false),
3880                    Field::new("b", DataType::Utf8, false),
3881                ],
3882            )
3883            .unwrap();
3884            let type_ids = ScalarBuffer::from(vec![1i8, 2, 3]); // 2 is invalid
3885            let offsets = None;
3886            let children: Vec<ArrayRef> = vec![
3887                Arc::new(Int32Array::from(vec![10, 20, 30])),
3888                Arc::new(StringArray::from(vec![Some("a"), Some("b"), Some("c")])),
3889            ];
3890
3891            UnionArray::new_unchecked(fields, type_ids, offsets, children)
3892        };
3893
3894        expect_ipc_validation_error(
3895            Arc::new(array),
3896            "Invalid argument error: Type Ids values must match one of the field type ids",
3897        );
3898    }
3899
3900    /// Invalid Utf-8 sequence in the first character
3901    /// <https://stackoverflow.com/questions/1301402/example-invalid-utf8-string>
3902    const INVALID_UTF8_FIRST_CHAR: &[u8] = &[0xa0, 0xa1, 0x20, 0x20];
3903
3904    /// Expect an error when reading the record batch using IPC or IPC Streams
3905    fn expect_ipc_validation_error(array: ArrayRef, expected_err: &str) {
3906        let rb = RecordBatch::try_from_iter([("a", array)]).unwrap();
3907
3908        // IPC Stream format
3909        let buf = write_stream(&rb); // write is ok
3910        read_stream_skip_validation(&buf).unwrap();
3911        let err = read_stream(&buf).unwrap_err();
3912        assert_eq!(err.to_string(), expected_err);
3913
3914        // IPC File format
3915        let buf = write_ipc(&rb); // write is ok
3916        read_ipc_skip_validation(&buf).unwrap();
3917        let err = read_ipc(&buf).unwrap_err();
3918        assert_eq!(err.to_string(), expected_err);
3919
3920        // IPC Format with FileDecoder
3921        read_ipc_with_decoder_skip_validation(buf.clone()).unwrap();
3922        let err = read_ipc_with_decoder(buf).unwrap_err();
3923        assert_eq!(err.to_string(), expected_err);
3924    }
3925
3926    #[test]
3927    fn test_roundtrip_schema() {
3928        let schema = Schema::new(vec![
3929            Field::new(
3930                "a",
3931                DataType::Dictionary(Box::new(DataType::UInt16), Box::new(DataType::Utf8)),
3932                false,
3933            ),
3934            Field::new(
3935                "b",
3936                DataType::Dictionary(Box::new(DataType::UInt16), Box::new(DataType::Utf8)),
3937                false,
3938            ),
3939        ]);
3940
3941        let options = IpcWriteOptions::default();
3942        let data_gen = IpcDataGenerator::default();
3943        let mut dict_tracker = DictionaryTracker::new(false);
3944        let encoded_data =
3945            data_gen.schema_to_bytes_with_dictionary_tracker(&schema, &mut dict_tracker, &options);
3946        let mut schema_bytes = vec![];
3947        write_message(&mut schema_bytes, encoded_data, &options).expect("write_message");
3948
3949        let begin_offset: usize = if schema_bytes[0..4].eq(&CONTINUATION_MARKER) {
3950            4
3951        } else {
3952            0
3953        };
3954
3955        size_prefixed_root_as_message(&schema_bytes[begin_offset..])
3956            .expect_err("size_prefixed_root_as_message");
3957
3958        let msg = parse_message(&schema_bytes).expect("parse_message");
3959        let ipc_schema = msg.header_as_schema().expect("header_as_schema");
3960        let new_schema = try_fb_to_schema(ipc_schema).unwrap();
3961
3962        assert_eq!(schema, new_schema);
3963    }
3964
3965    #[test]
3966    fn test_negative_meta_len() {
3967        let bytes = i32::to_le_bytes(-1);
3968        let mut buf = vec![];
3969        buf.extend(CONTINUATION_MARKER);
3970        buf.extend(bytes);
3971
3972        let reader = StreamReader::try_new(Cursor::new(buf), None);
3973        assert!(reader.is_err());
3974    }
3975
3976    /// Per the IPC specification, dictionary batches may be omitted for
3977    /// dictionary-encoded columns where all values are null.  The C++
3978    /// implementation relies on this and does not emit a dictionary batch
3979    /// in that case.  Verify that the Rust reader handles such streams
3980    /// by synthesizing an empty dictionary instead of returning an error.
3981    #[test]
3982    fn test_read_null_dict_without_dictionary_batch() {
3983        // Build an all-null dictionary-encoded column.
3984        let keys = Int32Array::new_null(4);
3985        let values: ArrayRef = new_empty_array(&DataType::Utf8);
3986        let dict_array = DictionaryArray::new(keys, values);
3987
3988        let schema = Arc::new(Schema::new(vec![Field::new(
3989            "d",
3990            dict_array.data_type().clone(),
3991            true,
3992        )]));
3993        let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(dict_array)]).unwrap();
3994
3995        // Write a normal IPC stream (which includes the dictionary batch).
3996        let full_stream = write_stream(&batch);
3997
3998        // Parse the stream into individual messages and reconstruct it
3999        // without the DictionaryBatch message, simulating what C++ emits
4000        // for an all-null dictionary column.
4001        let mut stripped = Vec::new();
4002        let mut cursor = Cursor::new(&full_stream);
4003        loop {
4004            // Each message is: [continuation (4 bytes)] [meta_len (4 bytes)]
4005            //                   [metadata (meta_len bytes)] [body (bodyLength bytes)]
4006            let mut header = [0u8; 4];
4007            if cursor.read_exact(&mut header).is_err() {
4008                break;
4009            }
4010            if header == CONTINUATION_MARKER && cursor.read_exact(&mut header).is_err() {
4011                break;
4012            }
4013            let meta_len = u32::from_le_bytes(header) as usize;
4014            if meta_len == 0 {
4015                // EOS marker — write it through.
4016                stripped.extend_from_slice(&CONTINUATION_MARKER);
4017                stripped.extend_from_slice(&0u32.to_le_bytes());
4018                break;
4019            }
4020            let mut meta_buf = vec![0u8; meta_len];
4021            cursor.read_exact(&mut meta_buf).unwrap();
4022
4023            let message = root_as_message(&meta_buf).unwrap();
4024            let body_len = message.bodyLength() as usize;
4025            let mut body_buf = vec![0u8; body_len];
4026            cursor.read_exact(&mut body_buf).unwrap();
4027
4028            if message.header_type() == crate::MessageHeader::DictionaryBatch {
4029                // Skip the dictionary batch — this is what C++ does for
4030                // all-null dictionary columns.
4031                continue;
4032            }
4033            stripped.extend_from_slice(&CONTINUATION_MARKER);
4034            stripped.extend_from_slice(&(meta_len as u32).to_le_bytes());
4035            stripped.extend_from_slice(&meta_buf);
4036            stripped.extend_from_slice(&body_buf);
4037        }
4038
4039        // Reading the stripped stream must succeed.
4040        let result = read_stream(&stripped).unwrap();
4041        assert_eq!(result.num_rows(), 4);
4042        assert_eq!(result.num_columns(), 1);
4043
4044        let col = result.column(0);
4045        assert_eq!(col.null_count(), 4);
4046        assert_eq!(col.len(), 4);
4047        // The result must be a dictionary-typed array.
4048        assert!(matches!(col.data_type(), DataType::Dictionary(_, _)));
4049    }
4050
4051    // Tests projected reads where a ListView column is skipped before another column.
4052    // This catches cases where skipping the ListView consumes the wrong number of buffers.
4053    #[test]
4054    fn test_projection_skip_list_view() {
4055        use crate::reader::FileReader;
4056        use crate::writer::FileWriter;
4057        use arrow_array::{
4058            GenericListViewArray, Int32Array, RecordBatch,
4059            builder::{GenericListViewBuilder, UInt32Builder},
4060        };
4061        use arrow_schema::{DataType, Field, Schema};
4062        use std::sync::Arc;
4063
4064        // Build a small ListView column with a mix of valid and null entries
4065        let mut builder = GenericListViewBuilder::<i32, _>::new(UInt32Builder::new());
4066
4067        builder.values().append_value(1);
4068        builder.values().append_value(2);
4069        builder.append(true);
4070
4071        builder.append(false);
4072
4073        builder.values().append_value(3);
4074        builder.values().append_value(4);
4075        builder.append(true);
4076
4077        let list_view: GenericListViewArray<i32> = builder.finish();
4078
4079        // Second column with simple values
4080        let values = Int32Array::from(vec![10, 20, 30]);
4081
4082        // Schema: first column is ListView, second is Int32
4083        let schema = Arc::new(Schema::new(vec![
4084            Field::new("a", list_view.data_type().clone(), true),
4085            Field::new("b", DataType::Int32, false),
4086        ]));
4087        // Create a batch with both columns
4088        let batch =
4089            RecordBatch::try_new(schema, vec![Arc::new(list_view), Arc::new(values.clone())])
4090                .unwrap();
4091
4092        // Write the batch to IPC
4093        let mut buf = Vec::new();
4094        {
4095            let mut writer = FileWriter::try_new(&mut buf, &batch.schema()).unwrap();
4096            writer.write(&batch).unwrap();
4097            writer.finish().unwrap();
4098        }
4099
4100        // Skip ListView column and Project only column "b"
4101        let mut reader = FileReader::try_new(std::io::Cursor::new(buf), Some(vec![1])).unwrap();
4102        let read_batch = reader.next().unwrap().unwrap();
4103
4104        // Verify that the projected column is read correctly
4105        assert_eq!(read_batch.num_columns(), 1);
4106        assert_eq!(read_batch.column(0).as_ref(), &values);
4107    }
4108
4109    // Tests reading a column when a preceding V4 Union column is skipped.
4110    // V4 Union columns include a null buffer and type ids (and offsets for dense unions).
4111    #[test]
4112    fn test_projection_skip_union_v4() {
4113        use crate::MetadataVersion;
4114        use crate::reader::FileReader;
4115        use crate::writer::{FileWriter, IpcWriteOptions};
4116        use arrow_array::{
4117            ArrayRef, Int32Array, RecordBatch, builder::UnionBuilder, types::Int32Type,
4118        };
4119        use arrow_schema::{DataType, Field, Schema};
4120        use std::sync::Arc;
4121
4122        // Build a dense Union column with simple Int32 values
4123        let mut builder = UnionBuilder::new_dense();
4124        builder.append::<Int32Type>("a", 1).unwrap();
4125        builder.append::<Int32Type>("a", 2).unwrap();
4126        builder.append::<Int32Type>("a", 3).unwrap();
4127        let union = builder.build().unwrap();
4128
4129        // Second column with known values to verify correctness after projection
4130        let values = Int32Array::from(vec![10, 20, 30]);
4131
4132        // Schema: first column is Union (to be skipped), second is Int32 (to be read)
4133        let schema = Arc::new(Schema::new(vec![
4134            Field::new("union", union.data_type().clone(), false),
4135            Field::new("values", DataType::Int32, false),
4136        ]));
4137
4138        // Create a batch containing both columns
4139        let batch = RecordBatch::try_new(
4140            schema,
4141            vec![Arc::new(union) as ArrayRef, Arc::new(values.clone())],
4142        )
4143        .unwrap();
4144
4145        // Write IPC using V4 metadata to trigger Union null buffer behavior
4146        let mut buf = Vec::new();
4147        {
4148            let options = IpcWriteOptions::try_new(8, false, MetadataVersion::V4).unwrap();
4149            let mut writer =
4150                FileWriter::try_new_with_options(&mut buf, &batch.schema(), options).unwrap();
4151            writer.write(&batch).unwrap();
4152            writer.finish().unwrap();
4153        }
4154        // Read only the second column (skip the Union column)
4155        let mut reader = FileReader::try_new(std::io::Cursor::new(buf), Some(vec![1])).unwrap();
4156        let read_batch = reader.next().unwrap().unwrap();
4157
4158        // Verify that the projected column is read correctly after skipping Union
4159        assert_eq!(read_batch.num_columns(), 1);
4160        assert_eq!(read_batch.column(0).as_ref(), &values);
4161    }
4162
4163    // Tests reading a column when preceding fixed-width and boolean columns are skipped.
4164    // Covers all types that use the same two-buffer layout (null + values).
4165    // Verifies that skipping these types does not affect subsequent column decoding.
4166    #[test]
4167    fn test_projection_skip_fixed_width_types() {
4168        use std::sync::Arc;
4169
4170        use arrow_array::{ArrayRef, BooleanArray, Int32Array, RecordBatch, make_array};
4171        use arrow_buffer::{Buffer, MutableBuffer};
4172        use arrow_data::ArrayData;
4173        use arrow_schema::{DataType, Field, IntervalUnit, Schema, TimeUnit};
4174
4175        use crate::reader::FileReader;
4176        use crate::writer::FileWriter;
4177
4178        // Create a minimal array for a given fixed-width or boolean type
4179        fn make_array_for_type(data_type: DataType) -> ArrayRef {
4180            let len = 3;
4181
4182            if matches!(data_type, DataType::Boolean) {
4183                return Arc::new(BooleanArray::from(vec![true, false, true]));
4184            }
4185
4186            let width = data_type.primitive_width().unwrap();
4187            let data = ArrayData::builder(data_type)
4188                .len(len)
4189                .add_buffer(Buffer::from(MutableBuffer::from_len_zeroed(len * width)))
4190                .build()
4191                .unwrap();
4192
4193            make_array(data)
4194        }
4195
4196        // List of types that follow the same two-buffer layout (null + values)
4197        let data_types = vec![
4198            DataType::Boolean,
4199            DataType::Int8,
4200            DataType::Int16,
4201            DataType::Int32,
4202            DataType::Int64,
4203            DataType::UInt8,
4204            DataType::UInt16,
4205            DataType::UInt32,
4206            DataType::UInt64,
4207            DataType::Float16,
4208            DataType::Float32,
4209            DataType::Float64,
4210            DataType::Timestamp(TimeUnit::Second, None),
4211            DataType::Date32,
4212            DataType::Date64,
4213            DataType::Time32(TimeUnit::Second),
4214            DataType::Time64(TimeUnit::Microsecond),
4215            DataType::Duration(TimeUnit::Second),
4216            DataType::Interval(IntervalUnit::YearMonth),
4217            DataType::Interval(IntervalUnit::DayTime),
4218            DataType::Interval(IntervalUnit::MonthDayNano),
4219            DataType::Decimal32(9, 2),
4220            DataType::Decimal64(18, 2),
4221            DataType::Decimal128(38, 2),
4222            DataType::Decimal256(76, 2),
4223        ];
4224
4225        // For each type:
4226        // - write a batch with [skipped_column, values]
4227        // - read only the second column
4228        // - verify the result is correct
4229        for data_type in data_types {
4230            let skipped = make_array_for_type(data_type.clone());
4231            let values = Int32Array::from(vec![10, 20, 30]);
4232
4233            let schema = Arc::new(Schema::new(vec![
4234                Field::new("skipped", data_type, false),
4235                Field::new("values", DataType::Int32, false),
4236            ]));
4237
4238            let batch =
4239                RecordBatch::try_new(schema, vec![skipped, Arc::new(values.clone())]).unwrap();
4240
4241            // Serialize the batch into IPC format
4242            let mut buf = Vec::new();
4243            {
4244                let mut writer = FileWriter::try_new(&mut buf, &batch.schema()).unwrap();
4245                writer.write(&batch).unwrap();
4246                writer.finish().unwrap();
4247            }
4248
4249            // Read back only the second column (skip the first)
4250            let mut reader = FileReader::try_new(std::io::Cursor::new(buf), Some(vec![1])).unwrap();
4251            let read_batch = reader.next().unwrap().unwrap();
4252
4253            // Verify that the returned column matches the original values column
4254            assert_eq!(read_batch.num_columns(), 1);
4255            assert_eq!(read_batch.column(0).as_ref(), &values);
4256        }
4257    }
4258
4259    /// Builds a stream whose single message declares `body_length` but is followed by only
4260    /// `body_bytes` bytes of body. The body length is consumed before the header is
4261    /// interpreted, so the header itself is immaterial.
4262    fn stream_with_declared_body(body_length: i64, body_bytes: usize) -> Vec<u8> {
4263        let mut fbb = flatbuffers::FlatBufferBuilder::new();
4264        let mut message = crate::MessageBuilder::new(&mut fbb);
4265        message.add_version(crate::MetadataVersion::V5);
4266        message.add_header_type(crate::MessageHeader::NONE);
4267        message.add_bodyLength(body_length);
4268        let root = message.finish();
4269        fbb.finish(root, None);
4270        let metadata = fbb.finished_data();
4271
4272        let mut stream = Vec::new();
4273        stream.extend_from_slice(&CONTINUATION_MARKER);
4274        stream.extend_from_slice(&(metadata.len() as i32).to_le_bytes());
4275        stream.extend_from_slice(metadata);
4276        stream.resize(stream.len() + body_bytes, 0xAB);
4277        stream
4278    }
4279
4280    /// A message body length is read from the stream before any of the bytes it describes,
4281    /// so it cannot be trusted. Declaring an implausible one used to reserve it outright,
4282    /// and the resulting allocation failure aborts the process instead of surfacing an error
4283    /// the caller can handle.
4284    #[test]
4285    fn test_stream_reader_rejects_implausible_body_length() {
4286        for body_length in [i64::MAX, 1 << 50, -1] {
4287            let stream = stream_with_declared_body(body_length, 0);
4288            let err = StreamReader::try_new(std::io::Cursor::new(stream), None).expect_err(
4289                &format!("a message declaring {body_length} body bytes must not be accepted"),
4290            );
4291            let err = err.to_string();
4292            assert!(
4293                err.contains("Invalid IPC message body length")
4294                    || err.contains("Unexpected end of stream")
4295                    || err.contains("failed to fill whole buffer"),
4296                "unexpected error for body_length {body_length}: {err}"
4297            );
4298        }
4299    }
4300
4301    /// A plausible body length whose bytes end early must fail as a short read: before the
4302    /// first read for a small body, and after at least one growth step for a body larger
4303    /// than `MAX_PREALLOC_BYTES`.
4304    #[test]
4305    #[cfg_attr(miri, ignore)] // Takes too long
4306    fn test_stream_reader_rejects_truncated_body() {
4307        let over_prealloc = MAX_PREALLOC_BYTES as i64 + 1;
4308        for (body_length, body_bytes) in [(1024, 10), (over_prealloc, MAX_PREALLOC_BYTES)] {
4309            let stream = stream_with_declared_body(body_length, body_bytes);
4310            let err = StreamReader::try_new(std::io::Cursor::new(stream), None)
4311                .expect_err("a body backed by fewer bytes than declared must not be accepted");
4312            assert!(
4313                err.to_string().contains("failed to fill whole buffer"),
4314                "unexpected error for body_length {body_length}: {err}"
4315            );
4316        }
4317    }
4318
4319    /// The metadata length is untrusted for the same reason as the body length.
4320    #[test]
4321    fn test_stream_reader_rejects_unbacked_metadata_length() {
4322        let mut stream = Vec::new();
4323        stream.extend_from_slice(&CONTINUATION_MARKER);
4324        stream.extend_from_slice(&i32::MAX.to_le_bytes());
4325        stream.extend_from_slice(b"not this many bytes");
4326
4327        let err = StreamReader::try_new(std::io::Cursor::new(stream), None)
4328            .expect_err("metadata length must be backed by the stream");
4329        assert!(
4330            err.to_string().contains("Unexpected end of stream"),
4331            "unexpected error: {err}"
4332        );
4333    }
4334}