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