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