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 if run_array.is_empty() {
1284 return Ok(ArrayData::new_empty(run_array.data_type()).into());
1285 }
1286
1287 let run_ends = run_array.run_ends();
1288 if run_ends.offset() == 0 && run_ends.max_value() == run_ends.len() {
1289 return Ok(run_array);
1290 }
1291
1292 let start_physical_index = run_ends.get_start_physical_index();
1294
1295 let end_physical_index = run_ends.get_end_physical_index();
1297
1298 let physical_length = end_physical_index - start_physical_index + 1;
1299
1300 let offset = R::Native::usize_as(run_ends.offset());
1302 let mut run_ends_values = Vec::<R::Native>::with_capacity(physical_length);
1303 for run_end_value in &run_ends.values()[start_physical_index..end_physical_index] {
1304 run_ends_values.push(run_end_value.sub_wrapping(offset));
1305 }
1306 run_ends_values.push(R::Native::from_usize(run_array.len()).unwrap());
1307 let offset_buffer = Buffer::from_vec(run_ends_values);
1308 let new_run_ends = unsafe {
1309 ArrayDataBuilder::new(R::DATA_TYPE)
1312 .len(physical_length)
1313 .add_buffer(offset_buffer)
1314 .build_unchecked()
1315 };
1316
1317 let new_values = run_array
1319 .values()
1320 .slice(start_physical_index, physical_length)
1321 .into_data();
1322
1323 let builder = ArrayDataBuilder::new(run_array.data_type().clone())
1324 .len(run_array.len())
1325 .add_child_data(new_run_ends)
1326 .add_child_data(new_values);
1327 let array_data = unsafe {
1328 builder.build_unchecked()
1331 };
1332 Ok(array_data.into())
1333}
1334
1335#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
1337pub enum DictionaryHandling {
1338 #[default]
1340 Resend,
1341 Delta,
1347}
1348
1349#[derive(Debug, Clone)]
1351pub enum DictionaryUpdate {
1352 None,
1355 New,
1357 Replaced,
1359 Delta(ArrayData),
1361}
1362
1363#[derive(Debug)]
1369pub struct DictionaryTracker {
1370 written: HashMap<i64, ArrayData>,
1372 dict_ids: Vec<i64>,
1373 error_on_replacement: bool,
1374}
1375
1376impl DictionaryTracker {
1377 pub fn new(error_on_replacement: bool) -> Self {
1383 Self {
1384 written: HashMap::new(),
1385 dict_ids: Vec::new(),
1386 error_on_replacement,
1387 }
1388 }
1389
1390 pub fn next_dict_id(&mut self) -> i64 {
1392 let next = self
1393 .dict_ids
1394 .last()
1395 .copied()
1396 .map(|i| i + 1)
1397 .unwrap_or_default();
1398
1399 self.dict_ids.push(next);
1400 next
1401 }
1402
1403 pub fn dict_id(&mut self) -> &[i64] {
1406 &self.dict_ids
1407 }
1408
1409 #[deprecated(since = "56.1.0", note = "Use `insert_column` instead")]
1419 pub fn insert(&mut self, dict_id: i64, column: &ArrayRef) -> Result<bool, ArrowError> {
1420 let dict_data = column.to_data();
1421 let dict_values = &dict_data.child_data()[0];
1422
1423 if let Some(last) = self.written.get(&dict_id) {
1425 if ArrayData::ptr_eq(&last.child_data()[0], dict_values) {
1426 return Ok(false);
1428 }
1429 if self.error_on_replacement {
1430 if last.child_data()[0] == *dict_values {
1432 return Ok(false);
1434 }
1435 return Err(ArrowError::InvalidArgumentError(
1436 "Dictionary replacement detected when writing IPC file format. \
1437 Arrow IPC files only support a single dictionary for a given field \
1438 across all batches."
1439 .to_string(),
1440 ));
1441 }
1442 }
1443
1444 self.written.insert(dict_id, dict_data);
1445 Ok(true)
1446 }
1447
1448 pub fn insert_column(
1464 &mut self,
1465 dict_id: i64,
1466 column: &ArrayRef,
1467 dict_handling: DictionaryHandling,
1468 ) -> Result<DictionaryUpdate, ArrowError> {
1469 let new_data = column.to_data();
1470 let new_values = &new_data.child_data()[0];
1471
1472 let Some(old) = self.written.get(&dict_id) else {
1474 self.written.insert(dict_id, new_data);
1475 return Ok(DictionaryUpdate::New);
1476 };
1477
1478 let old_values = &old.child_data()[0];
1481 if ArrayData::ptr_eq(old_values, new_values) {
1482 return Ok(DictionaryUpdate::None);
1483 }
1484
1485 let comparison = compare_dictionaries(old_values, new_values);
1487 if matches!(comparison, DictionaryComparison::Equal) {
1488 return Ok(DictionaryUpdate::None);
1489 }
1490
1491 const REPLACEMENT_ERROR: &str = "Dictionary replacement detected when writing IPC file format. \
1492 Arrow IPC files only support a single dictionary for a given field \
1493 across all batches.";
1494
1495 match comparison {
1496 DictionaryComparison::NotEqual => {
1497 if self.error_on_replacement {
1498 return Err(ArrowError::InvalidArgumentError(
1499 REPLACEMENT_ERROR.to_string(),
1500 ));
1501 }
1502
1503 self.written.insert(dict_id, new_data);
1504 Ok(DictionaryUpdate::Replaced)
1505 }
1506 DictionaryComparison::Delta => match dict_handling {
1507 DictionaryHandling::Resend => {
1508 if self.error_on_replacement {
1509 return Err(ArrowError::InvalidArgumentError(
1510 REPLACEMENT_ERROR.to_string(),
1511 ));
1512 }
1513
1514 self.written.insert(dict_id, new_data);
1515 Ok(DictionaryUpdate::Replaced)
1516 }
1517 DictionaryHandling::Delta => {
1518 let delta =
1519 new_values.slice(old_values.len(), new_values.len() - old_values.len());
1520 self.written.insert(dict_id, new_data);
1521 Ok(DictionaryUpdate::Delta(delta))
1522 }
1523 },
1524 DictionaryComparison::Equal => unreachable!("Already checked equal case"),
1525 }
1526 }
1527
1528 pub fn clear(&mut self) {
1534 self.dict_ids.clear();
1535 self.written.clear();
1536 }
1537}
1538
1539#[derive(Debug, Clone)]
1541enum DictionaryComparison {
1542 NotEqual,
1544 Equal,
1546 Delta,
1549}
1550
1551fn compare_dictionaries(old: &ArrayData, new: &ArrayData) -> DictionaryComparison {
1553 let existing_len = old.len();
1555 let new_len = new.len();
1556 if existing_len == new_len {
1557 return if *old == *new {
1558 DictionaryComparison::Equal
1559 } else {
1560 DictionaryComparison::NotEqual
1561 };
1562 }
1563
1564 if new_len < existing_len {
1566 return DictionaryComparison::NotEqual;
1567 }
1568
1569 if new.slice(0, existing_len) == *old {
1571 return DictionaryComparison::Delta;
1572 }
1573
1574 DictionaryComparison::NotEqual
1575}
1576
1577pub struct FileWriter<W> {
1600 writer: W,
1602 write_options: IpcWriteOptions,
1604 schema: SchemaRef,
1606 block_offsets: usize,
1608 dictionary_blocks: Vec<crate::Block>,
1610 record_blocks: Vec<crate::Block>,
1612 finished: bool,
1614 dictionary_tracker: DictionaryTracker,
1616 custom_metadata: Metadata,
1618
1619 data_gen: IpcDataGenerator,
1620
1621 ipc_write_context: IpcWriteContext,
1622}
1623
1624impl<W: Write> FileWriter<BufWriter<W>> {
1625 pub fn try_new_buffered(writer: W, schema: &Schema) -> Result<Self, ArrowError> {
1629 Self::try_new(BufWriter::new(writer), schema)
1630 }
1631}
1632
1633impl<W: Write> FileWriter<W> {
1634 pub fn try_new(writer: W, schema: &Schema) -> Result<Self, ArrowError> {
1642 let write_options = IpcWriteOptions::default();
1643 Self::try_new_with_options(writer, schema, write_options)
1644 }
1645
1646 pub fn try_new_with_options(
1654 mut writer: W,
1655 schema: &Schema,
1656 write_options: IpcWriteOptions,
1657 ) -> Result<Self, ArrowError> {
1658 ensure_supported_ipc_schema(schema)?;
1659
1660 let data_gen = IpcDataGenerator::default();
1661 let pad_len = pad_to_alignment(write_options.alignment, super::ARROW_MAGIC.len());
1663 let header_size = super::ARROW_MAGIC.len() + pad_len;
1664 writer.write_all(&super::ARROW_MAGIC)?;
1665 writer.write_all(&PADDING[..pad_len])?;
1666 let mut dictionary_tracker = DictionaryTracker::new(true);
1668 let encoded_message = data_gen.schema_to_bytes_with_dictionary_tracker(
1669 schema,
1670 &mut dictionary_tracker,
1671 &write_options,
1672 );
1673 let (meta, data) = write_message(&mut writer, encoded_message, &write_options)?;
1674 Ok(Self {
1675 writer,
1676 write_options,
1677 schema: Arc::new(schema.clone()),
1678 block_offsets: meta + data + header_size,
1679 dictionary_blocks: vec![],
1680 record_blocks: vec![],
1681 finished: false,
1682 dictionary_tracker,
1683 custom_metadata: Default::default(),
1684 data_gen,
1685 ipc_write_context: IpcWriteContext::default(),
1686 })
1687 }
1688
1689 pub fn write_metadata(&mut self, key: impl Into<String>, value: impl Into<String>) {
1691 self.custom_metadata.insert(key.into(), value.into());
1692 }
1693
1694 pub fn write(&mut self, batch: &RecordBatch) -> Result<(), ArrowError> {
1696 if self.finished {
1697 return Err(ArrowError::IpcError(
1698 "Cannot write record batch to file writer as it is closed".to_string(),
1699 ));
1700 }
1701
1702 let meta = self.data_gen.write(
1703 batch,
1704 &mut self.dictionary_tracker,
1705 &self.write_options,
1706 &mut self.ipc_write_context,
1707 &mut self.writer,
1708 )?;
1709
1710 for (header_len, body_len) in meta.dictionary_block_sizes {
1711 let block = crate::Block::new(
1712 self.block_offsets as i64,
1713 header_len as i32,
1714 body_len as i64,
1715 );
1716 self.dictionary_blocks.push(block);
1717 self.block_offsets += header_len + body_len;
1718 }
1719
1720 let block = crate::Block::new(
1722 self.block_offsets as i64,
1723 meta.padded_header_len as i32,
1724 meta.body_len as i64,
1725 );
1726 self.record_blocks.push(block);
1727 self.block_offsets += meta.padded_header_len + meta.body_len;
1728 Ok(())
1729 }
1730
1731 pub fn finish(&mut self) -> Result<(), ArrowError> {
1733 if self.finished {
1734 return Err(ArrowError::IpcError(
1735 "Cannot write footer to file writer as it is closed".to_string(),
1736 ));
1737 }
1738
1739 {
1741 self.writer.write_eos(&self.write_options)?;
1742 }
1743
1744 let mut fbb = FlatBufferBuilder::new();
1745 let dictionaries = fbb.create_vector(&self.dictionary_blocks);
1746 let record_batches = fbb.create_vector(&self.record_blocks);
1747
1748 self.dictionary_tracker.clear();
1750 let schema = IpcSchemaEncoder::new()
1751 .with_dictionary_tracker(&mut self.dictionary_tracker)
1752 .schema_to_fb_offset(&mut fbb, &self.schema);
1753 let fb_custom_metadata = (!self.custom_metadata.is_empty())
1754 .then(|| crate::convert::metadata_to_fb(&mut fbb, &self.custom_metadata));
1755
1756 let root = {
1757 let mut footer_builder = crate::FooterBuilder::new(&mut fbb);
1758 footer_builder.add_version(self.write_options.metadata_version);
1759 footer_builder.add_schema(schema);
1760 footer_builder.add_dictionaries(dictionaries);
1761 footer_builder.add_recordBatches(record_batches);
1762 if let Some(fb_custom_metadata) = fb_custom_metadata {
1763 footer_builder.add_custom_metadata(fb_custom_metadata);
1764 }
1765 footer_builder.finish()
1766 };
1767 fbb.finish(root, None);
1768 let footer_data = fbb.finished_data();
1769 self.writer.write_all(footer_data)?;
1770 self.writer
1771 .write_all(&(footer_data.len() as i32).to_le_bytes())?;
1772 self.writer.write_all(&super::ARROW_MAGIC)?;
1773 self.writer.flush()?;
1774 self.finished = true;
1775
1776 Ok(())
1777 }
1778
1779 pub fn schema(&self) -> &SchemaRef {
1781 &self.schema
1782 }
1783
1784 pub fn get_ref(&self) -> &W {
1786 &self.writer
1787 }
1788
1789 pub fn get_mut(&mut self) -> &mut W {
1793 &mut self.writer
1794 }
1795
1796 pub fn flush(&mut self) -> Result<(), ArrowError> {
1800 self.writer.flush()?;
1801 Ok(())
1802 }
1803
1804 pub fn into_inner(mut self) -> Result<W, ArrowError> {
1813 if !self.finished {
1814 self.finish()?;
1816 }
1817 Ok(self.writer)
1818 }
1819}
1820
1821impl<W: Write> RecordBatchWriter for FileWriter<W> {
1822 fn write(&mut self, batch: &RecordBatch) -> Result<(), ArrowError> {
1823 self.write(batch)
1824 }
1825
1826 fn close(mut self) -> Result<(), ArrowError> {
1827 self.finish()
1828 }
1829}
1830
1831pub struct StreamEncoder {
1860 schema: Schema,
1861 write_options: IpcWriteOptions,
1863 schema_encoded: bool,
1865 dictionary_tracker: DictionaryTracker,
1867 data_gen: IpcDataGenerator,
1868 ipc_write_context: IpcWriteContext,
1869}
1870
1871impl StreamEncoder {
1872 pub fn try_new(schema: &Schema) -> Result<Self, ArrowError> {
1874 let write_options = IpcWriteOptions::default();
1875 Self::try_new_with_options(schema, write_options)
1876 }
1877
1878 pub fn try_new_with_options(
1880 schema: &Schema,
1881 write_options: IpcWriteOptions,
1882 ) -> Result<Self, ArrowError> {
1883 ensure_supported_ipc_schema(schema)?;
1884
1885 Ok(Self {
1886 schema: schema.clone(),
1887 write_options,
1888 schema_encoded: false,
1889 dictionary_tracker: DictionaryTracker::new(false),
1890 data_gen: IpcDataGenerator::default(),
1891 ipc_write_context: IpcWriteContext::default(),
1892 })
1893 }
1894
1895 pub fn encode(&mut self, batch: &RecordBatch) -> Result<Vec<Buffer>, ArrowError> {
1905 let mut out = vec![];
1906 self.encode_schema(&mut out)?;
1907 self.data_gen.encode_to_buffers(
1908 batch,
1909 &mut self.dictionary_tracker,
1910 &self.write_options,
1911 &mut self.ipc_write_context,
1912 &mut out,
1913 )?;
1914 Ok(out)
1915 }
1916
1917 pub fn finish(mut self) -> Result<Vec<Buffer>, ArrowError> {
1926 let mut out = vec![];
1927 self.encode_schema(&mut out)?;
1928 let mut sink = Buffers { out: &mut out };
1929 sink.write_eos(&self.write_options)?;
1930 Ok(out)
1931 }
1932
1933 fn encode_schema(&mut self, out: &mut Vec<Buffer>) -> Result<(), ArrowError> {
1934 if !self.schema_encoded {
1935 let encoded_message = self.data_gen.schema_to_bytes_with_dictionary_tracker(
1936 &self.schema,
1937 &mut self.dictionary_tracker,
1938 &self.write_options,
1939 );
1940 let mut sink = Buffers { out };
1941 sink.write_encoded_data(encoded_message, &self.write_options)?;
1942 self.schema_encoded = true;
1943 }
1944 Ok(())
1945 }
1946}
1947
1948pub struct StreamWriter<W> {
2022 writer: W,
2024 write_options: IpcWriteOptions,
2026 finished: bool,
2028 dictionary_tracker: DictionaryTracker,
2030
2031 data_gen: IpcDataGenerator,
2032
2033 ipc_write_context: IpcWriteContext,
2034}
2035
2036impl<W: Write> StreamWriter<BufWriter<W>> {
2037 pub fn try_new_buffered(writer: W, schema: &Schema) -> Result<Self, ArrowError> {
2041 Self::try_new(BufWriter::new(writer), schema)
2042 }
2043}
2044
2045impl<W: Write> StreamWriter<W> {
2046 pub fn try_new(writer: W, schema: &Schema) -> Result<Self, ArrowError> {
2054 let write_options = IpcWriteOptions::default();
2055 Self::try_new_with_options(writer, schema, write_options)
2056 }
2057
2058 pub fn try_new_with_options(
2064 mut writer: W,
2065 schema: &Schema,
2066 write_options: IpcWriteOptions,
2067 ) -> Result<Self, ArrowError> {
2068 ensure_supported_ipc_schema(schema)?;
2069
2070 let data_gen = IpcDataGenerator::default();
2071 let mut dictionary_tracker = DictionaryTracker::new(false);
2072
2073 let encoded_message = data_gen.schema_to_bytes_with_dictionary_tracker(
2075 schema,
2076 &mut dictionary_tracker,
2077 &write_options,
2078 );
2079 write_message(&mut writer, encoded_message, &write_options)?;
2080 Ok(Self {
2081 writer,
2082 write_options,
2083 finished: false,
2084 dictionary_tracker,
2085 data_gen,
2086 ipc_write_context: IpcWriteContext::default(),
2087 })
2088 }
2089
2090 pub fn write(&mut self, batch: &RecordBatch) -> Result<(), ArrowError> {
2092 if self.finished {
2093 return Err(ArrowError::IpcError(
2094 "Cannot write record batch to stream writer as it is closed".to_string(),
2095 ));
2096 }
2097
2098 self.data_gen.write(
2099 batch,
2100 &mut self.dictionary_tracker,
2101 &self.write_options,
2102 &mut self.ipc_write_context,
2103 &mut self.writer,
2104 )?;
2105 Ok(())
2106 }
2107
2108 pub fn finish(&mut self) -> Result<(), ArrowError> {
2110 if self.finished {
2111 return Err(ArrowError::IpcError(
2112 "Cannot write footer to stream writer as it is closed".to_string(),
2113 ));
2114 }
2115
2116 {
2117 self.writer.write_eos(&self.write_options)?;
2118 }
2119 self.writer.flush()?;
2120
2121 self.finished = true;
2122
2123 Ok(())
2124 }
2125
2126 pub fn get_ref(&self) -> &W {
2128 &self.writer
2129 }
2130
2131 pub fn get_mut(&mut self) -> &mut W {
2135 &mut self.writer
2136 }
2137
2138 pub fn flush(&mut self) -> Result<(), ArrowError> {
2142 self.writer.flush()?;
2143 Ok(())
2144 }
2145
2146 pub fn into_inner(mut self) -> Result<W, ArrowError> {
2184 if !self.finished {
2185 self.finish()?;
2187 }
2188 Ok(self.writer)
2189 }
2190}
2191
2192impl<W: Write> RecordBatchWriter for StreamWriter<W> {
2193 fn write(&mut self, batch: &RecordBatch) -> Result<(), ArrowError> {
2194 self.write(batch)
2195 }
2196
2197 fn close(mut self) -> Result<(), ArrowError> {
2198 self.finish()
2199 }
2200}
2201
2202pub struct EncodedData {
2204 pub ipc_message: Vec<u8>,
2206 pub arrow_data: Vec<u8>,
2208}
2209
2210pub fn write_message<W: Write>(
2212 mut writer: W,
2213 encoded: EncodedData,
2214 write_options: &IpcWriteOptions,
2215) -> Result<(usize, usize), ArrowError> {
2216 writer.write_encoded_data(encoded, write_options)
2217}
2218
2219fn has_validity_bitmap(data_type: &DataType, write_options: &IpcWriteOptions) -> bool {
2223 if write_options.metadata_version < crate::MetadataVersion::V5 {
2224 !matches!(data_type, DataType::Null)
2225 } else {
2226 !matches!(
2227 data_type,
2228 DataType::Null | DataType::Union(_, _) | DataType::RunEndEncoded(_, _)
2229 )
2230 }
2231}
2232
2233#[inline]
2235fn buffer_need_truncate(
2236 array_offset: usize,
2237 buffer: &Buffer,
2238 spec: &BufferSpec,
2239 min_length: usize,
2240) -> bool {
2241 spec != &BufferSpec::AlwaysNull && (array_offset != 0 || min_length < buffer.len())
2242}
2243
2244#[inline]
2246fn get_buffer_element_width(spec: &BufferSpec) -> usize {
2247 match spec {
2248 BufferSpec::FixedWidth { byte_width, .. } => *byte_width,
2249 _ => 0,
2250 }
2251}
2252
2253fn reencode_offsets<O: OffsetSizeTrait>(
2256 offsets: &Buffer,
2257 data: &ArrayData,
2258) -> (Buffer, usize, usize) {
2259 let offsets_slice: &[O] = offsets.typed_data::<O>();
2260 let offset_slice = &offsets_slice[data.offset()..data.offset() + data.len() + 1];
2261
2262 let start_offset = offset_slice.first().unwrap();
2263 let end_offset = offset_slice.last().unwrap();
2264
2265 let offsets = match start_offset.as_usize() {
2266 0 => {
2267 let size = size_of::<O>();
2268 offsets.slice_with_length(data.offset() * size, (data.len() + 1) * size)
2269 }
2270 _ => offset_slice.iter().map(|x| *x - *start_offset).collect(),
2271 };
2272
2273 let start_offset = start_offset.as_usize();
2274 let end_offset = end_offset.as_usize();
2275
2276 (offsets, start_offset, end_offset - start_offset)
2277}
2278
2279fn get_byte_array_buffers<O: OffsetSizeTrait>(data: &ArrayData) -> [Buffer; 2] {
2286 if data.is_empty() {
2287 let mut offsets = MutableBuffer::new(size_of::<O>());
2290 offsets.extend_from_slice(O::usize_as(0).to_byte_slice());
2291 return [offsets.into(), MutableBuffer::new(0).into()];
2292 }
2293
2294 let (offsets, original_start_offset, len) = reencode_offsets::<O>(&data.buffers()[0], data);
2295 let values = data.buffers()[1].slice_with_length(original_start_offset, len);
2296 [offsets, values]
2297}
2298
2299fn get_list_array_buffers<O: OffsetSizeTrait>(data: &ArrayData) -> (Buffer, ArrayData) {
2302 if data.is_empty() {
2303 let mut offsets = MutableBuffer::new(size_of::<O>());
2306 offsets.extend_from_slice(O::usize_as(0).to_byte_slice());
2307 return (offsets.into(), data.child_data()[0].slice(0, 0));
2308 }
2309
2310 let (offsets, original_start_offset, len) = reencode_offsets::<O>(&data.buffers()[0], data);
2311 let child_data = data.child_data()[0].slice(original_start_offset, len);
2312 (offsets, child_data)
2313}
2314
2315fn get_list_view_array_buffers<O: OffsetSizeTrait>(
2321 data: &ArrayData,
2322) -> (Buffer, Buffer, ArrayData) {
2323 if data.is_empty() {
2324 return (
2325 MutableBuffer::new(0).into(),
2326 MutableBuffer::new(0).into(),
2327 data.child_data()[0].slice(0, 0),
2328 );
2329 }
2330
2331 let offsets = &data.buffers()[0];
2332 let sizes = &data.buffers()[1];
2333
2334 let element_size = std::mem::size_of::<O>();
2335 let offsets_slice =
2336 offsets.slice_with_length(data.offset() * element_size, data.len() * element_size);
2337 let sizes_slice =
2338 sizes.slice_with_length(data.offset() * element_size, data.len() * element_size);
2339
2340 let child_data = data.child_data()[0].clone();
2341
2342 (offsets_slice, sizes_slice, child_data)
2343}
2344
2345fn get_or_truncate_buffer(array_data: &ArrayData) -> Buffer {
2352 let buffer = &array_data.buffers()[0];
2353 let layout = layout(array_data.data_type());
2354 let spec = &layout.buffers[0];
2355
2356 let byte_width = get_buffer_element_width(spec);
2357 let min_length = array_data.len() * byte_width;
2358 if buffer_need_truncate(array_data.offset(), buffer, spec, min_length) {
2359 let byte_offset = array_data.offset() * byte_width;
2360 let buffer_length = min(min_length, buffer.len() - byte_offset);
2361 buffer.slice_with_length(byte_offset, buffer_length)
2362 } else {
2363 buffer.clone()
2364 }
2365}
2366
2367fn write_array_data(
2373 array_data: &ArrayData,
2374 meta: &mut IpcMetadataBuilder,
2375 sink: &mut IpcBodySink<'_>,
2376 offset: i64,
2377 compression_codec: Option<CompressionCodec>,
2378 ipc_write_context: &mut IpcWriteContext,
2379 write_options: &IpcWriteOptions,
2380) -> Result<i64, ArrowError> {
2381 let mut offset = offset;
2382 let num_rows = array_data.len();
2383 if !matches!(array_data.data_type(), DataType::Null) {
2384 meta.nodes.push(crate::FieldNode::new(
2385 num_rows as i64,
2386 array_data.null_count() as i64,
2387 ));
2388 } else {
2389 meta.nodes
2391 .push(crate::FieldNode::new(num_rows as i64, num_rows as i64));
2392 }
2393 if has_validity_bitmap(array_data.data_type(), write_options) {
2394 let null_buffer = match array_data.nulls() {
2396 None => {
2397 let num_bytes = bit_util::ceil(num_rows, 8);
2399 let buffer = MutableBuffer::new(num_bytes);
2400 let buffer = buffer.with_bitset(num_bytes, true);
2401 buffer.into()
2402 }
2403 Some(buffer) => buffer.inner().sliced(),
2404 };
2405
2406 offset = encode_sink_buffer(
2407 null_buffer,
2408 meta,
2409 sink,
2410 offset,
2411 compression_codec,
2412 ipc_write_context,
2413 write_options.alignment,
2414 )?;
2415 }
2416
2417 let data_type = array_data.data_type();
2418 if matches!(data_type, DataType::Binary | DataType::Utf8) {
2419 for buffer in get_byte_array_buffers::<i32>(array_data) {
2420 offset = encode_sink_buffer(
2421 buffer,
2422 meta,
2423 sink,
2424 offset,
2425 compression_codec,
2426 ipc_write_context,
2427 write_options.alignment,
2428 )?;
2429 }
2430 } else if matches!(data_type, DataType::BinaryView | DataType::Utf8View) {
2431 let views = get_or_truncate_buffer(array_data);
2438 offset = encode_sink_buffer(
2439 views,
2440 meta,
2441 sink,
2442 offset,
2443 compression_codec,
2444 ipc_write_context,
2445 write_options.alignment,
2446 )?;
2447
2448 for buffer in array_data.buffers().iter().skip(1) {
2449 offset = encode_sink_buffer(
2450 buffer.clone(),
2451 meta,
2452 sink,
2453 offset,
2454 compression_codec,
2455 ipc_write_context,
2456 write_options.alignment,
2457 )?;
2458 }
2459 } else if matches!(data_type, DataType::LargeBinary | DataType::LargeUtf8) {
2460 for buffer in get_byte_array_buffers::<i64>(array_data) {
2461 offset = encode_sink_buffer(
2462 buffer,
2463 meta,
2464 sink,
2465 offset,
2466 compression_codec,
2467 ipc_write_context,
2468 write_options.alignment,
2469 )?;
2470 }
2471 } else if DataType::is_numeric(data_type)
2472 || DataType::is_temporal(data_type)
2473 || matches!(
2474 array_data.data_type(),
2475 DataType::FixedSizeBinary(_) | DataType::Dictionary(_, _)
2476 )
2477 {
2478 assert_eq!(array_data.buffers().len(), 1);
2480
2481 let buffer = get_or_truncate_buffer(array_data);
2482 offset = encode_sink_buffer(
2483 buffer,
2484 meta,
2485 sink,
2486 offset,
2487 compression_codec,
2488 ipc_write_context,
2489 write_options.alignment,
2490 )?;
2491 } else if matches!(data_type, DataType::Boolean) {
2492 assert_eq!(array_data.buffers().len(), 1);
2495
2496 let buffer = &array_data.buffers()[0];
2497 let buffer = buffer.bit_slice(array_data.offset(), array_data.len());
2498 offset = encode_sink_buffer(
2499 buffer,
2500 meta,
2501 sink,
2502 offset,
2503 compression_codec,
2504 ipc_write_context,
2505 write_options.alignment,
2506 )?;
2507 } else if matches!(
2508 data_type,
2509 DataType::List(_) | DataType::LargeList(_) | DataType::Map(_, _)
2510 ) {
2511 assert_eq!(array_data.buffers().len(), 1);
2512 assert_eq!(array_data.child_data().len(), 1);
2513
2514 let (offsets, sliced_child_data) = match data_type {
2516 DataType::List(_) => get_list_array_buffers::<i32>(array_data),
2517 DataType::Map(_, _) => get_list_array_buffers::<i32>(array_data),
2518 DataType::LargeList(_) => get_list_array_buffers::<i64>(array_data),
2519 _ => unreachable!(),
2520 };
2521 offset = encode_sink_buffer(
2522 offsets,
2523 meta,
2524 sink,
2525 offset,
2526 compression_codec,
2527 ipc_write_context,
2528 write_options.alignment,
2529 )?;
2530 offset = write_array_data(
2531 &sliced_child_data,
2532 meta,
2533 sink,
2534 offset,
2535 compression_codec,
2536 ipc_write_context,
2537 write_options,
2538 )?;
2539 return Ok(offset);
2540 } else if matches!(
2541 data_type,
2542 DataType::ListView(_) | DataType::LargeListView(_)
2543 ) {
2544 assert_eq!(array_data.buffers().len(), 2); assert_eq!(array_data.child_data().len(), 1);
2546
2547 let (offsets, sizes, child_data) = match data_type {
2548 DataType::ListView(_) => get_list_view_array_buffers::<i32>(array_data),
2549 DataType::LargeListView(_) => get_list_view_array_buffers::<i64>(array_data),
2550 _ => unreachable!(),
2551 };
2552
2553 offset = encode_sink_buffer(
2554 offsets,
2555 meta,
2556 sink,
2557 offset,
2558 compression_codec,
2559 ipc_write_context,
2560 write_options.alignment,
2561 )?;
2562 offset = encode_sink_buffer(
2563 sizes,
2564 meta,
2565 sink,
2566 offset,
2567 compression_codec,
2568 ipc_write_context,
2569 write_options.alignment,
2570 )?;
2571
2572 offset = write_array_data(
2573 &child_data,
2574 meta,
2575 sink,
2576 offset,
2577 compression_codec,
2578 ipc_write_context,
2579 write_options,
2580 )?;
2581 return Ok(offset);
2582 } else if let DataType::FixedSizeList(_, fixed_size) = data_type {
2583 assert_eq!(array_data.child_data().len(), 1);
2584 let fixed_size = *fixed_size as usize;
2585
2586 let child_offset = array_data.offset() * fixed_size;
2587 let child_length = array_data.len() * fixed_size;
2588 let child_data = array_data.child_data()[0].slice(child_offset, child_length);
2589
2590 offset = write_array_data(
2591 &child_data,
2592 meta,
2593 sink,
2594 offset,
2595 compression_codec,
2596 ipc_write_context,
2597 write_options,
2598 )?;
2599 return Ok(offset);
2600 } else {
2601 for buffer in array_data.buffers() {
2602 offset = encode_sink_buffer(
2603 buffer.clone(),
2604 meta,
2605 sink,
2606 offset,
2607 compression_codec,
2608 ipc_write_context,
2609 write_options.alignment,
2610 )?;
2611 }
2612 }
2613
2614 match array_data.data_type() {
2615 DataType::Dictionary(_, _) => {}
2616 DataType::RunEndEncoded(_, _) => {
2617 let arr = unslice_run_array(array_data.clone())?;
2619 for data_ref in arr.child_data() {
2621 offset = write_array_data(
2623 data_ref,
2624 meta,
2625 sink,
2626 offset,
2627 compression_codec,
2628 ipc_write_context,
2629 write_options,
2630 )?;
2631 }
2632 }
2633 _ => {
2634 for data_ref in array_data.child_data() {
2636 offset = write_array_data(
2638 data_ref,
2639 meta,
2640 sink,
2641 offset,
2642 compression_codec,
2643 ipc_write_context,
2644 write_options,
2645 )?;
2646 }
2647 }
2648 }
2649 Ok(offset)
2650}
2651
2652fn encode_sink_buffer(
2666 buffer: Buffer,
2667 ipc_meta_data: &mut IpcMetadataBuilder,
2668 sink: &mut IpcBodySink<'_>,
2669 offset: i64,
2670 compression_codec: Option<CompressionCodec>,
2671 ipc_write_context: &mut IpcWriteContext,
2672 alignment: u8,
2673) -> Result<i64, ArrowError> {
2674 let (encoded, len) = match compression_codec {
2675 None => {
2676 let len = buffer.len() as i64;
2677 (EncodedBuffer::Raw(buffer), len)
2678 }
2679 Some(codec) => {
2680 let mut scratch = Vec::new();
2681 let written =
2682 codec.compress_to_vec(buffer.as_slice(), &mut scratch, ipc_write_context)?;
2683 let len = i64::try_from(written)
2684 .map_err(|e| ArrowError::InvalidArgumentError(format!("{e}")))?;
2685 (EncodedBuffer::Compressed(scratch), len)
2686 }
2687 };
2688
2689 let pad_len = pad_to_alignment(alignment, len as usize);
2690 sink.write(pad_len, encoded);
2691 ipc_meta_data.buffers.push(crate::Buffer::new(offset, len));
2692 Ok(offset + len + pad_len as i64)
2693}
2694
2695const PADDING: [u8; 64] = [0; 64];
2696
2697#[inline]
2703fn estimate_encoded_buffer_count(dt: &DataType) -> usize {
2704 match dt {
2705 DataType::Null => 0,
2706
2707 DataType::Binary | DataType::Utf8 | DataType::LargeBinary | DataType::LargeUtf8 => 3,
2708
2709 DataType::BinaryView | DataType::Utf8View => 3,
2710
2711 DataType::List(f) | DataType::LargeList(f) | DataType::Map(f, _) => {
2712 2 + estimate_encoded_buffer_count(f.data_type())
2713 }
2714
2715 DataType::ListView(f) | DataType::LargeListView(f) => {
2716 3 + estimate_encoded_buffer_count(f.data_type())
2717 }
2718
2719 DataType::FixedSizeList(f, _) => 1 + estimate_encoded_buffer_count(f.data_type()),
2720
2721 DataType::Struct(fields) => {
2722 1 + fields
2723 .iter()
2724 .map(|f| estimate_encoded_buffer_count(f.data_type()))
2725 .sum::<usize>()
2726 }
2727
2728 DataType::Dictionary(_, _) => 2,
2730
2731 DataType::Union(fields, UnionMode::Sparse) => {
2732 1 + fields
2733 .iter()
2734 .map(|(_, f)| estimate_encoded_buffer_count(f.data_type()))
2735 .sum::<usize>()
2736 }
2737 DataType::Union(fields, UnionMode::Dense) => {
2738 2 + fields
2739 .iter()
2740 .map(|(_, f)| estimate_encoded_buffer_count(f.data_type()))
2741 .sum::<usize>()
2742 }
2743
2744 DataType::RunEndEncoded(run_ends, values) => {
2745 estimate_encoded_buffer_count(run_ends.data_type())
2746 + estimate_encoded_buffer_count(values.data_type())
2747 }
2748 _ => 2,
2750 }
2751}
2752
2753#[inline]
2755fn pad_to_alignment(alignment: u8, len: usize) -> usize {
2756 let a = usize::from(alignment - 1);
2757 ((len + a) & !a) - len
2758}
2759
2760#[cfg(test)]
2761mod tests {
2762 use std::hash::Hasher;
2763 use std::io::Cursor;
2764 use std::io::Seek;
2765
2766 use arrow_array::builder::FixedSizeListBuilder;
2767 use arrow_array::builder::Float32Builder;
2768 use arrow_array::builder::Int64Builder;
2769 use arrow_array::builder::MapBuilder;
2770 use arrow_array::builder::StringViewBuilder;
2771 use arrow_array::builder::UnionBuilder;
2772 use arrow_array::builder::{
2773 GenericListBuilder, GenericListViewBuilder, ListBuilder, StringBuilder,
2774 };
2775 use arrow_array::builder::{PrimitiveRunBuilder, UInt32Builder};
2776 use arrow_array::types::*;
2777 use arrow_buffer::ScalarBuffer;
2778
2779 use crate::MetadataVersion;
2780 use crate::convert::try_fb_to_schema;
2781 use crate::reader::*;
2782 use crate::root_as_footer;
2783
2784 use super::*;
2785
2786 fn serialize_file(rb: &RecordBatch) -> Vec<u8> {
2787 let mut writer = FileWriter::try_new(vec![], rb.schema_ref()).unwrap();
2788 writer.write(rb).unwrap();
2789 writer.finish().unwrap();
2790 writer.into_inner().unwrap()
2791 }
2792
2793 fn deserialize_file(bytes: Vec<u8>) -> RecordBatch {
2794 let mut reader = FileReader::try_new(Cursor::new(bytes), None).unwrap();
2795 reader.next().unwrap().unwrap()
2796 }
2797
2798 fn serialize_stream(record: &RecordBatch) -> Vec<u8> {
2799 const IPC_ALIGNMENT: usize = 8;
2803
2804 let mut stream_writer = StreamWriter::try_new_with_options(
2805 vec![],
2806 record.schema_ref(),
2807 IpcWriteOptions::try_new(IPC_ALIGNMENT, false, MetadataVersion::V5).unwrap(),
2808 )
2809 .unwrap();
2810 stream_writer.write(record).unwrap();
2811 stream_writer.finish().unwrap();
2812 stream_writer.into_inner().unwrap()
2813 }
2814
2815 fn deserialize_stream(bytes: Vec<u8>) -> RecordBatch {
2816 let mut stream_reader = StreamReader::try_new(Cursor::new(bytes), None).unwrap();
2817 stream_reader.next().unwrap().unwrap()
2818 }
2819
2820 fn encode_stream(
2825 schema: &Schema,
2826 batches: &[RecordBatch],
2827 options: IpcWriteOptions,
2828 ) -> Vec<u8> {
2829 let mut encoder = StreamEncoder::try_new_with_options(schema, options).unwrap();
2830 let mut bytes = Vec::new();
2831 for batch in batches {
2832 for buffer in encoder.encode(batch).unwrap() {
2833 bytes.write_all(buffer.as_slice()).unwrap();
2834 }
2835 }
2836 for buffer in encoder.finish().unwrap() {
2837 bytes.write_all(buffer.as_slice()).unwrap();
2838 }
2839 bytes
2840 }
2841
2842 fn write_stream(schema: &Schema, batches: &[RecordBatch], options: IpcWriteOptions) -> Vec<u8> {
2843 let mut bytes = Vec::new();
2844 let mut writer = StreamWriter::try_new_with_options(&mut bytes, schema, options).unwrap();
2845 for batch in batches {
2846 writer.write(batch).unwrap();
2847 }
2848 writer.finish().unwrap();
2849 bytes
2850 }
2851
2852 #[tokio::test]
2855 async fn test_stream_encoder_async_writer_matches_stream_writer() {
2856 use tokio::io::{AsyncReadExt, AsyncWriteExt};
2857
2858 let batch = record_batch!(("a", Int32, [1, 2, 3]), ("b", Utf8, ["x", "y", "z"])).unwrap();
2859 let options = IpcWriteOptions::default();
2860 let expected = write_stream(
2861 batch.schema_ref(),
2862 std::slice::from_ref(&batch),
2863 options.clone(),
2864 );
2865
2866 let (mut sink, mut source) = tokio::io::duplex(64);
2867 let read = tokio::spawn(async move {
2868 let mut bytes = Vec::new();
2869 source.read_to_end(&mut bytes).await.unwrap();
2870 bytes
2871 });
2872
2873 let mut encoder = StreamEncoder::try_new_with_options(batch.schema_ref(), options).unwrap();
2874 for buffer in encoder.encode(&batch).unwrap() {
2875 sink.write_all(buffer.as_slice()).await.unwrap();
2876 }
2877 for buffer in encoder.finish().unwrap() {
2878 sink.write_all(buffer.as_slice()).await.unwrap();
2879 }
2880 sink.shutdown().await.unwrap();
2881
2882 let encoded = read.await.unwrap();
2883 assert_eq!(encoded, expected);
2884
2885 let mut reader = StreamReader::try_new(Cursor::new(encoded), None).unwrap();
2886 assert_eq!(reader.next().unwrap().unwrap(), batch);
2887 assert!(reader.next().is_none());
2888 }
2889
2890 #[test]
2891 fn test_stream_encoder_empty_stream_matches_stream_writer() {
2892 let schema = Schema::new(vec![Field::new("a", DataType::Int32, true)]);
2893 let options = IpcWriteOptions::default();
2894 let encoded = encode_stream(&schema, &[], options.clone());
2895 let written = write_stream(&schema, &[], options);
2896
2897 assert_eq!(encoded, written);
2898
2899 let mut reader = StreamReader::try_new(Cursor::new(encoded), None).unwrap();
2900 assert!(reader.next().is_none());
2901 }
2902
2903 #[test]
2904 fn test_stream_encoder_dictionary_batches_match_stream_writer() {
2905 let schema = Arc::new(Schema::new(vec![Field::new(
2906 "a",
2907 DataType::Dictionary(Box::new(DataType::UInt8), Box::new(DataType::Utf8)),
2908 false,
2909 )]));
2910 let batch = RecordBatch::try_new(
2911 schema.clone(),
2912 vec![Arc::new(DictionaryArray::new(
2913 UInt8Array::from_iter_values([0, 1, 0]),
2914 Arc::new(StringArray::from_iter_values(["a", "b"])),
2915 ))],
2916 )
2917 .unwrap();
2918 let options = IpcWriteOptions::default();
2919 let encoded = encode_stream(&schema, std::slice::from_ref(&batch), options.clone());
2920 let written = write_stream(&schema, std::slice::from_ref(&batch), options);
2921
2922 assert_eq!(encoded, written);
2923
2924 let mut reader = StreamReader::try_new(Cursor::new(encoded), None).unwrap();
2925 assert_eq!(reader.next().unwrap().unwrap(), batch);
2926 assert!(reader.next().is_none());
2927 }
2928
2929 #[test]
2930 #[cfg(feature = "lz4")]
2931 fn test_write_empty_record_batch_lz4_compression() {
2932 let schema = Schema::new(vec![Field::new("field1", DataType::Int32, true)]);
2933 let values: Vec<Option<i32>> = vec![];
2934 let array = Int32Array::from(values);
2935 let record_batch =
2936 RecordBatch::try_new(Arc::new(schema.clone()), vec![Arc::new(array)]).unwrap();
2937
2938 let mut file = tempfile::tempfile().unwrap();
2939
2940 {
2941 let write_option = IpcWriteOptions::try_new(8, false, crate::MetadataVersion::V5)
2942 .unwrap()
2943 .try_with_compression(Some(crate::CompressionType::LZ4_FRAME))
2944 .unwrap();
2945
2946 let mut writer =
2947 FileWriter::try_new_with_options(&mut file, &schema, write_option).unwrap();
2948 writer.write(&record_batch).unwrap();
2949 writer.finish().unwrap();
2950 }
2951 file.rewind().unwrap();
2952 {
2953 let reader = FileReader::try_new(file, None).unwrap();
2955 for read_batch in reader {
2956 read_batch
2957 .unwrap()
2958 .columns()
2959 .iter()
2960 .zip(record_batch.columns())
2961 .for_each(|(a, b)| {
2962 assert_eq!(a.data_type(), b.data_type());
2963 assert_eq!(a.len(), b.len());
2964 assert_eq!(a.null_count(), b.null_count());
2965 });
2966 }
2967 }
2968 }
2969
2970 #[test]
2971 #[cfg(feature = "lz4")]
2972 fn test_write_file_with_lz4_compression() {
2973 let schema = Schema::new(vec![Field::new("field1", DataType::Int32, true)]);
2974 let values: Vec<Option<i32>> = vec![Some(12), Some(1)];
2975 let array = Int32Array::from(values);
2976 let record_batch =
2977 RecordBatch::try_new(Arc::new(schema.clone()), vec![Arc::new(array)]).unwrap();
2978
2979 let mut file = tempfile::tempfile().unwrap();
2980 {
2981 let write_option = IpcWriteOptions::try_new(8, false, crate::MetadataVersion::V5)
2982 .unwrap()
2983 .try_with_compression(Some(crate::CompressionType::LZ4_FRAME))
2984 .unwrap();
2985
2986 let mut writer =
2987 FileWriter::try_new_with_options(&mut file, &schema, write_option).unwrap();
2988 writer.write(&record_batch).unwrap();
2989 writer.finish().unwrap();
2990 }
2991 file.rewind().unwrap();
2992 {
2993 let reader = FileReader::try_new(file, None).unwrap();
2995 for read_batch in reader {
2996 read_batch
2997 .unwrap()
2998 .columns()
2999 .iter()
3000 .zip(record_batch.columns())
3001 .for_each(|(a, b)| {
3002 assert_eq!(a.data_type(), b.data_type());
3003 assert_eq!(a.len(), b.len());
3004 assert_eq!(a.null_count(), b.null_count());
3005 });
3006 }
3007 }
3008 }
3009
3010 #[test]
3011 #[cfg(feature = "zstd")]
3012 fn test_write_file_with_zstd_compression() {
3013 let schema = Schema::new(vec![Field::new("field1", DataType::Int32, true)]);
3014 let values: Vec<Option<i32>> = vec![Some(12), Some(1)];
3015 let array = Int32Array::from(values);
3016 let record_batch =
3017 RecordBatch::try_new(Arc::new(schema.clone()), vec![Arc::new(array)]).unwrap();
3018 let mut file = tempfile::tempfile().unwrap();
3019 {
3020 let write_option = IpcWriteOptions::try_new(8, false, crate::MetadataVersion::V5)
3021 .unwrap()
3022 .try_with_compression(Some(crate::CompressionType::ZSTD))
3023 .unwrap()
3024 .try_with_compression_level(Some(1))
3025 .unwrap();
3026
3027 let mut writer =
3028 FileWriter::try_new_with_options(&mut file, &schema, write_option).unwrap();
3029 writer.write(&record_batch).unwrap();
3030 writer.finish().unwrap();
3031 }
3032 file.rewind().unwrap();
3033 {
3034 let reader = FileReader::try_new(file, None).unwrap();
3036 for read_batch in reader {
3037 read_batch
3038 .unwrap()
3039 .columns()
3040 .iter()
3041 .zip(record_batch.columns())
3042 .for_each(|(a, b)| {
3043 assert_eq!(a.data_type(), b.data_type());
3044 assert_eq!(a.len(), b.len());
3045 assert_eq!(a.null_count(), b.null_count());
3046 });
3047 }
3048 }
3049 }
3050
3051 #[test]
3052 fn test_write_file() {
3053 let schema = Schema::new(vec![Field::new("field1", DataType::UInt32, true)]);
3054 let values: Vec<Option<u32>> = vec![
3055 Some(999),
3056 None,
3057 Some(235),
3058 Some(123),
3059 None,
3060 None,
3061 None,
3062 None,
3063 None,
3064 ];
3065 let array1 = UInt32Array::from(values);
3066 let batch =
3067 RecordBatch::try_new(Arc::new(schema.clone()), vec![Arc::new(array1) as ArrayRef])
3068 .unwrap();
3069 let mut file = tempfile::tempfile().unwrap();
3070 {
3071 let mut writer = FileWriter::try_new(&mut file, &schema).unwrap();
3072
3073 writer.write(&batch).unwrap();
3074 writer.finish().unwrap();
3075 }
3076 file.rewind().unwrap();
3077
3078 {
3079 let mut reader = FileReader::try_new(file, None).unwrap();
3080 while let Some(Ok(read_batch)) = reader.next() {
3081 read_batch
3082 .columns()
3083 .iter()
3084 .zip(batch.columns())
3085 .for_each(|(a, b)| {
3086 assert_eq!(a.data_type(), b.data_type());
3087 assert_eq!(a.len(), b.len());
3088 assert_eq!(a.null_count(), b.null_count());
3089 });
3090 }
3091 }
3092 }
3093
3094 #[test]
3095 fn test_empty_utf8_ipc_writes_nonempty_offsets_buffer() {
3096 let name = StringArray::from(Vec::<String>::new());
3097 let [offsets, values] = get_byte_array_buffers::<i32>(&name.to_data());
3098
3099 assert_eq!(name.len(), 0);
3100 assert_eq!(
3101 offsets.len(),
3102 std::mem::size_of::<i32>(),
3103 "offsets buffer should contain one zero i32 offset"
3104 );
3105 assert_eq!(values.len(), 0, "values buffer should remain empty");
3106 }
3107
3108 #[test]
3109 fn test_empty_large_utf8_ipc_writes_nonempty_offsets_buffer() {
3110 let name = LargeStringArray::from(Vec::<String>::new());
3111 let [offsets, values] = get_byte_array_buffers::<i64>(&name.to_data());
3112
3113 assert_eq!(name.len(), 0);
3114 assert_eq!(
3115 offsets.len(),
3116 std::mem::size_of::<i64>(),
3117 "offsets buffer should contain one zero i64 offset"
3118 );
3119 assert_eq!(values.len(), 0, "values buffer should remain empty");
3120 }
3121
3122 #[test]
3123 fn test_empty_list_ipc_writes_nonempty_offsets_buffer() {
3124 let list = GenericListBuilder::<i32, _>::new(UInt32Builder::new()).finish();
3125 let (offsets, child_data) = get_list_array_buffers::<i32>(&list.to_data());
3126
3127 assert_eq!(list.len(), 0);
3128 assert_eq!(
3129 offsets.len(),
3130 std::mem::size_of::<i32>(),
3131 "offsets buffer should contain one zero i32 offset"
3132 );
3133 assert_eq!(child_data.len(), 0, "child data should remain empty");
3134 }
3135
3136 #[test]
3137 fn test_empty_large_list_ipc_writes_nonempty_offsets_buffer() {
3138 let list = GenericListBuilder::<i64, _>::new(UInt32Builder::new()).finish();
3139 let (offsets, child_data) = get_list_array_buffers::<i64>(&list.to_data());
3140
3141 assert_eq!(list.len(), 0);
3142 assert_eq!(
3143 offsets.len(),
3144 std::mem::size_of::<i64>(),
3145 "offsets buffer should contain one zero i64 offset"
3146 );
3147 assert_eq!(child_data.len(), 0, "child data should remain empty");
3148 }
3149
3150 fn write_null_file(options: IpcWriteOptions) {
3151 let schema = Schema::new(vec![
3152 Field::new("nulls", DataType::Null, true),
3153 Field::new("int32s", DataType::Int32, false),
3154 Field::new("nulls2", DataType::Null, true),
3155 Field::new("f64s", DataType::Float64, false),
3156 ]);
3157 let array1 = NullArray::new(32);
3158 let array2 = Int32Array::from(vec![1; 32]);
3159 let array3 = NullArray::new(32);
3160 let array4 = Float64Array::from(vec![f64::NAN; 32]);
3161 let batch = RecordBatch::try_new(
3162 Arc::new(schema.clone()),
3163 vec![
3164 Arc::new(array1) as ArrayRef,
3165 Arc::new(array2) as ArrayRef,
3166 Arc::new(array3) as ArrayRef,
3167 Arc::new(array4) as ArrayRef,
3168 ],
3169 )
3170 .unwrap();
3171 let mut file = tempfile::tempfile().unwrap();
3172 {
3173 let mut writer = FileWriter::try_new_with_options(&mut file, &schema, options).unwrap();
3174
3175 writer.write(&batch).unwrap();
3176 writer.finish().unwrap();
3177 }
3178
3179 file.rewind().unwrap();
3180
3181 {
3182 let reader = FileReader::try_new(file, None).unwrap();
3183 reader.for_each(|maybe_batch| {
3184 maybe_batch
3185 .unwrap()
3186 .columns()
3187 .iter()
3188 .zip(batch.columns())
3189 .for_each(|(a, b)| {
3190 assert_eq!(a.data_type(), b.data_type());
3191 assert_eq!(a.len(), b.len());
3192 assert_eq!(a.null_count(), b.null_count());
3193 });
3194 });
3195 }
3196 }
3197 #[test]
3198 fn test_write_null_file_v4() {
3199 write_null_file(IpcWriteOptions::try_new(8, false, MetadataVersion::V4).unwrap());
3200 write_null_file(IpcWriteOptions::try_new(8, true, MetadataVersion::V4).unwrap());
3201 write_null_file(IpcWriteOptions::try_new(64, false, MetadataVersion::V4).unwrap());
3202 write_null_file(IpcWriteOptions::try_new(64, true, MetadataVersion::V4).unwrap());
3203 }
3204
3205 #[test]
3206 fn test_write_null_file_v5() {
3207 write_null_file(IpcWriteOptions::try_new(8, false, MetadataVersion::V5).unwrap());
3208 write_null_file(IpcWriteOptions::try_new(64, false, MetadataVersion::V5).unwrap());
3209 }
3210
3211 #[test]
3212 fn track_union_nested_dict() {
3213 let inner: DictionaryArray<Int32Type> = vec!["a", "b", "a"].into_iter().collect();
3214
3215 let array = Arc::new(inner) as ArrayRef;
3216
3217 #[expect(deprecated)]
3219 let dctfield = Field::new_dict("dict", array.data_type().clone(), false, 0, false);
3220 let union_fields = std::iter::once((0, Arc::new(dctfield))).collect();
3221
3222 let types = [0, 0, 0].into_iter().collect::<ScalarBuffer<i8>>();
3223 let offsets = [0, 1, 2].into_iter().collect::<ScalarBuffer<i32>>();
3224
3225 let union = UnionArray::try_new(union_fields, types, Some(offsets), vec![array]).unwrap();
3226
3227 let schema = Arc::new(Schema::new(vec![Field::new(
3228 "union",
3229 union.data_type().clone(),
3230 false,
3231 )]));
3232
3233 let r#gen = IpcDataGenerator::default();
3234 let mut dict_tracker = DictionaryTracker::new(false);
3235 r#gen.schema_to_bytes_with_dictionary_tracker(
3236 &schema,
3237 &mut dict_tracker,
3238 &IpcWriteOptions::default(),
3239 );
3240
3241 let batch = RecordBatch::try_new(schema, vec![Arc::new(union)]).unwrap();
3242
3243 r#gen
3244 .encode(
3245 &batch,
3246 &mut dict_tracker,
3247 &Default::default(),
3248 &mut Default::default(),
3249 )
3250 .unwrap();
3251
3252 assert!(dict_tracker.written.contains_key(&0));
3255 }
3256
3257 #[test]
3258 fn track_struct_nested_dict() {
3259 let inner: DictionaryArray<Int32Type> = vec!["a", "b", "a"].into_iter().collect();
3260
3261 let array = Arc::new(inner) as ArrayRef;
3262
3263 #[expect(deprecated)]
3265 let dctfield = Arc::new(Field::new_dict(
3266 "dict",
3267 array.data_type().clone(),
3268 false,
3269 2,
3270 false,
3271 ));
3272
3273 let s = StructArray::from(vec![(dctfield, array)]);
3274 let struct_array = Arc::new(s) as ArrayRef;
3275
3276 let schema = Arc::new(Schema::new(vec![Field::new(
3277 "struct",
3278 struct_array.data_type().clone(),
3279 false,
3280 )]));
3281
3282 let r#gen = IpcDataGenerator::default();
3283 let mut dict_tracker = DictionaryTracker::new(false);
3284 r#gen.schema_to_bytes_with_dictionary_tracker(
3285 &schema,
3286 &mut dict_tracker,
3287 &IpcWriteOptions::default(),
3288 );
3289
3290 let batch = RecordBatch::try_new(schema, vec![struct_array]).unwrap();
3291
3292 r#gen
3293 .encode(
3294 &batch,
3295 &mut dict_tracker,
3296 &Default::default(),
3297 &mut Default::default(),
3298 )
3299 .unwrap();
3300
3301 assert!(dict_tracker.written.contains_key(&0));
3302 }
3303
3304 fn write_union_file(options: IpcWriteOptions) {
3305 let schema = Schema::new(vec![Field::new_union(
3306 "union",
3307 vec![0, 1],
3308 vec![
3309 Field::new("a", DataType::Int32, false),
3310 Field::new("c", DataType::Float64, false),
3311 ],
3312 UnionMode::Sparse,
3313 )]);
3314 let mut builder = UnionBuilder::with_capacity_sparse(5);
3315 builder.append::<Int32Type>("a", 1).unwrap();
3316 builder.append_null::<Int32Type>("a").unwrap();
3317 builder.append::<Float64Type>("c", 3.0).unwrap();
3318 builder.append_null::<Float64Type>("c").unwrap();
3319 builder.append::<Int32Type>("a", 4).unwrap();
3320 let union = builder.build().unwrap();
3321
3322 let batch =
3323 RecordBatch::try_new(Arc::new(schema.clone()), vec![Arc::new(union) as ArrayRef])
3324 .unwrap();
3325
3326 let mut file = tempfile::tempfile().unwrap();
3327 {
3328 let mut writer = FileWriter::try_new_with_options(&mut file, &schema, options).unwrap();
3329
3330 writer.write(&batch).unwrap();
3331 writer.finish().unwrap();
3332 }
3333 file.rewind().unwrap();
3334
3335 {
3336 let reader = FileReader::try_new(file, None).unwrap();
3337 reader.for_each(|maybe_batch| {
3338 maybe_batch
3339 .unwrap()
3340 .columns()
3341 .iter()
3342 .zip(batch.columns())
3343 .for_each(|(a, b)| {
3344 assert_eq!(a.data_type(), b.data_type());
3345 assert_eq!(a.len(), b.len());
3346 assert_eq!(a.null_count(), b.null_count());
3347 });
3348 });
3349 }
3350 }
3351
3352 #[test]
3353 fn test_write_union_file_v4_v5() {
3354 write_union_file(IpcWriteOptions::try_new(8, false, MetadataVersion::V4).unwrap());
3355 write_union_file(IpcWriteOptions::try_new(8, false, MetadataVersion::V5).unwrap());
3356 }
3357
3358 #[test]
3359 fn test_write_view_types() {
3360 const LONG_TEST_STRING: &str =
3361 "This is a long string to make sure binary view array handles it";
3362 let schema = Schema::new(vec![
3363 Field::new("field1", DataType::BinaryView, true),
3364 Field::new("field2", DataType::Utf8View, true),
3365 ]);
3366 let values: Vec<Option<&[u8]>> = vec![
3367 Some(b"foo"),
3368 Some(b"bar"),
3369 Some(LONG_TEST_STRING.as_bytes()),
3370 ];
3371 let binary_array = BinaryViewArray::from_iter(values);
3372 let utf8_array =
3373 StringViewArray::from_iter(vec![Some("foo"), Some("bar"), Some(LONG_TEST_STRING)]);
3374 let record_batch = RecordBatch::try_new(
3375 Arc::new(schema.clone()),
3376 vec![Arc::new(binary_array), Arc::new(utf8_array)],
3377 )
3378 .unwrap();
3379
3380 let mut file = tempfile::tempfile().unwrap();
3381 {
3382 let mut writer = FileWriter::try_new(&mut file, &schema).unwrap();
3383 writer.write(&record_batch).unwrap();
3384 writer.finish().unwrap();
3385 }
3386 file.rewind().unwrap();
3387 {
3388 let mut reader = FileReader::try_new(&file, None).unwrap();
3389 let read_batch = reader.next().unwrap().unwrap();
3390 read_batch
3391 .columns()
3392 .iter()
3393 .zip(record_batch.columns())
3394 .for_each(|(a, b)| {
3395 assert_eq!(a, b);
3396 });
3397 }
3398 file.rewind().unwrap();
3399 {
3400 let mut reader = FileReader::try_new(&file, Some(vec![0])).unwrap();
3401 let read_batch = reader.next().unwrap().unwrap();
3402 assert_eq!(read_batch.num_columns(), 1);
3403 let read_array = read_batch.column(0);
3404 let write_array = record_batch.column(0);
3405 assert_eq!(read_array, write_array);
3406 }
3407 }
3408
3409 #[test]
3410 #[cfg_attr(miri, ignore)] fn truncate_ipc_record_batch() {
3412 fn create_batch(rows: usize) -> RecordBatch {
3413 let schema = Schema::new(vec![
3414 Field::new("a", DataType::Int32, false),
3415 Field::new("b", DataType::Utf8, false),
3416 ]);
3417
3418 let a = Int32Array::from_iter_values(0..rows as i32);
3419 let b = StringArray::from_iter_values((0..rows).map(|i| i.to_string()));
3420
3421 RecordBatch::try_new(Arc::new(schema), vec![Arc::new(a), Arc::new(b)]).unwrap()
3422 }
3423
3424 let big_record_batch = create_batch(65536);
3425
3426 let length = 5;
3427 let small_record_batch = create_batch(length);
3428
3429 let offset = 2;
3430 let record_batch_slice = big_record_batch.slice(offset, length);
3431 assert!(
3432 serialize_stream(&big_record_batch).len() > serialize_stream(&small_record_batch).len()
3433 );
3434 assert_eq!(
3435 serialize_stream(&small_record_batch).len(),
3436 serialize_stream(&record_batch_slice).len()
3437 );
3438
3439 assert_eq!(
3440 deserialize_stream(serialize_stream(&record_batch_slice)),
3441 record_batch_slice
3442 );
3443 }
3444
3445 #[test]
3446 fn truncate_ipc_record_batch_with_nulls() {
3447 fn create_batch() -> RecordBatch {
3448 let schema = Schema::new(vec![
3449 Field::new("a", DataType::Int32, true),
3450 Field::new("b", DataType::Utf8, true),
3451 ]);
3452
3453 let a = Int32Array::from(vec![Some(1), None, Some(1), None, Some(1)]);
3454 let b = StringArray::from(vec![None, Some("a"), Some("a"), None, Some("a")]);
3455
3456 RecordBatch::try_new(Arc::new(schema), vec![Arc::new(a), Arc::new(b)]).unwrap()
3457 }
3458
3459 let record_batch = create_batch();
3460 let record_batch_slice = record_batch.slice(1, 2);
3461 let deserialized_batch = deserialize_stream(serialize_stream(&record_batch_slice));
3462
3463 assert!(
3464 serialize_stream(&record_batch).len() > serialize_stream(&record_batch_slice).len()
3465 );
3466
3467 assert!(deserialized_batch.column(0).is_null(0));
3468 assert!(deserialized_batch.column(0).is_valid(1));
3469 assert!(deserialized_batch.column(1).is_valid(0));
3470 assert!(deserialized_batch.column(1).is_valid(1));
3471
3472 assert_eq!(record_batch_slice, deserialized_batch);
3473 }
3474
3475 #[test]
3476 fn truncate_ipc_dictionary_array() {
3477 fn create_batch() -> RecordBatch {
3478 let values: StringArray = [Some("foo"), Some("bar"), Some("baz")]
3479 .into_iter()
3480 .collect();
3481 let keys: Int32Array = [Some(0), Some(2), None, Some(1)].into_iter().collect();
3482
3483 let array = DictionaryArray::new(keys, Arc::new(values));
3484
3485 let schema = Schema::new(vec![Field::new("dict", array.data_type().clone(), true)]);
3486
3487 RecordBatch::try_new(Arc::new(schema), vec![Arc::new(array)]).unwrap()
3488 }
3489
3490 let record_batch = create_batch();
3491 let record_batch_slice = record_batch.slice(1, 2);
3492 let deserialized_batch = deserialize_stream(serialize_stream(&record_batch_slice));
3493
3494 assert!(
3495 serialize_stream(&record_batch).len() > serialize_stream(&record_batch_slice).len()
3496 );
3497
3498 assert!(deserialized_batch.column(0).is_valid(0));
3499 assert!(deserialized_batch.column(0).is_null(1));
3500
3501 assert_eq!(record_batch_slice, deserialized_batch);
3502 }
3503
3504 #[test]
3505 fn truncate_ipc_struct_array() {
3506 fn create_batch() -> RecordBatch {
3507 let strings: StringArray = [Some("foo"), None, Some("bar"), Some("baz")]
3508 .into_iter()
3509 .collect();
3510 let ints: Int32Array = [Some(0), Some(2), None, Some(1)].into_iter().collect();
3511
3512 let struct_array = StructArray::from(vec![
3513 (
3514 Arc::new(Field::new("s", DataType::Utf8, true)),
3515 Arc::new(strings) as ArrayRef,
3516 ),
3517 (
3518 Arc::new(Field::new("c", DataType::Int32, true)),
3519 Arc::new(ints) as ArrayRef,
3520 ),
3521 ]);
3522
3523 let schema = Schema::new(vec![Field::new(
3524 "struct_array",
3525 struct_array.data_type().clone(),
3526 true,
3527 )]);
3528
3529 RecordBatch::try_new(Arc::new(schema), vec![Arc::new(struct_array)]).unwrap()
3530 }
3531
3532 let record_batch = create_batch();
3533 let record_batch_slice = record_batch.slice(1, 2);
3534 let deserialized_batch = deserialize_stream(serialize_stream(&record_batch_slice));
3535
3536 assert!(
3537 serialize_stream(&record_batch).len() > serialize_stream(&record_batch_slice).len()
3538 );
3539
3540 let structs = deserialized_batch
3541 .column(0)
3542 .as_any()
3543 .downcast_ref::<StructArray>()
3544 .unwrap();
3545
3546 assert!(structs.column(0).is_null(0));
3547 assert!(structs.column(0).is_valid(1));
3548 assert!(structs.column(1).is_valid(0));
3549 assert!(structs.column(1).is_null(1));
3550 assert_eq!(record_batch_slice, deserialized_batch);
3551 }
3552
3553 #[test]
3554 fn truncate_ipc_string_array_with_all_empty_string() {
3555 fn create_batch() -> RecordBatch {
3556 let schema = Schema::new(vec![Field::new("a", DataType::Utf8, true)]);
3557 let a = StringArray::from(vec![Some(""), Some(""), Some(""), Some(""), Some("")]);
3558 RecordBatch::try_new(Arc::new(schema), vec![Arc::new(a)]).unwrap()
3559 }
3560
3561 let record_batch = create_batch();
3562 let record_batch_slice = record_batch.slice(0, 1);
3563 let deserialized_batch = deserialize_stream(serialize_stream(&record_batch_slice));
3564
3565 assert!(
3566 serialize_stream(&record_batch).len() > serialize_stream(&record_batch_slice).len()
3567 );
3568 assert_eq!(record_batch_slice, deserialized_batch);
3569 }
3570
3571 #[test]
3572 fn test_stream_writer_writes_array_slice() {
3573 let array = UInt32Array::from(vec![Some(1), Some(2), Some(3)]);
3574 assert_eq!(
3575 vec![Some(1), Some(2), Some(3)],
3576 array.iter().collect::<Vec<_>>()
3577 );
3578
3579 let sliced = array.slice(1, 2);
3580 assert_eq!(vec![Some(2), Some(3)], sliced.iter().collect::<Vec<_>>());
3581
3582 let batch = RecordBatch::try_new(
3583 Arc::new(Schema::new(vec![Field::new("a", DataType::UInt32, true)])),
3584 vec![Arc::new(sliced)],
3585 )
3586 .expect("new batch");
3587
3588 let mut writer = StreamWriter::try_new(vec![], batch.schema_ref()).expect("new writer");
3589 writer.write(&batch).expect("write");
3590 let outbuf = writer.into_inner().expect("inner");
3591
3592 let mut reader = StreamReader::try_new(&outbuf[..], None).expect("new reader");
3593 let read_batch = reader.next().unwrap().expect("read batch");
3594
3595 let read_array: &UInt32Array = read_batch.column(0).as_primitive();
3596 assert_eq!(
3597 vec![Some(2), Some(3)],
3598 read_array.iter().collect::<Vec<_>>()
3599 );
3600 }
3601
3602 #[test]
3603 fn test_large_slice_uint32() {
3604 ensure_roundtrip(Arc::new(UInt32Array::from_iter(
3605 (0..8000).map(|i| if i % 2 == 0 { Some(i) } else { None }),
3606 )));
3607 }
3608
3609 #[test]
3610 #[cfg_attr(miri, ignore)] fn test_large_slice_string() {
3612 let strings: Vec<_> = (0..8000)
3613 .map(|i| {
3614 if i % 2 == 0 {
3615 Some(format!("value{i}"))
3616 } else {
3617 None
3618 }
3619 })
3620 .collect();
3621
3622 ensure_roundtrip(Arc::new(StringArray::from(strings)));
3623 }
3624
3625 #[test]
3626 #[cfg_attr(miri, ignore)] fn test_large_slice_string_list() {
3628 let mut ls = ListBuilder::new(StringBuilder::new());
3629
3630 let mut s = String::new();
3631 for row_number in 0..8000 {
3632 if row_number % 2 == 0 {
3633 for list_element in 0..1000 {
3634 s.clear();
3635 use std::fmt::Write;
3636 write!(&mut s, "value{row_number}-{list_element}").unwrap();
3637 ls.values().append_value(&s);
3638 }
3639 ls.append(true)
3640 } else {
3641 ls.append(false); }
3643 }
3644
3645 ensure_roundtrip(Arc::new(ls.finish()));
3646 }
3647
3648 #[test]
3649 #[cfg_attr(miri, ignore)] fn test_large_slice_string_list_of_lists() {
3651 let mut ls = ListBuilder::new(ListBuilder::new(StringBuilder::new()));
3655
3656 for _ in 0..4000 {
3657 ls.values().append(true);
3658 ls.append(true)
3659 }
3660
3661 let mut s = String::new();
3662 for row_number in 0..4000 {
3663 if row_number % 2 == 0 {
3664 for list_element in 0..1000 {
3665 s.clear();
3666 use std::fmt::Write;
3667 write!(&mut s, "value{row_number}-{list_element}").unwrap();
3668 ls.values().values().append_value(&s);
3669 }
3670 ls.values().append(true);
3671 ls.append(true)
3672 } else {
3673 ls.append(false); }
3675 }
3676
3677 ensure_roundtrip(Arc::new(ls.finish()));
3678 }
3679
3680 fn ensure_roundtrip(array: ArrayRef) {
3682 let num_rows = array.len();
3683 let orig_batch = RecordBatch::try_from_iter(vec![("a", array)]).unwrap();
3684 let sliced_batch = orig_batch.slice(1, num_rows - 1);
3686
3687 let schema = orig_batch.schema();
3688 let stream_data = {
3689 let mut writer = StreamWriter::try_new(vec![], &schema).unwrap();
3690 writer.write(&sliced_batch).unwrap();
3691 writer.into_inner().unwrap()
3692 };
3693 let read_batch = {
3694 let projection = None;
3695 let mut reader = StreamReader::try_new(Cursor::new(stream_data), projection).unwrap();
3696 reader
3697 .next()
3698 .expect("expect no errors reading batch")
3699 .expect("expect batch")
3700 };
3701 assert_eq!(sliced_batch, read_batch);
3702
3703 let file_data = {
3704 let mut writer = FileWriter::try_new_buffered(vec![], &schema).unwrap();
3705 writer.write(&sliced_batch).unwrap();
3706 writer.into_inner().unwrap().into_inner().unwrap()
3707 };
3708 let read_batch = {
3709 let projection = None;
3710 let mut reader = FileReader::try_new(Cursor::new(file_data), projection).unwrap();
3711 reader
3712 .next()
3713 .expect("expect no errors reading batch")
3714 .expect("expect batch")
3715 };
3716 assert_eq!(sliced_batch, read_batch);
3717
3718 }
3720
3721 #[test]
3722 fn encode_bools_slice() {
3723 assert_bool_roundtrip([true, false], 1, 1);
3725
3726 assert_bool_roundtrip(
3728 [
3729 true, false, true, true, false, false, true, true, true, false, false, false, true,
3730 true, true, true, false, false, false, false, true, true, true, true, true, false,
3731 false, false, false, false,
3732 ],
3733 13,
3734 17,
3735 );
3736
3737 assert_bool_roundtrip(
3739 [
3740 true, false, true, true, false, false, true, true, true, false, false, false,
3741 ],
3742 8,
3743 2,
3744 );
3745
3746 assert_bool_roundtrip(
3748 [
3749 true, false, true, true, false, false, true, true, true, false, false, false, true,
3750 true, true, true, true, false, false, false, false, false,
3751 ],
3752 8,
3753 8,
3754 );
3755 }
3756
3757 fn assert_bool_roundtrip<const N: usize>(bools: [bool; N], offset: usize, length: usize) {
3758 let val_bool_field = Field::new("val", DataType::Boolean, false);
3759
3760 let schema = Arc::new(Schema::new(vec![val_bool_field]));
3761
3762 let bools = BooleanArray::from(bools.to_vec());
3763
3764 let batch = RecordBatch::try_new(Arc::clone(&schema), vec![Arc::new(bools)]).unwrap();
3765 let batch = batch.slice(offset, length);
3766
3767 let data = serialize_stream(&batch);
3768 let batch2 = deserialize_stream(data);
3769 assert_eq!(batch, batch2);
3770 }
3771
3772 fn assert_empty_run_array_roundtrip<R: RunEndIndexType>() {
3773 let run_ends = PrimitiveArray::<R>::from_iter_values(
3774 [2, 5]
3775 .into_iter()
3776 .map(|v| R::Native::from_usize(v).unwrap()),
3777 );
3778 let values = Int32Array::from(vec![10, 20]);
3779 let array = RunArray::<R>::try_new(&run_ends, &values).unwrap();
3780 let empty = RunArray::<R>::from(ArrayData::new_empty(array.data_type()));
3781
3782 for source in [&array, &empty] {
3783 for offset in 0..=source.len() {
3784 let sliced = source.slice(offset, 0);
3785 let batch = RecordBatch::try_from_iter(vec![("run", Arc::new(sliced) as ArrayRef)])
3786 .unwrap();
3787 for decoded in [
3788 deserialize_stream(serialize_stream(&batch)),
3789 deserialize_file(serialize_file(&batch)),
3790 ] {
3791 assert_eq!(decoded, batch);
3792 let data = decoded.column(0).to_data();
3793 data.validate_full().unwrap();
3794 assert_eq!(data.offset(), 0);
3795 assert!(data.child_data().iter().all(ArrayData::is_empty));
3796 }
3797
3798 let normalized = into_zero_offset_run_array(source.slice(offset, 0)).unwrap();
3799 normalized.to_data().validate_full().unwrap();
3800 assert_eq!(normalized.offset(), 0);
3801 assert!(normalized.run_ends().values().is_empty());
3802 assert!(normalized.values().is_empty());
3803 }
3804 }
3805 }
3806
3807 #[test]
3808 fn test_empty_run_array_roundtrip_i16() {
3809 assert_empty_run_array_roundtrip::<Int16Type>();
3810 }
3811
3812 #[test]
3813 fn test_empty_run_array_roundtrip_i32() {
3814 assert_empty_run_array_roundtrip::<Int32Type>();
3815 }
3816
3817 #[test]
3818 fn test_empty_run_array_roundtrip_i64() {
3819 assert_empty_run_array_roundtrip::<Int64Type>();
3820 }
3821
3822 #[test]
3823 fn test_run_array_unslice() {
3824 let total_len = 80;
3825 let vals: Vec<Option<i32>> = vec![Some(1), None, Some(2), Some(3), Some(4), None, Some(5)];
3826 let repeats: Vec<usize> = vec![3, 4, 1, 2];
3827 let mut input_array: Vec<Option<i32>> = Vec::with_capacity(total_len);
3828 for ix in 0_usize..32 {
3829 let repeat: usize = repeats[ix % repeats.len()];
3830 let val: Option<i32> = vals[ix % vals.len()];
3831 input_array.resize(input_array.len() + repeat, val);
3832 }
3833
3834 let mut builder =
3836 PrimitiveRunBuilder::<Int16Type, Int32Type>::with_capacity(input_array.len());
3837 builder.extend(input_array.iter().copied());
3838 let run_array = builder.finish();
3839
3840 for slice_len in 1..=total_len {
3842 let sliced_run_array: RunArray<Int16Type> =
3844 run_array.slice(0, slice_len).into_data().into();
3845
3846 let unsliced_run_array = into_zero_offset_run_array(sliced_run_array).unwrap();
3848 let typed = unsliced_run_array
3849 .downcast::<PrimitiveArray<Int32Type>>()
3850 .unwrap();
3851 let expected: Vec<Option<i32>> = input_array.iter().take(slice_len).copied().collect();
3852 let actual: Vec<Option<i32>> = typed.into_iter().collect();
3853 assert_eq!(expected, actual);
3854
3855 let sliced_run_array: RunArray<Int16Type> = run_array
3857 .slice(total_len - slice_len, slice_len)
3858 .into_data()
3859 .into();
3860
3861 let unsliced_run_array = into_zero_offset_run_array(sliced_run_array).unwrap();
3863 let typed = unsliced_run_array
3864 .downcast::<PrimitiveArray<Int32Type>>()
3865 .unwrap();
3866 let expected: Vec<Option<i32>> = input_array
3867 .iter()
3868 .skip(total_len - slice_len)
3869 .copied()
3870 .collect();
3871 let actual: Vec<Option<i32>> = typed.into_iter().collect();
3872 assert_eq!(expected, actual);
3873 }
3874 }
3875
3876 fn generate_list_data<O: OffsetSizeTrait>() -> GenericListArray<O> {
3877 let mut ls = GenericListBuilder::<O, _>::new(UInt32Builder::new());
3878
3879 for i in 0..100_000 {
3880 for value in [i, i, i] {
3881 ls.values().append_value(value);
3882 }
3883 ls.append(true)
3884 }
3885
3886 ls.finish()
3887 }
3888
3889 fn generate_utf8view_list_data<O: OffsetSizeTrait>() -> GenericListArray<O> {
3890 let mut ls = GenericListBuilder::<O, _>::new(StringViewBuilder::new());
3891
3892 for i in 0..100_000 {
3893 for value in [
3894 format!("value{i}"),
3895 format!("value{i}"),
3896 format!("value{i}"),
3897 ] {
3898 ls.values().append_value(&value);
3899 }
3900 ls.append(true)
3901 }
3902
3903 ls.finish()
3904 }
3905
3906 fn generate_string_list_data<O: OffsetSizeTrait>() -> GenericListArray<O> {
3907 let mut ls = GenericListBuilder::<O, _>::new(StringBuilder::new());
3908
3909 for i in 0..100_000 {
3910 for value in [
3911 format!("value{i}"),
3912 format!("value{i}"),
3913 format!("value{i}"),
3914 ] {
3915 ls.values().append_value(&value);
3916 }
3917 ls.append(true)
3918 }
3919
3920 ls.finish()
3921 }
3922
3923 fn generate_nested_list_data<O: OffsetSizeTrait>() -> GenericListArray<O> {
3924 let mut ls =
3925 GenericListBuilder::<O, _>::new(GenericListBuilder::<O, _>::new(UInt32Builder::new()));
3926
3927 for _i in 0..10_000 {
3928 for j in 0..10 {
3929 for value in [j, j, j, j] {
3930 ls.values().values().append_value(value);
3931 }
3932 ls.values().append(true)
3933 }
3934 ls.append(true);
3935 }
3936
3937 ls.finish()
3938 }
3939
3940 fn generate_nested_list_data_starting_at_zero<O: OffsetSizeTrait>() -> GenericListArray<O> {
3941 let mut ls =
3942 GenericListBuilder::<O, _>::new(GenericListBuilder::<O, _>::new(UInt32Builder::new()));
3943
3944 for _i in 0..999 {
3945 ls.values().append(true);
3946 ls.append(true);
3947 }
3948
3949 for j in 0..10 {
3950 for value in [j, j, j, j] {
3951 ls.values().values().append_value(value);
3952 }
3953 ls.values().append(true)
3954 }
3955 ls.append(true);
3956
3957 for i in 0..9_000 {
3958 for j in 0..10 {
3959 for value in [i + j, i + j, i + j, i + j] {
3960 ls.values().values().append_value(value);
3961 }
3962 ls.values().append(true)
3963 }
3964 ls.append(true);
3965 }
3966
3967 ls.finish()
3968 }
3969
3970 fn generate_map_array_data() -> MapArray {
3971 let keys_builder = UInt32Builder::new();
3972 let values_builder = UInt32Builder::new();
3973
3974 let mut builder = MapBuilder::new(None, keys_builder, values_builder);
3975
3976 for i in 0..100_000 {
3977 for _j in 0..3 {
3978 builder.keys().append_value(i);
3979 builder.values().append_value(i * 2);
3980 }
3981 builder.append(true).unwrap();
3982 }
3983
3984 builder.finish()
3985 }
3986
3987 #[test]
3988 #[cfg_attr(miri, ignore)] fn reencode_offsets_when_first_offset_is_not_zero() {
3990 let original_list = generate_list_data::<i32>();
3991 let original_data = original_list.into_data();
3992 let slice_data = original_data.slice(75, 7);
3993 let (new_offsets, original_start, length) =
3994 reencode_offsets::<i32>(&slice_data.buffers()[0], &slice_data);
3995 assert_eq!(
3996 vec![0, 3, 6, 9, 12, 15, 18, 21],
3997 new_offsets.typed_data::<i32>()
3998 );
3999 assert_eq!(225, original_start);
4000 assert_eq!(21, length);
4001 }
4002
4003 #[test]
4004 fn reencode_offsets_when_first_offset_is_zero() {
4005 let mut ls = GenericListBuilder::<i32, _>::new(UInt32Builder::new());
4006 ls.append(true);
4008 ls.values().append_value(35);
4009 ls.values().append_value(42);
4010 ls.append(true);
4011 let original_list = ls.finish();
4012 let original_data = original_list.into_data();
4013
4014 let slice_data = original_data.slice(1, 1);
4015 let (new_offsets, original_start, length) =
4016 reencode_offsets::<i32>(&slice_data.buffers()[0], &slice_data);
4017 assert_eq!(vec![0, 2], new_offsets.typed_data::<i32>());
4018 assert_eq!(0, original_start);
4019 assert_eq!(2, length);
4020 }
4021
4022 fn roundtrip_ensure_sliced_smaller(in_batch: RecordBatch, expected_size_factor: usize) {
4025 let in_sliced = in_batch.slice(999, 1);
4027
4028 let bytes_batch = serialize_file(&in_batch);
4029 let bytes_sliced = serialize_file(&in_sliced);
4030
4031 assert!(bytes_sliced.len() < (bytes_batch.len() / expected_size_factor));
4033
4034 let out_batch = deserialize_file(bytes_batch);
4036 assert_eq!(in_batch, out_batch);
4037
4038 let out_sliced = deserialize_file(bytes_sliced);
4039 assert_eq!(in_sliced, out_sliced);
4040 }
4041
4042 #[test]
4043 #[cfg_attr(miri, ignore)] fn encode_lists() {
4045 let val_inner = Field::new_list_field(DataType::UInt32, true);
4046 let val_list_field = Field::new("val", DataType::List(Arc::new(val_inner)), false);
4047 let schema = Arc::new(Schema::new(vec![val_list_field]));
4048
4049 let values = Arc::new(generate_list_data::<i32>());
4050
4051 let in_batch = RecordBatch::try_new(schema, vec![values]).unwrap();
4052 roundtrip_ensure_sliced_smaller(in_batch, 1000);
4053 }
4054
4055 #[test]
4056 #[cfg_attr(miri, ignore)] fn encode_empty_list() {
4058 let val_inner = Field::new_list_field(DataType::UInt32, true);
4059 let val_list_field = Field::new("val", DataType::List(Arc::new(val_inner)), false);
4060 let schema = Arc::new(Schema::new(vec![val_list_field]));
4061
4062 let values = Arc::new(generate_list_data::<i32>());
4063
4064 let in_batch = RecordBatch::try_new(schema, vec![values])
4065 .unwrap()
4066 .slice(999, 0);
4067 let out_batch = deserialize_file(serialize_file(&in_batch));
4068 assert_eq!(in_batch, out_batch);
4069 }
4070
4071 #[test]
4072 #[cfg_attr(miri, ignore)] fn encode_large_lists() {
4074 let val_inner = Field::new_list_field(DataType::UInt32, true);
4075 let val_list_field = Field::new("val", DataType::LargeList(Arc::new(val_inner)), false);
4076 let schema = Arc::new(Schema::new(vec![val_list_field]));
4077
4078 let values = Arc::new(generate_list_data::<i64>());
4079
4080 let in_batch = RecordBatch::try_new(schema, vec![values]).unwrap();
4083 roundtrip_ensure_sliced_smaller(in_batch, 1000);
4084 }
4085
4086 #[test]
4087 #[cfg_attr(miri, ignore)] fn encode_large_lists_non_zero_offset() {
4089 let val_inner = Field::new_list_field(DataType::UInt32, true);
4090 let val_list_field = Field::new("val", DataType::LargeList(Arc::new(val_inner)), false);
4091 let schema = Arc::new(Schema::new(vec![val_list_field]));
4092
4093 let values = Arc::new(generate_list_data::<i64>());
4094
4095 check_sliced_list_array(schema, values);
4096 }
4097
4098 #[test]
4099 #[cfg_attr(miri, ignore)] fn encode_large_lists_string_non_zero_offset() {
4101 let val_inner = Field::new_list_field(DataType::Utf8, true);
4102 let val_list_field = Field::new("val", DataType::LargeList(Arc::new(val_inner)), false);
4103 let schema = Arc::new(Schema::new(vec![val_list_field]));
4104
4105 let values = Arc::new(generate_string_list_data::<i64>());
4106
4107 check_sliced_list_array(schema, values);
4108 }
4109
4110 #[test]
4111 #[cfg_attr(miri, ignore)] fn encode_large_list_string_view_non_zero_offset() {
4113 let val_inner = Field::new_list_field(DataType::Utf8View, true);
4114 let val_list_field = Field::new("val", DataType::LargeList(Arc::new(val_inner)), false);
4115 let schema = Arc::new(Schema::new(vec![val_list_field]));
4116
4117 let values = Arc::new(generate_utf8view_list_data::<i64>());
4118
4119 check_sliced_list_array(schema, values);
4120 }
4121
4122 fn check_sliced_list_array(schema: Arc<Schema>, values: Arc<GenericListArray<i64>>) {
4123 for (offset, len) in [(999, 1), (0, 13), (47, 12), (values.len() - 13, 13)] {
4124 let in_batch = RecordBatch::try_new(schema.clone(), vec![values.clone()])
4125 .unwrap()
4126 .slice(offset, len);
4127 let out_batch = deserialize_file(serialize_file(&in_batch));
4128 assert_eq!(in_batch, out_batch);
4129 }
4130 }
4131
4132 #[test]
4133 #[cfg_attr(miri, ignore)] fn encode_nested_lists() {
4135 let inner_int = Arc::new(Field::new_list_field(DataType::UInt32, true));
4136 let inner_list_field = Arc::new(Field::new_list_field(DataType::List(inner_int), true));
4137 let list_field = Field::new("val", DataType::List(inner_list_field), true);
4138 let schema = Arc::new(Schema::new(vec![list_field]));
4139
4140 let values = Arc::new(generate_nested_list_data::<i32>());
4141
4142 let in_batch = RecordBatch::try_new(schema, vec![values]).unwrap();
4143 roundtrip_ensure_sliced_smaller(in_batch, 1000);
4144 }
4145
4146 #[test]
4147 #[cfg_attr(miri, ignore)] fn encode_nested_lists_starting_at_zero() {
4149 let inner_int = Arc::new(Field::new("item", DataType::UInt32, true));
4150 let inner_list_field = Arc::new(Field::new("item", DataType::List(inner_int), true));
4151 let list_field = Field::new("val", DataType::List(inner_list_field), true);
4152 let schema = Arc::new(Schema::new(vec![list_field]));
4153
4154 let values = Arc::new(generate_nested_list_data_starting_at_zero::<i32>());
4155
4156 let in_batch = RecordBatch::try_new(schema, vec![values]).unwrap();
4157 roundtrip_ensure_sliced_smaller(in_batch, 1);
4158 }
4159
4160 #[test]
4161 #[cfg_attr(miri, ignore)] fn encode_map_array() {
4163 let keys = Arc::new(Field::new(
4164 Field::MAP_KEY_FIELD_DEFAULT_NAME,
4165 DataType::UInt32,
4166 false,
4167 ));
4168 let values = Arc::new(Field::new(
4169 Field::MAP_VALUE_FIELD_DEFAULT_NAME,
4170 DataType::UInt32,
4171 true,
4172 ));
4173 let map_field = Field::new_map(
4174 "map",
4175 Field::MAP_ENTRIES_FIELD_DEFAULT_NAME,
4176 keys,
4177 values,
4178 false,
4179 true,
4180 );
4181 let schema = Arc::new(Schema::new(vec![map_field]));
4182
4183 let values = Arc::new(generate_map_array_data());
4184
4185 let in_batch = RecordBatch::try_new(schema, vec![values]).unwrap();
4186 roundtrip_ensure_sliced_smaller(in_batch, 1000);
4187 }
4188
4189 fn generate_list_view_data<O: OffsetSizeTrait>() -> GenericListViewArray<O> {
4190 let mut builder = GenericListViewBuilder::<O, _>::new(UInt32Builder::new());
4191
4192 for i in 0u32..100_000 {
4193 if i.is_multiple_of(10_000) {
4194 builder.append(false);
4195 continue;
4196 }
4197 for value in [i, i, i] {
4198 builder.values().append_value(value);
4199 }
4200 builder.append(true);
4201 }
4202
4203 builder.finish()
4204 }
4205
4206 #[test]
4207 #[cfg_attr(miri, ignore)] fn encode_list_view_arrays() {
4209 let val_inner = Field::new_list_field(DataType::UInt32, true);
4210 let val_field = Field::new("val", DataType::ListView(Arc::new(val_inner)), true);
4211 let schema = Arc::new(Schema::new(vec![val_field]));
4212
4213 let values = Arc::new(generate_list_view_data::<i32>());
4214
4215 let in_batch = RecordBatch::try_new(schema, vec![values]).unwrap();
4216 let out_batch = deserialize_file(serialize_file(&in_batch));
4217 assert_eq!(in_batch, out_batch);
4218 }
4219
4220 #[test]
4221 #[cfg_attr(miri, ignore)] fn encode_large_list_view_arrays() {
4223 let val_inner = Field::new_list_field(DataType::UInt32, true);
4224 let val_field = Field::new("val", DataType::LargeListView(Arc::new(val_inner)), true);
4225 let schema = Arc::new(Schema::new(vec![val_field]));
4226
4227 let values = Arc::new(generate_list_view_data::<i64>());
4228
4229 let in_batch = RecordBatch::try_new(schema, vec![values]).unwrap();
4230 let out_batch = deserialize_file(serialize_file(&in_batch));
4231 assert_eq!(in_batch, out_batch);
4232 }
4233
4234 #[test]
4235 #[cfg_attr(miri, ignore)] fn check_sliced_list_view_array() {
4237 let inner = Field::new_list_field(DataType::UInt32, true);
4238 let field = Field::new("val", DataType::ListView(Arc::new(inner)), true);
4239 let schema = Arc::new(Schema::new(vec![field]));
4240 let values = Arc::new(generate_list_view_data::<i32>());
4241
4242 for (offset, len) in [(999, 1), (0, 13), (47, 12), (values.len() - 13, 13)] {
4243 let in_batch = RecordBatch::try_new(schema.clone(), vec![values.clone()])
4244 .unwrap()
4245 .slice(offset, len);
4246 let out_batch = deserialize_file(serialize_file(&in_batch));
4247 assert_eq!(in_batch, out_batch);
4248 }
4249 }
4250
4251 #[test]
4252 #[cfg_attr(miri, ignore)] fn check_sliced_large_list_view_array() {
4254 let inner = Field::new_list_field(DataType::UInt32, true);
4255 let field = Field::new("val", DataType::LargeListView(Arc::new(inner)), true);
4256 let schema = Arc::new(Schema::new(vec![field]));
4257 let values = Arc::new(generate_list_view_data::<i64>());
4258
4259 for (offset, len) in [(999, 1), (0, 13), (47, 12), (values.len() - 13, 13)] {
4260 let in_batch = RecordBatch::try_new(schema.clone(), vec![values.clone()])
4261 .unwrap()
4262 .slice(offset, len);
4263 let out_batch = deserialize_file(serialize_file(&in_batch));
4264 assert_eq!(in_batch, out_batch);
4265 }
4266 }
4267
4268 fn generate_nested_list_view_data<O: OffsetSizeTrait>() -> GenericListViewArray<O> {
4269 let inner_builder = UInt32Builder::new();
4270 let middle_builder = GenericListViewBuilder::<O, _>::new(inner_builder);
4271 let mut outer_builder = GenericListViewBuilder::<O, _>::new(middle_builder);
4272
4273 for i in 0u32..10_000 {
4274 if i.is_multiple_of(1_000) {
4275 outer_builder.append(false);
4276 continue;
4277 }
4278
4279 for _ in 0..3 {
4280 for value in [i, i + 1, i + 2] {
4281 outer_builder.values().values().append_value(value);
4282 }
4283 outer_builder.values().append(true);
4284 }
4285 outer_builder.append(true);
4286 }
4287
4288 outer_builder.finish()
4289 }
4290
4291 #[test]
4292 #[cfg_attr(miri, ignore)] fn encode_nested_list_views() {
4294 let inner_int = Arc::new(Field::new_list_field(DataType::UInt32, true));
4295 let inner_list_field = Arc::new(Field::new_list_field(DataType::ListView(inner_int), true));
4296 let list_field = Field::new("val", DataType::ListView(inner_list_field), true);
4297 let schema = Arc::new(Schema::new(vec![list_field]));
4298
4299 let values = Arc::new(generate_nested_list_view_data::<i32>());
4300
4301 let in_batch = RecordBatch::try_new(schema, vec![values]).unwrap();
4302 let out_batch = deserialize_file(serialize_file(&in_batch));
4303 assert_eq!(in_batch, out_batch);
4304 }
4305
4306 fn test_roundtrip_list_view_of_dict_impl<OffsetSize: OffsetSizeTrait, U: ArrowNativeType>(
4307 list_data_type: DataType,
4308 offsets: &[U; 5],
4309 sizes: &[U; 4],
4310 ) {
4311 let values = StringArray::from(vec![Some("alpha"), None, Some("beta"), Some("gamma")]);
4312 let keys = Int32Array::from_iter_values([0, 0, 1, 2, 3, 0, 2]);
4313 let dict_array = DictionaryArray::new(keys, Arc::new(values));
4314 let dict_data = dict_array.to_data();
4315
4316 let value_offsets = Buffer::from_slice_ref(offsets);
4317 let value_sizes = Buffer::from_slice_ref(sizes);
4318
4319 let list_data = ArrayData::builder(list_data_type)
4320 .len(4)
4321 .add_buffer(value_offsets)
4322 .add_buffer(value_sizes)
4323 .add_child_data(dict_data)
4324 .build()
4325 .unwrap();
4326 let list_view_array = GenericListViewArray::<OffsetSize>::from(list_data);
4327
4328 let schema = Arc::new(Schema::new(vec![Field::new(
4329 "f1",
4330 list_view_array.data_type().clone(),
4331 false,
4332 )]));
4333 let input_batch = RecordBatch::try_new(schema, vec![Arc::new(list_view_array)]).unwrap();
4334
4335 let output_batch = deserialize_file(serialize_file(&input_batch));
4336 assert_eq!(input_batch, output_batch);
4337
4338 let output_batch = deserialize_stream(serialize_stream(&input_batch));
4339 assert_eq!(input_batch, output_batch);
4340 }
4341
4342 #[test]
4343 fn test_roundtrip_list_view_of_dict() {
4344 #[expect(deprecated)]
4345 let list_data_type = DataType::ListView(Arc::new(Field::new_dict(
4346 "item",
4347 DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
4348 true,
4349 1,
4350 false,
4351 )));
4352 let offsets: &[i32; 5] = &[0, 2, 4, 4, 7];
4353 let sizes: &[i32; 4] = &[2, 2, 0, 3];
4354 test_roundtrip_list_view_of_dict_impl::<i32, i32>(list_data_type, offsets, sizes);
4355 }
4356
4357 #[test]
4358 fn test_roundtrip_large_list_view_of_dict() {
4359 #[expect(deprecated)]
4360 let list_data_type = DataType::LargeListView(Arc::new(Field::new_dict(
4361 "item",
4362 DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
4363 true,
4364 2,
4365 false,
4366 )));
4367 let offsets: &[i64; 5] = &[0, 2, 4, 4, 7];
4368 let sizes: &[i64; 4] = &[2, 2, 0, 3];
4369 test_roundtrip_list_view_of_dict_impl::<i64, i64>(list_data_type, offsets, sizes);
4370 }
4371
4372 #[test]
4373 fn test_roundtrip_sliced_list_view_of_dict() {
4374 #[expect(deprecated)]
4375 let list_data_type = DataType::ListView(Arc::new(Field::new_dict(
4376 "item",
4377 DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
4378 true,
4379 3,
4380 false,
4381 )));
4382
4383 let values = StringArray::from(vec![Some("alpha"), None, Some("beta"), Some("gamma")]);
4384 let keys = Int32Array::from_iter_values([0, 0, 1, 2, 3, 0, 2, 1, 0, 3, 2, 1]);
4385 let dict_array = DictionaryArray::new(keys, Arc::new(values));
4386 let dict_data = dict_array.to_data();
4387
4388 let offsets: &[i32; 7] = &[0, 2, 4, 4, 7, 9, 12];
4389 let sizes: &[i32; 6] = &[2, 2, 0, 3, 2, 3];
4390 let value_offsets = Buffer::from_slice_ref(offsets);
4391 let value_sizes = Buffer::from_slice_ref(sizes);
4392
4393 let list_data = ArrayData::builder(list_data_type)
4394 .len(6)
4395 .add_buffer(value_offsets)
4396 .add_buffer(value_sizes)
4397 .add_child_data(dict_data)
4398 .build()
4399 .unwrap();
4400 let list_view_array = GenericListViewArray::<i32>::from(list_data);
4401
4402 let schema = Arc::new(Schema::new(vec![Field::new(
4403 "f1",
4404 list_view_array.data_type().clone(),
4405 false,
4406 )]));
4407 let input_batch = RecordBatch::try_new(schema, vec![Arc::new(list_view_array)]).unwrap();
4408
4409 let sliced_batch = input_batch.slice(1, 4);
4410
4411 let output_batch = deserialize_file(serialize_file(&sliced_batch));
4412 assert_eq!(sliced_batch, output_batch);
4413
4414 let output_batch = deserialize_stream(serialize_stream(&sliced_batch));
4415 assert_eq!(sliced_batch, output_batch);
4416 }
4417
4418 #[test]
4419 fn test_roundtrip_dense_union_of_dict() {
4420 let values = StringArray::from(vec![Some("alpha"), None, Some("beta"), Some("gamma")]);
4421 let keys = Int32Array::from_iter_values([0, 0, 1, 2, 3, 0, 2]);
4422 let dict_array = DictionaryArray::new(keys, Arc::new(values));
4423
4424 #[expect(deprecated)]
4425 let dict_field = Arc::new(Field::new_dict(
4426 "dict",
4427 DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
4428 true,
4429 1,
4430 false,
4431 ));
4432 let int_field = Arc::new(Field::new("int", DataType::Int32, false));
4433 let union_fields = UnionFields::try_new(vec![0, 1], vec![dict_field, int_field]).unwrap();
4434
4435 let types = ScalarBuffer::from(vec![0i8, 0, 1, 0, 1, 0, 0]);
4436 let offsets = ScalarBuffer::from(vec![0i32, 1, 0, 2, 1, 3, 4]);
4437
4438 let int_array = Int32Array::from(vec![100, 200]);
4439
4440 let union = UnionArray::try_new(
4441 union_fields.clone(),
4442 types,
4443 Some(offsets),
4444 vec![Arc::new(dict_array), Arc::new(int_array)],
4445 )
4446 .unwrap();
4447
4448 let schema = Arc::new(Schema::new(vec![Field::new(
4449 "union",
4450 DataType::Union(union_fields, UnionMode::Dense),
4451 false,
4452 )]));
4453 let input_batch = RecordBatch::try_new(schema, vec![Arc::new(union)]).unwrap();
4454
4455 let output_batch = deserialize_file(serialize_file(&input_batch));
4456 assert_eq!(input_batch, output_batch);
4457
4458 let output_batch = deserialize_stream(serialize_stream(&input_batch));
4459 assert_eq!(input_batch, output_batch);
4460 }
4461
4462 #[test]
4463 fn test_roundtrip_sparse_union_of_dict() {
4464 let values = StringArray::from(vec![Some("alpha"), None, Some("beta"), Some("gamma")]);
4465 let keys = Int32Array::from_iter_values([0, 0, 1, 2, 3, 0, 2]);
4466 let dict_array = DictionaryArray::new(keys, Arc::new(values));
4467
4468 #[expect(deprecated)]
4469 let dict_field = Arc::new(Field::new_dict(
4470 "dict",
4471 DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
4472 true,
4473 2,
4474 false,
4475 ));
4476 let int_field = Arc::new(Field::new("int", DataType::Int32, false));
4477 let union_fields = UnionFields::try_new(vec![0, 1], vec![dict_field, int_field]).unwrap();
4478
4479 let types = ScalarBuffer::from(vec![0i8, 0, 1, 0, 1, 0, 0]);
4480
4481 let int_array = Int32Array::from(vec![0, 0, 100, 0, 200, 0, 0]);
4482
4483 let union = UnionArray::try_new(
4484 union_fields.clone(),
4485 types,
4486 None,
4487 vec![Arc::new(dict_array), Arc::new(int_array)],
4488 )
4489 .unwrap();
4490
4491 let schema = Arc::new(Schema::new(vec![Field::new(
4492 "union",
4493 DataType::Union(union_fields, UnionMode::Sparse),
4494 false,
4495 )]));
4496 let input_batch = RecordBatch::try_new(schema, vec![Arc::new(union)]).unwrap();
4497
4498 let output_batch = deserialize_file(serialize_file(&input_batch));
4499 assert_eq!(input_batch, output_batch);
4500
4501 let output_batch = deserialize_stream(serialize_stream(&input_batch));
4502 assert_eq!(input_batch, output_batch);
4503 }
4504
4505 #[test]
4506 fn test_roundtrip_map_with_dict_keys() {
4507 let key_values = StringArray::from(vec!["key_a", "key_b", "key_c"]);
4510 let keys = Int32Array::from_iter_values([0, 1, 2, 0, 1, 0]);
4511 let dict_keys = DictionaryArray::new(keys, Arc::new(key_values));
4512
4513 let values = Int32Array::from(vec![1, 2, 3, 4, 5, 6]);
4514
4515 #[expect(deprecated)]
4516 let entries_field = Arc::new(Field::new(
4517 Field::MAP_ENTRIES_FIELD_DEFAULT_NAME,
4518 DataType::Struct(
4519 vec![
4520 Field::new_dict(
4521 Field::MAP_KEY_FIELD_DEFAULT_NAME,
4522 DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
4523 false,
4524 1,
4525 false,
4526 ),
4527 Field::new(Field::MAP_VALUE_FIELD_DEFAULT_NAME, DataType::Int32, true),
4528 ]
4529 .into(),
4530 ),
4531 false,
4532 ));
4533
4534 let entries = StructArray::from(vec![
4535 (
4536 Arc::new(Field::new(
4537 Field::MAP_KEY_FIELD_DEFAULT_NAME,
4538 DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
4539 false,
4540 )),
4541 Arc::new(dict_keys) as ArrayRef,
4542 ),
4543 (
4544 Arc::new(Field::new(
4545 Field::MAP_VALUE_FIELD_DEFAULT_NAME,
4546 DataType::Int32,
4547 true,
4548 )),
4549 Arc::new(values) as ArrayRef,
4550 ),
4551 ]);
4552
4553 let offsets = Buffer::from_slice_ref([0i32, 2, 4, 6]);
4554
4555 let map_data = ArrayData::builder(DataType::Map(entries_field, false))
4556 .len(3)
4557 .add_buffer(offsets)
4558 .add_child_data(entries.into_data())
4559 .build()
4560 .unwrap();
4561 let map_array = MapArray::from(map_data);
4562
4563 let schema = Arc::new(Schema::new(vec![Field::new(
4564 "map",
4565 map_array.data_type().clone(),
4566 false,
4567 )]));
4568 let input_batch = RecordBatch::try_new(schema, vec![Arc::new(map_array)]).unwrap();
4569
4570 let output_batch = deserialize_file(serialize_file(&input_batch));
4571 assert_eq!(input_batch, output_batch);
4572
4573 let output_batch = deserialize_stream(serialize_stream(&input_batch));
4574 assert_eq!(input_batch, output_batch);
4575 }
4576
4577 #[test]
4578 fn test_roundtrip_map_with_dict_values() {
4579 let keys = StringArray::from(vec!["a", "b", "c", "d", "e", "f"]);
4582
4583 let value_values = StringArray::from(vec!["val_x", "val_y", "val_z"]);
4584 let value_keys = Int32Array::from_iter_values([0, 1, 2, 0, 1, 0]);
4585 let dict_values = DictionaryArray::new(value_keys, Arc::new(value_values));
4586
4587 #[expect(deprecated)]
4588 let entries_field = Arc::new(Field::new(
4589 Field::MAP_ENTRIES_FIELD_DEFAULT_NAME,
4590 DataType::Struct(
4591 vec![
4592 Field::new(Field::MAP_KEY_FIELD_DEFAULT_NAME, DataType::Utf8, false),
4593 Field::new_dict(
4594 Field::MAP_VALUE_FIELD_DEFAULT_NAME,
4595 DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
4596 true,
4597 2,
4598 false,
4599 ),
4600 ]
4601 .into(),
4602 ),
4603 false,
4604 ));
4605
4606 let entries = StructArray::from(vec![
4607 (
4608 Arc::new(Field::new(
4609 Field::MAP_KEY_FIELD_DEFAULT_NAME,
4610 DataType::Utf8,
4611 false,
4612 )),
4613 Arc::new(keys) as ArrayRef,
4614 ),
4615 (
4616 Arc::new(Field::new(
4617 Field::MAP_VALUE_FIELD_DEFAULT_NAME,
4618 DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
4619 true,
4620 )),
4621 Arc::new(dict_values) as ArrayRef,
4622 ),
4623 ]);
4624
4625 let offsets = Buffer::from_slice_ref([0i32, 2, 4, 6]);
4626
4627 let map_data = ArrayData::builder(DataType::Map(entries_field, false))
4628 .len(3)
4629 .add_buffer(offsets)
4630 .add_child_data(entries.into_data())
4631 .build()
4632 .unwrap();
4633 let map_array = MapArray::from(map_data);
4634
4635 let schema = Arc::new(Schema::new(vec![Field::new(
4636 "map",
4637 map_array.data_type().clone(),
4638 false,
4639 )]));
4640 let input_batch = RecordBatch::try_new(schema, vec![Arc::new(map_array)]).unwrap();
4641
4642 let output_batch = deserialize_file(serialize_file(&input_batch));
4643 assert_eq!(input_batch, output_batch);
4644
4645 let output_batch = deserialize_stream(serialize_stream(&input_batch));
4646 assert_eq!(input_batch, output_batch);
4647 }
4648
4649 #[test]
4650 fn test_decimal128_alignment16_is_sufficient() {
4651 const IPC_ALIGNMENT: usize = 16;
4652
4653 for num_cols in [1, 2, 3, 17, 50, 73, 99] {
4658 let num_rows = (num_cols * 7 + 11) % 100; 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 =
4685 read_footer_length(buffer[trailer_start..].try_into().unwrap()).unwrap();
4686 let footer =
4687 root_as_footer(&buffer[trailer_start - footer_len..trailer_start]).unwrap();
4688
4689 let schema = try_fb_to_schema(footer.schema().unwrap()).unwrap();
4690
4691 let decoder =
4694 FileDecoder::new(Arc::new(schema), footer.version()).with_require_alignment(true);
4695
4696 let batches = footer.recordBatches().unwrap();
4697
4698 let block = batches.get(0);
4699 let block_len = block.bodyLength() as usize + block.metaDataLength() as usize;
4700 let data = buffer.slice_with_length(block.offset() as _, block_len);
4701
4702 let batch2 = decoder.read_record_batch(block, &data).unwrap().unwrap();
4703
4704 assert_eq!(batch, batch2);
4705 }
4706 }
4707
4708 #[test]
4709 fn test_decimal128_alignment8_is_unaligned() {
4710 const IPC_ALIGNMENT: usize = 8;
4711
4712 let num_cols = 2;
4713 let num_rows = 1;
4714
4715 let mut fields = Vec::new();
4716 let mut arrays = Vec::new();
4717 for i in 0..num_cols {
4718 let field = Field::new(format!("col_{i}"), DataType::Decimal128(38, 10), true);
4719 let array = Decimal128Array::from(vec![num_cols as i128; num_rows]);
4720 fields.push(field);
4721 arrays.push(Arc::new(array) as Arc<dyn Array>);
4722 }
4723 let schema = Schema::new(fields);
4724 let batch = RecordBatch::try_new(Arc::new(schema), arrays).unwrap();
4725
4726 let mut writer = FileWriter::try_new_with_options(
4727 Vec::new(),
4728 batch.schema_ref(),
4729 IpcWriteOptions::try_new(IPC_ALIGNMENT, false, MetadataVersion::V5).unwrap(),
4730 )
4731 .unwrap();
4732 writer.write(&batch).unwrap();
4733 writer.finish().unwrap();
4734
4735 let out: Vec<u8> = writer.into_inner().unwrap();
4736
4737 let buffer = Buffer::from_slice_ref(out);
4738 let trailer_start = buffer.len() - 10;
4739 let footer_len = read_footer_length(buffer[trailer_start..].try_into().unwrap()).unwrap();
4740 let footer = root_as_footer(&buffer[trailer_start - footer_len..trailer_start]).unwrap();
4741 let schema = try_fb_to_schema(footer.schema().unwrap()).unwrap();
4742
4743 let decoder =
4746 FileDecoder::new(Arc::new(schema), footer.version()).with_require_alignment(true);
4747
4748 let batches = footer.recordBatches().unwrap();
4749
4750 let block = batches.get(0);
4751 let block_len = block.bodyLength() as usize + block.metaDataLength() as usize;
4752 let data = buffer.slice_with_length(block.offset() as _, block_len);
4753
4754 let result = decoder.read_record_batch(block, &data);
4755
4756 let error = result.unwrap_err();
4757 assert_eq!(
4758 error.to_string(),
4759 "Invalid argument error: Misaligned buffers[0] in array of type Decimal128(38, 10), \
4760 offset from expected alignment of 16 by 8"
4761 );
4762 }
4763
4764 #[test]
4765 fn test_flush() {
4766 let num_cols = 2;
4769 let mut fields = Vec::new();
4770 let options = IpcWriteOptions::try_new(8, false, MetadataVersion::V5).unwrap();
4771 for i in 0..num_cols {
4772 let field = Field::new(format!("col_{i}"), DataType::Decimal128(38, 10), true);
4773 fields.push(field);
4774 }
4775 let schema = Schema::new(fields);
4776 let inner_stream_writer = BufWriter::with_capacity(1024, Vec::new());
4777 let inner_file_writer = BufWriter::with_capacity(1024, Vec::new());
4778 let mut stream_writer =
4779 StreamWriter::try_new_with_options(inner_stream_writer, &schema, options.clone())
4780 .unwrap();
4781 let mut file_writer =
4782 FileWriter::try_new_with_options(inner_file_writer, &schema, options).unwrap();
4783
4784 let stream_bytes_written_on_new = stream_writer.get_ref().get_ref().len();
4785 let file_bytes_written_on_new = file_writer.get_ref().get_ref().len();
4786 stream_writer.flush().unwrap();
4787 file_writer.flush().unwrap();
4788 let stream_bytes_written_on_flush = stream_writer.get_ref().get_ref().len();
4789 let file_bytes_written_on_flush = file_writer.get_ref().get_ref().len();
4790 let stream_out = stream_writer.into_inner().unwrap().into_inner().unwrap();
4791 let expected_stream_flushed_bytes = stream_out.len() - 8;
4795 let expected_file_flushed_bytes = expected_stream_flushed_bytes + 8;
4798
4799 assert!(
4800 stream_bytes_written_on_new < stream_bytes_written_on_flush,
4801 "this test makes no sense if flush is not actually required"
4802 );
4803 assert!(
4804 file_bytes_written_on_new < file_bytes_written_on_flush,
4805 "this test makes no sense if flush is not actually required"
4806 );
4807 assert_eq!(stream_bytes_written_on_flush, expected_stream_flushed_bytes);
4808 assert_eq!(file_bytes_written_on_flush, expected_file_flushed_bytes);
4809 }
4810
4811 #[test]
4812 fn test_roundtrip_list_of_fixed_list() -> Result<(), ArrowError> {
4813 let l1_type =
4814 DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Float32, false)), 3);
4815 let l2_type = DataType::List(Arc::new(Field::new("item", l1_type.clone(), false)));
4816
4817 let l0_builder = Float32Builder::new();
4818 let l1_builder = FixedSizeListBuilder::new(l0_builder, 3).with_field(Arc::new(Field::new(
4819 "item",
4820 DataType::Float32,
4821 false,
4822 )));
4823 let mut l2_builder =
4824 ListBuilder::new(l1_builder).with_field(Arc::new(Field::new("item", l1_type, false)));
4825
4826 for point in [[1.0, 2.0, 3.0], [4.0, 5.0, 6.0], [7.0, 8.0, 9.0]] {
4827 l2_builder.values().values().append_value(point[0]);
4828 l2_builder.values().values().append_value(point[1]);
4829 l2_builder.values().values().append_value(point[2]);
4830
4831 l2_builder.values().append(true);
4832 }
4833 l2_builder.append(true);
4834
4835 let point = [10., 11., 12.];
4836 l2_builder.values().values().append_value(point[0]);
4837 l2_builder.values().values().append_value(point[1]);
4838 l2_builder.values().values().append_value(point[2]);
4839
4840 l2_builder.values().append(true);
4841 l2_builder.append(true);
4842
4843 let array = Arc::new(l2_builder.finish()) as ArrayRef;
4844
4845 let schema = Arc::new(Schema::new_with_metadata(
4846 vec![Field::new("points", l2_type, false)],
4847 HashMap::default(),
4848 ));
4849
4850 test_slices(&array, &schema, 0, 1)?;
4853 test_slices(&array, &schema, 0, 2)?;
4854 test_slices(&array, &schema, 1, 1)?;
4855
4856 Ok(())
4857 }
4858
4859 #[test]
4860 fn test_roundtrip_list_of_fixed_list_w_nulls() -> Result<(), ArrowError> {
4861 let l0_builder = Float32Builder::new();
4862 let l1_builder = FixedSizeListBuilder::new(l0_builder, 3);
4863 let mut l2_builder = ListBuilder::new(l1_builder);
4864
4865 for point in [
4866 [Some(1.0), Some(2.0), None],
4867 [Some(4.0), Some(5.0), Some(6.0)],
4868 [None, Some(8.0), Some(9.0)],
4869 ] {
4870 for p in point {
4871 match p {
4872 Some(p) => l2_builder.values().values().append_value(p),
4873 None => l2_builder.values().values().append_null(),
4874 }
4875 }
4876
4877 l2_builder.values().append(true);
4878 }
4879 l2_builder.append(true);
4880
4881 let point = [Some(10.), None, None];
4882 for p in point {
4883 match p {
4884 Some(p) => l2_builder.values().values().append_value(p),
4885 None => l2_builder.values().values().append_null(),
4886 }
4887 }
4888
4889 l2_builder.values().append(true);
4890 l2_builder.append(true);
4891
4892 let array = Arc::new(l2_builder.finish()) as ArrayRef;
4893
4894 let schema = Arc::new(Schema::new_with_metadata(
4895 vec![Field::new(
4896 "points",
4897 DataType::List(Arc::new(Field::new(
4898 "item",
4899 DataType::FixedSizeList(
4900 Arc::new(Field::new("item", DataType::Float32, true)),
4901 3,
4902 ),
4903 true,
4904 ))),
4905 true,
4906 )],
4907 HashMap::default(),
4908 ));
4909
4910 test_slices(&array, &schema, 0, 1)?;
4913 test_slices(&array, &schema, 0, 2)?;
4914 test_slices(&array, &schema, 1, 1)?;
4915
4916 Ok(())
4917 }
4918
4919 fn test_slices(
4920 parent_array: &ArrayRef,
4921 schema: &SchemaRef,
4922 offset: usize,
4923 length: usize,
4924 ) -> Result<(), ArrowError> {
4925 let subarray = parent_array.slice(offset, length);
4926 let original_batch = RecordBatch::try_new(schema.clone(), vec![subarray])?;
4927
4928 let mut bytes = Vec::new();
4929 let mut writer = StreamWriter::try_new(&mut bytes, schema)?;
4930 writer.write(&original_batch)?;
4931 writer.finish()?;
4932
4933 let mut cursor = std::io::Cursor::new(bytes);
4934 let mut reader = StreamReader::try_new(&mut cursor, None)?;
4935 let returned_batch = reader.next().unwrap()?;
4936
4937 assert_eq!(original_batch, returned_batch);
4938
4939 Ok(())
4940 }
4941
4942 #[test]
4943 fn test_roundtrip_fixed_list() -> Result<(), ArrowError> {
4944 let int_builder = Int64Builder::new();
4945 let mut fixed_list_builder = FixedSizeListBuilder::new(int_builder, 3)
4946 .with_field(Arc::new(Field::new("item", DataType::Int64, false)));
4947
4948 for point in [[1, 2, 3], [4, 5, 6], [7, 8, 9], [10, 11, 12]] {
4949 fixed_list_builder.values().append_value(point[0]);
4950 fixed_list_builder.values().append_value(point[1]);
4951 fixed_list_builder.values().append_value(point[2]);
4952
4953 fixed_list_builder.append(true);
4954 }
4955
4956 let array = Arc::new(fixed_list_builder.finish()) as ArrayRef;
4957
4958 let schema = Arc::new(Schema::new_with_metadata(
4959 vec![Field::new(
4960 "points",
4961 DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Int64, false)), 3),
4962 false,
4963 )],
4964 HashMap::default(),
4965 ));
4966
4967 test_slices(&array, &schema, 0, 4)?;
4970 test_slices(&array, &schema, 0, 2)?;
4971 test_slices(&array, &schema, 1, 3)?;
4972 test_slices(&array, &schema, 2, 1)?;
4973
4974 Ok(())
4975 }
4976
4977 #[test]
4978 fn test_roundtrip_fixed_list_w_nulls() -> Result<(), ArrowError> {
4979 let int_builder = Int64Builder::new();
4980 let mut fixed_list_builder = FixedSizeListBuilder::new(int_builder, 3);
4981
4982 for point in [
4983 [Some(1), Some(2), None],
4984 [Some(4), Some(5), Some(6)],
4985 [None, Some(8), Some(9)],
4986 [Some(10), None, None],
4987 ] {
4988 for p in point {
4989 match p {
4990 Some(p) => fixed_list_builder.values().append_value(p),
4991 None => fixed_list_builder.values().append_null(),
4992 }
4993 }
4994
4995 fixed_list_builder.append(true);
4996 }
4997
4998 let array = Arc::new(fixed_list_builder.finish()) as ArrayRef;
4999
5000 let schema = Arc::new(Schema::new_with_metadata(
5001 vec![Field::new(
5002 "points",
5003 DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Int64, true)), 3),
5004 true,
5005 )],
5006 HashMap::default(),
5007 ));
5008
5009 test_slices(&array, &schema, 0, 4)?;
5012 test_slices(&array, &schema, 0, 2)?;
5013 test_slices(&array, &schema, 1, 3)?;
5014 test_slices(&array, &schema, 2, 1)?;
5015
5016 Ok(())
5017 }
5018
5019 #[test]
5020 fn test_metadata_encoding_ordering() {
5021 fn create_hash() -> u64 {
5022 let metadata: HashMap<String, String> = [
5023 ("a", "1"), ("b", "2"), ("c", "3"), ("d", "4"), ("e", "5"), ]
5029 .into_iter()
5030 .map(|(k, v)| (k.to_owned(), v.to_owned()))
5031 .collect();
5032
5033 let schema = Arc::new(
5035 Schema::new(vec![
5036 Field::new("a", DataType::Int64, true).with_metadata(metadata.clone()),
5037 ])
5038 .with_metadata(metadata)
5039 .clone(),
5040 );
5041 let batch = RecordBatch::new_empty(schema.clone());
5042
5043 let mut bytes = Vec::new();
5044 let mut w = StreamWriter::try_new(&mut bytes, batch.schema_ref()).unwrap();
5045 w.write(&batch).unwrap();
5046 w.finish().unwrap();
5047
5048 let mut h = std::hash::DefaultHasher::new();
5049 h.write(&bytes);
5050 h.finish()
5051 }
5052
5053 let expected = create_hash();
5054
5055 let all_passed = (0..20).all(|_| create_hash() == expected);
5060 assert!(all_passed);
5061 }
5062
5063 #[test]
5064 fn test_dictionary_tracker_reset() {
5065 let data_gen = IpcDataGenerator::default();
5066 let mut dictionary_tracker = DictionaryTracker::new(false);
5067 let writer_options = IpcWriteOptions::default();
5068 let mut compression_ctx = IpcWriteContext::default();
5069
5070 let schema = Arc::new(Schema::new(vec![Field::new(
5071 "a",
5072 DataType::Dictionary(Box::new(DataType::UInt8), Box::new(DataType::Utf8)),
5073 false,
5074 )]));
5075
5076 let mut write_single_batch_stream =
5077 |batch: RecordBatch, dict_tracker: &mut DictionaryTracker| -> Vec<u8> {
5078 let mut buffer = Vec::new();
5079
5080 let stream_header = data_gen.schema_to_bytes_with_dictionary_tracker(
5082 &schema,
5083 dict_tracker,
5084 &writer_options,
5085 );
5086 _ = write_message(&mut buffer, stream_header, &writer_options).unwrap();
5087
5088 let (encoded_dicts, encoded_batch) = data_gen
5089 .encode(&batch, dict_tracker, &writer_options, &mut compression_ctx)
5090 .unwrap();
5091 for encoded_dict in encoded_dicts {
5092 _ = write_message(&mut buffer, encoded_dict, &writer_options).unwrap();
5093 }
5094 _ = write_message(&mut buffer, encoded_batch, &writer_options).unwrap();
5095
5096 buffer
5097 };
5098
5099 let batch1 = RecordBatch::try_new(
5100 schema.clone(),
5101 vec![Arc::new(DictionaryArray::new(
5102 UInt8Array::from_iter_values([0]),
5103 Arc::new(StringArray::from_iter_values(["a"])),
5104 ))],
5105 )
5106 .unwrap();
5107 let buffer = write_single_batch_stream(batch1.clone(), &mut dictionary_tracker);
5108
5109 let mut reader = StreamReader::try_new(Cursor::new(buffer), None).unwrap();
5111 let read_batch = reader.next().unwrap().unwrap();
5112 assert_eq!(read_batch, batch1);
5113
5114 dictionary_tracker.clear();
5116
5117 let batch2 = RecordBatch::try_new(
5119 schema.clone(),
5120 vec![Arc::new(DictionaryArray::new(
5121 UInt8Array::from_iter_values([0]),
5122 Arc::new(StringArray::from_iter_values(["a"])),
5123 ))],
5124 )
5125 .unwrap();
5126 let buffer = write_single_batch_stream(batch2.clone(), &mut dictionary_tracker);
5127 let mut reader = StreamReader::try_new(Cursor::new(buffer), None).unwrap();
5128 let read_batch = reader.next().unwrap().unwrap();
5129 assert_eq!(read_batch, batch2);
5130 }
5131}