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 % usize::from(write_options.alignment) != 0 {
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 as zstd::zstd_safe::CompressionLevel)) {
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: HashMap<String, String>,
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: HashMap::new(),
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, Buffer) {
2279 if data.is_empty() {
2280 let mut offsets = MutableBuffer::new(size_of::<O>());
2283 offsets.extend_from_slice(O::usize_as(0).to_byte_slice());
2284 return (offsets.into(), MutableBuffer::new(0).into());
2285 }
2286
2287 let (offsets, original_start_offset, len) = reencode_offsets::<O>(&data.buffers()[0], data);
2288 let values = data.buffers()[1].slice_with_length(original_start_offset, len);
2289 (offsets, values)
2290}
2291
2292fn get_list_array_buffers<O: OffsetSizeTrait>(data: &ArrayData) -> (Buffer, ArrayData) {
2295 if data.is_empty() {
2296 let mut offsets = MutableBuffer::new(size_of::<O>());
2299 offsets.extend_from_slice(O::usize_as(0).to_byte_slice());
2300 return (offsets.into(), data.child_data()[0].slice(0, 0));
2301 }
2302
2303 let (offsets, original_start_offset, len) = reencode_offsets::<O>(&data.buffers()[0], data);
2304 let child_data = data.child_data()[0].slice(original_start_offset, len);
2305 (offsets, child_data)
2306}
2307
2308fn get_list_view_array_buffers<O: OffsetSizeTrait>(
2314 data: &ArrayData,
2315) -> (Buffer, Buffer, ArrayData) {
2316 if data.is_empty() {
2317 return (
2318 MutableBuffer::new(0).into(),
2319 MutableBuffer::new(0).into(),
2320 data.child_data()[0].slice(0, 0),
2321 );
2322 }
2323
2324 let offsets = &data.buffers()[0];
2325 let sizes = &data.buffers()[1];
2326
2327 let element_size = std::mem::size_of::<O>();
2328 let offsets_slice =
2329 offsets.slice_with_length(data.offset() * element_size, data.len() * element_size);
2330 let sizes_slice =
2331 sizes.slice_with_length(data.offset() * element_size, data.len() * element_size);
2332
2333 let child_data = data.child_data()[0].clone();
2334
2335 (offsets_slice, sizes_slice, child_data)
2336}
2337
2338fn get_or_truncate_buffer(array_data: &ArrayData) -> Buffer {
2345 let buffer = &array_data.buffers()[0];
2346 let layout = layout(array_data.data_type());
2347 let spec = &layout.buffers[0];
2348
2349 let byte_width = get_buffer_element_width(spec);
2350 let min_length = array_data.len() * byte_width;
2351 if buffer_need_truncate(array_data.offset(), buffer, spec, min_length) {
2352 let byte_offset = array_data.offset() * byte_width;
2353 let buffer_length = min(min_length, buffer.len() - byte_offset);
2354 buffer.slice_with_length(byte_offset, buffer_length)
2355 } else {
2356 buffer.clone()
2357 }
2358}
2359
2360fn write_array_data(
2366 array_data: &ArrayData,
2367 meta: &mut IpcMetadataBuilder,
2368 sink: &mut IpcBodySink<'_>,
2369 offset: i64,
2370 compression_codec: Option<CompressionCodec>,
2371 ipc_write_context: &mut IpcWriteContext,
2372 write_options: &IpcWriteOptions,
2373) -> Result<i64, ArrowError> {
2374 let mut offset = offset;
2375 let num_rows = array_data.len();
2376 if !matches!(array_data.data_type(), DataType::Null) {
2377 meta.nodes.push(crate::FieldNode::new(
2378 num_rows as i64,
2379 array_data.null_count() as i64,
2380 ));
2381 } else {
2382 meta.nodes
2384 .push(crate::FieldNode::new(num_rows as i64, num_rows as i64));
2385 }
2386 if has_validity_bitmap(array_data.data_type(), write_options) {
2387 let null_buffer = match array_data.nulls() {
2389 None => {
2390 let num_bytes = bit_util::ceil(num_rows, 8);
2392 let buffer = MutableBuffer::new(num_bytes);
2393 let buffer = buffer.with_bitset(num_bytes, true);
2394 buffer.into()
2395 }
2396 Some(buffer) => buffer.inner().sliced(),
2397 };
2398
2399 offset = encode_sink_buffer(
2400 null_buffer,
2401 meta,
2402 sink,
2403 offset,
2404 compression_codec,
2405 ipc_write_context,
2406 write_options.alignment,
2407 )?;
2408 }
2409
2410 let data_type = array_data.data_type();
2411 if matches!(data_type, DataType::Binary | DataType::Utf8) {
2412 let (offsets, values) = get_byte_array_buffers::<i32>(array_data);
2413 for buffer in [offsets, values] {
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 let (offsets, values) = get_byte_array_buffers::<i64>(array_data);
2455 for buffer in [offsets, values] {
2456 offset = encode_sink_buffer(
2457 buffer,
2458 meta,
2459 sink,
2460 offset,
2461 compression_codec,
2462 ipc_write_context,
2463 write_options.alignment,
2464 )?;
2465 }
2466 } else if DataType::is_numeric(data_type)
2467 || DataType::is_temporal(data_type)
2468 || matches!(
2469 array_data.data_type(),
2470 DataType::FixedSizeBinary(_) | DataType::Dictionary(_, _)
2471 )
2472 {
2473 assert_eq!(array_data.buffers().len(), 1);
2475
2476 let buffer = get_or_truncate_buffer(array_data);
2477 offset = encode_sink_buffer(
2478 buffer,
2479 meta,
2480 sink,
2481 offset,
2482 compression_codec,
2483 ipc_write_context,
2484 write_options.alignment,
2485 )?;
2486 } else if matches!(data_type, DataType::Boolean) {
2487 assert_eq!(array_data.buffers().len(), 1);
2490
2491 let buffer = &array_data.buffers()[0];
2492 let buffer = buffer.bit_slice(array_data.offset(), array_data.len());
2493 offset = encode_sink_buffer(
2494 buffer,
2495 meta,
2496 sink,
2497 offset,
2498 compression_codec,
2499 ipc_write_context,
2500 write_options.alignment,
2501 )?;
2502 } else if matches!(
2503 data_type,
2504 DataType::List(_) | DataType::LargeList(_) | DataType::Map(_, _)
2505 ) {
2506 assert_eq!(array_data.buffers().len(), 1);
2507 assert_eq!(array_data.child_data().len(), 1);
2508
2509 let (offsets, sliced_child_data) = match data_type {
2511 DataType::List(_) => get_list_array_buffers::<i32>(array_data),
2512 DataType::Map(_, _) => get_list_array_buffers::<i32>(array_data),
2513 DataType::LargeList(_) => get_list_array_buffers::<i64>(array_data),
2514 _ => unreachable!(),
2515 };
2516 offset = encode_sink_buffer(
2517 offsets,
2518 meta,
2519 sink,
2520 offset,
2521 compression_codec,
2522 ipc_write_context,
2523 write_options.alignment,
2524 )?;
2525 offset = write_array_data(
2526 &sliced_child_data,
2527 meta,
2528 sink,
2529 offset,
2530 compression_codec,
2531 ipc_write_context,
2532 write_options,
2533 )?;
2534 return Ok(offset);
2535 } else if matches!(
2536 data_type,
2537 DataType::ListView(_) | DataType::LargeListView(_)
2538 ) {
2539 assert_eq!(array_data.buffers().len(), 2); assert_eq!(array_data.child_data().len(), 1);
2541
2542 let (offsets, sizes, child_data) = match data_type {
2543 DataType::ListView(_) => get_list_view_array_buffers::<i32>(array_data),
2544 DataType::LargeListView(_) => get_list_view_array_buffers::<i64>(array_data),
2545 _ => unreachable!(),
2546 };
2547
2548 offset = encode_sink_buffer(
2549 offsets,
2550 meta,
2551 sink,
2552 offset,
2553 compression_codec,
2554 ipc_write_context,
2555 write_options.alignment,
2556 )?;
2557 offset = encode_sink_buffer(
2558 sizes,
2559 meta,
2560 sink,
2561 offset,
2562 compression_codec,
2563 ipc_write_context,
2564 write_options.alignment,
2565 )?;
2566
2567 offset = write_array_data(
2568 &child_data,
2569 meta,
2570 sink,
2571 offset,
2572 compression_codec,
2573 ipc_write_context,
2574 write_options,
2575 )?;
2576 return Ok(offset);
2577 } else if let DataType::FixedSizeList(_, fixed_size) = data_type {
2578 assert_eq!(array_data.child_data().len(), 1);
2579 let fixed_size = *fixed_size as usize;
2580
2581 let child_offset = array_data.offset() * fixed_size;
2582 let child_length = array_data.len() * fixed_size;
2583 let child_data = array_data.child_data()[0].slice(child_offset, child_length);
2584
2585 offset = write_array_data(
2586 &child_data,
2587 meta,
2588 sink,
2589 offset,
2590 compression_codec,
2591 ipc_write_context,
2592 write_options,
2593 )?;
2594 return Ok(offset);
2595 } else {
2596 for buffer in array_data.buffers() {
2597 offset = encode_sink_buffer(
2598 buffer.clone(),
2599 meta,
2600 sink,
2601 offset,
2602 compression_codec,
2603 ipc_write_context,
2604 write_options.alignment,
2605 )?;
2606 }
2607 }
2608
2609 match array_data.data_type() {
2610 DataType::Dictionary(_, _) => {}
2611 DataType::RunEndEncoded(_, _) => {
2612 let arr = unslice_run_array(array_data.clone())?;
2614 for data_ref in arr.child_data() {
2616 offset = write_array_data(
2618 data_ref,
2619 meta,
2620 sink,
2621 offset,
2622 compression_codec,
2623 ipc_write_context,
2624 write_options,
2625 )?;
2626 }
2627 }
2628 _ => {
2629 for data_ref in array_data.child_data() {
2631 offset = write_array_data(
2633 data_ref,
2634 meta,
2635 sink,
2636 offset,
2637 compression_codec,
2638 ipc_write_context,
2639 write_options,
2640 )?;
2641 }
2642 }
2643 }
2644 Ok(offset)
2645}
2646
2647fn encode_sink_buffer(
2661 buffer: Buffer,
2662 ipc_meta_data: &mut IpcMetadataBuilder,
2663 sink: &mut IpcBodySink<'_>,
2664 offset: i64,
2665 compression_codec: Option<CompressionCodec>,
2666 ipc_write_context: &mut IpcWriteContext,
2667 alignment: u8,
2668) -> Result<i64, ArrowError> {
2669 let (encoded, len) = match compression_codec {
2670 None => {
2671 let len = buffer.len() as i64;
2672 (EncodedBuffer::Raw(buffer), len)
2673 }
2674 Some(codec) => {
2675 let mut scratch = Vec::new();
2676 let written =
2677 codec.compress_to_vec(buffer.as_slice(), &mut scratch, ipc_write_context)?;
2678 let len = i64::try_from(written)
2679 .map_err(|e| ArrowError::InvalidArgumentError(format!("{e}")))?;
2680 (EncodedBuffer::Compressed(scratch), len)
2681 }
2682 };
2683
2684 let pad_len = pad_to_alignment(alignment, len as usize);
2685 sink.write(pad_len, encoded);
2686 ipc_meta_data.buffers.push(crate::Buffer::new(offset, len));
2687 Ok(offset + len + pad_len as i64)
2688}
2689
2690const PADDING: [u8; 64] = [0; 64];
2691
2692#[inline]
2698fn estimate_encoded_buffer_count(dt: &DataType) -> usize {
2699 match dt {
2700 DataType::Null => 0,
2701
2702 DataType::Binary | DataType::Utf8 | DataType::LargeBinary | DataType::LargeUtf8 => 3,
2703
2704 DataType::BinaryView | DataType::Utf8View => 3,
2705
2706 DataType::List(f) | DataType::LargeList(f) | DataType::Map(f, _) => {
2707 2 + estimate_encoded_buffer_count(f.data_type())
2708 }
2709
2710 DataType::ListView(f) | DataType::LargeListView(f) => {
2711 3 + estimate_encoded_buffer_count(f.data_type())
2712 }
2713
2714 DataType::FixedSizeList(f, _) => 1 + estimate_encoded_buffer_count(f.data_type()),
2715
2716 DataType::Struct(fields) => {
2717 1 + fields
2718 .iter()
2719 .map(|f| estimate_encoded_buffer_count(f.data_type()))
2720 .sum::<usize>()
2721 }
2722
2723 DataType::Dictionary(_, _) => 2,
2725
2726 DataType::Union(fields, UnionMode::Sparse) => {
2727 1 + fields
2728 .iter()
2729 .map(|(_, f)| estimate_encoded_buffer_count(f.data_type()))
2730 .sum::<usize>()
2731 }
2732 DataType::Union(fields, UnionMode::Dense) => {
2733 2 + fields
2734 .iter()
2735 .map(|(_, f)| estimate_encoded_buffer_count(f.data_type()))
2736 .sum::<usize>()
2737 }
2738
2739 DataType::RunEndEncoded(run_ends, values) => {
2740 estimate_encoded_buffer_count(run_ends.data_type())
2741 + estimate_encoded_buffer_count(values.data_type())
2742 }
2743 _ => 2,
2745 }
2746}
2747
2748#[inline]
2750fn pad_to_alignment(alignment: u8, len: usize) -> usize {
2751 let a = usize::from(alignment - 1);
2752 ((len + a) & !a) - len
2753}
2754
2755#[cfg(test)]
2756mod tests {
2757 use std::hash::Hasher;
2758 use std::io::Cursor;
2759 use std::io::Seek;
2760
2761 use arrow_array::builder::FixedSizeListBuilder;
2762 use arrow_array::builder::Float32Builder;
2763 use arrow_array::builder::Int64Builder;
2764 use arrow_array::builder::MapBuilder;
2765 use arrow_array::builder::StringViewBuilder;
2766 use arrow_array::builder::UnionBuilder;
2767 use arrow_array::builder::{
2768 GenericListBuilder, GenericListViewBuilder, ListBuilder, StringBuilder,
2769 };
2770 use arrow_array::builder::{PrimitiveRunBuilder, UInt32Builder};
2771 use arrow_array::types::*;
2772 use arrow_buffer::ScalarBuffer;
2773
2774 use crate::MetadataVersion;
2775 use crate::convert::fb_to_schema;
2776 use crate::reader::*;
2777 use crate::root_as_footer;
2778
2779 use super::*;
2780
2781 fn serialize_file(rb: &RecordBatch) -> Vec<u8> {
2782 let mut writer = FileWriter::try_new(vec![], rb.schema_ref()).unwrap();
2783 writer.write(rb).unwrap();
2784 writer.finish().unwrap();
2785 writer.into_inner().unwrap()
2786 }
2787
2788 fn deserialize_file(bytes: Vec<u8>) -> RecordBatch {
2789 let mut reader = FileReader::try_new(Cursor::new(bytes), None).unwrap();
2790 reader.next().unwrap().unwrap()
2791 }
2792
2793 fn serialize_stream(record: &RecordBatch) -> Vec<u8> {
2794 const IPC_ALIGNMENT: usize = 8;
2798
2799 let mut stream_writer = StreamWriter::try_new_with_options(
2800 vec![],
2801 record.schema_ref(),
2802 IpcWriteOptions::try_new(IPC_ALIGNMENT, false, MetadataVersion::V5).unwrap(),
2803 )
2804 .unwrap();
2805 stream_writer.write(record).unwrap();
2806 stream_writer.finish().unwrap();
2807 stream_writer.into_inner().unwrap()
2808 }
2809
2810 fn deserialize_stream(bytes: Vec<u8>) -> RecordBatch {
2811 let mut stream_reader = StreamReader::try_new(Cursor::new(bytes), None).unwrap();
2812 stream_reader.next().unwrap().unwrap()
2813 }
2814
2815 fn encode_stream(
2820 schema: &Schema,
2821 batches: &[RecordBatch],
2822 options: IpcWriteOptions,
2823 ) -> Vec<u8> {
2824 let mut encoder = StreamEncoder::try_new_with_options(schema, options).unwrap();
2825 let mut bytes = Vec::new();
2826 for batch in batches {
2827 for buffer in encoder.encode(batch).unwrap() {
2828 bytes.write_all(buffer.as_slice()).unwrap();
2829 }
2830 }
2831 for buffer in encoder.finish().unwrap() {
2832 bytes.write_all(buffer.as_slice()).unwrap();
2833 }
2834 bytes
2835 }
2836
2837 fn write_stream(schema: &Schema, batches: &[RecordBatch], options: IpcWriteOptions) -> Vec<u8> {
2838 let mut bytes = Vec::new();
2839 let mut writer = StreamWriter::try_new_with_options(&mut bytes, schema, options).unwrap();
2840 for batch in batches {
2841 writer.write(batch).unwrap();
2842 }
2843 writer.finish().unwrap();
2844 bytes
2845 }
2846
2847 #[tokio::test]
2850 async fn test_stream_encoder_async_writer_matches_stream_writer() {
2851 use tokio::io::{AsyncReadExt, AsyncWriteExt};
2852
2853 let batch = record_batch!(("a", Int32, [1, 2, 3]), ("b", Utf8, ["x", "y", "z"])).unwrap();
2854 let options = IpcWriteOptions::default();
2855 let expected = write_stream(
2856 batch.schema_ref(),
2857 std::slice::from_ref(&batch),
2858 options.clone(),
2859 );
2860
2861 let (mut sink, mut source) = tokio::io::duplex(64);
2862 let read = tokio::spawn(async move {
2863 let mut bytes = Vec::new();
2864 source.read_to_end(&mut bytes).await.unwrap();
2865 bytes
2866 });
2867
2868 let mut encoder = StreamEncoder::try_new_with_options(batch.schema_ref(), options).unwrap();
2869 for buffer in encoder.encode(&batch).unwrap() {
2870 sink.write_all(buffer.as_slice()).await.unwrap();
2871 }
2872 for buffer in encoder.finish().unwrap() {
2873 sink.write_all(buffer.as_slice()).await.unwrap();
2874 }
2875 sink.shutdown().await.unwrap();
2876
2877 let encoded = read.await.unwrap();
2878 assert_eq!(encoded, expected);
2879
2880 let mut reader = StreamReader::try_new(Cursor::new(encoded), None).unwrap();
2881 assert_eq!(reader.next().unwrap().unwrap(), batch);
2882 assert!(reader.next().is_none());
2883 }
2884
2885 #[test]
2886 fn test_stream_encoder_empty_stream_matches_stream_writer() {
2887 let schema = Schema::new(vec![Field::new("a", DataType::Int32, true)]);
2888 let options = IpcWriteOptions::default();
2889 let encoded = encode_stream(&schema, &[], options.clone());
2890 let written = write_stream(&schema, &[], options);
2891
2892 assert_eq!(encoded, written);
2893
2894 let mut reader = StreamReader::try_new(Cursor::new(encoded), None).unwrap();
2895 assert!(reader.next().is_none());
2896 }
2897
2898 #[test]
2899 fn test_stream_encoder_dictionary_batches_match_stream_writer() {
2900 let schema = Arc::new(Schema::new(vec![Field::new(
2901 "a",
2902 DataType::Dictionary(Box::new(DataType::UInt8), Box::new(DataType::Utf8)),
2903 false,
2904 )]));
2905 let batch = RecordBatch::try_new(
2906 schema.clone(),
2907 vec![Arc::new(DictionaryArray::new(
2908 UInt8Array::from_iter_values([0, 1, 0]),
2909 Arc::new(StringArray::from_iter_values(["a", "b"])),
2910 ))],
2911 )
2912 .unwrap();
2913 let options = IpcWriteOptions::default();
2914 let encoded = encode_stream(&schema, std::slice::from_ref(&batch), options.clone());
2915 let written = write_stream(&schema, std::slice::from_ref(&batch), options);
2916
2917 assert_eq!(encoded, written);
2918
2919 let mut reader = StreamReader::try_new(Cursor::new(encoded), None).unwrap();
2920 assert_eq!(reader.next().unwrap().unwrap(), batch);
2921 assert!(reader.next().is_none());
2922 }
2923
2924 #[test]
2925 #[cfg(feature = "lz4")]
2926 fn test_write_empty_record_batch_lz4_compression() {
2927 let schema = Schema::new(vec![Field::new("field1", DataType::Int32, true)]);
2928 let values: Vec<Option<i32>> = vec![];
2929 let array = Int32Array::from(values);
2930 let record_batch =
2931 RecordBatch::try_new(Arc::new(schema.clone()), vec![Arc::new(array)]).unwrap();
2932
2933 let mut file = tempfile::tempfile().unwrap();
2934
2935 {
2936 let write_option = IpcWriteOptions::try_new(8, false, crate::MetadataVersion::V5)
2937 .unwrap()
2938 .try_with_compression(Some(crate::CompressionType::LZ4_FRAME))
2939 .unwrap();
2940
2941 let mut writer =
2942 FileWriter::try_new_with_options(&mut file, &schema, write_option).unwrap();
2943 writer.write(&record_batch).unwrap();
2944 writer.finish().unwrap();
2945 }
2946 file.rewind().unwrap();
2947 {
2948 let reader = FileReader::try_new(file, None).unwrap();
2950 for read_batch in reader {
2951 read_batch
2952 .unwrap()
2953 .columns()
2954 .iter()
2955 .zip(record_batch.columns())
2956 .for_each(|(a, b)| {
2957 assert_eq!(a.data_type(), b.data_type());
2958 assert_eq!(a.len(), b.len());
2959 assert_eq!(a.null_count(), b.null_count());
2960 });
2961 }
2962 }
2963 }
2964
2965 #[test]
2966 #[cfg(feature = "lz4")]
2967 fn test_write_file_with_lz4_compression() {
2968 let schema = Schema::new(vec![Field::new("field1", DataType::Int32, true)]);
2969 let values: Vec<Option<i32>> = vec![Some(12), Some(1)];
2970 let array = Int32Array::from(values);
2971 let record_batch =
2972 RecordBatch::try_new(Arc::new(schema.clone()), vec![Arc::new(array)]).unwrap();
2973
2974 let mut file = tempfile::tempfile().unwrap();
2975 {
2976 let write_option = IpcWriteOptions::try_new(8, false, crate::MetadataVersion::V5)
2977 .unwrap()
2978 .try_with_compression(Some(crate::CompressionType::LZ4_FRAME))
2979 .unwrap();
2980
2981 let mut writer =
2982 FileWriter::try_new_with_options(&mut file, &schema, write_option).unwrap();
2983 writer.write(&record_batch).unwrap();
2984 writer.finish().unwrap();
2985 }
2986 file.rewind().unwrap();
2987 {
2988 let reader = FileReader::try_new(file, None).unwrap();
2990 for read_batch in reader {
2991 read_batch
2992 .unwrap()
2993 .columns()
2994 .iter()
2995 .zip(record_batch.columns())
2996 .for_each(|(a, b)| {
2997 assert_eq!(a.data_type(), b.data_type());
2998 assert_eq!(a.len(), b.len());
2999 assert_eq!(a.null_count(), b.null_count());
3000 });
3001 }
3002 }
3003 }
3004
3005 #[test]
3006 #[cfg(feature = "zstd")]
3007 fn test_write_file_with_zstd_compression() {
3008 let schema = Schema::new(vec![Field::new("field1", DataType::Int32, true)]);
3009 let values: Vec<Option<i32>> = vec![Some(12), Some(1)];
3010 let array = Int32Array::from(values);
3011 let record_batch =
3012 RecordBatch::try_new(Arc::new(schema.clone()), vec![Arc::new(array)]).unwrap();
3013 let mut file = tempfile::tempfile().unwrap();
3014 {
3015 let write_option = IpcWriteOptions::try_new(8, false, crate::MetadataVersion::V5)
3016 .unwrap()
3017 .try_with_compression(Some(crate::CompressionType::ZSTD))
3018 .unwrap()
3019 .try_with_compression_level(Some(1))
3020 .unwrap();
3021
3022 let mut writer =
3023 FileWriter::try_new_with_options(&mut file, &schema, write_option).unwrap();
3024 writer.write(&record_batch).unwrap();
3025 writer.finish().unwrap();
3026 }
3027 file.rewind().unwrap();
3028 {
3029 let reader = FileReader::try_new(file, None).unwrap();
3031 for read_batch in reader {
3032 read_batch
3033 .unwrap()
3034 .columns()
3035 .iter()
3036 .zip(record_batch.columns())
3037 .for_each(|(a, b)| {
3038 assert_eq!(a.data_type(), b.data_type());
3039 assert_eq!(a.len(), b.len());
3040 assert_eq!(a.null_count(), b.null_count());
3041 });
3042 }
3043 }
3044 }
3045
3046 #[test]
3047 fn test_write_file() {
3048 let schema = Schema::new(vec![Field::new("field1", DataType::UInt32, true)]);
3049 let values: Vec<Option<u32>> = vec![
3050 Some(999),
3051 None,
3052 Some(235),
3053 Some(123),
3054 None,
3055 None,
3056 None,
3057 None,
3058 None,
3059 ];
3060 let array1 = UInt32Array::from(values);
3061 let batch =
3062 RecordBatch::try_new(Arc::new(schema.clone()), vec![Arc::new(array1) as ArrayRef])
3063 .unwrap();
3064 let mut file = tempfile::tempfile().unwrap();
3065 {
3066 let mut writer = FileWriter::try_new(&mut file, &schema).unwrap();
3067
3068 writer.write(&batch).unwrap();
3069 writer.finish().unwrap();
3070 }
3071 file.rewind().unwrap();
3072
3073 {
3074 let mut reader = FileReader::try_new(file, None).unwrap();
3075 while let Some(Ok(read_batch)) = reader.next() {
3076 read_batch
3077 .columns()
3078 .iter()
3079 .zip(batch.columns())
3080 .for_each(|(a, b)| {
3081 assert_eq!(a.data_type(), b.data_type());
3082 assert_eq!(a.len(), b.len());
3083 assert_eq!(a.null_count(), b.null_count());
3084 });
3085 }
3086 }
3087 }
3088
3089 #[test]
3090 fn test_empty_utf8_ipc_writes_nonempty_offsets_buffer() {
3091 let name = StringArray::from(Vec::<String>::new());
3092 let (offsets, values) = get_byte_array_buffers::<i32>(&name.to_data());
3093
3094 assert_eq!(name.len(), 0);
3095 assert_eq!(
3096 offsets.len(),
3097 std::mem::size_of::<i32>(),
3098 "offsets buffer should contain one zero i32 offset"
3099 );
3100 assert_eq!(values.len(), 0, "values buffer should remain empty");
3101 }
3102
3103 #[test]
3104 fn test_empty_large_utf8_ipc_writes_nonempty_offsets_buffer() {
3105 let name = LargeStringArray::from(Vec::<String>::new());
3106 let (offsets, values) = get_byte_array_buffers::<i64>(&name.to_data());
3107
3108 assert_eq!(name.len(), 0);
3109 assert_eq!(
3110 offsets.len(),
3111 std::mem::size_of::<i64>(),
3112 "offsets buffer should contain one zero i64 offset"
3113 );
3114 assert_eq!(values.len(), 0, "values buffer should remain empty");
3115 }
3116
3117 #[test]
3118 fn test_empty_list_ipc_writes_nonempty_offsets_buffer() {
3119 let list = GenericListBuilder::<i32, _>::new(UInt32Builder::new()).finish();
3120 let (offsets, child_data) = get_list_array_buffers::<i32>(&list.to_data());
3121
3122 assert_eq!(list.len(), 0);
3123 assert_eq!(
3124 offsets.len(),
3125 std::mem::size_of::<i32>(),
3126 "offsets buffer should contain one zero i32 offset"
3127 );
3128 assert_eq!(child_data.len(), 0, "child data should remain empty");
3129 }
3130
3131 #[test]
3132 fn test_empty_large_list_ipc_writes_nonempty_offsets_buffer() {
3133 let list = GenericListBuilder::<i64, _>::new(UInt32Builder::new()).finish();
3134 let (offsets, child_data) = get_list_array_buffers::<i64>(&list.to_data());
3135
3136 assert_eq!(list.len(), 0);
3137 assert_eq!(
3138 offsets.len(),
3139 std::mem::size_of::<i64>(),
3140 "offsets buffer should contain one zero i64 offset"
3141 );
3142 assert_eq!(child_data.len(), 0, "child data should remain empty");
3143 }
3144
3145 fn write_null_file(options: IpcWriteOptions) {
3146 let schema = Schema::new(vec![
3147 Field::new("nulls", DataType::Null, true),
3148 Field::new("int32s", DataType::Int32, false),
3149 Field::new("nulls2", DataType::Null, true),
3150 Field::new("f64s", DataType::Float64, false),
3151 ]);
3152 let array1 = NullArray::new(32);
3153 let array2 = Int32Array::from(vec![1; 32]);
3154 let array3 = NullArray::new(32);
3155 let array4 = Float64Array::from(vec![f64::NAN; 32]);
3156 let batch = RecordBatch::try_new(
3157 Arc::new(schema.clone()),
3158 vec![
3159 Arc::new(array1) as ArrayRef,
3160 Arc::new(array2) as ArrayRef,
3161 Arc::new(array3) as ArrayRef,
3162 Arc::new(array4) as ArrayRef,
3163 ],
3164 )
3165 .unwrap();
3166 let mut file = tempfile::tempfile().unwrap();
3167 {
3168 let mut writer = FileWriter::try_new_with_options(&mut file, &schema, options).unwrap();
3169
3170 writer.write(&batch).unwrap();
3171 writer.finish().unwrap();
3172 }
3173
3174 file.rewind().unwrap();
3175
3176 {
3177 let reader = FileReader::try_new(file, None).unwrap();
3178 reader.for_each(|maybe_batch| {
3179 maybe_batch
3180 .unwrap()
3181 .columns()
3182 .iter()
3183 .zip(batch.columns())
3184 .for_each(|(a, b)| {
3185 assert_eq!(a.data_type(), b.data_type());
3186 assert_eq!(a.len(), b.len());
3187 assert_eq!(a.null_count(), b.null_count());
3188 });
3189 });
3190 }
3191 }
3192 #[test]
3193 fn test_write_null_file_v4() {
3194 write_null_file(IpcWriteOptions::try_new(8, false, MetadataVersion::V4).unwrap());
3195 write_null_file(IpcWriteOptions::try_new(8, true, MetadataVersion::V4).unwrap());
3196 write_null_file(IpcWriteOptions::try_new(64, false, MetadataVersion::V4).unwrap());
3197 write_null_file(IpcWriteOptions::try_new(64, true, MetadataVersion::V4).unwrap());
3198 }
3199
3200 #[test]
3201 fn test_write_null_file_v5() {
3202 write_null_file(IpcWriteOptions::try_new(8, false, MetadataVersion::V5).unwrap());
3203 write_null_file(IpcWriteOptions::try_new(64, false, MetadataVersion::V5).unwrap());
3204 }
3205
3206 #[test]
3207 fn track_union_nested_dict() {
3208 let inner: DictionaryArray<Int32Type> = vec!["a", "b", "a"].into_iter().collect();
3209
3210 let array = Arc::new(inner) as ArrayRef;
3211
3212 #[allow(deprecated)]
3214 let dctfield = Field::new_dict("dict", array.data_type().clone(), false, 0, false);
3215 let union_fields = [(0, Arc::new(dctfield))].into_iter().collect();
3216
3217 let types = [0, 0, 0].into_iter().collect::<ScalarBuffer<i8>>();
3218 let offsets = [0, 1, 2].into_iter().collect::<ScalarBuffer<i32>>();
3219
3220 let union = UnionArray::try_new(union_fields, types, Some(offsets), vec![array]).unwrap();
3221
3222 let schema = Arc::new(Schema::new(vec![Field::new(
3223 "union",
3224 union.data_type().clone(),
3225 false,
3226 )]));
3227
3228 let r#gen = IpcDataGenerator::default();
3229 let mut dict_tracker = DictionaryTracker::new(false);
3230 r#gen.schema_to_bytes_with_dictionary_tracker(
3231 &schema,
3232 &mut dict_tracker,
3233 &IpcWriteOptions::default(),
3234 );
3235
3236 let batch = RecordBatch::try_new(schema, vec![Arc::new(union)]).unwrap();
3237
3238 r#gen
3239 .encode(
3240 &batch,
3241 &mut dict_tracker,
3242 &Default::default(),
3243 &mut Default::default(),
3244 )
3245 .unwrap();
3246
3247 assert!(dict_tracker.written.contains_key(&0));
3250 }
3251
3252 #[test]
3253 fn track_struct_nested_dict() {
3254 let inner: DictionaryArray<Int32Type> = vec!["a", "b", "a"].into_iter().collect();
3255
3256 let array = Arc::new(inner) as ArrayRef;
3257
3258 #[allow(deprecated)]
3260 let dctfield = Arc::new(Field::new_dict(
3261 "dict",
3262 array.data_type().clone(),
3263 false,
3264 2,
3265 false,
3266 ));
3267
3268 let s = StructArray::from(vec![(dctfield, array)]);
3269 let struct_array = Arc::new(s) as ArrayRef;
3270
3271 let schema = Arc::new(Schema::new(vec![Field::new(
3272 "struct",
3273 struct_array.data_type().clone(),
3274 false,
3275 )]));
3276
3277 let r#gen = IpcDataGenerator::default();
3278 let mut dict_tracker = DictionaryTracker::new(false);
3279 r#gen.schema_to_bytes_with_dictionary_tracker(
3280 &schema,
3281 &mut dict_tracker,
3282 &IpcWriteOptions::default(),
3283 );
3284
3285 let batch = RecordBatch::try_new(schema, vec![struct_array]).unwrap();
3286
3287 r#gen
3288 .encode(
3289 &batch,
3290 &mut dict_tracker,
3291 &Default::default(),
3292 &mut Default::default(),
3293 )
3294 .unwrap();
3295
3296 assert!(dict_tracker.written.contains_key(&0));
3297 }
3298
3299 fn write_union_file(options: IpcWriteOptions) {
3300 let schema = Schema::new(vec![Field::new_union(
3301 "union",
3302 vec![0, 1],
3303 vec![
3304 Field::new("a", DataType::Int32, false),
3305 Field::new("c", DataType::Float64, false),
3306 ],
3307 UnionMode::Sparse,
3308 )]);
3309 let mut builder = UnionBuilder::with_capacity_sparse(5);
3310 builder.append::<Int32Type>("a", 1).unwrap();
3311 builder.append_null::<Int32Type>("a").unwrap();
3312 builder.append::<Float64Type>("c", 3.0).unwrap();
3313 builder.append_null::<Float64Type>("c").unwrap();
3314 builder.append::<Int32Type>("a", 4).unwrap();
3315 let union = builder.build().unwrap();
3316
3317 let batch =
3318 RecordBatch::try_new(Arc::new(schema.clone()), vec![Arc::new(union) as ArrayRef])
3319 .unwrap();
3320
3321 let mut file = tempfile::tempfile().unwrap();
3322 {
3323 let mut writer = FileWriter::try_new_with_options(&mut file, &schema, options).unwrap();
3324
3325 writer.write(&batch).unwrap();
3326 writer.finish().unwrap();
3327 }
3328 file.rewind().unwrap();
3329
3330 {
3331 let reader = FileReader::try_new(file, None).unwrap();
3332 reader.for_each(|maybe_batch| {
3333 maybe_batch
3334 .unwrap()
3335 .columns()
3336 .iter()
3337 .zip(batch.columns())
3338 .for_each(|(a, b)| {
3339 assert_eq!(a.data_type(), b.data_type());
3340 assert_eq!(a.len(), b.len());
3341 assert_eq!(a.null_count(), b.null_count());
3342 });
3343 });
3344 }
3345 }
3346
3347 #[test]
3348 fn test_write_union_file_v4_v5() {
3349 write_union_file(IpcWriteOptions::try_new(8, false, MetadataVersion::V4).unwrap());
3350 write_union_file(IpcWriteOptions::try_new(8, false, MetadataVersion::V5).unwrap());
3351 }
3352
3353 #[test]
3354 fn test_write_view_types() {
3355 const LONG_TEST_STRING: &str =
3356 "This is a long string to make sure binary view array handles it";
3357 let schema = Schema::new(vec![
3358 Field::new("field1", DataType::BinaryView, true),
3359 Field::new("field2", DataType::Utf8View, true),
3360 ]);
3361 let values: Vec<Option<&[u8]>> = vec![
3362 Some(b"foo"),
3363 Some(b"bar"),
3364 Some(LONG_TEST_STRING.as_bytes()),
3365 ];
3366 let binary_array = BinaryViewArray::from_iter(values);
3367 let utf8_array =
3368 StringViewArray::from_iter(vec![Some("foo"), Some("bar"), Some(LONG_TEST_STRING)]);
3369 let record_batch = RecordBatch::try_new(
3370 Arc::new(schema.clone()),
3371 vec![Arc::new(binary_array), Arc::new(utf8_array)],
3372 )
3373 .unwrap();
3374
3375 let mut file = tempfile::tempfile().unwrap();
3376 {
3377 let mut writer = FileWriter::try_new(&mut file, &schema).unwrap();
3378 writer.write(&record_batch).unwrap();
3379 writer.finish().unwrap();
3380 }
3381 file.rewind().unwrap();
3382 {
3383 let mut reader = FileReader::try_new(&file, None).unwrap();
3384 let read_batch = reader.next().unwrap().unwrap();
3385 read_batch
3386 .columns()
3387 .iter()
3388 .zip(record_batch.columns())
3389 .for_each(|(a, b)| {
3390 assert_eq!(a, b);
3391 });
3392 }
3393 file.rewind().unwrap();
3394 {
3395 let mut reader = FileReader::try_new(&file, Some(vec![0])).unwrap();
3396 let read_batch = reader.next().unwrap().unwrap();
3397 assert_eq!(read_batch.num_columns(), 1);
3398 let read_array = read_batch.column(0);
3399 let write_array = record_batch.column(0);
3400 assert_eq!(read_array, write_array);
3401 }
3402 }
3403
3404 #[test]
3405 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 fn test_large_slice_string() {
3605 let strings: Vec<_> = (0..8000)
3606 .map(|i| {
3607 if i % 2 == 0 {
3608 Some(format!("value{i}"))
3609 } else {
3610 None
3611 }
3612 })
3613 .collect();
3614
3615 ensure_roundtrip(Arc::new(StringArray::from(strings)));
3616 }
3617
3618 #[test]
3619 fn test_large_slice_string_list() {
3620 let mut ls = ListBuilder::new(StringBuilder::new());
3621
3622 let mut s = String::new();
3623 for row_number in 0..8000 {
3624 if row_number % 2 == 0 {
3625 for list_element in 0..1000 {
3626 s.clear();
3627 use std::fmt::Write;
3628 write!(&mut s, "value{row_number}-{list_element}").unwrap();
3629 ls.values().append_value(&s);
3630 }
3631 ls.append(true)
3632 } else {
3633 ls.append(false); }
3635 }
3636
3637 ensure_roundtrip(Arc::new(ls.finish()));
3638 }
3639
3640 #[test]
3641 fn test_large_slice_string_list_of_lists() {
3642 let mut ls = ListBuilder::new(ListBuilder::new(StringBuilder::new()));
3646
3647 for _ in 0..4000 {
3648 ls.values().append(true);
3649 ls.append(true)
3650 }
3651
3652 let mut s = String::new();
3653 for row_number in 0..4000 {
3654 if row_number % 2 == 0 {
3655 for list_element in 0..1000 {
3656 s.clear();
3657 use std::fmt::Write;
3658 write!(&mut s, "value{row_number}-{list_element}").unwrap();
3659 ls.values().values().append_value(&s);
3660 }
3661 ls.values().append(true);
3662 ls.append(true)
3663 } else {
3664 ls.append(false); }
3666 }
3667
3668 ensure_roundtrip(Arc::new(ls.finish()));
3669 }
3670
3671 fn ensure_roundtrip(array: ArrayRef) {
3673 let num_rows = array.len();
3674 let orig_batch = RecordBatch::try_from_iter(vec![("a", array)]).unwrap();
3675 let sliced_batch = orig_batch.slice(1, num_rows - 1);
3677
3678 let schema = orig_batch.schema();
3679 let stream_data = {
3680 let mut writer = StreamWriter::try_new(vec![], &schema).unwrap();
3681 writer.write(&sliced_batch).unwrap();
3682 writer.into_inner().unwrap()
3683 };
3684 let read_batch = {
3685 let projection = None;
3686 let mut reader = StreamReader::try_new(Cursor::new(stream_data), projection).unwrap();
3687 reader
3688 .next()
3689 .expect("expect no errors reading batch")
3690 .expect("expect batch")
3691 };
3692 assert_eq!(sliced_batch, read_batch);
3693
3694 let file_data = {
3695 let mut writer = FileWriter::try_new_buffered(vec![], &schema).unwrap();
3696 writer.write(&sliced_batch).unwrap();
3697 writer.into_inner().unwrap().into_inner().unwrap()
3698 };
3699 let read_batch = {
3700 let projection = None;
3701 let mut reader = FileReader::try_new(Cursor::new(file_data), projection).unwrap();
3702 reader
3703 .next()
3704 .expect("expect no errors reading batch")
3705 .expect("expect batch")
3706 };
3707 assert_eq!(sliced_batch, read_batch);
3708
3709 }
3711
3712 #[test]
3713 fn encode_bools_slice() {
3714 assert_bool_roundtrip([true, false], 1, 1);
3716
3717 assert_bool_roundtrip(
3719 [
3720 true, false, true, true, false, false, true, true, true, false, false, false, true,
3721 true, true, true, false, false, false, false, true, true, true, true, true, false,
3722 false, false, false, false,
3723 ],
3724 13,
3725 17,
3726 );
3727
3728 assert_bool_roundtrip(
3730 [
3731 true, false, true, true, false, false, true, true, true, false, false, false,
3732 ],
3733 8,
3734 2,
3735 );
3736
3737 assert_bool_roundtrip(
3739 [
3740 true, false, true, true, false, false, true, true, true, false, false, false, true,
3741 true, true, true, true, false, false, false, false, false,
3742 ],
3743 8,
3744 8,
3745 );
3746 }
3747
3748 fn assert_bool_roundtrip<const N: usize>(bools: [bool; N], offset: usize, length: usize) {
3749 let val_bool_field = Field::new("val", DataType::Boolean, false);
3750
3751 let schema = Arc::new(Schema::new(vec![val_bool_field]));
3752
3753 let bools = BooleanArray::from(bools.to_vec());
3754
3755 let batch = RecordBatch::try_new(Arc::clone(&schema), vec![Arc::new(bools)]).unwrap();
3756 let batch = batch.slice(offset, length);
3757
3758 let data = serialize_stream(&batch);
3759 let batch2 = deserialize_stream(data);
3760 assert_eq!(batch, batch2);
3761 }
3762
3763 #[test]
3764 fn test_run_array_unslice() {
3765 let total_len = 80;
3766 let vals: Vec<Option<i32>> = vec![Some(1), None, Some(2), Some(3), Some(4), None, Some(5)];
3767 let repeats: Vec<usize> = vec![3, 4, 1, 2];
3768 let mut input_array: Vec<Option<i32>> = Vec::with_capacity(total_len);
3769 for ix in 0_usize..32 {
3770 let repeat: usize = repeats[ix % repeats.len()];
3771 let val: Option<i32> = vals[ix % vals.len()];
3772 input_array.resize(input_array.len() + repeat, val);
3773 }
3774
3775 let mut builder =
3777 PrimitiveRunBuilder::<Int16Type, Int32Type>::with_capacity(input_array.len());
3778 builder.extend(input_array.iter().copied());
3779 let run_array = builder.finish();
3780
3781 for slice_len in 1..=total_len {
3783 let sliced_run_array: RunArray<Int16Type> =
3785 run_array.slice(0, slice_len).into_data().into();
3786
3787 let unsliced_run_array = into_zero_offset_run_array(sliced_run_array).unwrap();
3789 let typed = unsliced_run_array
3790 .downcast::<PrimitiveArray<Int32Type>>()
3791 .unwrap();
3792 let expected: Vec<Option<i32>> = input_array.iter().take(slice_len).copied().collect();
3793 let actual: Vec<Option<i32>> = typed.into_iter().collect();
3794 assert_eq!(expected, actual);
3795
3796 let sliced_run_array: RunArray<Int16Type> = run_array
3798 .slice(total_len - slice_len, slice_len)
3799 .into_data()
3800 .into();
3801
3802 let unsliced_run_array = into_zero_offset_run_array(sliced_run_array).unwrap();
3804 let typed = unsliced_run_array
3805 .downcast::<PrimitiveArray<Int32Type>>()
3806 .unwrap();
3807 let expected: Vec<Option<i32>> = input_array
3808 .iter()
3809 .skip(total_len - slice_len)
3810 .copied()
3811 .collect();
3812 let actual: Vec<Option<i32>> = typed.into_iter().collect();
3813 assert_eq!(expected, actual);
3814 }
3815 }
3816
3817 fn generate_list_data<O: OffsetSizeTrait>() -> GenericListArray<O> {
3818 let mut ls = GenericListBuilder::<O, _>::new(UInt32Builder::new());
3819
3820 for i in 0..100_000 {
3821 for value in [i, i, i] {
3822 ls.values().append_value(value);
3823 }
3824 ls.append(true)
3825 }
3826
3827 ls.finish()
3828 }
3829
3830 fn generate_utf8view_list_data<O: OffsetSizeTrait>() -> GenericListArray<O> {
3831 let mut ls = GenericListBuilder::<O, _>::new(StringViewBuilder::new());
3832
3833 for i in 0..100_000 {
3834 for value in [
3835 format!("value{}", i),
3836 format!("value{}", i),
3837 format!("value{}", i),
3838 ] {
3839 ls.values().append_value(&value);
3840 }
3841 ls.append(true)
3842 }
3843
3844 ls.finish()
3845 }
3846
3847 fn generate_string_list_data<O: OffsetSizeTrait>() -> GenericListArray<O> {
3848 let mut ls = GenericListBuilder::<O, _>::new(StringBuilder::new());
3849
3850 for i in 0..100_000 {
3851 for value in [
3852 format!("value{}", i),
3853 format!("value{}", i),
3854 format!("value{}", i),
3855 ] {
3856 ls.values().append_value(&value);
3857 }
3858 ls.append(true)
3859 }
3860
3861 ls.finish()
3862 }
3863
3864 fn generate_nested_list_data<O: OffsetSizeTrait>() -> GenericListArray<O> {
3865 let mut ls =
3866 GenericListBuilder::<O, _>::new(GenericListBuilder::<O, _>::new(UInt32Builder::new()));
3867
3868 for _i in 0..10_000 {
3869 for j in 0..10 {
3870 for value in [j, j, j, j] {
3871 ls.values().values().append_value(value);
3872 }
3873 ls.values().append(true)
3874 }
3875 ls.append(true);
3876 }
3877
3878 ls.finish()
3879 }
3880
3881 fn generate_nested_list_data_starting_at_zero<O: OffsetSizeTrait>() -> GenericListArray<O> {
3882 let mut ls =
3883 GenericListBuilder::<O, _>::new(GenericListBuilder::<O, _>::new(UInt32Builder::new()));
3884
3885 for _i in 0..999 {
3886 ls.values().append(true);
3887 ls.append(true);
3888 }
3889
3890 for j in 0..10 {
3891 for value in [j, j, j, j] {
3892 ls.values().values().append_value(value);
3893 }
3894 ls.values().append(true)
3895 }
3896 ls.append(true);
3897
3898 for i in 0..9_000 {
3899 for j in 0..10 {
3900 for value in [i + j, i + j, i + j, i + j] {
3901 ls.values().values().append_value(value);
3902 }
3903 ls.values().append(true)
3904 }
3905 ls.append(true);
3906 }
3907
3908 ls.finish()
3909 }
3910
3911 fn generate_map_array_data() -> MapArray {
3912 let keys_builder = UInt32Builder::new();
3913 let values_builder = UInt32Builder::new();
3914
3915 let mut builder = MapBuilder::new(None, keys_builder, values_builder);
3916
3917 for i in 0..100_000 {
3918 for _j in 0..3 {
3919 builder.keys().append_value(i);
3920 builder.values().append_value(i * 2);
3921 }
3922 builder.append(true).unwrap();
3923 }
3924
3925 builder.finish()
3926 }
3927
3928 #[test]
3929 fn reencode_offsets_when_first_offset_is_not_zero() {
3930 let original_list = generate_list_data::<i32>();
3931 let original_data = original_list.into_data();
3932 let slice_data = original_data.slice(75, 7);
3933 let (new_offsets, original_start, length) =
3934 reencode_offsets::<i32>(&slice_data.buffers()[0], &slice_data);
3935 assert_eq!(
3936 vec![0, 3, 6, 9, 12, 15, 18, 21],
3937 new_offsets.typed_data::<i32>()
3938 );
3939 assert_eq!(225, original_start);
3940 assert_eq!(21, length);
3941 }
3942
3943 #[test]
3944 fn reencode_offsets_when_first_offset_is_zero() {
3945 let mut ls = GenericListBuilder::<i32, _>::new(UInt32Builder::new());
3946 ls.append(true);
3948 ls.values().append_value(35);
3949 ls.values().append_value(42);
3950 ls.append(true);
3951 let original_list = ls.finish();
3952 let original_data = original_list.into_data();
3953
3954 let slice_data = original_data.slice(1, 1);
3955 let (new_offsets, original_start, length) =
3956 reencode_offsets::<i32>(&slice_data.buffers()[0], &slice_data);
3957 assert_eq!(vec![0, 2], new_offsets.typed_data::<i32>());
3958 assert_eq!(0, original_start);
3959 assert_eq!(2, length);
3960 }
3961
3962 fn roundtrip_ensure_sliced_smaller(in_batch: RecordBatch, expected_size_factor: usize) {
3965 let in_sliced = in_batch.slice(999, 1);
3967
3968 let bytes_batch = serialize_file(&in_batch);
3969 let bytes_sliced = serialize_file(&in_sliced);
3970
3971 assert!(bytes_sliced.len() < (bytes_batch.len() / expected_size_factor));
3973
3974 let out_batch = deserialize_file(bytes_batch);
3976 assert_eq!(in_batch, out_batch);
3977
3978 let out_sliced = deserialize_file(bytes_sliced);
3979 assert_eq!(in_sliced, out_sliced);
3980 }
3981
3982 #[test]
3983 fn encode_lists() {
3984 let val_inner = Field::new_list_field(DataType::UInt32, true);
3985 let val_list_field = Field::new("val", DataType::List(Arc::new(val_inner)), false);
3986 let schema = Arc::new(Schema::new(vec![val_list_field]));
3987
3988 let values = Arc::new(generate_list_data::<i32>());
3989
3990 let in_batch = RecordBatch::try_new(schema, vec![values]).unwrap();
3991 roundtrip_ensure_sliced_smaller(in_batch, 1000);
3992 }
3993
3994 #[test]
3995 fn encode_empty_list() {
3996 let val_inner = Field::new_list_field(DataType::UInt32, true);
3997 let val_list_field = Field::new("val", DataType::List(Arc::new(val_inner)), false);
3998 let schema = Arc::new(Schema::new(vec![val_list_field]));
3999
4000 let values = Arc::new(generate_list_data::<i32>());
4001
4002 let in_batch = RecordBatch::try_new(schema, vec![values])
4003 .unwrap()
4004 .slice(999, 0);
4005 let out_batch = deserialize_file(serialize_file(&in_batch));
4006 assert_eq!(in_batch, out_batch);
4007 }
4008
4009 #[test]
4010 fn encode_large_lists() {
4011 let val_inner = Field::new_list_field(DataType::UInt32, true);
4012 let val_list_field = Field::new("val", DataType::LargeList(Arc::new(val_inner)), false);
4013 let schema = Arc::new(Schema::new(vec![val_list_field]));
4014
4015 let values = Arc::new(generate_list_data::<i64>());
4016
4017 let in_batch = RecordBatch::try_new(schema, vec![values]).unwrap();
4020 roundtrip_ensure_sliced_smaller(in_batch, 1000);
4021 }
4022
4023 #[test]
4024 fn encode_large_lists_non_zero_offset() {
4025 let val_inner = Field::new_list_field(DataType::UInt32, true);
4026 let val_list_field = Field::new("val", DataType::LargeList(Arc::new(val_inner)), false);
4027 let schema = Arc::new(Schema::new(vec![val_list_field]));
4028
4029 let values = Arc::new(generate_list_data::<i64>());
4030
4031 check_sliced_list_array(schema, values);
4032 }
4033
4034 #[test]
4035 fn encode_large_lists_string_non_zero_offset() {
4036 let val_inner = Field::new_list_field(DataType::Utf8, true);
4037 let val_list_field = Field::new("val", DataType::LargeList(Arc::new(val_inner)), false);
4038 let schema = Arc::new(Schema::new(vec![val_list_field]));
4039
4040 let values = Arc::new(generate_string_list_data::<i64>());
4041
4042 check_sliced_list_array(schema, values);
4043 }
4044
4045 #[test]
4046 fn encode_large_list_string_view_non_zero_offset() {
4047 let val_inner = Field::new_list_field(DataType::Utf8View, true);
4048 let val_list_field = Field::new("val", DataType::LargeList(Arc::new(val_inner)), false);
4049 let schema = Arc::new(Schema::new(vec![val_list_field]));
4050
4051 let values = Arc::new(generate_utf8view_list_data::<i64>());
4052
4053 check_sliced_list_array(schema, values);
4054 }
4055
4056 fn check_sliced_list_array(schema: Arc<Schema>, values: Arc<GenericListArray<i64>>) {
4057 for (offset, len) in [(999, 1), (0, 13), (47, 12), (values.len() - 13, 13)] {
4058 let in_batch = RecordBatch::try_new(schema.clone(), vec![values.clone()])
4059 .unwrap()
4060 .slice(offset, len);
4061 let out_batch = deserialize_file(serialize_file(&in_batch));
4062 assert_eq!(in_batch, out_batch);
4063 }
4064 }
4065
4066 #[test]
4067 fn encode_nested_lists() {
4068 let inner_int = Arc::new(Field::new_list_field(DataType::UInt32, true));
4069 let inner_list_field = Arc::new(Field::new_list_field(DataType::List(inner_int), true));
4070 let list_field = Field::new("val", DataType::List(inner_list_field), true);
4071 let schema = Arc::new(Schema::new(vec![list_field]));
4072
4073 let values = Arc::new(generate_nested_list_data::<i32>());
4074
4075 let in_batch = RecordBatch::try_new(schema, vec![values]).unwrap();
4076 roundtrip_ensure_sliced_smaller(in_batch, 1000);
4077 }
4078
4079 #[test]
4080 fn encode_nested_lists_starting_at_zero() {
4081 let inner_int = Arc::new(Field::new("item", DataType::UInt32, true));
4082 let inner_list_field = Arc::new(Field::new("item", DataType::List(inner_int), true));
4083 let list_field = Field::new("val", DataType::List(inner_list_field), true);
4084 let schema = Arc::new(Schema::new(vec![list_field]));
4085
4086 let values = Arc::new(generate_nested_list_data_starting_at_zero::<i32>());
4087
4088 let in_batch = RecordBatch::try_new(schema, vec![values]).unwrap();
4089 roundtrip_ensure_sliced_smaller(in_batch, 1);
4090 }
4091
4092 #[test]
4093 fn encode_map_array() {
4094 let keys = Arc::new(Field::new("keys", DataType::UInt32, false));
4095 let values = Arc::new(Field::new("values", DataType::UInt32, true));
4096 let map_field = Field::new_map("map", "entries", keys, values, false, true);
4097 let schema = Arc::new(Schema::new(vec![map_field]));
4098
4099 let values = Arc::new(generate_map_array_data());
4100
4101 let in_batch = RecordBatch::try_new(schema, vec![values]).unwrap();
4102 roundtrip_ensure_sliced_smaller(in_batch, 1000);
4103 }
4104
4105 fn generate_list_view_data<O: OffsetSizeTrait>() -> GenericListViewArray<O> {
4106 let mut builder = GenericListViewBuilder::<O, _>::new(UInt32Builder::new());
4107
4108 for i in 0u32..100_000 {
4109 if i.is_multiple_of(10_000) {
4110 builder.append(false);
4111 continue;
4112 }
4113 for value in [i, i, i] {
4114 builder.values().append_value(value);
4115 }
4116 builder.append(true);
4117 }
4118
4119 builder.finish()
4120 }
4121
4122 #[test]
4123 fn encode_list_view_arrays() {
4124 let val_inner = Field::new_list_field(DataType::UInt32, true);
4125 let val_field = Field::new("val", DataType::ListView(Arc::new(val_inner)), true);
4126 let schema = Arc::new(Schema::new(vec![val_field]));
4127
4128 let values = Arc::new(generate_list_view_data::<i32>());
4129
4130 let in_batch = RecordBatch::try_new(schema, vec![values]).unwrap();
4131 let out_batch = deserialize_file(serialize_file(&in_batch));
4132 assert_eq!(in_batch, out_batch);
4133 }
4134
4135 #[test]
4136 fn encode_large_list_view_arrays() {
4137 let val_inner = Field::new_list_field(DataType::UInt32, true);
4138 let val_field = Field::new("val", DataType::LargeListView(Arc::new(val_inner)), true);
4139 let schema = Arc::new(Schema::new(vec![val_field]));
4140
4141 let values = Arc::new(generate_list_view_data::<i64>());
4142
4143 let in_batch = RecordBatch::try_new(schema, vec![values]).unwrap();
4144 let out_batch = deserialize_file(serialize_file(&in_batch));
4145 assert_eq!(in_batch, out_batch);
4146 }
4147
4148 #[test]
4149 fn check_sliced_list_view_array() {
4150 let inner = Field::new_list_field(DataType::UInt32, true);
4151 let field = Field::new("val", DataType::ListView(Arc::new(inner)), true);
4152 let schema = Arc::new(Schema::new(vec![field]));
4153 let values = Arc::new(generate_list_view_data::<i32>());
4154
4155 for (offset, len) in [(999, 1), (0, 13), (47, 12), (values.len() - 13, 13)] {
4156 let in_batch = RecordBatch::try_new(schema.clone(), vec![values.clone()])
4157 .unwrap()
4158 .slice(offset, len);
4159 let out_batch = deserialize_file(serialize_file(&in_batch));
4160 assert_eq!(in_batch, out_batch);
4161 }
4162 }
4163
4164 #[test]
4165 fn check_sliced_large_list_view_array() {
4166 let inner = Field::new_list_field(DataType::UInt32, true);
4167 let field = Field::new("val", DataType::LargeListView(Arc::new(inner)), true);
4168 let schema = Arc::new(Schema::new(vec![field]));
4169 let values = Arc::new(generate_list_view_data::<i64>());
4170
4171 for (offset, len) in [(999, 1), (0, 13), (47, 12), (values.len() - 13, 13)] {
4172 let in_batch = RecordBatch::try_new(schema.clone(), vec![values.clone()])
4173 .unwrap()
4174 .slice(offset, len);
4175 let out_batch = deserialize_file(serialize_file(&in_batch));
4176 assert_eq!(in_batch, out_batch);
4177 }
4178 }
4179
4180 fn generate_nested_list_view_data<O: OffsetSizeTrait>() -> GenericListViewArray<O> {
4181 let inner_builder = UInt32Builder::new();
4182 let middle_builder = GenericListViewBuilder::<O, _>::new(inner_builder);
4183 let mut outer_builder = GenericListViewBuilder::<O, _>::new(middle_builder);
4184
4185 for i in 0u32..10_000 {
4186 if i.is_multiple_of(1_000) {
4187 outer_builder.append(false);
4188 continue;
4189 }
4190
4191 for _ in 0..3 {
4192 for value in [i, i + 1, i + 2] {
4193 outer_builder.values().values().append_value(value);
4194 }
4195 outer_builder.values().append(true);
4196 }
4197 outer_builder.append(true);
4198 }
4199
4200 outer_builder.finish()
4201 }
4202
4203 #[test]
4204 fn encode_nested_list_views() {
4205 let inner_int = Arc::new(Field::new_list_field(DataType::UInt32, true));
4206 let inner_list_field = Arc::new(Field::new_list_field(DataType::ListView(inner_int), true));
4207 let list_field = Field::new("val", DataType::ListView(inner_list_field), true);
4208 let schema = Arc::new(Schema::new(vec![list_field]));
4209
4210 let values = Arc::new(generate_nested_list_view_data::<i32>());
4211
4212 let in_batch = RecordBatch::try_new(schema, vec![values]).unwrap();
4213 let out_batch = deserialize_file(serialize_file(&in_batch));
4214 assert_eq!(in_batch, out_batch);
4215 }
4216
4217 fn test_roundtrip_list_view_of_dict_impl<OffsetSize: OffsetSizeTrait, U: ArrowNativeType>(
4218 list_data_type: DataType,
4219 offsets: &[U; 5],
4220 sizes: &[U; 4],
4221 ) {
4222 let values = StringArray::from(vec![Some("alpha"), None, Some("beta"), Some("gamma")]);
4223 let keys = Int32Array::from_iter_values([0, 0, 1, 2, 3, 0, 2]);
4224 let dict_array = DictionaryArray::new(keys, Arc::new(values));
4225 let dict_data = dict_array.to_data();
4226
4227 let value_offsets = Buffer::from_slice_ref(offsets);
4228 let value_sizes = Buffer::from_slice_ref(sizes);
4229
4230 let list_data = ArrayData::builder(list_data_type)
4231 .len(4)
4232 .add_buffer(value_offsets)
4233 .add_buffer(value_sizes)
4234 .add_child_data(dict_data)
4235 .build()
4236 .unwrap();
4237 let list_view_array = GenericListViewArray::<OffsetSize>::from(list_data);
4238
4239 let schema = Arc::new(Schema::new(vec![Field::new(
4240 "f1",
4241 list_view_array.data_type().clone(),
4242 false,
4243 )]));
4244 let input_batch = RecordBatch::try_new(schema, vec![Arc::new(list_view_array)]).unwrap();
4245
4246 let output_batch = deserialize_file(serialize_file(&input_batch));
4247 assert_eq!(input_batch, output_batch);
4248
4249 let output_batch = deserialize_stream(serialize_stream(&input_batch));
4250 assert_eq!(input_batch, output_batch);
4251 }
4252
4253 #[test]
4254 fn test_roundtrip_list_view_of_dict() {
4255 #[allow(deprecated)]
4256 let list_data_type = DataType::ListView(Arc::new(Field::new_dict(
4257 "item",
4258 DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
4259 true,
4260 1,
4261 false,
4262 )));
4263 let offsets: &[i32; 5] = &[0, 2, 4, 4, 7];
4264 let sizes: &[i32; 4] = &[2, 2, 0, 3];
4265 test_roundtrip_list_view_of_dict_impl::<i32, i32>(list_data_type, offsets, sizes);
4266 }
4267
4268 #[test]
4269 fn test_roundtrip_large_list_view_of_dict() {
4270 #[allow(deprecated)]
4271 let list_data_type = DataType::LargeListView(Arc::new(Field::new_dict(
4272 "item",
4273 DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
4274 true,
4275 2,
4276 false,
4277 )));
4278 let offsets: &[i64; 5] = &[0, 2, 4, 4, 7];
4279 let sizes: &[i64; 4] = &[2, 2, 0, 3];
4280 test_roundtrip_list_view_of_dict_impl::<i64, i64>(list_data_type, offsets, sizes);
4281 }
4282
4283 #[test]
4284 fn test_roundtrip_sliced_list_view_of_dict() {
4285 #[allow(deprecated)]
4286 let list_data_type = DataType::ListView(Arc::new(Field::new_dict(
4287 "item",
4288 DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
4289 true,
4290 3,
4291 false,
4292 )));
4293
4294 let values = StringArray::from(vec![Some("alpha"), None, Some("beta"), Some("gamma")]);
4295 let keys = Int32Array::from_iter_values([0, 0, 1, 2, 3, 0, 2, 1, 0, 3, 2, 1]);
4296 let dict_array = DictionaryArray::new(keys, Arc::new(values));
4297 let dict_data = dict_array.to_data();
4298
4299 let offsets: &[i32; 7] = &[0, 2, 4, 4, 7, 9, 12];
4300 let sizes: &[i32; 6] = &[2, 2, 0, 3, 2, 3];
4301 let value_offsets = Buffer::from_slice_ref(offsets);
4302 let value_sizes = Buffer::from_slice_ref(sizes);
4303
4304 let list_data = ArrayData::builder(list_data_type)
4305 .len(6)
4306 .add_buffer(value_offsets)
4307 .add_buffer(value_sizes)
4308 .add_child_data(dict_data)
4309 .build()
4310 .unwrap();
4311 let list_view_array = GenericListViewArray::<i32>::from(list_data);
4312
4313 let schema = Arc::new(Schema::new(vec![Field::new(
4314 "f1",
4315 list_view_array.data_type().clone(),
4316 false,
4317 )]));
4318 let input_batch = RecordBatch::try_new(schema, vec![Arc::new(list_view_array)]).unwrap();
4319
4320 let sliced_batch = input_batch.slice(1, 4);
4321
4322 let output_batch = deserialize_file(serialize_file(&sliced_batch));
4323 assert_eq!(sliced_batch, output_batch);
4324
4325 let output_batch = deserialize_stream(serialize_stream(&sliced_batch));
4326 assert_eq!(sliced_batch, output_batch);
4327 }
4328
4329 #[test]
4330 fn test_roundtrip_dense_union_of_dict() {
4331 let values = StringArray::from(vec![Some("alpha"), None, Some("beta"), Some("gamma")]);
4332 let keys = Int32Array::from_iter_values([0, 0, 1, 2, 3, 0, 2]);
4333 let dict_array = DictionaryArray::new(keys, Arc::new(values));
4334
4335 #[allow(deprecated)]
4336 let dict_field = Arc::new(Field::new_dict(
4337 "dict",
4338 DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
4339 true,
4340 1,
4341 false,
4342 ));
4343 let int_field = Arc::new(Field::new("int", DataType::Int32, false));
4344 let union_fields = UnionFields::try_new(vec![0, 1], vec![dict_field, int_field]).unwrap();
4345
4346 let types = ScalarBuffer::from(vec![0i8, 0, 1, 0, 1, 0, 0]);
4347 let offsets = ScalarBuffer::from(vec![0i32, 1, 0, 2, 1, 3, 4]);
4348
4349 let int_array = Int32Array::from(vec![100, 200]);
4350
4351 let union = UnionArray::try_new(
4352 union_fields.clone(),
4353 types,
4354 Some(offsets),
4355 vec![Arc::new(dict_array), Arc::new(int_array)],
4356 )
4357 .unwrap();
4358
4359 let schema = Arc::new(Schema::new(vec![Field::new(
4360 "union",
4361 DataType::Union(union_fields, UnionMode::Dense),
4362 false,
4363 )]));
4364 let input_batch = RecordBatch::try_new(schema, vec![Arc::new(union)]).unwrap();
4365
4366 let output_batch = deserialize_file(serialize_file(&input_batch));
4367 assert_eq!(input_batch, output_batch);
4368
4369 let output_batch = deserialize_stream(serialize_stream(&input_batch));
4370 assert_eq!(input_batch, output_batch);
4371 }
4372
4373 #[test]
4374 fn test_roundtrip_sparse_union_of_dict() {
4375 let values = StringArray::from(vec![Some("alpha"), None, Some("beta"), Some("gamma")]);
4376 let keys = Int32Array::from_iter_values([0, 0, 1, 2, 3, 0, 2]);
4377 let dict_array = DictionaryArray::new(keys, Arc::new(values));
4378
4379 #[allow(deprecated)]
4380 let dict_field = Arc::new(Field::new_dict(
4381 "dict",
4382 DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
4383 true,
4384 2,
4385 false,
4386 ));
4387 let int_field = Arc::new(Field::new("int", DataType::Int32, false));
4388 let union_fields = UnionFields::try_new(vec![0, 1], vec![dict_field, int_field]).unwrap();
4389
4390 let types = ScalarBuffer::from(vec![0i8, 0, 1, 0, 1, 0, 0]);
4391
4392 let int_array = Int32Array::from(vec![0, 0, 100, 0, 200, 0, 0]);
4393
4394 let union = UnionArray::try_new(
4395 union_fields.clone(),
4396 types,
4397 None,
4398 vec![Arc::new(dict_array), Arc::new(int_array)],
4399 )
4400 .unwrap();
4401
4402 let schema = Arc::new(Schema::new(vec![Field::new(
4403 "union",
4404 DataType::Union(union_fields, UnionMode::Sparse),
4405 false,
4406 )]));
4407 let input_batch = RecordBatch::try_new(schema, vec![Arc::new(union)]).unwrap();
4408
4409 let output_batch = deserialize_file(serialize_file(&input_batch));
4410 assert_eq!(input_batch, output_batch);
4411
4412 let output_batch = deserialize_stream(serialize_stream(&input_batch));
4413 assert_eq!(input_batch, output_batch);
4414 }
4415
4416 #[test]
4417 fn test_roundtrip_map_with_dict_keys() {
4418 let key_values = StringArray::from(vec!["key_a", "key_b", "key_c"]);
4421 let keys = Int32Array::from_iter_values([0, 1, 2, 0, 1, 0]);
4422 let dict_keys = DictionaryArray::new(keys, Arc::new(key_values));
4423
4424 let values = Int32Array::from(vec![1, 2, 3, 4, 5, 6]);
4425
4426 #[allow(deprecated)]
4427 let entries_field = Arc::new(Field::new(
4428 "entries",
4429 DataType::Struct(
4430 vec![
4431 Field::new_dict(
4432 "key",
4433 DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
4434 false,
4435 1,
4436 false,
4437 ),
4438 Field::new("value", DataType::Int32, true),
4439 ]
4440 .into(),
4441 ),
4442 false,
4443 ));
4444
4445 let entries = StructArray::from(vec![
4446 (
4447 Arc::new(Field::new(
4448 "key",
4449 DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
4450 false,
4451 )),
4452 Arc::new(dict_keys) as ArrayRef,
4453 ),
4454 (
4455 Arc::new(Field::new("value", DataType::Int32, true)),
4456 Arc::new(values) as ArrayRef,
4457 ),
4458 ]);
4459
4460 let offsets = Buffer::from_slice_ref([0i32, 2, 4, 6]);
4461
4462 let map_data = ArrayData::builder(DataType::Map(entries_field, false))
4463 .len(3)
4464 .add_buffer(offsets)
4465 .add_child_data(entries.into_data())
4466 .build()
4467 .unwrap();
4468 let map_array = MapArray::from(map_data);
4469
4470 let schema = Arc::new(Schema::new(vec![Field::new(
4471 "map",
4472 map_array.data_type().clone(),
4473 false,
4474 )]));
4475 let input_batch = RecordBatch::try_new(schema, vec![Arc::new(map_array)]).unwrap();
4476
4477 let output_batch = deserialize_file(serialize_file(&input_batch));
4478 assert_eq!(input_batch, output_batch);
4479
4480 let output_batch = deserialize_stream(serialize_stream(&input_batch));
4481 assert_eq!(input_batch, output_batch);
4482 }
4483
4484 #[test]
4485 fn test_roundtrip_map_with_dict_values() {
4486 let keys = StringArray::from(vec!["a", "b", "c", "d", "e", "f"]);
4489
4490 let value_values = StringArray::from(vec!["val_x", "val_y", "val_z"]);
4491 let value_keys = Int32Array::from_iter_values([0, 1, 2, 0, 1, 0]);
4492 let dict_values = DictionaryArray::new(value_keys, Arc::new(value_values));
4493
4494 #[allow(deprecated)]
4495 let entries_field = Arc::new(Field::new(
4496 "entries",
4497 DataType::Struct(
4498 vec![
4499 Field::new("key", DataType::Utf8, false),
4500 Field::new_dict(
4501 "value",
4502 DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
4503 true,
4504 2,
4505 false,
4506 ),
4507 ]
4508 .into(),
4509 ),
4510 false,
4511 ));
4512
4513 let entries = StructArray::from(vec![
4514 (
4515 Arc::new(Field::new("key", DataType::Utf8, false)),
4516 Arc::new(keys) as ArrayRef,
4517 ),
4518 (
4519 Arc::new(Field::new(
4520 "value",
4521 DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
4522 true,
4523 )),
4524 Arc::new(dict_values) as ArrayRef,
4525 ),
4526 ]);
4527
4528 let offsets = Buffer::from_slice_ref([0i32, 2, 4, 6]);
4529
4530 let map_data = ArrayData::builder(DataType::Map(entries_field, false))
4531 .len(3)
4532 .add_buffer(offsets)
4533 .add_child_data(entries.into_data())
4534 .build()
4535 .unwrap();
4536 let map_array = MapArray::from(map_data);
4537
4538 let schema = Arc::new(Schema::new(vec![Field::new(
4539 "map",
4540 map_array.data_type().clone(),
4541 false,
4542 )]));
4543 let input_batch = RecordBatch::try_new(schema, vec![Arc::new(map_array)]).unwrap();
4544
4545 let output_batch = deserialize_file(serialize_file(&input_batch));
4546 assert_eq!(input_batch, output_batch);
4547
4548 let output_batch = deserialize_stream(serialize_stream(&input_batch));
4549 assert_eq!(input_batch, output_batch);
4550 }
4551
4552 #[test]
4553 fn test_decimal128_alignment16_is_sufficient() {
4554 const IPC_ALIGNMENT: usize = 16;
4555
4556 for num_cols in [1, 2, 3, 17, 50, 73, 99] {
4561 let num_rows = (num_cols * 7 + 11) % 100; let mut fields = Vec::new();
4564 let mut arrays = Vec::new();
4565 for i in 0..num_cols {
4566 let field = Field::new(format!("col_{i}"), DataType::Decimal128(38, 10), true);
4567 let array = Decimal128Array::from(vec![num_cols as i128; num_rows]);
4568 fields.push(field);
4569 arrays.push(Arc::new(array) as Arc<dyn Array>);
4570 }
4571 let schema = Schema::new(fields);
4572 let batch = RecordBatch::try_new(Arc::new(schema), arrays).unwrap();
4573
4574 let mut writer = FileWriter::try_new_with_options(
4575 Vec::new(),
4576 batch.schema_ref(),
4577 IpcWriteOptions::try_new(IPC_ALIGNMENT, false, MetadataVersion::V5).unwrap(),
4578 )
4579 .unwrap();
4580 writer.write(&batch).unwrap();
4581 writer.finish().unwrap();
4582
4583 let out: Vec<u8> = writer.into_inner().unwrap();
4584
4585 let buffer = Buffer::from_vec(out);
4586 let trailer_start = buffer.len() - 10;
4587 let footer_len =
4588 read_footer_length(buffer[trailer_start..].try_into().unwrap()).unwrap();
4589 let footer =
4590 root_as_footer(&buffer[trailer_start - footer_len..trailer_start]).unwrap();
4591
4592 let schema = fb_to_schema(footer.schema().unwrap());
4593
4594 let decoder =
4597 FileDecoder::new(Arc::new(schema), footer.version()).with_require_alignment(true);
4598
4599 let batches = footer.recordBatches().unwrap();
4600
4601 let block = batches.get(0);
4602 let block_len = block.bodyLength() as usize + block.metaDataLength() as usize;
4603 let data = buffer.slice_with_length(block.offset() as _, block_len);
4604
4605 let batch2 = decoder.read_record_batch(block, &data).unwrap().unwrap();
4606
4607 assert_eq!(batch, batch2);
4608 }
4609 }
4610
4611 #[test]
4612 fn test_decimal128_alignment8_is_unaligned() {
4613 const IPC_ALIGNMENT: usize = 8;
4614
4615 let num_cols = 2;
4616 let num_rows = 1;
4617
4618 let mut fields = Vec::new();
4619 let mut arrays = Vec::new();
4620 for i in 0..num_cols {
4621 let field = Field::new(format!("col_{i}"), DataType::Decimal128(38, 10), true);
4622 let array = Decimal128Array::from(vec![num_cols as i128; num_rows]);
4623 fields.push(field);
4624 arrays.push(Arc::new(array) as Arc<dyn Array>);
4625 }
4626 let schema = Schema::new(fields);
4627 let batch = RecordBatch::try_new(Arc::new(schema), arrays).unwrap();
4628
4629 let mut writer = FileWriter::try_new_with_options(
4630 Vec::new(),
4631 batch.schema_ref(),
4632 IpcWriteOptions::try_new(IPC_ALIGNMENT, false, MetadataVersion::V5).unwrap(),
4633 )
4634 .unwrap();
4635 writer.write(&batch).unwrap();
4636 writer.finish().unwrap();
4637
4638 let out: Vec<u8> = writer.into_inner().unwrap();
4639
4640 let buffer = Buffer::from_vec(out);
4641 let trailer_start = buffer.len() - 10;
4642 let footer_len = read_footer_length(buffer[trailer_start..].try_into().unwrap()).unwrap();
4643 let footer = root_as_footer(&buffer[trailer_start - footer_len..trailer_start]).unwrap();
4644 let schema = fb_to_schema(footer.schema().unwrap());
4645
4646 let decoder =
4649 FileDecoder::new(Arc::new(schema), footer.version()).with_require_alignment(true);
4650
4651 let batches = footer.recordBatches().unwrap();
4652
4653 let block = batches.get(0);
4654 let block_len = block.bodyLength() as usize + block.metaDataLength() as usize;
4655 let data = buffer.slice_with_length(block.offset() as _, block_len);
4656
4657 let result = decoder.read_record_batch(block, &data);
4658
4659 let error = result.unwrap_err();
4660 assert_eq!(
4661 error.to_string(),
4662 "Invalid argument error: Misaligned buffers[0] in array of type Decimal128(38, 10), \
4663 offset from expected alignment of 16 by 8"
4664 );
4665 }
4666
4667 #[test]
4668 fn test_flush() {
4669 let num_cols = 2;
4672 let mut fields = Vec::new();
4673 let options = IpcWriteOptions::try_new(8, false, MetadataVersion::V5).unwrap();
4674 for i in 0..num_cols {
4675 let field = Field::new(format!("col_{i}"), DataType::Decimal128(38, 10), true);
4676 fields.push(field);
4677 }
4678 let schema = Schema::new(fields);
4679 let inner_stream_writer = BufWriter::with_capacity(1024, Vec::new());
4680 let inner_file_writer = BufWriter::with_capacity(1024, Vec::new());
4681 let mut stream_writer =
4682 StreamWriter::try_new_with_options(inner_stream_writer, &schema, options.clone())
4683 .unwrap();
4684 let mut file_writer =
4685 FileWriter::try_new_with_options(inner_file_writer, &schema, options).unwrap();
4686
4687 let stream_bytes_written_on_new = stream_writer.get_ref().get_ref().len();
4688 let file_bytes_written_on_new = file_writer.get_ref().get_ref().len();
4689 stream_writer.flush().unwrap();
4690 file_writer.flush().unwrap();
4691 let stream_bytes_written_on_flush = stream_writer.get_ref().get_ref().len();
4692 let file_bytes_written_on_flush = file_writer.get_ref().get_ref().len();
4693 let stream_out = stream_writer.into_inner().unwrap().into_inner().unwrap();
4694 let expected_stream_flushed_bytes = stream_out.len() - 8;
4698 let expected_file_flushed_bytes = expected_stream_flushed_bytes + 8;
4701
4702 assert!(
4703 stream_bytes_written_on_new < stream_bytes_written_on_flush,
4704 "this test makes no sense if flush is not actually required"
4705 );
4706 assert!(
4707 file_bytes_written_on_new < file_bytes_written_on_flush,
4708 "this test makes no sense if flush is not actually required"
4709 );
4710 assert_eq!(stream_bytes_written_on_flush, expected_stream_flushed_bytes);
4711 assert_eq!(file_bytes_written_on_flush, expected_file_flushed_bytes);
4712 }
4713
4714 #[test]
4715 fn test_roundtrip_list_of_fixed_list() -> Result<(), ArrowError> {
4716 let l1_type =
4717 DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Float32, false)), 3);
4718 let l2_type = DataType::List(Arc::new(Field::new("item", l1_type.clone(), false)));
4719
4720 let l0_builder = Float32Builder::new();
4721 let l1_builder = FixedSizeListBuilder::new(l0_builder, 3).with_field(Arc::new(Field::new(
4722 "item",
4723 DataType::Float32,
4724 false,
4725 )));
4726 let mut l2_builder =
4727 ListBuilder::new(l1_builder).with_field(Arc::new(Field::new("item", l1_type, false)));
4728
4729 for point in [[1.0, 2.0, 3.0], [4.0, 5.0, 6.0], [7.0, 8.0, 9.0]] {
4730 l2_builder.values().values().append_value(point[0]);
4731 l2_builder.values().values().append_value(point[1]);
4732 l2_builder.values().values().append_value(point[2]);
4733
4734 l2_builder.values().append(true);
4735 }
4736 l2_builder.append(true);
4737
4738 let point = [10., 11., 12.];
4739 l2_builder.values().values().append_value(point[0]);
4740 l2_builder.values().values().append_value(point[1]);
4741 l2_builder.values().values().append_value(point[2]);
4742
4743 l2_builder.values().append(true);
4744 l2_builder.append(true);
4745
4746 let array = Arc::new(l2_builder.finish()) as ArrayRef;
4747
4748 let schema = Arc::new(Schema::new_with_metadata(
4749 vec![Field::new("points", l2_type, false)],
4750 HashMap::default(),
4751 ));
4752
4753 test_slices(&array, &schema, 0, 1)?;
4756 test_slices(&array, &schema, 0, 2)?;
4757 test_slices(&array, &schema, 1, 1)?;
4758
4759 Ok(())
4760 }
4761
4762 #[test]
4763 fn test_roundtrip_list_of_fixed_list_w_nulls() -> Result<(), ArrowError> {
4764 let l0_builder = Float32Builder::new();
4765 let l1_builder = FixedSizeListBuilder::new(l0_builder, 3);
4766 let mut l2_builder = ListBuilder::new(l1_builder);
4767
4768 for point in [
4769 [Some(1.0), Some(2.0), None],
4770 [Some(4.0), Some(5.0), Some(6.0)],
4771 [None, Some(8.0), Some(9.0)],
4772 ] {
4773 for p in point {
4774 match p {
4775 Some(p) => l2_builder.values().values().append_value(p),
4776 None => l2_builder.values().values().append_null(),
4777 }
4778 }
4779
4780 l2_builder.values().append(true);
4781 }
4782 l2_builder.append(true);
4783
4784 let point = [Some(10.), None, None];
4785 for p in point {
4786 match p {
4787 Some(p) => l2_builder.values().values().append_value(p),
4788 None => l2_builder.values().values().append_null(),
4789 }
4790 }
4791
4792 l2_builder.values().append(true);
4793 l2_builder.append(true);
4794
4795 let array = Arc::new(l2_builder.finish()) as ArrayRef;
4796
4797 let schema = Arc::new(Schema::new_with_metadata(
4798 vec![Field::new(
4799 "points",
4800 DataType::List(Arc::new(Field::new(
4801 "item",
4802 DataType::FixedSizeList(
4803 Arc::new(Field::new("item", DataType::Float32, true)),
4804 3,
4805 ),
4806 true,
4807 ))),
4808 true,
4809 )],
4810 HashMap::default(),
4811 ));
4812
4813 test_slices(&array, &schema, 0, 1)?;
4816 test_slices(&array, &schema, 0, 2)?;
4817 test_slices(&array, &schema, 1, 1)?;
4818
4819 Ok(())
4820 }
4821
4822 fn test_slices(
4823 parent_array: &ArrayRef,
4824 schema: &SchemaRef,
4825 offset: usize,
4826 length: usize,
4827 ) -> Result<(), ArrowError> {
4828 let subarray = parent_array.slice(offset, length);
4829 let original_batch = RecordBatch::try_new(schema.clone(), vec![subarray])?;
4830
4831 let mut bytes = Vec::new();
4832 let mut writer = StreamWriter::try_new(&mut bytes, schema)?;
4833 writer.write(&original_batch)?;
4834 writer.finish()?;
4835
4836 let mut cursor = std::io::Cursor::new(bytes);
4837 let mut reader = StreamReader::try_new(&mut cursor, None)?;
4838 let returned_batch = reader.next().unwrap()?;
4839
4840 assert_eq!(original_batch, returned_batch);
4841
4842 Ok(())
4843 }
4844
4845 #[test]
4846 fn test_roundtrip_fixed_list() -> Result<(), ArrowError> {
4847 let int_builder = Int64Builder::new();
4848 let mut fixed_list_builder = FixedSizeListBuilder::new(int_builder, 3)
4849 .with_field(Arc::new(Field::new("item", DataType::Int64, false)));
4850
4851 for point in [[1, 2, 3], [4, 5, 6], [7, 8, 9], [10, 11, 12]] {
4852 fixed_list_builder.values().append_value(point[0]);
4853 fixed_list_builder.values().append_value(point[1]);
4854 fixed_list_builder.values().append_value(point[2]);
4855
4856 fixed_list_builder.append(true);
4857 }
4858
4859 let array = Arc::new(fixed_list_builder.finish()) as ArrayRef;
4860
4861 let schema = Arc::new(Schema::new_with_metadata(
4862 vec![Field::new(
4863 "points",
4864 DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Int64, false)), 3),
4865 false,
4866 )],
4867 HashMap::default(),
4868 ));
4869
4870 test_slices(&array, &schema, 0, 4)?;
4873 test_slices(&array, &schema, 0, 2)?;
4874 test_slices(&array, &schema, 1, 3)?;
4875 test_slices(&array, &schema, 2, 1)?;
4876
4877 Ok(())
4878 }
4879
4880 #[test]
4881 fn test_roundtrip_fixed_list_w_nulls() -> Result<(), ArrowError> {
4882 let int_builder = Int64Builder::new();
4883 let mut fixed_list_builder = FixedSizeListBuilder::new(int_builder, 3);
4884
4885 for point in [
4886 [Some(1), Some(2), None],
4887 [Some(4), Some(5), Some(6)],
4888 [None, Some(8), Some(9)],
4889 [Some(10), None, None],
4890 ] {
4891 for p in point {
4892 match p {
4893 Some(p) => fixed_list_builder.values().append_value(p),
4894 None => fixed_list_builder.values().append_null(),
4895 }
4896 }
4897
4898 fixed_list_builder.append(true);
4899 }
4900
4901 let array = Arc::new(fixed_list_builder.finish()) as ArrayRef;
4902
4903 let schema = Arc::new(Schema::new_with_metadata(
4904 vec![Field::new(
4905 "points",
4906 DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Int64, true)), 3),
4907 true,
4908 )],
4909 HashMap::default(),
4910 ));
4911
4912 test_slices(&array, &schema, 0, 4)?;
4915 test_slices(&array, &schema, 0, 2)?;
4916 test_slices(&array, &schema, 1, 3)?;
4917 test_slices(&array, &schema, 2, 1)?;
4918
4919 Ok(())
4920 }
4921
4922 #[test]
4923 fn test_metadata_encoding_ordering() {
4924 fn create_hash() -> u64 {
4925 let metadata: HashMap<String, String> = [
4926 ("a", "1"), ("b", "2"), ("c", "3"), ("d", "4"), ("e", "5"), ]
4932 .into_iter()
4933 .map(|(k, v)| (k.to_owned(), v.to_owned()))
4934 .collect();
4935
4936 let schema = Arc::new(
4938 Schema::new(vec![
4939 Field::new("a", DataType::Int64, true).with_metadata(metadata.clone()),
4940 ])
4941 .with_metadata(metadata)
4942 .clone(),
4943 );
4944 let batch = RecordBatch::new_empty(schema.clone());
4945
4946 let mut bytes = Vec::new();
4947 let mut w = StreamWriter::try_new(&mut bytes, batch.schema_ref()).unwrap();
4948 w.write(&batch).unwrap();
4949 w.finish().unwrap();
4950
4951 let mut h = std::hash::DefaultHasher::new();
4952 h.write(&bytes);
4953 h.finish()
4954 }
4955
4956 let expected = create_hash();
4957
4958 let all_passed = (0..20).all(|_| create_hash() == expected);
4963 assert!(all_passed);
4964 }
4965
4966 #[test]
4967 fn test_dictionary_tracker_reset() {
4968 let data_gen = IpcDataGenerator::default();
4969 let mut dictionary_tracker = DictionaryTracker::new(false);
4970 let writer_options = IpcWriteOptions::default();
4971 let mut compression_ctx = IpcWriteContext::default();
4972
4973 let schema = Arc::new(Schema::new(vec![Field::new(
4974 "a",
4975 DataType::Dictionary(Box::new(DataType::UInt8), Box::new(DataType::Utf8)),
4976 false,
4977 )]));
4978
4979 let mut write_single_batch_stream =
4980 |batch: RecordBatch, dict_tracker: &mut DictionaryTracker| -> Vec<u8> {
4981 let mut buffer = Vec::new();
4982
4983 let stream_header = data_gen.schema_to_bytes_with_dictionary_tracker(
4985 &schema,
4986 dict_tracker,
4987 &writer_options,
4988 );
4989 _ = write_message(&mut buffer, stream_header, &writer_options).unwrap();
4990
4991 let (encoded_dicts, encoded_batch) = data_gen
4992 .encode(&batch, dict_tracker, &writer_options, &mut compression_ctx)
4993 .unwrap();
4994 for encoded_dict in encoded_dicts {
4995 _ = write_message(&mut buffer, encoded_dict, &writer_options).unwrap();
4996 }
4997 _ = write_message(&mut buffer, encoded_batch, &writer_options).unwrap();
4998
4999 buffer
5000 };
5001
5002 let batch1 = RecordBatch::try_new(
5003 schema.clone(),
5004 vec![Arc::new(DictionaryArray::new(
5005 UInt8Array::from_iter_values([0]),
5006 Arc::new(StringArray::from_iter_values(["a"])),
5007 ))],
5008 )
5009 .unwrap();
5010 let buffer = write_single_batch_stream(batch1.clone(), &mut dictionary_tracker);
5011
5012 let mut reader = StreamReader::try_new(Cursor::new(buffer), None).unwrap();
5014 let read_batch = reader.next().unwrap().unwrap();
5015 assert_eq!(read_batch, batch1);
5016
5017 dictionary_tracker.clear();
5019
5020 let batch2 = RecordBatch::try_new(
5022 schema.clone(),
5023 vec![Arc::new(DictionaryArray::new(
5024 UInt8Array::from_iter_values([0]),
5025 Arc::new(StringArray::from_iter_values(["a"])),
5026 ))],
5027 )
5028 .unwrap();
5029 let buffer = write_single_batch_stream(batch2.clone(), &mut dictionary_tracker);
5030 let mut reader = StreamReader::try_new(Cursor::new(buffer), None).unwrap();
5031 let read_batch = reader.next().unwrap().unwrap();
5032 assert_eq!(read_batch, batch2);
5033 }
5034}