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