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