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