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