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(
4095 Field::MAP_KEY_FIELD_DEFAULT_NAME,
4096 DataType::UInt32,
4097 false,
4098 ));
4099 let values = Arc::new(Field::new(
4100 Field::MAP_VALUE_FIELD_DEFAULT_NAME,
4101 DataType::UInt32,
4102 true,
4103 ));
4104 let map_field = Field::new_map(
4105 "map",
4106 Field::MAP_ENTRIES_FIELD_DEFAULT_NAME,
4107 keys,
4108 values,
4109 false,
4110 true,
4111 );
4112 let schema = Arc::new(Schema::new(vec![map_field]));
4113
4114 let values = Arc::new(generate_map_array_data());
4115
4116 let in_batch = RecordBatch::try_new(schema, vec![values]).unwrap();
4117 roundtrip_ensure_sliced_smaller(in_batch, 1000);
4118 }
4119
4120 fn generate_list_view_data<O: OffsetSizeTrait>() -> GenericListViewArray<O> {
4121 let mut builder = GenericListViewBuilder::<O, _>::new(UInt32Builder::new());
4122
4123 for i in 0u32..100_000 {
4124 if i.is_multiple_of(10_000) {
4125 builder.append(false);
4126 continue;
4127 }
4128 for value in [i, i, i] {
4129 builder.values().append_value(value);
4130 }
4131 builder.append(true);
4132 }
4133
4134 builder.finish()
4135 }
4136
4137 #[test]
4138 fn encode_list_view_arrays() {
4139 let val_inner = Field::new_list_field(DataType::UInt32, true);
4140 let val_field = Field::new("val", DataType::ListView(Arc::new(val_inner)), true);
4141 let schema = Arc::new(Schema::new(vec![val_field]));
4142
4143 let values = Arc::new(generate_list_view_data::<i32>());
4144
4145 let in_batch = RecordBatch::try_new(schema, vec![values]).unwrap();
4146 let out_batch = deserialize_file(serialize_file(&in_batch));
4147 assert_eq!(in_batch, out_batch);
4148 }
4149
4150 #[test]
4151 fn encode_large_list_view_arrays() {
4152 let val_inner = Field::new_list_field(DataType::UInt32, true);
4153 let val_field = Field::new("val", DataType::LargeListView(Arc::new(val_inner)), true);
4154 let schema = Arc::new(Schema::new(vec![val_field]));
4155
4156 let values = Arc::new(generate_list_view_data::<i64>());
4157
4158 let in_batch = RecordBatch::try_new(schema, vec![values]).unwrap();
4159 let out_batch = deserialize_file(serialize_file(&in_batch));
4160 assert_eq!(in_batch, out_batch);
4161 }
4162
4163 #[test]
4164 fn check_sliced_list_view_array() {
4165 let inner = Field::new_list_field(DataType::UInt32, true);
4166 let field = Field::new("val", DataType::ListView(Arc::new(inner)), true);
4167 let schema = Arc::new(Schema::new(vec![field]));
4168 let values = Arc::new(generate_list_view_data::<i32>());
4169
4170 for (offset, len) in [(999, 1), (0, 13), (47, 12), (values.len() - 13, 13)] {
4171 let in_batch = RecordBatch::try_new(schema.clone(), vec![values.clone()])
4172 .unwrap()
4173 .slice(offset, len);
4174 let out_batch = deserialize_file(serialize_file(&in_batch));
4175 assert_eq!(in_batch, out_batch);
4176 }
4177 }
4178
4179 #[test]
4180 fn check_sliced_large_list_view_array() {
4181 let inner = Field::new_list_field(DataType::UInt32, true);
4182 let field = Field::new("val", DataType::LargeListView(Arc::new(inner)), true);
4183 let schema = Arc::new(Schema::new(vec![field]));
4184 let values = Arc::new(generate_list_view_data::<i64>());
4185
4186 for (offset, len) in [(999, 1), (0, 13), (47, 12), (values.len() - 13, 13)] {
4187 let in_batch = RecordBatch::try_new(schema.clone(), vec![values.clone()])
4188 .unwrap()
4189 .slice(offset, len);
4190 let out_batch = deserialize_file(serialize_file(&in_batch));
4191 assert_eq!(in_batch, out_batch);
4192 }
4193 }
4194
4195 fn generate_nested_list_view_data<O: OffsetSizeTrait>() -> GenericListViewArray<O> {
4196 let inner_builder = UInt32Builder::new();
4197 let middle_builder = GenericListViewBuilder::<O, _>::new(inner_builder);
4198 let mut outer_builder = GenericListViewBuilder::<O, _>::new(middle_builder);
4199
4200 for i in 0u32..10_000 {
4201 if i.is_multiple_of(1_000) {
4202 outer_builder.append(false);
4203 continue;
4204 }
4205
4206 for _ in 0..3 {
4207 for value in [i, i + 1, i + 2] {
4208 outer_builder.values().values().append_value(value);
4209 }
4210 outer_builder.values().append(true);
4211 }
4212 outer_builder.append(true);
4213 }
4214
4215 outer_builder.finish()
4216 }
4217
4218 #[test]
4219 fn encode_nested_list_views() {
4220 let inner_int = Arc::new(Field::new_list_field(DataType::UInt32, true));
4221 let inner_list_field = Arc::new(Field::new_list_field(DataType::ListView(inner_int), true));
4222 let list_field = Field::new("val", DataType::ListView(inner_list_field), true);
4223 let schema = Arc::new(Schema::new(vec![list_field]));
4224
4225 let values = Arc::new(generate_nested_list_view_data::<i32>());
4226
4227 let in_batch = RecordBatch::try_new(schema, vec![values]).unwrap();
4228 let out_batch = deserialize_file(serialize_file(&in_batch));
4229 assert_eq!(in_batch, out_batch);
4230 }
4231
4232 fn test_roundtrip_list_view_of_dict_impl<OffsetSize: OffsetSizeTrait, U: ArrowNativeType>(
4233 list_data_type: DataType,
4234 offsets: &[U; 5],
4235 sizes: &[U; 4],
4236 ) {
4237 let values = StringArray::from(vec![Some("alpha"), None, Some("beta"), Some("gamma")]);
4238 let keys = Int32Array::from_iter_values([0, 0, 1, 2, 3, 0, 2]);
4239 let dict_array = DictionaryArray::new(keys, Arc::new(values));
4240 let dict_data = dict_array.to_data();
4241
4242 let value_offsets = Buffer::from_slice_ref(offsets);
4243 let value_sizes = Buffer::from_slice_ref(sizes);
4244
4245 let list_data = ArrayData::builder(list_data_type)
4246 .len(4)
4247 .add_buffer(value_offsets)
4248 .add_buffer(value_sizes)
4249 .add_child_data(dict_data)
4250 .build()
4251 .unwrap();
4252 let list_view_array = GenericListViewArray::<OffsetSize>::from(list_data);
4253
4254 let schema = Arc::new(Schema::new(vec![Field::new(
4255 "f1",
4256 list_view_array.data_type().clone(),
4257 false,
4258 )]));
4259 let input_batch = RecordBatch::try_new(schema, vec![Arc::new(list_view_array)]).unwrap();
4260
4261 let output_batch = deserialize_file(serialize_file(&input_batch));
4262 assert_eq!(input_batch, output_batch);
4263
4264 let output_batch = deserialize_stream(serialize_stream(&input_batch));
4265 assert_eq!(input_batch, output_batch);
4266 }
4267
4268 #[test]
4269 fn test_roundtrip_list_view_of_dict() {
4270 #[allow(deprecated)]
4271 let list_data_type = DataType::ListView(Arc::new(Field::new_dict(
4272 "item",
4273 DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
4274 true,
4275 1,
4276 false,
4277 )));
4278 let offsets: &[i32; 5] = &[0, 2, 4, 4, 7];
4279 let sizes: &[i32; 4] = &[2, 2, 0, 3];
4280 test_roundtrip_list_view_of_dict_impl::<i32, i32>(list_data_type, offsets, sizes);
4281 }
4282
4283 #[test]
4284 fn test_roundtrip_large_list_view_of_dict() {
4285 #[allow(deprecated)]
4286 let list_data_type = DataType::LargeListView(Arc::new(Field::new_dict(
4287 "item",
4288 DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
4289 true,
4290 2,
4291 false,
4292 )));
4293 let offsets: &[i64; 5] = &[0, 2, 4, 4, 7];
4294 let sizes: &[i64; 4] = &[2, 2, 0, 3];
4295 test_roundtrip_list_view_of_dict_impl::<i64, i64>(list_data_type, offsets, sizes);
4296 }
4297
4298 #[test]
4299 fn test_roundtrip_sliced_list_view_of_dict() {
4300 #[allow(deprecated)]
4301 let list_data_type = DataType::ListView(Arc::new(Field::new_dict(
4302 "item",
4303 DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
4304 true,
4305 3,
4306 false,
4307 )));
4308
4309 let values = StringArray::from(vec![Some("alpha"), None, Some("beta"), Some("gamma")]);
4310 let keys = Int32Array::from_iter_values([0, 0, 1, 2, 3, 0, 2, 1, 0, 3, 2, 1]);
4311 let dict_array = DictionaryArray::new(keys, Arc::new(values));
4312 let dict_data = dict_array.to_data();
4313
4314 let offsets: &[i32; 7] = &[0, 2, 4, 4, 7, 9, 12];
4315 let sizes: &[i32; 6] = &[2, 2, 0, 3, 2, 3];
4316 let value_offsets = Buffer::from_slice_ref(offsets);
4317 let value_sizes = Buffer::from_slice_ref(sizes);
4318
4319 let list_data = ArrayData::builder(list_data_type)
4320 .len(6)
4321 .add_buffer(value_offsets)
4322 .add_buffer(value_sizes)
4323 .add_child_data(dict_data)
4324 .build()
4325 .unwrap();
4326 let list_view_array = GenericListViewArray::<i32>::from(list_data);
4327
4328 let schema = Arc::new(Schema::new(vec![Field::new(
4329 "f1",
4330 list_view_array.data_type().clone(),
4331 false,
4332 )]));
4333 let input_batch = RecordBatch::try_new(schema, vec![Arc::new(list_view_array)]).unwrap();
4334
4335 let sliced_batch = input_batch.slice(1, 4);
4336
4337 let output_batch = deserialize_file(serialize_file(&sliced_batch));
4338 assert_eq!(sliced_batch, output_batch);
4339
4340 let output_batch = deserialize_stream(serialize_stream(&sliced_batch));
4341 assert_eq!(sliced_batch, output_batch);
4342 }
4343
4344 #[test]
4345 fn test_roundtrip_dense_union_of_dict() {
4346 let values = StringArray::from(vec![Some("alpha"), None, Some("beta"), Some("gamma")]);
4347 let keys = Int32Array::from_iter_values([0, 0, 1, 2, 3, 0, 2]);
4348 let dict_array = DictionaryArray::new(keys, Arc::new(values));
4349
4350 #[allow(deprecated)]
4351 let dict_field = Arc::new(Field::new_dict(
4352 "dict",
4353 DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
4354 true,
4355 1,
4356 false,
4357 ));
4358 let int_field = Arc::new(Field::new("int", DataType::Int32, false));
4359 let union_fields = UnionFields::try_new(vec![0, 1], vec![dict_field, int_field]).unwrap();
4360
4361 let types = ScalarBuffer::from(vec![0i8, 0, 1, 0, 1, 0, 0]);
4362 let offsets = ScalarBuffer::from(vec![0i32, 1, 0, 2, 1, 3, 4]);
4363
4364 let int_array = Int32Array::from(vec![100, 200]);
4365
4366 let union = UnionArray::try_new(
4367 union_fields.clone(),
4368 types,
4369 Some(offsets),
4370 vec![Arc::new(dict_array), Arc::new(int_array)],
4371 )
4372 .unwrap();
4373
4374 let schema = Arc::new(Schema::new(vec![Field::new(
4375 "union",
4376 DataType::Union(union_fields, UnionMode::Dense),
4377 false,
4378 )]));
4379 let input_batch = RecordBatch::try_new(schema, vec![Arc::new(union)]).unwrap();
4380
4381 let output_batch = deserialize_file(serialize_file(&input_batch));
4382 assert_eq!(input_batch, output_batch);
4383
4384 let output_batch = deserialize_stream(serialize_stream(&input_batch));
4385 assert_eq!(input_batch, output_batch);
4386 }
4387
4388 #[test]
4389 fn test_roundtrip_sparse_union_of_dict() {
4390 let values = StringArray::from(vec![Some("alpha"), None, Some("beta"), Some("gamma")]);
4391 let keys = Int32Array::from_iter_values([0, 0, 1, 2, 3, 0, 2]);
4392 let dict_array = DictionaryArray::new(keys, Arc::new(values));
4393
4394 #[allow(deprecated)]
4395 let dict_field = Arc::new(Field::new_dict(
4396 "dict",
4397 DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
4398 true,
4399 2,
4400 false,
4401 ));
4402 let int_field = Arc::new(Field::new("int", DataType::Int32, false));
4403 let union_fields = UnionFields::try_new(vec![0, 1], vec![dict_field, int_field]).unwrap();
4404
4405 let types = ScalarBuffer::from(vec![0i8, 0, 1, 0, 1, 0, 0]);
4406
4407 let int_array = Int32Array::from(vec![0, 0, 100, 0, 200, 0, 0]);
4408
4409 let union = UnionArray::try_new(
4410 union_fields.clone(),
4411 types,
4412 None,
4413 vec![Arc::new(dict_array), Arc::new(int_array)],
4414 )
4415 .unwrap();
4416
4417 let schema = Arc::new(Schema::new(vec![Field::new(
4418 "union",
4419 DataType::Union(union_fields, UnionMode::Sparse),
4420 false,
4421 )]));
4422 let input_batch = RecordBatch::try_new(schema, vec![Arc::new(union)]).unwrap();
4423
4424 let output_batch = deserialize_file(serialize_file(&input_batch));
4425 assert_eq!(input_batch, output_batch);
4426
4427 let output_batch = deserialize_stream(serialize_stream(&input_batch));
4428 assert_eq!(input_batch, output_batch);
4429 }
4430
4431 #[test]
4432 fn test_roundtrip_map_with_dict_keys() {
4433 let key_values = StringArray::from(vec!["key_a", "key_b", "key_c"]);
4436 let keys = Int32Array::from_iter_values([0, 1, 2, 0, 1, 0]);
4437 let dict_keys = DictionaryArray::new(keys, Arc::new(key_values));
4438
4439 let values = Int32Array::from(vec![1, 2, 3, 4, 5, 6]);
4440
4441 #[allow(deprecated)]
4442 let entries_field = Arc::new(Field::new(
4443 Field::MAP_ENTRIES_FIELD_DEFAULT_NAME,
4444 DataType::Struct(
4445 vec![
4446 Field::new_dict(
4447 Field::MAP_KEY_FIELD_DEFAULT_NAME,
4448 DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
4449 false,
4450 1,
4451 false,
4452 ),
4453 Field::new(Field::MAP_VALUE_FIELD_DEFAULT_NAME, DataType::Int32, true),
4454 ]
4455 .into(),
4456 ),
4457 false,
4458 ));
4459
4460 let entries = StructArray::from(vec![
4461 (
4462 Arc::new(Field::new(
4463 Field::MAP_KEY_FIELD_DEFAULT_NAME,
4464 DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
4465 false,
4466 )),
4467 Arc::new(dict_keys) as ArrayRef,
4468 ),
4469 (
4470 Arc::new(Field::new(
4471 Field::MAP_VALUE_FIELD_DEFAULT_NAME,
4472 DataType::Int32,
4473 true,
4474 )),
4475 Arc::new(values) as ArrayRef,
4476 ),
4477 ]);
4478
4479 let offsets = Buffer::from_slice_ref([0i32, 2, 4, 6]);
4480
4481 let map_data = ArrayData::builder(DataType::Map(entries_field, false))
4482 .len(3)
4483 .add_buffer(offsets)
4484 .add_child_data(entries.into_data())
4485 .build()
4486 .unwrap();
4487 let map_array = MapArray::from(map_data);
4488
4489 let schema = Arc::new(Schema::new(vec![Field::new(
4490 "map",
4491 map_array.data_type().clone(),
4492 false,
4493 )]));
4494 let input_batch = RecordBatch::try_new(schema, vec![Arc::new(map_array)]).unwrap();
4495
4496 let output_batch = deserialize_file(serialize_file(&input_batch));
4497 assert_eq!(input_batch, output_batch);
4498
4499 let output_batch = deserialize_stream(serialize_stream(&input_batch));
4500 assert_eq!(input_batch, output_batch);
4501 }
4502
4503 #[test]
4504 fn test_roundtrip_map_with_dict_values() {
4505 let keys = StringArray::from(vec!["a", "b", "c", "d", "e", "f"]);
4508
4509 let value_values = StringArray::from(vec!["val_x", "val_y", "val_z"]);
4510 let value_keys = Int32Array::from_iter_values([0, 1, 2, 0, 1, 0]);
4511 let dict_values = DictionaryArray::new(value_keys, Arc::new(value_values));
4512
4513 #[allow(deprecated)]
4514 let entries_field = Arc::new(Field::new(
4515 Field::MAP_ENTRIES_FIELD_DEFAULT_NAME,
4516 DataType::Struct(
4517 vec![
4518 Field::new(Field::MAP_KEY_FIELD_DEFAULT_NAME, DataType::Utf8, false),
4519 Field::new_dict(
4520 Field::MAP_VALUE_FIELD_DEFAULT_NAME,
4521 DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
4522 true,
4523 2,
4524 false,
4525 ),
4526 ]
4527 .into(),
4528 ),
4529 false,
4530 ));
4531
4532 let entries = StructArray::from(vec![
4533 (
4534 Arc::new(Field::new(
4535 Field::MAP_KEY_FIELD_DEFAULT_NAME,
4536 DataType::Utf8,
4537 false,
4538 )),
4539 Arc::new(keys) as ArrayRef,
4540 ),
4541 (
4542 Arc::new(Field::new(
4543 Field::MAP_VALUE_FIELD_DEFAULT_NAME,
4544 DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
4545 true,
4546 )),
4547 Arc::new(dict_values) as ArrayRef,
4548 ),
4549 ]);
4550
4551 let offsets = Buffer::from_slice_ref([0i32, 2, 4, 6]);
4552
4553 let map_data = ArrayData::builder(DataType::Map(entries_field, false))
4554 .len(3)
4555 .add_buffer(offsets)
4556 .add_child_data(entries.into_data())
4557 .build()
4558 .unwrap();
4559 let map_array = MapArray::from(map_data);
4560
4561 let schema = Arc::new(Schema::new(vec![Field::new(
4562 "map",
4563 map_array.data_type().clone(),
4564 false,
4565 )]));
4566 let input_batch = RecordBatch::try_new(schema, vec![Arc::new(map_array)]).unwrap();
4567
4568 let output_batch = deserialize_file(serialize_file(&input_batch));
4569 assert_eq!(input_batch, output_batch);
4570
4571 let output_batch = deserialize_stream(serialize_stream(&input_batch));
4572 assert_eq!(input_batch, output_batch);
4573 }
4574
4575 #[test]
4576 fn test_decimal128_alignment16_is_sufficient() {
4577 const IPC_ALIGNMENT: usize = 16;
4578
4579 for num_cols in [1, 2, 3, 17, 50, 73, 99] {
4584 let num_rows = (num_cols * 7 + 11) % 100; let mut fields = Vec::new();
4587 let mut arrays = Vec::new();
4588 for i in 0..num_cols {
4589 let field = Field::new(format!("col_{i}"), DataType::Decimal128(38, 10), true);
4590 let array = Decimal128Array::from(vec![num_cols as i128; num_rows]);
4591 fields.push(field);
4592 arrays.push(Arc::new(array) as Arc<dyn Array>);
4593 }
4594 let schema = Schema::new(fields);
4595 let batch = RecordBatch::try_new(Arc::new(schema), arrays).unwrap();
4596
4597 let mut writer = FileWriter::try_new_with_options(
4598 Vec::new(),
4599 batch.schema_ref(),
4600 IpcWriteOptions::try_new(IPC_ALIGNMENT, false, MetadataVersion::V5).unwrap(),
4601 )
4602 .unwrap();
4603 writer.write(&batch).unwrap();
4604 writer.finish().unwrap();
4605
4606 let out: Vec<u8> = writer.into_inner().unwrap();
4607
4608 let buffer = Buffer::from_vec(out);
4609 let trailer_start = buffer.len() - 10;
4610 let footer_len =
4611 read_footer_length(buffer[trailer_start..].try_into().unwrap()).unwrap();
4612 let footer =
4613 root_as_footer(&buffer[trailer_start - footer_len..trailer_start]).unwrap();
4614
4615 let schema = fb_to_schema(footer.schema().unwrap());
4616
4617 let decoder =
4620 FileDecoder::new(Arc::new(schema), footer.version()).with_require_alignment(true);
4621
4622 let batches = footer.recordBatches().unwrap();
4623
4624 let block = batches.get(0);
4625 let block_len = block.bodyLength() as usize + block.metaDataLength() as usize;
4626 let data = buffer.slice_with_length(block.offset() as _, block_len);
4627
4628 let batch2 = decoder.read_record_batch(block, &data).unwrap().unwrap();
4629
4630 assert_eq!(batch, batch2);
4631 }
4632 }
4633
4634 #[test]
4635 fn test_decimal128_alignment8_is_unaligned() {
4636 const IPC_ALIGNMENT: usize = 8;
4637
4638 let num_cols = 2;
4639 let num_rows = 1;
4640
4641 let mut fields = Vec::new();
4642 let mut arrays = Vec::new();
4643 for i in 0..num_cols {
4644 let field = Field::new(format!("col_{i}"), DataType::Decimal128(38, 10), true);
4645 let array = Decimal128Array::from(vec![num_cols as i128; num_rows]);
4646 fields.push(field);
4647 arrays.push(Arc::new(array) as Arc<dyn Array>);
4648 }
4649 let schema = Schema::new(fields);
4650 let batch = RecordBatch::try_new(Arc::new(schema), arrays).unwrap();
4651
4652 let mut writer = FileWriter::try_new_with_options(
4653 Vec::new(),
4654 batch.schema_ref(),
4655 IpcWriteOptions::try_new(IPC_ALIGNMENT, false, MetadataVersion::V5).unwrap(),
4656 )
4657 .unwrap();
4658 writer.write(&batch).unwrap();
4659 writer.finish().unwrap();
4660
4661 let out: Vec<u8> = writer.into_inner().unwrap();
4662
4663 let buffer = Buffer::from_vec(out);
4664 let trailer_start = buffer.len() - 10;
4665 let footer_len = read_footer_length(buffer[trailer_start..].try_into().unwrap()).unwrap();
4666 let footer = root_as_footer(&buffer[trailer_start - footer_len..trailer_start]).unwrap();
4667 let schema = fb_to_schema(footer.schema().unwrap());
4668
4669 let decoder =
4672 FileDecoder::new(Arc::new(schema), footer.version()).with_require_alignment(true);
4673
4674 let batches = footer.recordBatches().unwrap();
4675
4676 let block = batches.get(0);
4677 let block_len = block.bodyLength() as usize + block.metaDataLength() as usize;
4678 let data = buffer.slice_with_length(block.offset() as _, block_len);
4679
4680 let result = decoder.read_record_batch(block, &data);
4681
4682 let error = result.unwrap_err();
4683 assert_eq!(
4684 error.to_string(),
4685 "Invalid argument error: Misaligned buffers[0] in array of type Decimal128(38, 10), \
4686 offset from expected alignment of 16 by 8"
4687 );
4688 }
4689
4690 #[test]
4691 fn test_flush() {
4692 let num_cols = 2;
4695 let mut fields = Vec::new();
4696 let options = IpcWriteOptions::try_new(8, false, MetadataVersion::V5).unwrap();
4697 for i in 0..num_cols {
4698 let field = Field::new(format!("col_{i}"), DataType::Decimal128(38, 10), true);
4699 fields.push(field);
4700 }
4701 let schema = Schema::new(fields);
4702 let inner_stream_writer = BufWriter::with_capacity(1024, Vec::new());
4703 let inner_file_writer = BufWriter::with_capacity(1024, Vec::new());
4704 let mut stream_writer =
4705 StreamWriter::try_new_with_options(inner_stream_writer, &schema, options.clone())
4706 .unwrap();
4707 let mut file_writer =
4708 FileWriter::try_new_with_options(inner_file_writer, &schema, options).unwrap();
4709
4710 let stream_bytes_written_on_new = stream_writer.get_ref().get_ref().len();
4711 let file_bytes_written_on_new = file_writer.get_ref().get_ref().len();
4712 stream_writer.flush().unwrap();
4713 file_writer.flush().unwrap();
4714 let stream_bytes_written_on_flush = stream_writer.get_ref().get_ref().len();
4715 let file_bytes_written_on_flush = file_writer.get_ref().get_ref().len();
4716 let stream_out = stream_writer.into_inner().unwrap().into_inner().unwrap();
4717 let expected_stream_flushed_bytes = stream_out.len() - 8;
4721 let expected_file_flushed_bytes = expected_stream_flushed_bytes + 8;
4724
4725 assert!(
4726 stream_bytes_written_on_new < stream_bytes_written_on_flush,
4727 "this test makes no sense if flush is not actually required"
4728 );
4729 assert!(
4730 file_bytes_written_on_new < file_bytes_written_on_flush,
4731 "this test makes no sense if flush is not actually required"
4732 );
4733 assert_eq!(stream_bytes_written_on_flush, expected_stream_flushed_bytes);
4734 assert_eq!(file_bytes_written_on_flush, expected_file_flushed_bytes);
4735 }
4736
4737 #[test]
4738 fn test_roundtrip_list_of_fixed_list() -> Result<(), ArrowError> {
4739 let l1_type =
4740 DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Float32, false)), 3);
4741 let l2_type = DataType::List(Arc::new(Field::new("item", l1_type.clone(), false)));
4742
4743 let l0_builder = Float32Builder::new();
4744 let l1_builder = FixedSizeListBuilder::new(l0_builder, 3).with_field(Arc::new(Field::new(
4745 "item",
4746 DataType::Float32,
4747 false,
4748 )));
4749 let mut l2_builder =
4750 ListBuilder::new(l1_builder).with_field(Arc::new(Field::new("item", l1_type, false)));
4751
4752 for point in [[1.0, 2.0, 3.0], [4.0, 5.0, 6.0], [7.0, 8.0, 9.0]] {
4753 l2_builder.values().values().append_value(point[0]);
4754 l2_builder.values().values().append_value(point[1]);
4755 l2_builder.values().values().append_value(point[2]);
4756
4757 l2_builder.values().append(true);
4758 }
4759 l2_builder.append(true);
4760
4761 let point = [10., 11., 12.];
4762 l2_builder.values().values().append_value(point[0]);
4763 l2_builder.values().values().append_value(point[1]);
4764 l2_builder.values().values().append_value(point[2]);
4765
4766 l2_builder.values().append(true);
4767 l2_builder.append(true);
4768
4769 let array = Arc::new(l2_builder.finish()) as ArrayRef;
4770
4771 let schema = Arc::new(Schema::new_with_metadata(
4772 vec![Field::new("points", l2_type, false)],
4773 HashMap::default(),
4774 ));
4775
4776 test_slices(&array, &schema, 0, 1)?;
4779 test_slices(&array, &schema, 0, 2)?;
4780 test_slices(&array, &schema, 1, 1)?;
4781
4782 Ok(())
4783 }
4784
4785 #[test]
4786 fn test_roundtrip_list_of_fixed_list_w_nulls() -> Result<(), ArrowError> {
4787 let l0_builder = Float32Builder::new();
4788 let l1_builder = FixedSizeListBuilder::new(l0_builder, 3);
4789 let mut l2_builder = ListBuilder::new(l1_builder);
4790
4791 for point in [
4792 [Some(1.0), Some(2.0), None],
4793 [Some(4.0), Some(5.0), Some(6.0)],
4794 [None, Some(8.0), Some(9.0)],
4795 ] {
4796 for p in point {
4797 match p {
4798 Some(p) => l2_builder.values().values().append_value(p),
4799 None => l2_builder.values().values().append_null(),
4800 }
4801 }
4802
4803 l2_builder.values().append(true);
4804 }
4805 l2_builder.append(true);
4806
4807 let point = [Some(10.), None, None];
4808 for p in point {
4809 match p {
4810 Some(p) => l2_builder.values().values().append_value(p),
4811 None => l2_builder.values().values().append_null(),
4812 }
4813 }
4814
4815 l2_builder.values().append(true);
4816 l2_builder.append(true);
4817
4818 let array = Arc::new(l2_builder.finish()) as ArrayRef;
4819
4820 let schema = Arc::new(Schema::new_with_metadata(
4821 vec![Field::new(
4822 "points",
4823 DataType::List(Arc::new(Field::new(
4824 "item",
4825 DataType::FixedSizeList(
4826 Arc::new(Field::new("item", DataType::Float32, true)),
4827 3,
4828 ),
4829 true,
4830 ))),
4831 true,
4832 )],
4833 HashMap::default(),
4834 ));
4835
4836 test_slices(&array, &schema, 0, 1)?;
4839 test_slices(&array, &schema, 0, 2)?;
4840 test_slices(&array, &schema, 1, 1)?;
4841
4842 Ok(())
4843 }
4844
4845 fn test_slices(
4846 parent_array: &ArrayRef,
4847 schema: &SchemaRef,
4848 offset: usize,
4849 length: usize,
4850 ) -> Result<(), ArrowError> {
4851 let subarray = parent_array.slice(offset, length);
4852 let original_batch = RecordBatch::try_new(schema.clone(), vec![subarray])?;
4853
4854 let mut bytes = Vec::new();
4855 let mut writer = StreamWriter::try_new(&mut bytes, schema)?;
4856 writer.write(&original_batch)?;
4857 writer.finish()?;
4858
4859 let mut cursor = std::io::Cursor::new(bytes);
4860 let mut reader = StreamReader::try_new(&mut cursor, None)?;
4861 let returned_batch = reader.next().unwrap()?;
4862
4863 assert_eq!(original_batch, returned_batch);
4864
4865 Ok(())
4866 }
4867
4868 #[test]
4869 fn test_roundtrip_fixed_list() -> Result<(), ArrowError> {
4870 let int_builder = Int64Builder::new();
4871 let mut fixed_list_builder = FixedSizeListBuilder::new(int_builder, 3)
4872 .with_field(Arc::new(Field::new("item", DataType::Int64, false)));
4873
4874 for point in [[1, 2, 3], [4, 5, 6], [7, 8, 9], [10, 11, 12]] {
4875 fixed_list_builder.values().append_value(point[0]);
4876 fixed_list_builder.values().append_value(point[1]);
4877 fixed_list_builder.values().append_value(point[2]);
4878
4879 fixed_list_builder.append(true);
4880 }
4881
4882 let array = Arc::new(fixed_list_builder.finish()) as ArrayRef;
4883
4884 let schema = Arc::new(Schema::new_with_metadata(
4885 vec![Field::new(
4886 "points",
4887 DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Int64, false)), 3),
4888 false,
4889 )],
4890 HashMap::default(),
4891 ));
4892
4893 test_slices(&array, &schema, 0, 4)?;
4896 test_slices(&array, &schema, 0, 2)?;
4897 test_slices(&array, &schema, 1, 3)?;
4898 test_slices(&array, &schema, 2, 1)?;
4899
4900 Ok(())
4901 }
4902
4903 #[test]
4904 fn test_roundtrip_fixed_list_w_nulls() -> Result<(), ArrowError> {
4905 let int_builder = Int64Builder::new();
4906 let mut fixed_list_builder = FixedSizeListBuilder::new(int_builder, 3);
4907
4908 for point in [
4909 [Some(1), Some(2), None],
4910 [Some(4), Some(5), Some(6)],
4911 [None, Some(8), Some(9)],
4912 [Some(10), None, None],
4913 ] {
4914 for p in point {
4915 match p {
4916 Some(p) => fixed_list_builder.values().append_value(p),
4917 None => fixed_list_builder.values().append_null(),
4918 }
4919 }
4920
4921 fixed_list_builder.append(true);
4922 }
4923
4924 let array = Arc::new(fixed_list_builder.finish()) as ArrayRef;
4925
4926 let schema = Arc::new(Schema::new_with_metadata(
4927 vec![Field::new(
4928 "points",
4929 DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Int64, true)), 3),
4930 true,
4931 )],
4932 HashMap::default(),
4933 ));
4934
4935 test_slices(&array, &schema, 0, 4)?;
4938 test_slices(&array, &schema, 0, 2)?;
4939 test_slices(&array, &schema, 1, 3)?;
4940 test_slices(&array, &schema, 2, 1)?;
4941
4942 Ok(())
4943 }
4944
4945 #[test]
4946 fn test_metadata_encoding_ordering() {
4947 fn create_hash() -> u64 {
4948 let metadata: HashMap<String, String> = [
4949 ("a", "1"), ("b", "2"), ("c", "3"), ("d", "4"), ("e", "5"), ]
4955 .into_iter()
4956 .map(|(k, v)| (k.to_owned(), v.to_owned()))
4957 .collect();
4958
4959 let schema = Arc::new(
4961 Schema::new(vec![
4962 Field::new("a", DataType::Int64, true).with_metadata(metadata.clone()),
4963 ])
4964 .with_metadata(metadata)
4965 .clone(),
4966 );
4967 let batch = RecordBatch::new_empty(schema.clone());
4968
4969 let mut bytes = Vec::new();
4970 let mut w = StreamWriter::try_new(&mut bytes, batch.schema_ref()).unwrap();
4971 w.write(&batch).unwrap();
4972 w.finish().unwrap();
4973
4974 let mut h = std::hash::DefaultHasher::new();
4975 h.write(&bytes);
4976 h.finish()
4977 }
4978
4979 let expected = create_hash();
4980
4981 let all_passed = (0..20).all(|_| create_hash() == expected);
4986 assert!(all_passed);
4987 }
4988
4989 #[test]
4990 fn test_dictionary_tracker_reset() {
4991 let data_gen = IpcDataGenerator::default();
4992 let mut dictionary_tracker = DictionaryTracker::new(false);
4993 let writer_options = IpcWriteOptions::default();
4994 let mut compression_ctx = IpcWriteContext::default();
4995
4996 let schema = Arc::new(Schema::new(vec![Field::new(
4997 "a",
4998 DataType::Dictionary(Box::new(DataType::UInt8), Box::new(DataType::Utf8)),
4999 false,
5000 )]));
5001
5002 let mut write_single_batch_stream =
5003 |batch: RecordBatch, dict_tracker: &mut DictionaryTracker| -> Vec<u8> {
5004 let mut buffer = Vec::new();
5005
5006 let stream_header = data_gen.schema_to_bytes_with_dictionary_tracker(
5008 &schema,
5009 dict_tracker,
5010 &writer_options,
5011 );
5012 _ = write_message(&mut buffer, stream_header, &writer_options).unwrap();
5013
5014 let (encoded_dicts, encoded_batch) = data_gen
5015 .encode(&batch, dict_tracker, &writer_options, &mut compression_ctx)
5016 .unwrap();
5017 for encoded_dict in encoded_dicts {
5018 _ = write_message(&mut buffer, encoded_dict, &writer_options).unwrap();
5019 }
5020 _ = write_message(&mut buffer, encoded_batch, &writer_options).unwrap();
5021
5022 buffer
5023 };
5024
5025 let batch1 = RecordBatch::try_new(
5026 schema.clone(),
5027 vec![Arc::new(DictionaryArray::new(
5028 UInt8Array::from_iter_values([0]),
5029 Arc::new(StringArray::from_iter_values(["a"])),
5030 ))],
5031 )
5032 .unwrap();
5033 let buffer = write_single_batch_stream(batch1.clone(), &mut dictionary_tracker);
5034
5035 let mut reader = StreamReader::try_new(Cursor::new(buffer), None).unwrap();
5037 let read_batch = reader.next().unwrap().unwrap();
5038 assert_eq!(read_batch, batch1);
5039
5040 dictionary_tracker.clear();
5042
5043 let batch2 = RecordBatch::try_new(
5045 schema.clone(),
5046 vec![Arc::new(DictionaryArray::new(
5047 UInt8Array::from_iter_values([0]),
5048 Arc::new(StringArray::from_iter_values(["a"])),
5049 ))],
5050 )
5051 .unwrap();
5052 let buffer = write_single_batch_stream(batch2.clone(), &mut dictionary_tracker);
5053 let mut reader = StreamReader::try_new(Cursor::new(buffer), None).unwrap();
5054 let read_batch = reader.next().unwrap().unwrap();
5055 assert_eq!(read_batch, batch2);
5056 }
5057}