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