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