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