Skip to main content

arrow_ipc/
writer.rs

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