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