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