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("run_ends", DataType::Int16, false)),
2054            Arc::new(Field::new("values", DataType::Int32, true)),
2055        );
2056
2057        // define schema
2058        Schema::new(vec![
2059            Field::new("f0", DataType::UInt32, false),
2060            Field::new("f1", DataType::Utf8, false),
2061            Field::new("f2", DataType::Boolean, false),
2062            Field::new("f3", union_data_type, true),
2063            Field::new("f4", DataType::Null, true),
2064            Field::new("f5", DataType::Float64, true),
2065            Field::new("f6", list_data_type, false),
2066            Field::new("f7", DataType::FixedSizeBinary(3), true),
2067            Field::new("f8", fixed_size_list_data_type, false),
2068            Field::new("f9", struct_data_type, false),
2069            Field::new("f10", run_encoded_data_type, false),
2070            Field::new("f11", DataType::Boolean, false),
2071            Field::new_dictionary("f12", DataType::Int8, DataType::Utf8, false),
2072            Field::new("f13", DataType::Utf8, false),
2073        ])
2074    }
2075
2076    fn create_test_projection_batch_data(schema: &Schema) -> RecordBatch {
2077        // set test data for each column
2078        let array0 = UInt32Array::from(vec![1, 2, 3]);
2079        let array1 = StringArray::from(vec!["foo", "bar", "baz"]);
2080        let array2 = BooleanArray::from(vec![true, false, true]);
2081
2082        let mut union_builder = UnionBuilder::new_dense();
2083        union_builder.append::<Int32Type>("a", 1).unwrap();
2084        union_builder.append::<Float64Type>("b", 10.1).unwrap();
2085        union_builder.append_null::<Float64Type>("b").unwrap();
2086        let array3 = union_builder.build().unwrap();
2087
2088        let array4 = NullArray::new(3);
2089        let array5 = Float64Array::from(vec![Some(1.1), None, Some(3.3)]);
2090        let array6_values = vec![
2091            Some(vec![Some(10), Some(10), Some(10)]),
2092            Some(vec![Some(20), Some(20), Some(20)]),
2093            Some(vec![Some(30), Some(30)]),
2094        ];
2095        let array6 = ListArray::from_iter_primitive::<Int32Type, _, _>(array6_values);
2096        let array7_values = vec![vec![11, 12, 13], vec![22, 23, 24], vec![33, 34, 35]];
2097        let array7 = FixedSizeBinaryArray::try_from_iter(array7_values.into_iter()).unwrap();
2098
2099        let array8_values = ArrayData::builder(DataType::Int32)
2100            .len(9)
2101            .add_buffer(Buffer::from_slice_ref([40, 41, 42, 43, 44, 45, 46, 47, 48]))
2102            .build()
2103            .unwrap();
2104        let array8_data = ArrayData::builder(schema.field(8).data_type().clone())
2105            .len(3)
2106            .add_child_data(array8_values)
2107            .build()
2108            .unwrap();
2109        let array8 = FixedSizeListArray::from(array8_data);
2110
2111        let array9_id: ArrayRef = Arc::new(Int32Array::from(vec![1001, 1002, 1003]));
2112        let array9_list: ArrayRef =
2113            Arc::new(ListArray::from_iter_primitive::<Int8Type, _, _>(vec![
2114                Some(vec![Some(-10)]),
2115                Some(vec![Some(-20), Some(-20), Some(-20)]),
2116                Some(vec![Some(-30)]),
2117            ]));
2118        let array9 = ArrayDataBuilder::new(schema.field(9).data_type().clone())
2119            .add_child_data(array9_id.into_data())
2120            .add_child_data(array9_list.into_data())
2121            .len(3)
2122            .build()
2123            .unwrap();
2124        let array9 = StructArray::from(array9);
2125
2126        let array10_input = vec![Some(1_i32), None, None];
2127        let mut array10_builder = PrimitiveRunBuilder::<Int16Type, Int32Type>::new();
2128        array10_builder.extend(array10_input);
2129        let array10 = array10_builder.finish();
2130
2131        let array11 = BooleanArray::from(vec![false, false, true]);
2132
2133        let array12_values = StringArray::from(vec!["x", "yy", "zzz"]);
2134        let array12_keys = Int8Array::from_iter_values([1, 1, 2]);
2135        let array12 = DictionaryArray::new(array12_keys, Arc::new(array12_values));
2136
2137        let array13 = StringArray::from(vec!["a", "bb", "ccc"]);
2138
2139        // create record batch
2140        RecordBatch::try_new(
2141            Arc::new(schema.clone()),
2142            vec![
2143                Arc::new(array0),
2144                Arc::new(array1),
2145                Arc::new(array2),
2146                Arc::new(array3),
2147                Arc::new(array4),
2148                Arc::new(array5),
2149                Arc::new(array6),
2150                Arc::new(array7),
2151                Arc::new(array8),
2152                Arc::new(array9),
2153                Arc::new(array10),
2154                Arc::new(array11),
2155                Arc::new(array12),
2156                Arc::new(array13),
2157            ],
2158        )
2159        .unwrap()
2160    }
2161
2162    #[test]
2163    fn test_negative_meta_len_start_stream() {
2164        let bytes = i32::to_le_bytes(-1);
2165        let mut buf = vec![];
2166        buf.extend(CONTINUATION_MARKER);
2167        buf.extend(bytes);
2168
2169        let reader_err = StreamReader::try_new(Cursor::new(buf), None).err();
2170        assert!(reader_err.is_some());
2171        assert_eq!(
2172            reader_err.unwrap().to_string(),
2173            "Parser error: Invalid metadata length: -1"
2174        );
2175    }
2176
2177    #[test]
2178    fn test_negative_meta_len_mid_stream() {
2179        let schema = Schema::new(vec![Field::new("a", DataType::Int32, false)]);
2180        let mut buf = Vec::new();
2181        {
2182            let mut writer = crate::writer::StreamWriter::try_new(&mut buf, &schema).unwrap();
2183            let batch =
2184                RecordBatch::try_new(Arc::new(schema), vec![Arc::new(Int32Array::from(vec![1]))])
2185                    .unwrap();
2186            writer.write(&batch).unwrap();
2187        }
2188
2189        let bytes = i32::to_le_bytes(-1);
2190        buf.extend(CONTINUATION_MARKER);
2191        buf.extend(bytes);
2192
2193        let mut reader = StreamReader::try_new(Cursor::new(buf), None).unwrap();
2194        // Read the valid value
2195        assert!(reader.maybe_next().is_ok());
2196        // Read the invalid meta len
2197        let batch_err = reader.maybe_next().err();
2198        assert!(batch_err.is_some());
2199        assert_eq!(
2200            batch_err.unwrap().to_string(),
2201            "Parser error: Invalid metadata length: -1"
2202        );
2203    }
2204
2205    #[test]
2206    fn test_invalid_dictionary_batch_without_data() {
2207        use crate::r#gen::Message::*;
2208        use flatbuffers::FlatBufferBuilder;
2209
2210        let schema = Schema::new(vec![Field::new(
2211            "col",
2212            DataType::Dictionary(Box::new(DataType::Int8), Box::new(DataType::Utf8)),
2213            true,
2214        )]);
2215
2216        // DictionaryBatch.data is optional in the flatbuffer grammar and
2217        // required by the format
2218        let mut fbb = FlatBufferBuilder::new();
2219        let batch_offset = DictionaryBatch::create(
2220            &mut fbb,
2221            &DictionaryBatchArgs {
2222                id: 0,
2223                data: None,
2224                isDelta: false,
2225            },
2226        );
2227        fbb.finish_minimal(batch_offset);
2228        let batch_bytes = fbb.finished_data().to_vec();
2229        let batch = flatbuffers::root::<DictionaryBatch>(&batch_bytes).unwrap();
2230
2231        let data_buffer = Buffer::from(vec![0u8; 0]);
2232        let mut dictionaries: HashMap<i64, ArrayRef> = HashMap::new();
2233
2234        let err = read_dictionary(
2235            &data_buffer,
2236            batch,
2237            &schema,
2238            &mut dictionaries,
2239            &MetadataVersion::V5,
2240        )
2241        .unwrap_err();
2242
2243        assert_eq!(
2244            err.to_string(),
2245            "Parser error: Dictionary batch is missing its data"
2246        );
2247    }
2248
2249    #[test]
2250    fn test_missing_buffer_metadata_error() {
2251        use crate::r#gen::Message::*;
2252        use flatbuffers::FlatBufferBuilder;
2253
2254        let schema = Arc::new(Schema::new(vec![Field::new("col", DataType::Int32, true)]));
2255
2256        // create RecordBatch buffer metadata with invalid buffer count
2257        // Int32Array needs 2 buffers (validity + data) but we provide only 1
2258        let mut fbb = FlatBufferBuilder::new();
2259        let nodes = fbb.create_vector(&[FieldNode::new(2, 0)]);
2260        let buffers = fbb.create_vector(&[crate::Buffer::new(0, 8)]);
2261        let batch_offset = RecordBatch::create(
2262            &mut fbb,
2263            &RecordBatchArgs {
2264                length: 2,
2265                nodes: Some(nodes),
2266                buffers: Some(buffers),
2267                compression: None,
2268                variadicBufferCounts: None,
2269            },
2270        );
2271        fbb.finish_minimal(batch_offset);
2272        let batch_bytes = fbb.finished_data().to_vec();
2273        let batch = flatbuffers::root::<RecordBatch>(&batch_bytes).unwrap();
2274
2275        let data_buffer = Buffer::from(vec![0u8; 8]);
2276        let dictionaries: HashMap<i64, ArrayRef> = HashMap::new();
2277        let metadata = MetadataVersion::V5;
2278
2279        let decoder = RecordBatchDecoder::try_new(
2280            &data_buffer,
2281            batch,
2282            schema.clone(),
2283            &dictionaries,
2284            &metadata,
2285        )
2286        .unwrap();
2287
2288        let result = decoder.read_record_batch();
2289
2290        match result {
2291            Err(ArrowError::IpcError(msg)) => {
2292                assert_eq!(msg, "Buffer count mismatched with metadata");
2293            }
2294            other => panic!("unexpected error: {other:?}"),
2295        }
2296    }
2297
2298    /// A `Utf8View` batch whose variadic buffer count is `count`, with two buffers
2299    /// in the message. Returns the error from reading it with the given projection.
2300    fn read_batch_with_variadic_count(count: i64, projection: Option<&[usize]>) -> ArrowError {
2301        use crate::r#gen::Message::*;
2302        use flatbuffers::FlatBufferBuilder;
2303
2304        let schema = Arc::new(Schema::new(vec![Field::new(
2305            "col",
2306            DataType::Utf8View,
2307            true,
2308        )]));
2309
2310        let mut fbb = FlatBufferBuilder::new();
2311        let nodes = fbb.create_vector(&[FieldNode::new(1, 0)]);
2312        let buffers = fbb.create_vector(&[crate::Buffer::new(0, 8), crate::Buffer::new(8, 8)]);
2313        let variadic_buffer_counts = fbb.create_vector(&[count]);
2314        let batch_offset = RecordBatch::create(
2315            &mut fbb,
2316            &RecordBatchArgs {
2317                length: 1,
2318                nodes: Some(nodes),
2319                buffers: Some(buffers),
2320                compression: None,
2321                variadicBufferCounts: Some(variadic_buffer_counts),
2322            },
2323        );
2324        fbb.finish_minimal(batch_offset);
2325        let batch_bytes = fbb.finished_data().to_vec();
2326        let batch = flatbuffers::root::<RecordBatch>(&batch_bytes).unwrap();
2327
2328        let data_buffer = Buffer::from(vec![0u8; 16]);
2329        let dictionaries: HashMap<i64, ArrayRef> = HashMap::new();
2330
2331        RecordBatchDecoder::try_new(
2332            &data_buffer,
2333            batch,
2334            schema,
2335            &dictionaries,
2336            &MetadataVersion::V5,
2337        )
2338        .unwrap()
2339        .with_projection(projection)
2340        .read_record_batch()
2341        .expect_err("should get error")
2342    }
2343
2344    /// A variadic count the message cannot honour used to panic while slicing the
2345    /// buffers it did not read, both when reading the column and when skipping it.
2346    #[test]
2347    fn test_invalid_variadic_buffer_count_error() {
2348        // -2 leaves no buffers at all, -1 leaves too few, and 1 asks for more than the
2349        // message has. The projection selects nothing, so the column is skipped instead.
2350        for count in [-2, -1, 1, i64::MAX] {
2351            for projection in [None, Some([].as_slice())] {
2352                let err = read_batch_with_variadic_count(count, projection);
2353                assert_eq!(
2354                    err.to_string(),
2355                    format!(
2356                        "Ipc error: Invalid variadic count {count} for Utf8View column, \
2357                         with 2 buffer(s) left in the message"
2358                    ),
2359                    "count {count}, projection {projection:?}"
2360                );
2361            }
2362        }
2363    }
2364
2365    /// The valid count for a message with two buffers is zero.
2366    #[test]
2367    fn test_valid_variadic_buffer_count_is_accepted() {
2368        let err = read_batch_with_variadic_count(0, None);
2369        assert!(!err.to_string().contains("Invalid variadic count"), "{err}");
2370    }
2371
2372    #[test]
2373    fn test_missing_footer_schema_error() {
2374        use crate::r#gen::File::{Footer, FooterArgs};
2375        use flatbuffers::FlatBufferBuilder;
2376
2377        // a footer that verifies but has no schema table. record batches present
2378        // (so the earlier ok_or_else passes) but schema absent, which used to panic.
2379        let mut fbb = FlatBufferBuilder::new();
2380        let record_batches = fbb.create_vector::<Block>(&[]);
2381        let footer = Footer::create(
2382            &mut fbb,
2383            &FooterArgs {
2384                version: MetadataVersion::V5,
2385                schema: None,
2386                dictionaries: None,
2387                recordBatches: Some(record_batches),
2388                custom_metadata: None,
2389            },
2390        );
2391        fbb.finish(footer, None);
2392        let footer_data = fbb.finished_data();
2393
2394        // assemble a minimal IPC file: magic header, footer, footer length, magic trailer
2395        let mut buf = Vec::new();
2396        buf.extend_from_slice(&crate::ARROW_MAGIC);
2397        buf.extend_from_slice(footer_data);
2398        buf.extend_from_slice(&(footer_data.len() as i32).to_le_bytes());
2399        buf.extend_from_slice(&crate::ARROW_MAGIC);
2400
2401        let err = FileReader::try_new(Cursor::new(buf), None)
2402            .expect_err("expected an error, not a panic");
2403        assert!(
2404            matches!(err, ArrowError::ParseError(_)),
2405            "expected ParseError, got {err:?}"
2406        );
2407    }
2408
2409    /// Test that the reader can read legacy files where empty list arrays were written with a 0-byte offsets buffer.
2410    #[test]
2411    fn test_read_legacy_empty_list_without_offsets_buffer() {
2412        use crate::r#gen::Message::*;
2413        use flatbuffers::FlatBufferBuilder;
2414
2415        let schema = Arc::new(Schema::new(vec![Field::new_list(
2416            "items",
2417            Field::new_list_field(DataType::Int32, true),
2418            true,
2419        )]));
2420
2421        // Legacy arrow-rs versions wrote empty offsets buffers for empty list arrays.
2422        // Keep reader compatibility with such files by accepting a 0-byte offsets buffer.
2423        let mut fbb = FlatBufferBuilder::new();
2424        let nodes = fbb.create_vector(&[
2425            FieldNode::new(0, 0), // list node
2426            FieldNode::new(0, 0), // child int32 node
2427        ]);
2428        let buffers = fbb.create_vector(&[
2429            crate::Buffer::new(0, 0), // list validity
2430            crate::Buffer::new(0, 0), // list offsets (legacy empty buffer)
2431            crate::Buffer::new(0, 0), // child validity
2432            crate::Buffer::new(0, 0), // child values
2433        ]);
2434        let batch_offset = RecordBatch::create(
2435            &mut fbb,
2436            &RecordBatchArgs {
2437                length: 0,
2438                nodes: Some(nodes),
2439                buffers: Some(buffers),
2440                compression: None,
2441                variadicBufferCounts: None,
2442            },
2443        );
2444        fbb.finish_minimal(batch_offset);
2445        let batch_bytes = fbb.finished_data().to_vec();
2446        let batch = flatbuffers::root::<RecordBatch>(&batch_bytes).unwrap();
2447
2448        let body = Buffer::from(Vec::<u8>::new());
2449        let dictionaries: HashMap<i64, ArrayRef> = HashMap::new();
2450        let metadata = MetadataVersion::V5;
2451
2452        let decoder =
2453            RecordBatchDecoder::try_new(&body, batch, schema.clone(), &dictionaries, &metadata)
2454                .unwrap();
2455
2456        let read_batch = decoder.read_record_batch().unwrap();
2457        assert_eq!(read_batch.num_rows(), 0);
2458
2459        let list = read_batch
2460            .column(0)
2461            .as_any()
2462            .downcast_ref::<ListArray>()
2463            .unwrap();
2464        assert_eq!(list.len(), 0);
2465        assert_eq!(list.values().len(), 0);
2466    }
2467
2468    /// Test that the reader can read legacy files where empty Utf8/Binary arrays were written with a 0-byte offsets buffer.
2469    #[test]
2470    fn test_read_legacy_empty_utf8_and_binary_without_offsets_buffer() {
2471        use crate::r#gen::Message::*;
2472        use flatbuffers::FlatBufferBuilder;
2473
2474        let schema = Arc::new(Schema::new(vec![
2475            Field::new("name", DataType::Utf8, true),
2476            Field::new("payload", DataType::Binary, true),
2477        ]));
2478
2479        // Legacy arrow-rs versions wrote empty offsets buffers for empty Utf8/Binary arrays.
2480        // Keep reader compatibility with such files by accepting 0-byte offsets buffers.
2481        let mut fbb = FlatBufferBuilder::new();
2482        let nodes = fbb.create_vector(&[
2483            FieldNode::new(0, 0), // utf8 node
2484            FieldNode::new(0, 0), // binary node
2485        ]);
2486        let buffers = fbb.create_vector(&[
2487            crate::Buffer::new(0, 0), // utf8 validity
2488            crate::Buffer::new(0, 0), // utf8 offsets (legacy empty buffer)
2489            crate::Buffer::new(0, 0), // utf8 values
2490            crate::Buffer::new(0, 0), // binary validity
2491            crate::Buffer::new(0, 0), // binary offsets (legacy empty buffer)
2492            crate::Buffer::new(0, 0), // binary values
2493        ]);
2494        let batch_offset = RecordBatch::create(
2495            &mut fbb,
2496            &RecordBatchArgs {
2497                length: 0,
2498                nodes: Some(nodes),
2499                buffers: Some(buffers),
2500                compression: None,
2501                variadicBufferCounts: None,
2502            },
2503        );
2504        fbb.finish_minimal(batch_offset);
2505        let batch_bytes = fbb.finished_data().to_vec();
2506        let batch = flatbuffers::root::<RecordBatch>(&batch_bytes).unwrap();
2507
2508        let body = Buffer::from(Vec::<u8>::new());
2509        let dictionaries: HashMap<i64, ArrayRef> = HashMap::new();
2510        let metadata = MetadataVersion::V5;
2511
2512        let decoder =
2513            RecordBatchDecoder::try_new(&body, batch, schema.clone(), &dictionaries, &metadata)
2514                .unwrap();
2515
2516        let read_batch = decoder.read_record_batch().unwrap();
2517        assert_eq!(read_batch.num_rows(), 0);
2518
2519        let utf8 = read_batch
2520            .column(0)
2521            .as_any()
2522            .downcast_ref::<StringArray>()
2523            .unwrap();
2524        assert_eq!(utf8.len(), 0);
2525        assert_eq!(utf8.value_offsets(), [0]);
2526
2527        let binary = read_batch
2528            .column(1)
2529            .as_any()
2530            .downcast_ref::<BinaryArray>()
2531            .unwrap();
2532        assert_eq!(binary.len(), 0);
2533        assert_eq!(binary.value_offsets(), [0]);
2534    }
2535
2536    #[test]
2537    fn test_projection_array_values() {
2538        // define schema
2539        let schema = create_test_projection_schema();
2540
2541        // create record batch with test data
2542        let batch = create_test_projection_batch_data(&schema);
2543
2544        // write record batch in IPC format
2545        let mut buf = Vec::new();
2546        {
2547            let mut writer = crate::writer::FileWriter::try_new(&mut buf, &schema).unwrap();
2548            writer.write(&batch).unwrap();
2549            writer.finish().unwrap();
2550        }
2551
2552        // read record batch with projection
2553        for index in 0..12 {
2554            let projection = vec![index];
2555            let reader = FileReader::try_new(std::io::Cursor::new(buf.clone()), Some(projection));
2556            let read_batch = reader.unwrap().next().unwrap().unwrap();
2557            let projected_column = read_batch.column(0);
2558            let expected_column = batch.column(index);
2559
2560            // check the projected column equals the expected column
2561            assert_eq!(projected_column.as_ref(), expected_column.as_ref());
2562        }
2563
2564        {
2565            // read record batch with reversed projection
2566            let reader =
2567                FileReader::try_new(std::io::Cursor::new(buf.clone()), Some(vec![3, 2, 1]));
2568            let read_batch = reader.unwrap().next().unwrap().unwrap();
2569            let expected_batch = batch.project(&[3, 2, 1]).unwrap();
2570            assert_eq!(read_batch, expected_batch);
2571        }
2572    }
2573
2574    #[test]
2575    fn test_file_reader_projected_schema_matches_batch_schema() {
2576        let schema = create_test_projection_schema();
2577        let batch = create_test_projection_batch_data(&schema);
2578
2579        let mut buf = Vec::new();
2580        {
2581            let mut writer = crate::writer::FileWriter::try_new(&mut buf, &schema).unwrap();
2582            writer.write(&batch).unwrap();
2583            writer.finish().unwrap();
2584        }
2585
2586        let projection = vec![3, 2, 1];
2587        let mut reader = FileReader::try_new(Cursor::new(buf), Some(projection)).unwrap();
2588        let reader_schema = RecordBatchReader::schema(&reader);
2589        let read_batch = reader.next().unwrap().unwrap();
2590
2591        assert_eq!(reader_schema, read_batch.schema());
2592    }
2593
2594    #[test]
2595    fn test_stream_reader_projected_schema_matches_batch_schema() {
2596        let schema = create_test_projection_schema();
2597        let batch = create_test_projection_batch_data(&schema);
2598
2599        let mut buf = Vec::new();
2600        {
2601            let mut writer = crate::writer::StreamWriter::try_new(&mut buf, &schema).unwrap();
2602            writer.write(&batch).unwrap();
2603            writer.finish().unwrap();
2604        }
2605
2606        let projection = vec![3, 2, 1];
2607        let mut reader = StreamReader::try_new(Cursor::new(buf), Some(projection)).unwrap();
2608        let reader_schema = RecordBatchReader::schema(&reader);
2609        let read_batch = reader.next().unwrap().unwrap();
2610
2611        assert_eq!(reader_schema, read_batch.schema());
2612    }
2613
2614    #[test]
2615    fn test_file_reader_rejects_invalid_projection() {
2616        let schema = create_test_projection_schema();
2617        let batch = create_test_projection_batch_data(&schema);
2618
2619        let mut buf = Vec::new();
2620        {
2621            let mut writer = crate::writer::FileWriter::try_new(&mut buf, &schema).unwrap();
2622            writer.write(&batch).unwrap();
2623            writer.finish().unwrap();
2624        }
2625
2626        let result = FileReader::try_new(Cursor::new(buf), Some(vec![schema.fields().len()]));
2627
2628        assert!(matches!(result, Err(ArrowError::SchemaError(_))));
2629    }
2630
2631    #[test]
2632    fn test_projection_duplicate_indices() {
2633        let schema = create_test_projection_schema();
2634        let batch = create_test_projection_batch_data(&schema);
2635
2636        // Write the batch to IPC
2637        let mut buf = Vec::new();
2638        {
2639            let mut writer = crate::writer::FileWriter::try_new(&mut buf, &schema).unwrap();
2640            writer.write(&batch).unwrap();
2641            writer.finish().unwrap();
2642        }
2643
2644        // Verify duplicate([1, 1]) and reordered([2, 0, 2]) projection indices
2645        for projection in [vec![1, 1], vec![2, 0, 2]] {
2646            let reader =
2647                FileReader::try_new(std::io::Cursor::new(buf.clone()), Some(projection.clone()));
2648            let read_batch = reader.unwrap().next().unwrap().unwrap();
2649
2650            let expected_batch = batch.project(&projection).unwrap();
2651            assert_eq!(read_batch, expected_batch);
2652        }
2653    }
2654
2655    #[test]
2656    fn test_arrow_single_float_row() {
2657        let schema = Schema::new(vec![
2658            Field::new("a", DataType::Float32, false),
2659            Field::new("b", DataType::Float32, false),
2660            Field::new("c", DataType::Int32, false),
2661            Field::new("d", DataType::Int32, false),
2662        ]);
2663        let arrays = vec![
2664            Arc::new(Float32Array::from(vec![1.23])) as ArrayRef,
2665            Arc::new(Float32Array::from(vec![-6.50])) as ArrayRef,
2666            Arc::new(Int32Array::from(vec![2])) as ArrayRef,
2667            Arc::new(Int32Array::from(vec![1])) as ArrayRef,
2668        ];
2669        let batch = RecordBatch::try_new(Arc::new(schema.clone()), arrays).unwrap();
2670        // create stream writer
2671        let mut file = tempfile::tempfile().unwrap();
2672        let mut stream_writer = crate::writer::StreamWriter::try_new(&mut file, &schema).unwrap();
2673        stream_writer.write(&batch).unwrap();
2674        stream_writer.finish().unwrap();
2675
2676        drop(stream_writer);
2677
2678        file.rewind().unwrap();
2679
2680        // read stream back
2681        let reader = StreamReader::try_new(&mut file, None).unwrap();
2682
2683        reader.for_each(|batch| {
2684            let batch = batch.unwrap();
2685            assert!(
2686                batch
2687                    .column(0)
2688                    .as_any()
2689                    .downcast_ref::<Float32Array>()
2690                    .unwrap()
2691                    .value(0)
2692                    != 0.0
2693            );
2694            assert!(
2695                batch
2696                    .column(1)
2697                    .as_any()
2698                    .downcast_ref::<Float32Array>()
2699                    .unwrap()
2700                    .value(0)
2701                    != 0.0
2702            );
2703        });
2704
2705        file.rewind().unwrap();
2706
2707        // Read with projection
2708        let reader = StreamReader::try_new(file, Some(vec![0, 3])).unwrap();
2709
2710        reader.for_each(|batch| {
2711            let batch = batch.unwrap();
2712            assert_eq!(batch.schema().fields().len(), 2);
2713            assert_eq!(batch.schema().fields()[0].data_type(), &DataType::Float32);
2714            assert_eq!(batch.schema().fields()[1].data_type(), &DataType::Int32);
2715        });
2716    }
2717
2718    /// Write the record batch to an in-memory buffer in IPC File format
2719    fn write_ipc(rb: &RecordBatch) -> Vec<u8> {
2720        let mut buf = Vec::new();
2721        let mut writer = crate::writer::FileWriter::try_new(&mut buf, rb.schema_ref()).unwrap();
2722        writer.write(rb).unwrap();
2723        writer.finish().unwrap();
2724        buf
2725    }
2726
2727    /// Return the first record batch read from the IPC File buffer
2728    fn read_ipc(buf: &[u8]) -> Result<RecordBatch, ArrowError> {
2729        let mut reader = FileReader::try_new(std::io::Cursor::new(buf), None)?;
2730        reader.next().unwrap()
2731    }
2732
2733    /// Return the first record batch read from the IPC File buffer, disabling
2734    /// validation
2735    fn read_ipc_skip_validation(buf: &[u8]) -> Result<RecordBatch, ArrowError> {
2736        let mut reader = unsafe {
2737            FileReader::try_new(std::io::Cursor::new(buf), None)?.with_skip_validation(true)
2738        };
2739        reader.next().unwrap()
2740    }
2741
2742    fn roundtrip_ipc(rb: &RecordBatch) -> RecordBatch {
2743        let buf = write_ipc(rb);
2744        read_ipc(&buf).unwrap()
2745    }
2746
2747    /// Return the first record batch read from the IPC File buffer
2748    /// using the FileDecoder API
2749    fn read_ipc_with_decoder(buf: Vec<u8>) -> Result<RecordBatch, ArrowError> {
2750        read_ipc_with_decoder_inner(buf, false)
2751    }
2752
2753    /// Return the first record batch read from the IPC File buffer
2754    /// using the FileDecoder API, disabling validation
2755    fn read_ipc_with_decoder_skip_validation(buf: Vec<u8>) -> Result<RecordBatch, ArrowError> {
2756        read_ipc_with_decoder_inner(buf, true)
2757    }
2758
2759    fn read_ipc_with_decoder_inner(
2760        buf: Vec<u8>,
2761        skip_validation: bool,
2762    ) -> Result<RecordBatch, ArrowError> {
2763        let buffer = Buffer::from_vec(buf);
2764        let trailer_start = buffer.len() - 10;
2765        let footer_len = read_footer_length(buffer[trailer_start..].try_into().unwrap())?;
2766        let footer = root_as_footer(&buffer[trailer_start - footer_len..trailer_start])
2767            .map_err(|e| ArrowError::InvalidArgumentError(format!("Invalid footer: {e}")))?;
2768
2769        let schema = try_fb_to_schema(footer.schema().unwrap()).unwrap();
2770
2771        let mut decoder = unsafe {
2772            FileDecoder::new(Arc::new(schema), footer.version())
2773                .with_skip_validation(skip_validation)
2774        };
2775        // Read dictionaries
2776        for block in footer.dictionaries().iter().flatten() {
2777            let block_len = block.bodyLength() as usize + block.metaDataLength() as usize;
2778            let data = buffer.slice_with_length(block.offset() as _, block_len);
2779            decoder.read_dictionary(block, &data)?
2780        }
2781
2782        // Read record batch
2783        let batches = footer.recordBatches().unwrap();
2784        assert_eq!(batches.len(), 1); // Only wrote a single batch
2785
2786        let block = batches.get(0);
2787        let block_len = block.bodyLength() as usize + block.metaDataLength() as usize;
2788        let data = buffer.slice_with_length(block.offset() as _, block_len);
2789        Ok(decoder.read_record_batch(block, &data)?.unwrap())
2790    }
2791
2792    /// Write the record batch to an in-memory buffer in IPC Stream format
2793    fn write_stream(rb: &RecordBatch) -> Vec<u8> {
2794        let mut buf = Vec::new();
2795        let mut writer = crate::writer::StreamWriter::try_new(&mut buf, rb.schema_ref()).unwrap();
2796        writer.write(rb).unwrap();
2797        writer.finish().unwrap();
2798        buf
2799    }
2800
2801    /// Return the first record batch read from the IPC Stream buffer
2802    fn read_stream(buf: &[u8]) -> Result<RecordBatch, ArrowError> {
2803        let mut reader = StreamReader::try_new(std::io::Cursor::new(buf), None)?;
2804        reader.next().unwrap()
2805    }
2806
2807    /// Return the first record batch read from the IPC Stream buffer,
2808    /// disabling validation
2809    fn read_stream_skip_validation(buf: &[u8]) -> Result<RecordBatch, ArrowError> {
2810        let mut reader = unsafe {
2811            StreamReader::try_new(std::io::Cursor::new(buf), None)?.with_skip_validation(true)
2812        };
2813        reader.next().unwrap()
2814    }
2815
2816    fn roundtrip_ipc_stream(rb: &RecordBatch) -> RecordBatch {
2817        let buf = write_stream(rb);
2818        read_stream(&buf).unwrap()
2819    }
2820
2821    #[test]
2822    fn test_roundtrip_with_custom_metadata() {
2823        let schema = Schema::new(vec![Field::new("dummy", DataType::Float64, false)]);
2824        let mut buf = Vec::new();
2825        let mut writer = crate::writer::FileWriter::try_new(&mut buf, &schema).unwrap();
2826        let mut test_metadata = HashMap::new();
2827        test_metadata.insert("abc".to_string(), "abc".to_string());
2828        test_metadata.insert("def".to_string(), "def".to_string());
2829        for (k, v) in &test_metadata {
2830            writer.write_metadata(k, v);
2831        }
2832        writer.finish().unwrap();
2833        drop(writer);
2834
2835        let reader = crate::reader::FileReader::try_new(std::io::Cursor::new(buf), None).unwrap();
2836        assert_eq!(reader.custom_metadata(), &test_metadata);
2837    }
2838
2839    #[test]
2840    fn test_roundtrip_nested_dict() {
2841        let inner: DictionaryArray<Int32Type> = vec!["a", "b", "a"].into_iter().collect();
2842
2843        let array = Arc::new(inner) as ArrayRef;
2844
2845        let dctfield = Arc::new(Field::new("dict", array.data_type().clone(), false));
2846
2847        let s = StructArray::from(vec![(dctfield, array)]);
2848        let struct_array = Arc::new(s) as ArrayRef;
2849
2850        let schema = Arc::new(Schema::new(vec![Field::new(
2851            "struct",
2852            struct_array.data_type().clone(),
2853            false,
2854        )]));
2855
2856        let batch = RecordBatch::try_new(schema, vec![struct_array]).unwrap();
2857
2858        assert_eq!(batch, roundtrip_ipc(&batch));
2859    }
2860
2861    #[test]
2862    fn test_roundtrip_nested_dict_no_preserve_dict_id() {
2863        let inner: DictionaryArray<Int32Type> = vec!["a", "b", "a"].into_iter().collect();
2864
2865        let array = Arc::new(inner) as ArrayRef;
2866
2867        let dctfield = Arc::new(Field::new("dict", array.data_type().clone(), false));
2868
2869        let s = StructArray::from(vec![(dctfield, array)]);
2870        let struct_array = Arc::new(s) as ArrayRef;
2871
2872        let schema = Arc::new(Schema::new(vec![Field::new(
2873            "struct",
2874            struct_array.data_type().clone(),
2875            false,
2876        )]));
2877
2878        let batch = RecordBatch::try_new(schema, vec![struct_array]).unwrap();
2879
2880        let mut buf = Vec::new();
2881        let mut writer = crate::writer::FileWriter::try_new_with_options(
2882            &mut buf,
2883            batch.schema_ref(),
2884            IpcWriteOptions::default(),
2885        )
2886        .unwrap();
2887        writer.write(&batch).unwrap();
2888        writer.finish().unwrap();
2889        drop(writer);
2890
2891        let mut reader = FileReader::try_new(std::io::Cursor::new(buf), None).unwrap();
2892
2893        assert_eq!(batch, reader.next().unwrap().unwrap());
2894    }
2895
2896    fn check_union_with_builder(mut builder: UnionBuilder) {
2897        builder.append::<Int32Type>("a", 1).unwrap();
2898        builder.append_null::<Int32Type>("a").unwrap();
2899        builder.append::<Float64Type>("c", 3.0).unwrap();
2900        builder.append::<Int32Type>("a", 4).unwrap();
2901        builder.append::<Int64Type>("d", 11).unwrap();
2902        let union = builder.build().unwrap();
2903
2904        let schema = Arc::new(Schema::new(vec![Field::new(
2905            "union",
2906            union.data_type().clone(),
2907            false,
2908        )]));
2909
2910        let union_array = Arc::new(union) as ArrayRef;
2911
2912        let rb = RecordBatch::try_new(schema, vec![union_array]).unwrap();
2913        let rb2 = roundtrip_ipc(&rb);
2914        // TODO: equality not yet implemented for union, so we check that the length of the array is
2915        // the same and that all of the buffers are the same instead.
2916        assert_eq!(rb.schema(), rb2.schema());
2917        assert_eq!(rb.num_columns(), rb2.num_columns());
2918        assert_eq!(rb.num_rows(), rb2.num_rows());
2919        let union1 = rb.column(0);
2920        let union2 = rb2.column(0);
2921
2922        assert_eq!(union1, union2);
2923    }
2924
2925    #[test]
2926    fn test_roundtrip_dense_union() {
2927        check_union_with_builder(UnionBuilder::new_dense());
2928    }
2929
2930    #[test]
2931    fn test_roundtrip_sparse_union() {
2932        check_union_with_builder(UnionBuilder::new_sparse());
2933    }
2934
2935    #[test]
2936    fn test_roundtrip_struct_empty_fields() {
2937        let nulls = NullBuffer::from(&[true, true, false]);
2938        let rb = RecordBatch::try_from_iter([(
2939            "",
2940            Arc::new(StructArray::new_empty_fields(nulls.len(), Some(nulls))) as _,
2941        )])
2942        .unwrap();
2943        let rb2 = roundtrip_ipc(&rb);
2944        assert_eq!(rb, rb2);
2945    }
2946
2947    #[test]
2948    fn test_roundtrip_stream_run_array_sliced() {
2949        let run_array_1: Int32RunArray = vec!["a", "a", "a", "b", "b", "c", "c", "c"]
2950            .into_iter()
2951            .collect();
2952        let run_array_1_sliced = run_array_1.slice(2, 5);
2953
2954        let run_array_2_inupt = vec![Some(1_i32), None, None, Some(2), Some(2)];
2955        let mut run_array_2_builder = PrimitiveRunBuilder::<Int16Type, Int32Type>::new();
2956        run_array_2_builder.extend(run_array_2_inupt);
2957        let run_array_2 = run_array_2_builder.finish();
2958
2959        let schema = Arc::new(Schema::new(vec![
2960            Field::new(
2961                "run_array_1_sliced",
2962                run_array_1_sliced.data_type().clone(),
2963                false,
2964            ),
2965            Field::new("run_array_2", run_array_2.data_type().clone(), false),
2966        ]));
2967        let input_batch = RecordBatch::try_new(
2968            schema,
2969            vec![Arc::new(run_array_1_sliced.clone()), Arc::new(run_array_2)],
2970        )
2971        .unwrap();
2972        let output_batch = roundtrip_ipc_stream(&input_batch);
2973
2974        // As partial comparison not yet supported for run arrays, the sliced run array
2975        // has to be unsliced before comparing with the output. the second run array
2976        // can be compared as such.
2977        assert_eq!(input_batch.column(1), output_batch.column(1));
2978
2979        let run_array_1_unsliced = unslice_run_array(run_array_1_sliced.into_data()).unwrap();
2980        assert_eq!(run_array_1_unsliced, output_batch.column(0).into_data());
2981    }
2982
2983    #[test]
2984    fn test_roundtrip_stream_nested_dict() {
2985        let xs = vec!["AA", "BB", "AA", "CC", "BB"];
2986        let dict = Arc::new(
2987            xs.clone()
2988                .into_iter()
2989                .collect::<DictionaryArray<Int8Type>>(),
2990        );
2991        let string_array: ArrayRef = Arc::new(StringArray::from(xs.clone()));
2992        let struct_array = StructArray::from(vec![
2993            (
2994                Arc::new(Field::new("f2.1", DataType::Utf8, false)),
2995                string_array,
2996            ),
2997            (
2998                Arc::new(Field::new("f2.2_struct", dict.data_type().clone(), false)),
2999                dict.clone() as ArrayRef,
3000            ),
3001        ]);
3002        let schema = Arc::new(Schema::new(vec![
3003            Field::new("f1_string", DataType::Utf8, false),
3004            Field::new("f2_struct", struct_array.data_type().clone(), false),
3005        ]));
3006        let input_batch = RecordBatch::try_new(
3007            schema,
3008            vec![
3009                Arc::new(StringArray::from(xs.clone())),
3010                Arc::new(struct_array),
3011            ],
3012        )
3013        .unwrap();
3014        let output_batch = roundtrip_ipc_stream(&input_batch);
3015        assert_eq!(input_batch, output_batch);
3016    }
3017
3018    #[test]
3019    fn test_ipc_writers_reject_dictionary_of_dictionary_schema() {
3020        let values = Arc::new(StringArray::from(vec![Some("a"), Some("b")])) as ArrayRef;
3021        let inner = Arc::new(DictionaryArray::new(
3022            UInt32Array::from_iter_values([0, 1]),
3023            values,
3024        )) as ArrayRef;
3025        let outer = Arc::new(DictionaryArray::new(
3026            UInt32Array::from_iter_values([0, 1, 0]),
3027            inner,
3028        )) as ArrayRef;
3029
3030        let schema = Arc::new(Schema::new(vec![Field::new(
3031            "f1",
3032            outer.data_type().clone(),
3033            false,
3034        )]));
3035        let batch = RecordBatch::try_new(schema, vec![outer]).unwrap();
3036
3037        let mut stream = Vec::new();
3038        let Err(err) = crate::writer::StreamWriter::try_new(&mut stream, batch.schema_ref()) else {
3039            panic!("IPC stream writer should reject dictionary-of-dictionary schemas");
3040        };
3041        assert!(stream.is_empty());
3042
3043        assert!(
3044            err.to_string().contains("dictionary-of-dictionary values"),
3045            "unexpected error: {err}"
3046        );
3047
3048        let mut file = Vec::new();
3049        let Err(err) = crate::writer::FileWriter::try_new(&mut file, batch.schema_ref()) else {
3050            panic!("IPC file writer should reject dictionary-of-dictionary schemas");
3051        };
3052        assert!(file.is_empty());
3053
3054        assert!(
3055            err.to_string().contains("dictionary-of-dictionary values"),
3056            "unexpected error: {err}"
3057        );
3058    }
3059
3060    #[test]
3061    fn test_roundtrip_stream_nested_dict_of_map_of_dict() {
3062        let values = StringArray::from(vec![Some("a"), None, Some("b"), Some("c")]);
3063        let values = Arc::new(values) as ArrayRef;
3064        let value_dict_keys = Int8Array::from_iter_values([0, 1, 1, 2, 3, 1]);
3065        let value_dict_array = DictionaryArray::new(value_dict_keys, values.clone());
3066
3067        let key_dict_keys = Int8Array::from_iter_values([0, 0, 2, 2, 2, 3]);
3068        let key_dict_array = DictionaryArray::new(key_dict_keys, values);
3069
3070        #[expect(deprecated)]
3071        let keys_field = Arc::new(Field::new_dict(
3072            Field::MAP_KEY_FIELD_DEFAULT_NAME,
3073            DataType::Dictionary(Box::new(DataType::Int8), Box::new(DataType::Utf8)),
3074            false,
3075            1,
3076            false,
3077        ));
3078        #[expect(deprecated)]
3079        let values_field = Arc::new(Field::new_dict(
3080            Field::MAP_VALUE_FIELD_DEFAULT_NAME,
3081            DataType::Dictionary(Box::new(DataType::Int8), Box::new(DataType::Utf8)),
3082            true,
3083            2,
3084            false,
3085        ));
3086        let entry_struct = StructArray::from(vec![
3087            (keys_field, make_array(key_dict_array.into_data())),
3088            (values_field, make_array(value_dict_array.into_data())),
3089        ]);
3090        let map_data_type = DataType::Map(
3091            Arc::new(Field::new(
3092                Field::MAP_ENTRIES_FIELD_DEFAULT_NAME,
3093                entry_struct.data_type().clone(),
3094                false,
3095            )),
3096            false,
3097        );
3098
3099        let entry_offsets = Buffer::from_slice_ref([0, 2, 4, 6]);
3100        let map_data = ArrayData::builder(map_data_type)
3101            .len(3)
3102            .add_buffer(entry_offsets)
3103            .add_child_data(entry_struct.into_data())
3104            .build()
3105            .unwrap();
3106        let map_array = MapArray::from(map_data);
3107
3108        let dict_keys = Int8Array::from_iter_values([0, 1, 1, 2, 2, 1]);
3109        let dict_dict_array = DictionaryArray::new(dict_keys, Arc::new(map_array));
3110
3111        let schema = Arc::new(Schema::new(vec![Field::new(
3112            "f1",
3113            dict_dict_array.data_type().clone(),
3114            false,
3115        )]));
3116        let input_batch = RecordBatch::try_new(schema, vec![Arc::new(dict_dict_array)]).unwrap();
3117        let output_batch = roundtrip_ipc_stream(&input_batch);
3118        assert_eq!(input_batch, output_batch);
3119    }
3120
3121    fn test_roundtrip_stream_dict_of_list_of_dict_impl<
3122        OffsetSize: OffsetSizeTrait,
3123        U: ArrowNativeType,
3124    >(
3125        list_data_type: DataType,
3126        offsets: &[U; 5],
3127    ) {
3128        let values = StringArray::from(vec![Some("a"), None, Some("c"), None]);
3129        let keys = Int8Array::from_iter_values([0, 0, 1, 2, 0, 1, 3]);
3130        let dict_array = DictionaryArray::new(keys, Arc::new(values));
3131        let dict_data = dict_array.to_data();
3132
3133        let value_offsets = Buffer::from_slice_ref(offsets);
3134
3135        let list_data = ArrayData::builder(list_data_type)
3136            .len(4)
3137            .add_buffer(value_offsets)
3138            .add_child_data(dict_data)
3139            .build()
3140            .unwrap();
3141        let list_array = GenericListArray::<OffsetSize>::from(list_data);
3142
3143        let keys_for_dict = Int8Array::from_iter_values([0, 3, 0, 1, 1, 2, 0, 1, 3]);
3144        let dict_dict_array = DictionaryArray::new(keys_for_dict, Arc::new(list_array));
3145
3146        let schema = Arc::new(Schema::new(vec![Field::new(
3147            "f1",
3148            dict_dict_array.data_type().clone(),
3149            false,
3150        )]));
3151        let input_batch = RecordBatch::try_new(schema, vec![Arc::new(dict_dict_array)]).unwrap();
3152        let output_batch = roundtrip_ipc_stream(&input_batch);
3153        assert_eq!(input_batch, output_batch);
3154    }
3155
3156    #[test]
3157    fn test_roundtrip_stream_dict_of_list_of_dict() {
3158        // list
3159        #[expect(deprecated)]
3160        let list_data_type = DataType::List(Arc::new(Field::new_dict(
3161            "item",
3162            DataType::Dictionary(Box::new(DataType::Int8), Box::new(DataType::Utf8)),
3163            true,
3164            1,
3165            false,
3166        )));
3167        let offsets: &[i32; 5] = &[0, 2, 4, 4, 6];
3168        test_roundtrip_stream_dict_of_list_of_dict_impl::<i32, i32>(list_data_type, offsets);
3169
3170        // large list
3171        #[expect(deprecated)]
3172        let list_data_type = DataType::LargeList(Arc::new(Field::new_dict(
3173            "item",
3174            DataType::Dictionary(Box::new(DataType::Int8), Box::new(DataType::Utf8)),
3175            true,
3176            1,
3177            false,
3178        )));
3179        let offsets: &[i64; 5] = &[0, 2, 4, 4, 7];
3180        test_roundtrip_stream_dict_of_list_of_dict_impl::<i64, i64>(list_data_type, offsets);
3181    }
3182
3183    #[test]
3184    fn test_roundtrip_stream_dict_of_fixed_size_list_of_dict() {
3185        let values = StringArray::from(vec![Some("a"), None, Some("c"), None]);
3186        let keys = Int8Array::from_iter_values([0, 0, 1, 2, 0, 1, 3, 1, 2]);
3187        let dict_array = DictionaryArray::new(keys, Arc::new(values));
3188        let dict_data = dict_array.into_data();
3189
3190        #[expect(deprecated)]
3191        let list_data_type = DataType::FixedSizeList(
3192            Arc::new(Field::new_dict(
3193                "item",
3194                DataType::Dictionary(Box::new(DataType::Int8), Box::new(DataType::Utf8)),
3195                true,
3196                1,
3197                false,
3198            )),
3199            3,
3200        );
3201        let list_data = ArrayData::builder(list_data_type)
3202            .len(3)
3203            .add_child_data(dict_data)
3204            .build()
3205            .unwrap();
3206        let list_array = FixedSizeListArray::from(list_data);
3207
3208        let keys_for_dict = Int8Array::from_iter_values([0, 1, 0, 1, 1, 2, 0, 1, 2]);
3209        let dict_dict_array = DictionaryArray::new(keys_for_dict, Arc::new(list_array));
3210
3211        let schema = Arc::new(Schema::new(vec![Field::new(
3212            "f1",
3213            dict_dict_array.data_type().clone(),
3214            false,
3215        )]));
3216        let input_batch = RecordBatch::try_new(schema, vec![Arc::new(dict_dict_array)]).unwrap();
3217        let output_batch = roundtrip_ipc_stream(&input_batch);
3218        assert_eq!(input_batch, output_batch);
3219    }
3220
3221    const LONG_TEST_STRING: &str =
3222        "This is a long string to make sure binary view array handles it";
3223
3224    #[test]
3225    fn test_roundtrip_view_types() {
3226        let schema = Schema::new(vec![
3227            Field::new("field_1", DataType::BinaryView, true),
3228            Field::new("field_2", DataType::Utf8, true),
3229            Field::new("field_3", DataType::Utf8View, true),
3230        ]);
3231        let bin_values: Vec<Option<&[u8]>> = vec![
3232            Some(b"foo"),
3233            None,
3234            Some(b"bar"),
3235            Some(LONG_TEST_STRING.as_bytes()),
3236        ];
3237        let utf8_values: Vec<Option<&str>> =
3238            vec![Some("foo"), None, Some("bar"), Some(LONG_TEST_STRING)];
3239        let bin_view_array = BinaryViewArray::from_iter(bin_values);
3240        let utf8_array = StringArray::from_iter(utf8_values.iter());
3241        let utf8_view_array = StringViewArray::from_iter(utf8_values);
3242        let record_batch = RecordBatch::try_new(
3243            Arc::new(schema.clone()),
3244            vec![
3245                Arc::new(bin_view_array),
3246                Arc::new(utf8_array),
3247                Arc::new(utf8_view_array),
3248            ],
3249        )
3250        .unwrap();
3251
3252        assert_eq!(record_batch, roundtrip_ipc(&record_batch));
3253        assert_eq!(record_batch, roundtrip_ipc_stream(&record_batch));
3254
3255        let sliced_batch = record_batch.slice(1, 2);
3256        assert_eq!(sliced_batch, roundtrip_ipc(&sliced_batch));
3257        assert_eq!(sliced_batch, roundtrip_ipc_stream(&sliced_batch));
3258    }
3259
3260    #[test]
3261    fn test_roundtrip_view_types_nested_dict() {
3262        let bin_values: Vec<Option<&[u8]>> = vec![
3263            Some(b"foo"),
3264            None,
3265            Some(b"bar"),
3266            Some(LONG_TEST_STRING.as_bytes()),
3267            Some(b"field"),
3268        ];
3269        let utf8_values: Vec<Option<&str>> = vec![
3270            Some("foo"),
3271            None,
3272            Some("bar"),
3273            Some(LONG_TEST_STRING),
3274            Some("field"),
3275        ];
3276        let bin_view_array = Arc::new(BinaryViewArray::from_iter(bin_values));
3277        let utf8_view_array = Arc::new(StringViewArray::from_iter(utf8_values));
3278
3279        let key_dict_keys = Int8Array::from_iter_values([0, 0, 2, 2, 0, 2, 3]);
3280        let key_dict_array = DictionaryArray::new(key_dict_keys, utf8_view_array.clone());
3281        #[expect(deprecated)]
3282        let keys_field = Arc::new(Field::new_dict(
3283            Field::MAP_KEY_FIELD_DEFAULT_NAME,
3284            DataType::Dictionary(Box::new(DataType::Int8), Box::new(DataType::Utf8View)),
3285            false,
3286            1,
3287            false,
3288        ));
3289
3290        let value_dict_keys = Int8Array::from_iter_values([0, 3, 0, 1, 2, 0, 1]);
3291        let value_dict_array = DictionaryArray::new(value_dict_keys, bin_view_array);
3292        #[expect(deprecated)]
3293        let values_field = Arc::new(Field::new_dict(
3294            Field::MAP_VALUE_FIELD_DEFAULT_NAME,
3295            DataType::Dictionary(Box::new(DataType::Int8), Box::new(DataType::BinaryView)),
3296            true,
3297            2,
3298            false,
3299        ));
3300        let entry_struct = StructArray::from(vec![
3301            (keys_field, make_array(key_dict_array.into_data())),
3302            (values_field, make_array(value_dict_array.into_data())),
3303        ]);
3304
3305        let map_data_type = DataType::Map(
3306            Arc::new(Field::new(
3307                Field::MAP_ENTRIES_FIELD_DEFAULT_NAME,
3308                entry_struct.data_type().clone(),
3309                false,
3310            )),
3311            false,
3312        );
3313        let entry_offsets = Buffer::from_slice_ref([0, 2, 4, 7]);
3314        let map_data = ArrayData::builder(map_data_type)
3315            .len(3)
3316            .add_buffer(entry_offsets)
3317            .add_child_data(entry_struct.into_data())
3318            .build()
3319            .unwrap();
3320        let map_array = MapArray::from(map_data);
3321
3322        let dict_keys = Int8Array::from_iter_values([0, 1, 0, 1, 1, 2, 0, 1, 2]);
3323        let dict_dict_array = DictionaryArray::new(dict_keys, Arc::new(map_array));
3324        let schema = Arc::new(Schema::new(vec![Field::new(
3325            "f1",
3326            dict_dict_array.data_type().clone(),
3327            false,
3328        )]));
3329        let batch = RecordBatch::try_new(schema, vec![Arc::new(dict_dict_array)]).unwrap();
3330        assert_eq!(batch, roundtrip_ipc(&batch));
3331        assert_eq!(batch, roundtrip_ipc_stream(&batch));
3332
3333        let sliced_batch = batch.slice(1, 2);
3334        assert_eq!(sliced_batch, roundtrip_ipc(&sliced_batch));
3335        assert_eq!(sliced_batch, roundtrip_ipc_stream(&sliced_batch));
3336    }
3337
3338    #[test]
3339    fn test_no_columns_batch() {
3340        let schema = Arc::new(Schema::empty());
3341        let options = RecordBatchOptions::new()
3342            .with_match_field_names(true)
3343            .with_row_count(Some(10));
3344        let input_batch = RecordBatch::try_new_with_options(schema, vec![], &options).unwrap();
3345        let output_batch = roundtrip_ipc_stream(&input_batch);
3346        assert_eq!(input_batch, output_batch);
3347    }
3348
3349    #[test]
3350    fn test_unaligned() {
3351        let batch = RecordBatch::try_from_iter(vec![(
3352            "i32",
3353            Arc::new(Int32Array::from(vec![1, 2, 3, 4])) as _,
3354        )])
3355        .unwrap();
3356
3357        let r#gen = IpcDataGenerator {};
3358        let mut dict_tracker = DictionaryTracker::new(false);
3359        let (_, encoded) = r#gen
3360            .encode(
3361                &batch,
3362                &mut dict_tracker,
3363                &Default::default(),
3364                &mut Default::default(),
3365            )
3366            .unwrap();
3367
3368        let message = root_as_message(&encoded.ipc_message).unwrap();
3369
3370        // Construct an unaligned buffer
3371        let mut buffer = MutableBuffer::with_capacity(encoded.arrow_data.len() + 1);
3372        buffer.push(0_u8);
3373        buffer.extend_from_slice(&encoded.arrow_data);
3374        let b = Buffer::from(buffer).slice(1);
3375        assert_ne!(b.as_ptr().align_offset(8), 0);
3376
3377        let ipc_batch = message.header_as_record_batch().unwrap();
3378        let roundtrip = RecordBatchDecoder::try_new(
3379            &b,
3380            ipc_batch,
3381            batch.schema(),
3382            &Default::default(),
3383            &message.version(),
3384        )
3385        .unwrap()
3386        .with_require_alignment(false)
3387        .read_record_batch()
3388        .unwrap();
3389        assert_eq!(batch, roundtrip);
3390    }
3391
3392    #[test]
3393    fn test_unaligned_throws_error_with_require_alignment() {
3394        let batch = RecordBatch::try_from_iter(vec![(
3395            "i32",
3396            Arc::new(Int32Array::from(vec![1, 2, 3, 4])) as _,
3397        )])
3398        .unwrap();
3399
3400        let r#gen = IpcDataGenerator {};
3401        let mut dict_tracker = DictionaryTracker::new(false);
3402        let (_, encoded) = r#gen
3403            .encode(
3404                &batch,
3405                &mut dict_tracker,
3406                &Default::default(),
3407                &mut Default::default(),
3408            )
3409            .unwrap();
3410
3411        let message = root_as_message(&encoded.ipc_message).unwrap();
3412
3413        // Construct an unaligned buffer
3414        let mut buffer = MutableBuffer::with_capacity(encoded.arrow_data.len() + 1);
3415        buffer.push(0_u8);
3416        buffer.extend_from_slice(&encoded.arrow_data);
3417        let b = Buffer::from(buffer).slice(1);
3418        assert_ne!(b.as_ptr().align_offset(8), 0);
3419
3420        let ipc_batch = message.header_as_record_batch().unwrap();
3421        let result = RecordBatchDecoder::try_new(
3422            &b,
3423            ipc_batch,
3424            batch.schema(),
3425            &Default::default(),
3426            &message.version(),
3427        )
3428        .unwrap()
3429        .with_require_alignment(true)
3430        .read_record_batch();
3431
3432        let error = result.unwrap_err();
3433        assert_eq!(
3434            error.to_string(),
3435            "Invalid argument error: Misaligned buffers[0] in array of type Int32, \
3436             offset from expected alignment of 4 by 1"
3437        );
3438    }
3439
3440    /// Verify that misaligned IPC buffers are caught by `require_alignment = true`.
3441    ///
3442    /// For each array type we shift the IPC body by every byte offset in 0..OFFSET_RANGE and
3443    /// assert that the decoder errors exactly when the offset violates the type's
3444    /// minimum alignment requirement.  The Arrow columnar spec permits alignment to
3445    /// any multiple of 8 or 64 bytes, so multiples of 8 (which include every multiple
3446    /// of 64) must succeed; everything else must return a "Misaligned buffers" error.
3447    /// See <https://arrow.apache.org/docs/format/Columnar.html#buffer-alignment-and-padding>.
3448    #[test]
3449    fn test_misaligned_buffers_error() {
3450        const OFFSET_RANGE: usize = 128;
3451
3452        // (array, minimum required alignment in bytes)
3453        // Fixed-width: alignment == element byte width.
3454        // Variable-width (e.g. StringArray): the offsets buffer drives alignment (Int32 → 4).
3455        let cases: Vec<(ArrayRef, usize)> = vec![
3456            (Arc::new(Int32Array::from_iter(0i32..100)) as _, 4),
3457            (Arc::new(Int64Array::from_iter(0i64..100)) as _, 8),
3458            (
3459                Arc::new(StringArray::from_iter_values(
3460                    (0..100).map(|i| i.to_string()),
3461                )) as _,
3462                4,
3463            ),
3464            (
3465                Arc::new(LargeStringArray::from_iter_values(
3466                    (0..100).map(|i| i.to_string()),
3467                )) as _,
3468                8,
3469            ),
3470        ];
3471
3472        for (array, alignment) in cases {
3473            let batch = RecordBatch::try_from_iter(vec![("col", Arc::clone(&array))]).unwrap();
3474            let encoder = IpcDataGenerator {};
3475            let mut dict_tracker = DictionaryTracker::new(false);
3476            let (_, encoded) = encoder
3477                .encode(
3478                    &batch,
3479                    &mut dict_tracker,
3480                    &Default::default(),
3481                    &mut Default::default(),
3482                )
3483                .unwrap();
3484            let message = root_as_message(&encoded.ipc_message).unwrap();
3485            let ipc_batch = message.header_as_record_batch().unwrap();
3486
3487            for offset in 0..OFFSET_RANGE {
3488                // MutableBuffer always allocates at a 64-byte aligned base address.
3489                // Slicing by `offset` bytes yields a pointer at `base_ptr + offset`,
3490                // whose alignment is gcd(64, offset).  That makes `offset` the sole
3491                // determinant of the resulting alignment — without this guarantee the
3492                // test would be non-deterministic depending on what the allocator returns.
3493                let mut storage = MutableBuffer::with_capacity(encoded.arrow_data.len() + offset);
3494                for _ in 0..offset {
3495                    storage.push(0_u8);
3496                }
3497                storage.extend_from_slice(&encoded.arrow_data);
3498                let buf = Buffer::from(storage).slice(offset);
3499
3500                let result = RecordBatchDecoder::try_new(
3501                    &buf,
3502                    ipc_batch,
3503                    batch.schema(),
3504                    &Default::default(),
3505                    &message.version(),
3506                )
3507                .unwrap()
3508                .with_require_alignment(true)
3509                .read_record_batch();
3510
3511                if offset % alignment == 0 {
3512                    assert!(
3513                        result.is_ok(),
3514                        "type={} offset={offset}: expected Ok but got {:?}",
3515                        array.data_type(),
3516                        result.unwrap_err(),
3517                    );
3518                } else {
3519                    let err = result
3520                        .expect_err(&format!(
3521                            "type={} offset={offset}: expected Err for misaligned buffer",
3522                            array.data_type()
3523                        ))
3524                        .to_string();
3525                    assert!(
3526                        err.contains("Misaligned buffers"),
3527                        "type={} offset={offset}: unexpected error: {err}",
3528                        array.data_type(),
3529                    );
3530                }
3531            }
3532        }
3533    }
3534
3535    #[test]
3536    #[cfg_attr(miri, ignore)] // Takes too long
3537    fn test_file_with_massive_column_count() {
3538        // 499_999 is upper limit for default settings (1_000_000)
3539        let limit = 600_000;
3540
3541        let fields = (0..limit)
3542            .map(|i| Field::new(format!("{i}"), DataType::Boolean, false))
3543            .collect::<Vec<_>>();
3544        let schema = Arc::new(Schema::new(fields));
3545        let batch = RecordBatch::new_empty(schema);
3546
3547        let mut buf = Vec::new();
3548        let mut writer = crate::writer::FileWriter::try_new(&mut buf, batch.schema_ref()).unwrap();
3549        writer.write(&batch).unwrap();
3550        writer.finish().unwrap();
3551        drop(writer);
3552
3553        let mut reader = FileReaderBuilder::new()
3554            .with_max_footer_fb_tables(1_500_000)
3555            .build(std::io::Cursor::new(buf))
3556            .unwrap();
3557        let roundtrip_batch = reader.next().unwrap().unwrap();
3558
3559        assert_eq!(batch, roundtrip_batch);
3560    }
3561
3562    #[test]
3563    fn test_file_with_deeply_nested_columns() {
3564        // 60 is upper limit for default settings (64)
3565        let limit = 61;
3566
3567        let fields = (0..limit).fold(
3568            vec![Field::new("leaf", DataType::Boolean, false)],
3569            |field, index| vec![Field::new_struct(format!("{index}"), field, false)],
3570        );
3571        let schema = Arc::new(Schema::new(fields));
3572        let batch = RecordBatch::new_empty(schema);
3573
3574        let mut buf = Vec::new();
3575        let mut writer = crate::writer::FileWriter::try_new(&mut buf, batch.schema_ref()).unwrap();
3576        writer.write(&batch).unwrap();
3577        writer.finish().unwrap();
3578        drop(writer);
3579
3580        let mut reader = FileReaderBuilder::new()
3581            .with_max_footer_fb_depth(65)
3582            .build(std::io::Cursor::new(buf))
3583            .unwrap();
3584        let roundtrip_batch = reader.next().unwrap().unwrap();
3585
3586        assert_eq!(batch, roundtrip_batch);
3587    }
3588
3589    #[test]
3590    fn test_invalid_struct_array_ipc_read_errors() {
3591        let a_field = Field::new("a", DataType::Int32, false);
3592        let b_field = Field::new("b", DataType::Int32, false);
3593        let struct_fields = Fields::from(vec![a_field.clone(), b_field.clone()]);
3594
3595        let a_array_data = ArrayData::builder(a_field.data_type().clone())
3596            .len(4)
3597            .add_buffer(Buffer::from_slice_ref([1, 2, 3, 4]))
3598            .build()
3599            .unwrap();
3600        let b_array_data = ArrayData::builder(b_field.data_type().clone())
3601            .len(3)
3602            .add_buffer(Buffer::from_slice_ref([5, 6, 7]))
3603            .build()
3604            .unwrap();
3605
3606        let invalid_struct_arr = unsafe {
3607            StructArray::new_unchecked(
3608                struct_fields,
3609                vec![make_array(a_array_data), make_array(b_array_data)],
3610                None,
3611            )
3612        };
3613
3614        expect_ipc_validation_error(
3615            Arc::new(invalid_struct_arr),
3616            "Invalid argument error: Incorrect array length for StructArray field \"b\", expected 4 got 3",
3617        );
3618    }
3619
3620    #[test]
3621    fn test_invalid_nested_array_ipc_read_errors() {
3622        // one of the nested arrays has invalid data
3623        let a_field = Field::new("a", DataType::Int32, false);
3624        let b_field = Field::new("b", DataType::Utf8, false);
3625
3626        let schema = Arc::new(Schema::new(vec![Field::new_struct(
3627            "s",
3628            vec![a_field.clone(), b_field.clone()],
3629            false,
3630        )]));
3631
3632        let a_array_data = ArrayData::builder(a_field.data_type().clone())
3633            .len(4)
3634            .add_buffer(Buffer::from_slice_ref([1, 2, 3, 4]))
3635            .build()
3636            .unwrap();
3637        // invalid nested child array -- length is correct, but has invalid utf8 data
3638        let b_array_data = {
3639            let valid: &[u8] = b"   ";
3640            let mut invalid = vec![];
3641            invalid.extend_from_slice(b"ValidString");
3642            invalid.extend_from_slice(INVALID_UTF8_FIRST_CHAR);
3643            let binary_array =
3644                BinaryArray::from_iter(vec![None, Some(valid), None, Some(&invalid)]);
3645            let array = unsafe {
3646                StringArray::new_unchecked(
3647                    binary_array.offsets().clone(),
3648                    binary_array.values().clone(),
3649                    binary_array.nulls().cloned(),
3650                )
3651            };
3652            array.into_data()
3653        };
3654        let struct_data_type = schema.field(0).data_type();
3655
3656        let invalid_struct_arr = unsafe {
3657            make_array(
3658                ArrayData::builder(struct_data_type.clone())
3659                    .len(4)
3660                    .add_child_data(a_array_data)
3661                    .add_child_data(b_array_data)
3662                    .build_unchecked(),
3663            )
3664        };
3665        expect_ipc_validation_error(
3666            invalid_struct_arr,
3667            "Invalid argument error: Invalid UTF8 sequence at string index 3 (3..18): invalid utf-8 sequence of 1 bytes from index 11",
3668        );
3669    }
3670
3671    #[test]
3672    fn test_same_dict_id_without_preserve() {
3673        let batch = RecordBatch::try_new(
3674            Arc::new(Schema::new(
3675                ["a", "b"]
3676                    .iter()
3677                    .map(|name| {
3678                        #[expect(deprecated)]
3679                        Field::new_dict(
3680                            name.to_string(),
3681                            DataType::Dictionary(
3682                                Box::new(DataType::Int32),
3683                                Box::new(DataType::Utf8),
3684                            ),
3685                            true,
3686                            0,
3687                            false,
3688                        )
3689                    })
3690                    .collect::<Vec<Field>>(),
3691            )),
3692            vec![
3693                Arc::new(
3694                    vec![Some("c"), Some("d")]
3695                        .into_iter()
3696                        .collect::<DictionaryArray<Int32Type>>(),
3697                ) as ArrayRef,
3698                Arc::new(
3699                    vec![Some("e"), Some("f")]
3700                        .into_iter()
3701                        .collect::<DictionaryArray<Int32Type>>(),
3702                ) as ArrayRef,
3703            ],
3704        )
3705        .expect("Failed to create RecordBatch");
3706
3707        // serialize the record batch as an IPC stream
3708        let mut buf = vec![];
3709        {
3710            let mut writer = crate::writer::StreamWriter::try_new_with_options(
3711                &mut buf,
3712                batch.schema().as_ref(),
3713                crate::writer::IpcWriteOptions::default(),
3714            )
3715            .expect("Failed to create StreamWriter");
3716            writer.write(&batch).expect("Failed to write RecordBatch");
3717            writer.finish().expect("Failed to finish StreamWriter");
3718        }
3719
3720        StreamReader::try_new(std::io::Cursor::new(buf), None)
3721            .expect("Failed to create StreamReader")
3722            .for_each(|decoded_batch| {
3723                assert_eq!(decoded_batch.expect("Failed to read RecordBatch"), batch);
3724            });
3725    }
3726
3727    #[test]
3728    fn test_validation_of_invalid_list_array() {
3729        // ListArray with invalid offsets
3730        let array = unsafe {
3731            let values = Int32Array::from(vec![1, 2, 3]);
3732            let bad_offsets = ScalarBuffer::<i32>::from(vec![0, 2, 4, 2]); // offsets can't go backwards
3733            let offsets = OffsetBuffer::new_unchecked(bad_offsets); // INVALID array created
3734            let field = Field::new_list_field(DataType::Int32, true);
3735            let nulls = None;
3736            ListArray::new(Arc::new(field), offsets, Arc::new(values), nulls)
3737        };
3738
3739        expect_ipc_validation_error(
3740            Arc::new(array),
3741            "Invalid argument error: Offset invariant failure: offset at position 2 out of bounds: 4 > 2",
3742        );
3743    }
3744
3745    #[test]
3746    fn test_validation_of_invalid_string_array() {
3747        let valid: &[u8] = b"   ";
3748        let mut invalid = vec![];
3749        invalid.extend_from_slice(b"ThisStringIsCertainlyLongerThan12Bytes");
3750        invalid.extend_from_slice(INVALID_UTF8_FIRST_CHAR);
3751        let binary_array = BinaryArray::from_iter(vec![None, Some(valid), None, Some(&invalid)]);
3752        // data is not valid utf8 we can not construct a correct StringArray
3753        // safely, so purposely create an invalid StringArray
3754        let array = unsafe {
3755            StringArray::new_unchecked(
3756                binary_array.offsets().clone(),
3757                binary_array.values().clone(),
3758                binary_array.nulls().cloned(),
3759            )
3760        };
3761        expect_ipc_validation_error(
3762            Arc::new(array),
3763            "Invalid argument error: Invalid UTF8 sequence at string index 3 (3..45): invalid utf-8 sequence of 1 bytes from index 38",
3764        );
3765    }
3766
3767    #[test]
3768    fn test_validation_of_invalid_string_view_array() {
3769        let valid: &[u8] = b"   ";
3770        let mut invalid = vec![];
3771        invalid.extend_from_slice(b"ThisStringIsCertainlyLongerThan12Bytes");
3772        invalid.extend_from_slice(INVALID_UTF8_FIRST_CHAR);
3773        let binary_view_array =
3774            BinaryViewArray::from_iter(vec![None, Some(valid), None, Some(&invalid)]);
3775        // data is not valid utf8 we can not construct a correct StringArray
3776        // safely, so purposely create an invalid StringArray
3777        let array = unsafe {
3778            StringViewArray::new_unchecked(
3779                binary_view_array.views().clone(),
3780                Arc::clone(binary_view_array.data_buffers()),
3781                binary_view_array.nulls().cloned(),
3782            )
3783        };
3784        expect_ipc_validation_error(
3785            Arc::new(array),
3786            "Invalid argument error: Encountered non-UTF-8 data at index 3: invalid utf-8 sequence of 1 bytes from index 38",
3787        );
3788    }
3789
3790    /// return an invalid dictionary array (key is larger than values)
3791    /// ListArray with invalid offsets
3792    #[test]
3793    fn test_validation_of_invalid_dictionary_array() {
3794        let array = unsafe {
3795            let values = StringArray::from_iter_values(["a", "b", "c"]);
3796            let keys = Int32Array::from(vec![1, 200]); // keys are not valid for values
3797            DictionaryArray::new_unchecked(keys, Arc::new(values))
3798        };
3799
3800        expect_ipc_validation_error(
3801            Arc::new(array),
3802            "Invalid argument error: Value at position 1 out of bounds: 200 (should be in [0, 2])",
3803        );
3804    }
3805
3806    #[test]
3807    fn test_validation_of_invalid_union_array() {
3808        let array = unsafe {
3809            let fields = UnionFields::try_new(
3810                vec![1, 3], // typeids : type id 2 is not valid
3811                vec![
3812                    Field::new("a", DataType::Int32, false),
3813                    Field::new("b", DataType::Utf8, false),
3814                ],
3815            )
3816            .unwrap();
3817            let type_ids = ScalarBuffer::from(vec![1i8, 2, 3]); // 2 is invalid
3818            let offsets = None;
3819            let children: Vec<ArrayRef> = vec![
3820                Arc::new(Int32Array::from(vec![10, 20, 30])),
3821                Arc::new(StringArray::from(vec![Some("a"), Some("b"), Some("c")])),
3822            ];
3823
3824            UnionArray::new_unchecked(fields, type_ids, offsets, children)
3825        };
3826
3827        expect_ipc_validation_error(
3828            Arc::new(array),
3829            "Invalid argument error: Type Ids values must match one of the field type ids",
3830        );
3831    }
3832
3833    /// Invalid Utf-8 sequence in the first character
3834    /// <https://stackoverflow.com/questions/1301402/example-invalid-utf8-string>
3835    const INVALID_UTF8_FIRST_CHAR: &[u8] = &[0xa0, 0xa1, 0x20, 0x20];
3836
3837    /// Expect an error when reading the record batch using IPC or IPC Streams
3838    fn expect_ipc_validation_error(array: ArrayRef, expected_err: &str) {
3839        let rb = RecordBatch::try_from_iter([("a", array)]).unwrap();
3840
3841        // IPC Stream format
3842        let buf = write_stream(&rb); // write is ok
3843        read_stream_skip_validation(&buf).unwrap();
3844        let err = read_stream(&buf).unwrap_err();
3845        assert_eq!(err.to_string(), expected_err);
3846
3847        // IPC File format
3848        let buf = write_ipc(&rb); // write is ok
3849        read_ipc_skip_validation(&buf).unwrap();
3850        let err = read_ipc(&buf).unwrap_err();
3851        assert_eq!(err.to_string(), expected_err);
3852
3853        // IPC Format with FileDecoder
3854        read_ipc_with_decoder_skip_validation(buf.clone()).unwrap();
3855        let err = read_ipc_with_decoder(buf).unwrap_err();
3856        assert_eq!(err.to_string(), expected_err);
3857    }
3858
3859    #[test]
3860    fn test_roundtrip_schema() {
3861        let schema = Schema::new(vec![
3862            Field::new(
3863                "a",
3864                DataType::Dictionary(Box::new(DataType::UInt16), Box::new(DataType::Utf8)),
3865                false,
3866            ),
3867            Field::new(
3868                "b",
3869                DataType::Dictionary(Box::new(DataType::UInt16), Box::new(DataType::Utf8)),
3870                false,
3871            ),
3872        ]);
3873
3874        let options = IpcWriteOptions::default();
3875        let data_gen = IpcDataGenerator::default();
3876        let mut dict_tracker = DictionaryTracker::new(false);
3877        let encoded_data =
3878            data_gen.schema_to_bytes_with_dictionary_tracker(&schema, &mut dict_tracker, &options);
3879        let mut schema_bytes = vec![];
3880        write_message(&mut schema_bytes, encoded_data, &options).expect("write_message");
3881
3882        let begin_offset: usize = if schema_bytes[0..4].eq(&CONTINUATION_MARKER) {
3883            4
3884        } else {
3885            0
3886        };
3887
3888        size_prefixed_root_as_message(&schema_bytes[begin_offset..])
3889            .expect_err("size_prefixed_root_as_message");
3890
3891        let msg = parse_message(&schema_bytes).expect("parse_message");
3892        let ipc_schema = msg.header_as_schema().expect("header_as_schema");
3893        let new_schema = try_fb_to_schema(ipc_schema).unwrap();
3894
3895        assert_eq!(schema, new_schema);
3896    }
3897
3898    #[test]
3899    fn test_negative_meta_len() {
3900        let bytes = i32::to_le_bytes(-1);
3901        let mut buf = vec![];
3902        buf.extend(CONTINUATION_MARKER);
3903        buf.extend(bytes);
3904
3905        let reader = StreamReader::try_new(Cursor::new(buf), None);
3906        assert!(reader.is_err());
3907    }
3908
3909    /// Per the IPC specification, dictionary batches may be omitted for
3910    /// dictionary-encoded columns where all values are null.  The C++
3911    /// implementation relies on this and does not emit a dictionary batch
3912    /// in that case.  Verify that the Rust reader handles such streams
3913    /// by synthesizing an empty dictionary instead of returning an error.
3914    #[test]
3915    fn test_read_null_dict_without_dictionary_batch() {
3916        // Build an all-null dictionary-encoded column.
3917        let keys = Int32Array::new_null(4);
3918        let values: ArrayRef = new_empty_array(&DataType::Utf8);
3919        let dict_array = DictionaryArray::new(keys, values);
3920
3921        let schema = Arc::new(Schema::new(vec![Field::new(
3922            "d",
3923            dict_array.data_type().clone(),
3924            true,
3925        )]));
3926        let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(dict_array)]).unwrap();
3927
3928        // Write a normal IPC stream (which includes the dictionary batch).
3929        let full_stream = write_stream(&batch);
3930
3931        // Parse the stream into individual messages and reconstruct it
3932        // without the DictionaryBatch message, simulating what C++ emits
3933        // for an all-null dictionary column.
3934        let mut stripped = Vec::new();
3935        let mut cursor = Cursor::new(&full_stream);
3936        loop {
3937            // Each message is: [continuation (4 bytes)] [meta_len (4 bytes)]
3938            //                   [metadata (meta_len bytes)] [body (bodyLength bytes)]
3939            let mut header = [0u8; 4];
3940            if cursor.read_exact(&mut header).is_err() {
3941                break;
3942            }
3943            if header == CONTINUATION_MARKER && cursor.read_exact(&mut header).is_err() {
3944                break;
3945            }
3946            let meta_len = u32::from_le_bytes(header) as usize;
3947            if meta_len == 0 {
3948                // EOS marker — write it through.
3949                stripped.extend_from_slice(&CONTINUATION_MARKER);
3950                stripped.extend_from_slice(&0u32.to_le_bytes());
3951                break;
3952            }
3953            let mut meta_buf = vec![0u8; meta_len];
3954            cursor.read_exact(&mut meta_buf).unwrap();
3955
3956            let message = root_as_message(&meta_buf).unwrap();
3957            let body_len = message.bodyLength() as usize;
3958            let mut body_buf = vec![0u8; body_len];
3959            cursor.read_exact(&mut body_buf).unwrap();
3960
3961            if message.header_type() == crate::MessageHeader::DictionaryBatch {
3962                // Skip the dictionary batch — this is what C++ does for
3963                // all-null dictionary columns.
3964                continue;
3965            }
3966            stripped.extend_from_slice(&CONTINUATION_MARKER);
3967            stripped.extend_from_slice(&(meta_len as u32).to_le_bytes());
3968            stripped.extend_from_slice(&meta_buf);
3969            stripped.extend_from_slice(&body_buf);
3970        }
3971
3972        // Reading the stripped stream must succeed.
3973        let result = read_stream(&stripped).unwrap();
3974        assert_eq!(result.num_rows(), 4);
3975        assert_eq!(result.num_columns(), 1);
3976
3977        let col = result.column(0);
3978        assert_eq!(col.null_count(), 4);
3979        assert_eq!(col.len(), 4);
3980        // The result must be a dictionary-typed array.
3981        assert!(matches!(col.data_type(), DataType::Dictionary(_, _)));
3982    }
3983
3984    // Tests projected reads where a ListView column is skipped before another column.
3985    // This catches cases where skipping the ListView consumes the wrong number of buffers.
3986    #[test]
3987    fn test_projection_skip_list_view() {
3988        use crate::reader::FileReader;
3989        use crate::writer::FileWriter;
3990        use arrow_array::{
3991            GenericListViewArray, Int32Array, RecordBatch,
3992            builder::{GenericListViewBuilder, UInt32Builder},
3993        };
3994        use arrow_schema::{DataType, Field, Schema};
3995        use std::sync::Arc;
3996
3997        // Build a small ListView column with a mix of valid and null entries
3998        let mut builder = GenericListViewBuilder::<i32, _>::new(UInt32Builder::new());
3999
4000        builder.values().append_value(1);
4001        builder.values().append_value(2);
4002        builder.append(true);
4003
4004        builder.append(false);
4005
4006        builder.values().append_value(3);
4007        builder.values().append_value(4);
4008        builder.append(true);
4009
4010        let list_view: GenericListViewArray<i32> = builder.finish();
4011
4012        // Second column with simple values
4013        let values = Int32Array::from(vec![10, 20, 30]);
4014
4015        // Schema: first column is ListView, second is Int32
4016        let schema = Arc::new(Schema::new(vec![
4017            Field::new("a", list_view.data_type().clone(), true),
4018            Field::new("b", DataType::Int32, false),
4019        ]));
4020        // Create a batch with both columns
4021        let batch =
4022            RecordBatch::try_new(schema, vec![Arc::new(list_view), Arc::new(values.clone())])
4023                .unwrap();
4024
4025        // Write the batch to IPC
4026        let mut buf = Vec::new();
4027        {
4028            let mut writer = FileWriter::try_new(&mut buf, &batch.schema()).unwrap();
4029            writer.write(&batch).unwrap();
4030            writer.finish().unwrap();
4031        }
4032
4033        // Skip ListView column and Project only column "b"
4034        let mut reader = FileReader::try_new(std::io::Cursor::new(buf), Some(vec![1])).unwrap();
4035        let read_batch = reader.next().unwrap().unwrap();
4036
4037        // Verify that the projected column is read correctly
4038        assert_eq!(read_batch.num_columns(), 1);
4039        assert_eq!(read_batch.column(0).as_ref(), &values);
4040    }
4041
4042    // Tests reading a column when a preceding V4 Union column is skipped.
4043    // V4 Union columns include a null buffer and type ids (and offsets for dense unions).
4044    #[test]
4045    fn test_projection_skip_union_v4() {
4046        use crate::MetadataVersion;
4047        use crate::reader::FileReader;
4048        use crate::writer::{FileWriter, IpcWriteOptions};
4049        use arrow_array::{
4050            ArrayRef, Int32Array, RecordBatch, builder::UnionBuilder, types::Int32Type,
4051        };
4052        use arrow_schema::{DataType, Field, Schema};
4053        use std::sync::Arc;
4054
4055        // Build a dense Union column with simple Int32 values
4056        let mut builder = UnionBuilder::new_dense();
4057        builder.append::<Int32Type>("a", 1).unwrap();
4058        builder.append::<Int32Type>("a", 2).unwrap();
4059        builder.append::<Int32Type>("a", 3).unwrap();
4060        let union = builder.build().unwrap();
4061
4062        // Second column with known values to verify correctness after projection
4063        let values = Int32Array::from(vec![10, 20, 30]);
4064
4065        // Schema: first column is Union (to be skipped), second is Int32 (to be read)
4066        let schema = Arc::new(Schema::new(vec![
4067            Field::new("union", union.data_type().clone(), false),
4068            Field::new("values", DataType::Int32, false),
4069        ]));
4070
4071        // Create a batch containing both columns
4072        let batch = RecordBatch::try_new(
4073            schema,
4074            vec![Arc::new(union) as ArrayRef, Arc::new(values.clone())],
4075        )
4076        .unwrap();
4077
4078        // Write IPC using V4 metadata to trigger Union null buffer behavior
4079        let mut buf = Vec::new();
4080        {
4081            let options = IpcWriteOptions::try_new(8, false, MetadataVersion::V4).unwrap();
4082            let mut writer =
4083                FileWriter::try_new_with_options(&mut buf, &batch.schema(), options).unwrap();
4084            writer.write(&batch).unwrap();
4085            writer.finish().unwrap();
4086        }
4087        // Read only the second column (skip the Union column)
4088        let mut reader = FileReader::try_new(std::io::Cursor::new(buf), Some(vec![1])).unwrap();
4089        let read_batch = reader.next().unwrap().unwrap();
4090
4091        // Verify that the projected column is read correctly after skipping Union
4092        assert_eq!(read_batch.num_columns(), 1);
4093        assert_eq!(read_batch.column(0).as_ref(), &values);
4094    }
4095
4096    // Tests reading a column when preceding fixed-width and boolean columns are skipped.
4097    // Covers all types that use the same two-buffer layout (null + values).
4098    // Verifies that skipping these types does not affect subsequent column decoding.
4099    #[test]
4100    fn test_projection_skip_fixed_width_types() {
4101        use std::sync::Arc;
4102
4103        use arrow_array::{ArrayRef, BooleanArray, Int32Array, RecordBatch, make_array};
4104        use arrow_buffer::{Buffer, MutableBuffer};
4105        use arrow_data::ArrayData;
4106        use arrow_schema::{DataType, Field, IntervalUnit, Schema, TimeUnit};
4107
4108        use crate::reader::FileReader;
4109        use crate::writer::FileWriter;
4110
4111        // Create a minimal array for a given fixed-width or boolean type
4112        fn make_array_for_type(data_type: DataType) -> ArrayRef {
4113            let len = 3;
4114
4115            if matches!(data_type, DataType::Boolean) {
4116                return Arc::new(BooleanArray::from(vec![true, false, true]));
4117            }
4118
4119            let width = data_type.primitive_width().unwrap();
4120            let data = ArrayData::builder(data_type)
4121                .len(len)
4122                .add_buffer(Buffer::from(MutableBuffer::from_len_zeroed(len * width)))
4123                .build()
4124                .unwrap();
4125
4126            make_array(data)
4127        }
4128
4129        // List of types that follow the same two-buffer layout (null + values)
4130        let data_types = vec![
4131            DataType::Boolean,
4132            DataType::Int8,
4133            DataType::Int16,
4134            DataType::Int32,
4135            DataType::Int64,
4136            DataType::UInt8,
4137            DataType::UInt16,
4138            DataType::UInt32,
4139            DataType::UInt64,
4140            DataType::Float16,
4141            DataType::Float32,
4142            DataType::Float64,
4143            DataType::Timestamp(TimeUnit::Second, None),
4144            DataType::Date32,
4145            DataType::Date64,
4146            DataType::Time32(TimeUnit::Second),
4147            DataType::Time64(TimeUnit::Microsecond),
4148            DataType::Duration(TimeUnit::Second),
4149            DataType::Interval(IntervalUnit::YearMonth),
4150            DataType::Interval(IntervalUnit::DayTime),
4151            DataType::Interval(IntervalUnit::MonthDayNano),
4152            DataType::Decimal32(9, 2),
4153            DataType::Decimal64(18, 2),
4154            DataType::Decimal128(38, 2),
4155            DataType::Decimal256(76, 2),
4156        ];
4157
4158        // For each type:
4159        // - write a batch with [skipped_column, values]
4160        // - read only the second column
4161        // - verify the result is correct
4162        for data_type in data_types {
4163            let skipped = make_array_for_type(data_type.clone());
4164            let values = Int32Array::from(vec![10, 20, 30]);
4165
4166            let schema = Arc::new(Schema::new(vec![
4167                Field::new("skipped", data_type, false),
4168                Field::new("values", DataType::Int32, false),
4169            ]));
4170
4171            let batch =
4172                RecordBatch::try_new(schema, vec![skipped, Arc::new(values.clone())]).unwrap();
4173
4174            // Serialize the batch into IPC format
4175            let mut buf = Vec::new();
4176            {
4177                let mut writer = FileWriter::try_new(&mut buf, &batch.schema()).unwrap();
4178                writer.write(&batch).unwrap();
4179                writer.finish().unwrap();
4180            }
4181
4182            // Read back only the second column (skip the first)
4183            let mut reader = FileReader::try_new(std::io::Cursor::new(buf), Some(vec![1])).unwrap();
4184            let read_batch = reader.next().unwrap().unwrap();
4185
4186            // Verify that the returned column matches the original values column
4187            assert_eq!(read_batch.num_columns(), 1);
4188            assert_eq!(read_batch.column(0).as_ref(), &values);
4189        }
4190    }
4191
4192    /// Builds a stream whose single message declares `body_length` but is followed by only
4193    /// `body_bytes` bytes of body. The body length is consumed before the header is
4194    /// interpreted, so the header itself is immaterial.
4195    fn stream_with_declared_body(body_length: i64, body_bytes: usize) -> Vec<u8> {
4196        let mut fbb = flatbuffers::FlatBufferBuilder::new();
4197        let mut message = crate::MessageBuilder::new(&mut fbb);
4198        message.add_version(crate::MetadataVersion::V5);
4199        message.add_header_type(crate::MessageHeader::NONE);
4200        message.add_bodyLength(body_length);
4201        let root = message.finish();
4202        fbb.finish(root, None);
4203        let metadata = fbb.finished_data();
4204
4205        let mut stream = Vec::new();
4206        stream.extend_from_slice(&CONTINUATION_MARKER);
4207        stream.extend_from_slice(&(metadata.len() as i32).to_le_bytes());
4208        stream.extend_from_slice(metadata);
4209        stream.resize(stream.len() + body_bytes, 0xAB);
4210        stream
4211    }
4212
4213    /// A message body length is read from the stream before any of the bytes it describes,
4214    /// so it cannot be trusted. Declaring an implausible one used to reserve it outright,
4215    /// and the resulting allocation failure aborts the process instead of surfacing an error
4216    /// the caller can handle.
4217    #[test]
4218    fn test_stream_reader_rejects_implausible_body_length() {
4219        for body_length in [i64::MAX, 1 << 50, -1] {
4220            let stream = stream_with_declared_body(body_length, 0);
4221            let err = StreamReader::try_new(std::io::Cursor::new(stream), None).expect_err(
4222                &format!("a message declaring {body_length} body bytes must not be accepted"),
4223            );
4224            let err = err.to_string();
4225            assert!(
4226                err.contains("Invalid IPC message body length")
4227                    || err.contains("Unexpected end of stream")
4228                    || err.contains("failed to fill whole buffer"),
4229                "unexpected error for body_length {body_length}: {err}"
4230            );
4231        }
4232    }
4233
4234    /// A plausible body length whose bytes end early must fail as a short read: before the
4235    /// first read for a small body, and after at least one growth step for a body larger
4236    /// than `MAX_PREALLOC_BYTES`.
4237    #[test]
4238    #[cfg_attr(miri, ignore)] // Takes too long
4239    fn test_stream_reader_rejects_truncated_body() {
4240        let over_prealloc = MAX_PREALLOC_BYTES as i64 + 1;
4241        for (body_length, body_bytes) in [(1024, 10), (over_prealloc, MAX_PREALLOC_BYTES)] {
4242            let stream = stream_with_declared_body(body_length, body_bytes);
4243            let err = StreamReader::try_new(std::io::Cursor::new(stream), None)
4244                .expect_err("a body backed by fewer bytes than declared must not be accepted");
4245            assert!(
4246                err.to_string().contains("failed to fill whole buffer"),
4247                "unexpected error for body_length {body_length}: {err}"
4248            );
4249        }
4250    }
4251
4252    /// The metadata length is untrusted for the same reason as the body length.
4253    #[test]
4254    fn test_stream_reader_rejects_unbacked_metadata_length() {
4255        let mut stream = Vec::new();
4256        stream.extend_from_slice(&CONTINUATION_MARKER);
4257        stream.extend_from_slice(&i32::MAX.to_le_bytes());
4258        stream.extend_from_slice(b"not this many bytes");
4259
4260        let err = StreamReader::try_new(std::io::Cursor::new(stream), None)
4261            .expect_err("metadata length must be backed by the stream");
4262        assert!(
4263            err.to_string().contains("Unexpected end of stream"),
4264            "unexpected error: {err}"
4265        );
4266    }
4267}