Skip to main content

arrow_ipc/
reader.rs

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