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 % usize::from(write_options.alignment) != 0 {
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 as zstd::zstd_safe::CompressionLevel)) {
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: HashMap<String, String>,
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'](Result::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'](Result::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: HashMap::new(),
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'](Result::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'](Result::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'](Result::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'](Result::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
2278fn get_byte_array_buffers<O: OffsetSizeTrait>(data: &ArrayData) -> (Buffer, Buffer) {
2279    if data.is_empty() {
2280        // As per specification, offsets buffer has N+1 elements.
2281        // So an empty array should still be encoded with a single 0 offset.
2282        let mut offsets = MutableBuffer::new(size_of::<O>());
2283        offsets.extend_from_slice(O::usize_as(0).to_byte_slice());
2284        return (offsets.into(), MutableBuffer::new(0).into());
2285    }
2286
2287    let (offsets, original_start_offset, len) = reencode_offsets::<O>(&data.buffers()[0], data);
2288    let values = data.buffers()[1].slice_with_length(original_start_offset, len);
2289    (offsets, values)
2290}
2291
2292/// Similar logic as [`get_byte_array_buffers()`] but slices the child array instead
2293/// of a values buffer.
2294fn get_list_array_buffers<O: OffsetSizeTrait>(data: &ArrayData) -> (Buffer, ArrayData) {
2295    if data.is_empty() {
2296        // As per specification, offsets buffer has N+1 elements.
2297        // So an empty array should still be encoded with a single 0 offset.
2298        let mut offsets = MutableBuffer::new(size_of::<O>());
2299        offsets.extend_from_slice(O::usize_as(0).to_byte_slice());
2300        return (offsets.into(), data.child_data()[0].slice(0, 0));
2301    }
2302
2303    let (offsets, original_start_offset, len) = reencode_offsets::<O>(&data.buffers()[0], data);
2304    let child_data = data.child_data()[0].slice(original_start_offset, len);
2305    (offsets, child_data)
2306}
2307
2308/// Returns the offsets, sizes, and child data buffers for a ListView array.
2309///
2310/// Unlike List arrays, ListView arrays store both offsets and sizes explicitly,
2311/// and offsets can be non-monotonic. When slicing, we simply pass through the
2312/// offsets and sizes without re-encoding, and do not slice the child data.
2313fn get_list_view_array_buffers<O: OffsetSizeTrait>(
2314    data: &ArrayData,
2315) -> (Buffer, Buffer, ArrayData) {
2316    if data.is_empty() {
2317        return (
2318            MutableBuffer::new(0).into(),
2319            MutableBuffer::new(0).into(),
2320            data.child_data()[0].slice(0, 0),
2321        );
2322    }
2323
2324    let offsets = &data.buffers()[0];
2325    let sizes = &data.buffers()[1];
2326
2327    let element_size = std::mem::size_of::<O>();
2328    let offsets_slice =
2329        offsets.slice_with_length(data.offset() * element_size, data.len() * element_size);
2330    let sizes_slice =
2331        sizes.slice_with_length(data.offset() * element_size, data.len() * element_size);
2332
2333    let child_data = data.child_data()[0].clone();
2334
2335    (offsets_slice, sizes_slice, child_data)
2336}
2337
2338/// Returns the sliced views [`Buffer`] for a BinaryView/Utf8View array.
2339///
2340/// The views buffer is sliced to only include views in the valid range based on
2341/// the array's offset and length. This helps reduce the encoded size of sliced
2342/// arrays
2343///
2344fn get_or_truncate_buffer(array_data: &ArrayData) -> Buffer {
2345    let buffer = &array_data.buffers()[0];
2346    let layout = layout(array_data.data_type());
2347    let spec = &layout.buffers[0];
2348
2349    let byte_width = get_buffer_element_width(spec);
2350    let min_length = array_data.len() * byte_width;
2351    if buffer_need_truncate(array_data.offset(), buffer, spec, min_length) {
2352        let byte_offset = array_data.offset() * byte_width;
2353        let buffer_length = min(min_length, buffer.len() - byte_offset);
2354        buffer.slice_with_length(byte_offset, buffer_length)
2355    } else {
2356        buffer.clone()
2357    }
2358}
2359
2360/// Recursively encodes `array_data` into its IPC representation.
2361///
2362/// Output goes to two separate channels:
2363/// - `meta`: accumulates IPC metadata (`nodes` and `buffers`) for the flatbuffer header.
2364/// - `sink`: the raw Arrow data bytes that form the IPC message body.
2365fn write_array_data(
2366    array_data: &ArrayData,
2367    meta: &mut IpcMetadataBuilder,
2368    sink: &mut IpcBodySink<'_>,
2369    offset: i64,
2370    compression_codec: Option<CompressionCodec>,
2371    ipc_write_context: &mut IpcWriteContext,
2372    write_options: &IpcWriteOptions,
2373) -> Result<i64, ArrowError> {
2374    let mut offset = offset;
2375    let num_rows = array_data.len();
2376    if !matches!(array_data.data_type(), DataType::Null) {
2377        meta.nodes.push(crate::FieldNode::new(
2378            num_rows as i64,
2379            array_data.null_count() as i64,
2380        ));
2381    } else {
2382        // NullArray's null_count equals to len, but ArrayData null_count is always 0.
2383        meta.nodes
2384            .push(crate::FieldNode::new(num_rows as i64, num_rows as i64));
2385    }
2386    if has_validity_bitmap(array_data.data_type(), write_options) {
2387        // write null buffer if exists
2388        let null_buffer = match array_data.nulls() {
2389            None => {
2390                // create a buffer and fill it with valid bits
2391                let num_bytes = bit_util::ceil(num_rows, 8);
2392                let buffer = MutableBuffer::new(num_bytes);
2393                let buffer = buffer.with_bitset(num_bytes, true);
2394                buffer.into()
2395            }
2396            Some(buffer) => buffer.inner().sliced(),
2397        };
2398
2399        offset = encode_sink_buffer(
2400            null_buffer,
2401            meta,
2402            sink,
2403            offset,
2404            compression_codec,
2405            ipc_write_context,
2406            write_options.alignment,
2407        )?;
2408    }
2409
2410    let data_type = array_data.data_type();
2411    if matches!(data_type, DataType::Binary | DataType::Utf8) {
2412        let (offsets, values) = get_byte_array_buffers::<i32>(array_data);
2413        for buffer in [offsets, values] {
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        let (offsets, values) = get_byte_array_buffers::<i64>(array_data);
2455        for buffer in [offsets, values] {
2456            offset = encode_sink_buffer(
2457                buffer,
2458                meta,
2459                sink,
2460                offset,
2461                compression_codec,
2462                ipc_write_context,
2463                write_options.alignment,
2464            )?;
2465        }
2466    } else if DataType::is_numeric(data_type)
2467        || DataType::is_temporal(data_type)
2468        || matches!(
2469            array_data.data_type(),
2470            DataType::FixedSizeBinary(_) | DataType::Dictionary(_, _)
2471        )
2472    {
2473        // Truncate values
2474        assert_eq!(array_data.buffers().len(), 1);
2475
2476        let buffer = get_or_truncate_buffer(array_data);
2477        offset = encode_sink_buffer(
2478            buffer,
2479            meta,
2480            sink,
2481            offset,
2482            compression_codec,
2483            ipc_write_context,
2484            write_options.alignment,
2485        )?;
2486    } else if matches!(data_type, DataType::Boolean) {
2487        // Bools are special because the payload (= 1 bit) is smaller than the physical container elements (= bytes).
2488        // The array data may not start at the physical boundary of the underlying buffer, so we need to shift bits around.
2489        assert_eq!(array_data.buffers().len(), 1);
2490
2491        let buffer = &array_data.buffers()[0];
2492        let buffer = buffer.bit_slice(array_data.offset(), array_data.len());
2493        offset = encode_sink_buffer(
2494            buffer,
2495            meta,
2496            sink,
2497            offset,
2498            compression_codec,
2499            ipc_write_context,
2500            write_options.alignment,
2501        )?;
2502    } else if matches!(
2503        data_type,
2504        DataType::List(_) | DataType::LargeList(_) | DataType::Map(_, _)
2505    ) {
2506        assert_eq!(array_data.buffers().len(), 1);
2507        assert_eq!(array_data.child_data().len(), 1);
2508
2509        // Truncate offsets and the child data to avoid writing unnecessary data
2510        let (offsets, sliced_child_data) = match data_type {
2511            DataType::List(_) => get_list_array_buffers::<i32>(array_data),
2512            DataType::Map(_, _) => get_list_array_buffers::<i32>(array_data),
2513            DataType::LargeList(_) => get_list_array_buffers::<i64>(array_data),
2514            _ => unreachable!(),
2515        };
2516        offset = encode_sink_buffer(
2517            offsets,
2518            meta,
2519            sink,
2520            offset,
2521            compression_codec,
2522            ipc_write_context,
2523            write_options.alignment,
2524        )?;
2525        offset = write_array_data(
2526            &sliced_child_data,
2527            meta,
2528            sink,
2529            offset,
2530            compression_codec,
2531            ipc_write_context,
2532            write_options,
2533        )?;
2534        return Ok(offset);
2535    } else if matches!(
2536        data_type,
2537        DataType::ListView(_) | DataType::LargeListView(_)
2538    ) {
2539        assert_eq!(array_data.buffers().len(), 2); // offsets + sizes
2540        assert_eq!(array_data.child_data().len(), 1);
2541
2542        let (offsets, sizes, child_data) = match data_type {
2543            DataType::ListView(_) => get_list_view_array_buffers::<i32>(array_data),
2544            DataType::LargeListView(_) => get_list_view_array_buffers::<i64>(array_data),
2545            _ => unreachable!(),
2546        };
2547
2548        offset = encode_sink_buffer(
2549            offsets,
2550            meta,
2551            sink,
2552            offset,
2553            compression_codec,
2554            ipc_write_context,
2555            write_options.alignment,
2556        )?;
2557        offset = encode_sink_buffer(
2558            sizes,
2559            meta,
2560            sink,
2561            offset,
2562            compression_codec,
2563            ipc_write_context,
2564            write_options.alignment,
2565        )?;
2566
2567        offset = write_array_data(
2568            &child_data,
2569            meta,
2570            sink,
2571            offset,
2572            compression_codec,
2573            ipc_write_context,
2574            write_options,
2575        )?;
2576        return Ok(offset);
2577    } else if let DataType::FixedSizeList(_, fixed_size) = data_type {
2578        assert_eq!(array_data.child_data().len(), 1);
2579        let fixed_size = *fixed_size as usize;
2580
2581        let child_offset = array_data.offset() * fixed_size;
2582        let child_length = array_data.len() * fixed_size;
2583        let child_data = array_data.child_data()[0].slice(child_offset, child_length);
2584
2585        offset = write_array_data(
2586            &child_data,
2587            meta,
2588            sink,
2589            offset,
2590            compression_codec,
2591            ipc_write_context,
2592            write_options,
2593        )?;
2594        return Ok(offset);
2595    } else {
2596        for buffer in array_data.buffers() {
2597            offset = encode_sink_buffer(
2598                buffer.clone(),
2599                meta,
2600                sink,
2601                offset,
2602                compression_codec,
2603                ipc_write_context,
2604                write_options.alignment,
2605            )?;
2606        }
2607    }
2608
2609    match array_data.data_type() {
2610        DataType::Dictionary(_, _) => {}
2611        DataType::RunEndEncoded(_, _) => {
2612            // unslice the run encoded array.
2613            let arr = unslice_run_array(array_data.clone())?;
2614            // recursively write out nested structures
2615            for data_ref in arr.child_data() {
2616                // write the nested data (e.g list data)
2617                offset = write_array_data(
2618                    data_ref,
2619                    meta,
2620                    sink,
2621                    offset,
2622                    compression_codec,
2623                    ipc_write_context,
2624                    write_options,
2625                )?;
2626            }
2627        }
2628        _ => {
2629            // recursively write out nested structures
2630            for data_ref in array_data.child_data() {
2631                // write the nested data (e.g list data)
2632                offset = write_array_data(
2633                    data_ref,
2634                    meta,
2635                    sink,
2636                    offset,
2637                    compression_codec,
2638                    ipc_write_context,
2639                    write_options,
2640                )?;
2641            }
2642        }
2643    }
2644    Ok(offset)
2645}
2646
2647/// Encodes a single Arrow [`Buffer`] into the IPC body and records its metadata.
2648///
2649/// - `buffer`: the Arrow data buffer to encode (validity bitmap, offsets, values, etc.)
2650/// - `buffers`: in-progress list of IPC `Buffer` metadata entries (body offset + length) that
2651///   will eventually be serialised into the flatbuffer `RecordBatch` header.
2652/// - `sink`: destination for the actual encoded bytes; either a contiguous `Vec<u8>` for
2653///   in-memory writes, or a list of [`EncodedBuffer`] segments for deferred zero-copy streaming.
2654/// - `offset`: running byte offset into the IPC message body, used to compute the metadata entry.
2655/// - `compression_codec` / `ipc_write_context`: if `Some`, the buffer is compressed before
2656///   writing; `ipc_write_context` provides reusable scratch space across calls.
2657/// - `alignment`: each buffer is padded to this many bytes so the next buffer starts aligned.
2658///
2659/// Returns the updated `offset` (advanced by the encoded length plus any alignment padding).
2660fn encode_sink_buffer(
2661    buffer: Buffer,
2662    ipc_meta_data: &mut IpcMetadataBuilder,
2663    sink: &mut IpcBodySink<'_>,
2664    offset: i64,
2665    compression_codec: Option<CompressionCodec>,
2666    ipc_write_context: &mut IpcWriteContext,
2667    alignment: u8,
2668) -> Result<i64, ArrowError> {
2669    let (encoded, len) = match compression_codec {
2670        None => {
2671            let len = buffer.len() as i64;
2672            (EncodedBuffer::Raw(buffer), len)
2673        }
2674        Some(codec) => {
2675            let mut scratch = Vec::new();
2676            let written =
2677                codec.compress_to_vec(buffer.as_slice(), &mut scratch, ipc_write_context)?;
2678            let len = i64::try_from(written)
2679                .map_err(|e| ArrowError::InvalidArgumentError(format!("{e}")))?;
2680            (EncodedBuffer::Compressed(scratch), len)
2681        }
2682    };
2683
2684    let pad_len = pad_to_alignment(alignment, len as usize);
2685    sink.write(pad_len, encoded);
2686    ipc_meta_data.buffers.push(crate::Buffer::new(offset, len));
2687    Ok(offset + len + pad_len as i64)
2688}
2689
2690const PADDING: [u8; 64] = [0; 64];
2691
2692/// Estimates the number of [`EncodedBuffer`] segments that [`write_array_data`]
2693/// will produce for a column of the given type.
2694///
2695/// Based on the Arrow IPC buffer layout
2696/// (<https://arrow.apache.org/docs/format/Columnar.html#recordbatch-message>):
2697#[inline]
2698fn estimate_encoded_buffer_count(dt: &DataType) -> usize {
2699    match dt {
2700        DataType::Null => 0,
2701
2702        DataType::Binary | DataType::Utf8 | DataType::LargeBinary | DataType::LargeUtf8 => 3,
2703
2704        DataType::BinaryView | DataType::Utf8View => 3,
2705
2706        DataType::List(f) | DataType::LargeList(f) | DataType::Map(f, _) => {
2707            2 + estimate_encoded_buffer_count(f.data_type())
2708        }
2709
2710        DataType::ListView(f) | DataType::LargeListView(f) => {
2711            3 + estimate_encoded_buffer_count(f.data_type())
2712        }
2713
2714        DataType::FixedSizeList(f, _) => 1 + estimate_encoded_buffer_count(f.data_type()),
2715
2716        DataType::Struct(fields) => {
2717            1 + fields
2718                .iter()
2719                .map(|f| estimate_encoded_buffer_count(f.data_type()))
2720                .sum::<usize>()
2721        }
2722
2723        // Dictionary indices only; dictionary body is a separate IPC message.
2724        DataType::Dictionary(_, _) => 2,
2725
2726        DataType::Union(fields, UnionMode::Sparse) => {
2727            1 + fields
2728                .iter()
2729                .map(|(_, f)| estimate_encoded_buffer_count(f.data_type()))
2730                .sum::<usize>()
2731        }
2732        DataType::Union(fields, UnionMode::Dense) => {
2733            2 + fields
2734                .iter()
2735                .map(|(_, f)| estimate_encoded_buffer_count(f.data_type()))
2736                .sum::<usize>()
2737        }
2738
2739        DataType::RunEndEncoded(run_ends, values) => {
2740            estimate_encoded_buffer_count(run_ends.data_type())
2741                + estimate_encoded_buffer_count(values.data_type())
2742        }
2743        // Primitive, Bool, temporal, Decimal*, FixedSizeBinary: validity + values.
2744        _ => 2,
2745    }
2746}
2747
2748/// Calculate an alignment boundary and return the number of bytes needed to pad to the alignment boundary
2749#[inline]
2750fn pad_to_alignment(alignment: u8, len: usize) -> usize {
2751    let a = usize::from(alignment - 1);
2752    ((len + a) & !a) - len
2753}
2754
2755#[cfg(test)]
2756mod tests {
2757    use std::hash::Hasher;
2758    use std::io::Cursor;
2759    use std::io::Seek;
2760
2761    use arrow_array::builder::FixedSizeListBuilder;
2762    use arrow_array::builder::Float32Builder;
2763    use arrow_array::builder::Int64Builder;
2764    use arrow_array::builder::MapBuilder;
2765    use arrow_array::builder::StringViewBuilder;
2766    use arrow_array::builder::UnionBuilder;
2767    use arrow_array::builder::{
2768        GenericListBuilder, GenericListViewBuilder, ListBuilder, StringBuilder,
2769    };
2770    use arrow_array::builder::{PrimitiveRunBuilder, UInt32Builder};
2771    use arrow_array::types::*;
2772    use arrow_buffer::ScalarBuffer;
2773
2774    use crate::MetadataVersion;
2775    use crate::convert::fb_to_schema;
2776    use crate::reader::*;
2777    use crate::root_as_footer;
2778
2779    use super::*;
2780
2781    fn serialize_file(rb: &RecordBatch) -> Vec<u8> {
2782        let mut writer = FileWriter::try_new(vec![], rb.schema_ref()).unwrap();
2783        writer.write(rb).unwrap();
2784        writer.finish().unwrap();
2785        writer.into_inner().unwrap()
2786    }
2787
2788    fn deserialize_file(bytes: Vec<u8>) -> RecordBatch {
2789        let mut reader = FileReader::try_new(Cursor::new(bytes), None).unwrap();
2790        reader.next().unwrap().unwrap()
2791    }
2792
2793    fn serialize_stream(record: &RecordBatch) -> Vec<u8> {
2794        // Use 8-byte alignment so that the various `truncate_*` tests can be compactly written,
2795        // without needing to construct a giant array to spill over the 64-byte default alignment
2796        // boundary.
2797        const IPC_ALIGNMENT: usize = 8;
2798
2799        let mut stream_writer = StreamWriter::try_new_with_options(
2800            vec![],
2801            record.schema_ref(),
2802            IpcWriteOptions::try_new(IPC_ALIGNMENT, false, MetadataVersion::V5).unwrap(),
2803        )
2804        .unwrap();
2805        stream_writer.write(record).unwrap();
2806        stream_writer.finish().unwrap();
2807        stream_writer.into_inner().unwrap()
2808    }
2809
2810    fn deserialize_stream(bytes: Vec<u8>) -> RecordBatch {
2811        let mut stream_reader = StreamReader::try_new(Cursor::new(bytes), None).unwrap();
2812        stream_reader.next().unwrap().unwrap()
2813    }
2814
2815    /// Encodes record batches with [`StreamEncoder`] into one contiguous byte vector.
2816    ///
2817    /// This still exercises the sans-IO encoder path; `Vec<u8>` is only used
2818    /// as a convenient [`Write`] sink for comparing the resulting byte stream.
2819    fn encode_stream(
2820        schema: &Schema,
2821        batches: &[RecordBatch],
2822        options: IpcWriteOptions,
2823    ) -> Vec<u8> {
2824        let mut encoder = StreamEncoder::try_new_with_options(schema, options).unwrap();
2825        let mut bytes = Vec::new();
2826        for batch in batches {
2827            for buffer in encoder.encode(batch).unwrap() {
2828                bytes.write_all(buffer.as_slice()).unwrap();
2829            }
2830        }
2831        for buffer in encoder.finish().unwrap() {
2832            bytes.write_all(buffer.as_slice()).unwrap();
2833        }
2834        bytes
2835    }
2836
2837    fn write_stream(schema: &Schema, batches: &[RecordBatch], options: IpcWriteOptions) -> Vec<u8> {
2838        let mut bytes = Vec::new();
2839        let mut writer = StreamWriter::try_new_with_options(&mut bytes, schema, options).unwrap();
2840        for batch in batches {
2841            writer.write(batch).unwrap();
2842        }
2843        writer.finish().unwrap();
2844        bytes
2845    }
2846
2847    // StreamEncoder and StreamWriter currently use separate encoding paths, so these tests
2848    // verify the new sans-IO API preserves the existing IPC stream byte layout.
2849    #[tokio::test]
2850    async fn test_stream_encoder_async_writer_matches_stream_writer() {
2851        use tokio::io::{AsyncReadExt, AsyncWriteExt};
2852
2853        let batch = record_batch!(("a", Int32, [1, 2, 3]), ("b", Utf8, ["x", "y", "z"])).unwrap();
2854        let options = IpcWriteOptions::default();
2855        let expected = write_stream(
2856            batch.schema_ref(),
2857            std::slice::from_ref(&batch),
2858            options.clone(),
2859        );
2860
2861        let (mut sink, mut source) = tokio::io::duplex(64);
2862        let read = tokio::spawn(async move {
2863            let mut bytes = Vec::new();
2864            source.read_to_end(&mut bytes).await.unwrap();
2865            bytes
2866        });
2867
2868        let mut encoder = StreamEncoder::try_new_with_options(batch.schema_ref(), options).unwrap();
2869        for buffer in encoder.encode(&batch).unwrap() {
2870            sink.write_all(buffer.as_slice()).await.unwrap();
2871        }
2872        for buffer in encoder.finish().unwrap() {
2873            sink.write_all(buffer.as_slice()).await.unwrap();
2874        }
2875        sink.shutdown().await.unwrap();
2876
2877        let encoded = read.await.unwrap();
2878        assert_eq!(encoded, expected);
2879
2880        let mut reader = StreamReader::try_new(Cursor::new(encoded), None).unwrap();
2881        assert_eq!(reader.next().unwrap().unwrap(), batch);
2882        assert!(reader.next().is_none());
2883    }
2884
2885    #[test]
2886    fn test_stream_encoder_empty_stream_matches_stream_writer() {
2887        let schema = Schema::new(vec![Field::new("a", DataType::Int32, true)]);
2888        let options = IpcWriteOptions::default();
2889        let encoded = encode_stream(&schema, &[], options.clone());
2890        let written = write_stream(&schema, &[], options);
2891
2892        assert_eq!(encoded, written);
2893
2894        let mut reader = StreamReader::try_new(Cursor::new(encoded), None).unwrap();
2895        assert!(reader.next().is_none());
2896    }
2897
2898    #[test]
2899    fn test_stream_encoder_dictionary_batches_match_stream_writer() {
2900        let schema = Arc::new(Schema::new(vec![Field::new(
2901            "a",
2902            DataType::Dictionary(Box::new(DataType::UInt8), Box::new(DataType::Utf8)),
2903            false,
2904        )]));
2905        let batch = RecordBatch::try_new(
2906            schema.clone(),
2907            vec![Arc::new(DictionaryArray::new(
2908                UInt8Array::from_iter_values([0, 1, 0]),
2909                Arc::new(StringArray::from_iter_values(["a", "b"])),
2910            ))],
2911        )
2912        .unwrap();
2913        let options = IpcWriteOptions::default();
2914        let encoded = encode_stream(&schema, std::slice::from_ref(&batch), options.clone());
2915        let written = write_stream(&schema, std::slice::from_ref(&batch), options);
2916
2917        assert_eq!(encoded, written);
2918
2919        let mut reader = StreamReader::try_new(Cursor::new(encoded), None).unwrap();
2920        assert_eq!(reader.next().unwrap().unwrap(), batch);
2921        assert!(reader.next().is_none());
2922    }
2923
2924    #[test]
2925    #[cfg(feature = "lz4")]
2926    fn test_write_empty_record_batch_lz4_compression() {
2927        let schema = Schema::new(vec![Field::new("field1", DataType::Int32, true)]);
2928        let values: Vec<Option<i32>> = vec![];
2929        let array = Int32Array::from(values);
2930        let record_batch =
2931            RecordBatch::try_new(Arc::new(schema.clone()), vec![Arc::new(array)]).unwrap();
2932
2933        let mut file = tempfile::tempfile().unwrap();
2934
2935        {
2936            let write_option = IpcWriteOptions::try_new(8, false, crate::MetadataVersion::V5)
2937                .unwrap()
2938                .try_with_compression(Some(crate::CompressionType::LZ4_FRAME))
2939                .unwrap();
2940
2941            let mut writer =
2942                FileWriter::try_new_with_options(&mut file, &schema, write_option).unwrap();
2943            writer.write(&record_batch).unwrap();
2944            writer.finish().unwrap();
2945        }
2946        file.rewind().unwrap();
2947        {
2948            // read file
2949            let reader = FileReader::try_new(file, None).unwrap();
2950            for read_batch in reader {
2951                read_batch
2952                    .unwrap()
2953                    .columns()
2954                    .iter()
2955                    .zip(record_batch.columns())
2956                    .for_each(|(a, b)| {
2957                        assert_eq!(a.data_type(), b.data_type());
2958                        assert_eq!(a.len(), b.len());
2959                        assert_eq!(a.null_count(), b.null_count());
2960                    });
2961            }
2962        }
2963    }
2964
2965    #[test]
2966    #[cfg(feature = "lz4")]
2967    fn test_write_file_with_lz4_compression() {
2968        let schema = Schema::new(vec![Field::new("field1", DataType::Int32, true)]);
2969        let values: Vec<Option<i32>> = vec![Some(12), Some(1)];
2970        let array = Int32Array::from(values);
2971        let record_batch =
2972            RecordBatch::try_new(Arc::new(schema.clone()), vec![Arc::new(array)]).unwrap();
2973
2974        let mut file = tempfile::tempfile().unwrap();
2975        {
2976            let write_option = IpcWriteOptions::try_new(8, false, crate::MetadataVersion::V5)
2977                .unwrap()
2978                .try_with_compression(Some(crate::CompressionType::LZ4_FRAME))
2979                .unwrap();
2980
2981            let mut writer =
2982                FileWriter::try_new_with_options(&mut file, &schema, write_option).unwrap();
2983            writer.write(&record_batch).unwrap();
2984            writer.finish().unwrap();
2985        }
2986        file.rewind().unwrap();
2987        {
2988            // read file
2989            let reader = FileReader::try_new(file, None).unwrap();
2990            for read_batch in reader {
2991                read_batch
2992                    .unwrap()
2993                    .columns()
2994                    .iter()
2995                    .zip(record_batch.columns())
2996                    .for_each(|(a, b)| {
2997                        assert_eq!(a.data_type(), b.data_type());
2998                        assert_eq!(a.len(), b.len());
2999                        assert_eq!(a.null_count(), b.null_count());
3000                    });
3001            }
3002        }
3003    }
3004
3005    #[test]
3006    #[cfg(feature = "zstd")]
3007    fn test_write_file_with_zstd_compression() {
3008        let schema = Schema::new(vec![Field::new("field1", DataType::Int32, true)]);
3009        let values: Vec<Option<i32>> = vec![Some(12), Some(1)];
3010        let array = Int32Array::from(values);
3011        let record_batch =
3012            RecordBatch::try_new(Arc::new(schema.clone()), vec![Arc::new(array)]).unwrap();
3013        let mut file = tempfile::tempfile().unwrap();
3014        {
3015            let write_option = IpcWriteOptions::try_new(8, false, crate::MetadataVersion::V5)
3016                .unwrap()
3017                .try_with_compression(Some(crate::CompressionType::ZSTD))
3018                .unwrap()
3019                .try_with_compression_level(Some(1))
3020                .unwrap();
3021
3022            let mut writer =
3023                FileWriter::try_new_with_options(&mut file, &schema, write_option).unwrap();
3024            writer.write(&record_batch).unwrap();
3025            writer.finish().unwrap();
3026        }
3027        file.rewind().unwrap();
3028        {
3029            // read file
3030            let reader = FileReader::try_new(file, None).unwrap();
3031            for read_batch in reader {
3032                read_batch
3033                    .unwrap()
3034                    .columns()
3035                    .iter()
3036                    .zip(record_batch.columns())
3037                    .for_each(|(a, b)| {
3038                        assert_eq!(a.data_type(), b.data_type());
3039                        assert_eq!(a.len(), b.len());
3040                        assert_eq!(a.null_count(), b.null_count());
3041                    });
3042            }
3043        }
3044    }
3045
3046    #[test]
3047    fn test_write_file() {
3048        let schema = Schema::new(vec![Field::new("field1", DataType::UInt32, true)]);
3049        let values: Vec<Option<u32>> = vec![
3050            Some(999),
3051            None,
3052            Some(235),
3053            Some(123),
3054            None,
3055            None,
3056            None,
3057            None,
3058            None,
3059        ];
3060        let array1 = UInt32Array::from(values);
3061        let batch =
3062            RecordBatch::try_new(Arc::new(schema.clone()), vec![Arc::new(array1) as ArrayRef])
3063                .unwrap();
3064        let mut file = tempfile::tempfile().unwrap();
3065        {
3066            let mut writer = FileWriter::try_new(&mut file, &schema).unwrap();
3067
3068            writer.write(&batch).unwrap();
3069            writer.finish().unwrap();
3070        }
3071        file.rewind().unwrap();
3072
3073        {
3074            let mut reader = FileReader::try_new(file, None).unwrap();
3075            while let Some(Ok(read_batch)) = reader.next() {
3076                read_batch
3077                    .columns()
3078                    .iter()
3079                    .zip(batch.columns())
3080                    .for_each(|(a, b)| {
3081                        assert_eq!(a.data_type(), b.data_type());
3082                        assert_eq!(a.len(), b.len());
3083                        assert_eq!(a.null_count(), b.null_count());
3084                    });
3085            }
3086        }
3087    }
3088
3089    #[test]
3090    fn test_empty_utf8_ipc_writes_nonempty_offsets_buffer() {
3091        let name = StringArray::from(Vec::<String>::new());
3092        let (offsets, values) = get_byte_array_buffers::<i32>(&name.to_data());
3093
3094        assert_eq!(name.len(), 0);
3095        assert_eq!(
3096            offsets.len(),
3097            std::mem::size_of::<i32>(),
3098            "offsets buffer should contain one zero i32 offset"
3099        );
3100        assert_eq!(values.len(), 0, "values buffer should remain empty");
3101    }
3102
3103    #[test]
3104    fn test_empty_large_utf8_ipc_writes_nonempty_offsets_buffer() {
3105        let name = LargeStringArray::from(Vec::<String>::new());
3106        let (offsets, values) = get_byte_array_buffers::<i64>(&name.to_data());
3107
3108        assert_eq!(name.len(), 0);
3109        assert_eq!(
3110            offsets.len(),
3111            std::mem::size_of::<i64>(),
3112            "offsets buffer should contain one zero i64 offset"
3113        );
3114        assert_eq!(values.len(), 0, "values buffer should remain empty");
3115    }
3116
3117    #[test]
3118    fn test_empty_list_ipc_writes_nonempty_offsets_buffer() {
3119        let list = GenericListBuilder::<i32, _>::new(UInt32Builder::new()).finish();
3120        let (offsets, child_data) = get_list_array_buffers::<i32>(&list.to_data());
3121
3122        assert_eq!(list.len(), 0);
3123        assert_eq!(
3124            offsets.len(),
3125            std::mem::size_of::<i32>(),
3126            "offsets buffer should contain one zero i32 offset"
3127        );
3128        assert_eq!(child_data.len(), 0, "child data should remain empty");
3129    }
3130
3131    #[test]
3132    fn test_empty_large_list_ipc_writes_nonempty_offsets_buffer() {
3133        let list = GenericListBuilder::<i64, _>::new(UInt32Builder::new()).finish();
3134        let (offsets, child_data) = get_list_array_buffers::<i64>(&list.to_data());
3135
3136        assert_eq!(list.len(), 0);
3137        assert_eq!(
3138            offsets.len(),
3139            std::mem::size_of::<i64>(),
3140            "offsets buffer should contain one zero i64 offset"
3141        );
3142        assert_eq!(child_data.len(), 0, "child data should remain empty");
3143    }
3144
3145    fn write_null_file(options: IpcWriteOptions) {
3146        let schema = Schema::new(vec![
3147            Field::new("nulls", DataType::Null, true),
3148            Field::new("int32s", DataType::Int32, false),
3149            Field::new("nulls2", DataType::Null, true),
3150            Field::new("f64s", DataType::Float64, false),
3151        ]);
3152        let array1 = NullArray::new(32);
3153        let array2 = Int32Array::from(vec![1; 32]);
3154        let array3 = NullArray::new(32);
3155        let array4 = Float64Array::from(vec![f64::NAN; 32]);
3156        let batch = RecordBatch::try_new(
3157            Arc::new(schema.clone()),
3158            vec![
3159                Arc::new(array1) as ArrayRef,
3160                Arc::new(array2) as ArrayRef,
3161                Arc::new(array3) as ArrayRef,
3162                Arc::new(array4) as ArrayRef,
3163            ],
3164        )
3165        .unwrap();
3166        let mut file = tempfile::tempfile().unwrap();
3167        {
3168            let mut writer = FileWriter::try_new_with_options(&mut file, &schema, options).unwrap();
3169
3170            writer.write(&batch).unwrap();
3171            writer.finish().unwrap();
3172        }
3173
3174        file.rewind().unwrap();
3175
3176        {
3177            let reader = FileReader::try_new(file, None).unwrap();
3178            reader.for_each(|maybe_batch| {
3179                maybe_batch
3180                    .unwrap()
3181                    .columns()
3182                    .iter()
3183                    .zip(batch.columns())
3184                    .for_each(|(a, b)| {
3185                        assert_eq!(a.data_type(), b.data_type());
3186                        assert_eq!(a.len(), b.len());
3187                        assert_eq!(a.null_count(), b.null_count());
3188                    });
3189            });
3190        }
3191    }
3192    #[test]
3193    fn test_write_null_file_v4() {
3194        write_null_file(IpcWriteOptions::try_new(8, false, MetadataVersion::V4).unwrap());
3195        write_null_file(IpcWriteOptions::try_new(8, true, MetadataVersion::V4).unwrap());
3196        write_null_file(IpcWriteOptions::try_new(64, false, MetadataVersion::V4).unwrap());
3197        write_null_file(IpcWriteOptions::try_new(64, true, MetadataVersion::V4).unwrap());
3198    }
3199
3200    #[test]
3201    fn test_write_null_file_v5() {
3202        write_null_file(IpcWriteOptions::try_new(8, false, MetadataVersion::V5).unwrap());
3203        write_null_file(IpcWriteOptions::try_new(64, false, MetadataVersion::V5).unwrap());
3204    }
3205
3206    #[test]
3207    fn track_union_nested_dict() {
3208        let inner: DictionaryArray<Int32Type> = vec!["a", "b", "a"].into_iter().collect();
3209
3210        let array = Arc::new(inner) as ArrayRef;
3211
3212        // Dict field with id 2
3213        #[allow(deprecated)]
3214        let dctfield = Field::new_dict("dict", array.data_type().clone(), false, 0, false);
3215        let union_fields = [(0, Arc::new(dctfield))].into_iter().collect();
3216
3217        let types = [0, 0, 0].into_iter().collect::<ScalarBuffer<i8>>();
3218        let offsets = [0, 1, 2].into_iter().collect::<ScalarBuffer<i32>>();
3219
3220        let union = UnionArray::try_new(union_fields, types, Some(offsets), vec![array]).unwrap();
3221
3222        let schema = Arc::new(Schema::new(vec![Field::new(
3223            "union",
3224            union.data_type().clone(),
3225            false,
3226        )]));
3227
3228        let r#gen = IpcDataGenerator::default();
3229        let mut dict_tracker = DictionaryTracker::new(false);
3230        r#gen.schema_to_bytes_with_dictionary_tracker(
3231            &schema,
3232            &mut dict_tracker,
3233            &IpcWriteOptions::default(),
3234        );
3235
3236        let batch = RecordBatch::try_new(schema, vec![Arc::new(union)]).unwrap();
3237
3238        r#gen
3239            .encode(
3240                &batch,
3241                &mut dict_tracker,
3242                &Default::default(),
3243                &mut Default::default(),
3244            )
3245            .unwrap();
3246
3247        // The encoder will assign dict IDs itself to ensure uniqueness and ignore the dict ID in the schema
3248        // so we expect the dict will be keyed to 0
3249        assert!(dict_tracker.written.contains_key(&0));
3250    }
3251
3252    #[test]
3253    fn track_struct_nested_dict() {
3254        let inner: DictionaryArray<Int32Type> = vec!["a", "b", "a"].into_iter().collect();
3255
3256        let array = Arc::new(inner) as ArrayRef;
3257
3258        // Dict field with id 2
3259        #[allow(deprecated)]
3260        let dctfield = Arc::new(Field::new_dict(
3261            "dict",
3262            array.data_type().clone(),
3263            false,
3264            2,
3265            false,
3266        ));
3267
3268        let s = StructArray::from(vec![(dctfield, array)]);
3269        let struct_array = Arc::new(s) as ArrayRef;
3270
3271        let schema = Arc::new(Schema::new(vec![Field::new(
3272            "struct",
3273            struct_array.data_type().clone(),
3274            false,
3275        )]));
3276
3277        let r#gen = IpcDataGenerator::default();
3278        let mut dict_tracker = DictionaryTracker::new(false);
3279        r#gen.schema_to_bytes_with_dictionary_tracker(
3280            &schema,
3281            &mut dict_tracker,
3282            &IpcWriteOptions::default(),
3283        );
3284
3285        let batch = RecordBatch::try_new(schema, vec![struct_array]).unwrap();
3286
3287        r#gen
3288            .encode(
3289                &batch,
3290                &mut dict_tracker,
3291                &Default::default(),
3292                &mut Default::default(),
3293            )
3294            .unwrap();
3295
3296        assert!(dict_tracker.written.contains_key(&0));
3297    }
3298
3299    fn write_union_file(options: IpcWriteOptions) {
3300        let schema = Schema::new(vec![Field::new_union(
3301            "union",
3302            vec![0, 1],
3303            vec![
3304                Field::new("a", DataType::Int32, false),
3305                Field::new("c", DataType::Float64, false),
3306            ],
3307            UnionMode::Sparse,
3308        )]);
3309        let mut builder = UnionBuilder::with_capacity_sparse(5);
3310        builder.append::<Int32Type>("a", 1).unwrap();
3311        builder.append_null::<Int32Type>("a").unwrap();
3312        builder.append::<Float64Type>("c", 3.0).unwrap();
3313        builder.append_null::<Float64Type>("c").unwrap();
3314        builder.append::<Int32Type>("a", 4).unwrap();
3315        let union = builder.build().unwrap();
3316
3317        let batch =
3318            RecordBatch::try_new(Arc::new(schema.clone()), vec![Arc::new(union) as ArrayRef])
3319                .unwrap();
3320
3321        let mut file = tempfile::tempfile().unwrap();
3322        {
3323            let mut writer = FileWriter::try_new_with_options(&mut file, &schema, options).unwrap();
3324
3325            writer.write(&batch).unwrap();
3326            writer.finish().unwrap();
3327        }
3328        file.rewind().unwrap();
3329
3330        {
3331            let reader = FileReader::try_new(file, None).unwrap();
3332            reader.for_each(|maybe_batch| {
3333                maybe_batch
3334                    .unwrap()
3335                    .columns()
3336                    .iter()
3337                    .zip(batch.columns())
3338                    .for_each(|(a, b)| {
3339                        assert_eq!(a.data_type(), b.data_type());
3340                        assert_eq!(a.len(), b.len());
3341                        assert_eq!(a.null_count(), b.null_count());
3342                    });
3343            });
3344        }
3345    }
3346
3347    #[test]
3348    fn test_write_union_file_v4_v5() {
3349        write_union_file(IpcWriteOptions::try_new(8, false, MetadataVersion::V4).unwrap());
3350        write_union_file(IpcWriteOptions::try_new(8, false, MetadataVersion::V5).unwrap());
3351    }
3352
3353    #[test]
3354    fn test_write_view_types() {
3355        const LONG_TEST_STRING: &str =
3356            "This is a long string to make sure binary view array handles it";
3357        let schema = Schema::new(vec![
3358            Field::new("field1", DataType::BinaryView, true),
3359            Field::new("field2", DataType::Utf8View, true),
3360        ]);
3361        let values: Vec<Option<&[u8]>> = vec![
3362            Some(b"foo"),
3363            Some(b"bar"),
3364            Some(LONG_TEST_STRING.as_bytes()),
3365        ];
3366        let binary_array = BinaryViewArray::from_iter(values);
3367        let utf8_array =
3368            StringViewArray::from_iter(vec![Some("foo"), Some("bar"), Some(LONG_TEST_STRING)]);
3369        let record_batch = RecordBatch::try_new(
3370            Arc::new(schema.clone()),
3371            vec![Arc::new(binary_array), Arc::new(utf8_array)],
3372        )
3373        .unwrap();
3374
3375        let mut file = tempfile::tempfile().unwrap();
3376        {
3377            let mut writer = FileWriter::try_new(&mut file, &schema).unwrap();
3378            writer.write(&record_batch).unwrap();
3379            writer.finish().unwrap();
3380        }
3381        file.rewind().unwrap();
3382        {
3383            let mut reader = FileReader::try_new(&file, None).unwrap();
3384            let read_batch = reader.next().unwrap().unwrap();
3385            read_batch
3386                .columns()
3387                .iter()
3388                .zip(record_batch.columns())
3389                .for_each(|(a, b)| {
3390                    assert_eq!(a, b);
3391                });
3392        }
3393        file.rewind().unwrap();
3394        {
3395            let mut reader = FileReader::try_new(&file, Some(vec![0])).unwrap();
3396            let read_batch = reader.next().unwrap().unwrap();
3397            assert_eq!(read_batch.num_columns(), 1);
3398            let read_array = read_batch.column(0);
3399            let write_array = record_batch.column(0);
3400            assert_eq!(read_array, write_array);
3401        }
3402    }
3403
3404    #[test]
3405    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    fn test_large_slice_string() {
3605        let strings: Vec<_> = (0..8000)
3606            .map(|i| {
3607                if i % 2 == 0 {
3608                    Some(format!("value{i}"))
3609                } else {
3610                    None
3611                }
3612            })
3613            .collect();
3614
3615        ensure_roundtrip(Arc::new(StringArray::from(strings)));
3616    }
3617
3618    #[test]
3619    fn test_large_slice_string_list() {
3620        let mut ls = ListBuilder::new(StringBuilder::new());
3621
3622        let mut s = String::new();
3623        for row_number in 0..8000 {
3624            if row_number % 2 == 0 {
3625                for list_element in 0..1000 {
3626                    s.clear();
3627                    use std::fmt::Write;
3628                    write!(&mut s, "value{row_number}-{list_element}").unwrap();
3629                    ls.values().append_value(&s);
3630                }
3631                ls.append(true)
3632            } else {
3633                ls.append(false); // null
3634            }
3635        }
3636
3637        ensure_roundtrip(Arc::new(ls.finish()));
3638    }
3639
3640    #[test]
3641    fn test_large_slice_string_list_of_lists() {
3642        // The reason for the special test is to verify reencode_offsets which looks both at
3643        // the starting offset and the data offset.  So need a dataset where the starting_offset
3644        // is zero but the data offset is not.
3645        let mut ls = ListBuilder::new(ListBuilder::new(StringBuilder::new()));
3646
3647        for _ in 0..4000 {
3648            ls.values().append(true);
3649            ls.append(true)
3650        }
3651
3652        let mut s = String::new();
3653        for row_number in 0..4000 {
3654            if row_number % 2 == 0 {
3655                for list_element in 0..1000 {
3656                    s.clear();
3657                    use std::fmt::Write;
3658                    write!(&mut s, "value{row_number}-{list_element}").unwrap();
3659                    ls.values().values().append_value(&s);
3660                }
3661                ls.values().append(true);
3662                ls.append(true)
3663            } else {
3664                ls.append(false); // null
3665            }
3666        }
3667
3668        ensure_roundtrip(Arc::new(ls.finish()));
3669    }
3670
3671    /// Read/write a record batch to a File and Stream and ensure it is the same at the outout
3672    fn ensure_roundtrip(array: ArrayRef) {
3673        let num_rows = array.len();
3674        let orig_batch = RecordBatch::try_from_iter(vec![("a", array)]).unwrap();
3675        // take off the first element
3676        let sliced_batch = orig_batch.slice(1, num_rows - 1);
3677
3678        let schema = orig_batch.schema();
3679        let stream_data = {
3680            let mut writer = StreamWriter::try_new(vec![], &schema).unwrap();
3681            writer.write(&sliced_batch).unwrap();
3682            writer.into_inner().unwrap()
3683        };
3684        let read_batch = {
3685            let projection = None;
3686            let mut reader = StreamReader::try_new(Cursor::new(stream_data), projection).unwrap();
3687            reader
3688                .next()
3689                .expect("expect no errors reading batch")
3690                .expect("expect batch")
3691        };
3692        assert_eq!(sliced_batch, read_batch);
3693
3694        let file_data = {
3695            let mut writer = FileWriter::try_new_buffered(vec![], &schema).unwrap();
3696            writer.write(&sliced_batch).unwrap();
3697            writer.into_inner().unwrap().into_inner().unwrap()
3698        };
3699        let read_batch = {
3700            let projection = None;
3701            let mut reader = FileReader::try_new(Cursor::new(file_data), projection).unwrap();
3702            reader
3703                .next()
3704                .expect("expect no errors reading batch")
3705                .expect("expect batch")
3706        };
3707        assert_eq!(sliced_batch, read_batch);
3708
3709        // TODO test file writer/reader
3710    }
3711
3712    #[test]
3713    fn encode_bools_slice() {
3714        // Test case for https://github.com/apache/arrow-rs/issues/3496
3715        assert_bool_roundtrip([true, false], 1, 1);
3716
3717        // slice somewhere in the middle
3718        assert_bool_roundtrip(
3719            [
3720                true, false, true, true, false, false, true, true, true, false, false, false, true,
3721                true, true, true, false, false, false, false, true, true, true, true, true, false,
3722                false, false, false, false,
3723            ],
3724            13,
3725            17,
3726        );
3727
3728        // start at byte boundary, end in the middle
3729        assert_bool_roundtrip(
3730            [
3731                true, false, true, true, false, false, true, true, true, false, false, false,
3732            ],
3733            8,
3734            2,
3735        );
3736
3737        // start and stop and byte boundary
3738        assert_bool_roundtrip(
3739            [
3740                true, false, true, true, false, false, true, true, true, false, false, false, true,
3741                true, true, true, true, false, false, false, false, false,
3742            ],
3743            8,
3744            8,
3745        );
3746    }
3747
3748    fn assert_bool_roundtrip<const N: usize>(bools: [bool; N], offset: usize, length: usize) {
3749        let val_bool_field = Field::new("val", DataType::Boolean, false);
3750
3751        let schema = Arc::new(Schema::new(vec![val_bool_field]));
3752
3753        let bools = BooleanArray::from(bools.to_vec());
3754
3755        let batch = RecordBatch::try_new(Arc::clone(&schema), vec![Arc::new(bools)]).unwrap();
3756        let batch = batch.slice(offset, length);
3757
3758        let data = serialize_stream(&batch);
3759        let batch2 = deserialize_stream(data);
3760        assert_eq!(batch, batch2);
3761    }
3762
3763    #[test]
3764    fn test_run_array_unslice() {
3765        let total_len = 80;
3766        let vals: Vec<Option<i32>> = vec![Some(1), None, Some(2), Some(3), Some(4), None, Some(5)];
3767        let repeats: Vec<usize> = vec![3, 4, 1, 2];
3768        let mut input_array: Vec<Option<i32>> = Vec::with_capacity(total_len);
3769        for ix in 0_usize..32 {
3770            let repeat: usize = repeats[ix % repeats.len()];
3771            let val: Option<i32> = vals[ix % vals.len()];
3772            input_array.resize(input_array.len() + repeat, val);
3773        }
3774
3775        // Encode the input_array to run array
3776        let mut builder =
3777            PrimitiveRunBuilder::<Int16Type, Int32Type>::with_capacity(input_array.len());
3778        builder.extend(input_array.iter().copied());
3779        let run_array = builder.finish();
3780
3781        // test for all slice lengths.
3782        for slice_len in 1..=total_len {
3783            // test for offset = 0, slice length = slice_len
3784            let sliced_run_array: RunArray<Int16Type> =
3785                run_array.slice(0, slice_len).into_data().into();
3786
3787            // Create unsliced run array.
3788            let unsliced_run_array = into_zero_offset_run_array(sliced_run_array).unwrap();
3789            let typed = unsliced_run_array
3790                .downcast::<PrimitiveArray<Int32Type>>()
3791                .unwrap();
3792            let expected: Vec<Option<i32>> = input_array.iter().take(slice_len).copied().collect();
3793            let actual: Vec<Option<i32>> = typed.into_iter().collect();
3794            assert_eq!(expected, actual);
3795
3796            // test for offset = total_len - slice_len, length = slice_len
3797            let sliced_run_array: RunArray<Int16Type> = run_array
3798                .slice(total_len - slice_len, slice_len)
3799                .into_data()
3800                .into();
3801
3802            // Create unsliced run array.
3803            let unsliced_run_array = into_zero_offset_run_array(sliced_run_array).unwrap();
3804            let typed = unsliced_run_array
3805                .downcast::<PrimitiveArray<Int32Type>>()
3806                .unwrap();
3807            let expected: Vec<Option<i32>> = input_array
3808                .iter()
3809                .skip(total_len - slice_len)
3810                .copied()
3811                .collect();
3812            let actual: Vec<Option<i32>> = typed.into_iter().collect();
3813            assert_eq!(expected, actual);
3814        }
3815    }
3816
3817    fn generate_list_data<O: OffsetSizeTrait>() -> GenericListArray<O> {
3818        let mut ls = GenericListBuilder::<O, _>::new(UInt32Builder::new());
3819
3820        for i in 0..100_000 {
3821            for value in [i, i, i] {
3822                ls.values().append_value(value);
3823            }
3824            ls.append(true)
3825        }
3826
3827        ls.finish()
3828    }
3829
3830    fn generate_utf8view_list_data<O: OffsetSizeTrait>() -> GenericListArray<O> {
3831        let mut ls = GenericListBuilder::<O, _>::new(StringViewBuilder::new());
3832
3833        for i in 0..100_000 {
3834            for value in [
3835                format!("value{}", i),
3836                format!("value{}", i),
3837                format!("value{}", i),
3838            ] {
3839                ls.values().append_value(&value);
3840            }
3841            ls.append(true)
3842        }
3843
3844        ls.finish()
3845    }
3846
3847    fn generate_string_list_data<O: OffsetSizeTrait>() -> GenericListArray<O> {
3848        let mut ls = GenericListBuilder::<O, _>::new(StringBuilder::new());
3849
3850        for i in 0..100_000 {
3851            for value in [
3852                format!("value{}", i),
3853                format!("value{}", i),
3854                format!("value{}", i),
3855            ] {
3856                ls.values().append_value(&value);
3857            }
3858            ls.append(true)
3859        }
3860
3861        ls.finish()
3862    }
3863
3864    fn generate_nested_list_data<O: OffsetSizeTrait>() -> GenericListArray<O> {
3865        let mut ls =
3866            GenericListBuilder::<O, _>::new(GenericListBuilder::<O, _>::new(UInt32Builder::new()));
3867
3868        for _i in 0..10_000 {
3869            for j in 0..10 {
3870                for value in [j, j, j, j] {
3871                    ls.values().values().append_value(value);
3872                }
3873                ls.values().append(true)
3874            }
3875            ls.append(true);
3876        }
3877
3878        ls.finish()
3879    }
3880
3881    fn generate_nested_list_data_starting_at_zero<O: OffsetSizeTrait>() -> GenericListArray<O> {
3882        let mut ls =
3883            GenericListBuilder::<O, _>::new(GenericListBuilder::<O, _>::new(UInt32Builder::new()));
3884
3885        for _i in 0..999 {
3886            ls.values().append(true);
3887            ls.append(true);
3888        }
3889
3890        for j in 0..10 {
3891            for value in [j, j, j, j] {
3892                ls.values().values().append_value(value);
3893            }
3894            ls.values().append(true)
3895        }
3896        ls.append(true);
3897
3898        for i in 0..9_000 {
3899            for j in 0..10 {
3900                for value in [i + j, i + j, i + j, i + j] {
3901                    ls.values().values().append_value(value);
3902                }
3903                ls.values().append(true)
3904            }
3905            ls.append(true);
3906        }
3907
3908        ls.finish()
3909    }
3910
3911    fn generate_map_array_data() -> MapArray {
3912        let keys_builder = UInt32Builder::new();
3913        let values_builder = UInt32Builder::new();
3914
3915        let mut builder = MapBuilder::new(None, keys_builder, values_builder);
3916
3917        for i in 0..100_000 {
3918            for _j in 0..3 {
3919                builder.keys().append_value(i);
3920                builder.values().append_value(i * 2);
3921            }
3922            builder.append(true).unwrap();
3923        }
3924
3925        builder.finish()
3926    }
3927
3928    #[test]
3929    fn reencode_offsets_when_first_offset_is_not_zero() {
3930        let original_list = generate_list_data::<i32>();
3931        let original_data = original_list.into_data();
3932        let slice_data = original_data.slice(75, 7);
3933        let (new_offsets, original_start, length) =
3934            reencode_offsets::<i32>(&slice_data.buffers()[0], &slice_data);
3935        assert_eq!(
3936            vec![0, 3, 6, 9, 12, 15, 18, 21],
3937            new_offsets.typed_data::<i32>()
3938        );
3939        assert_eq!(225, original_start);
3940        assert_eq!(21, length);
3941    }
3942
3943    #[test]
3944    fn reencode_offsets_when_first_offset_is_zero() {
3945        let mut ls = GenericListBuilder::<i32, _>::new(UInt32Builder::new());
3946        // ls = [[], [35, 42]
3947        ls.append(true);
3948        ls.values().append_value(35);
3949        ls.values().append_value(42);
3950        ls.append(true);
3951        let original_list = ls.finish();
3952        let original_data = original_list.into_data();
3953
3954        let slice_data = original_data.slice(1, 1);
3955        let (new_offsets, original_start, length) =
3956            reencode_offsets::<i32>(&slice_data.buffers()[0], &slice_data);
3957        assert_eq!(vec![0, 2], new_offsets.typed_data::<i32>());
3958        assert_eq!(0, original_start);
3959        assert_eq!(2, length);
3960    }
3961
3962    /// Ensure when serde full & sliced versions they are equal to original input.
3963    /// Also ensure serialized sliced version is significantly smaller than serialized full.
3964    fn roundtrip_ensure_sliced_smaller(in_batch: RecordBatch, expected_size_factor: usize) {
3965        // test both full and sliced versions
3966        let in_sliced = in_batch.slice(999, 1);
3967
3968        let bytes_batch = serialize_file(&in_batch);
3969        let bytes_sliced = serialize_file(&in_sliced);
3970
3971        // serializing 1 row should be significantly smaller than serializing 100,000
3972        assert!(bytes_sliced.len() < (bytes_batch.len() / expected_size_factor));
3973
3974        // ensure both are still valid and equal to originals
3975        let out_batch = deserialize_file(bytes_batch);
3976        assert_eq!(in_batch, out_batch);
3977
3978        let out_sliced = deserialize_file(bytes_sliced);
3979        assert_eq!(in_sliced, out_sliced);
3980    }
3981
3982    #[test]
3983    fn encode_lists() {
3984        let val_inner = Field::new_list_field(DataType::UInt32, true);
3985        let val_list_field = Field::new("val", DataType::List(Arc::new(val_inner)), false);
3986        let schema = Arc::new(Schema::new(vec![val_list_field]));
3987
3988        let values = Arc::new(generate_list_data::<i32>());
3989
3990        let in_batch = RecordBatch::try_new(schema, vec![values]).unwrap();
3991        roundtrip_ensure_sliced_smaller(in_batch, 1000);
3992    }
3993
3994    #[test]
3995    fn encode_empty_list() {
3996        let val_inner = Field::new_list_field(DataType::UInt32, true);
3997        let val_list_field = Field::new("val", DataType::List(Arc::new(val_inner)), false);
3998        let schema = Arc::new(Schema::new(vec![val_list_field]));
3999
4000        let values = Arc::new(generate_list_data::<i32>());
4001
4002        let in_batch = RecordBatch::try_new(schema, vec![values])
4003            .unwrap()
4004            .slice(999, 0);
4005        let out_batch = deserialize_file(serialize_file(&in_batch));
4006        assert_eq!(in_batch, out_batch);
4007    }
4008
4009    #[test]
4010    fn encode_large_lists() {
4011        let val_inner = Field::new_list_field(DataType::UInt32, true);
4012        let val_list_field = Field::new("val", DataType::LargeList(Arc::new(val_inner)), false);
4013        let schema = Arc::new(Schema::new(vec![val_list_field]));
4014
4015        let values = Arc::new(generate_list_data::<i64>());
4016
4017        // ensure when serde full & sliced versions they are equal to original input
4018        // also ensure serialized sliced version is significantly smaller than serialized full
4019        let in_batch = RecordBatch::try_new(schema, vec![values]).unwrap();
4020        roundtrip_ensure_sliced_smaller(in_batch, 1000);
4021    }
4022
4023    #[test]
4024    fn encode_large_lists_non_zero_offset() {
4025        let val_inner = Field::new_list_field(DataType::UInt32, true);
4026        let val_list_field = Field::new("val", DataType::LargeList(Arc::new(val_inner)), false);
4027        let schema = Arc::new(Schema::new(vec![val_list_field]));
4028
4029        let values = Arc::new(generate_list_data::<i64>());
4030
4031        check_sliced_list_array(schema, values);
4032    }
4033
4034    #[test]
4035    fn encode_large_lists_string_non_zero_offset() {
4036        let val_inner = Field::new_list_field(DataType::Utf8, true);
4037        let val_list_field = Field::new("val", DataType::LargeList(Arc::new(val_inner)), false);
4038        let schema = Arc::new(Schema::new(vec![val_list_field]));
4039
4040        let values = Arc::new(generate_string_list_data::<i64>());
4041
4042        check_sliced_list_array(schema, values);
4043    }
4044
4045    #[test]
4046    fn encode_large_list_string_view_non_zero_offset() {
4047        let val_inner = Field::new_list_field(DataType::Utf8View, true);
4048        let val_list_field = Field::new("val", DataType::LargeList(Arc::new(val_inner)), false);
4049        let schema = Arc::new(Schema::new(vec![val_list_field]));
4050
4051        let values = Arc::new(generate_utf8view_list_data::<i64>());
4052
4053        check_sliced_list_array(schema, values);
4054    }
4055
4056    fn check_sliced_list_array(schema: Arc<Schema>, values: Arc<GenericListArray<i64>>) {
4057        for (offset, len) in [(999, 1), (0, 13), (47, 12), (values.len() - 13, 13)] {
4058            let in_batch = RecordBatch::try_new(schema.clone(), vec![values.clone()])
4059                .unwrap()
4060                .slice(offset, len);
4061            let out_batch = deserialize_file(serialize_file(&in_batch));
4062            assert_eq!(in_batch, out_batch);
4063        }
4064    }
4065
4066    #[test]
4067    fn encode_nested_lists() {
4068        let inner_int = Arc::new(Field::new_list_field(DataType::UInt32, true));
4069        let inner_list_field = Arc::new(Field::new_list_field(DataType::List(inner_int), true));
4070        let list_field = Field::new("val", DataType::List(inner_list_field), true);
4071        let schema = Arc::new(Schema::new(vec![list_field]));
4072
4073        let values = Arc::new(generate_nested_list_data::<i32>());
4074
4075        let in_batch = RecordBatch::try_new(schema, vec![values]).unwrap();
4076        roundtrip_ensure_sliced_smaller(in_batch, 1000);
4077    }
4078
4079    #[test]
4080    fn encode_nested_lists_starting_at_zero() {
4081        let inner_int = Arc::new(Field::new("item", DataType::UInt32, true));
4082        let inner_list_field = Arc::new(Field::new("item", DataType::List(inner_int), true));
4083        let list_field = Field::new("val", DataType::List(inner_list_field), true);
4084        let schema = Arc::new(Schema::new(vec![list_field]));
4085
4086        let values = Arc::new(generate_nested_list_data_starting_at_zero::<i32>());
4087
4088        let in_batch = RecordBatch::try_new(schema, vec![values]).unwrap();
4089        roundtrip_ensure_sliced_smaller(in_batch, 1);
4090    }
4091
4092    #[test]
4093    fn encode_map_array() {
4094        let keys = Arc::new(Field::new(
4095            Field::MAP_KEY_FIELD_DEFAULT_NAME,
4096            DataType::UInt32,
4097            false,
4098        ));
4099        let values = Arc::new(Field::new(
4100            Field::MAP_VALUE_FIELD_DEFAULT_NAME,
4101            DataType::UInt32,
4102            true,
4103        ));
4104        let map_field = Field::new_map(
4105            "map",
4106            Field::MAP_ENTRIES_FIELD_DEFAULT_NAME,
4107            keys,
4108            values,
4109            false,
4110            true,
4111        );
4112        let schema = Arc::new(Schema::new(vec![map_field]));
4113
4114        let values = Arc::new(generate_map_array_data());
4115
4116        let in_batch = RecordBatch::try_new(schema, vec![values]).unwrap();
4117        roundtrip_ensure_sliced_smaller(in_batch, 1000);
4118    }
4119
4120    fn generate_list_view_data<O: OffsetSizeTrait>() -> GenericListViewArray<O> {
4121        let mut builder = GenericListViewBuilder::<O, _>::new(UInt32Builder::new());
4122
4123        for i in 0u32..100_000 {
4124            if i.is_multiple_of(10_000) {
4125                builder.append(false);
4126                continue;
4127            }
4128            for value in [i, i, i] {
4129                builder.values().append_value(value);
4130            }
4131            builder.append(true);
4132        }
4133
4134        builder.finish()
4135    }
4136
4137    #[test]
4138    fn encode_list_view_arrays() {
4139        let val_inner = Field::new_list_field(DataType::UInt32, true);
4140        let val_field = Field::new("val", DataType::ListView(Arc::new(val_inner)), true);
4141        let schema = Arc::new(Schema::new(vec![val_field]));
4142
4143        let values = Arc::new(generate_list_view_data::<i32>());
4144
4145        let in_batch = RecordBatch::try_new(schema, vec![values]).unwrap();
4146        let out_batch = deserialize_file(serialize_file(&in_batch));
4147        assert_eq!(in_batch, out_batch);
4148    }
4149
4150    #[test]
4151    fn encode_large_list_view_arrays() {
4152        let val_inner = Field::new_list_field(DataType::UInt32, true);
4153        let val_field = Field::new("val", DataType::LargeListView(Arc::new(val_inner)), true);
4154        let schema = Arc::new(Schema::new(vec![val_field]));
4155
4156        let values = Arc::new(generate_list_view_data::<i64>());
4157
4158        let in_batch = RecordBatch::try_new(schema, vec![values]).unwrap();
4159        let out_batch = deserialize_file(serialize_file(&in_batch));
4160        assert_eq!(in_batch, out_batch);
4161    }
4162
4163    #[test]
4164    fn check_sliced_list_view_array() {
4165        let inner = Field::new_list_field(DataType::UInt32, true);
4166        let field = Field::new("val", DataType::ListView(Arc::new(inner)), true);
4167        let schema = Arc::new(Schema::new(vec![field]));
4168        let values = Arc::new(generate_list_view_data::<i32>());
4169
4170        for (offset, len) in [(999, 1), (0, 13), (47, 12), (values.len() - 13, 13)] {
4171            let in_batch = RecordBatch::try_new(schema.clone(), vec![values.clone()])
4172                .unwrap()
4173                .slice(offset, len);
4174            let out_batch = deserialize_file(serialize_file(&in_batch));
4175            assert_eq!(in_batch, out_batch);
4176        }
4177    }
4178
4179    #[test]
4180    fn check_sliced_large_list_view_array() {
4181        let inner = Field::new_list_field(DataType::UInt32, true);
4182        let field = Field::new("val", DataType::LargeListView(Arc::new(inner)), true);
4183        let schema = Arc::new(Schema::new(vec![field]));
4184        let values = Arc::new(generate_list_view_data::<i64>());
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    fn generate_nested_list_view_data<O: OffsetSizeTrait>() -> GenericListViewArray<O> {
4196        let inner_builder = UInt32Builder::new();
4197        let middle_builder = GenericListViewBuilder::<O, _>::new(inner_builder);
4198        let mut outer_builder = GenericListViewBuilder::<O, _>::new(middle_builder);
4199
4200        for i in 0u32..10_000 {
4201            if i.is_multiple_of(1_000) {
4202                outer_builder.append(false);
4203                continue;
4204            }
4205
4206            for _ in 0..3 {
4207                for value in [i, i + 1, i + 2] {
4208                    outer_builder.values().values().append_value(value);
4209                }
4210                outer_builder.values().append(true);
4211            }
4212            outer_builder.append(true);
4213        }
4214
4215        outer_builder.finish()
4216    }
4217
4218    #[test]
4219    fn encode_nested_list_views() {
4220        let inner_int = Arc::new(Field::new_list_field(DataType::UInt32, true));
4221        let inner_list_field = Arc::new(Field::new_list_field(DataType::ListView(inner_int), true));
4222        let list_field = Field::new("val", DataType::ListView(inner_list_field), true);
4223        let schema = Arc::new(Schema::new(vec![list_field]));
4224
4225        let values = Arc::new(generate_nested_list_view_data::<i32>());
4226
4227        let in_batch = RecordBatch::try_new(schema, vec![values]).unwrap();
4228        let out_batch = deserialize_file(serialize_file(&in_batch));
4229        assert_eq!(in_batch, out_batch);
4230    }
4231
4232    fn test_roundtrip_list_view_of_dict_impl<OffsetSize: OffsetSizeTrait, U: ArrowNativeType>(
4233        list_data_type: DataType,
4234        offsets: &[U; 5],
4235        sizes: &[U; 4],
4236    ) {
4237        let values = StringArray::from(vec![Some("alpha"), None, Some("beta"), Some("gamma")]);
4238        let keys = Int32Array::from_iter_values([0, 0, 1, 2, 3, 0, 2]);
4239        let dict_array = DictionaryArray::new(keys, Arc::new(values));
4240        let dict_data = dict_array.to_data();
4241
4242        let value_offsets = Buffer::from_slice_ref(offsets);
4243        let value_sizes = Buffer::from_slice_ref(sizes);
4244
4245        let list_data = ArrayData::builder(list_data_type)
4246            .len(4)
4247            .add_buffer(value_offsets)
4248            .add_buffer(value_sizes)
4249            .add_child_data(dict_data)
4250            .build()
4251            .unwrap();
4252        let list_view_array = GenericListViewArray::<OffsetSize>::from(list_data);
4253
4254        let schema = Arc::new(Schema::new(vec![Field::new(
4255            "f1",
4256            list_view_array.data_type().clone(),
4257            false,
4258        )]));
4259        let input_batch = RecordBatch::try_new(schema, vec![Arc::new(list_view_array)]).unwrap();
4260
4261        let output_batch = deserialize_file(serialize_file(&input_batch));
4262        assert_eq!(input_batch, output_batch);
4263
4264        let output_batch = deserialize_stream(serialize_stream(&input_batch));
4265        assert_eq!(input_batch, output_batch);
4266    }
4267
4268    #[test]
4269    fn test_roundtrip_list_view_of_dict() {
4270        #[allow(deprecated)]
4271        let list_data_type = DataType::ListView(Arc::new(Field::new_dict(
4272            "item",
4273            DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
4274            true,
4275            1,
4276            false,
4277        )));
4278        let offsets: &[i32; 5] = &[0, 2, 4, 4, 7];
4279        let sizes: &[i32; 4] = &[2, 2, 0, 3];
4280        test_roundtrip_list_view_of_dict_impl::<i32, i32>(list_data_type, offsets, sizes);
4281    }
4282
4283    #[test]
4284    fn test_roundtrip_large_list_view_of_dict() {
4285        #[allow(deprecated)]
4286        let list_data_type = DataType::LargeListView(Arc::new(Field::new_dict(
4287            "item",
4288            DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
4289            true,
4290            2,
4291            false,
4292        )));
4293        let offsets: &[i64; 5] = &[0, 2, 4, 4, 7];
4294        let sizes: &[i64; 4] = &[2, 2, 0, 3];
4295        test_roundtrip_list_view_of_dict_impl::<i64, i64>(list_data_type, offsets, sizes);
4296    }
4297
4298    #[test]
4299    fn test_roundtrip_sliced_list_view_of_dict() {
4300        #[allow(deprecated)]
4301        let list_data_type = DataType::ListView(Arc::new(Field::new_dict(
4302            "item",
4303            DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
4304            true,
4305            3,
4306            false,
4307        )));
4308
4309        let values = StringArray::from(vec![Some("alpha"), None, Some("beta"), Some("gamma")]);
4310        let keys = Int32Array::from_iter_values([0, 0, 1, 2, 3, 0, 2, 1, 0, 3, 2, 1]);
4311        let dict_array = DictionaryArray::new(keys, Arc::new(values));
4312        let dict_data = dict_array.to_data();
4313
4314        let offsets: &[i32; 7] = &[0, 2, 4, 4, 7, 9, 12];
4315        let sizes: &[i32; 6] = &[2, 2, 0, 3, 2, 3];
4316        let value_offsets = Buffer::from_slice_ref(offsets);
4317        let value_sizes = Buffer::from_slice_ref(sizes);
4318
4319        let list_data = ArrayData::builder(list_data_type)
4320            .len(6)
4321            .add_buffer(value_offsets)
4322            .add_buffer(value_sizes)
4323            .add_child_data(dict_data)
4324            .build()
4325            .unwrap();
4326        let list_view_array = GenericListViewArray::<i32>::from(list_data);
4327
4328        let schema = Arc::new(Schema::new(vec![Field::new(
4329            "f1",
4330            list_view_array.data_type().clone(),
4331            false,
4332        )]));
4333        let input_batch = RecordBatch::try_new(schema, vec![Arc::new(list_view_array)]).unwrap();
4334
4335        let sliced_batch = input_batch.slice(1, 4);
4336
4337        let output_batch = deserialize_file(serialize_file(&sliced_batch));
4338        assert_eq!(sliced_batch, output_batch);
4339
4340        let output_batch = deserialize_stream(serialize_stream(&sliced_batch));
4341        assert_eq!(sliced_batch, output_batch);
4342    }
4343
4344    #[test]
4345    fn test_roundtrip_dense_union_of_dict() {
4346        let values = StringArray::from(vec![Some("alpha"), None, Some("beta"), Some("gamma")]);
4347        let keys = Int32Array::from_iter_values([0, 0, 1, 2, 3, 0, 2]);
4348        let dict_array = DictionaryArray::new(keys, Arc::new(values));
4349
4350        #[allow(deprecated)]
4351        let dict_field = Arc::new(Field::new_dict(
4352            "dict",
4353            DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
4354            true,
4355            1,
4356            false,
4357        ));
4358        let int_field = Arc::new(Field::new("int", DataType::Int32, false));
4359        let union_fields = UnionFields::try_new(vec![0, 1], vec![dict_field, int_field]).unwrap();
4360
4361        let types = ScalarBuffer::from(vec![0i8, 0, 1, 0, 1, 0, 0]);
4362        let offsets = ScalarBuffer::from(vec![0i32, 1, 0, 2, 1, 3, 4]);
4363
4364        let int_array = Int32Array::from(vec![100, 200]);
4365
4366        let union = UnionArray::try_new(
4367            union_fields.clone(),
4368            types,
4369            Some(offsets),
4370            vec![Arc::new(dict_array), Arc::new(int_array)],
4371        )
4372        .unwrap();
4373
4374        let schema = Arc::new(Schema::new(vec![Field::new(
4375            "union",
4376            DataType::Union(union_fields, UnionMode::Dense),
4377            false,
4378        )]));
4379        let input_batch = RecordBatch::try_new(schema, vec![Arc::new(union)]).unwrap();
4380
4381        let output_batch = deserialize_file(serialize_file(&input_batch));
4382        assert_eq!(input_batch, output_batch);
4383
4384        let output_batch = deserialize_stream(serialize_stream(&input_batch));
4385        assert_eq!(input_batch, output_batch);
4386    }
4387
4388    #[test]
4389    fn test_roundtrip_sparse_union_of_dict() {
4390        let values = StringArray::from(vec![Some("alpha"), None, Some("beta"), Some("gamma")]);
4391        let keys = Int32Array::from_iter_values([0, 0, 1, 2, 3, 0, 2]);
4392        let dict_array = DictionaryArray::new(keys, Arc::new(values));
4393
4394        #[allow(deprecated)]
4395        let dict_field = Arc::new(Field::new_dict(
4396            "dict",
4397            DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
4398            true,
4399            2,
4400            false,
4401        ));
4402        let int_field = Arc::new(Field::new("int", DataType::Int32, false));
4403        let union_fields = UnionFields::try_new(vec![0, 1], vec![dict_field, int_field]).unwrap();
4404
4405        let types = ScalarBuffer::from(vec![0i8, 0, 1, 0, 1, 0, 0]);
4406
4407        let int_array = Int32Array::from(vec![0, 0, 100, 0, 200, 0, 0]);
4408
4409        let union = UnionArray::try_new(
4410            union_fields.clone(),
4411            types,
4412            None,
4413            vec![Arc::new(dict_array), Arc::new(int_array)],
4414        )
4415        .unwrap();
4416
4417        let schema = Arc::new(Schema::new(vec![Field::new(
4418            "union",
4419            DataType::Union(union_fields, UnionMode::Sparse),
4420            false,
4421        )]));
4422        let input_batch = RecordBatch::try_new(schema, vec![Arc::new(union)]).unwrap();
4423
4424        let output_batch = deserialize_file(serialize_file(&input_batch));
4425        assert_eq!(input_batch, output_batch);
4426
4427        let output_batch = deserialize_stream(serialize_stream(&input_batch));
4428        assert_eq!(input_batch, output_batch);
4429    }
4430
4431    #[test]
4432    fn test_roundtrip_map_with_dict_keys() {
4433        // Building a map array is a bit involved. We first build a struct arary that has a key and
4434        // value field and then use that to build the actual map array.
4435        let key_values = StringArray::from(vec!["key_a", "key_b", "key_c"]);
4436        let keys = Int32Array::from_iter_values([0, 1, 2, 0, 1, 0]);
4437        let dict_keys = DictionaryArray::new(keys, Arc::new(key_values));
4438
4439        let values = Int32Array::from(vec![1, 2, 3, 4, 5, 6]);
4440
4441        #[allow(deprecated)]
4442        let entries_field = Arc::new(Field::new(
4443            Field::MAP_ENTRIES_FIELD_DEFAULT_NAME,
4444            DataType::Struct(
4445                vec![
4446                    Field::new_dict(
4447                        Field::MAP_KEY_FIELD_DEFAULT_NAME,
4448                        DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
4449                        false,
4450                        1,
4451                        false,
4452                    ),
4453                    Field::new(Field::MAP_VALUE_FIELD_DEFAULT_NAME, DataType::Int32, true),
4454                ]
4455                .into(),
4456            ),
4457            false,
4458        ));
4459
4460        let entries = StructArray::from(vec![
4461            (
4462                Arc::new(Field::new(
4463                    Field::MAP_KEY_FIELD_DEFAULT_NAME,
4464                    DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
4465                    false,
4466                )),
4467                Arc::new(dict_keys) as ArrayRef,
4468            ),
4469            (
4470                Arc::new(Field::new(
4471                    Field::MAP_VALUE_FIELD_DEFAULT_NAME,
4472                    DataType::Int32,
4473                    true,
4474                )),
4475                Arc::new(values) as ArrayRef,
4476            ),
4477        ]);
4478
4479        let offsets = Buffer::from_slice_ref([0i32, 2, 4, 6]);
4480
4481        let map_data = ArrayData::builder(DataType::Map(entries_field, false))
4482            .len(3)
4483            .add_buffer(offsets)
4484            .add_child_data(entries.into_data())
4485            .build()
4486            .unwrap();
4487        let map_array = MapArray::from(map_data);
4488
4489        let schema = Arc::new(Schema::new(vec![Field::new(
4490            "map",
4491            map_array.data_type().clone(),
4492            false,
4493        )]));
4494        let input_batch = RecordBatch::try_new(schema, vec![Arc::new(map_array)]).unwrap();
4495
4496        let output_batch = deserialize_file(serialize_file(&input_batch));
4497        assert_eq!(input_batch, output_batch);
4498
4499        let output_batch = deserialize_stream(serialize_stream(&input_batch));
4500        assert_eq!(input_batch, output_batch);
4501    }
4502
4503    #[test]
4504    fn test_roundtrip_map_with_dict_values() {
4505        // Building a map array is a bit involved. We first build a struct arary that has a key and
4506        // value field and then use that to build the actual map array.
4507        let keys = StringArray::from(vec!["a", "b", "c", "d", "e", "f"]);
4508
4509        let value_values = StringArray::from(vec!["val_x", "val_y", "val_z"]);
4510        let value_keys = Int32Array::from_iter_values([0, 1, 2, 0, 1, 0]);
4511        let dict_values = DictionaryArray::new(value_keys, Arc::new(value_values));
4512
4513        #[allow(deprecated)]
4514        let entries_field = Arc::new(Field::new(
4515            Field::MAP_ENTRIES_FIELD_DEFAULT_NAME,
4516            DataType::Struct(
4517                vec![
4518                    Field::new(Field::MAP_KEY_FIELD_DEFAULT_NAME, DataType::Utf8, false),
4519                    Field::new_dict(
4520                        Field::MAP_VALUE_FIELD_DEFAULT_NAME,
4521                        DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
4522                        true,
4523                        2,
4524                        false,
4525                    ),
4526                ]
4527                .into(),
4528            ),
4529            false,
4530        ));
4531
4532        let entries = StructArray::from(vec![
4533            (
4534                Arc::new(Field::new(
4535                    Field::MAP_KEY_FIELD_DEFAULT_NAME,
4536                    DataType::Utf8,
4537                    false,
4538                )),
4539                Arc::new(keys) as ArrayRef,
4540            ),
4541            (
4542                Arc::new(Field::new(
4543                    Field::MAP_VALUE_FIELD_DEFAULT_NAME,
4544                    DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
4545                    true,
4546                )),
4547                Arc::new(dict_values) as ArrayRef,
4548            ),
4549        ]);
4550
4551        let offsets = Buffer::from_slice_ref([0i32, 2, 4, 6]);
4552
4553        let map_data = ArrayData::builder(DataType::Map(entries_field, false))
4554            .len(3)
4555            .add_buffer(offsets)
4556            .add_child_data(entries.into_data())
4557            .build()
4558            .unwrap();
4559        let map_array = MapArray::from(map_data);
4560
4561        let schema = Arc::new(Schema::new(vec![Field::new(
4562            "map",
4563            map_array.data_type().clone(),
4564            false,
4565        )]));
4566        let input_batch = RecordBatch::try_new(schema, vec![Arc::new(map_array)]).unwrap();
4567
4568        let output_batch = deserialize_file(serialize_file(&input_batch));
4569        assert_eq!(input_batch, output_batch);
4570
4571        let output_batch = deserialize_stream(serialize_stream(&input_batch));
4572        assert_eq!(input_batch, output_batch);
4573    }
4574
4575    #[test]
4576    fn test_decimal128_alignment16_is_sufficient() {
4577        const IPC_ALIGNMENT: usize = 16;
4578
4579        // Test a bunch of different dimensions to ensure alignment is never an issue.
4580        // For example, if we only test `num_cols = 1` then even with alignment 8 this
4581        // test would _happen_ to pass, even though for different dimensions like
4582        // `num_cols = 2` it would fail.
4583        for num_cols in [1, 2, 3, 17, 50, 73, 99] {
4584            let num_rows = (num_cols * 7 + 11) % 100; // Deterministic swizzle
4585
4586            let mut fields = Vec::new();
4587            let mut arrays = Vec::new();
4588            for i in 0..num_cols {
4589                let field = Field::new(format!("col_{i}"), DataType::Decimal128(38, 10), true);
4590                let array = Decimal128Array::from(vec![num_cols as i128; num_rows]);
4591                fields.push(field);
4592                arrays.push(Arc::new(array) as Arc<dyn Array>);
4593            }
4594            let schema = Schema::new(fields);
4595            let batch = RecordBatch::try_new(Arc::new(schema), arrays).unwrap();
4596
4597            let mut writer = FileWriter::try_new_with_options(
4598                Vec::new(),
4599                batch.schema_ref(),
4600                IpcWriteOptions::try_new(IPC_ALIGNMENT, false, MetadataVersion::V5).unwrap(),
4601            )
4602            .unwrap();
4603            writer.write(&batch).unwrap();
4604            writer.finish().unwrap();
4605
4606            let out: Vec<u8> = writer.into_inner().unwrap();
4607
4608            let buffer = Buffer::from_vec(out);
4609            let trailer_start = buffer.len() - 10;
4610            let footer_len =
4611                read_footer_length(buffer[trailer_start..].try_into().unwrap()).unwrap();
4612            let footer =
4613                root_as_footer(&buffer[trailer_start - footer_len..trailer_start]).unwrap();
4614
4615            let schema = fb_to_schema(footer.schema().unwrap());
4616
4617            // Importantly we set `require_alignment`, checking that 16-byte alignment is sufficient
4618            // for `read_record_batch` later on to read the data in a zero-copy manner.
4619            let decoder =
4620                FileDecoder::new(Arc::new(schema), footer.version()).with_require_alignment(true);
4621
4622            let batches = footer.recordBatches().unwrap();
4623
4624            let block = batches.get(0);
4625            let block_len = block.bodyLength() as usize + block.metaDataLength() as usize;
4626            let data = buffer.slice_with_length(block.offset() as _, block_len);
4627
4628            let batch2 = decoder.read_record_batch(block, &data).unwrap().unwrap();
4629
4630            assert_eq!(batch, batch2);
4631        }
4632    }
4633
4634    #[test]
4635    fn test_decimal128_alignment8_is_unaligned() {
4636        const IPC_ALIGNMENT: usize = 8;
4637
4638        let num_cols = 2;
4639        let num_rows = 1;
4640
4641        let mut fields = Vec::new();
4642        let mut arrays = Vec::new();
4643        for i in 0..num_cols {
4644            let field = Field::new(format!("col_{i}"), DataType::Decimal128(38, 10), true);
4645            let array = Decimal128Array::from(vec![num_cols as i128; num_rows]);
4646            fields.push(field);
4647            arrays.push(Arc::new(array) as Arc<dyn Array>);
4648        }
4649        let schema = Schema::new(fields);
4650        let batch = RecordBatch::try_new(Arc::new(schema), arrays).unwrap();
4651
4652        let mut writer = FileWriter::try_new_with_options(
4653            Vec::new(),
4654            batch.schema_ref(),
4655            IpcWriteOptions::try_new(IPC_ALIGNMENT, false, MetadataVersion::V5).unwrap(),
4656        )
4657        .unwrap();
4658        writer.write(&batch).unwrap();
4659        writer.finish().unwrap();
4660
4661        let out: Vec<u8> = writer.into_inner().unwrap();
4662
4663        let buffer = Buffer::from_vec(out);
4664        let trailer_start = buffer.len() - 10;
4665        let footer_len = read_footer_length(buffer[trailer_start..].try_into().unwrap()).unwrap();
4666        let footer = root_as_footer(&buffer[trailer_start - footer_len..trailer_start]).unwrap();
4667        let schema = fb_to_schema(footer.schema().unwrap());
4668
4669        // Importantly we set `require_alignment`, otherwise the error later is suppressed due to copying
4670        // to an aligned buffer in `ArrayDataBuilder.build_aligned`.
4671        let decoder =
4672            FileDecoder::new(Arc::new(schema), footer.version()).with_require_alignment(true);
4673
4674        let batches = footer.recordBatches().unwrap();
4675
4676        let block = batches.get(0);
4677        let block_len = block.bodyLength() as usize + block.metaDataLength() as usize;
4678        let data = buffer.slice_with_length(block.offset() as _, block_len);
4679
4680        let result = decoder.read_record_batch(block, &data);
4681
4682        let error = result.unwrap_err();
4683        assert_eq!(
4684            error.to_string(),
4685            "Invalid argument error: Misaligned buffers[0] in array of type Decimal128(38, 10), \
4686             offset from expected alignment of 16 by 8"
4687        );
4688    }
4689
4690    #[test]
4691    fn test_flush() {
4692        // We write a schema which is small enough to fit into a buffer and not get flushed,
4693        // and then force the write with .flush().
4694        let num_cols = 2;
4695        let mut fields = Vec::new();
4696        let options = IpcWriteOptions::try_new(8, false, MetadataVersion::V5).unwrap();
4697        for i in 0..num_cols {
4698            let field = Field::new(format!("col_{i}"), DataType::Decimal128(38, 10), true);
4699            fields.push(field);
4700        }
4701        let schema = Schema::new(fields);
4702        let inner_stream_writer = BufWriter::with_capacity(1024, Vec::new());
4703        let inner_file_writer = BufWriter::with_capacity(1024, Vec::new());
4704        let mut stream_writer =
4705            StreamWriter::try_new_with_options(inner_stream_writer, &schema, options.clone())
4706                .unwrap();
4707        let mut file_writer =
4708            FileWriter::try_new_with_options(inner_file_writer, &schema, options).unwrap();
4709
4710        let stream_bytes_written_on_new = stream_writer.get_ref().get_ref().len();
4711        let file_bytes_written_on_new = file_writer.get_ref().get_ref().len();
4712        stream_writer.flush().unwrap();
4713        file_writer.flush().unwrap();
4714        let stream_bytes_written_on_flush = stream_writer.get_ref().get_ref().len();
4715        let file_bytes_written_on_flush = file_writer.get_ref().get_ref().len();
4716        let stream_out = stream_writer.into_inner().unwrap().into_inner().unwrap();
4717        // Finishing a stream writes the continuation bytes in MetadataVersion::V5 (4 bytes)
4718        // and then a length of 0 (4 bytes) for a total of 8 bytes.
4719        // Everything before that should have been flushed in the .flush() call.
4720        let expected_stream_flushed_bytes = stream_out.len() - 8;
4721        // A file write is the same as the stream write except for the leading magic string
4722        // ARROW1 plus padding, which is 8 bytes.
4723        let expected_file_flushed_bytes = expected_stream_flushed_bytes + 8;
4724
4725        assert!(
4726            stream_bytes_written_on_new < stream_bytes_written_on_flush,
4727            "this test makes no sense if flush is not actually required"
4728        );
4729        assert!(
4730            file_bytes_written_on_new < file_bytes_written_on_flush,
4731            "this test makes no sense if flush is not actually required"
4732        );
4733        assert_eq!(stream_bytes_written_on_flush, expected_stream_flushed_bytes);
4734        assert_eq!(file_bytes_written_on_flush, expected_file_flushed_bytes);
4735    }
4736
4737    #[test]
4738    fn test_roundtrip_list_of_fixed_list() -> Result<(), ArrowError> {
4739        let l1_type =
4740            DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Float32, false)), 3);
4741        let l2_type = DataType::List(Arc::new(Field::new("item", l1_type.clone(), false)));
4742
4743        let l0_builder = Float32Builder::new();
4744        let l1_builder = FixedSizeListBuilder::new(l0_builder, 3).with_field(Arc::new(Field::new(
4745            "item",
4746            DataType::Float32,
4747            false,
4748        )));
4749        let mut l2_builder =
4750            ListBuilder::new(l1_builder).with_field(Arc::new(Field::new("item", l1_type, false)));
4751
4752        for point in [[1.0, 2.0, 3.0], [4.0, 5.0, 6.0], [7.0, 8.0, 9.0]] {
4753            l2_builder.values().values().append_value(point[0]);
4754            l2_builder.values().values().append_value(point[1]);
4755            l2_builder.values().values().append_value(point[2]);
4756
4757            l2_builder.values().append(true);
4758        }
4759        l2_builder.append(true);
4760
4761        let point = [10., 11., 12.];
4762        l2_builder.values().values().append_value(point[0]);
4763        l2_builder.values().values().append_value(point[1]);
4764        l2_builder.values().values().append_value(point[2]);
4765
4766        l2_builder.values().append(true);
4767        l2_builder.append(true);
4768
4769        let array = Arc::new(l2_builder.finish()) as ArrayRef;
4770
4771        let schema = Arc::new(Schema::new_with_metadata(
4772            vec![Field::new("points", l2_type, false)],
4773            HashMap::default(),
4774        ));
4775
4776        // Test a variety of combinations that include 0 and non-zero offsets
4777        // and also portions or the rest of the array
4778        test_slices(&array, &schema, 0, 1)?;
4779        test_slices(&array, &schema, 0, 2)?;
4780        test_slices(&array, &schema, 1, 1)?;
4781
4782        Ok(())
4783    }
4784
4785    #[test]
4786    fn test_roundtrip_list_of_fixed_list_w_nulls() -> Result<(), ArrowError> {
4787        let l0_builder = Float32Builder::new();
4788        let l1_builder = FixedSizeListBuilder::new(l0_builder, 3);
4789        let mut l2_builder = ListBuilder::new(l1_builder);
4790
4791        for point in [
4792            [Some(1.0), Some(2.0), None],
4793            [Some(4.0), Some(5.0), Some(6.0)],
4794            [None, Some(8.0), Some(9.0)],
4795        ] {
4796            for p in point {
4797                match p {
4798                    Some(p) => l2_builder.values().values().append_value(p),
4799                    None => l2_builder.values().values().append_null(),
4800                }
4801            }
4802
4803            l2_builder.values().append(true);
4804        }
4805        l2_builder.append(true);
4806
4807        let point = [Some(10.), None, None];
4808        for p in point {
4809            match p {
4810                Some(p) => l2_builder.values().values().append_value(p),
4811                None => l2_builder.values().values().append_null(),
4812            }
4813        }
4814
4815        l2_builder.values().append(true);
4816        l2_builder.append(true);
4817
4818        let array = Arc::new(l2_builder.finish()) as ArrayRef;
4819
4820        let schema = Arc::new(Schema::new_with_metadata(
4821            vec![Field::new(
4822                "points",
4823                DataType::List(Arc::new(Field::new(
4824                    "item",
4825                    DataType::FixedSizeList(
4826                        Arc::new(Field::new("item", DataType::Float32, true)),
4827                        3,
4828                    ),
4829                    true,
4830                ))),
4831                true,
4832            )],
4833            HashMap::default(),
4834        ));
4835
4836        // Test a variety of combinations that include 0 and non-zero offsets
4837        // and also portions or the rest of the array
4838        test_slices(&array, &schema, 0, 1)?;
4839        test_slices(&array, &schema, 0, 2)?;
4840        test_slices(&array, &schema, 1, 1)?;
4841
4842        Ok(())
4843    }
4844
4845    fn test_slices(
4846        parent_array: &ArrayRef,
4847        schema: &SchemaRef,
4848        offset: usize,
4849        length: usize,
4850    ) -> Result<(), ArrowError> {
4851        let subarray = parent_array.slice(offset, length);
4852        let original_batch = RecordBatch::try_new(schema.clone(), vec![subarray])?;
4853
4854        let mut bytes = Vec::new();
4855        let mut writer = StreamWriter::try_new(&mut bytes, schema)?;
4856        writer.write(&original_batch)?;
4857        writer.finish()?;
4858
4859        let mut cursor = std::io::Cursor::new(bytes);
4860        let mut reader = StreamReader::try_new(&mut cursor, None)?;
4861        let returned_batch = reader.next().unwrap()?;
4862
4863        assert_eq!(original_batch, returned_batch);
4864
4865        Ok(())
4866    }
4867
4868    #[test]
4869    fn test_roundtrip_fixed_list() -> Result<(), ArrowError> {
4870        let int_builder = Int64Builder::new();
4871        let mut fixed_list_builder = FixedSizeListBuilder::new(int_builder, 3)
4872            .with_field(Arc::new(Field::new("item", DataType::Int64, false)));
4873
4874        for point in [[1, 2, 3], [4, 5, 6], [7, 8, 9], [10, 11, 12]] {
4875            fixed_list_builder.values().append_value(point[0]);
4876            fixed_list_builder.values().append_value(point[1]);
4877            fixed_list_builder.values().append_value(point[2]);
4878
4879            fixed_list_builder.append(true);
4880        }
4881
4882        let array = Arc::new(fixed_list_builder.finish()) as ArrayRef;
4883
4884        let schema = Arc::new(Schema::new_with_metadata(
4885            vec![Field::new(
4886                "points",
4887                DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Int64, false)), 3),
4888                false,
4889            )],
4890            HashMap::default(),
4891        ));
4892
4893        // Test a variety of combinations that include 0 and non-zero offsets
4894        // and also portions or the rest of the array
4895        test_slices(&array, &schema, 0, 4)?;
4896        test_slices(&array, &schema, 0, 2)?;
4897        test_slices(&array, &schema, 1, 3)?;
4898        test_slices(&array, &schema, 2, 1)?;
4899
4900        Ok(())
4901    }
4902
4903    #[test]
4904    fn test_roundtrip_fixed_list_w_nulls() -> Result<(), ArrowError> {
4905        let int_builder = Int64Builder::new();
4906        let mut fixed_list_builder = FixedSizeListBuilder::new(int_builder, 3);
4907
4908        for point in [
4909            [Some(1), Some(2), None],
4910            [Some(4), Some(5), Some(6)],
4911            [None, Some(8), Some(9)],
4912            [Some(10), None, None],
4913        ] {
4914            for p in point {
4915                match p {
4916                    Some(p) => fixed_list_builder.values().append_value(p),
4917                    None => fixed_list_builder.values().append_null(),
4918                }
4919            }
4920
4921            fixed_list_builder.append(true);
4922        }
4923
4924        let array = Arc::new(fixed_list_builder.finish()) as ArrayRef;
4925
4926        let schema = Arc::new(Schema::new_with_metadata(
4927            vec![Field::new(
4928                "points",
4929                DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Int64, true)), 3),
4930                true,
4931            )],
4932            HashMap::default(),
4933        ));
4934
4935        // Test a variety of combinations that include 0 and non-zero offsets
4936        // and also portions or the rest of the array
4937        test_slices(&array, &schema, 0, 4)?;
4938        test_slices(&array, &schema, 0, 2)?;
4939        test_slices(&array, &schema, 1, 3)?;
4940        test_slices(&array, &schema, 2, 1)?;
4941
4942        Ok(())
4943    }
4944
4945    #[test]
4946    fn test_metadata_encoding_ordering() {
4947        fn create_hash() -> u64 {
4948            let metadata: HashMap<String, String> = [
4949                ("a", "1"), //
4950                ("b", "2"), //
4951                ("c", "3"), //
4952                ("d", "4"), //
4953                ("e", "5"), //
4954            ]
4955            .into_iter()
4956            .map(|(k, v)| (k.to_owned(), v.to_owned()))
4957            .collect();
4958
4959            // Set metadata on both the schema and a field within it.
4960            let schema = Arc::new(
4961                Schema::new(vec![
4962                    Field::new("a", DataType::Int64, true).with_metadata(metadata.clone()),
4963                ])
4964                .with_metadata(metadata)
4965                .clone(),
4966            );
4967            let batch = RecordBatch::new_empty(schema.clone());
4968
4969            let mut bytes = Vec::new();
4970            let mut w = StreamWriter::try_new(&mut bytes, batch.schema_ref()).unwrap();
4971            w.write(&batch).unwrap();
4972            w.finish().unwrap();
4973
4974            let mut h = std::hash::DefaultHasher::new();
4975            h.write(&bytes);
4976            h.finish()
4977        }
4978
4979        let expected = create_hash();
4980
4981        // Since there is randomness in the HashMap and we cannot specify our
4982        // own Hasher for the implementation used for metadata, run the above
4983        // code 20x and verify it does not change. This is not perfect but it
4984        // should be good enough.
4985        let all_passed = (0..20).all(|_| create_hash() == expected);
4986        assert!(all_passed);
4987    }
4988
4989    #[test]
4990    fn test_dictionary_tracker_reset() {
4991        let data_gen = IpcDataGenerator::default();
4992        let mut dictionary_tracker = DictionaryTracker::new(false);
4993        let writer_options = IpcWriteOptions::default();
4994        let mut compression_ctx = IpcWriteContext::default();
4995
4996        let schema = Arc::new(Schema::new(vec![Field::new(
4997            "a",
4998            DataType::Dictionary(Box::new(DataType::UInt8), Box::new(DataType::Utf8)),
4999            false,
5000        )]));
5001
5002        let mut write_single_batch_stream =
5003            |batch: RecordBatch, dict_tracker: &mut DictionaryTracker| -> Vec<u8> {
5004                let mut buffer = Vec::new();
5005
5006                // create a new IPC stream:
5007                let stream_header = data_gen.schema_to_bytes_with_dictionary_tracker(
5008                    &schema,
5009                    dict_tracker,
5010                    &writer_options,
5011                );
5012                _ = write_message(&mut buffer, stream_header, &writer_options).unwrap();
5013
5014                let (encoded_dicts, encoded_batch) = data_gen
5015                    .encode(&batch, dict_tracker, &writer_options, &mut compression_ctx)
5016                    .unwrap();
5017                for encoded_dict in encoded_dicts {
5018                    _ = write_message(&mut buffer, encoded_dict, &writer_options).unwrap();
5019                }
5020                _ = write_message(&mut buffer, encoded_batch, &writer_options).unwrap();
5021
5022                buffer
5023            };
5024
5025        let batch1 = RecordBatch::try_new(
5026            schema.clone(),
5027            vec![Arc::new(DictionaryArray::new(
5028                UInt8Array::from_iter_values([0]),
5029                Arc::new(StringArray::from_iter_values(["a"])),
5030            ))],
5031        )
5032        .unwrap();
5033        let buffer = write_single_batch_stream(batch1.clone(), &mut dictionary_tracker);
5034
5035        // ensure we can read the stream back
5036        let mut reader = StreamReader::try_new(Cursor::new(buffer), None).unwrap();
5037        let read_batch = reader.next().unwrap().unwrap();
5038        assert_eq!(read_batch, batch1);
5039
5040        // reset the dictionary tracker so it can be used for next stream
5041        dictionary_tracker.clear();
5042
5043        // now write a 2nd stream and ensure we can also read it:
5044        let batch2 = RecordBatch::try_new(
5045            schema.clone(),
5046            vec![Arc::new(DictionaryArray::new(
5047                UInt8Array::from_iter_values([0]),
5048                Arc::new(StringArray::from_iter_values(["a"])),
5049            ))],
5050        )
5051        .unwrap();
5052        let buffer = write_single_batch_stream(batch2.clone(), &mut dictionary_tracker);
5053        let mut reader = StreamReader::try_new(Cursor::new(buffer), None).unwrap();
5054        let read_batch = reader.next().unwrap().unwrap();
5055        assert_eq!(read_batch, batch2);
5056    }
5057}