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