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