Skip to main content

arrow_ipc/
reader.rs

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