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