Skip to main content

arrow_ipc/
writer.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: See the NOTICE file distributed with this work for additional information regarding copyright ownership
3// SPDX-FileContributor: Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements
4
5//! Arrow IPC File and Stream Writers
6//!
7//! # Notes
8//!
9//! [`FileWriter`] and [`StreamWriter`] have similar interfaces,
10//! however the [`FileWriter`] expects a reader that supports [`Seek`]ing
11//!
12//! [`Seek`]: std::io::Seek
13
14use std::cmp::min;
15use std::collections::HashMap;
16use std::io::{BufWriter, Write};
17use std::mem::size_of;
18use std::sync::Arc;
19
20use flatbuffers::FlatBufferBuilder;
21
22use arrow_array::cast::*;
23use arrow_array::types::{Int16Type, Int32Type, Int64Type, RunEndIndexType};
24use arrow_array::*;
25use arrow_buffer::{ArrowNativeType, Buffer, MutableBuffer, ToByteSlice};
26use arrow_data::{ArrayData, ArrayDataBuilder, BufferSpec, layout};
27use arrow_schema::*;
28
29use crate::CONTINUATION_MARKER;
30use crate::compression::CompressionCodec;
31#[expect(deprecated)]
32pub use crate::compression::{CompressionContext, IpcWriteContext};
33use crate::convert::IpcSchemaEncoder;
34
35/// IPC write options used to control the behaviour of the [`IpcDataGenerator`]
36#[derive(Debug, Clone)]
37pub struct IpcWriteOptions {
38    /// Write padding after memory buffers to this multiple of bytes.
39    /// Must be 8, 16, 32, or 64 - defaults to 64.
40    alignment: u8,
41    /// The legacy format is for releases before 0.15.0, and uses metadata V4
42    write_legacy_ipc_format: bool,
43    /// The metadata version to write. The Rust IPC writer supports V4+
44    ///
45    /// *Default versions per crate*
46    ///
47    /// When creating the default IpcWriteOptions, the following metadata versions are used:
48    ///
49    /// version 2.0.0: V4, with legacy format enabled
50    /// version 4.0.0: V5
51    metadata_version: crate::MetadataVersion,
52    /// Compression, if desired. Will result in a runtime error
53    /// if the corresponding feature is not enabled
54    batch_compression_type: Option<crate::CompressionType>,
55    // Compression level
56    batch_compression_level: Option<i32>,
57    /// How to handle updating dictionaries in IPC messages
58    dictionary_handling: DictionaryHandling,
59}
60
61/// A single buffer segment ready to be written to the output stream.
62///
63/// For the uncompressed path the original Arc-backed [`Buffer`] is stored
64/// directly (zero copy). For the compressed path the compressed bytes are
65/// owned by a scratch `Vec<u8>`.
66enum EncodedBuffer {
67    /// Uncompressed: Arc-backed reference to the original array buffer.
68    Raw(Buffer),
69    /// Compressed: owned scratch bytes produced by the codec.
70    Compressed(Vec<u8>),
71}
72
73impl EncodedBuffer {
74    fn as_slice(&self) -> &[u8] {
75        match self {
76            EncodedBuffer::Raw(b) => b.as_slice(),
77            EncodedBuffer::Compressed(v) => v.as_slice(),
78        }
79    }
80
81    fn len(&self) -> usize {
82        match self {
83            EncodedBuffer::Raw(b) => b.len(),
84            EncodedBuffer::Compressed(v) => v.len(),
85        }
86    }
87}
88/// Accumulates the IPC metadata produced by [`write_array_data`].
89///
90/// `nodes` and `buffers` are serialised into the flatbuffer `RecordBatch` (or `DictionaryBatch`)
91/// header. The companion [`IpcBodySink`] holds the actual encoded bytes.
92#[derive(Default)]
93struct IpcMetadataBuilder {
94    nodes: Vec<crate::FieldNode>,
95    buffers: Vec<crate::Buffer>,
96}
97
98/// Destination for the raw Arrow data bytes (the IPC message body) produced by [`write_array_data`].
99///
100/// The companion [`IpcMetadataBuilder`] accumulates the flatbuffer metadata
101/// (offset + length of each buffer in the body); together they form a complete IPC message.
102enum IpcBodySink<'a> {
103    /// Serialize buffer bytes (with padding) into a contiguous byte vec.
104    Write(&'a mut Vec<u8>),
105    /// Accumulate pre-encoded buffer segments for deferred zero-copy streaming.
106    Collect(&'a mut Vec<EncodedBuffer>),
107}
108impl IpcBodySink<'_> {
109    /// Writes the encoded buffer to the sink.
110    pub fn write(&mut self, pad_len: usize, buffer: EncodedBuffer) {
111        match self {
112            IpcBodySink::Write(vec) => {
113                vec.extend_from_slice(buffer.as_slice());
114                vec.extend_from_slice(&PADDING[..pad_len]);
115            }
116            IpcBodySink::Collect(vec) => {
117                vec.push(buffer);
118            }
119        }
120    }
121}
122
123struct MetadataLayout {
124    padded_header_len: usize,
125    padded_metadata_len: usize,
126    metadata_padding: usize,
127}
128
129impl MetadataLayout {
130    fn new(metadata_len: usize, write_options: &IpcWriteOptions) -> MetadataLayout {
131        let prefix_size = if write_options.write_legacy_ipc_format {
132            4
133        } else {
134            8
135        };
136        let alignment_mask = usize::from(write_options.alignment - 1);
137        let padded_header_len = (metadata_len + prefix_size + alignment_mask) & !alignment_mask;
138        let padded_metadata_len = padded_header_len - prefix_size;
139        let metadata_padding = padded_metadata_len - metadata_len;
140
141        MetadataLayout {
142            padded_header_len,
143            padded_metadata_len,
144            metadata_padding,
145        }
146    }
147}
148
149/// Destination-specific byte writing for a framed IPC message.
150///
151/// The default owned-buffer methods borrow their bytes and call `write_slice`,
152/// which can copy the data. Sink implementations that can take ownership of
153/// generated bytes should override `write_vec` or `write_encoded_buffer`.
154trait IpcMessageSink {
155    fn write_slice(&mut self, bytes: &[u8]) -> Result<(), ArrowError>;
156
157    fn write_vec(&mut self, bytes: Vec<u8>) -> Result<(), ArrowError> {
158        self.write_slice(&bytes)
159    }
160
161    fn write_encoded_buffer(&mut self, buffer: EncodedBuffer) -> Result<(), ArrowError> {
162        self.write_slice(buffer.as_slice())
163    }
164}
165
166/// Shared IPC framing helpers for [`IpcMessageSink`].
167trait IpcMessageSinkExt: IpcMessageSink {
168    fn write_padding(&mut self, len: usize) -> Result<(), ArrowError> {
169        self.write_slice(&PADDING[..len])
170    }
171
172    /// Writes the IPC continuation marker and metadata length prefix.
173    fn write_continuation(
174        &mut self,
175        write_options: &IpcWriteOptions,
176        metadata_len: i32,
177    ) -> Result<(), ArrowError> {
178        let mut buffer = [0; 8];
179        let len = match write_options.metadata_version {
180            crate::MetadataVersion::V1
181            | crate::MetadataVersion::V2
182            | crate::MetadataVersion::V3 => {
183                unreachable!("Options with the metadata version cannot be created")
184            }
185            crate::MetadataVersion::V4 => {
186                let metadata_len_bytes = metadata_len.to_le_bytes();
187                if !write_options.write_legacy_ipc_format {
188                    // v0.15.0 format
189                    buffer[..4].copy_from_slice(&CONTINUATION_MARKER);
190                    buffer[4..].copy_from_slice(&metadata_len_bytes);
191                    8
192                } else {
193                    buffer[..4].copy_from_slice(&metadata_len_bytes);
194                    4
195                }
196            }
197            crate::MetadataVersion::V5 => {
198                buffer[..4].copy_from_slice(&CONTINUATION_MARKER);
199                buffer[4..].copy_from_slice(&metadata_len.to_le_bytes());
200                8
201            }
202            z => panic!("Unsupported crate::MetadataVersion {z:?}"),
203        };
204        self.write_slice(&buffer[..len])
205    }
206
207    fn write_body_data(&mut self, data: Vec<u8>, alignment: u8) -> Result<usize, ArrowError> {
208        let len = data.len();
209        let pad_len = pad_to_alignment(alignment, len);
210        self.write_vec(data)?;
211        self.write_padding(pad_len)?;
212        Ok(len + pad_len)
213    }
214
215    /// Writes an already encoded IPC message with optional contiguous body data.
216    ///
217    /// This is used for schema and dictionary messages represented by [`EncodedData`].
218    /// Returns the padded metadata length and body length written.
219    fn write_encoded_data(
220        &mut self,
221        encoded: EncodedData,
222        write_options: &IpcWriteOptions,
223    ) -> Result<(usize, usize), ArrowError> {
224        let arrow_data_len = encoded.arrow_data.len();
225        if !arrow_data_len.is_multiple_of(usize::from(write_options.alignment)) {
226            return Err(ArrowError::MemoryError(
227                "Arrow data not aligned".to_string(),
228            ));
229        }
230
231        let metadata = encoded.ipc_message;
232        let metadata_len = metadata.len();
233        let layout = MetadataLayout::new(metadata_len, write_options);
234
235        self.write_continuation(write_options, layout.padded_metadata_len as i32)?;
236        self.write_vec(metadata)?;
237        self.write_padding(layout.metadata_padding)?;
238
239        let body_len = if arrow_data_len > 0 {
240            self.write_body_data(encoded.arrow_data, write_options.alignment)?
241        } else {
242            0
243        };
244
245        Ok((layout.padded_header_len, body_len))
246    }
247
248    /// Writes the IPC end-of-stream marker.
249    fn write_eos(&mut self, write_options: &IpcWriteOptions) -> Result<(), ArrowError> {
250        self.write_continuation(write_options, 0)?;
251        Ok(())
252    }
253}
254
255impl<T: IpcMessageSink + ?Sized> IpcMessageSinkExt for T {}
256
257/// Optional hot-path hook for record batch messages.
258trait IpcRecordBatchSink: IpcMessageSinkExt {
259    /// Writes a record batch (or dictionary batch) message from its encoded
260    /// metadata and body buffers.
261    ///
262    /// `metadata` is the raw flatbuffer [`crate::Message`] (without continuation
263    /// prefix), borrowed from the reused [`FlatBufferBuilder`]. The body buffers are
264    /// already materialized as [`EncodedBuffer`] segments, allowing buffer output to
265    /// preserve uncompressed Arrow buffers. They are drained out of
266    /// `encoded_buffers` so the caller can reuse its allocation.
267    ///
268    /// Each body buffer is padded to the alignment as it is written, so the body
269    /// needs no trailing padding.
270    ///
271    /// Returns the padded metadata length and body length written.
272    fn write_record_batch(
273        &mut self,
274        metadata: &[u8],
275        encoded_buffers: &mut Vec<EncodedBuffer>,
276        body_len: usize,
277        write_options: &IpcWriteOptions,
278    ) -> Result<(usize, usize), ArrowError> {
279        let alignment = write_options.alignment;
280        let layout = MetadataLayout::new(metadata.len(), write_options);
281
282        self.write_continuation(write_options, layout.padded_metadata_len as i32)?;
283        self.write_slice(metadata)?;
284        self.write_padding(layout.metadata_padding)?;
285        for enc in encoded_buffers.drain(..) {
286            let len = enc.len();
287            self.write_encoded_buffer(enc)?;
288            self.write_padding(pad_to_alignment(alignment, len))?;
289        }
290
291        Ok((layout.padded_header_len, body_len))
292    }
293}
294
295impl<W> IpcMessageSink for W
296where
297    W: Write,
298{
299    fn write_slice(&mut self, bytes: &[u8]) -> Result<(), ArrowError> {
300        if !bytes.is_empty() {
301            self.write_all(bytes)?;
302        }
303        Ok(())
304    }
305}
306
307impl<W> IpcRecordBatchSink for W
308where
309    W: Write,
310{
311    fn write_record_batch(
312        &mut self,
313        metadata: &[u8],
314        encoded_buffers: &mut Vec<EncodedBuffer>,
315        body_len: usize,
316        write_options: &IpcWriteOptions,
317    ) -> Result<(usize, usize), ArrowError> {
318        let alignment = write_options.alignment;
319        let layout = MetadataLayout::new(metadata.len(), write_options);
320
321        self.write_continuation(write_options, layout.padded_metadata_len as i32)?;
322        self.write_all(metadata)?;
323        self.write_all(&PADDING[..layout.metadata_padding])?;
324        for enc in encoded_buffers.iter_mut() {
325            self.write_all(enc.as_slice())?;
326            self.write_all(&PADDING[..pad_to_alignment(alignment, enc.len())])?;
327        }
328        // Clearing the buffers after the loop instead of draining the vec is
329        // a deliberate choice. It seems to perform much better on certain
330        // ARM chips.
331        //
332        // See: https://github.com/apache/arrow-rs/pull/10128#issuecomment-5406019272
333        encoded_buffers.clear();
334
335        Ok((layout.padded_header_len, body_len))
336    }
337}
338
339/// Accumulates complete framed IPC messages as ordered buffers.
340struct Buffers<'a> {
341    out: &'a mut Vec<Buffer>,
342}
343
344impl IpcMessageSink for Buffers<'_> {
345    fn write_slice(&mut self, bytes: &[u8]) -> Result<(), ArrowError> {
346        if !bytes.is_empty() {
347            self.out.push(Buffer::from(bytes));
348        }
349        Ok(())
350    }
351
352    fn write_vec(&mut self, bytes: Vec<u8>) -> Result<(), ArrowError> {
353        if !bytes.is_empty() {
354            self.out.push(Buffer::from(bytes));
355        }
356        Ok(())
357    }
358
359    fn write_encoded_buffer(&mut self, buffer: EncodedBuffer) -> Result<(), ArrowError> {
360        match buffer {
361            EncodedBuffer::Raw(buffer) => self.out.push(buffer),
362            EncodedBuffer::Compressed(bytes) => self.out.push(Buffer::from(bytes)),
363        }
364        Ok(())
365    }
366}
367
368impl IpcRecordBatchSink for Buffers<'_> {}
369
370/// Per-message sizes produced by [`IpcDataGenerator::write`].
371///
372/// [`FileWriter`] uses these to build the Block index entries required by the IPC footer for
373/// random-access reads.
374struct IpcWriteMetadata {
375    /// Per-dictionary `(padded_header_len, body_len)` for each dictionary batch written
376    /// before the record batch.
377    dictionary_block_sizes: Vec<(usize, usize)>,
378    /// Flatbuffer header size including continuation prefix and alignment padding.
379    padded_header_len: usize,
380    /// Total length of the record-batch body including alignment padding
381    body_len: usize,
382}
383
384impl IpcWriteOptions {
385    /// Configures compression when writing IPC files.
386    ///
387    /// Will result in a runtime error if the corresponding feature
388    /// is not enabled
389    pub fn try_with_compression(
390        mut self,
391        batch_compression_type: Option<crate::CompressionType>,
392    ) -> Result<Self, ArrowError> {
393        self.batch_compression_type = batch_compression_type;
394
395        if self.batch_compression_type.is_some()
396            && self.metadata_version < crate::MetadataVersion::V5
397        {
398            return Err(ArrowError::InvalidArgumentError(
399                "Compression only supported in metadata v5 and above".to_string(),
400            ));
401        }
402        Ok(self)
403    }
404
405    /// Configures the compression level used when writing compressed IPC batches.
406    ///
407    /// Compression levels require metadata V5 or newer and are currently only
408    /// supported for ZSTD compression.
409    pub fn try_with_compression_level(
410        mut self,
411        batch_compression_level: Option<i32>,
412    ) -> Result<Self, ArrowError> {
413        self.batch_compression_level = batch_compression_level;
414
415        if self.batch_compression_level.is_some()
416            && self.metadata_version < crate::MetadataVersion::V5
417        {
418            return Err(ArrowError::InvalidArgumentError(
419                "Compression only supported in metadata v5 and above".to_string(),
420            ));
421        }
422
423        match (self.batch_compression_type, self.batch_compression_level) {
424            (Some(crate::CompressionType::ZSTD), Some(level)) => {
425                return self.check_zstd_level(level);
426            }
427            (Some(crate::CompressionType::LZ4_FRAME), Some(_)) => {
428                return Err(ArrowError::InvalidArgumentError(
429                    "LZ4 Frame compression does not support configurable compression levels"
430                        .to_string(),
431                ));
432            }
433            _ => {}
434        }
435
436        Ok(self)
437    }
438
439    #[cfg(not(feature = "zstd"))]
440    fn check_zstd_level(self, _level: i32) -> Result<Self, ArrowError> {
441        Err(ArrowError::InvalidArgumentError(
442            "zstd IPC compression requires the zstd feature".to_string(),
443        ))
444    }
445
446    #[cfg(feature = "zstd")]
447    fn check_zstd_level(self, level: i32) -> Result<Self, ArrowError> {
448        let range = zstd::compression_level_range();
449        if !range.contains(&level) {
450            return Err(ArrowError::InvalidArgumentError(format!(
451                "ZSTD compression level must be between {} and {}, got {}",
452                range.start(),
453                range.end(),
454                level,
455            )));
456        }
457
458        Ok(self)
459    }
460
461    /// Try to create IpcWriteOptions, checking for incompatible settings
462    pub fn try_new(
463        alignment: usize,
464        write_legacy_ipc_format: bool,
465        metadata_version: crate::MetadataVersion,
466    ) -> Result<Self, ArrowError> {
467        let alignment: u8 = match alignment {
468            8 => 8,
469            16 => 16,
470            32 => 32,
471            64 => 64,
472            _ => {
473                return Err(ArrowError::InvalidArgumentError(
474                    "Alignment should be 8, 16, 32, or 64.".to_string(),
475                ));
476            }
477        };
478        match metadata_version {
479            crate::MetadataVersion::V1
480            | crate::MetadataVersion::V2
481            | crate::MetadataVersion::V3 => Err(ArrowError::InvalidArgumentError(
482                "Writing IPC metadata version 3 and lower not supported".to_string(),
483            )),
484            crate::MetadataVersion::V4 => Ok(Self {
485                alignment,
486                write_legacy_ipc_format,
487                metadata_version,
488                batch_compression_type: None,
489                batch_compression_level: None,
490                dictionary_handling: DictionaryHandling::default(),
491            }),
492            crate::MetadataVersion::V5 => {
493                if write_legacy_ipc_format {
494                    Err(ArrowError::InvalidArgumentError(
495                        "Legacy IPC format only supported on metadata version 4".to_string(),
496                    ))
497                } else {
498                    Ok(Self {
499                        alignment,
500                        write_legacy_ipc_format,
501                        metadata_version,
502                        batch_compression_type: None,
503                        batch_compression_level: None,
504                        dictionary_handling: DictionaryHandling::default(),
505                    })
506                }
507            }
508            z => Err(ArrowError::InvalidArgumentError(format!(
509                "Unsupported crate::MetadataVersion {z:?}"
510            ))),
511        }
512    }
513
514    /// Configure how dictionaries are handled in IPC messages
515    pub fn with_dictionary_handling(mut self, dictionary_handling: DictionaryHandling) -> Self {
516        self.dictionary_handling = dictionary_handling;
517        self
518    }
519}
520
521impl Default for IpcWriteOptions {
522    fn default() -> Self {
523        Self {
524            alignment: 64,
525            write_legacy_ipc_format: false,
526            metadata_version: crate::MetadataVersion::V5,
527            batch_compression_type: None,
528            batch_compression_level: None,
529            dictionary_handling: DictionaryHandling::default(),
530        }
531    }
532}
533
534#[derive(Debug, Default)]
535/// Handles low level details of encoding [`Array`] and [`Schema`] into the
536/// [Arrow IPC Format].
537///
538/// # Example
539/// ```
540/// # fn run() {
541/// # use std::sync::Arc;
542/// # use arrow_array::UInt64Array;
543/// # use arrow_array::RecordBatch;
544/// # use arrow_ipc::writer::{IpcWriteContext, DictionaryTracker, IpcDataGenerator, IpcWriteOptions};
545///
546/// // Create a record batch
547/// let batch = RecordBatch::try_from_iter(vec![
548///  ("col2", Arc::new(UInt64Array::from_iter([10, 23, 33])) as _)
549/// ]).unwrap();
550///
551/// // Error of dictionary ids are replaced.
552/// let error_on_replacement = true;
553/// let options = IpcWriteOptions::default();
554/// let mut dictionary_tracker = DictionaryTracker::new(error_on_replacement);
555///
556/// let mut ipc_write_context = IpcWriteContext::default();
557///
558/// // encode the batch into zero or more encoded dictionaries
559/// // and the data for the actual array.
560/// let data_gen = IpcDataGenerator::default();
561/// let (encoded_dictionaries, encoded_message) = data_gen
562///   .encode(&batch, &mut dictionary_tracker, &options, &mut ipc_write_context)
563///   .unwrap();
564/// # }
565/// ```
566///
567/// [Arrow IPC Format]: https://arrow.apache.org/docs/format/Columnar.html#serialization-and-interprocess-communication-ipc
568pub struct IpcDataGenerator {}
569
570/// A dictionary determined to need sending as a part of a record batch.
571struct DictionaryToEncode {
572    id: i64,
573    data: ArrayData,
574    is_delta: bool,
575}
576
577/// The result of encoding a record batch's body and flatbuffer `RecordBatch` table
578/// via [`IpcDataGenerator::encode_record_batch_data`].
579///
580/// The `RecordBatch` table is left in-progress in the caller's [`FlatBufferBuilder`];
581/// the caller wraps `record_batch` in a `Message` (directly for a record batch
582/// message, or inside a `DictionaryBatch` for a dictionary message).
583struct EncodedRecordBatchMeta {
584    /// In-progress flatbuffer offset of the `RecordBatch` table.
585    fb_offset: flatbuffers::WIPOffset<crate::RecordBatch<'static>>,
586    /// Total body length written to the sink (the sum of each buffer's encoded
587    /// length and its alignment padding).
588    body_len: usize,
589}
590
591impl IpcDataGenerator {
592    /// Converts a schema to an IPC message along with `dictionary_tracker`
593    /// and returns it encoded inside [EncodedData] as a flatbuffer.
594    pub fn schema_to_bytes_with_dictionary_tracker(
595        &self,
596        schema: &Schema,
597        dictionary_tracker: &mut DictionaryTracker,
598        write_options: &IpcWriteOptions,
599    ) -> EncodedData {
600        let mut fbb = FlatBufferBuilder::new();
601        let schema = {
602            let fb = IpcSchemaEncoder::new()
603                .with_dictionary_tracker(dictionary_tracker)
604                .schema_to_fb_offset(&mut fbb, schema);
605            fb.as_union_value()
606        };
607
608        let mut message = crate::MessageBuilder::new(&mut fbb);
609        message.add_version(write_options.metadata_version);
610        message.add_header_type(crate::MessageHeader::Schema);
611        message.add_bodyLength(0);
612        message.add_header(schema);
613        // TODO: custom metadata
614        let root = message.finish();
615        fbb.finish(root, None);
616
617        let metadata = fbb.finished_data();
618        EncodedData {
619            ipc_message: metadata.to_vec(),
620            arrow_data: vec![],
621        }
622    }
623
624    fn _collect_dict_updates<I: Iterator<Item = i64>>(
625        &self,
626        column: &ArrayRef,
627        dictionaries_to_encode: &mut Vec<DictionaryToEncode>,
628        dictionary_tracker: &mut DictionaryTracker,
629        write_options: &IpcWriteOptions,
630        dict_id: &mut I,
631    ) -> Result<(), ArrowError> {
632        match column.data_type() {
633            DataType::Struct(fields) => {
634                let s = as_struct_array(column);
635                for (field, column) in fields.iter().zip(s.columns()) {
636                    self.collect_dict_updates(
637                        field,
638                        column,
639                        dictionaries_to_encode,
640                        dictionary_tracker,
641                        write_options,
642                        dict_id,
643                    )?;
644                }
645            }
646            DataType::RunEndEncoded(_, values) => {
647                let data = column.to_data();
648                if data.child_data().len() != 2 {
649                    return Err(ArrowError::InvalidArgumentError(format!(
650                        "The run encoded array should have exactly two child arrays. Found {}",
651                        data.child_data().len()
652                    )));
653                }
654                // The run_ends array is not expected to be dictionary encoded. Hence encode dictionaries
655                // only for values array.
656                let values_array = make_array(data.child_data()[1].clone());
657                self.collect_dict_updates(
658                    values,
659                    &values_array,
660                    dictionaries_to_encode,
661                    dictionary_tracker,
662                    write_options,
663                    dict_id,
664                )?;
665            }
666            DataType::List(field) => {
667                let list = as_list_array(column);
668                self.collect_dict_updates(
669                    field,
670                    list.values(),
671                    dictionaries_to_encode,
672                    dictionary_tracker,
673                    write_options,
674                    dict_id,
675                )?;
676            }
677            DataType::LargeList(field) => {
678                let list = as_large_list_array(column);
679                self.collect_dict_updates(
680                    field,
681                    list.values(),
682                    dictionaries_to_encode,
683                    dictionary_tracker,
684                    write_options,
685                    dict_id,
686                )?;
687            }
688            DataType::ListView(field) => {
689                let list = column.as_list_view::<i32>();
690                self.collect_dict_updates(
691                    field,
692                    list.values(),
693                    dictionaries_to_encode,
694                    dictionary_tracker,
695                    write_options,
696                    dict_id,
697                )?;
698            }
699            DataType::LargeListView(field) => {
700                let list = column.as_list_view::<i64>();
701                self.collect_dict_updates(
702                    field,
703                    list.values(),
704                    dictionaries_to_encode,
705                    dictionary_tracker,
706                    write_options,
707                    dict_id,
708                )?;
709            }
710            DataType::FixedSizeList(field, _) => {
711                let list = column
712                    .as_any()
713                    .downcast_ref::<FixedSizeListArray>()
714                    .expect("Unable to downcast to fixed size list array");
715                self.collect_dict_updates(
716                    field,
717                    list.values(),
718                    dictionaries_to_encode,
719                    dictionary_tracker,
720                    write_options,
721                    dict_id,
722                )?;
723            }
724            DataType::Map(field, _) => {
725                let map_array = as_map_array(column);
726
727                let (keys, values) = match field.data_type() {
728                    DataType::Struct(fields) if fields.len() == 2 => (&fields[0], &fields[1]),
729                    _ => panic!("Incorrect field data type {:?}", field.data_type()),
730                };
731
732                // keys
733                self.collect_dict_updates(
734                    keys,
735                    map_array.keys(),
736                    dictionaries_to_encode,
737                    dictionary_tracker,
738                    write_options,
739                    dict_id,
740                )?;
741
742                // values
743                self.collect_dict_updates(
744                    values,
745                    map_array.values(),
746                    dictionaries_to_encode,
747                    dictionary_tracker,
748                    write_options,
749                    dict_id,
750                )?;
751            }
752            DataType::Union(fields, _) => {
753                let union = as_union_array(column);
754                for (type_id, field) in fields.iter() {
755                    let column = union.child(type_id);
756                    self.collect_dict_updates(
757                        field,
758                        column,
759                        dictionaries_to_encode,
760                        dictionary_tracker,
761                        write_options,
762                        dict_id,
763                    )?;
764                }
765            }
766            _ => (),
767        }
768
769        Ok(())
770    }
771
772    // Collect all dicts that need a dictionary message i.e. ones that were
773    // either not in the tracker previously or ones that were but need a
774    // replacement or delta.
775    fn collect_dict_updates<I: Iterator<Item = i64>>(
776        &self,
777        field: &Field,
778        column: &ArrayRef,
779        dictionaries_to_encode: &mut Vec<DictionaryToEncode>,
780        dictionary_tracker: &mut DictionaryTracker,
781        write_options: &IpcWriteOptions,
782        dict_id_seq: &mut I,
783    ) -> Result<(), ArrowError> {
784        match column.data_type() {
785            DataType::Dictionary(_key_type, value_type) => {
786                if matches!(value_type.as_ref(), DataType::Dictionary(_, _)) {
787                    return Err(ArrowError::InvalidArgumentError(format!(
788                        "Arrow IPC field metadata cannot encode direct dictionary-of-dictionary values for field {:?}",
789                        field.name()
790                    )));
791                }
792
793                let dict_data = column.to_data();
794                let dict_values = &dict_data.child_data()[0];
795
796                let values = make_array(dict_data.child_data()[0].clone());
797
798                self._collect_dict_updates(
799                    &values,
800                    dictionaries_to_encode,
801                    dictionary_tracker,
802                    write_options,
803                    dict_id_seq,
804                )?;
805
806                // It's important to only take the dict_id at this point, because the dict ID
807                // sequence is assigned depth-first, so we need to first encode children and have
808                // them take their assigned dict IDs before we take the dict ID for this field.
809                let dict_id = dict_id_seq.next().ok_or_else(|| {
810                    ArrowError::IpcError(format!(
811                        "no dict id for field {:?}: field.data_type={:?}, column.data_type={:?}",
812                        field.name(),
813                        field.data_type(),
814                        column.data_type()
815                    ))
816                })?;
817
818                match dictionary_tracker.insert_column(
819                    dict_id,
820                    column,
821                    write_options.dictionary_handling,
822                )? {
823                    DictionaryUpdate::None => {}
824                    DictionaryUpdate::New | DictionaryUpdate::Replaced => {
825                        dictionaries_to_encode.push(DictionaryToEncode {
826                            id: dict_id,
827                            data: dict_values.clone(),
828                            is_delta: false,
829                        });
830                    }
831                    DictionaryUpdate::Delta(data) => {
832                        dictionaries_to_encode.push(DictionaryToEncode {
833                            id: dict_id,
834                            data,
835                            is_delta: true,
836                        });
837                    }
838                }
839            }
840            _ => self._collect_dict_updates(
841                column,
842                dictionaries_to_encode,
843                dictionary_tracker,
844                write_options,
845                dict_id_seq,
846            )?,
847        }
848
849        Ok(())
850    }
851
852    /// Encodes a batch to a number of [EncodedData] items (dictionary batches + the record batch).
853    /// The [DictionaryTracker] keeps track of dictionaries with new `dict_id`s  (so they are only sent once)
854    /// Make sure the [DictionaryTracker] is initialized at the start of the stream.
855    pub fn encode(
856        &self,
857        batch: &RecordBatch,
858        dictionary_tracker: &mut DictionaryTracker,
859        write_options: &IpcWriteOptions,
860        ipc_write_context: &mut IpcWriteContext,
861    ) -> Result<(Vec<EncodedData>, EncodedData), ArrowError> {
862        let dictionaries = self.collect_all_dicts(batch, dictionary_tracker, write_options)?;
863        let mut encoded_dictionaries = Vec::with_capacity(dictionaries.len());
864        for dict in dictionaries {
865            encoded_dictionaries.push(self.dictionary_batch_to_bytes(
866                dict,
867                write_options,
868                ipc_write_context,
869            )?);
870        }
871        let mut arrow_data = ipc_write_context.scratch();
872        self.record_batch_to_bytes(
873            batch,
874            write_options,
875            ipc_write_context,
876            &mut IpcBodySink::Write(&mut arrow_data),
877        )?;
878        ipc_write_context.reserve_scratch_with_capacity(arrow_data.capacity());
879        Ok((
880            encoded_dictionaries,
881            EncodedData {
882                ipc_message: ipc_write_context.mut_fbb().finished_data().to_vec(),
883                arrow_data,
884            },
885        ))
886    }
887
888    /// Walk the record batch and collect dictionaries for later encoding.
889    fn collect_all_dicts(
890        &self,
891        batch: &RecordBatch,
892        dictionary_tracker: &mut DictionaryTracker,
893        write_options: &IpcWriteOptions,
894    ) -> Result<Vec<DictionaryToEncode>, ArrowError> {
895        let schema = batch.schema();
896        let mut dictionaries = Vec::with_capacity(schema.flattened_fields().len());
897        let mut dict_id = dictionary_tracker.dict_ids.clone().into_iter();
898        for (i, field) in schema.fields().iter().enumerate() {
899            self.collect_dict_updates(
900                field,
901                batch.column(i),
902                &mut dictionaries,
903                dictionary_tracker,
904                write_options,
905                &mut dict_id,
906            )?;
907        }
908        Ok(dictionaries)
909    }
910
911    /// Write dictionary batches and the record batch directly to `writer`, skipping the
912    /// intermediate body `Vec<u8>` allocations
913    /// Returns [`IpcWriteMetadata`] with the sizes needed to build footer blocks.
914    fn write<W: Write>(
915        &self,
916        batch: &RecordBatch,
917        dictionary_tracker: &mut DictionaryTracker,
918        write_options: &IpcWriteOptions,
919        ipc_write_context: &mut IpcWriteContext,
920        writer: &mut W,
921    ) -> Result<IpcWriteMetadata, ArrowError> {
922        self.write_to_sink(
923            batch,
924            dictionary_tracker,
925            write_options,
926            ipc_write_context,
927            writer,
928        )
929    }
930
931    /// Encode dictionary batches and the record batch to output buffers, skipping the
932    /// intermediate body `Vec<u8>` allocations for uncompressed record batch buffers.
933    fn encode_to_buffers(
934        &self,
935        batch: &RecordBatch,
936        dictionary_tracker: &mut DictionaryTracker,
937        write_options: &IpcWriteOptions,
938        ipc_write_context: &mut IpcWriteContext,
939        out: &mut Vec<Buffer>,
940    ) -> Result<IpcWriteMetadata, ArrowError> {
941        let mut sink = Buffers { out };
942        self.write_to_sink(
943            batch,
944            dictionary_tracker,
945            write_options,
946            ipc_write_context,
947            &mut sink,
948        )
949    }
950
951    fn write_to_sink<S: IpcRecordBatchSink>(
952        &self,
953        batch: &RecordBatch,
954        dictionary_tracker: &mut DictionaryTracker,
955        write_options: &IpcWriteOptions,
956        ipc_write_context: &mut IpcWriteContext,
957        sink: &mut S,
958    ) -> Result<IpcWriteMetadata, ArrowError> {
959        let dictionaries = self.collect_all_dicts(batch, dictionary_tracker, write_options)?;
960        let mut dictionary_block_sizes = Vec::with_capacity(dictionaries.len());
961
962        let capacity = batch
963            .columns()
964            .iter()
965            .map(|a| estimate_encoded_buffer_count(a.data_type()))
966            .sum();
967        let mut encoded_buffers: Vec<EncodedBuffer> = Vec::with_capacity(capacity);
968
969        for dict in &dictionaries {
970            encoded_buffers.clear();
971
972            let body_len = self.dictionary_batch_to_sink(
973                dict,
974                write_options,
975                ipc_write_context,
976                &mut IpcBodySink::Collect(&mut encoded_buffers),
977            )?;
978
979            dictionary_block_sizes.push(sink.write_record_batch(
980                ipc_write_context.mut_fbb().finished_data(),
981                &mut encoded_buffers,
982                body_len,
983                write_options,
984            )?);
985        }
986
987        encoded_buffers.clear();
988        let body_len = self.record_batch_to_bytes(
989            batch,
990            write_options,
991            ipc_write_context,
992            &mut IpcBodySink::Collect(&mut encoded_buffers),
993        )?;
994
995        let (padded_header_len, body_len) = sink.write_record_batch(
996            ipc_write_context.mut_fbb().finished_data(),
997            &mut encoded_buffers,
998            body_len,
999            write_options,
1000        )?;
1001
1002        Ok(IpcWriteMetadata {
1003            dictionary_block_sizes,
1004            padded_header_len,
1005            body_len,
1006        })
1007    }
1008
1009    /// Encodes a batch to a number of [EncodedData] items (dictionary batches + the record batch).
1010    /// The [DictionaryTracker] keeps track of dictionaries with new `dict_id`s  (so they are only sent once)
1011    /// Make sure the [DictionaryTracker] is initialized at the start of the stream.
1012    #[deprecated(since = "57.0.0", note = "Use `encode` instead")]
1013    pub fn encoded_batch(
1014        &self,
1015        batch: &RecordBatch,
1016        dictionary_tracker: &mut DictionaryTracker,
1017        write_options: &IpcWriteOptions,
1018    ) -> Result<(Vec<EncodedData>, EncodedData), ArrowError> {
1019        self.encode(
1020            batch,
1021            dictionary_tracker,
1022            write_options,
1023            &mut Default::default(),
1024        )
1025    }
1026
1027    /// Encodes a `RecordBatch` into a flatbuffer IPC message and fills `sink` with the
1028    /// serialised buffer data.
1029    ///
1030    /// Returns the total body length written to `sink` (including per-buffer alignment
1031    /// padding).
1032    ///
1033    /// The FlatBuffer [`crate::Message`] is located in `ipc_write_context`'s
1034    /// [`FlatBufferBuilder`] finished bytes. A successful Result from this function
1035    /// guarantees the builder is in a finished state to call
1036    /// [`FlatBufferBuilder::finished_data`].
1037    fn record_batch_to_bytes(
1038        &self,
1039        batch: &RecordBatch,
1040        write_options: &IpcWriteOptions,
1041        ipc_write_context: &mut IpcWriteContext,
1042        sink: &mut IpcBodySink<'_>,
1043    ) -> Result<usize, ArrowError> {
1044        // Reset the fbb
1045        ipc_write_context.mut_fbb().reset();
1046
1047        let EncodedRecordBatchMeta {
1048            fb_offset: record_batch,
1049            body_len,
1050        } = self.encode_record_batch_data(
1051            batch.columns().iter().map(|array| array.to_data()),
1052            batch.num_rows() as i64,
1053            write_options,
1054            ipc_write_context,
1055            sink,
1056        )?;
1057
1058        let fbb = ipc_write_context.mut_fbb();
1059        let mut message = crate::MessageBuilder::new(fbb);
1060        message.add_version(write_options.metadata_version);
1061        message.add_header_type(crate::MessageHeader::RecordBatch);
1062        message.add_bodyLength(body_len as i64);
1063        message.add_header(record_batch.as_union_value());
1064        let root = message.finish();
1065        fbb.finish(root, None);
1066
1067        Ok(body_len)
1068    }
1069
1070    /// Write dictionary values into two sets of bytes, one for the header (crate::Message) and the
1071    /// other for the data
1072    fn dictionary_batch_to_bytes(
1073        &self,
1074        dict: DictionaryToEncode,
1075        write_options: &IpcWriteOptions,
1076        ipc_write_context: &mut IpcWriteContext,
1077    ) -> Result<EncodedData, ArrowError> {
1078        let mut arrow_data: Vec<u8> = vec![];
1079        self.dictionary_batch_to_sink(
1080            &dict,
1081            write_options,
1082            ipc_write_context,
1083            &mut IpcBodySink::Write(&mut arrow_data),
1084        )?;
1085
1086        Ok(EncodedData {
1087            ipc_message: ipc_write_context.mut_fbb().finished_data().to_vec(),
1088            arrow_data,
1089        })
1090    }
1091
1092    /// Encodes a dictionary batch's flatbuffer metadata into `ipc_write_context`'s
1093    /// builder (read back via `ipc_write_context.mut_fbb().finished_data()`) and fills
1094    /// `sink` with its body buffers.
1095    ///
1096    /// Returns the total body length written to `sink` (including per-buffer alignment
1097    /// padding).
1098    fn dictionary_batch_to_sink(
1099        &self,
1100        dict: &DictionaryToEncode,
1101        write_options: &IpcWriteOptions,
1102        ipc_write_context: &mut IpcWriteContext,
1103        sink: &mut IpcBodySink<'_>,
1104    ) -> Result<usize, ArrowError> {
1105        // Reset fbb
1106        ipc_write_context.mut_fbb().reset();
1107
1108        // A dictionary batch is a record batch (the single column of dictionary
1109        // values) wrapped in a DictionaryBatch, so we share the record batch body
1110        // and table encoding and only differ in the message framing.
1111        let EncodedRecordBatchMeta {
1112            fb_offset: record_batch,
1113            body_len,
1114        } = self.encode_record_batch_data(
1115            std::iter::once(dict.data.clone()),
1116            dict.data.len() as i64,
1117            write_options,
1118            ipc_write_context,
1119            sink,
1120        )?;
1121
1122        let fbb = ipc_write_context.mut_fbb();
1123        let root = {
1124            let mut batch_builder = crate::DictionaryBatchBuilder::new(fbb);
1125            batch_builder.add_id(dict.id);
1126            batch_builder.add_data(record_batch);
1127            batch_builder.add_isDelta(dict.is_delta);
1128            batch_builder.finish().as_union_value()
1129        };
1130
1131        let root = {
1132            let mut message_builder = crate::MessageBuilder::new(fbb);
1133            message_builder.add_version(write_options.metadata_version);
1134            message_builder.add_header_type(crate::MessageHeader::DictionaryBatch);
1135            message_builder.add_bodyLength(body_len as i64);
1136            message_builder.add_header(root);
1137            message_builder.finish()
1138        };
1139
1140        fbb.finish(root, None);
1141
1142        Ok(body_len)
1143    }
1144
1145    /// Encode the body buffers of a single record batch into `sink` and build the
1146    /// flatbuffer `RecordBatch` table into `ipc_write_context`'s builder.
1147    ///
1148    /// This is the shared core of [`Self::record_batch_to_bytes`] and
1149    /// [`Self::dictionary_batch_to_sink`]: a dictionary message embeds a record
1150    /// batch (its single column of values), so both encode the same `RecordBatch`
1151    /// table and only differ in how they wrap the returned offset in a message.
1152    ///
1153    /// `columns` yields the column data in IPC buffer order and `row_count` is the
1154    /// logical length recorded in the `RecordBatch` header. The caller is responsible
1155    /// for wrapping the returned [`EncodedRecordBatchMeta::fb_offset`] in a `Message`
1156    /// and calling [`FlatBufferBuilder::finish`].
1157    #[inline]
1158    fn encode_record_batch_data(
1159        &self,
1160        columns: impl IntoIterator<Item = ArrayData>,
1161        row_count: i64,
1162        write_options: &IpcWriteOptions,
1163        ipc_write_context: &mut IpcWriteContext,
1164        sink: &mut IpcBodySink<'_>,
1165    ) -> Result<EncodedRecordBatchMeta, ArrowError> {
1166        let batch_compression_type = write_options.batch_compression_type;
1167
1168        let compression = batch_compression_type.map(|batch_compression_type| {
1169            let fbb = ipc_write_context.mut_fbb();
1170            let mut c = crate::BodyCompressionBuilder::new(fbb);
1171            c.add_method(crate::BodyCompressionMethod::BUFFER);
1172            c.add_codec(batch_compression_type);
1173            c.finish()
1174        });
1175
1176        let batch_compression_level = write_options.batch_compression_level;
1177        let compression_codec: Option<CompressionCodec> = batch_compression_type
1178            .map(|compression_type| match batch_compression_level {
1179                Some(level) => {
1180                    CompressionCodec::try_new_with_compression_level(compression_type, level)
1181                }
1182                None => compression_type.try_into(),
1183            })
1184            .transpose()?;
1185
1186        let mut meta = IpcMetadataBuilder::default();
1187        let mut variadic_buffer_counts: Vec<i64> = vec![];
1188        let mut offset = 0i64;
1189
1190        for array_data in columns {
1191            offset = write_array_data(
1192                &array_data,
1193                &mut meta,
1194                sink,
1195                offset,
1196                compression_codec,
1197                ipc_write_context,
1198                write_options,
1199            )?;
1200            append_variadic_buffer_counts(&mut variadic_buffer_counts, &array_data);
1201        }
1202
1203        // Each buffer is padded to the alignment as it is written, so `offset` is
1204        // already a multiple of the alignment -- the body needs no trailing padding.
1205        debug_assert_eq!(offset % write_options.alignment as i64, 0);
1206        let body_len = offset as usize;
1207
1208        let fbb = ipc_write_context.mut_fbb();
1209        let buffers = fbb.create_vector(&meta.buffers);
1210        let nodes = fbb.create_vector(&meta.nodes);
1211        let variadic_buffer = if variadic_buffer_counts.is_empty() {
1212            None
1213        } else {
1214            Some(fbb.create_vector(&variadic_buffer_counts))
1215        };
1216
1217        let mut batch_builder = crate::RecordBatchBuilder::new(fbb);
1218        batch_builder.add_length(row_count);
1219        batch_builder.add_nodes(nodes);
1220        batch_builder.add_buffers(buffers);
1221        if let Some(c) = compression {
1222            batch_builder.add_compression(c);
1223        }
1224        if let Some(v) = variadic_buffer {
1225            batch_builder.add_variadicBufferCounts(v);
1226        }
1227        let record_batch = batch_builder.finish();
1228
1229        Ok(EncodedRecordBatchMeta {
1230            fb_offset: record_batch,
1231            body_len,
1232        })
1233    }
1234}
1235
1236fn ensure_supported_ipc_schema(schema: &Schema) -> Result<(), ArrowError> {
1237    schema
1238        .fields()
1239        .iter()
1240        .try_for_each(|field| ensure_supported_ipc_data_type(field.name(), field.data_type()))
1241}
1242
1243fn ensure_supported_ipc_data_type(
1244    field_name: &str,
1245    data_type: &DataType,
1246) -> Result<(), ArrowError> {
1247    match data_type {
1248        DataType::Dictionary(_, value_type)
1249            if matches!(value_type.as_ref(), DataType::Dictionary(_, _)) =>
1250        {
1251            Err(ArrowError::InvalidArgumentError(format!(
1252                "Arrow IPC field metadata cannot encode direct dictionary-of-dictionary values for field {field_name:?}"
1253            )))
1254        }
1255        DataType::Dictionary(_, value_type) => {
1256            ensure_supported_ipc_data_type(field_name, value_type)
1257        }
1258        DataType::Struct(fields) => fields
1259            .iter()
1260            .try_for_each(|field| ensure_supported_ipc_data_type(field.name(), field.data_type())),
1261        DataType::RunEndEncoded(_, field)
1262        | DataType::List(field)
1263        | DataType::LargeList(field)
1264        | DataType::ListView(field)
1265        | DataType::LargeListView(field)
1266        | DataType::FixedSizeList(field, _)
1267        | DataType::Map(field, _) => {
1268            ensure_supported_ipc_data_type(field.name(), field.data_type())
1269        }
1270        DataType::Union(fields, _) => fields.iter().try_for_each(|(_, field)| {
1271            ensure_supported_ipc_data_type(field.name(), field.data_type())
1272        }),
1273        _ => Ok(()),
1274    }
1275}
1276
1277fn append_variadic_buffer_counts(counts: &mut Vec<i64>, array: &ArrayData) {
1278    match array.data_type() {
1279        DataType::BinaryView | DataType::Utf8View => {
1280            // The spec documents the counts only includes the variadic buffers, not the view/null buffers.
1281            // https://arrow.apache.org/docs/format/Columnar.html#variadic-buffers
1282            counts.push(array.buffers().len() as i64 - 1);
1283        }
1284        DataType::Dictionary(_, _) => {
1285            // Do nothing
1286            // Dictionary values are handled in special dictionary messages
1287        }
1288        _ => {
1289            for child in array.child_data() {
1290                append_variadic_buffer_counts(counts, child)
1291            }
1292        }
1293    }
1294}
1295
1296pub(crate) fn unslice_run_array(arr: ArrayData) -> Result<ArrayData, ArrowError> {
1297    match arr.data_type() {
1298        DataType::RunEndEncoded(k, _) => match k.data_type() {
1299            DataType::Int16 => {
1300                Ok(into_zero_offset_run_array(RunArray::<Int16Type>::from(arr))?.into_data())
1301            }
1302            DataType::Int32 => {
1303                Ok(into_zero_offset_run_array(RunArray::<Int32Type>::from(arr))?.into_data())
1304            }
1305            DataType::Int64 => {
1306                Ok(into_zero_offset_run_array(RunArray::<Int64Type>::from(arr))?.into_data())
1307            }
1308            d => unreachable!("Unexpected data type {d}"),
1309        },
1310        d => Err(ArrowError::InvalidArgumentError(format!(
1311            "The given array is not a run array. Data type of given array: {d}"
1312        ))),
1313    }
1314}
1315
1316// Returns a `RunArray` with zero offset and length matching the last value
1317// in run_ends array.
1318fn into_zero_offset_run_array<R: RunEndIndexType>(
1319    run_array: RunArray<R>,
1320) -> Result<RunArray<R>, ArrowError> {
1321    // Empty slices have no physical runs, regardless of their offset.
1322    if run_array.is_empty() {
1323        return Ok(ArrayData::new_empty(run_array.data_type()).into());
1324    }
1325
1326    let run_ends = run_array.run_ends();
1327    if run_ends.offset() == 0 && run_ends.max_value() == run_ends.len() {
1328        return Ok(run_array);
1329    }
1330
1331    // The physical index of original run_ends array from which the `ArrayData`is sliced.
1332    let start_physical_index = run_ends.get_start_physical_index();
1333
1334    // The physical index of original run_ends array until which the `ArrayData`is sliced.
1335    let end_physical_index = run_ends.get_end_physical_index();
1336
1337    let physical_length = end_physical_index - start_physical_index + 1;
1338
1339    // build new run_ends array by subtracting offset from run ends.
1340    let offset = R::Native::usize_as(run_ends.offset());
1341    let mut run_ends_values = Vec::<R::Native>::with_capacity(physical_length);
1342    for run_end_value in &run_ends.values()[start_physical_index..end_physical_index] {
1343        run_ends_values.push(run_end_value.sub_wrapping(offset));
1344    }
1345    run_ends_values.push(R::Native::from_usize(run_array.len()).unwrap());
1346    let offset_buffer = Buffer::from_vec(run_ends_values);
1347    let new_run_ends = unsafe {
1348        // Safety:
1349        // The function builds a valid run_ends array and hence need not be validated.
1350        ArrayDataBuilder::new(R::DATA_TYPE)
1351            .len(physical_length)
1352            .add_buffer(offset_buffer)
1353            .build_unchecked()
1354    };
1355
1356    // build new values by slicing physical indices.
1357    let new_values = run_array
1358        .values()
1359        .slice(start_physical_index, physical_length)
1360        .into_data();
1361
1362    let builder = ArrayDataBuilder::new(run_array.data_type().clone())
1363        .len(run_array.len())
1364        .add_child_data(new_run_ends)
1365        .add_child_data(new_values);
1366    let array_data = unsafe {
1367        // Safety:
1368        //  This function builds a valid run array and hence can skip validation.
1369        builder.build_unchecked()
1370    };
1371    Ok(array_data.into())
1372}
1373
1374/// Controls how dictionaries are handled in Arrow IPC messages
1375#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
1376pub enum DictionaryHandling {
1377    /// Send the entire dictionary every time it is encountered (default)
1378    #[default]
1379    Resend,
1380    /// Send only new dictionary values since the last batch (delta encoding)
1381    ///
1382    /// When a dictionary is first encountered, the entire dictionary is sent.
1383    /// For subsequent batches, only values that are new (not previously sent)
1384    /// are transmitted with the `isDelta` flag set to true.
1385    Delta,
1386}
1387
1388/// Describes what kind of update took place after a call to [`DictionaryTracker::insert`].
1389#[derive(Debug, Clone)]
1390pub enum DictionaryUpdate {
1391    /// No dictionary was written, the dictionary was identical to what was already
1392    /// in the tracker.
1393    None,
1394    /// No dictionary was present in the tracker
1395    New,
1396    /// Dictionary was replaced with the new data
1397    Replaced,
1398    /// Dictionary was updated, ArrayData is the delta between old and new
1399    Delta(ArrayData),
1400}
1401
1402/// Keeps track of dictionaries that have been written, to avoid emitting the same dictionary
1403/// multiple times.
1404///
1405/// Can optionally error if an update to an existing dictionary is attempted, which
1406/// isn't allowed in the `FileWriter`.
1407#[derive(Debug)]
1408pub struct DictionaryTracker {
1409    // NOTE: When adding fields, update the clear() method accordingly.
1410    written: HashMap<i64, ArrayData>,
1411    dict_ids: Vec<i64>,
1412    error_on_replacement: bool,
1413}
1414
1415impl DictionaryTracker {
1416    /// Create a new [`DictionaryTracker`].
1417    ///
1418    /// If `error_on_replacement`
1419    /// is true, an error will be generated if an update to an
1420    /// existing dictionary is attempted.
1421    pub fn new(error_on_replacement: bool) -> Self {
1422        Self {
1423            written: HashMap::new(),
1424            dict_ids: Vec::new(),
1425            error_on_replacement,
1426        }
1427    }
1428
1429    /// Record and return the next dictionary ID.
1430    pub fn next_dict_id(&mut self) -> i64 {
1431        let next = self
1432            .dict_ids
1433            .last()
1434            .copied()
1435            .map(|i| i + 1)
1436            .unwrap_or_default();
1437
1438        self.dict_ids.push(next);
1439        next
1440    }
1441
1442    /// Return the sequence of dictionary IDs in the order they should be observed while
1443    /// traversing the schema
1444    pub fn dict_id(&mut self) -> &[i64] {
1445        &self.dict_ids
1446    }
1447
1448    /// Keep track of the dictionary with the given ID and values. Behavior:
1449    ///
1450    /// * If this ID has been written already and has the same data, return `Ok(false)` to indicate
1451    ///   that the dictionary was not actually inserted (because it's already been seen).
1452    /// * If this ID has been written already but with different data, and this tracker is
1453    ///   configured to return an error, return an error.
1454    /// * If the tracker has not been configured to error on replacement or this dictionary
1455    ///   has never been seen before, return `Ok(true)` to indicate that the dictionary was just
1456    ///   inserted.
1457    #[deprecated(since = "56.1.0", note = "Use `insert_column` instead")]
1458    pub fn insert(&mut self, dict_id: i64, column: &ArrayRef) -> Result<bool, ArrowError> {
1459        let dict_data = column.to_data();
1460        let dict_values = &dict_data.child_data()[0];
1461
1462        // If a dictionary with this id was already emitted, check if it was the same.
1463        if let Some(last) = self.written.get(&dict_id) {
1464            if ArrayData::ptr_eq(&last.child_data()[0], dict_values) {
1465                // Same dictionary values => no need to emit it again
1466                return Ok(false);
1467            }
1468            if self.error_on_replacement {
1469                // If error on replacement perform a logical comparison
1470                if last.child_data()[0] == *dict_values {
1471                    // Same dictionary values => no need to emit it again
1472                    return Ok(false);
1473                }
1474                return Err(ArrowError::InvalidArgumentError(
1475                    "Dictionary replacement detected when writing IPC file format. \
1476                     Arrow IPC files only support a single dictionary for a given field \
1477                     across all batches."
1478                        .to_string(),
1479                ));
1480            }
1481        }
1482
1483        self.written.insert(dict_id, dict_data);
1484        Ok(true)
1485    }
1486
1487    /// Keep track of the dictionary with the given ID and values. The return
1488    /// value indicates what, if any, update to the internal map took place
1489    /// and how it should be interpreted based on the `dict_handling` parameter.
1490    ///
1491    /// # Returns
1492    ///
1493    /// * `Ok(Dictionary::New)` - If the dictionary was not previously written
1494    /// * `Ok(Dictionary::Replaced)` - If the dictionary was previously written
1495    ///   with completely different data, or if the data is a delta of the existing,
1496    ///   but with `dict_handling` set to `DictionaryHandling::Resend`
1497    /// * `Ok(Dictionary::Delta)` - If the dictionary was previously written, but
1498    ///   the new data is a delta of the old and the `dict_handling` is set to
1499    ///   `DictionaryHandling::Delta`
1500    /// * `Err(e)` - If the dictionary was previously written with different data,
1501    ///   and `error_on_replacement` is set to `true`.
1502    pub fn insert_column(
1503        &mut self,
1504        dict_id: i64,
1505        column: &ArrayRef,
1506        dict_handling: DictionaryHandling,
1507    ) -> Result<DictionaryUpdate, ArrowError> {
1508        let new_data = column.to_data();
1509        let new_values = &new_data.child_data()[0];
1510
1511        // If there is no existing dictionary with this ID, we always insert
1512        let Some(old) = self.written.get(&dict_id) else {
1513            self.written.insert(dict_id, new_data);
1514            return Ok(DictionaryUpdate::New);
1515        };
1516
1517        // Fast path - If the array data points to the same buffer as the
1518        // existing then they're the same.
1519        let old_values = &old.child_data()[0];
1520        if ArrayData::ptr_eq(old_values, new_values) {
1521            return Ok(DictionaryUpdate::None);
1522        }
1523
1524        // Slow path - Compare the dictionaries value by value
1525        let comparison = compare_dictionaries(old_values, new_values);
1526        if matches!(comparison, DictionaryComparison::Equal) {
1527            return Ok(DictionaryUpdate::None);
1528        }
1529
1530        const REPLACEMENT_ERROR: &str = "Dictionary replacement detected when writing IPC file format. \
1531                 Arrow IPC files only support a single dictionary for a given field \
1532                 across all batches.";
1533
1534        match comparison {
1535            DictionaryComparison::NotEqual => {
1536                if self.error_on_replacement {
1537                    return Err(ArrowError::InvalidArgumentError(
1538                        REPLACEMENT_ERROR.to_string(),
1539                    ));
1540                }
1541
1542                self.written.insert(dict_id, new_data);
1543                Ok(DictionaryUpdate::Replaced)
1544            }
1545            DictionaryComparison::Delta => match dict_handling {
1546                DictionaryHandling::Resend => {
1547                    if self.error_on_replacement {
1548                        return Err(ArrowError::InvalidArgumentError(
1549                            REPLACEMENT_ERROR.to_string(),
1550                        ));
1551                    }
1552
1553                    self.written.insert(dict_id, new_data);
1554                    Ok(DictionaryUpdate::Replaced)
1555                }
1556                DictionaryHandling::Delta => {
1557                    let delta =
1558                        new_values.slice(old_values.len(), new_values.len() - old_values.len());
1559                    self.written.insert(dict_id, new_data);
1560                    Ok(DictionaryUpdate::Delta(delta))
1561                }
1562            },
1563            DictionaryComparison::Equal => unreachable!("Already checked equal case"),
1564        }
1565    }
1566
1567    /// Clears the state of the dictionary tracker.
1568    ///
1569    /// This allows the dictionary tracker to be reused for a new IPC stream while avoiding the
1570    /// allocation cost of creating a new instance. This method should not be called if
1571    /// the dictionary tracker will be used to continue writing to an existing IPC stream.
1572    pub fn clear(&mut self) {
1573        self.dict_ids.clear();
1574        self.written.clear();
1575    }
1576}
1577
1578/// Describes how two dictionary arrays compare to each other.
1579#[derive(Debug, Clone)]
1580enum DictionaryComparison {
1581    /// Neither a delta, nor an exact match
1582    NotEqual,
1583    /// Exact element-wise match
1584    Equal,
1585    /// The two arrays are dictionary deltas of each other, meaning the first
1586    /// is a prefix of the second.
1587    Delta,
1588}
1589
1590// Compares two dictionaries and returns a [`DictionaryComparison`].
1591fn compare_dictionaries(old: &ArrayData, new: &ArrayData) -> DictionaryComparison {
1592    // Check for exact match
1593    let existing_len = old.len();
1594    let new_len = new.len();
1595    if existing_len == new_len {
1596        return if *old == *new {
1597            DictionaryComparison::Equal
1598        } else {
1599            DictionaryComparison::NotEqual
1600        };
1601    }
1602
1603    // Can't be a delta if the new is shorter than the existing
1604    if new_len < existing_len {
1605        return DictionaryComparison::NotEqual;
1606    }
1607
1608    // Check for delta
1609    if new.slice(0, existing_len) == *old {
1610        return DictionaryComparison::Delta;
1611    }
1612
1613    DictionaryComparison::NotEqual
1614}
1615
1616/// Arrow File Writer
1617///
1618/// Writes Arrow [`RecordBatch`]es in the [IPC File Format].
1619///
1620/// # See Also
1621///
1622/// * [`StreamWriter`] for writing IPC Streams
1623///
1624/// # Example
1625/// ```
1626/// # use arrow_array::record_batch;
1627/// # use arrow_ipc::writer::FileWriter;
1628/// # let mut file = vec![]; // mimic a file for the example
1629/// let batch = record_batch!(("a", Int32, [1, 2, 3])).unwrap();
1630/// // create a new writer, the schema must be known in advance
1631/// let mut writer = FileWriter::try_new(&mut file, &batch.schema()).unwrap();
1632/// // write each batch to the underlying writer
1633/// writer.write(&batch).unwrap();
1634/// // When all batches are written, call finish to flush all buffers
1635/// writer.finish().unwrap();
1636/// ```
1637/// [IPC File Format]: https://arrow.apache.org/docs/format/Columnar.html#ipc-file-format
1638pub struct FileWriter<W> {
1639    /// The object to write to
1640    writer: W,
1641    /// IPC write options
1642    write_options: IpcWriteOptions,
1643    /// A reference to the schema, used in validating record batches
1644    schema: SchemaRef,
1645    /// The number of bytes between each block of bytes, as an offset for random access
1646    block_offsets: usize,
1647    /// Dictionary blocks that will be written as part of the IPC footer
1648    dictionary_blocks: Vec<crate::Block>,
1649    /// Record blocks that will be written as part of the IPC footer
1650    record_blocks: Vec<crate::Block>,
1651    /// Whether the writer footer has been written, and the writer is finished
1652    finished: bool,
1653    /// Keeps track of dictionaries that have been written
1654    dictionary_tracker: DictionaryTracker,
1655    /// User level customized metadata
1656    custom_metadata: Metadata,
1657
1658    data_gen: IpcDataGenerator,
1659
1660    /// Holds reusable scratch state shared across all messages -- including the
1661    /// [`FlatBufferBuilder`] whose internal buffer is reused to avoid reallocating
1662    /// on every batch.
1663    ipc_write_context: IpcWriteContext,
1664}
1665
1666impl<W: Write> FileWriter<BufWriter<W>> {
1667    /// Try to create a new file writer with the writer wrapped in a BufWriter.
1668    ///
1669    /// See [`FileWriter::try_new`] for an unbuffered version.
1670    pub fn try_new_buffered(writer: W, schema: &Schema) -> Result<Self, ArrowError> {
1671        Self::try_new(BufWriter::new(writer), schema)
1672    }
1673}
1674
1675impl<W: Write> FileWriter<W> {
1676    /// Try to create a new writer, with the schema written as part of the header
1677    ///
1678    /// Note the created writer is not buffered. See [`FileWriter::try_new_buffered`] for details.
1679    ///
1680    /// # Errors
1681    ///
1682    /// An [`Err`] may be returned if writing the header to the writer fails.
1683    pub fn try_new(writer: W, schema: &Schema) -> Result<Self, ArrowError> {
1684        let write_options = IpcWriteOptions::default();
1685        Self::try_new_with_options(writer, schema, write_options)
1686    }
1687
1688    /// Try to create a new writer with IpcWriteOptions
1689    ///
1690    /// Note the created writer is not buffered. See [`FileWriter::try_new_buffered`] for details.
1691    ///
1692    /// # Errors
1693    ///
1694    /// An [`Err`] may be returned if writing the header to the writer fails.
1695    pub fn try_new_with_options(
1696        mut writer: W,
1697        schema: &Schema,
1698        write_options: IpcWriteOptions,
1699    ) -> Result<Self, ArrowError> {
1700        ensure_supported_ipc_schema(schema)?;
1701
1702        let data_gen = IpcDataGenerator::default();
1703        // write magic to header aligned on alignment boundary
1704        let pad_len = pad_to_alignment(write_options.alignment, super::ARROW_MAGIC.len());
1705        let header_size = super::ARROW_MAGIC.len() + pad_len;
1706        writer.write_all(&super::ARROW_MAGIC)?;
1707        writer.write_all(&PADDING[..pad_len])?;
1708        // write the schema, set the written bytes to the schema + header
1709        let mut dictionary_tracker = DictionaryTracker::new(true);
1710        let encoded_message = data_gen.schema_to_bytes_with_dictionary_tracker(
1711            schema,
1712            &mut dictionary_tracker,
1713            &write_options,
1714        );
1715        let (meta, data) = write_message(&mut writer, encoded_message, &write_options)?;
1716        Ok(Self {
1717            writer,
1718            write_options,
1719            schema: Arc::new(schema.clone()),
1720            block_offsets: meta + data + header_size,
1721            dictionary_blocks: vec![],
1722            record_blocks: vec![],
1723            finished: false,
1724            dictionary_tracker,
1725            custom_metadata: Default::default(),
1726            data_gen,
1727            ipc_write_context: IpcWriteContext::default(),
1728        })
1729    }
1730
1731    /// Adds a key-value pair to the [FileWriter]'s custom metadata
1732    pub fn write_metadata(&mut self, key: impl Into<String>, value: impl Into<String>) {
1733        self.custom_metadata.insert(key.into(), value.into());
1734    }
1735
1736    /// Write a record batch to the file
1737    pub fn write(&mut self, batch: &RecordBatch) -> Result<(), ArrowError> {
1738        if self.finished {
1739            return Err(ArrowError::IpcError(
1740                "Cannot write record batch to file writer as it is closed".to_string(),
1741            ));
1742        }
1743
1744        let meta = self.data_gen.write(
1745            batch,
1746            &mut self.dictionary_tracker,
1747            &self.write_options,
1748            &mut self.ipc_write_context,
1749            &mut self.writer,
1750        )?;
1751
1752        for (header_len, body_len) in meta.dictionary_block_sizes {
1753            let block = crate::Block::new(
1754                self.block_offsets as i64,
1755                header_len as i32,
1756                body_len as i64,
1757            );
1758            self.dictionary_blocks.push(block);
1759            self.block_offsets += header_len + body_len;
1760        }
1761
1762        // add a record block for the footer
1763        let block = crate::Block::new(
1764            self.block_offsets as i64,
1765            meta.padded_header_len as i32,
1766            meta.body_len as i64,
1767        );
1768        self.record_blocks.push(block);
1769        self.block_offsets += meta.padded_header_len + meta.body_len;
1770        Ok(())
1771    }
1772
1773    /// Write footer and closing tag, then mark the writer as done
1774    pub fn finish(&mut self) -> Result<(), ArrowError> {
1775        if self.finished {
1776            return Err(ArrowError::IpcError(
1777                "Cannot write footer to file writer as it is closed".to_string(),
1778            ));
1779        }
1780
1781        // write EOS
1782        self.writer.write_eos(&self.write_options)?;
1783
1784        let mut fbb = FlatBufferBuilder::new();
1785        let dictionaries = fbb.create_vector(&self.dictionary_blocks);
1786        let record_batches = fbb.create_vector(&self.record_blocks);
1787
1788        // dictionaries are already written, so we can reset dictionary tracker to reuse for schema
1789        self.dictionary_tracker.clear();
1790        let schema = IpcSchemaEncoder::new()
1791            .with_dictionary_tracker(&mut self.dictionary_tracker)
1792            .schema_to_fb_offset(&mut fbb, &self.schema);
1793        let fb_custom_metadata = (!self.custom_metadata.is_empty())
1794            .then(|| crate::convert::metadata_to_fb(&mut fbb, &self.custom_metadata));
1795
1796        let root = {
1797            let mut footer_builder = crate::FooterBuilder::new(&mut fbb);
1798            footer_builder.add_version(self.write_options.metadata_version);
1799            footer_builder.add_schema(schema);
1800            footer_builder.add_dictionaries(dictionaries);
1801            footer_builder.add_recordBatches(record_batches);
1802            if let Some(fb_custom_metadata) = fb_custom_metadata {
1803                footer_builder.add_custom_metadata(fb_custom_metadata);
1804            }
1805            footer_builder.finish()
1806        };
1807        fbb.finish(root, None);
1808        let footer_data = fbb.finished_data();
1809        self.writer.write_all(footer_data)?;
1810        self.writer
1811            .write_all(&(footer_data.len() as i32).to_le_bytes())?;
1812        self.writer.write_all(&super::ARROW_MAGIC)?;
1813        self.writer.flush()?;
1814        self.finished = true;
1815
1816        Ok(())
1817    }
1818
1819    /// Returns the arrow [`SchemaRef`] for this arrow file.
1820    pub fn schema(&self) -> &SchemaRef {
1821        &self.schema
1822    }
1823
1824    /// Gets a reference to the underlying writer.
1825    pub fn get_ref(&self) -> &W {
1826        &self.writer
1827    }
1828
1829    /// Gets a mutable reference to the underlying writer.
1830    ///
1831    /// It is inadvisable to directly write to the underlying writer.
1832    pub fn get_mut(&mut self) -> &mut W {
1833        &mut self.writer
1834    }
1835
1836    /// Flush the underlying writer.
1837    ///
1838    /// Both the BufWriter and the underlying writer are flushed.
1839    pub fn flush(&mut self) -> Result<(), ArrowError> {
1840        self.writer.flush()?;
1841        Ok(())
1842    }
1843
1844    /// Unwraps the underlying writer.
1845    ///
1846    /// The writer is flushed and the FileWriter is finished before returning.
1847    ///
1848    /// # Errors
1849    ///
1850    /// An [`Err`] may be returned if an error occurs while finishing the StreamWriter
1851    /// or while flushing the writer.
1852    pub fn into_inner(mut self) -> Result<W, ArrowError> {
1853        if !self.finished {
1854            // `finish` flushes the writer.
1855            self.finish()?;
1856        }
1857        Ok(self.writer)
1858    }
1859}
1860
1861impl<W: Write> RecordBatchWriter for FileWriter<W> {
1862    fn write(&mut self, batch: &RecordBatch) -> Result<(), ArrowError> {
1863        self.write(batch)
1864    }
1865
1866    fn close(mut self) -> Result<(), ArrowError> {
1867        self.finish()
1868    }
1869}
1870
1871/// Arrow IPC stream encoder.
1872///
1873/// Encodes Arrow [`RecordBatch`]es to byte buffers using the [IPC Streaming Format],
1874/// without performing any IO.
1875///
1876/// The returned [`Buffer`]s are ordered and should be written to the destination
1877/// stream in order. Uncompressed record batch body buffers can share the original
1878/// Arrow buffers instead of being copied into an intermediate contiguous buffer.
1879///
1880/// # Example
1881/// ```
1882/// # use arrow_array::record_batch;
1883/// # use arrow_ipc::writer::StreamEncoder;
1884/// # use arrow_schema::ArrowError;
1885/// # fn main() -> Result<(), ArrowError> {
1886/// let batch = record_batch!(("a", Int32, [1, 2, 3]))?;
1887///
1888/// let mut encoder = StreamEncoder::try_new(&batch.schema())?;
1889/// let mut stream = vec![];
1890/// for buffer in encoder.encode(&batch)? {
1891///     stream.extend_from_slice(buffer.as_slice());
1892/// }
1893/// for buffer in encoder.finish()? {
1894///     stream.extend_from_slice(buffer.as_slice());
1895/// }
1896/// # Ok(())
1897/// # }
1898/// ```
1899pub struct StreamEncoder {
1900    schema: Schema,
1901    /// IPC write options
1902    write_options: IpcWriteOptions,
1903    /// Whether the stream schema has been encoded
1904    schema_encoded: bool,
1905    /// Keeps track of dictionaries that have been encoded
1906    dictionary_tracker: DictionaryTracker,
1907    data_gen: IpcDataGenerator,
1908    ipc_write_context: IpcWriteContext,
1909}
1910
1911impl StreamEncoder {
1912    /// Try to create a new stream encoder.
1913    pub fn try_new(schema: &Schema) -> Result<Self, ArrowError> {
1914        let write_options = IpcWriteOptions::default();
1915        Self::try_new_with_options(schema, write_options)
1916    }
1917
1918    /// Try to create a new stream encoder with [`IpcWriteOptions`].
1919    pub fn try_new_with_options(
1920        schema: &Schema,
1921        write_options: IpcWriteOptions,
1922    ) -> Result<Self, ArrowError> {
1923        ensure_supported_ipc_schema(schema)?;
1924
1925        Ok(Self {
1926            schema: schema.clone(),
1927            write_options,
1928            schema_encoded: false,
1929            dictionary_tracker: DictionaryTracker::new(false),
1930            data_gen: IpcDataGenerator::default(),
1931            ipc_write_context: IpcWriteContext::default(),
1932        })
1933    }
1934
1935    /// Encode a [`RecordBatch`] into buffers.
1936    ///
1937    /// The first call also includes the IPC stream schema message before the
1938    /// record batch message. Later calls only include dictionary and record
1939    /// batch messages.
1940    ///
1941    /// # Errors
1942    ///
1943    /// Returns an error if encoding fails.
1944    pub fn encode(&mut self, batch: &RecordBatch) -> Result<Vec<Buffer>, ArrowError> {
1945        let mut out = vec![];
1946        self.encode_schema(&mut out)?;
1947        self.data_gen.encode_to_buffers(
1948            batch,
1949            &mut self.dictionary_tracker,
1950            &self.write_options,
1951            &mut self.ipc_write_context,
1952            &mut out,
1953        )?;
1954        Ok(out)
1955    }
1956
1957    /// Encode the end-of-stream marker.
1958    ///
1959    /// If no batches have been encoded, this also emits the IPC stream schema
1960    /// message so the returned buffers form a valid empty IPC stream.
1961    ///
1962    /// # Errors
1963    ///
1964    /// Returns an error if encoding the schema or end-of-stream marker fails.
1965    pub fn finish(mut self) -> Result<Vec<Buffer>, ArrowError> {
1966        let mut out = vec![];
1967        self.encode_schema(&mut out)?;
1968        let mut sink = Buffers { out: &mut out };
1969        sink.write_eos(&self.write_options)?;
1970        Ok(out)
1971    }
1972
1973    fn encode_schema(&mut self, out: &mut Vec<Buffer>) -> Result<(), ArrowError> {
1974        if !self.schema_encoded {
1975            let encoded_message = self.data_gen.schema_to_bytes_with_dictionary_tracker(
1976                &self.schema,
1977                &mut self.dictionary_tracker,
1978                &self.write_options,
1979            );
1980            let mut sink = Buffers { out };
1981            sink.write_encoded_data(encoded_message, &self.write_options)?;
1982            self.schema_encoded = true;
1983        }
1984        Ok(())
1985    }
1986}
1987
1988/// Arrow Stream Writer
1989///
1990/// Writes Arrow [`RecordBatch`]es to bytes using the [IPC Streaming Format].
1991///
1992/// # See Also
1993///
1994/// * [`FileWriter`] for writing IPC Files
1995///
1996/// # Example - Basic usage
1997/// ```
1998/// # use arrow_array::record_batch;
1999/// # use arrow_ipc::writer::StreamWriter;
2000/// # let mut stream = vec![]; // mimic a stream for the example
2001/// let batch = record_batch!(("a", Int32, [1, 2, 3])).unwrap();
2002/// // create a new writer, the schema must be known in advance
2003/// let mut writer = StreamWriter::try_new(&mut stream, &batch.schema()).unwrap();
2004/// // write each batch to the underlying stream
2005/// writer.write(&batch).unwrap();
2006/// // When all batches are written, call finish to flush all buffers
2007/// writer.finish().unwrap();
2008/// ```
2009/// # Example - Efficient delta dictionaries
2010/// ```
2011/// # use arrow_array::record_batch;
2012/// # use arrow_ipc::writer::{StreamWriter, IpcWriteOptions};
2013/// # use arrow_ipc::writer::DictionaryHandling;
2014/// # use arrow_schema::{DataType, Field, Schema, SchemaRef};
2015/// # use arrow_array::{
2016/// #    builder::StringDictionaryBuilder, types::Int32Type, Array, ArrayRef, DictionaryArray,
2017/// #    RecordBatch, StringArray,
2018/// # };
2019/// # use std::sync::Arc;
2020///
2021/// let schema = Arc::new(Schema::new(vec![Field::new(
2022///    "col1",
2023///    DataType::Dictionary(Box::from(DataType::Int32), Box::from(DataType::Utf8)),
2024///    true,
2025/// )]));
2026///
2027/// let mut builder = StringDictionaryBuilder::<arrow_array::types::Int32Type>::new();
2028///
2029/// // `finish_preserve_values` will keep the dictionary values along with their
2030/// // key assignments so that they can be re-used in the next batch.
2031/// builder.append("a").unwrap();
2032/// builder.append("b").unwrap();
2033/// let array1 = builder.finish_preserve_values();
2034/// let batch1 = RecordBatch::try_new(schema.clone(), vec![Arc::new(array1) as ArrayRef]).unwrap();
2035///
2036/// // In this batch, 'a' will have the same dictionary key as 'a' in the previous batch,
2037/// // and 'd' will take the next available key.
2038/// builder.append("a").unwrap();
2039/// builder.append("d").unwrap();
2040/// let array2 = builder.finish_preserve_values();
2041/// let batch2 = RecordBatch::try_new(schema.clone(), vec![Arc::new(array2) as ArrayRef]).unwrap();
2042///
2043/// let mut stream = vec![];
2044/// // You must set `.with_dictionary_handling(DictionaryHandling::Delta)` to
2045/// // enable delta dictionaries in the writer
2046/// let options = IpcWriteOptions::default().with_dictionary_handling(DictionaryHandling::Delta);
2047/// let mut writer = StreamWriter::try_new_with_options(&mut stream, &schema, options).unwrap();
2048///
2049/// // When writing the first batch, a dictionary message with 'a' and 'b' will be written
2050/// // prior to the record batch.
2051/// writer.write(&batch1).unwrap();
2052/// // With the second batch only a delta dictionary with 'd' will be written
2053/// // prior to the record batch. This is only possible with `finish_preserve_values`.
2054/// // Without it, 'a' and 'd' in this batch would have different keys than the
2055/// // first batch and so we'd have to send a replacement dictionary with new keys
2056/// // for both.
2057/// writer.write(&batch2).unwrap();
2058/// writer.finish().unwrap();
2059/// ```
2060/// [IPC Streaming Format]: https://arrow.apache.org/docs/format/Columnar.html#ipc-streaming-format
2061pub struct StreamWriter<W> {
2062    /// The object to write to
2063    writer: W,
2064    /// IPC write options
2065    write_options: IpcWriteOptions,
2066    /// Whether the writer footer has been written, and the writer is finished
2067    finished: bool,
2068    /// Keeps track of dictionaries that have been written
2069    dictionary_tracker: DictionaryTracker,
2070
2071    data_gen: IpcDataGenerator,
2072
2073    /// Reusable scratch state shared across all messages.
2074    ipc_write_context: IpcWriteContext,
2075}
2076
2077impl<W: Write> StreamWriter<BufWriter<W>> {
2078    /// Try to create a new stream writer with the writer wrapped in a BufWriter.
2079    ///
2080    /// See [`StreamWriter::try_new`] for an unbuffered version.
2081    pub fn try_new_buffered(writer: W, schema: &Schema) -> Result<Self, ArrowError> {
2082        Self::try_new(BufWriter::new(writer), schema)
2083    }
2084}
2085
2086impl<W: Write> StreamWriter<W> {
2087    /// Try to create a new writer, with the schema written as part of the header.
2088    ///
2089    /// Note that there is no internal buffering. See also [`StreamWriter::try_new_buffered`].
2090    ///
2091    /// # Errors
2092    ///
2093    /// An [`Err`] may be returned if writing the header to the writer fails.
2094    pub fn try_new(writer: W, schema: &Schema) -> Result<Self, ArrowError> {
2095        let write_options = IpcWriteOptions::default();
2096        Self::try_new_with_options(writer, schema, write_options)
2097    }
2098
2099    /// Try to create a new writer with [`IpcWriteOptions`].
2100    ///
2101    /// # Errors
2102    ///
2103    /// An [`Err`] may be returned if writing the header to the writer fails.
2104    pub fn try_new_with_options(
2105        mut writer: W,
2106        schema: &Schema,
2107        write_options: IpcWriteOptions,
2108    ) -> Result<Self, ArrowError> {
2109        ensure_supported_ipc_schema(schema)?;
2110
2111        let data_gen = IpcDataGenerator::default();
2112        let mut dictionary_tracker = DictionaryTracker::new(false);
2113
2114        // write the schema, set the written bytes to the schema
2115        let encoded_message = data_gen.schema_to_bytes_with_dictionary_tracker(
2116            schema,
2117            &mut dictionary_tracker,
2118            &write_options,
2119        );
2120        write_message(&mut writer, encoded_message, &write_options)?;
2121        Ok(Self {
2122            writer,
2123            write_options,
2124            finished: false,
2125            dictionary_tracker,
2126            data_gen,
2127            ipc_write_context: IpcWriteContext::default(),
2128        })
2129    }
2130
2131    /// Write a record batch to the stream
2132    pub fn write(&mut self, batch: &RecordBatch) -> Result<(), ArrowError> {
2133        if self.finished {
2134            return Err(ArrowError::IpcError(
2135                "Cannot write record batch to stream writer as it is closed".to_string(),
2136            ));
2137        }
2138
2139        self.data_gen.write(
2140            batch,
2141            &mut self.dictionary_tracker,
2142            &self.write_options,
2143            &mut self.ipc_write_context,
2144            &mut self.writer,
2145        )?;
2146        Ok(())
2147    }
2148
2149    /// Write continuation bytes, and mark the stream as done
2150    pub fn finish(&mut self) -> Result<(), ArrowError> {
2151        if self.finished {
2152            return Err(ArrowError::IpcError(
2153                "Cannot write footer to stream writer as it is closed".to_string(),
2154            ));
2155        }
2156
2157        self.writer.write_eos(&self.write_options)?;
2158        self.writer.flush()?;
2159
2160        self.finished = true;
2161
2162        Ok(())
2163    }
2164
2165    /// Gets a reference to the underlying writer.
2166    pub fn get_ref(&self) -> &W {
2167        &self.writer
2168    }
2169
2170    /// Gets a mutable reference to the underlying writer.
2171    ///
2172    /// It is inadvisable to directly write to the underlying writer.
2173    pub fn get_mut(&mut self) -> &mut W {
2174        &mut self.writer
2175    }
2176
2177    /// Flush the underlying writer.
2178    ///
2179    /// Both the BufWriter and the underlying writer are flushed.
2180    pub fn flush(&mut self) -> Result<(), ArrowError> {
2181        self.writer.flush()?;
2182        Ok(())
2183    }
2184
2185    /// Unwraps the the underlying writer.
2186    ///
2187    /// The writer is flushed and the StreamWriter is finished before returning.
2188    ///
2189    /// # Errors
2190    ///
2191    /// An [`Err`] may be returned if an error occurs while finishing the StreamWriter
2192    /// or while flushing the writer.
2193    ///
2194    /// # Example
2195    ///
2196    /// ```
2197    /// # use arrow_ipc::writer::{StreamWriter, IpcWriteOptions};
2198    /// # use arrow_ipc::MetadataVersion;
2199    /// # use arrow_schema::{ArrowError, Schema};
2200    /// # fn main() -> Result<(), ArrowError> {
2201    /// // The result we expect from an empty schema
2202    /// let expected = vec![
2203    ///     255, 255, 255, 255,  48,   0,   0,   0,
2204    ///      16,   0,   0,   0,   0,   0,  10,   0,
2205    ///      12,   0,  10,   0,   9,   0,   4,   0,
2206    ///      10,   0,   0,   0,  16,   0,   0,   0,
2207    ///       0,   1,   4,   0,   8,   0,   8,   0,
2208    ///       0,   0,   4,   0,   8,   0,   0,   0,
2209    ///       4,   0,   0,   0,   0,   0,   0,   0,
2210    ///     255, 255, 255, 255,   0,   0,   0,   0
2211    /// ];
2212    ///
2213    /// let schema = Schema::empty();
2214    /// let buffer: Vec<u8> = Vec::new();
2215    /// let options = IpcWriteOptions::try_new(8, false, MetadataVersion::V5)?;
2216    /// let stream_writer = StreamWriter::try_new_with_options(buffer, &schema, options)?;
2217    ///
2218    /// assert_eq!(stream_writer.into_inner()?, expected);
2219    /// # Ok(())
2220    /// # }
2221    /// ```
2222    pub fn into_inner(mut self) -> Result<W, ArrowError> {
2223        if !self.finished {
2224            // `finish` flushes.
2225            self.finish()?;
2226        }
2227        Ok(self.writer)
2228    }
2229}
2230
2231impl<W: Write> RecordBatchWriter for StreamWriter<W> {
2232    fn write(&mut self, batch: &RecordBatch) -> Result<(), ArrowError> {
2233        self.write(batch)
2234    }
2235
2236    fn close(mut self) -> Result<(), ArrowError> {
2237        self.finish()
2238    }
2239}
2240
2241/// Stores the encoded data, which is an crate::Message, and optional Arrow data
2242pub struct EncodedData {
2243    /// An encoded crate::Message
2244    pub ipc_message: Vec<u8>,
2245    /// Arrow buffers to be written, should be an empty vec for schema messages
2246    pub arrow_data: Vec<u8>,
2247}
2248
2249/// Write a message's IPC data and buffers, returning metadata and buffer data lengths written
2250pub fn write_message<W: Write>(
2251    mut writer: W,
2252    encoded: EncodedData,
2253    write_options: &IpcWriteOptions,
2254) -> Result<(usize, usize), ArrowError> {
2255    writer.write_encoded_data(encoded, write_options)
2256}
2257
2258/// In V4, null types have no validity bitmap
2259/// In V5 and later, null and union types have no validity bitmap
2260/// Run end encoded type has no validity bitmap.
2261fn has_validity_bitmap(data_type: &DataType, write_options: &IpcWriteOptions) -> bool {
2262    if write_options.metadata_version < crate::MetadataVersion::V5 {
2263        !matches!(data_type, DataType::Null)
2264    } else {
2265        !matches!(
2266            data_type,
2267            DataType::Null | DataType::Union(_, _) | DataType::RunEndEncoded(_, _)
2268        )
2269    }
2270}
2271
2272/// Whether to truncate the buffer
2273#[inline]
2274fn buffer_need_truncate(
2275    array_offset: usize,
2276    buffer: &Buffer,
2277    spec: &BufferSpec,
2278    min_length: usize,
2279) -> bool {
2280    spec != &BufferSpec::AlwaysNull && (array_offset != 0 || min_length < buffer.len())
2281}
2282
2283/// Returns byte width for a buffer spec. Only for `BufferSpec::FixedWidth`.
2284#[inline]
2285fn get_buffer_element_width(spec: &BufferSpec) -> usize {
2286    match spec {
2287        BufferSpec::FixedWidth { byte_width, .. } => *byte_width,
2288        _ => 0,
2289    }
2290}
2291
2292/// Common functionality for re-encoding offsets. Returns the new offsets as well as
2293/// original start offset and length for use in slicing child data.
2294fn reencode_offsets<O: OffsetSizeTrait>(
2295    offsets: &Buffer,
2296    data: &ArrayData,
2297) -> (Buffer, usize, usize) {
2298    let offsets_slice: &[O] = offsets.typed_data::<O>();
2299    let offset_slice = &offsets_slice[data.offset()..data.offset() + data.len() + 1];
2300
2301    let start_offset = offset_slice.first().unwrap();
2302    let end_offset = offset_slice.last().unwrap();
2303
2304    let offsets = match start_offset.as_usize() {
2305        0 => {
2306            let size = size_of::<O>();
2307            offsets.slice_with_length(data.offset() * size, (data.len() + 1) * size)
2308        }
2309        _ => offset_slice.iter().map(|x| *x - *start_offset).collect(),
2310    };
2311
2312    let start_offset = start_offset.as_usize();
2313    let end_offset = end_offset.as_usize();
2314
2315    (offsets, start_offset, end_offset - start_offset)
2316}
2317
2318/// Returns the values and offsets [`Buffer`] for a ByteArray with offset type `O`
2319///
2320/// In particular, this handles re-encoding the offsets if they don't start at `0`,
2321/// slicing the values buffer as appropriate. This helps reduce the encoded
2322/// size of sliced arrays, as values that have been sliced away are not encoded
2323/// Returns the offsets and values buffers, in that order.
2324fn get_byte_array_buffers<O: OffsetSizeTrait>(data: &ArrayData) -> [Buffer; 2] {
2325    if data.is_empty() {
2326        // As per specification, offsets buffer has N+1 elements.
2327        // So an empty array should still be encoded with a single 0 offset.
2328        let mut offsets = MutableBuffer::new(size_of::<O>());
2329        offsets.extend_from_slice(O::usize_as(0).to_byte_slice());
2330        return [offsets.into(), MutableBuffer::new(0).into()];
2331    }
2332
2333    let (offsets, original_start_offset, len) = reencode_offsets::<O>(&data.buffers()[0], data);
2334    let values = data.buffers()[1].slice_with_length(original_start_offset, len);
2335    [offsets, values]
2336}
2337
2338/// Similar logic as [`get_byte_array_buffers()`] but slices the child array instead
2339/// of a values buffer.
2340fn get_list_array_buffers<O: OffsetSizeTrait>(data: &ArrayData) -> (Buffer, ArrayData) {
2341    if data.is_empty() {
2342        // As per specification, offsets buffer has N+1 elements.
2343        // So an empty array should still be encoded with a single 0 offset.
2344        let mut offsets = MutableBuffer::new(size_of::<O>());
2345        offsets.extend_from_slice(O::usize_as(0).to_byte_slice());
2346        return (offsets.into(), data.child_data()[0].slice(0, 0));
2347    }
2348
2349    let (offsets, original_start_offset, len) = reencode_offsets::<O>(&data.buffers()[0], data);
2350    let child_data = data.child_data()[0].slice(original_start_offset, len);
2351    (offsets, child_data)
2352}
2353
2354/// Returns the offsets, sizes, and child data buffers for a ListView array.
2355///
2356/// Unlike List arrays, ListView arrays store both offsets and sizes explicitly,
2357/// and offsets can be non-monotonic. When slicing, we simply pass through the
2358/// offsets and sizes without re-encoding, and do not slice the child data.
2359fn get_list_view_array_buffers<O: OffsetSizeTrait>(
2360    data: &ArrayData,
2361) -> (Buffer, Buffer, ArrayData) {
2362    if data.is_empty() {
2363        return (
2364            MutableBuffer::new(0).into(),
2365            MutableBuffer::new(0).into(),
2366            data.child_data()[0].slice(0, 0),
2367        );
2368    }
2369
2370    let offsets = &data.buffers()[0];
2371    let sizes = &data.buffers()[1];
2372
2373    let element_size = std::mem::size_of::<O>();
2374    let offsets_slice =
2375        offsets.slice_with_length(data.offset() * element_size, data.len() * element_size);
2376    let sizes_slice =
2377        sizes.slice_with_length(data.offset() * element_size, data.len() * element_size);
2378
2379    let child_data = data.child_data()[0].clone();
2380
2381    (offsets_slice, sizes_slice, child_data)
2382}
2383
2384/// Returns the sliced views [`Buffer`] for a BinaryView/Utf8View array.
2385///
2386/// The views buffer is sliced to only include views in the valid range based on
2387/// the array's offset and length. This helps reduce the encoded size of sliced
2388/// arrays
2389///
2390fn get_or_truncate_buffer(array_data: &ArrayData) -> Buffer {
2391    let buffer = &array_data.buffers()[0];
2392    let layout = layout(array_data.data_type());
2393    let spec = &layout.buffers[0];
2394
2395    let byte_width = get_buffer_element_width(spec);
2396    let min_length = array_data.len() * byte_width;
2397    if buffer_need_truncate(array_data.offset(), buffer, spec, min_length) {
2398        let byte_offset = array_data.offset() * byte_width;
2399        let buffer_length = min(min_length, buffer.len() - byte_offset);
2400        buffer.slice_with_length(byte_offset, buffer_length)
2401    } else {
2402        buffer.clone()
2403    }
2404}
2405
2406/// Recursively encodes `array_data` into its IPC representation.
2407///
2408/// Output goes to two separate channels:
2409/// - `meta`: accumulates IPC metadata (`nodes` and `buffers`) for the flatbuffer header.
2410/// - `sink`: the raw Arrow data bytes that form the IPC message body.
2411fn write_array_data(
2412    array_data: &ArrayData,
2413    meta: &mut IpcMetadataBuilder,
2414    sink: &mut IpcBodySink<'_>,
2415    offset: i64,
2416    compression_codec: Option<CompressionCodec>,
2417    ipc_write_context: &mut IpcWriteContext,
2418    write_options: &IpcWriteOptions,
2419) -> Result<i64, ArrowError> {
2420    let mut offset = offset;
2421    let num_rows = array_data.len();
2422    let null_count = array_data.null_count();
2423    if !matches!(array_data.data_type(), DataType::Null) {
2424        meta.nodes
2425            .push(crate::FieldNode::new(num_rows as i64, null_count as i64));
2426    } else {
2427        // NullArray's null_count equals to len, but ArrayData null_count is always 0.
2428        meta.nodes
2429            .push(crate::FieldNode::new(num_rows as i64, num_rows as i64));
2430    }
2431    if has_validity_bitmap(array_data.data_type(), write_options) {
2432        // Fields with null_count == 0 may write a 0-length validity buffer
2433        // instead of a materialized bitmap (see FieldNode in Message.fbs)
2434        let null_buffer = match array_data.nulls() {
2435            Some(buffer) if null_count > 0 => buffer.inner().sliced(),
2436            _ => MutableBuffer::new(0).into(),
2437        };
2438
2439        offset = encode_sink_buffer(
2440            null_buffer,
2441            meta,
2442            sink,
2443            offset,
2444            compression_codec,
2445            ipc_write_context,
2446            write_options.alignment,
2447        )?;
2448    }
2449
2450    let data_type = array_data.data_type();
2451    if matches!(data_type, DataType::Binary | DataType::Utf8) {
2452        for buffer in get_byte_array_buffers::<i32>(array_data) {
2453            offset = encode_sink_buffer(
2454                buffer,
2455                meta,
2456                sink,
2457                offset,
2458                compression_codec,
2459                ipc_write_context,
2460                write_options.alignment,
2461            )?;
2462        }
2463    } else if matches!(data_type, DataType::BinaryView | DataType::Utf8View) {
2464        // Slicing the views buffer is safe and easy,
2465        // but pruning unneeded data buffers is much more nuanced since it's complicated to prove that no views reference the pruned buffers
2466        //
2467        // Current implementation just serialize the raw arrays as given and not try to optimize anything.
2468        // If users wants to "compact" the arrays prior to sending them over IPC,
2469        // they should consider the gc API suggested in #5513
2470        let views = get_or_truncate_buffer(array_data);
2471        offset = encode_sink_buffer(
2472            views,
2473            meta,
2474            sink,
2475            offset,
2476            compression_codec,
2477            ipc_write_context,
2478            write_options.alignment,
2479        )?;
2480
2481        for buffer in array_data.buffers().iter().skip(1) {
2482            offset = encode_sink_buffer(
2483                buffer.clone(),
2484                meta,
2485                sink,
2486                offset,
2487                compression_codec,
2488                ipc_write_context,
2489                write_options.alignment,
2490            )?;
2491        }
2492    } else if matches!(data_type, DataType::LargeBinary | DataType::LargeUtf8) {
2493        for buffer in get_byte_array_buffers::<i64>(array_data) {
2494            offset = encode_sink_buffer(
2495                buffer,
2496                meta,
2497                sink,
2498                offset,
2499                compression_codec,
2500                ipc_write_context,
2501                write_options.alignment,
2502            )?;
2503        }
2504    } else if DataType::is_numeric(data_type)
2505        || DataType::is_temporal(data_type)
2506        || matches!(
2507            array_data.data_type(),
2508            DataType::FixedSizeBinary(_) | DataType::Dictionary(_, _)
2509        )
2510    {
2511        // Truncate values
2512        assert_eq!(array_data.buffers().len(), 1);
2513
2514        let buffer = get_or_truncate_buffer(array_data);
2515        offset = encode_sink_buffer(
2516            buffer,
2517            meta,
2518            sink,
2519            offset,
2520            compression_codec,
2521            ipc_write_context,
2522            write_options.alignment,
2523        )?;
2524    } else if matches!(data_type, DataType::Boolean) {
2525        // Bools are special because the payload (= 1 bit) is smaller than the physical container elements (= bytes).
2526        // The array data may not start at the physical boundary of the underlying buffer, so we need to shift bits around.
2527        assert_eq!(array_data.buffers().len(), 1);
2528
2529        let buffer = &array_data.buffers()[0];
2530        let buffer = buffer.bit_slice(array_data.offset(), array_data.len());
2531        offset = encode_sink_buffer(
2532            buffer,
2533            meta,
2534            sink,
2535            offset,
2536            compression_codec,
2537            ipc_write_context,
2538            write_options.alignment,
2539        )?;
2540    } else if matches!(
2541        data_type,
2542        DataType::List(_) | DataType::LargeList(_) | DataType::Map(_, _)
2543    ) {
2544        assert_eq!(array_data.buffers().len(), 1);
2545        assert_eq!(array_data.child_data().len(), 1);
2546
2547        // Truncate offsets and the child data to avoid writing unnecessary data
2548        let (offsets, sliced_child_data) = match data_type {
2549            DataType::List(_) => get_list_array_buffers::<i32>(array_data),
2550            DataType::Map(_, _) => get_list_array_buffers::<i32>(array_data),
2551            DataType::LargeList(_) => get_list_array_buffers::<i64>(array_data),
2552            _ => unreachable!(),
2553        };
2554        offset = encode_sink_buffer(
2555            offsets,
2556            meta,
2557            sink,
2558            offset,
2559            compression_codec,
2560            ipc_write_context,
2561            write_options.alignment,
2562        )?;
2563        offset = write_array_data(
2564            &sliced_child_data,
2565            meta,
2566            sink,
2567            offset,
2568            compression_codec,
2569            ipc_write_context,
2570            write_options,
2571        )?;
2572        return Ok(offset);
2573    } else if matches!(
2574        data_type,
2575        DataType::ListView(_) | DataType::LargeListView(_)
2576    ) {
2577        assert_eq!(array_data.buffers().len(), 2); // offsets + sizes
2578        assert_eq!(array_data.child_data().len(), 1);
2579
2580        let (offsets, sizes, child_data) = match data_type {
2581            DataType::ListView(_) => get_list_view_array_buffers::<i32>(array_data),
2582            DataType::LargeListView(_) => get_list_view_array_buffers::<i64>(array_data),
2583            _ => unreachable!(),
2584        };
2585
2586        offset = encode_sink_buffer(
2587            offsets,
2588            meta,
2589            sink,
2590            offset,
2591            compression_codec,
2592            ipc_write_context,
2593            write_options.alignment,
2594        )?;
2595        offset = encode_sink_buffer(
2596            sizes,
2597            meta,
2598            sink,
2599            offset,
2600            compression_codec,
2601            ipc_write_context,
2602            write_options.alignment,
2603        )?;
2604
2605        offset = write_array_data(
2606            &child_data,
2607            meta,
2608            sink,
2609            offset,
2610            compression_codec,
2611            ipc_write_context,
2612            write_options,
2613        )?;
2614        return Ok(offset);
2615    } else if let DataType::FixedSizeList(_, fixed_size) = data_type {
2616        assert_eq!(array_data.child_data().len(), 1);
2617        let fixed_size = *fixed_size as usize;
2618
2619        let child_offset = array_data.offset() * fixed_size;
2620        let child_length = array_data.len() * fixed_size;
2621        let child_data = array_data.child_data()[0].slice(child_offset, child_length);
2622
2623        offset = write_array_data(
2624            &child_data,
2625            meta,
2626            sink,
2627            offset,
2628            compression_codec,
2629            ipc_write_context,
2630            write_options,
2631        )?;
2632        return Ok(offset);
2633    } else if matches!(data_type, DataType::Union(_, _)) {
2634        for (buffer, byte_width) in array_data
2635            .buffers()
2636            .iter()
2637            .zip([size_of::<i8>(), size_of::<i32>()])
2638        {
2639            let buffer =
2640                buffer.slice_with_length(array_data.offset() * byte_width, num_rows * byte_width);
2641            offset = encode_sink_buffer(
2642                buffer,
2643                meta,
2644                sink,
2645                offset,
2646                compression_codec,
2647                ipc_write_context,
2648                write_options.alignment,
2649            )?;
2650        }
2651    } else {
2652        for buffer in array_data.buffers() {
2653            offset = encode_sink_buffer(
2654                buffer.clone(),
2655                meta,
2656                sink,
2657                offset,
2658                compression_codec,
2659                ipc_write_context,
2660                write_options.alignment,
2661            )?;
2662        }
2663    }
2664
2665    match array_data.data_type() {
2666        DataType::Dictionary(_, _) => {}
2667        DataType::RunEndEncoded(_, _) => {
2668            // unslice the run encoded array.
2669            let arr = unslice_run_array(array_data.clone())?;
2670            // recursively write out nested structures
2671            for data_ref in arr.child_data() {
2672                // write the nested data (e.g list data)
2673                offset = write_array_data(
2674                    data_ref,
2675                    meta,
2676                    sink,
2677                    offset,
2678                    compression_codec,
2679                    ipc_write_context,
2680                    write_options,
2681                )?;
2682            }
2683        }
2684        DataType::Union(_, UnionMode::Sparse) => {
2685            for data_ref in array_data.child_data() {
2686                let child_data = data_ref.slice(array_data.offset(), num_rows);
2687                offset = write_array_data(
2688                    &child_data,
2689                    meta,
2690                    sink,
2691                    offset,
2692                    compression_codec,
2693                    ipc_write_context,
2694                    write_options,
2695                )?;
2696            }
2697        }
2698        _ => {
2699            // recursively write out nested structures
2700            for data_ref in array_data.child_data() {
2701                // write the nested data (e.g list data)
2702                offset = write_array_data(
2703                    data_ref,
2704                    meta,
2705                    sink,
2706                    offset,
2707                    compression_codec,
2708                    ipc_write_context,
2709                    write_options,
2710                )?;
2711            }
2712        }
2713    }
2714    Ok(offset)
2715}
2716
2717/// Encodes a single Arrow [`Buffer`] into the IPC body and records its metadata.
2718///
2719/// - `buffer`: the Arrow data buffer to encode (validity bitmap, offsets, values, etc.)
2720/// - `buffers`: in-progress list of IPC `Buffer` metadata entries (body offset + length) that
2721///   will eventually be serialised into the flatbuffer `RecordBatch` header.
2722/// - `sink`: destination for the actual encoded bytes; either a contiguous `Vec<u8>` for
2723///   in-memory writes, or a list of [`EncodedBuffer`] segments for deferred zero-copy streaming.
2724/// - `offset`: running byte offset into the IPC message body, used to compute the metadata entry.
2725/// - `compression_codec` / `ipc_write_context`: if `Some`, the buffer is compressed before
2726///   writing; `ipc_write_context` provides reusable scratch space across calls.
2727/// - `alignment`: each buffer is padded to this many bytes so the next buffer starts aligned.
2728///
2729/// Returns the updated `offset` (advanced by the encoded length plus any alignment padding).
2730fn encode_sink_buffer(
2731    buffer: Buffer,
2732    ipc_meta_data: &mut IpcMetadataBuilder,
2733    sink: &mut IpcBodySink<'_>,
2734    offset: i64,
2735    compression_codec: Option<CompressionCodec>,
2736    ipc_write_context: &mut IpcWriteContext,
2737    alignment: u8,
2738) -> Result<i64, ArrowError> {
2739    let (encoded, len) = match compression_codec {
2740        None => {
2741            let len = buffer.len() as i64;
2742            (EncodedBuffer::Raw(buffer), len)
2743        }
2744        Some(codec) => {
2745            let mut scratch = Vec::new();
2746            let written =
2747                codec.compress_to_vec(buffer.as_slice(), &mut scratch, ipc_write_context)?;
2748            let len = i64::try_from(written)
2749                .map_err(|e| ArrowError::InvalidArgumentError(format!("{e}")))?;
2750            (EncodedBuffer::Compressed(scratch), len)
2751        }
2752    };
2753
2754    let pad_len = pad_to_alignment(alignment, len as usize);
2755    sink.write(pad_len, encoded);
2756    ipc_meta_data.buffers.push(crate::Buffer::new(offset, len));
2757    Ok(offset + len + pad_len as i64)
2758}
2759
2760const PADDING: [u8; 64] = [0; 64];
2761
2762/// Estimates the number of [`EncodedBuffer`] segments that [`write_array_data`]
2763/// will produce for a column of the given type.
2764///
2765/// Based on the Arrow IPC buffer layout
2766/// (<https://arrow.apache.org/docs/format/Columnar.html#recordbatch-message>):
2767#[inline]
2768fn estimate_encoded_buffer_count(dt: &DataType) -> usize {
2769    match dt {
2770        DataType::Null => 0,
2771
2772        DataType::Binary | DataType::Utf8 | DataType::LargeBinary | DataType::LargeUtf8 => 3,
2773
2774        DataType::BinaryView | DataType::Utf8View => 3,
2775
2776        DataType::List(f) | DataType::LargeList(f) | DataType::Map(f, _) => {
2777            2 + estimate_encoded_buffer_count(f.data_type())
2778        }
2779
2780        DataType::ListView(f) | DataType::LargeListView(f) => {
2781            3 + estimate_encoded_buffer_count(f.data_type())
2782        }
2783
2784        DataType::FixedSizeList(f, _) => 1 + estimate_encoded_buffer_count(f.data_type()),
2785
2786        DataType::Struct(fields) => {
2787            1 + fields
2788                .iter()
2789                .map(|f| estimate_encoded_buffer_count(f.data_type()))
2790                .sum::<usize>()
2791        }
2792
2793        // Dictionary indices only; dictionary body is a separate IPC message.
2794        DataType::Dictionary(_, _) => 2,
2795
2796        DataType::Union(fields, UnionMode::Sparse) => {
2797            1 + fields
2798                .iter()
2799                .map(|(_, f)| estimate_encoded_buffer_count(f.data_type()))
2800                .sum::<usize>()
2801        }
2802        DataType::Union(fields, UnionMode::Dense) => {
2803            2 + fields
2804                .iter()
2805                .map(|(_, f)| estimate_encoded_buffer_count(f.data_type()))
2806                .sum::<usize>()
2807        }
2808
2809        DataType::RunEndEncoded(run_ends, values) => {
2810            estimate_encoded_buffer_count(run_ends.data_type())
2811                + estimate_encoded_buffer_count(values.data_type())
2812        }
2813        // Primitive, Bool, temporal, Decimal*, FixedSizeBinary: validity + values.
2814        _ => 2,
2815    }
2816}
2817
2818/// Calculate an alignment boundary and return the number of bytes needed to pad to the alignment boundary
2819#[inline]
2820fn pad_to_alignment(alignment: u8, len: usize) -> usize {
2821    let a = usize::from(alignment - 1);
2822    ((len + a) & !a) - len
2823}
2824
2825#[cfg(test)]
2826mod tests {
2827    use std::hash::Hasher;
2828    use std::io::Cursor;
2829    use std::io::Seek;
2830
2831    use arrow_array::builder::FixedSizeListBuilder;
2832    use arrow_array::builder::Float32Builder;
2833    use arrow_array::builder::Int64Builder;
2834    use arrow_array::builder::MapBuilder;
2835    use arrow_array::builder::StringViewBuilder;
2836    use arrow_array::builder::UnionBuilder;
2837    use arrow_array::builder::{
2838        GenericListBuilder, GenericListViewBuilder, ListBuilder, StringBuilder,
2839    };
2840    use arrow_array::builder::{PrimitiveRunBuilder, UInt32Builder};
2841    use arrow_array::types::*;
2842    use arrow_buffer::ScalarBuffer;
2843
2844    use crate::MetadataVersion;
2845    use crate::convert::try_fb_to_schema;
2846    use crate::reader::*;
2847    use crate::root_as_footer;
2848
2849    use super::*;
2850
2851    #[test]
2852    fn test_roundtrip_sliced_list_of_union() {
2853        fn decoded(array: &ArrayRef) -> Vec<Vec<Option<i32>>> {
2854            (0..array.len())
2855                .map(|row| {
2856                    let values = array.as_list::<i32>().value(row);
2857                    let union = values.as_union();
2858                    (0..union.len())
2859                        .map(|index| {
2860                            union
2861                                .value(index)
2862                                .as_primitive::<Int32Type>()
2863                                .iter()
2864                                .next()
2865                                .unwrap()
2866                        })
2867                        .collect()
2868                })
2869                .collect()
2870        }
2871
2872        for mode in [UnionMode::Dense, UnionMode::Sparse] {
2873            let fields = UnionFields::try_new(
2874                [3, 7],
2875                [
2876                    Field::new("a", DataType::Int32, true),
2877                    Field::new("b", DataType::Int32, true),
2878                ],
2879            )
2880            .unwrap();
2881            let offsets = (mode == UnionMode::Dense).then(|| vec![2, 0, 1].into());
2882            let union: ArrayRef = Arc::new(
2883                UnionArray::try_new(
2884                    fields,
2885                    vec![3_i8, 7, 3].into(),
2886                    offsets,
2887                    vec![
2888                        Arc::new(Int32Array::from(vec![Some(11), None, Some(33)])),
2889                        Arc::new(Int32Array::from(vec![Some(101), Some(202), None])),
2890                    ],
2891                )
2892                .unwrap(),
2893            );
2894            let item = Arc::new(Field::new("item", union.data_type().clone(), true));
2895            let list: ArrayRef = Arc::new(ListArray::new(
2896                item,
2897                arrow_buffer::OffsetBuffer::from_lengths([1, 2]),
2898                union,
2899                None,
2900            ));
2901            for (offset, len) in [(0, 2), (1, 1), (2, 0)] {
2902                let array = list.slice(offset, len);
2903                array.to_data().validate_full().unwrap();
2904                let expected = decoded(&array);
2905                let schema = Arc::new(Schema::new(vec![Field::new(
2906                    "v",
2907                    array.data_type().clone(),
2908                    true,
2909                )]));
2910                let batch = RecordBatch::try_new(schema, vec![array]).unwrap();
2911                for output in [
2912                    deserialize_file(serialize_file(&batch)),
2913                    deserialize_stream(serialize_stream(&batch)),
2914                ] {
2915                    output.column(0).to_data().validate_full().unwrap();
2916                    assert_eq!(decoded(output.column(0)), expected);
2917                }
2918            }
2919        }
2920    }
2921
2922    fn serialize_file(rb: &RecordBatch) -> Vec<u8> {
2923        let mut writer = FileWriter::try_new(vec![], rb.schema_ref()).unwrap();
2924        writer.write(rb).unwrap();
2925        writer.finish().unwrap();
2926        writer.into_inner().unwrap()
2927    }
2928
2929    fn deserialize_file(bytes: Vec<u8>) -> RecordBatch {
2930        let mut reader = FileReader::try_new(Cursor::new(bytes), None).unwrap();
2931        reader.next().unwrap().unwrap()
2932    }
2933
2934    fn serialize_stream(record: &RecordBatch) -> Vec<u8> {
2935        // Use 8-byte alignment so that the various `truncate_*` tests can be compactly written,
2936        // without needing to construct a giant array to spill over the 64-byte default alignment
2937        // boundary.
2938        const IPC_ALIGNMENT: usize = 8;
2939
2940        let mut stream_writer = StreamWriter::try_new_with_options(
2941            vec![],
2942            record.schema_ref(),
2943            IpcWriteOptions::try_new(IPC_ALIGNMENT, false, MetadataVersion::V5).unwrap(),
2944        )
2945        .unwrap();
2946        stream_writer.write(record).unwrap();
2947        stream_writer.finish().unwrap();
2948        stream_writer.into_inner().unwrap()
2949    }
2950
2951    fn deserialize_stream(bytes: Vec<u8>) -> RecordBatch {
2952        let mut stream_reader = StreamReader::try_new(Cursor::new(bytes), None).unwrap();
2953        stream_reader.next().unwrap().unwrap()
2954    }
2955
2956    /// Encodes record batches with [`StreamEncoder`] into one contiguous byte vector.
2957    ///
2958    /// This still exercises the sans-IO encoder path; `Vec<u8>` is only used
2959    /// as a convenient [`Write`] sink for comparing the resulting byte stream.
2960    fn encode_stream(
2961        schema: &Schema,
2962        batches: &[RecordBatch],
2963        options: IpcWriteOptions,
2964    ) -> Vec<u8> {
2965        let mut encoder = StreamEncoder::try_new_with_options(schema, options).unwrap();
2966        let mut bytes = Vec::new();
2967        for batch in batches {
2968            for buffer in encoder.encode(batch).unwrap() {
2969                bytes.write_all(buffer.as_slice()).unwrap();
2970            }
2971        }
2972        for buffer in encoder.finish().unwrap() {
2973            bytes.write_all(buffer.as_slice()).unwrap();
2974        }
2975        bytes
2976    }
2977
2978    fn write_stream(schema: &Schema, batches: &[RecordBatch], options: IpcWriteOptions) -> Vec<u8> {
2979        let mut bytes = Vec::new();
2980        let mut writer = StreamWriter::try_new_with_options(&mut bytes, schema, options).unwrap();
2981        for batch in batches {
2982            writer.write(batch).unwrap();
2983        }
2984        writer.finish().unwrap();
2985        bytes
2986    }
2987
2988    // StreamEncoder and StreamWriter currently use separate encoding paths, so these tests
2989    // verify the new sans-IO API preserves the existing IPC stream byte layout.
2990    #[tokio::test]
2991    async fn test_stream_encoder_async_writer_matches_stream_writer() {
2992        use tokio::io::{AsyncReadExt, AsyncWriteExt};
2993
2994        let batch = record_batch!(("a", Int32, [1, 2, 3]), ("b", Utf8, ["x", "y", "z"])).unwrap();
2995        let options = IpcWriteOptions::default();
2996        let expected = write_stream(
2997            batch.schema_ref(),
2998            std::slice::from_ref(&batch),
2999            options.clone(),
3000        );
3001
3002        let (mut sink, mut source) = tokio::io::duplex(64);
3003        let read = tokio::spawn(async move {
3004            let mut bytes = Vec::new();
3005            source.read_to_end(&mut bytes).await.unwrap();
3006            bytes
3007        });
3008
3009        let mut encoder = StreamEncoder::try_new_with_options(batch.schema_ref(), options).unwrap();
3010        for buffer in encoder.encode(&batch).unwrap() {
3011            sink.write_all(buffer.as_slice()).await.unwrap();
3012        }
3013        for buffer in encoder.finish().unwrap() {
3014            sink.write_all(buffer.as_slice()).await.unwrap();
3015        }
3016        sink.shutdown().await.unwrap();
3017
3018        let encoded = read.await.unwrap();
3019        assert_eq!(encoded, expected);
3020
3021        let mut reader = StreamReader::try_new(Cursor::new(encoded), None).unwrap();
3022        assert_eq!(reader.next().unwrap().unwrap(), batch);
3023        assert!(reader.next().is_none());
3024    }
3025
3026    #[test]
3027    fn test_stream_encoder_empty_stream_matches_stream_writer() {
3028        let schema = Schema::new(vec![Field::new("a", DataType::Int32, true)]);
3029        let options = IpcWriteOptions::default();
3030        let encoded = encode_stream(&schema, &[], options.clone());
3031        let written = write_stream(&schema, &[], options);
3032
3033        assert_eq!(encoded, written);
3034
3035        let mut reader = StreamReader::try_new(Cursor::new(encoded), None).unwrap();
3036        assert!(reader.next().is_none());
3037    }
3038
3039    #[test]
3040    fn test_stream_encoder_dictionary_batches_match_stream_writer() {
3041        let schema = Arc::new(Schema::new(vec![Field::new(
3042            "a",
3043            DataType::Dictionary(Box::new(DataType::UInt8), Box::new(DataType::Utf8)),
3044            false,
3045        )]));
3046        let batch = RecordBatch::try_new(
3047            schema.clone(),
3048            vec![Arc::new(DictionaryArray::new(
3049                UInt8Array::from_iter_values([0, 1, 0]),
3050                Arc::new(StringArray::from_iter_values(["a", "b"])),
3051            ))],
3052        )
3053        .unwrap();
3054        let options = IpcWriteOptions::default();
3055        let encoded = encode_stream(&schema, std::slice::from_ref(&batch), options.clone());
3056        let written = write_stream(&schema, std::slice::from_ref(&batch), options);
3057
3058        assert_eq!(encoded, written);
3059
3060        let mut reader = StreamReader::try_new(Cursor::new(encoded), None).unwrap();
3061        assert_eq!(reader.next().unwrap().unwrap(), batch);
3062        assert!(reader.next().is_none());
3063    }
3064
3065    #[test]
3066    #[cfg(feature = "lz4")]
3067    fn test_write_empty_record_batch_lz4_compression() {
3068        let schema = Schema::new(vec![Field::new("field1", DataType::Int32, true)]);
3069        let values: Vec<Option<i32>> = vec![];
3070        let array = Int32Array::from(values);
3071        let record_batch =
3072            RecordBatch::try_new(Arc::new(schema.clone()), vec![Arc::new(array)]).unwrap();
3073
3074        let mut file = tempfile::tempfile().unwrap();
3075
3076        {
3077            let write_option = IpcWriteOptions::try_new(8, false, crate::MetadataVersion::V5)
3078                .unwrap()
3079                .try_with_compression(Some(crate::CompressionType::LZ4_FRAME))
3080                .unwrap();
3081
3082            let mut writer =
3083                FileWriter::try_new_with_options(&mut file, &schema, write_option).unwrap();
3084            writer.write(&record_batch).unwrap();
3085            writer.finish().unwrap();
3086        }
3087        file.rewind().unwrap();
3088        {
3089            // read file
3090            let reader = FileReader::try_new(file, None).unwrap();
3091            for read_batch in reader {
3092                read_batch
3093                    .unwrap()
3094                    .columns()
3095                    .iter()
3096                    .zip(record_batch.columns())
3097                    .for_each(|(a, b)| {
3098                        assert_eq!(a.data_type(), b.data_type());
3099                        assert_eq!(a.len(), b.len());
3100                        assert_eq!(a.null_count(), b.null_count());
3101                    });
3102            }
3103        }
3104    }
3105
3106    #[test]
3107    #[cfg(feature = "lz4")]
3108    fn test_write_file_with_lz4_compression() {
3109        let schema = Schema::new(vec![Field::new("field1", DataType::Int32, true)]);
3110        let values: Vec<Option<i32>> = vec![Some(12), Some(1)];
3111        let array = Int32Array::from(values);
3112        let record_batch =
3113            RecordBatch::try_new(Arc::new(schema.clone()), vec![Arc::new(array)]).unwrap();
3114
3115        let mut file = tempfile::tempfile().unwrap();
3116        {
3117            let write_option = IpcWriteOptions::try_new(8, false, crate::MetadataVersion::V5)
3118                .unwrap()
3119                .try_with_compression(Some(crate::CompressionType::LZ4_FRAME))
3120                .unwrap();
3121
3122            let mut writer =
3123                FileWriter::try_new_with_options(&mut file, &schema, write_option).unwrap();
3124            writer.write(&record_batch).unwrap();
3125            writer.finish().unwrap();
3126        }
3127        file.rewind().unwrap();
3128        {
3129            // read file
3130            let reader = FileReader::try_new(file, None).unwrap();
3131            for read_batch in reader {
3132                read_batch
3133                    .unwrap()
3134                    .columns()
3135                    .iter()
3136                    .zip(record_batch.columns())
3137                    .for_each(|(a, b)| {
3138                        assert_eq!(a.data_type(), b.data_type());
3139                        assert_eq!(a.len(), b.len());
3140                        assert_eq!(a.null_count(), b.null_count());
3141                    });
3142            }
3143        }
3144    }
3145
3146    #[test]
3147    #[cfg(feature = "zstd")]
3148    fn test_write_file_with_zstd_compression() {
3149        let schema = Schema::new(vec![Field::new("field1", DataType::Int32, true)]);
3150        let values: Vec<Option<i32>> = vec![Some(12), Some(1)];
3151        let array = Int32Array::from(values);
3152        let record_batch =
3153            RecordBatch::try_new(Arc::new(schema.clone()), vec![Arc::new(array)]).unwrap();
3154        let mut file = tempfile::tempfile().unwrap();
3155        {
3156            let write_option = IpcWriteOptions::try_new(8, false, crate::MetadataVersion::V5)
3157                .unwrap()
3158                .try_with_compression(Some(crate::CompressionType::ZSTD))
3159                .unwrap()
3160                .try_with_compression_level(Some(1))
3161                .unwrap();
3162
3163            let mut writer =
3164                FileWriter::try_new_with_options(&mut file, &schema, write_option).unwrap();
3165            writer.write(&record_batch).unwrap();
3166            writer.finish().unwrap();
3167        }
3168        file.rewind().unwrap();
3169        {
3170            // read file
3171            let reader = FileReader::try_new(file, None).unwrap();
3172            for read_batch in reader {
3173                read_batch
3174                    .unwrap()
3175                    .columns()
3176                    .iter()
3177                    .zip(record_batch.columns())
3178                    .for_each(|(a, b)| {
3179                        assert_eq!(a.data_type(), b.data_type());
3180                        assert_eq!(a.len(), b.len());
3181                        assert_eq!(a.null_count(), b.null_count());
3182                    });
3183            }
3184        }
3185    }
3186
3187    #[test]
3188    fn test_write_file() {
3189        let schema = Schema::new(vec![Field::new("field1", DataType::UInt32, true)]);
3190        let values: Vec<Option<u32>> = vec![
3191            Some(999),
3192            None,
3193            Some(235),
3194            Some(123),
3195            None,
3196            None,
3197            None,
3198            None,
3199            None,
3200        ];
3201        let array1 = UInt32Array::from(values);
3202        let batch =
3203            RecordBatch::try_new(Arc::new(schema.clone()), vec![Arc::new(array1) as ArrayRef])
3204                .unwrap();
3205        let mut file = tempfile::tempfile().unwrap();
3206        {
3207            let mut writer = FileWriter::try_new(&mut file, &schema).unwrap();
3208
3209            writer.write(&batch).unwrap();
3210            writer.finish().unwrap();
3211        }
3212        file.rewind().unwrap();
3213
3214        {
3215            let mut reader = FileReader::try_new(file, None).unwrap();
3216            while let Some(Ok(read_batch)) = reader.next() {
3217                read_batch
3218                    .columns()
3219                    .iter()
3220                    .zip(batch.columns())
3221                    .for_each(|(a, b)| {
3222                        assert_eq!(a.data_type(), b.data_type());
3223                        assert_eq!(a.len(), b.len());
3224                        assert_eq!(a.null_count(), b.null_count());
3225                    });
3226            }
3227        }
3228    }
3229
3230    #[test]
3231    fn test_empty_utf8_ipc_writes_nonempty_offsets_buffer() {
3232        let name = StringArray::from(Vec::<String>::new());
3233        let [offsets, values] = get_byte_array_buffers::<i32>(&name.to_data());
3234
3235        assert_eq!(name.len(), 0);
3236        assert_eq!(
3237            offsets.len(),
3238            std::mem::size_of::<i32>(),
3239            "offsets buffer should contain one zero i32 offset"
3240        );
3241        assert_eq!(values.len(), 0, "values buffer should remain empty");
3242    }
3243
3244    #[test]
3245    fn test_empty_large_utf8_ipc_writes_nonempty_offsets_buffer() {
3246        let name = LargeStringArray::from(Vec::<String>::new());
3247        let [offsets, values] = get_byte_array_buffers::<i64>(&name.to_data());
3248
3249        assert_eq!(name.len(), 0);
3250        assert_eq!(
3251            offsets.len(),
3252            std::mem::size_of::<i64>(),
3253            "offsets buffer should contain one zero i64 offset"
3254        );
3255        assert_eq!(values.len(), 0, "values buffer should remain empty");
3256    }
3257
3258    #[test]
3259    fn test_empty_list_ipc_writes_nonempty_offsets_buffer() {
3260        let list = GenericListBuilder::<i32, _>::new(UInt32Builder::new()).finish();
3261        let (offsets, child_data) = get_list_array_buffers::<i32>(&list.to_data());
3262
3263        assert_eq!(list.len(), 0);
3264        assert_eq!(
3265            offsets.len(),
3266            std::mem::size_of::<i32>(),
3267            "offsets buffer should contain one zero i32 offset"
3268        );
3269        assert_eq!(child_data.len(), 0, "child data should remain empty");
3270    }
3271
3272    #[test]
3273    fn test_empty_large_list_ipc_writes_nonempty_offsets_buffer() {
3274        let list = GenericListBuilder::<i64, _>::new(UInt32Builder::new()).finish();
3275        let (offsets, child_data) = get_list_array_buffers::<i64>(&list.to_data());
3276
3277        assert_eq!(list.len(), 0);
3278        assert_eq!(
3279            offsets.len(),
3280            std::mem::size_of::<i64>(),
3281            "offsets buffer should contain one zero i64 offset"
3282        );
3283        assert_eq!(child_data.len(), 0, "child data should remain empty");
3284    }
3285
3286    fn write_null_file(options: IpcWriteOptions) {
3287        let schema = Schema::new(vec![
3288            Field::new("nulls", DataType::Null, true),
3289            Field::new("int32s", DataType::Int32, false),
3290            Field::new("nulls2", DataType::Null, true),
3291            Field::new("f64s", DataType::Float64, false),
3292        ]);
3293        let array1 = NullArray::new(32);
3294        let array2 = Int32Array::from(vec![1; 32]);
3295        let array3 = NullArray::new(32);
3296        let array4 = Float64Array::from(vec![f64::NAN; 32]);
3297        let batch = RecordBatch::try_new(
3298            Arc::new(schema.clone()),
3299            vec![
3300                Arc::new(array1) as ArrayRef,
3301                Arc::new(array2) as ArrayRef,
3302                Arc::new(array3) as ArrayRef,
3303                Arc::new(array4) as ArrayRef,
3304            ],
3305        )
3306        .unwrap();
3307        let mut file = tempfile::tempfile().unwrap();
3308        {
3309            let mut writer = FileWriter::try_new_with_options(&mut file, &schema, options).unwrap();
3310
3311            writer.write(&batch).unwrap();
3312            writer.finish().unwrap();
3313        }
3314
3315        file.rewind().unwrap();
3316
3317        {
3318            let reader = FileReader::try_new(file, None).unwrap();
3319            reader.for_each(|maybe_batch| {
3320                maybe_batch
3321                    .unwrap()
3322                    .columns()
3323                    .iter()
3324                    .zip(batch.columns())
3325                    .for_each(|(a, b)| {
3326                        assert_eq!(a.data_type(), b.data_type());
3327                        assert_eq!(a.len(), b.len());
3328                        assert_eq!(a.null_count(), b.null_count());
3329                    });
3330            });
3331        }
3332    }
3333    #[test]
3334    fn test_write_null_file_v4() {
3335        write_null_file(IpcWriteOptions::try_new(8, false, MetadataVersion::V4).unwrap());
3336        write_null_file(IpcWriteOptions::try_new(8, true, MetadataVersion::V4).unwrap());
3337        write_null_file(IpcWriteOptions::try_new(64, false, MetadataVersion::V4).unwrap());
3338        write_null_file(IpcWriteOptions::try_new(64, true, MetadataVersion::V4).unwrap());
3339    }
3340
3341    #[test]
3342    fn test_write_null_file_v5() {
3343        write_null_file(IpcWriteOptions::try_new(8, false, MetadataVersion::V5).unwrap());
3344        write_null_file(IpcWriteOptions::try_new(64, false, MetadataVersion::V5).unwrap());
3345    }
3346
3347    #[test]
3348    fn track_union_nested_dict() {
3349        let inner: DictionaryArray<Int32Type> = vec!["a", "b", "a"].into_iter().collect();
3350
3351        let array = Arc::new(inner) as ArrayRef;
3352
3353        // Dict field with id 2
3354        #[expect(deprecated)]
3355        let dctfield = Field::new_dict("dict", array.data_type().clone(), false, 0, false);
3356        let union_fields = std::iter::once((0, Arc::new(dctfield))).collect();
3357
3358        let types = [0, 0, 0].into_iter().collect::<ScalarBuffer<i8>>();
3359        let offsets = [0, 1, 2].into_iter().collect::<ScalarBuffer<i32>>();
3360
3361        let union = UnionArray::try_new(union_fields, types, Some(offsets), vec![array]).unwrap();
3362
3363        let schema = Arc::new(Schema::new(vec![Field::new(
3364            "union",
3365            union.data_type().clone(),
3366            false,
3367        )]));
3368
3369        let r#gen = IpcDataGenerator::default();
3370        let mut dict_tracker = DictionaryTracker::new(false);
3371        r#gen.schema_to_bytes_with_dictionary_tracker(
3372            &schema,
3373            &mut dict_tracker,
3374            &IpcWriteOptions::default(),
3375        );
3376
3377        let batch = RecordBatch::try_new(schema, vec![Arc::new(union)]).unwrap();
3378
3379        r#gen
3380            .encode(
3381                &batch,
3382                &mut dict_tracker,
3383                &Default::default(),
3384                &mut Default::default(),
3385            )
3386            .unwrap();
3387
3388        // The encoder will assign dict IDs itself to ensure uniqueness and ignore the dict ID in the schema
3389        // so we expect the dict will be keyed to 0
3390        assert!(dict_tracker.written.contains_key(&0));
3391    }
3392
3393    #[test]
3394    fn track_struct_nested_dict() {
3395        let inner: DictionaryArray<Int32Type> = vec!["a", "b", "a"].into_iter().collect();
3396
3397        let array = Arc::new(inner) as ArrayRef;
3398
3399        // Dict field with id 2
3400        #[expect(deprecated)]
3401        let dctfield = Arc::new(Field::new_dict(
3402            "dict",
3403            array.data_type().clone(),
3404            false,
3405            2,
3406            false,
3407        ));
3408
3409        let s = StructArray::from(vec![(dctfield, array)]);
3410        let struct_array = Arc::new(s) as ArrayRef;
3411
3412        let schema = Arc::new(Schema::new(vec![Field::new(
3413            "struct",
3414            struct_array.data_type().clone(),
3415            false,
3416        )]));
3417
3418        let r#gen = IpcDataGenerator::default();
3419        let mut dict_tracker = DictionaryTracker::new(false);
3420        r#gen.schema_to_bytes_with_dictionary_tracker(
3421            &schema,
3422            &mut dict_tracker,
3423            &IpcWriteOptions::default(),
3424        );
3425
3426        let batch = RecordBatch::try_new(schema, vec![struct_array]).unwrap();
3427
3428        r#gen
3429            .encode(
3430                &batch,
3431                &mut dict_tracker,
3432                &Default::default(),
3433                &mut Default::default(),
3434            )
3435            .unwrap();
3436
3437        assert!(dict_tracker.written.contains_key(&0));
3438    }
3439
3440    fn write_union_file(options: IpcWriteOptions) {
3441        let schema = Schema::new(vec![Field::new_union(
3442            "union",
3443            vec![0, 1],
3444            vec![
3445                Field::new("a", DataType::Int32, false),
3446                Field::new("c", DataType::Float64, false),
3447            ],
3448            UnionMode::Sparse,
3449        )]);
3450        let mut builder = UnionBuilder::with_capacity_sparse(5);
3451        builder.append::<Int32Type>("a", 1).unwrap();
3452        builder.append_null::<Int32Type>("a").unwrap();
3453        builder.append::<Float64Type>("c", 3.0).unwrap();
3454        builder.append_null::<Float64Type>("c").unwrap();
3455        builder.append::<Int32Type>("a", 4).unwrap();
3456        let union = builder.build().unwrap();
3457
3458        let batch =
3459            RecordBatch::try_new(Arc::new(schema.clone()), vec![Arc::new(union) as ArrayRef])
3460                .unwrap();
3461
3462        let mut file = tempfile::tempfile().unwrap();
3463        {
3464            let mut writer = FileWriter::try_new_with_options(&mut file, &schema, options).unwrap();
3465
3466            writer.write(&batch).unwrap();
3467            writer.finish().unwrap();
3468        }
3469        file.rewind().unwrap();
3470
3471        {
3472            let reader = FileReader::try_new(file, None).unwrap();
3473            reader.for_each(|maybe_batch| {
3474                maybe_batch
3475                    .unwrap()
3476                    .columns()
3477                    .iter()
3478                    .zip(batch.columns())
3479                    .for_each(|(a, b)| {
3480                        assert_eq!(a.data_type(), b.data_type());
3481                        assert_eq!(a.len(), b.len());
3482                        assert_eq!(a.null_count(), b.null_count());
3483                    });
3484            });
3485        }
3486    }
3487
3488    #[test]
3489    fn test_write_union_file_v4_v5() {
3490        write_union_file(IpcWriteOptions::try_new(8, false, MetadataVersion::V4).unwrap());
3491        write_union_file(IpcWriteOptions::try_new(8, false, MetadataVersion::V5).unwrap());
3492    }
3493
3494    #[test]
3495    fn test_write_view_types() {
3496        const LONG_TEST_STRING: &str =
3497            "This is a long string to make sure binary view array handles it";
3498        let schema = Schema::new(vec![
3499            Field::new("field1", DataType::BinaryView, true),
3500            Field::new("field2", DataType::Utf8View, true),
3501        ]);
3502        let values: Vec<Option<&[u8]>> = vec![
3503            Some(b"foo"),
3504            Some(b"bar"),
3505            Some(LONG_TEST_STRING.as_bytes()),
3506        ];
3507        let binary_array = BinaryViewArray::from_iter(values);
3508        let utf8_array =
3509            StringViewArray::from_iter(vec![Some("foo"), Some("bar"), Some(LONG_TEST_STRING)]);
3510        let record_batch = RecordBatch::try_new(
3511            Arc::new(schema.clone()),
3512            vec![Arc::new(binary_array), Arc::new(utf8_array)],
3513        )
3514        .unwrap();
3515
3516        let mut file = tempfile::tempfile().unwrap();
3517        {
3518            let mut writer = FileWriter::try_new(&mut file, &schema).unwrap();
3519            writer.write(&record_batch).unwrap();
3520            writer.finish().unwrap();
3521        }
3522        file.rewind().unwrap();
3523        {
3524            let mut reader = FileReader::try_new(&file, None).unwrap();
3525            let read_batch = reader.next().unwrap().unwrap();
3526            read_batch
3527                .columns()
3528                .iter()
3529                .zip(record_batch.columns())
3530                .for_each(|(a, b)| {
3531                    assert_eq!(a, b);
3532                });
3533        }
3534        file.rewind().unwrap();
3535        {
3536            let mut reader = FileReader::try_new(&file, Some(vec![0])).unwrap();
3537            let read_batch = reader.next().unwrap().unwrap();
3538            assert_eq!(read_batch.num_columns(), 1);
3539            let read_array = read_batch.column(0);
3540            let write_array = record_batch.column(0);
3541            assert_eq!(read_array, write_array);
3542        }
3543    }
3544
3545    #[test]
3546    #[cfg_attr(miri, ignore)] // Takes too long
3547    fn truncate_ipc_record_batch() {
3548        fn create_batch(rows: usize) -> RecordBatch {
3549            let schema = Schema::new(vec![
3550                Field::new("a", DataType::Int32, false),
3551                Field::new("b", DataType::Utf8, false),
3552            ]);
3553
3554            let a = Int32Array::from_iter_values(0..rows as i32);
3555            let b = StringArray::from_iter_values((0..rows).map(|i| i.to_string()));
3556
3557            RecordBatch::try_new(Arc::new(schema), vec![Arc::new(a), Arc::new(b)]).unwrap()
3558        }
3559
3560        let big_record_batch = create_batch(65536);
3561
3562        let length = 5;
3563        let small_record_batch = create_batch(length);
3564
3565        let offset = 2;
3566        let record_batch_slice = big_record_batch.slice(offset, length);
3567        assert!(
3568            serialize_stream(&big_record_batch).len() > serialize_stream(&small_record_batch).len()
3569        );
3570        assert_eq!(
3571            serialize_stream(&small_record_batch).len(),
3572            serialize_stream(&record_batch_slice).len()
3573        );
3574
3575        assert_eq!(
3576            deserialize_stream(serialize_stream(&record_batch_slice)),
3577            record_batch_slice
3578        );
3579    }
3580
3581    #[test]
3582    fn truncate_ipc_record_batch_with_nulls() {
3583        fn create_batch() -> RecordBatch {
3584            let schema = Schema::new(vec![
3585                Field::new("a", DataType::Int32, true),
3586                Field::new("b", DataType::Utf8, true),
3587            ]);
3588
3589            let a = Int32Array::from(vec![Some(1), None, Some(1), None, Some(1)]);
3590            let b = StringArray::from(vec![None, Some("a"), Some("a"), None, Some("a")]);
3591
3592            RecordBatch::try_new(Arc::new(schema), vec![Arc::new(a), Arc::new(b)]).unwrap()
3593        }
3594
3595        let record_batch = create_batch();
3596        let record_batch_slice = record_batch.slice(1, 2);
3597        let deserialized_batch = deserialize_stream(serialize_stream(&record_batch_slice));
3598
3599        assert!(
3600            serialize_stream(&record_batch).len() > serialize_stream(&record_batch_slice).len()
3601        );
3602
3603        assert!(deserialized_batch.column(0).is_null(0));
3604        assert!(deserialized_batch.column(0).is_valid(1));
3605        assert!(deserialized_batch.column(1).is_valid(0));
3606        assert!(deserialized_batch.column(1).is_valid(1));
3607
3608        assert_eq!(record_batch_slice, deserialized_batch);
3609    }
3610
3611    #[test]
3612    fn test_write_empty_validity_buffer_without_nulls() {
3613        let batch = RecordBatch::try_from_iter([
3614            ("a", Arc::new(Int32Array::from(vec![1, 2])) as _),
3615            // Keeps a null buffer, but the slice has no nulls
3616            (
3617                "b",
3618                Arc::new(Int32Array::from(vec![None, Some(2), Some(3)]).slice(1, 2)) as _,
3619            ),
3620            ("c", Arc::new(Int32Array::from(vec![Some(1), None])) as _),
3621        ])
3622        .unwrap();
3623        assert_eq!(batch.column(1).nulls().map(|n| n.null_count()), Some(0));
3624
3625        let (_, encoded) = IpcDataGenerator::default()
3626            .encode(
3627                &batch,
3628                &mut DictionaryTracker::new(false),
3629                &Default::default(),
3630                &mut Default::default(),
3631            )
3632            .unwrap();
3633        let message = crate::root_as_message(&encoded.ipc_message).unwrap();
3634        let buffers = message.header_as_record_batch().unwrap().buffers().unwrap();
3635        // Each column writes a validity buffer, then a values buffer
3636        let validity_lengths: Vec<_> = buffers.iter().step_by(2).map(|b| b.length()).collect();
3637        assert_eq!(validity_lengths, [0, 0, 1]);
3638
3639        assert_eq!(deserialize_stream(serialize_stream(&batch)), batch);
3640    }
3641
3642    #[test]
3643    fn truncate_ipc_dictionary_array() {
3644        fn create_batch() -> RecordBatch {
3645            let values: StringArray = [Some("foo"), Some("bar"), Some("baz")]
3646                .into_iter()
3647                .collect();
3648            let keys: Int32Array = [Some(0), Some(2), None, Some(1)].into_iter().collect();
3649
3650            let array = DictionaryArray::new(keys, Arc::new(values));
3651
3652            let schema = Schema::new(vec![Field::new("dict", array.data_type().clone(), true)]);
3653
3654            RecordBatch::try_new(Arc::new(schema), vec![Arc::new(array)]).unwrap()
3655        }
3656
3657        let record_batch = create_batch();
3658        let record_batch_slice = record_batch.slice(1, 2);
3659        let deserialized_batch = deserialize_stream(serialize_stream(&record_batch_slice));
3660
3661        assert!(
3662            serialize_stream(&record_batch).len() > serialize_stream(&record_batch_slice).len()
3663        );
3664
3665        assert!(deserialized_batch.column(0).is_valid(0));
3666        assert!(deserialized_batch.column(0).is_null(1));
3667
3668        assert_eq!(record_batch_slice, deserialized_batch);
3669    }
3670
3671    #[test]
3672    fn truncate_ipc_struct_array() {
3673        fn create_batch() -> RecordBatch {
3674            let strings: StringArray = [Some("foo"), None, Some("bar"), Some("baz")]
3675                .into_iter()
3676                .collect();
3677            let ints: Int32Array = [Some(0), Some(2), None, Some(1)].into_iter().collect();
3678
3679            let struct_array = StructArray::from(vec![
3680                (
3681                    Arc::new(Field::new("s", DataType::Utf8, true)),
3682                    Arc::new(strings) as ArrayRef,
3683                ),
3684                (
3685                    Arc::new(Field::new("c", DataType::Int32, true)),
3686                    Arc::new(ints) as ArrayRef,
3687                ),
3688            ]);
3689
3690            let schema = Schema::new(vec![Field::new(
3691                "struct_array",
3692                struct_array.data_type().clone(),
3693                true,
3694            )]);
3695
3696            RecordBatch::try_new(Arc::new(schema), vec![Arc::new(struct_array)]).unwrap()
3697        }
3698
3699        let record_batch = create_batch();
3700        let record_batch_slice = record_batch.slice(1, 2);
3701        let deserialized_batch = deserialize_stream(serialize_stream(&record_batch_slice));
3702
3703        assert!(
3704            serialize_stream(&record_batch).len() > serialize_stream(&record_batch_slice).len()
3705        );
3706
3707        let structs = deserialized_batch
3708            .column(0)
3709            .as_any()
3710            .downcast_ref::<StructArray>()
3711            .unwrap();
3712
3713        assert!(structs.column(0).is_null(0));
3714        assert!(structs.column(0).is_valid(1));
3715        assert!(structs.column(1).is_valid(0));
3716        assert!(structs.column(1).is_null(1));
3717        assert_eq!(record_batch_slice, deserialized_batch);
3718    }
3719
3720    #[test]
3721    fn truncate_ipc_string_array_with_all_empty_string() {
3722        fn create_batch() -> RecordBatch {
3723            let schema = Schema::new(vec![Field::new("a", DataType::Utf8, true)]);
3724            let a = StringArray::from(vec![Some(""), Some(""), Some(""), Some(""), Some("")]);
3725            RecordBatch::try_new(Arc::new(schema), vec![Arc::new(a)]).unwrap()
3726        }
3727
3728        let record_batch = create_batch();
3729        let record_batch_slice = record_batch.slice(0, 1);
3730        let deserialized_batch = deserialize_stream(serialize_stream(&record_batch_slice));
3731
3732        assert!(
3733            serialize_stream(&record_batch).len() > serialize_stream(&record_batch_slice).len()
3734        );
3735        assert_eq!(record_batch_slice, deserialized_batch);
3736    }
3737
3738    #[test]
3739    fn test_stream_writer_writes_array_slice() {
3740        let array = UInt32Array::from(vec![Some(1), Some(2), Some(3)]);
3741        assert_eq!(
3742            vec![Some(1), Some(2), Some(3)],
3743            array.iter().collect::<Vec<_>>()
3744        );
3745
3746        let sliced = array.slice(1, 2);
3747        assert_eq!(vec![Some(2), Some(3)], sliced.iter().collect::<Vec<_>>());
3748
3749        let batch = RecordBatch::try_new(
3750            Arc::new(Schema::new(vec![Field::new("a", DataType::UInt32, true)])),
3751            vec![Arc::new(sliced)],
3752        )
3753        .expect("new batch");
3754
3755        let mut writer = StreamWriter::try_new(vec![], batch.schema_ref()).expect("new writer");
3756        writer.write(&batch).expect("write");
3757        let outbuf = writer.into_inner().expect("inner");
3758
3759        let mut reader = StreamReader::try_new(&outbuf[..], None).expect("new reader");
3760        let read_batch = reader.next().unwrap().expect("read batch");
3761
3762        let read_array: &UInt32Array = read_batch.column(0).as_primitive();
3763        assert_eq!(
3764            vec![Some(2), Some(3)],
3765            read_array.iter().collect::<Vec<_>>()
3766        );
3767    }
3768
3769    #[test]
3770    fn test_large_slice_uint32() {
3771        ensure_roundtrip(Arc::new(UInt32Array::from_iter(
3772            (0..8000).map(|i| if i % 2 == 0 { Some(i) } else { None }),
3773        )));
3774    }
3775
3776    #[test]
3777    #[cfg_attr(miri, ignore)] // Takes too long
3778    fn test_large_slice_string() {
3779        let strings: Vec<_> = (0..8000)
3780            .map(|i| {
3781                if i % 2 == 0 {
3782                    Some(format!("value{i}"))
3783                } else {
3784                    None
3785                }
3786            })
3787            .collect();
3788
3789        ensure_roundtrip(Arc::new(StringArray::from(strings)));
3790    }
3791
3792    #[test]
3793    #[cfg_attr(miri, ignore)] // Takes too long
3794    fn test_large_slice_string_list() {
3795        let mut ls = ListBuilder::new(StringBuilder::new());
3796
3797        let mut s = String::new();
3798        for row_number in 0..8000 {
3799            if row_number % 2 == 0 {
3800                for list_element in 0..1000 {
3801                    s.clear();
3802                    use std::fmt::Write;
3803                    write!(&mut s, "value{row_number}-{list_element}").unwrap();
3804                    ls.values().append_value(&s);
3805                }
3806                ls.append(true)
3807            } else {
3808                ls.append(false); // null
3809            }
3810        }
3811
3812        ensure_roundtrip(Arc::new(ls.finish()));
3813    }
3814
3815    #[test]
3816    #[cfg_attr(miri, ignore)] // Takes too long
3817    fn test_large_slice_string_list_of_lists() {
3818        // The reason for the special test is to verify reencode_offsets which looks both at
3819        // the starting offset and the data offset.  So need a dataset where the starting_offset
3820        // is zero but the data offset is not.
3821        let mut ls = ListBuilder::new(ListBuilder::new(StringBuilder::new()));
3822
3823        for _ in 0..4000 {
3824            ls.values().append(true);
3825            ls.append(true)
3826        }
3827
3828        let mut s = String::new();
3829        for row_number in 0..4000 {
3830            if row_number % 2 == 0 {
3831                for list_element in 0..1000 {
3832                    s.clear();
3833                    use std::fmt::Write;
3834                    write!(&mut s, "value{row_number}-{list_element}").unwrap();
3835                    ls.values().values().append_value(&s);
3836                }
3837                ls.values().append(true);
3838                ls.append(true)
3839            } else {
3840                ls.append(false); // null
3841            }
3842        }
3843
3844        ensure_roundtrip(Arc::new(ls.finish()));
3845    }
3846
3847    /// Read/write a record batch to a File and Stream and ensure it is the same at the output
3848    fn ensure_roundtrip(array: ArrayRef) {
3849        let num_rows = array.len();
3850        let orig_batch = RecordBatch::try_from_iter(vec![("a", array)]).unwrap();
3851        // take off the first element
3852        let sliced_batch = orig_batch.slice(1, num_rows - 1);
3853
3854        let schema = orig_batch.schema();
3855        let stream_data = {
3856            let mut writer = StreamWriter::try_new(vec![], &schema).unwrap();
3857            writer.write(&sliced_batch).unwrap();
3858            writer.into_inner().unwrap()
3859        };
3860        let read_batch = {
3861            let projection = None;
3862            let mut reader = StreamReader::try_new(Cursor::new(stream_data), projection).unwrap();
3863            reader
3864                .next()
3865                .expect("expect no errors reading batch")
3866                .expect("expect batch")
3867        };
3868        assert_eq!(sliced_batch, read_batch);
3869
3870        let file_data = {
3871            let mut writer = FileWriter::try_new_buffered(vec![], &schema).unwrap();
3872            writer.write(&sliced_batch).unwrap();
3873            writer.into_inner().unwrap().into_inner().unwrap()
3874        };
3875        let read_batch = {
3876            let projection = None;
3877            let mut reader = FileReader::try_new(Cursor::new(file_data), projection).unwrap();
3878            reader
3879                .next()
3880                .expect("expect no errors reading batch")
3881                .expect("expect batch")
3882        };
3883        assert_eq!(sliced_batch, read_batch);
3884
3885        // TODO test file writer/reader
3886    }
3887
3888    #[test]
3889    fn encode_bools_slice() {
3890        // Test case for https://github.com/apache/arrow-rs/issues/3496
3891        assert_bool_roundtrip([true, false], 1, 1);
3892
3893        // slice somewhere in the middle
3894        assert_bool_roundtrip(
3895            [
3896                true, false, true, true, false, false, true, true, true, false, false, false, true,
3897                true, true, true, false, false, false, false, true, true, true, true, true, false,
3898                false, false, false, false,
3899            ],
3900            13,
3901            17,
3902        );
3903
3904        // start at byte boundary, end in the middle
3905        assert_bool_roundtrip(
3906            [
3907                true, false, true, true, false, false, true, true, true, false, false, false,
3908            ],
3909            8,
3910            2,
3911        );
3912
3913        // start and stop and byte boundary
3914        assert_bool_roundtrip(
3915            [
3916                true, false, true, true, false, false, true, true, true, false, false, false, true,
3917                true, true, true, true, false, false, false, false, false,
3918            ],
3919            8,
3920            8,
3921        );
3922    }
3923
3924    fn assert_bool_roundtrip<const N: usize>(bools: [bool; N], offset: usize, length: usize) {
3925        let val_bool_field = Field::new("val", DataType::Boolean, false);
3926
3927        let schema = Arc::new(Schema::new(vec![val_bool_field]));
3928
3929        let bools = BooleanArray::from(bools.to_vec());
3930
3931        let batch = RecordBatch::try_new(Arc::clone(&schema), vec![Arc::new(bools)]).unwrap();
3932        let batch = batch.slice(offset, length);
3933
3934        let data = serialize_stream(&batch);
3935        let batch2 = deserialize_stream(data);
3936        assert_eq!(batch, batch2);
3937    }
3938
3939    fn assert_empty_run_array_roundtrip<R: RunEndIndexType>() {
3940        let run_ends = PrimitiveArray::<R>::from_iter_values(
3941            [2, 5]
3942                .into_iter()
3943                .map(|v| R::Native::from_usize(v).unwrap()),
3944        );
3945        let values = Int32Array::from(vec![10, 20]);
3946        let array = RunArray::<R>::try_new(&run_ends, &values).unwrap();
3947        let empty = RunArray::<R>::from(ArrayData::new_empty(array.data_type()));
3948
3949        for source in [&array, &empty] {
3950            for offset in 0..=source.len() {
3951                let sliced = source.slice(offset, 0);
3952                let batch = RecordBatch::try_from_iter(vec![("run", Arc::new(sliced) as ArrayRef)])
3953                    .unwrap();
3954                for decoded in [
3955                    deserialize_stream(serialize_stream(&batch)),
3956                    deserialize_file(serialize_file(&batch)),
3957                ] {
3958                    assert_eq!(decoded, batch);
3959                    let data = decoded.column(0).to_data();
3960                    data.validate_full().unwrap();
3961                    assert_eq!(data.offset(), 0);
3962                    assert!(data.child_data().iter().all(ArrayData::is_empty));
3963                }
3964
3965                let normalized = into_zero_offset_run_array(source.slice(offset, 0)).unwrap();
3966                normalized.to_data().validate_full().unwrap();
3967                assert_eq!(normalized.offset(), 0);
3968                assert!(normalized.run_ends().values().is_empty());
3969                assert!(normalized.values().is_empty());
3970            }
3971        }
3972    }
3973
3974    #[test]
3975    fn test_empty_run_array_roundtrip_i16() {
3976        assert_empty_run_array_roundtrip::<Int16Type>();
3977    }
3978
3979    #[test]
3980    fn test_empty_run_array_roundtrip_i32() {
3981        assert_empty_run_array_roundtrip::<Int32Type>();
3982    }
3983
3984    #[test]
3985    fn test_empty_run_array_roundtrip_i64() {
3986        assert_empty_run_array_roundtrip::<Int64Type>();
3987    }
3988
3989    #[test]
3990    fn test_run_array_unslice() {
3991        let total_len = 80;
3992        let vals: Vec<Option<i32>> = vec![Some(1), None, Some(2), Some(3), Some(4), None, Some(5)];
3993        let repeats: Vec<usize> = vec![3, 4, 1, 2];
3994        let mut input_array: Vec<Option<i32>> = Vec::with_capacity(total_len);
3995        for ix in 0_usize..32 {
3996            let repeat: usize = repeats[ix % repeats.len()];
3997            let val: Option<i32> = vals[ix % vals.len()];
3998            input_array.resize(input_array.len() + repeat, val);
3999        }
4000
4001        // Encode the input_array to run array
4002        let mut builder =
4003            PrimitiveRunBuilder::<Int16Type, Int32Type>::with_capacity(input_array.len());
4004        builder.extend(input_array.iter().copied());
4005        let run_array = builder.finish();
4006
4007        // test for all slice lengths.
4008        for slice_len in 1..=total_len {
4009            // test for offset = 0, slice length = slice_len
4010            let sliced_run_array: RunArray<Int16Type> =
4011                run_array.slice(0, slice_len).into_data().into();
4012
4013            // Create unsliced run array.
4014            let unsliced_run_array = into_zero_offset_run_array(sliced_run_array).unwrap();
4015            let typed = unsliced_run_array
4016                .downcast::<PrimitiveArray<Int32Type>>()
4017                .unwrap();
4018            let expected: Vec<Option<i32>> = input_array.iter().take(slice_len).copied().collect();
4019            let actual: Vec<Option<i32>> = typed.into_iter().collect();
4020            assert_eq!(expected, actual);
4021
4022            // test for offset = total_len - slice_len, length = slice_len
4023            let sliced_run_array: RunArray<Int16Type> = run_array
4024                .slice(total_len - slice_len, slice_len)
4025                .into_data()
4026                .into();
4027
4028            // Create unsliced run array.
4029            let unsliced_run_array = into_zero_offset_run_array(sliced_run_array).unwrap();
4030            let typed = unsliced_run_array
4031                .downcast::<PrimitiveArray<Int32Type>>()
4032                .unwrap();
4033            let expected: Vec<Option<i32>> = input_array
4034                .iter()
4035                .skip(total_len - slice_len)
4036                .copied()
4037                .collect();
4038            let actual: Vec<Option<i32>> = typed.into_iter().collect();
4039            assert_eq!(expected, actual);
4040        }
4041    }
4042
4043    fn generate_list_data<O: OffsetSizeTrait>() -> GenericListArray<O> {
4044        let mut ls = GenericListBuilder::<O, _>::new(UInt32Builder::new());
4045
4046        for i in 0..100_000 {
4047            for value in [i, i, i] {
4048                ls.values().append_value(value);
4049            }
4050            ls.append(true)
4051        }
4052
4053        ls.finish()
4054    }
4055
4056    fn generate_utf8view_list_data<O: OffsetSizeTrait>() -> GenericListArray<O> {
4057        let mut ls = GenericListBuilder::<O, _>::new(StringViewBuilder::new());
4058
4059        for i in 0..100_000 {
4060            for value in [
4061                format!("value{i}"),
4062                format!("value{i}"),
4063                format!("value{i}"),
4064            ] {
4065                ls.values().append_value(&value);
4066            }
4067            ls.append(true)
4068        }
4069
4070        ls.finish()
4071    }
4072
4073    fn generate_string_list_data<O: OffsetSizeTrait>() -> GenericListArray<O> {
4074        let mut ls = GenericListBuilder::<O, _>::new(StringBuilder::new());
4075
4076        for i in 0..100_000 {
4077            for value in [
4078                format!("value{i}"),
4079                format!("value{i}"),
4080                format!("value{i}"),
4081            ] {
4082                ls.values().append_value(&value);
4083            }
4084            ls.append(true)
4085        }
4086
4087        ls.finish()
4088    }
4089
4090    fn generate_nested_list_data<O: OffsetSizeTrait>() -> GenericListArray<O> {
4091        let mut ls =
4092            GenericListBuilder::<O, _>::new(GenericListBuilder::<O, _>::new(UInt32Builder::new()));
4093
4094        for _i in 0..10_000 {
4095            for j in 0..10 {
4096                for value in [j, j, j, j] {
4097                    ls.values().values().append_value(value);
4098                }
4099                ls.values().append(true)
4100            }
4101            ls.append(true);
4102        }
4103
4104        ls.finish()
4105    }
4106
4107    fn generate_nested_list_data_starting_at_zero<O: OffsetSizeTrait>() -> GenericListArray<O> {
4108        let mut ls =
4109            GenericListBuilder::<O, _>::new(GenericListBuilder::<O, _>::new(UInt32Builder::new()));
4110
4111        for _i in 0..999 {
4112            ls.values().append(true);
4113            ls.append(true);
4114        }
4115
4116        for j in 0..10 {
4117            for value in [j, j, j, j] {
4118                ls.values().values().append_value(value);
4119            }
4120            ls.values().append(true)
4121        }
4122        ls.append(true);
4123
4124        for i in 0..9_000 {
4125            for j in 0..10 {
4126                for value in [i + j, i + j, i + j, i + j] {
4127                    ls.values().values().append_value(value);
4128                }
4129                ls.values().append(true)
4130            }
4131            ls.append(true);
4132        }
4133
4134        ls.finish()
4135    }
4136
4137    fn generate_map_array_data() -> MapArray {
4138        let keys_builder = UInt32Builder::new();
4139        let values_builder = UInt32Builder::new();
4140
4141        let mut builder = MapBuilder::new(None, keys_builder, values_builder);
4142
4143        for i in 0..100_000 {
4144            for _j in 0..3 {
4145                builder.keys().append_value(i);
4146                builder.values().append_value(i * 2);
4147            }
4148            builder.append(true).unwrap();
4149        }
4150
4151        builder.finish()
4152    }
4153
4154    #[test]
4155    #[cfg_attr(miri, ignore)] // Takes too long
4156    fn reencode_offsets_when_first_offset_is_not_zero() {
4157        let original_list = generate_list_data::<i32>();
4158        let original_data = original_list.into_data();
4159        let slice_data = original_data.slice(75, 7);
4160        let (new_offsets, original_start, length) =
4161            reencode_offsets::<i32>(&slice_data.buffers()[0], &slice_data);
4162        assert_eq!(
4163            vec![0, 3, 6, 9, 12, 15, 18, 21],
4164            new_offsets.typed_data::<i32>()
4165        );
4166        assert_eq!(225, original_start);
4167        assert_eq!(21, length);
4168    }
4169
4170    #[test]
4171    fn reencode_offsets_when_first_offset_is_zero() {
4172        let mut ls = GenericListBuilder::<i32, _>::new(UInt32Builder::new());
4173        // ls = [[], [35, 42]
4174        ls.append(true);
4175        ls.values().append_value(35);
4176        ls.values().append_value(42);
4177        ls.append(true);
4178        let original_list = ls.finish();
4179        let original_data = original_list.into_data();
4180
4181        let slice_data = original_data.slice(1, 1);
4182        let (new_offsets, original_start, length) =
4183            reencode_offsets::<i32>(&slice_data.buffers()[0], &slice_data);
4184        assert_eq!(vec![0, 2], new_offsets.typed_data::<i32>());
4185        assert_eq!(0, original_start);
4186        assert_eq!(2, length);
4187    }
4188
4189    /// Ensure when serde full & sliced versions they are equal to original input.
4190    /// Also ensure serialized sliced version is significantly smaller than serialized full.
4191    fn roundtrip_ensure_sliced_smaller(in_batch: RecordBatch, expected_size_factor: usize) {
4192        // test both full and sliced versions
4193        let in_sliced = in_batch.slice(999, 1);
4194
4195        let bytes_batch = serialize_file(&in_batch);
4196        let bytes_sliced = serialize_file(&in_sliced);
4197
4198        // serializing 1 row should be significantly smaller than serializing 100,000
4199        assert!(bytes_sliced.len() < (bytes_batch.len() / expected_size_factor));
4200
4201        // ensure both are still valid and equal to originals
4202        let out_batch = deserialize_file(bytes_batch);
4203        assert_eq!(in_batch, out_batch);
4204
4205        let out_sliced = deserialize_file(bytes_sliced);
4206        assert_eq!(in_sliced, out_sliced);
4207    }
4208
4209    #[test]
4210    #[cfg_attr(miri, ignore)] // Takes too long
4211    fn encode_lists() {
4212        let val_inner = Field::new_list_field(DataType::UInt32, true);
4213        let val_list_field = Field::new("val", DataType::List(Arc::new(val_inner)), false);
4214        let schema = Arc::new(Schema::new(vec![val_list_field]));
4215
4216        let values = Arc::new(generate_list_data::<i32>());
4217
4218        let in_batch = RecordBatch::try_new(schema, vec![values]).unwrap();
4219        roundtrip_ensure_sliced_smaller(in_batch, 1000);
4220    }
4221
4222    #[test]
4223    #[cfg_attr(miri, ignore)] // Takes too long
4224    fn encode_empty_list() {
4225        let val_inner = Field::new_list_field(DataType::UInt32, true);
4226        let val_list_field = Field::new("val", DataType::List(Arc::new(val_inner)), false);
4227        let schema = Arc::new(Schema::new(vec![val_list_field]));
4228
4229        let values = Arc::new(generate_list_data::<i32>());
4230
4231        let in_batch = RecordBatch::try_new(schema, vec![values])
4232            .unwrap()
4233            .slice(999, 0);
4234        let out_batch = deserialize_file(serialize_file(&in_batch));
4235        assert_eq!(in_batch, out_batch);
4236    }
4237
4238    #[test]
4239    #[cfg_attr(miri, ignore)] // Takes too long
4240    fn encode_large_lists() {
4241        let val_inner = Field::new_list_field(DataType::UInt32, true);
4242        let val_list_field = Field::new("val", DataType::LargeList(Arc::new(val_inner)), false);
4243        let schema = Arc::new(Schema::new(vec![val_list_field]));
4244
4245        let values = Arc::new(generate_list_data::<i64>());
4246
4247        // ensure when serde full & sliced versions they are equal to original input
4248        // also ensure serialized sliced version is significantly smaller than serialized full
4249        let in_batch = RecordBatch::try_new(schema, vec![values]).unwrap();
4250        roundtrip_ensure_sliced_smaller(in_batch, 1000);
4251    }
4252
4253    #[test]
4254    #[cfg_attr(miri, ignore)] // Takes too long
4255    fn encode_large_lists_non_zero_offset() {
4256        let val_inner = Field::new_list_field(DataType::UInt32, true);
4257        let val_list_field = Field::new("val", DataType::LargeList(Arc::new(val_inner)), false);
4258        let schema = Arc::new(Schema::new(vec![val_list_field]));
4259
4260        let values = Arc::new(generate_list_data::<i64>());
4261
4262        check_sliced_list_array(schema, values);
4263    }
4264
4265    #[test]
4266    #[cfg_attr(miri, ignore)] // Takes too long
4267    fn encode_large_lists_string_non_zero_offset() {
4268        let val_inner = Field::new_list_field(DataType::Utf8, true);
4269        let val_list_field = Field::new("val", DataType::LargeList(Arc::new(val_inner)), false);
4270        let schema = Arc::new(Schema::new(vec![val_list_field]));
4271
4272        let values = Arc::new(generate_string_list_data::<i64>());
4273
4274        check_sliced_list_array(schema, values);
4275    }
4276
4277    #[test]
4278    #[cfg_attr(miri, ignore)] // Takes too long
4279    fn encode_large_list_string_view_non_zero_offset() {
4280        let val_inner = Field::new_list_field(DataType::Utf8View, true);
4281        let val_list_field = Field::new("val", DataType::LargeList(Arc::new(val_inner)), false);
4282        let schema = Arc::new(Schema::new(vec![val_list_field]));
4283
4284        let values = Arc::new(generate_utf8view_list_data::<i64>());
4285
4286        check_sliced_list_array(schema, values);
4287    }
4288
4289    fn check_sliced_list_array(schema: Arc<Schema>, values: Arc<GenericListArray<i64>>) {
4290        for (offset, len) in [(999, 1), (0, 13), (47, 12), (values.len() - 13, 13)] {
4291            let in_batch = RecordBatch::try_new(schema.clone(), vec![values.clone()])
4292                .unwrap()
4293                .slice(offset, len);
4294            let out_batch = deserialize_file(serialize_file(&in_batch));
4295            assert_eq!(in_batch, out_batch);
4296        }
4297    }
4298
4299    #[test]
4300    #[cfg_attr(miri, ignore)] // Takes too long
4301    fn encode_nested_lists() {
4302        let inner_int = Arc::new(Field::new_list_field(DataType::UInt32, true));
4303        let inner_list_field = Arc::new(Field::new_list_field(DataType::List(inner_int), true));
4304        let list_field = Field::new("val", DataType::List(inner_list_field), true);
4305        let schema = Arc::new(Schema::new(vec![list_field]));
4306
4307        let values = Arc::new(generate_nested_list_data::<i32>());
4308
4309        let in_batch = RecordBatch::try_new(schema, vec![values]).unwrap();
4310        roundtrip_ensure_sliced_smaller(in_batch, 1000);
4311    }
4312
4313    #[test]
4314    #[cfg_attr(miri, ignore)] // Takes too long
4315    fn encode_nested_lists_starting_at_zero() {
4316        let inner_int = Arc::new(Field::new("item", DataType::UInt32, true));
4317        let inner_list_field = Arc::new(Field::new("item", DataType::List(inner_int), true));
4318        let list_field = Field::new("val", DataType::List(inner_list_field), true);
4319        let schema = Arc::new(Schema::new(vec![list_field]));
4320
4321        let values = Arc::new(generate_nested_list_data_starting_at_zero::<i32>());
4322
4323        let in_batch = RecordBatch::try_new(schema, vec![values]).unwrap();
4324        roundtrip_ensure_sliced_smaller(in_batch, 1);
4325    }
4326
4327    #[test]
4328    #[cfg_attr(miri, ignore)] // Takes too long
4329    fn encode_map_array() {
4330        let keys = Arc::new(Field::new(
4331            Field::MAP_KEY_FIELD_DEFAULT_NAME,
4332            DataType::UInt32,
4333            false,
4334        ));
4335        let values = Arc::new(Field::new(
4336            Field::MAP_VALUE_FIELD_DEFAULT_NAME,
4337            DataType::UInt32,
4338            true,
4339        ));
4340        let map_field = Field::new_map(
4341            "map",
4342            Field::MAP_ENTRIES_FIELD_DEFAULT_NAME,
4343            keys,
4344            values,
4345            false,
4346            true,
4347        );
4348        let schema = Arc::new(Schema::new(vec![map_field]));
4349
4350        let values = Arc::new(generate_map_array_data());
4351
4352        let in_batch = RecordBatch::try_new(schema, vec![values]).unwrap();
4353        roundtrip_ensure_sliced_smaller(in_batch, 1000);
4354    }
4355
4356    fn generate_list_view_data<O: OffsetSizeTrait>() -> GenericListViewArray<O> {
4357        let mut builder = GenericListViewBuilder::<O, _>::new(UInt32Builder::new());
4358
4359        for i in 0u32..100_000 {
4360            if i.is_multiple_of(10_000) {
4361                builder.append(false);
4362                continue;
4363            }
4364            for value in [i, i, i] {
4365                builder.values().append_value(value);
4366            }
4367            builder.append(true);
4368        }
4369
4370        builder.finish()
4371    }
4372
4373    #[test]
4374    #[cfg_attr(miri, ignore)] // Takes too long
4375    fn encode_list_view_arrays() {
4376        let val_inner = Field::new_list_field(DataType::UInt32, true);
4377        let val_field = Field::new("val", DataType::ListView(Arc::new(val_inner)), true);
4378        let schema = Arc::new(Schema::new(vec![val_field]));
4379
4380        let values = Arc::new(generate_list_view_data::<i32>());
4381
4382        let in_batch = RecordBatch::try_new(schema, vec![values]).unwrap();
4383        let out_batch = deserialize_file(serialize_file(&in_batch));
4384        assert_eq!(in_batch, out_batch);
4385    }
4386
4387    #[test]
4388    #[cfg_attr(miri, ignore)] // Takes too long
4389    fn encode_large_list_view_arrays() {
4390        let val_inner = Field::new_list_field(DataType::UInt32, true);
4391        let val_field = Field::new("val", DataType::LargeListView(Arc::new(val_inner)), true);
4392        let schema = Arc::new(Schema::new(vec![val_field]));
4393
4394        let values = Arc::new(generate_list_view_data::<i64>());
4395
4396        let in_batch = RecordBatch::try_new(schema, vec![values]).unwrap();
4397        let out_batch = deserialize_file(serialize_file(&in_batch));
4398        assert_eq!(in_batch, out_batch);
4399    }
4400
4401    #[test]
4402    #[cfg_attr(miri, ignore)] // Takes too long
4403    fn check_sliced_list_view_array() {
4404        let inner = Field::new_list_field(DataType::UInt32, true);
4405        let field = Field::new("val", DataType::ListView(Arc::new(inner)), true);
4406        let schema = Arc::new(Schema::new(vec![field]));
4407        let values = Arc::new(generate_list_view_data::<i32>());
4408
4409        for (offset, len) in [(999, 1), (0, 13), (47, 12), (values.len() - 13, 13)] {
4410            let in_batch = RecordBatch::try_new(schema.clone(), vec![values.clone()])
4411                .unwrap()
4412                .slice(offset, len);
4413            let out_batch = deserialize_file(serialize_file(&in_batch));
4414            assert_eq!(in_batch, out_batch);
4415        }
4416    }
4417
4418    #[test]
4419    #[cfg_attr(miri, ignore)] // Takes too long
4420    fn check_sliced_large_list_view_array() {
4421        let inner = Field::new_list_field(DataType::UInt32, true);
4422        let field = Field::new("val", DataType::LargeListView(Arc::new(inner)), true);
4423        let schema = Arc::new(Schema::new(vec![field]));
4424        let values = Arc::new(generate_list_view_data::<i64>());
4425
4426        for (offset, len) in [(999, 1), (0, 13), (47, 12), (values.len() - 13, 13)] {
4427            let in_batch = RecordBatch::try_new(schema.clone(), vec![values.clone()])
4428                .unwrap()
4429                .slice(offset, len);
4430            let out_batch = deserialize_file(serialize_file(&in_batch));
4431            assert_eq!(in_batch, out_batch);
4432        }
4433    }
4434
4435    fn generate_nested_list_view_data<O: OffsetSizeTrait>() -> GenericListViewArray<O> {
4436        let inner_builder = UInt32Builder::new();
4437        let middle_builder = GenericListViewBuilder::<O, _>::new(inner_builder);
4438        let mut outer_builder = GenericListViewBuilder::<O, _>::new(middle_builder);
4439
4440        for i in 0u32..10_000 {
4441            if i.is_multiple_of(1_000) {
4442                outer_builder.append(false);
4443                continue;
4444            }
4445
4446            for _ in 0..3 {
4447                for value in [i, i + 1, i + 2] {
4448                    outer_builder.values().values().append_value(value);
4449                }
4450                outer_builder.values().append(true);
4451            }
4452            outer_builder.append(true);
4453        }
4454
4455        outer_builder.finish()
4456    }
4457
4458    #[test]
4459    #[cfg_attr(miri, ignore)] // Takes too long
4460    fn encode_nested_list_views() {
4461        let inner_int = Arc::new(Field::new_list_field(DataType::UInt32, true));
4462        let inner_list_field = Arc::new(Field::new_list_field(DataType::ListView(inner_int), true));
4463        let list_field = Field::new("val", DataType::ListView(inner_list_field), true);
4464        let schema = Arc::new(Schema::new(vec![list_field]));
4465
4466        let values = Arc::new(generate_nested_list_view_data::<i32>());
4467
4468        let in_batch = RecordBatch::try_new(schema, vec![values]).unwrap();
4469        let out_batch = deserialize_file(serialize_file(&in_batch));
4470        assert_eq!(in_batch, out_batch);
4471    }
4472
4473    fn test_roundtrip_list_view_of_dict_impl<OffsetSize: OffsetSizeTrait, U: ArrowNativeType>(
4474        list_data_type: DataType,
4475        offsets: &[U; 5],
4476        sizes: &[U; 4],
4477    ) {
4478        let values = StringArray::from(vec![Some("alpha"), None, Some("beta"), Some("gamma")]);
4479        let keys = Int32Array::from_iter_values([0, 0, 1, 2, 3, 0, 2]);
4480        let dict_array = DictionaryArray::new(keys, Arc::new(values));
4481        let dict_data = dict_array.to_data();
4482
4483        let value_offsets = Buffer::from_slice_ref(offsets);
4484        let value_sizes = Buffer::from_slice_ref(sizes);
4485
4486        let list_data = ArrayData::builder(list_data_type)
4487            .len(4)
4488            .add_buffer(value_offsets)
4489            .add_buffer(value_sizes)
4490            .add_child_data(dict_data)
4491            .build()
4492            .unwrap();
4493        let list_view_array = GenericListViewArray::<OffsetSize>::from(list_data);
4494
4495        let schema = Arc::new(Schema::new(vec![Field::new(
4496            "f1",
4497            list_view_array.data_type().clone(),
4498            false,
4499        )]));
4500        let input_batch = RecordBatch::try_new(schema, vec![Arc::new(list_view_array)]).unwrap();
4501
4502        let output_batch = deserialize_file(serialize_file(&input_batch));
4503        assert_eq!(input_batch, output_batch);
4504
4505        let output_batch = deserialize_stream(serialize_stream(&input_batch));
4506        assert_eq!(input_batch, output_batch);
4507    }
4508
4509    #[test]
4510    fn test_roundtrip_list_view_of_dict() {
4511        #[expect(deprecated)]
4512        let list_data_type = DataType::ListView(Arc::new(Field::new_dict(
4513            "item",
4514            DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
4515            true,
4516            1,
4517            false,
4518        )));
4519        let offsets: &[i32; 5] = &[0, 2, 4, 4, 7];
4520        let sizes: &[i32; 4] = &[2, 2, 0, 3];
4521        test_roundtrip_list_view_of_dict_impl::<i32, i32>(list_data_type, offsets, sizes);
4522    }
4523
4524    #[test]
4525    fn test_roundtrip_large_list_view_of_dict() {
4526        #[expect(deprecated)]
4527        let list_data_type = DataType::LargeListView(Arc::new(Field::new_dict(
4528            "item",
4529            DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
4530            true,
4531            2,
4532            false,
4533        )));
4534        let offsets: &[i64; 5] = &[0, 2, 4, 4, 7];
4535        let sizes: &[i64; 4] = &[2, 2, 0, 3];
4536        test_roundtrip_list_view_of_dict_impl::<i64, i64>(list_data_type, offsets, sizes);
4537    }
4538
4539    #[test]
4540    fn test_roundtrip_sliced_list_view_of_dict() {
4541        #[expect(deprecated)]
4542        let list_data_type = DataType::ListView(Arc::new(Field::new_dict(
4543            "item",
4544            DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
4545            true,
4546            3,
4547            false,
4548        )));
4549
4550        let values = StringArray::from(vec![Some("alpha"), None, Some("beta"), Some("gamma")]);
4551        let keys = Int32Array::from_iter_values([0, 0, 1, 2, 3, 0, 2, 1, 0, 3, 2, 1]);
4552        let dict_array = DictionaryArray::new(keys, Arc::new(values));
4553        let dict_data = dict_array.to_data();
4554
4555        let offsets: &[i32; 7] = &[0, 2, 4, 4, 7, 9, 12];
4556        let sizes: &[i32; 6] = &[2, 2, 0, 3, 2, 3];
4557        let value_offsets = Buffer::from_slice_ref(offsets);
4558        let value_sizes = Buffer::from_slice_ref(sizes);
4559
4560        let list_data = ArrayData::builder(list_data_type)
4561            .len(6)
4562            .add_buffer(value_offsets)
4563            .add_buffer(value_sizes)
4564            .add_child_data(dict_data)
4565            .build()
4566            .unwrap();
4567        let list_view_array = GenericListViewArray::<i32>::from(list_data);
4568
4569        let schema = Arc::new(Schema::new(vec![Field::new(
4570            "f1",
4571            list_view_array.data_type().clone(),
4572            false,
4573        )]));
4574        let input_batch = RecordBatch::try_new(schema, vec![Arc::new(list_view_array)]).unwrap();
4575
4576        let sliced_batch = input_batch.slice(1, 4);
4577
4578        let output_batch = deserialize_file(serialize_file(&sliced_batch));
4579        assert_eq!(sliced_batch, output_batch);
4580
4581        let output_batch = deserialize_stream(serialize_stream(&sliced_batch));
4582        assert_eq!(sliced_batch, output_batch);
4583    }
4584
4585    #[test]
4586    fn test_roundtrip_dense_union_of_dict() {
4587        let values = StringArray::from(vec![Some("alpha"), None, Some("beta"), Some("gamma")]);
4588        let keys = Int32Array::from_iter_values([0, 0, 1, 2, 3, 0, 2]);
4589        let dict_array = DictionaryArray::new(keys, Arc::new(values));
4590
4591        #[expect(deprecated)]
4592        let dict_field = Arc::new(Field::new_dict(
4593            "dict",
4594            DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
4595            true,
4596            1,
4597            false,
4598        ));
4599        let int_field = Arc::new(Field::new("int", DataType::Int32, false));
4600        let union_fields = UnionFields::try_new(vec![0, 1], vec![dict_field, int_field]).unwrap();
4601
4602        let types = ScalarBuffer::from(vec![0i8, 0, 1, 0, 1, 0, 0]);
4603        let offsets = ScalarBuffer::from(vec![0i32, 1, 0, 2, 1, 3, 4]);
4604
4605        let int_array = Int32Array::from(vec![100, 200]);
4606
4607        let union = UnionArray::try_new(
4608            union_fields.clone(),
4609            types,
4610            Some(offsets),
4611            vec![Arc::new(dict_array), Arc::new(int_array)],
4612        )
4613        .unwrap();
4614
4615        let schema = Arc::new(Schema::new(vec![Field::new(
4616            "union",
4617            DataType::Union(union_fields, UnionMode::Dense),
4618            false,
4619        )]));
4620        let input_batch = RecordBatch::try_new(schema, vec![Arc::new(union)]).unwrap();
4621
4622        let output_batch = deserialize_file(serialize_file(&input_batch));
4623        assert_eq!(input_batch, output_batch);
4624
4625        let output_batch = deserialize_stream(serialize_stream(&input_batch));
4626        assert_eq!(input_batch, output_batch);
4627    }
4628
4629    #[test]
4630    fn test_roundtrip_sparse_union_of_dict() {
4631        let values = StringArray::from(vec![Some("alpha"), None, Some("beta"), Some("gamma")]);
4632        let keys = Int32Array::from_iter_values([0, 0, 1, 2, 3, 0, 2]);
4633        let dict_array = DictionaryArray::new(keys, Arc::new(values));
4634
4635        #[expect(deprecated)]
4636        let dict_field = Arc::new(Field::new_dict(
4637            "dict",
4638            DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
4639            true,
4640            2,
4641            false,
4642        ));
4643        let int_field = Arc::new(Field::new("int", DataType::Int32, false));
4644        let union_fields = UnionFields::try_new(vec![0, 1], vec![dict_field, int_field]).unwrap();
4645
4646        let types = ScalarBuffer::from(vec![0i8, 0, 1, 0, 1, 0, 0]);
4647
4648        let int_array = Int32Array::from(vec![0, 0, 100, 0, 200, 0, 0]);
4649
4650        let union = UnionArray::try_new(
4651            union_fields.clone(),
4652            types,
4653            None,
4654            vec![Arc::new(dict_array), Arc::new(int_array)],
4655        )
4656        .unwrap();
4657
4658        let schema = Arc::new(Schema::new(vec![Field::new(
4659            "union",
4660            DataType::Union(union_fields, UnionMode::Sparse),
4661            false,
4662        )]));
4663        let input_batch = RecordBatch::try_new(schema, vec![Arc::new(union)]).unwrap();
4664
4665        let output_batch = deserialize_file(serialize_file(&input_batch));
4666        assert_eq!(input_batch, output_batch);
4667
4668        let output_batch = deserialize_stream(serialize_stream(&input_batch));
4669        assert_eq!(input_batch, output_batch);
4670    }
4671
4672    #[test]
4673    fn test_roundtrip_map_with_dict_keys() {
4674        // Building a map array is a bit involved. We first build a struct arary that has a key and
4675        // value field and then use that to build the actual map array.
4676        let key_values = StringArray::from(vec!["key_a", "key_b", "key_c"]);
4677        let keys = Int32Array::from_iter_values([0, 1, 2, 0, 1, 0]);
4678        let dict_keys = DictionaryArray::new(keys, Arc::new(key_values));
4679
4680        let values = Int32Array::from(vec![1, 2, 3, 4, 5, 6]);
4681
4682        #[expect(deprecated)]
4683        let entries_field = Arc::new(Field::new(
4684            Field::MAP_ENTRIES_FIELD_DEFAULT_NAME,
4685            DataType::Struct(
4686                vec![
4687                    Field::new_dict(
4688                        Field::MAP_KEY_FIELD_DEFAULT_NAME,
4689                        DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
4690                        false,
4691                        1,
4692                        false,
4693                    ),
4694                    Field::new(Field::MAP_VALUE_FIELD_DEFAULT_NAME, DataType::Int32, true),
4695                ]
4696                .into(),
4697            ),
4698            false,
4699        ));
4700
4701        let entries = StructArray::from(vec![
4702            (
4703                Arc::new(Field::new(
4704                    Field::MAP_KEY_FIELD_DEFAULT_NAME,
4705                    DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
4706                    false,
4707                )),
4708                Arc::new(dict_keys) as ArrayRef,
4709            ),
4710            (
4711                Arc::new(Field::new(
4712                    Field::MAP_VALUE_FIELD_DEFAULT_NAME,
4713                    DataType::Int32,
4714                    true,
4715                )),
4716                Arc::new(values) as ArrayRef,
4717            ),
4718        ]);
4719
4720        let offsets = Buffer::from_slice_ref([0i32, 2, 4, 6]);
4721
4722        let map_data = ArrayData::builder(DataType::Map(entries_field, false))
4723            .len(3)
4724            .add_buffer(offsets)
4725            .add_child_data(entries.into_data())
4726            .build()
4727            .unwrap();
4728        let map_array = MapArray::from(map_data);
4729
4730        let schema = Arc::new(Schema::new(vec![Field::new(
4731            "map",
4732            map_array.data_type().clone(),
4733            false,
4734        )]));
4735        let input_batch = RecordBatch::try_new(schema, vec![Arc::new(map_array)]).unwrap();
4736
4737        let output_batch = deserialize_file(serialize_file(&input_batch));
4738        assert_eq!(input_batch, output_batch);
4739
4740        let output_batch = deserialize_stream(serialize_stream(&input_batch));
4741        assert_eq!(input_batch, output_batch);
4742    }
4743
4744    #[test]
4745    fn test_roundtrip_map_with_dict_values() {
4746        // Building a map array is a bit involved. We first build a struct arary that has a key and
4747        // value field and then use that to build the actual map array.
4748        let keys = StringArray::from(vec!["a", "b", "c", "d", "e", "f"]);
4749
4750        let value_values = StringArray::from(vec!["val_x", "val_y", "val_z"]);
4751        let value_keys = Int32Array::from_iter_values([0, 1, 2, 0, 1, 0]);
4752        let dict_values = DictionaryArray::new(value_keys, Arc::new(value_values));
4753
4754        #[expect(deprecated)]
4755        let entries_field = Arc::new(Field::new(
4756            Field::MAP_ENTRIES_FIELD_DEFAULT_NAME,
4757            DataType::Struct(
4758                vec![
4759                    Field::new(Field::MAP_KEY_FIELD_DEFAULT_NAME, DataType::Utf8, false),
4760                    Field::new_dict(
4761                        Field::MAP_VALUE_FIELD_DEFAULT_NAME,
4762                        DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
4763                        true,
4764                        2,
4765                        false,
4766                    ),
4767                ]
4768                .into(),
4769            ),
4770            false,
4771        ));
4772
4773        let entries = StructArray::from(vec![
4774            (
4775                Arc::new(Field::new(
4776                    Field::MAP_KEY_FIELD_DEFAULT_NAME,
4777                    DataType::Utf8,
4778                    false,
4779                )),
4780                Arc::new(keys) as ArrayRef,
4781            ),
4782            (
4783                Arc::new(Field::new(
4784                    Field::MAP_VALUE_FIELD_DEFAULT_NAME,
4785                    DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
4786                    true,
4787                )),
4788                Arc::new(dict_values) as ArrayRef,
4789            ),
4790        ]);
4791
4792        let offsets = Buffer::from_slice_ref([0i32, 2, 4, 6]);
4793
4794        let map_data = ArrayData::builder(DataType::Map(entries_field, false))
4795            .len(3)
4796            .add_buffer(offsets)
4797            .add_child_data(entries.into_data())
4798            .build()
4799            .unwrap();
4800        let map_array = MapArray::from(map_data);
4801
4802        let schema = Arc::new(Schema::new(vec![Field::new(
4803            "map",
4804            map_array.data_type().clone(),
4805            false,
4806        )]));
4807        let input_batch = RecordBatch::try_new(schema, vec![Arc::new(map_array)]).unwrap();
4808
4809        let output_batch = deserialize_file(serialize_file(&input_batch));
4810        assert_eq!(input_batch, output_batch);
4811
4812        let output_batch = deserialize_stream(serialize_stream(&input_batch));
4813        assert_eq!(input_batch, output_batch);
4814    }
4815
4816    #[test]
4817    fn test_decimal128_alignment16_is_sufficient() {
4818        const IPC_ALIGNMENT: usize = 16;
4819
4820        // Test a bunch of different dimensions to ensure alignment is never an issue.
4821        // For example, if we only test `num_cols = 1` then even with alignment 8 this
4822        // test would _happen_ to pass, even though for different dimensions like
4823        // `num_cols = 2` it would fail.
4824        for num_cols in [1, 2, 3, 17, 50, 73, 99] {
4825            let num_rows = (num_cols * 7 + 11) % 100; // Deterministic swizzle
4826
4827            let mut fields = Vec::new();
4828            let mut arrays = Vec::new();
4829            for i in 0..num_cols {
4830                let field = Field::new(format!("col_{i}"), DataType::Decimal128(38, 10), true);
4831                let array = Decimal128Array::from(vec![num_cols as i128; num_rows]);
4832                fields.push(field);
4833                arrays.push(Arc::new(array) as Arc<dyn Array>);
4834            }
4835            let schema = Schema::new(fields);
4836            let batch = RecordBatch::try_new(Arc::new(schema), arrays).unwrap();
4837
4838            let mut writer = FileWriter::try_new_with_options(
4839                Vec::new(),
4840                batch.schema_ref(),
4841                IpcWriteOptions::try_new(IPC_ALIGNMENT, false, MetadataVersion::V5).unwrap(),
4842            )
4843            .unwrap();
4844            writer.write(&batch).unwrap();
4845            writer.finish().unwrap();
4846
4847            let out: Vec<u8> = writer.into_inner().unwrap();
4848
4849            let buffer = Buffer::from_slice_ref(out);
4850            let trailer_start = buffer.len() - 10;
4851            let footer_len =
4852                read_footer_length(buffer[trailer_start..].try_into().unwrap()).unwrap();
4853            let footer =
4854                root_as_footer(&buffer[trailer_start - footer_len..trailer_start]).unwrap();
4855
4856            let schema = try_fb_to_schema(footer.schema().unwrap()).unwrap();
4857
4858            // Importantly we set `require_alignment`, checking that 16-byte alignment is sufficient
4859            // for `read_record_batch` later on to read the data in a zero-copy manner.
4860            let decoder =
4861                FileDecoder::new(Arc::new(schema), footer.version()).with_require_alignment(true);
4862
4863            let batches = footer.recordBatches().unwrap();
4864
4865            let block = batches.get(0);
4866            let block_len = block.bodyLength() as usize + block.metaDataLength() as usize;
4867            let data = buffer.slice_with_length(block.offset() as _, block_len);
4868
4869            let batch2 = decoder.read_record_batch(block, &data).unwrap().unwrap();
4870
4871            assert_eq!(batch, batch2);
4872        }
4873    }
4874
4875    #[test]
4876    fn test_decimal128_alignment8_is_unaligned() {
4877        const IPC_ALIGNMENT: usize = 8;
4878
4879        let num_cols = 2;
4880
4881        let mut fields = Vec::new();
4882        let mut arrays = Vec::new();
4883        for i in 0..num_cols {
4884            let field = Field::new(format!("col_{i}"), DataType::Decimal128(38, 10), true);
4885            // The null makes the writer emit a 1-byte validity buffer. With IPC_ALIGNMENT = 8,
4886            // the writer pads that buffer to 8 bytes, so the decimal values that follow start
4887            // at offset 8 of the body, which is not a multiple of 16.
4888            let array = Decimal128Array::from(vec![Some(num_cols as i128), None]);
4889            fields.push(field);
4890            arrays.push(Arc::new(array) as Arc<dyn Array>);
4891        }
4892        let schema = Schema::new(fields);
4893        let batch = RecordBatch::try_new(Arc::new(schema), arrays).unwrap();
4894
4895        let mut writer = FileWriter::try_new_with_options(
4896            Vec::new(),
4897            batch.schema_ref(),
4898            IpcWriteOptions::try_new(IPC_ALIGNMENT, false, MetadataVersion::V5).unwrap(),
4899        )
4900        .unwrap();
4901        writer.write(&batch).unwrap();
4902        writer.finish().unwrap();
4903
4904        let out: Vec<u8> = writer.into_inner().unwrap();
4905
4906        let buffer = Buffer::from_slice_ref(out);
4907        let trailer_start = buffer.len() - 10;
4908        let footer_len = read_footer_length(buffer[trailer_start..].try_into().unwrap()).unwrap();
4909        let footer = root_as_footer(&buffer[trailer_start - footer_len..trailer_start]).unwrap();
4910        let schema = try_fb_to_schema(footer.schema().unwrap()).unwrap();
4911
4912        // Importantly we set `require_alignment`, otherwise the error later is suppressed due to copying
4913        // to an aligned buffer in `ArrayDataBuilder.build_aligned`.
4914        let decoder =
4915            FileDecoder::new(Arc::new(schema), footer.version()).with_require_alignment(true);
4916
4917        let batches = footer.recordBatches().unwrap();
4918
4919        let block = batches.get(0);
4920        let block_len = block.bodyLength() as usize + block.metaDataLength() as usize;
4921        let data = buffer.slice_with_length(block.offset() as _, block_len);
4922
4923        let result = decoder.read_record_batch(block, &data);
4924
4925        let error = result.unwrap_err();
4926        assert_eq!(
4927            error.to_string(),
4928            "Invalid argument error: Misaligned buffers[0] in array of type Decimal128(38, 10), \
4929             offset from expected alignment of 16 by 8"
4930        );
4931    }
4932
4933    #[test]
4934    fn test_flush() {
4935        // We write a schema which is small enough to fit into a buffer and not get flushed,
4936        // and then force the write with .flush().
4937        let num_cols = 2;
4938        let mut fields = Vec::new();
4939        let options = IpcWriteOptions::try_new(8, false, MetadataVersion::V5).unwrap();
4940        for i in 0..num_cols {
4941            let field = Field::new(format!("col_{i}"), DataType::Decimal128(38, 10), true);
4942            fields.push(field);
4943        }
4944        let schema = Schema::new(fields);
4945        let inner_stream_writer = BufWriter::with_capacity(1024, Vec::new());
4946        let inner_file_writer = BufWriter::with_capacity(1024, Vec::new());
4947        let mut stream_writer =
4948            StreamWriter::try_new_with_options(inner_stream_writer, &schema, options.clone())
4949                .unwrap();
4950        let mut file_writer =
4951            FileWriter::try_new_with_options(inner_file_writer, &schema, options).unwrap();
4952
4953        let stream_bytes_written_on_new = stream_writer.get_ref().get_ref().len();
4954        let file_bytes_written_on_new = file_writer.get_ref().get_ref().len();
4955        stream_writer.flush().unwrap();
4956        file_writer.flush().unwrap();
4957        let stream_bytes_written_on_flush = stream_writer.get_ref().get_ref().len();
4958        let file_bytes_written_on_flush = file_writer.get_ref().get_ref().len();
4959        let stream_out = stream_writer.into_inner().unwrap().into_inner().unwrap();
4960        // Finishing a stream writes the continuation bytes in MetadataVersion::V5 (4 bytes)
4961        // and then a length of 0 (4 bytes) for a total of 8 bytes.
4962        // Everything before that should have been flushed in the .flush() call.
4963        let expected_stream_flushed_bytes = stream_out.len() - 8;
4964        // A file write is the same as the stream write except for the leading magic string
4965        // ARROW1 plus padding, which is 8 bytes.
4966        let expected_file_flushed_bytes = expected_stream_flushed_bytes + 8;
4967
4968        assert!(
4969            stream_bytes_written_on_new < stream_bytes_written_on_flush,
4970            "this test makes no sense if flush is not actually required"
4971        );
4972        assert!(
4973            file_bytes_written_on_new < file_bytes_written_on_flush,
4974            "this test makes no sense if flush is not actually required"
4975        );
4976        assert_eq!(stream_bytes_written_on_flush, expected_stream_flushed_bytes);
4977        assert_eq!(file_bytes_written_on_flush, expected_file_flushed_bytes);
4978    }
4979
4980    #[test]
4981    fn test_roundtrip_list_of_fixed_list() -> Result<(), ArrowError> {
4982        let l1_type =
4983            DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Float32, false)), 3);
4984        let l2_type = DataType::List(Arc::new(Field::new("item", l1_type.clone(), false)));
4985
4986        let l0_builder = Float32Builder::new();
4987        let l1_builder = FixedSizeListBuilder::new(l0_builder, 3).with_field(Arc::new(Field::new(
4988            "item",
4989            DataType::Float32,
4990            false,
4991        )));
4992        let mut l2_builder =
4993            ListBuilder::new(l1_builder).with_field(Arc::new(Field::new("item", l1_type, false)));
4994
4995        for point in [[1.0, 2.0, 3.0], [4.0, 5.0, 6.0], [7.0, 8.0, 9.0]] {
4996            l2_builder.values().values().append_value(point[0]);
4997            l2_builder.values().values().append_value(point[1]);
4998            l2_builder.values().values().append_value(point[2]);
4999
5000            l2_builder.values().append(true);
5001        }
5002        l2_builder.append(true);
5003
5004        let point = [10., 11., 12.];
5005        l2_builder.values().values().append_value(point[0]);
5006        l2_builder.values().values().append_value(point[1]);
5007        l2_builder.values().values().append_value(point[2]);
5008
5009        l2_builder.values().append(true);
5010        l2_builder.append(true);
5011
5012        let array = Arc::new(l2_builder.finish()) as ArrayRef;
5013
5014        let schema = Arc::new(Schema::new_with_metadata(
5015            vec![Field::new("points", l2_type, false)],
5016            HashMap::default(),
5017        ));
5018
5019        // Test a variety of combinations that include 0 and non-zero offsets
5020        // and also portions or the rest of the array
5021        test_slices(&array, &schema, 0, 1)?;
5022        test_slices(&array, &schema, 0, 2)?;
5023        test_slices(&array, &schema, 1, 1)?;
5024
5025        Ok(())
5026    }
5027
5028    #[test]
5029    fn test_roundtrip_list_of_fixed_list_w_nulls() -> Result<(), ArrowError> {
5030        let l0_builder = Float32Builder::new();
5031        let l1_builder = FixedSizeListBuilder::new(l0_builder, 3);
5032        let mut l2_builder = ListBuilder::new(l1_builder);
5033
5034        for point in [
5035            [Some(1.0), Some(2.0), None],
5036            [Some(4.0), Some(5.0), Some(6.0)],
5037            [None, Some(8.0), Some(9.0)],
5038        ] {
5039            for p in point {
5040                match p {
5041                    Some(p) => l2_builder.values().values().append_value(p),
5042                    None => l2_builder.values().values().append_null(),
5043                }
5044            }
5045
5046            l2_builder.values().append(true);
5047        }
5048        l2_builder.append(true);
5049
5050        let point = [Some(10.), None, None];
5051        for p in point {
5052            match p {
5053                Some(p) => l2_builder.values().values().append_value(p),
5054                None => l2_builder.values().values().append_null(),
5055            }
5056        }
5057
5058        l2_builder.values().append(true);
5059        l2_builder.append(true);
5060
5061        let array = Arc::new(l2_builder.finish()) as ArrayRef;
5062
5063        let schema = Arc::new(Schema::new_with_metadata(
5064            vec![Field::new(
5065                "points",
5066                DataType::List(Arc::new(Field::new(
5067                    "item",
5068                    DataType::FixedSizeList(
5069                        Arc::new(Field::new("item", DataType::Float32, true)),
5070                        3,
5071                    ),
5072                    true,
5073                ))),
5074                true,
5075            )],
5076            HashMap::default(),
5077        ));
5078
5079        // Test a variety of combinations that include 0 and non-zero offsets
5080        // and also portions or the rest of the array
5081        test_slices(&array, &schema, 0, 1)?;
5082        test_slices(&array, &schema, 0, 2)?;
5083        test_slices(&array, &schema, 1, 1)?;
5084
5085        Ok(())
5086    }
5087
5088    fn test_slices(
5089        parent_array: &ArrayRef,
5090        schema: &SchemaRef,
5091        offset: usize,
5092        length: usize,
5093    ) -> Result<(), ArrowError> {
5094        let subarray = parent_array.slice(offset, length);
5095        let original_batch = RecordBatch::try_new(schema.clone(), vec![subarray])?;
5096
5097        let mut bytes = Vec::new();
5098        let mut writer = StreamWriter::try_new(&mut bytes, schema)?;
5099        writer.write(&original_batch)?;
5100        writer.finish()?;
5101
5102        let mut cursor = std::io::Cursor::new(bytes);
5103        let mut reader = StreamReader::try_new(&mut cursor, None)?;
5104        let returned_batch = reader.next().unwrap()?;
5105
5106        assert_eq!(original_batch, returned_batch);
5107
5108        Ok(())
5109    }
5110
5111    #[test]
5112    fn test_roundtrip_fixed_list() -> Result<(), ArrowError> {
5113        let int_builder = Int64Builder::new();
5114        let mut fixed_list_builder = FixedSizeListBuilder::new(int_builder, 3)
5115            .with_field(Arc::new(Field::new("item", DataType::Int64, false)));
5116
5117        for point in [[1, 2, 3], [4, 5, 6], [7, 8, 9], [10, 11, 12]] {
5118            fixed_list_builder.values().append_value(point[0]);
5119            fixed_list_builder.values().append_value(point[1]);
5120            fixed_list_builder.values().append_value(point[2]);
5121
5122            fixed_list_builder.append(true);
5123        }
5124
5125        let array = Arc::new(fixed_list_builder.finish()) as ArrayRef;
5126
5127        let schema = Arc::new(Schema::new_with_metadata(
5128            vec![Field::new(
5129                "points",
5130                DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Int64, false)), 3),
5131                false,
5132            )],
5133            HashMap::default(),
5134        ));
5135
5136        // Test a variety of combinations that include 0 and non-zero offsets
5137        // and also portions or the rest of the array
5138        test_slices(&array, &schema, 0, 4)?;
5139        test_slices(&array, &schema, 0, 2)?;
5140        test_slices(&array, &schema, 1, 3)?;
5141        test_slices(&array, &schema, 2, 1)?;
5142
5143        Ok(())
5144    }
5145
5146    #[test]
5147    fn test_roundtrip_fixed_list_w_nulls() -> Result<(), ArrowError> {
5148        let int_builder = Int64Builder::new();
5149        let mut fixed_list_builder = FixedSizeListBuilder::new(int_builder, 3);
5150
5151        for point in [
5152            [Some(1), Some(2), None],
5153            [Some(4), Some(5), Some(6)],
5154            [None, Some(8), Some(9)],
5155            [Some(10), None, None],
5156        ] {
5157            for p in point {
5158                match p {
5159                    Some(p) => fixed_list_builder.values().append_value(p),
5160                    None => fixed_list_builder.values().append_null(),
5161                }
5162            }
5163
5164            fixed_list_builder.append(true);
5165        }
5166
5167        let array = Arc::new(fixed_list_builder.finish()) as ArrayRef;
5168
5169        let schema = Arc::new(Schema::new_with_metadata(
5170            vec![Field::new(
5171                "points",
5172                DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Int64, true)), 3),
5173                true,
5174            )],
5175            HashMap::default(),
5176        ));
5177
5178        // Test a variety of combinations that include 0 and non-zero offsets
5179        // and also portions or the rest of the array
5180        test_slices(&array, &schema, 0, 4)?;
5181        test_slices(&array, &schema, 0, 2)?;
5182        test_slices(&array, &schema, 1, 3)?;
5183        test_slices(&array, &schema, 2, 1)?;
5184
5185        Ok(())
5186    }
5187
5188    #[test]
5189    fn test_metadata_encoding_ordering() {
5190        fn create_hash() -> u64 {
5191            let metadata: HashMap<String, String> = [
5192                ("a", "1"), //
5193                ("b", "2"), //
5194                ("c", "3"), //
5195                ("d", "4"), //
5196                ("e", "5"), //
5197            ]
5198            .into_iter()
5199            .map(|(k, v)| (k.to_owned(), v.to_owned()))
5200            .collect();
5201
5202            // Set metadata on both the schema and a field within it.
5203            let schema = Arc::new(
5204                Schema::new(vec![
5205                    Field::new("a", DataType::Int64, true).with_metadata(metadata.clone()),
5206                ])
5207                .with_metadata(metadata)
5208                .clone(),
5209            );
5210            let batch = RecordBatch::new_empty(schema.clone());
5211
5212            let mut bytes = Vec::new();
5213            let mut w = StreamWriter::try_new(&mut bytes, batch.schema_ref()).unwrap();
5214            w.write(&batch).unwrap();
5215            w.finish().unwrap();
5216
5217            let mut h = std::hash::DefaultHasher::new();
5218            h.write(&bytes);
5219            h.finish()
5220        }
5221
5222        let expected = create_hash();
5223
5224        // Since there is randomness in the HashMap and we cannot specify our
5225        // own Hasher for the implementation used for metadata, run the above
5226        // code 20x and verify it does not change. This is not perfect but it
5227        // should be good enough.
5228        let all_passed = (0..20).all(|_| create_hash() == expected);
5229        assert!(all_passed);
5230    }
5231
5232    #[test]
5233    fn test_dictionary_tracker_reset() {
5234        let data_gen = IpcDataGenerator::default();
5235        let mut dictionary_tracker = DictionaryTracker::new(false);
5236        let writer_options = IpcWriteOptions::default();
5237        let mut compression_ctx = IpcWriteContext::default();
5238
5239        let schema = Arc::new(Schema::new(vec![Field::new(
5240            "a",
5241            DataType::Dictionary(Box::new(DataType::UInt8), Box::new(DataType::Utf8)),
5242            false,
5243        )]));
5244
5245        let mut write_single_batch_stream =
5246            |batch: RecordBatch, dict_tracker: &mut DictionaryTracker| -> Vec<u8> {
5247                let mut buffer = Vec::new();
5248
5249                // create a new IPC stream:
5250                let stream_header = data_gen.schema_to_bytes_with_dictionary_tracker(
5251                    &schema,
5252                    dict_tracker,
5253                    &writer_options,
5254                );
5255                _ = write_message(&mut buffer, stream_header, &writer_options).unwrap();
5256
5257                let (encoded_dicts, encoded_batch) = data_gen
5258                    .encode(&batch, dict_tracker, &writer_options, &mut compression_ctx)
5259                    .unwrap();
5260                for encoded_dict in encoded_dicts {
5261                    _ = write_message(&mut buffer, encoded_dict, &writer_options).unwrap();
5262                }
5263                _ = write_message(&mut buffer, encoded_batch, &writer_options).unwrap();
5264
5265                buffer
5266            };
5267
5268        let batch1 = RecordBatch::try_new(
5269            schema.clone(),
5270            vec![Arc::new(DictionaryArray::new(
5271                UInt8Array::from_iter_values([0]),
5272                Arc::new(StringArray::from_iter_values(["a"])),
5273            ))],
5274        )
5275        .unwrap();
5276        let buffer = write_single_batch_stream(batch1.clone(), &mut dictionary_tracker);
5277
5278        // ensure we can read the stream back
5279        let mut reader = StreamReader::try_new(Cursor::new(buffer), None).unwrap();
5280        let read_batch = reader.next().unwrap().unwrap();
5281        assert_eq!(read_batch, batch1);
5282
5283        // reset the dictionary tracker so it can be used for next stream
5284        dictionary_tracker.clear();
5285
5286        // now write a 2nd stream and ensure we can also read it:
5287        let batch2 = RecordBatch::try_new(
5288            schema.clone(),
5289            vec![Arc::new(DictionaryArray::new(
5290                UInt8Array::from_iter_values([0]),
5291                Arc::new(StringArray::from_iter_values(["a"])),
5292            ))],
5293        )
5294        .unwrap();
5295        let buffer = write_single_batch_stream(batch2.clone(), &mut dictionary_tracker);
5296        let mut reader = StreamReader::try_new(Cursor::new(buffer), None).unwrap();
5297        let read_batch = reader.next().unwrap().unwrap();
5298        assert_eq!(read_batch, batch2);
5299    }
5300}