1use std::cmp::min;
28use std::collections::HashMap;
29use std::io::{BufWriter, Write};
30use std::mem::size_of;
31use std::sync::Arc;
32
33use flatbuffers::FlatBufferBuilder;
34
35use arrow_array::builder::BufferBuilder;
36use arrow_array::cast::*;
37use arrow_array::types::{Int16Type, Int32Type, Int64Type, RunEndIndexType};
38use arrow_array::*;
39use arrow_buffer::bit_util;
40use arrow_buffer::{ArrowNativeType, Buffer, MutableBuffer, ToByteSlice};
41use arrow_data::{ArrayData, ArrayDataBuilder, BufferSpec, layout};
42use arrow_schema::*;
43
44use crate::CONTINUATION_MARKER;
45use crate::compression::CompressionCodec;
46#[expect(deprecated)]
47pub use crate::compression::{CompressionContext, IpcWriteContext};
48use crate::convert::IpcSchemaEncoder;
49
50#[derive(Debug, Clone)]
52pub struct IpcWriteOptions {
53 alignment: u8,
56 write_legacy_ipc_format: bool,
58 metadata_version: crate::MetadataVersion,
67 batch_compression_type: Option<crate::CompressionType>,
70 batch_compression_level: Option<i32>,
72 dictionary_handling: DictionaryHandling,
74}
75
76enum EncodedBuffer {
82 Raw(Buffer),
84 Compressed(Vec<u8>),
86}
87
88impl EncodedBuffer {
89 fn as_slice(&self) -> &[u8] {
90 match self {
91 EncodedBuffer::Raw(b) => b.as_slice(),
92 EncodedBuffer::Compressed(v) => v.as_slice(),
93 }
94 }
95
96 fn len(&self) -> usize {
97 match self {
98 EncodedBuffer::Raw(b) => b.len(),
99 EncodedBuffer::Compressed(v) => v.len(),
100 }
101 }
102}
103#[derive(Default)]
108struct IpcMetadataBuilder {
109 nodes: Vec<crate::FieldNode>,
110 buffers: Vec<crate::Buffer>,
111}
112
113enum IpcBodySink<'a> {
118 Write(&'a mut Vec<u8>),
120 Collect(&'a mut Vec<EncodedBuffer>),
122}
123impl<'a> IpcBodySink<'a> {
124 pub fn write(&mut self, pad_len: usize, buffer: EncodedBuffer) {
126 match self {
127 IpcBodySink::Write(vec) => {
128 vec.extend_from_slice(buffer.as_slice());
129 vec.extend_from_slice(&PADDING[..pad_len]);
130 }
131 IpcBodySink::Collect(vec) => {
132 vec.push(buffer);
133 }
134 }
135 }
136}
137
138struct IpcWriteMetadata {
143 dictionary_block_sizes: Vec<(usize, usize)>,
146 padded_header_len: usize,
148 body_len: usize,
150}
151
152impl IpcWriteOptions {
153 pub fn try_with_compression(
158 mut self,
159 batch_compression_type: Option<crate::CompressionType>,
160 ) -> Result<Self, ArrowError> {
161 self.batch_compression_type = batch_compression_type;
162
163 if self.batch_compression_type.is_some()
164 && self.metadata_version < crate::MetadataVersion::V5
165 {
166 return Err(ArrowError::InvalidArgumentError(
167 "Compression only supported in metadata v5 and above".to_string(),
168 ));
169 }
170 Ok(self)
171 }
172
173 pub fn try_with_compression_level(
178 mut self,
179 batch_compression_level: Option<i32>,
180 ) -> Result<Self, ArrowError> {
181 self.batch_compression_level = batch_compression_level;
182
183 if self.batch_compression_level.is_some()
184 && self.metadata_version < crate::MetadataVersion::V5
185 {
186 return Err(ArrowError::InvalidArgumentError(
187 "Compression only supported in metadata v5 and above".to_string(),
188 ));
189 }
190
191 match (self.batch_compression_type, self.batch_compression_level) {
192 (Some(crate::CompressionType::ZSTD), Some(level)) => {
193 return self.check_zstd_level(level);
194 }
195 (Some(crate::CompressionType::LZ4_FRAME), Some(_)) => {
196 return Err(ArrowError::InvalidArgumentError(
197 "LZ4 Frame compression does not support configurable compression levels"
198 .to_string(),
199 ));
200 }
201 _ => {}
202 }
203
204 Ok(self)
205 }
206
207 #[cfg(not(feature = "zstd"))]
208 fn check_zstd_level(self, _level: i32) -> Result<Self, ArrowError> {
209 Err(ArrowError::InvalidArgumentError(
210 "zstd IPC compression requires the zstd feature".to_string(),
211 ))
212 }
213
214 #[cfg(feature = "zstd")]
215 fn check_zstd_level(self, level: i32) -> Result<Self, ArrowError> {
216 let range = zstd::compression_level_range();
217 if !range.contains(&(level as zstd::zstd_safe::CompressionLevel)) {
218 return Err(ArrowError::InvalidArgumentError(format!(
219 "ZSTD compression level must be between {} and {}, got {}",
220 range.start(),
221 range.end(),
222 level,
223 )));
224 }
225
226 Ok(self)
227 }
228
229 pub fn try_new(
231 alignment: usize,
232 write_legacy_ipc_format: bool,
233 metadata_version: crate::MetadataVersion,
234 ) -> Result<Self, ArrowError> {
235 let is_alignment_valid =
236 alignment == 8 || alignment == 16 || alignment == 32 || alignment == 64;
237 if !is_alignment_valid {
238 return Err(ArrowError::InvalidArgumentError(
239 "Alignment should be 8, 16, 32, or 64.".to_string(),
240 ));
241 }
242 let alignment: u8 = u8::try_from(alignment).expect("range already checked");
243 match metadata_version {
244 crate::MetadataVersion::V1
245 | crate::MetadataVersion::V2
246 | crate::MetadataVersion::V3 => Err(ArrowError::InvalidArgumentError(
247 "Writing IPC metadata version 3 and lower not supported".to_string(),
248 )),
249 #[allow(deprecated)]
250 crate::MetadataVersion::V4 => Ok(Self {
251 alignment,
252 write_legacy_ipc_format,
253 metadata_version,
254 batch_compression_type: None,
255 batch_compression_level: None,
256 dictionary_handling: DictionaryHandling::default(),
257 }),
258 crate::MetadataVersion::V5 => {
259 if write_legacy_ipc_format {
260 Err(ArrowError::InvalidArgumentError(
261 "Legacy IPC format only supported on metadata version 4".to_string(),
262 ))
263 } else {
264 Ok(Self {
265 alignment,
266 write_legacy_ipc_format,
267 metadata_version,
268 batch_compression_type: None,
269 batch_compression_level: None,
270 dictionary_handling: DictionaryHandling::default(),
271 })
272 }
273 }
274 z => Err(ArrowError::InvalidArgumentError(format!(
275 "Unsupported crate::MetadataVersion {z:?}"
276 ))),
277 }
278 }
279
280 pub fn with_dictionary_handling(mut self, dictionary_handling: DictionaryHandling) -> Self {
282 self.dictionary_handling = dictionary_handling;
283 self
284 }
285}
286
287impl Default for IpcWriteOptions {
288 fn default() -> Self {
289 Self {
290 alignment: 64,
291 write_legacy_ipc_format: false,
292 metadata_version: crate::MetadataVersion::V5,
293 batch_compression_type: None,
294 batch_compression_level: None,
295 dictionary_handling: DictionaryHandling::default(),
296 }
297 }
298}
299
300#[derive(Debug, Default)]
301pub struct IpcDataGenerator {}
335
336impl IpcDataGenerator {
337 pub fn schema_to_bytes_with_dictionary_tracker(
340 &self,
341 schema: &Schema,
342 dictionary_tracker: &mut DictionaryTracker,
343 write_options: &IpcWriteOptions,
344 ) -> EncodedData {
345 let mut fbb = FlatBufferBuilder::new();
346 let schema = {
347 let fb = IpcSchemaEncoder::new()
348 .with_dictionary_tracker(dictionary_tracker)
349 .schema_to_fb_offset(&mut fbb, schema);
350 fb.as_union_value()
351 };
352
353 let mut message = crate::MessageBuilder::new(&mut fbb);
354 message.add_version(write_options.metadata_version);
355 message.add_header_type(crate::MessageHeader::Schema);
356 message.add_bodyLength(0);
357 message.add_header(schema);
358 let data = message.finish();
360 fbb.finish(data, None);
361
362 let data = fbb.finished_data();
363 EncodedData {
364 ipc_message: data.to_vec(),
365 arrow_data: vec![],
366 }
367 }
368
369 fn _encode_dictionaries<I: Iterator<Item = i64>>(
370 &self,
371 column: &ArrayRef,
372 encoded_dictionaries: &mut Vec<EncodedData>,
373 dictionary_tracker: &mut DictionaryTracker,
374 write_options: &IpcWriteOptions,
375 dict_id: &mut I,
376 ipc_write_context: &mut IpcWriteContext,
377 ) -> Result<(), ArrowError> {
378 match column.data_type() {
379 DataType::Struct(fields) => {
380 let s = as_struct_array(column);
381 for (field, column) in fields.iter().zip(s.columns()) {
382 self.encode_dictionaries(
383 field,
384 column,
385 encoded_dictionaries,
386 dictionary_tracker,
387 write_options,
388 dict_id,
389 ipc_write_context,
390 )?;
391 }
392 }
393 DataType::RunEndEncoded(_, values) => {
394 let data = column.to_data();
395 if data.child_data().len() != 2 {
396 return Err(ArrowError::InvalidArgumentError(format!(
397 "The run encoded array should have exactly two child arrays. Found {}",
398 data.child_data().len()
399 )));
400 }
401 let values_array = make_array(data.child_data()[1].clone());
404 self.encode_dictionaries(
405 values,
406 &values_array,
407 encoded_dictionaries,
408 dictionary_tracker,
409 write_options,
410 dict_id,
411 ipc_write_context,
412 )?;
413 }
414 DataType::List(field) => {
415 let list = as_list_array(column);
416 self.encode_dictionaries(
417 field,
418 list.values(),
419 encoded_dictionaries,
420 dictionary_tracker,
421 write_options,
422 dict_id,
423 ipc_write_context,
424 )?;
425 }
426 DataType::LargeList(field) => {
427 let list = as_large_list_array(column);
428 self.encode_dictionaries(
429 field,
430 list.values(),
431 encoded_dictionaries,
432 dictionary_tracker,
433 write_options,
434 dict_id,
435 ipc_write_context,
436 )?;
437 }
438 DataType::ListView(field) => {
439 let list = column.as_list_view::<i32>();
440 self.encode_dictionaries(
441 field,
442 list.values(),
443 encoded_dictionaries,
444 dictionary_tracker,
445 write_options,
446 dict_id,
447 ipc_write_context,
448 )?;
449 }
450 DataType::LargeListView(field) => {
451 let list = column.as_list_view::<i64>();
452 self.encode_dictionaries(
453 field,
454 list.values(),
455 encoded_dictionaries,
456 dictionary_tracker,
457 write_options,
458 dict_id,
459 ipc_write_context,
460 )?;
461 }
462 DataType::FixedSizeList(field, _) => {
463 let list = column
464 .as_any()
465 .downcast_ref::<FixedSizeListArray>()
466 .expect("Unable to downcast to fixed size list array");
467 self.encode_dictionaries(
468 field,
469 list.values(),
470 encoded_dictionaries,
471 dictionary_tracker,
472 write_options,
473 dict_id,
474 ipc_write_context,
475 )?;
476 }
477 DataType::Map(field, _) => {
478 let map_array = as_map_array(column);
479
480 let (keys, values) = match field.data_type() {
481 DataType::Struct(fields) if fields.len() == 2 => (&fields[0], &fields[1]),
482 _ => panic!("Incorrect field data type {:?}", field.data_type()),
483 };
484
485 self.encode_dictionaries(
487 keys,
488 map_array.keys(),
489 encoded_dictionaries,
490 dictionary_tracker,
491 write_options,
492 dict_id,
493 ipc_write_context,
494 )?;
495
496 self.encode_dictionaries(
498 values,
499 map_array.values(),
500 encoded_dictionaries,
501 dictionary_tracker,
502 write_options,
503 dict_id,
504 ipc_write_context,
505 )?;
506 }
507 DataType::Union(fields, _) => {
508 let union = as_union_array(column);
509 for (type_id, field) in fields.iter() {
510 let column = union.child(type_id);
511 self.encode_dictionaries(
512 field,
513 column,
514 encoded_dictionaries,
515 dictionary_tracker,
516 write_options,
517 dict_id,
518 ipc_write_context,
519 )?;
520 }
521 }
522 _ => (),
523 }
524
525 Ok(())
526 }
527
528 #[allow(clippy::too_many_arguments)]
529 fn encode_dictionaries<I: Iterator<Item = i64>>(
530 &self,
531 field: &Field,
532 column: &ArrayRef,
533 encoded_dictionaries: &mut Vec<EncodedData>,
534 dictionary_tracker: &mut DictionaryTracker,
535 write_options: &IpcWriteOptions,
536 dict_id_seq: &mut I,
537 ipc_write_context: &mut IpcWriteContext,
538 ) -> Result<(), ArrowError> {
539 match column.data_type() {
540 DataType::Dictionary(_key_type, value_type) => {
541 if matches!(value_type.as_ref(), DataType::Dictionary(_, _)) {
542 return Err(ArrowError::InvalidArgumentError(format!(
543 "Arrow IPC field metadata cannot encode direct dictionary-of-dictionary values for field {:?}",
544 field.name()
545 )));
546 }
547
548 let dict_data = column.to_data();
549 let dict_values = &dict_data.child_data()[0];
550
551 let values = make_array(dict_data.child_data()[0].clone());
552
553 self._encode_dictionaries(
554 &values,
555 encoded_dictionaries,
556 dictionary_tracker,
557 write_options,
558 dict_id_seq,
559 ipc_write_context,
560 )?;
561
562 let dict_id = dict_id_seq.next().ok_or_else(|| {
566 ArrowError::IpcError(format!(
567 "no dict id for field {:?}: field.data_type={:?}, column.data_type={:?}",
568 field.name(),
569 field.data_type(),
570 column.data_type()
571 ))
572 })?;
573
574 match dictionary_tracker.insert_column(
575 dict_id,
576 column,
577 write_options.dictionary_handling,
578 )? {
579 DictionaryUpdate::None => {}
580 DictionaryUpdate::New | DictionaryUpdate::Replaced => {
581 encoded_dictionaries.push(self.dictionary_batch_to_bytes(
582 dict_id,
583 dict_values,
584 write_options,
585 false,
586 ipc_write_context,
587 )?);
588 }
589 DictionaryUpdate::Delta(data) => {
590 encoded_dictionaries.push(self.dictionary_batch_to_bytes(
591 dict_id,
592 &data,
593 write_options,
594 true,
595 ipc_write_context,
596 )?);
597 }
598 }
599 }
600 _ => self._encode_dictionaries(
601 column,
602 encoded_dictionaries,
603 dictionary_tracker,
604 write_options,
605 dict_id_seq,
606 ipc_write_context,
607 )?,
608 }
609
610 Ok(())
611 }
612
613 pub fn encode(
617 &self,
618 batch: &RecordBatch,
619 dictionary_tracker: &mut DictionaryTracker,
620 write_options: &IpcWriteOptions,
621 ipc_write_context: &mut IpcWriteContext,
622 ) -> Result<(Vec<EncodedData>, EncodedData), ArrowError> {
623 let encoded_dictionaries =
624 self.encode_all_dicts(batch, dictionary_tracker, write_options, ipc_write_context)?;
625 let mut arrow_data = ipc_write_context.scratch();
626 let (ipc_message, _, tail_pad) = self.record_batch_to_bytes(
627 batch,
628 write_options,
629 ipc_write_context,
630 &mut IpcBodySink::Write(&mut arrow_data),
631 )?;
632 arrow_data.extend_from_slice(&PADDING[..tail_pad]);
633 ipc_write_context.reserve_scratch_with_capacity(arrow_data.capacity());
634 Ok((
635 encoded_dictionaries,
636 EncodedData {
637 ipc_message,
638 arrow_data,
639 },
640 ))
641 }
642
643 fn encode_all_dicts(
645 &self,
646 batch: &RecordBatch,
647 dictionary_tracker: &mut DictionaryTracker,
648 write_options: &IpcWriteOptions,
649 ipc_write_context: &mut IpcWriteContext,
650 ) -> Result<Vec<EncodedData>, ArrowError> {
651 let schema = batch.schema();
652 let mut encoded_dictionaries = Vec::with_capacity(schema.flattened_fields().len());
653 let mut dict_id = dictionary_tracker.dict_ids.clone().into_iter();
654 for (i, field) in schema.fields().iter().enumerate() {
655 self.encode_dictionaries(
656 field,
657 batch.column(i),
658 &mut encoded_dictionaries,
659 dictionary_tracker,
660 write_options,
661 &mut dict_id,
662 ipc_write_context,
663 )?;
664 }
665 Ok(encoded_dictionaries)
666 }
667
668 fn write<W: Write>(
672 &self,
673 batch: &RecordBatch,
674 dictionary_tracker: &mut DictionaryTracker,
675 write_options: &IpcWriteOptions,
676 ipc_write_context: &mut IpcWriteContext,
677 writer: &mut W,
678 ) -> Result<IpcWriteMetadata, ArrowError> {
679 let encoded_dictionaries =
680 self.encode_all_dicts(batch, dictionary_tracker, write_options, ipc_write_context)?;
681
682 let mut dictionary_block_sizes = Vec::with_capacity(encoded_dictionaries.len());
683 for dict in encoded_dictionaries {
684 dictionary_block_sizes.push(write_message(&mut *writer, dict, write_options)?);
685 }
686
687 let capacity = batch
688 .columns()
689 .iter()
690 .map(|a| estimate_encoded_buffer_count(a.data_type()))
691 .sum();
692 let mut encoded_buffers: Vec<EncodedBuffer> = Vec::with_capacity(capacity);
693 let (ipc_message, body_len, tail_pad) = self.record_batch_to_bytes(
694 batch,
695 write_options,
696 ipc_write_context,
697 &mut IpcBodySink::Collect(&mut encoded_buffers),
698 )?;
699
700 let alignment = write_options.alignment;
701 let a = usize::from(alignment - 1);
702 let prefix_size = if write_options.write_legacy_ipc_format {
703 4
704 } else {
705 8
706 };
707 let aligned_size = (ipc_message.len() + prefix_size + a) & !a;
708 write_continuation(
709 &mut *writer,
710 write_options,
711 (aligned_size - prefix_size) as i32,
712 )?;
713 writer.write_all(&ipc_message)?;
714 writer.write_all(&PADDING[..aligned_size - ipc_message.len() - prefix_size])?;
715 for enc in &encoded_buffers {
716 writer.write_all(enc.as_slice())?;
717 writer.write_all(&PADDING[..pad_to_alignment(alignment, enc.len())])?;
718 }
719 writer.write_all(&PADDING[..tail_pad])?;
720
721 Ok(IpcWriteMetadata {
722 dictionary_block_sizes,
723 padded_header_len: aligned_size,
724 body_len,
725 })
726 }
727
728 #[deprecated(since = "57.0.0", note = "Use `encode` instead")]
732 pub fn encoded_batch(
733 &self,
734 batch: &RecordBatch,
735 dictionary_tracker: &mut DictionaryTracker,
736 write_options: &IpcWriteOptions,
737 ) -> Result<(Vec<EncodedData>, EncodedData), ArrowError> {
738 self.encode(
739 batch,
740 dictionary_tracker,
741 write_options,
742 &mut Default::default(),
743 )
744 }
745
746 fn record_batch_to_bytes(
752 &self,
753 batch: &RecordBatch,
754 write_options: &IpcWriteOptions,
755 ipc_write_context: &mut IpcWriteContext,
756 sink: &mut IpcBodySink<'_>,
757 ) -> Result<(Vec<u8>, usize, usize), ArrowError> {
758 let batch_compression_type = write_options.batch_compression_type;
759
760 let compression = batch_compression_type.map(|batch_compression_type| {
761 let fbb = ipc_write_context.mut_fbb();
762 let mut c = crate::BodyCompressionBuilder::new(fbb);
763 c.add_method(crate::BodyCompressionMethod::BUFFER);
764 c.add_codec(batch_compression_type);
765 c.finish()
766 });
767
768 let batch_compression_level = write_options.batch_compression_level;
769 let compression_codec: Option<CompressionCodec> = batch_compression_type
770 .map(|compression_type| match batch_compression_level {
771 Some(level) => {
772 CompressionCodec::try_new_with_compression_level(compression_type, level)
773 }
774 None => compression_type.try_into(),
775 })
776 .transpose()?;
777
778 let alignment = write_options.alignment;
779 let mut variadic_buffer_counts = vec![];
780 let mut meta = IpcMetadataBuilder::default();
781 let mut offset = 0i64;
782
783 for array in batch.columns() {
784 let array_data = array.to_data();
785 offset = write_array_data(
786 &array_data,
787 &mut meta,
788 sink,
789 offset,
790 compression_codec,
791 ipc_write_context,
792 write_options,
793 )?;
794 append_variadic_buffer_counts(&mut variadic_buffer_counts, &array_data);
795 }
796
797 let tail_pad = pad_to_alignment(alignment, offset as usize);
798 let body_len = offset as usize + tail_pad;
799
800 let fbb = ipc_write_context.mut_fbb();
801 let buffers = fbb.create_vector(&meta.buffers);
802 let nodes = fbb.create_vector(&meta.nodes);
803 let variadic_buffer = if variadic_buffer_counts.is_empty() {
804 None
805 } else {
806 Some(fbb.create_vector(&variadic_buffer_counts))
807 };
808
809 let root = {
810 let mut batch_builder = crate::RecordBatchBuilder::new(fbb);
811 batch_builder.add_length(batch.num_rows() as i64);
812 batch_builder.add_nodes(nodes);
813 batch_builder.add_buffers(buffers);
814 if let Some(c) = compression {
815 batch_builder.add_compression(c);
816 }
817 if let Some(v) = variadic_buffer {
818 batch_builder.add_variadicBufferCounts(v);
819 }
820 batch_builder.finish().as_union_value()
821 };
822 let mut message = crate::MessageBuilder::new(fbb);
823 message.add_version(write_options.metadata_version);
824 message.add_header_type(crate::MessageHeader::RecordBatch);
825 message.add_bodyLength(body_len as i64);
826 message.add_header(root);
827 let root = message.finish();
828 fbb.finish(root, None);
829
830 let ipc_message = fbb.finished_data().to_vec();
831 fbb.reset();
832 Ok((ipc_message, body_len, tail_pad))
833 }
834
835 fn dictionary_batch_to_bytes(
838 &self,
839 dict_id: i64,
840 array_data: &ArrayData,
841 write_options: &IpcWriteOptions,
842 is_delta: bool,
843 ipc_write_context: &mut IpcWriteContext,
844 ) -> Result<EncodedData, ArrowError> {
845 let mut arrow_data: Vec<u8> = vec![];
846
847 let batch_compression_type = write_options.batch_compression_type;
849
850 let compression = batch_compression_type.map(|batch_compression_type| {
851 let fbb = ipc_write_context.mut_fbb();
852 let mut c = crate::BodyCompressionBuilder::new(fbb);
853 c.add_method(crate::BodyCompressionMethod::BUFFER);
854 c.add_codec(batch_compression_type);
855 c.finish()
856 });
857
858 let batch_compression_level = write_options.batch_compression_level;
859 let compression_codec: Option<CompressionCodec> = batch_compression_type
860 .map(|batch_compression_type| match batch_compression_level {
861 Some(level) => {
862 CompressionCodec::try_new_with_compression_level(batch_compression_type, level)
863 }
864 None => batch_compression_type.try_into(),
865 })
866 .transpose()?;
867
868 let alignment = write_options.alignment;
869 let mut meta = IpcMetadataBuilder::default();
870 let mut sink = IpcBodySink::Write(&mut arrow_data);
871 let offset = write_array_data(
872 array_data,
873 &mut meta,
874 &mut sink,
875 0,
876 compression_codec,
877 ipc_write_context,
878 write_options,
879 )?;
880
881 let mut variadic_buffer_counts = vec![];
882 append_variadic_buffer_counts(&mut variadic_buffer_counts, array_data);
883
884 let tail_pad = pad_to_alignment(alignment, offset as usize);
886 let body_len = offset as usize + tail_pad;
887 arrow_data.extend_from_slice(&PADDING[..tail_pad]);
888
889 let fbb = ipc_write_context.mut_fbb();
890 let buffers = fbb.create_vector(&meta.buffers);
891 let nodes = fbb.create_vector(&meta.nodes);
892 let variadic_buffer = if variadic_buffer_counts.is_empty() {
893 None
894 } else {
895 Some(fbb.create_vector(&variadic_buffer_counts))
896 };
897
898 let root = {
899 let mut batch_builder = crate::RecordBatchBuilder::new(fbb);
900 batch_builder.add_length(array_data.len() as i64);
901 batch_builder.add_nodes(nodes);
902 batch_builder.add_buffers(buffers);
903 if let Some(c) = compression {
904 batch_builder.add_compression(c);
905 }
906 if let Some(v) = variadic_buffer {
907 batch_builder.add_variadicBufferCounts(v);
908 }
909 batch_builder.finish()
910 };
911
912 let root = {
913 let mut batch_builder = crate::DictionaryBatchBuilder::new(fbb);
914 batch_builder.add_id(dict_id);
915 batch_builder.add_data(root);
916 batch_builder.add_isDelta(is_delta);
917 batch_builder.finish().as_union_value()
918 };
919
920 let root = {
921 let mut message_builder = crate::MessageBuilder::new(fbb);
922 message_builder.add_version(write_options.metadata_version);
923 message_builder.add_header_type(crate::MessageHeader::DictionaryBatch);
924 message_builder.add_bodyLength(body_len as i64);
925 message_builder.add_header(root);
926 message_builder.finish()
927 };
928
929 fbb.finish(root, None);
930 let ipc_message = fbb.finished_data().to_vec();
931 fbb.reset();
932
933 Ok(EncodedData {
934 ipc_message,
935 arrow_data,
936 })
937 }
938}
939
940fn ensure_supported_ipc_schema(schema: &Schema) -> Result<(), ArrowError> {
941 schema
942 .fields()
943 .iter()
944 .try_for_each(|field| ensure_supported_ipc_data_type(field.name(), field.data_type()))
945}
946
947fn ensure_supported_ipc_data_type(
948 field_name: &str,
949 data_type: &DataType,
950) -> Result<(), ArrowError> {
951 match data_type {
952 DataType::Dictionary(_, value_type)
953 if matches!(value_type.as_ref(), DataType::Dictionary(_, _)) =>
954 {
955 Err(ArrowError::InvalidArgumentError(format!(
956 "Arrow IPC field metadata cannot encode direct dictionary-of-dictionary values for field {field_name:?}"
957 )))
958 }
959 DataType::Dictionary(_, value_type) => {
960 ensure_supported_ipc_data_type(field_name, value_type)
961 }
962 DataType::Struct(fields) => fields
963 .iter()
964 .try_for_each(|field| ensure_supported_ipc_data_type(field.name(), field.data_type())),
965 DataType::RunEndEncoded(_, field)
966 | DataType::List(field)
967 | DataType::LargeList(field)
968 | DataType::ListView(field)
969 | DataType::LargeListView(field)
970 | DataType::FixedSizeList(field, _)
971 | DataType::Map(field, _) => {
972 ensure_supported_ipc_data_type(field.name(), field.data_type())
973 }
974 DataType::Union(fields, _) => fields.iter().try_for_each(|(_, field)| {
975 ensure_supported_ipc_data_type(field.name(), field.data_type())
976 }),
977 _ => Ok(()),
978 }
979}
980
981fn append_variadic_buffer_counts(counts: &mut Vec<i64>, array: &ArrayData) {
982 match array.data_type() {
983 DataType::BinaryView | DataType::Utf8View => {
984 counts.push(array.buffers().len() as i64 - 1);
987 }
988 DataType::Dictionary(_, _) => {
989 }
992 _ => {
993 for child in array.child_data() {
994 append_variadic_buffer_counts(counts, child)
995 }
996 }
997 }
998}
999
1000pub(crate) fn unslice_run_array(arr: ArrayData) -> Result<ArrayData, ArrowError> {
1001 match arr.data_type() {
1002 DataType::RunEndEncoded(k, _) => match k.data_type() {
1003 DataType::Int16 => {
1004 Ok(into_zero_offset_run_array(RunArray::<Int16Type>::from(arr))?.into_data())
1005 }
1006 DataType::Int32 => {
1007 Ok(into_zero_offset_run_array(RunArray::<Int32Type>::from(arr))?.into_data())
1008 }
1009 DataType::Int64 => {
1010 Ok(into_zero_offset_run_array(RunArray::<Int64Type>::from(arr))?.into_data())
1011 }
1012 d => unreachable!("Unexpected data type {d}"),
1013 },
1014 d => Err(ArrowError::InvalidArgumentError(format!(
1015 "The given array is not a run array. Data type of given array: {d}"
1016 ))),
1017 }
1018}
1019
1020fn into_zero_offset_run_array<R: RunEndIndexType>(
1023 run_array: RunArray<R>,
1024) -> Result<RunArray<R>, ArrowError> {
1025 let run_ends = run_array.run_ends();
1026 if run_ends.offset() == 0 && run_ends.max_value() == run_ends.len() {
1027 return Ok(run_array);
1028 }
1029
1030 let start_physical_index = run_ends.get_start_physical_index();
1032
1033 let end_physical_index = run_ends.get_end_physical_index();
1035
1036 let physical_length = end_physical_index - start_physical_index + 1;
1037
1038 let offset = R::Native::usize_as(run_ends.offset());
1040 let mut builder = BufferBuilder::<R::Native>::new(physical_length);
1041 for run_end_value in &run_ends.values()[start_physical_index..end_physical_index] {
1042 builder.append(run_end_value.sub_wrapping(offset));
1043 }
1044 builder.append(R::Native::from_usize(run_array.len()).unwrap());
1045 let new_run_ends = unsafe {
1046 ArrayDataBuilder::new(R::DATA_TYPE)
1049 .len(physical_length)
1050 .add_buffer(builder.finish())
1051 .build_unchecked()
1052 };
1053
1054 let new_values = run_array
1056 .values()
1057 .slice(start_physical_index, physical_length)
1058 .into_data();
1059
1060 let builder = ArrayDataBuilder::new(run_array.data_type().clone())
1061 .len(run_array.len())
1062 .add_child_data(new_run_ends)
1063 .add_child_data(new_values);
1064 let array_data = unsafe {
1065 builder.build_unchecked()
1068 };
1069 Ok(array_data.into())
1070}
1071
1072#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
1074pub enum DictionaryHandling {
1075 #[default]
1077 Resend,
1078 Delta,
1084}
1085
1086#[derive(Debug, Clone)]
1088pub enum DictionaryUpdate {
1089 None,
1092 New,
1094 Replaced,
1096 Delta(ArrayData),
1098}
1099
1100#[derive(Debug)]
1106pub struct DictionaryTracker {
1107 written: HashMap<i64, ArrayData>,
1109 dict_ids: Vec<i64>,
1110 error_on_replacement: bool,
1111}
1112
1113impl DictionaryTracker {
1114 pub fn new(error_on_replacement: bool) -> Self {
1120 #[allow(deprecated)]
1121 Self {
1122 written: HashMap::new(),
1123 dict_ids: Vec::new(),
1124 error_on_replacement,
1125 }
1126 }
1127
1128 pub fn next_dict_id(&mut self) -> i64 {
1130 let next = self
1131 .dict_ids
1132 .last()
1133 .copied()
1134 .map(|i| i + 1)
1135 .unwrap_or_default();
1136
1137 self.dict_ids.push(next);
1138 next
1139 }
1140
1141 pub fn dict_id(&mut self) -> &[i64] {
1144 &self.dict_ids
1145 }
1146
1147 #[deprecated(since = "56.1.0", note = "Use `insert_column` instead")]
1157 pub fn insert(&mut self, dict_id: i64, column: &ArrayRef) -> Result<bool, ArrowError> {
1158 let dict_data = column.to_data();
1159 let dict_values = &dict_data.child_data()[0];
1160
1161 if let Some(last) = self.written.get(&dict_id) {
1163 if ArrayData::ptr_eq(&last.child_data()[0], dict_values) {
1164 return Ok(false);
1166 }
1167 if self.error_on_replacement {
1168 if last.child_data()[0] == *dict_values {
1170 return Ok(false);
1172 }
1173 return Err(ArrowError::InvalidArgumentError(
1174 "Dictionary replacement detected when writing IPC file format. \
1175 Arrow IPC files only support a single dictionary for a given field \
1176 across all batches."
1177 .to_string(),
1178 ));
1179 }
1180 }
1181
1182 self.written.insert(dict_id, dict_data);
1183 Ok(true)
1184 }
1185
1186 pub fn insert_column(
1202 &mut self,
1203 dict_id: i64,
1204 column: &ArrayRef,
1205 dict_handling: DictionaryHandling,
1206 ) -> Result<DictionaryUpdate, ArrowError> {
1207 let new_data = column.to_data();
1208 let new_values = &new_data.child_data()[0];
1209
1210 let Some(old) = self.written.get(&dict_id) else {
1212 self.written.insert(dict_id, new_data);
1213 return Ok(DictionaryUpdate::New);
1214 };
1215
1216 let old_values = &old.child_data()[0];
1219 if ArrayData::ptr_eq(old_values, new_values) {
1220 return Ok(DictionaryUpdate::None);
1221 }
1222
1223 let comparison = compare_dictionaries(old_values, new_values);
1225 if matches!(comparison, DictionaryComparison::Equal) {
1226 return Ok(DictionaryUpdate::None);
1227 }
1228
1229 const REPLACEMENT_ERROR: &str = "Dictionary replacement detected when writing IPC file format. \
1230 Arrow IPC files only support a single dictionary for a given field \
1231 across all batches.";
1232
1233 match comparison {
1234 DictionaryComparison::NotEqual => {
1235 if self.error_on_replacement {
1236 return Err(ArrowError::InvalidArgumentError(
1237 REPLACEMENT_ERROR.to_string(),
1238 ));
1239 }
1240
1241 self.written.insert(dict_id, new_data);
1242 Ok(DictionaryUpdate::Replaced)
1243 }
1244 DictionaryComparison::Delta => match dict_handling {
1245 DictionaryHandling::Resend => {
1246 if self.error_on_replacement {
1247 return Err(ArrowError::InvalidArgumentError(
1248 REPLACEMENT_ERROR.to_string(),
1249 ));
1250 }
1251
1252 self.written.insert(dict_id, new_data);
1253 Ok(DictionaryUpdate::Replaced)
1254 }
1255 DictionaryHandling::Delta => {
1256 let delta =
1257 new_values.slice(old_values.len(), new_values.len() - old_values.len());
1258 self.written.insert(dict_id, new_data);
1259 Ok(DictionaryUpdate::Delta(delta))
1260 }
1261 },
1262 DictionaryComparison::Equal => unreachable!("Already checked equal case"),
1263 }
1264 }
1265
1266 pub fn clear(&mut self) {
1272 self.dict_ids.clear();
1273 self.written.clear();
1274 }
1275}
1276
1277#[derive(Debug, Clone)]
1279enum DictionaryComparison {
1280 NotEqual,
1282 Equal,
1284 Delta,
1287}
1288
1289fn compare_dictionaries(old: &ArrayData, new: &ArrayData) -> DictionaryComparison {
1291 let existing_len = old.len();
1293 let new_len = new.len();
1294 if existing_len == new_len {
1295 if *old == *new {
1296 return DictionaryComparison::Equal;
1297 } else {
1298 return DictionaryComparison::NotEqual;
1299 }
1300 }
1301
1302 if new_len < existing_len {
1304 return DictionaryComparison::NotEqual;
1305 }
1306
1307 if new.slice(0, existing_len) == *old {
1309 return DictionaryComparison::Delta;
1310 }
1311
1312 DictionaryComparison::NotEqual
1313}
1314
1315pub struct FileWriter<W> {
1338 writer: W,
1340 write_options: IpcWriteOptions,
1342 schema: SchemaRef,
1344 block_offsets: usize,
1346 dictionary_blocks: Vec<crate::Block>,
1348 record_blocks: Vec<crate::Block>,
1350 finished: bool,
1352 dictionary_tracker: DictionaryTracker,
1354 custom_metadata: HashMap<String, String>,
1356
1357 data_gen: IpcDataGenerator,
1358
1359 ipc_write_context: IpcWriteContext,
1360}
1361
1362impl<W: Write> FileWriter<BufWriter<W>> {
1363 pub fn try_new_buffered(writer: W, schema: &Schema) -> Result<Self, ArrowError> {
1367 Self::try_new(BufWriter::new(writer), schema)
1368 }
1369}
1370
1371impl<W: Write> FileWriter<W> {
1372 pub fn try_new(writer: W, schema: &Schema) -> Result<Self, ArrowError> {
1380 let write_options = IpcWriteOptions::default();
1381 Self::try_new_with_options(writer, schema, write_options)
1382 }
1383
1384 pub fn try_new_with_options(
1392 mut writer: W,
1393 schema: &Schema,
1394 write_options: IpcWriteOptions,
1395 ) -> Result<Self, ArrowError> {
1396 ensure_supported_ipc_schema(schema)?;
1397
1398 let data_gen = IpcDataGenerator::default();
1399 let pad_len = pad_to_alignment(write_options.alignment, super::ARROW_MAGIC.len());
1401 let header_size = super::ARROW_MAGIC.len() + pad_len;
1402 writer.write_all(&super::ARROW_MAGIC)?;
1403 writer.write_all(&PADDING[..pad_len])?;
1404 let mut dictionary_tracker = DictionaryTracker::new(true);
1406 let encoded_message = data_gen.schema_to_bytes_with_dictionary_tracker(
1407 schema,
1408 &mut dictionary_tracker,
1409 &write_options,
1410 );
1411 let (meta, data) = write_message(&mut writer, encoded_message, &write_options)?;
1412 Ok(Self {
1413 writer,
1414 write_options,
1415 schema: Arc::new(schema.clone()),
1416 block_offsets: meta + data + header_size,
1417 dictionary_blocks: vec![],
1418 record_blocks: vec![],
1419 finished: false,
1420 dictionary_tracker,
1421 custom_metadata: HashMap::new(),
1422 data_gen,
1423 ipc_write_context: IpcWriteContext::default(),
1424 })
1425 }
1426
1427 pub fn write_metadata(&mut self, key: impl Into<String>, value: impl Into<String>) {
1429 self.custom_metadata.insert(key.into(), value.into());
1430 }
1431
1432 pub fn write(&mut self, batch: &RecordBatch) -> Result<(), ArrowError> {
1434 if self.finished {
1435 return Err(ArrowError::IpcError(
1436 "Cannot write record batch to file writer as it is closed".to_string(),
1437 ));
1438 }
1439
1440 let meta = self.data_gen.write(
1441 batch,
1442 &mut self.dictionary_tracker,
1443 &self.write_options,
1444 &mut self.ipc_write_context,
1445 &mut self.writer,
1446 )?;
1447
1448 for (header_len, body_len) in meta.dictionary_block_sizes {
1449 let block = crate::Block::new(
1450 self.block_offsets as i64,
1451 header_len as i32,
1452 body_len as i64,
1453 );
1454 self.dictionary_blocks.push(block);
1455 self.block_offsets += header_len + body_len;
1456 }
1457
1458 let block = crate::Block::new(
1460 self.block_offsets as i64,
1461 meta.padded_header_len as i32,
1462 meta.body_len as i64,
1463 );
1464 self.record_blocks.push(block);
1465 self.block_offsets += meta.padded_header_len + meta.body_len;
1466 Ok(())
1467 }
1468
1469 pub fn finish(&mut self) -> Result<(), ArrowError> {
1471 if self.finished {
1472 return Err(ArrowError::IpcError(
1473 "Cannot write footer to file writer as it is closed".to_string(),
1474 ));
1475 }
1476
1477 write_continuation(&mut self.writer, &self.write_options, 0)?;
1479
1480 let mut fbb = FlatBufferBuilder::new();
1481 let dictionaries = fbb.create_vector(&self.dictionary_blocks);
1482 let record_batches = fbb.create_vector(&self.record_blocks);
1483
1484 self.dictionary_tracker.clear();
1486 let schema = IpcSchemaEncoder::new()
1487 .with_dictionary_tracker(&mut self.dictionary_tracker)
1488 .schema_to_fb_offset(&mut fbb, &self.schema);
1489 let fb_custom_metadata = (!self.custom_metadata.is_empty())
1490 .then(|| crate::convert::metadata_to_fb(&mut fbb, &self.custom_metadata));
1491
1492 let root = {
1493 let mut footer_builder = crate::FooterBuilder::new(&mut fbb);
1494 footer_builder.add_version(self.write_options.metadata_version);
1495 footer_builder.add_schema(schema);
1496 footer_builder.add_dictionaries(dictionaries);
1497 footer_builder.add_recordBatches(record_batches);
1498 if let Some(fb_custom_metadata) = fb_custom_metadata {
1499 footer_builder.add_custom_metadata(fb_custom_metadata);
1500 }
1501 footer_builder.finish()
1502 };
1503 fbb.finish(root, None);
1504 let footer_data = fbb.finished_data();
1505 self.writer.write_all(footer_data)?;
1506 self.writer
1507 .write_all(&(footer_data.len() as i32).to_le_bytes())?;
1508 self.writer.write_all(&super::ARROW_MAGIC)?;
1509 self.writer.flush()?;
1510 self.finished = true;
1511
1512 Ok(())
1513 }
1514
1515 pub fn schema(&self) -> &SchemaRef {
1517 &self.schema
1518 }
1519
1520 pub fn get_ref(&self) -> &W {
1522 &self.writer
1523 }
1524
1525 pub fn get_mut(&mut self) -> &mut W {
1529 &mut self.writer
1530 }
1531
1532 pub fn flush(&mut self) -> Result<(), ArrowError> {
1536 self.writer.flush()?;
1537 Ok(())
1538 }
1539
1540 pub fn into_inner(mut self) -> Result<W, ArrowError> {
1549 if !self.finished {
1550 self.finish()?;
1552 }
1553 Ok(self.writer)
1554 }
1555}
1556
1557impl<W: Write> RecordBatchWriter for FileWriter<W> {
1558 fn write(&mut self, batch: &RecordBatch) -> Result<(), ArrowError> {
1559 self.write(batch)
1560 }
1561
1562 fn close(mut self) -> Result<(), ArrowError> {
1563 self.finish()
1564 }
1565}
1566
1567pub struct StreamWriter<W> {
1641 writer: W,
1643 write_options: IpcWriteOptions,
1645 finished: bool,
1647 dictionary_tracker: DictionaryTracker,
1649
1650 data_gen: IpcDataGenerator,
1651
1652 ipc_write_context: IpcWriteContext,
1653}
1654
1655impl<W: Write> StreamWriter<BufWriter<W>> {
1656 pub fn try_new_buffered(writer: W, schema: &Schema) -> Result<Self, ArrowError> {
1660 Self::try_new(BufWriter::new(writer), schema)
1661 }
1662}
1663
1664impl<W: Write> StreamWriter<W> {
1665 pub fn try_new(writer: W, schema: &Schema) -> Result<Self, ArrowError> {
1673 let write_options = IpcWriteOptions::default();
1674 Self::try_new_with_options(writer, schema, write_options)
1675 }
1676
1677 pub fn try_new_with_options(
1683 mut writer: W,
1684 schema: &Schema,
1685 write_options: IpcWriteOptions,
1686 ) -> Result<Self, ArrowError> {
1687 ensure_supported_ipc_schema(schema)?;
1688
1689 let data_gen = IpcDataGenerator::default();
1690 let mut dictionary_tracker = DictionaryTracker::new(false);
1691
1692 let encoded_message = data_gen.schema_to_bytes_with_dictionary_tracker(
1694 schema,
1695 &mut dictionary_tracker,
1696 &write_options,
1697 );
1698 write_message(&mut writer, encoded_message, &write_options)?;
1699 Ok(Self {
1700 writer,
1701 write_options,
1702 finished: false,
1703 dictionary_tracker,
1704 data_gen,
1705 ipc_write_context: IpcWriteContext::default(),
1706 })
1707 }
1708
1709 pub fn write(&mut self, batch: &RecordBatch) -> Result<(), ArrowError> {
1711 if self.finished {
1712 return Err(ArrowError::IpcError(
1713 "Cannot write record batch to stream writer as it is closed".to_string(),
1714 ));
1715 }
1716
1717 self.data_gen.write(
1718 batch,
1719 &mut self.dictionary_tracker,
1720 &self.write_options,
1721 &mut self.ipc_write_context,
1722 &mut self.writer,
1723 )?;
1724 Ok(())
1725 }
1726
1727 pub fn finish(&mut self) -> Result<(), ArrowError> {
1729 if self.finished {
1730 return Err(ArrowError::IpcError(
1731 "Cannot write footer to stream writer as it is closed".to_string(),
1732 ));
1733 }
1734
1735 write_continuation(&mut self.writer, &self.write_options, 0)?;
1736 self.writer.flush()?;
1737
1738 self.finished = true;
1739
1740 Ok(())
1741 }
1742
1743 pub fn get_ref(&self) -> &W {
1745 &self.writer
1746 }
1747
1748 pub fn get_mut(&mut self) -> &mut W {
1752 &mut self.writer
1753 }
1754
1755 pub fn flush(&mut self) -> Result<(), ArrowError> {
1759 self.writer.flush()?;
1760 Ok(())
1761 }
1762
1763 pub fn into_inner(mut self) -> Result<W, ArrowError> {
1801 if !self.finished {
1802 self.finish()?;
1804 }
1805 Ok(self.writer)
1806 }
1807}
1808
1809impl<W: Write> RecordBatchWriter for StreamWriter<W> {
1810 fn write(&mut self, batch: &RecordBatch) -> Result<(), ArrowError> {
1811 self.write(batch)
1812 }
1813
1814 fn close(mut self) -> Result<(), ArrowError> {
1815 self.finish()
1816 }
1817}
1818
1819pub struct EncodedData {
1821 pub ipc_message: Vec<u8>,
1823 pub arrow_data: Vec<u8>,
1825}
1826pub fn write_message<W: Write>(
1828 mut writer: W,
1829 encoded: EncodedData,
1830 write_options: &IpcWriteOptions,
1831) -> Result<(usize, usize), ArrowError> {
1832 let arrow_data_len = encoded.arrow_data.len();
1833 if arrow_data_len % usize::from(write_options.alignment) != 0 {
1834 return Err(ArrowError::MemoryError(
1835 "Arrow data not aligned".to_string(),
1836 ));
1837 }
1838
1839 let a = usize::from(write_options.alignment - 1);
1840 let buffer = encoded.ipc_message;
1841 let flatbuf_size = buffer.len();
1842 let prefix_size = if write_options.write_legacy_ipc_format {
1843 4
1844 } else {
1845 8
1846 };
1847 let aligned_size = (flatbuf_size + prefix_size + a) & !a;
1848 let padding_bytes = aligned_size - flatbuf_size - prefix_size;
1849
1850 write_continuation(
1851 &mut writer,
1852 write_options,
1853 (aligned_size - prefix_size) as i32,
1854 )?;
1855
1856 if flatbuf_size > 0 {
1858 writer.write_all(&buffer)?;
1859 }
1860 writer.write_all(&PADDING[..padding_bytes])?;
1862
1863 let body_len = if arrow_data_len > 0 {
1865 write_body_buffers(&mut writer, &encoded.arrow_data, write_options.alignment)?
1866 } else {
1867 0
1868 };
1869
1870 Ok((aligned_size, body_len))
1871}
1872
1873fn write_body_buffers<W: Write>(
1874 mut writer: W,
1875 data: &[u8],
1876 alignment: u8,
1877) -> Result<usize, ArrowError> {
1878 let len = data.len();
1879 let pad_len = pad_to_alignment(alignment, len);
1880 let total_len = len + pad_len;
1881
1882 writer.write_all(data)?;
1884 if pad_len > 0 {
1885 writer.write_all(&PADDING[..pad_len])?;
1886 }
1887
1888 Ok(total_len)
1889}
1890
1891fn write_continuation<W: Write>(
1894 mut writer: W,
1895 write_options: &IpcWriteOptions,
1896 total_len: i32,
1897) -> Result<usize, ArrowError> {
1898 let mut written = 8;
1899
1900 match write_options.metadata_version {
1902 crate::MetadataVersion::V1 | crate::MetadataVersion::V2 | crate::MetadataVersion::V3 => {
1903 unreachable!("Options with the metadata version cannot be created")
1904 }
1905 crate::MetadataVersion::V4 => {
1906 if !write_options.write_legacy_ipc_format {
1907 writer.write_all(&CONTINUATION_MARKER)?;
1909 written = 4;
1910 }
1911 writer.write_all(&total_len.to_le_bytes()[..])?;
1912 }
1913 crate::MetadataVersion::V5 => {
1914 writer.write_all(&CONTINUATION_MARKER)?;
1916 writer.write_all(&total_len.to_le_bytes()[..])?;
1917 }
1918 z => panic!("Unsupported crate::MetadataVersion {z:?}"),
1919 };
1920
1921 Ok(written)
1922}
1923
1924fn has_validity_bitmap(data_type: &DataType, write_options: &IpcWriteOptions) -> bool {
1928 if write_options.metadata_version < crate::MetadataVersion::V5 {
1929 !matches!(data_type, DataType::Null)
1930 } else {
1931 !matches!(
1932 data_type,
1933 DataType::Null | DataType::Union(_, _) | DataType::RunEndEncoded(_, _)
1934 )
1935 }
1936}
1937
1938#[inline]
1940fn buffer_need_truncate(
1941 array_offset: usize,
1942 buffer: &Buffer,
1943 spec: &BufferSpec,
1944 min_length: usize,
1945) -> bool {
1946 spec != &BufferSpec::AlwaysNull && (array_offset != 0 || min_length < buffer.len())
1947}
1948
1949#[inline]
1951fn get_buffer_element_width(spec: &BufferSpec) -> usize {
1952 match spec {
1953 BufferSpec::FixedWidth { byte_width, .. } => *byte_width,
1954 _ => 0,
1955 }
1956}
1957
1958fn reencode_offsets<O: OffsetSizeTrait>(
1961 offsets: &Buffer,
1962 data: &ArrayData,
1963) -> (Buffer, usize, usize) {
1964 let offsets_slice: &[O] = offsets.typed_data::<O>();
1965 let offset_slice = &offsets_slice[data.offset()..data.offset() + data.len() + 1];
1966
1967 let start_offset = offset_slice.first().unwrap();
1968 let end_offset = offset_slice.last().unwrap();
1969
1970 let offsets = match start_offset.as_usize() {
1971 0 => {
1972 let size = size_of::<O>();
1973 offsets.slice_with_length(data.offset() * size, (data.len() + 1) * size)
1974 }
1975 _ => offset_slice.iter().map(|x| *x - *start_offset).collect(),
1976 };
1977
1978 let start_offset = start_offset.as_usize();
1979 let end_offset = end_offset.as_usize();
1980
1981 (offsets, start_offset, end_offset - start_offset)
1982}
1983
1984fn get_byte_array_buffers<O: OffsetSizeTrait>(data: &ArrayData) -> (Buffer, Buffer) {
1990 if data.is_empty() {
1991 let mut offsets = MutableBuffer::new(size_of::<O>());
1994 offsets.extend_from_slice(O::usize_as(0).to_byte_slice());
1995 return (offsets.into(), MutableBuffer::new(0).into());
1996 }
1997
1998 let (offsets, original_start_offset, len) = reencode_offsets::<O>(&data.buffers()[0], data);
1999 let values = data.buffers()[1].slice_with_length(original_start_offset, len);
2000 (offsets, values)
2001}
2002
2003fn get_list_array_buffers<O: OffsetSizeTrait>(data: &ArrayData) -> (Buffer, ArrayData) {
2006 if data.is_empty() {
2007 let mut offsets = MutableBuffer::new(size_of::<O>());
2010 offsets.extend_from_slice(O::usize_as(0).to_byte_slice());
2011 return (offsets.into(), data.child_data()[0].slice(0, 0));
2012 }
2013
2014 let (offsets, original_start_offset, len) = reencode_offsets::<O>(&data.buffers()[0], data);
2015 let child_data = data.child_data()[0].slice(original_start_offset, len);
2016 (offsets, child_data)
2017}
2018
2019fn get_list_view_array_buffers<O: OffsetSizeTrait>(
2025 data: &ArrayData,
2026) -> (Buffer, Buffer, ArrayData) {
2027 if data.is_empty() {
2028 return (
2029 MutableBuffer::new(0).into(),
2030 MutableBuffer::new(0).into(),
2031 data.child_data()[0].slice(0, 0),
2032 );
2033 }
2034
2035 let offsets = &data.buffers()[0];
2036 let sizes = &data.buffers()[1];
2037
2038 let element_size = std::mem::size_of::<O>();
2039 let offsets_slice =
2040 offsets.slice_with_length(data.offset() * element_size, data.len() * element_size);
2041 let sizes_slice =
2042 sizes.slice_with_length(data.offset() * element_size, data.len() * element_size);
2043
2044 let child_data = data.child_data()[0].clone();
2045
2046 (offsets_slice, sizes_slice, child_data)
2047}
2048
2049fn get_or_truncate_buffer(array_data: &ArrayData) -> Buffer {
2056 let buffer = &array_data.buffers()[0];
2057 let layout = layout(array_data.data_type());
2058 let spec = &layout.buffers[0];
2059
2060 let byte_width = get_buffer_element_width(spec);
2061 let min_length = array_data.len() * byte_width;
2062 if buffer_need_truncate(array_data.offset(), buffer, spec, min_length) {
2063 let byte_offset = array_data.offset() * byte_width;
2064 let buffer_length = min(min_length, buffer.len() - byte_offset);
2065 buffer.slice_with_length(byte_offset, buffer_length)
2066 } else {
2067 buffer.clone()
2068 }
2069}
2070
2071fn write_array_data(
2077 array_data: &ArrayData,
2078 meta: &mut IpcMetadataBuilder,
2079 sink: &mut IpcBodySink<'_>,
2080 offset: i64,
2081 compression_codec: Option<CompressionCodec>,
2082 ipc_write_context: &mut IpcWriteContext,
2083 write_options: &IpcWriteOptions,
2084) -> Result<i64, ArrowError> {
2085 let mut offset = offset;
2086 let num_rows = array_data.len();
2087 if !matches!(array_data.data_type(), DataType::Null) {
2088 meta.nodes.push(crate::FieldNode::new(
2089 num_rows as i64,
2090 array_data.null_count() as i64,
2091 ));
2092 } else {
2093 meta.nodes
2095 .push(crate::FieldNode::new(num_rows as i64, num_rows as i64));
2096 }
2097 if has_validity_bitmap(array_data.data_type(), write_options) {
2098 let null_buffer = match array_data.nulls() {
2100 None => {
2101 let num_bytes = bit_util::ceil(num_rows, 8);
2103 let buffer = MutableBuffer::new(num_bytes);
2104 let buffer = buffer.with_bitset(num_bytes, true);
2105 buffer.into()
2106 }
2107 Some(buffer) => buffer.inner().sliced(),
2108 };
2109
2110 offset = encode_sink_buffer(
2111 null_buffer,
2112 meta,
2113 sink,
2114 offset,
2115 compression_codec,
2116 ipc_write_context,
2117 write_options.alignment,
2118 )?;
2119 }
2120
2121 let data_type = array_data.data_type();
2122 if matches!(data_type, DataType::Binary | DataType::Utf8) {
2123 let (offsets, values) = get_byte_array_buffers::<i32>(array_data);
2124 for buffer in [offsets, values] {
2125 offset = encode_sink_buffer(
2126 buffer,
2127 meta,
2128 sink,
2129 offset,
2130 compression_codec,
2131 ipc_write_context,
2132 write_options.alignment,
2133 )?;
2134 }
2135 } else if matches!(data_type, DataType::BinaryView | DataType::Utf8View) {
2136 let views = get_or_truncate_buffer(array_data);
2143 offset = encode_sink_buffer(
2144 views,
2145 meta,
2146 sink,
2147 offset,
2148 compression_codec,
2149 ipc_write_context,
2150 write_options.alignment,
2151 )?;
2152
2153 for buffer in array_data.buffers().iter().skip(1) {
2154 offset = encode_sink_buffer(
2155 buffer.clone(),
2156 meta,
2157 sink,
2158 offset,
2159 compression_codec,
2160 ipc_write_context,
2161 write_options.alignment,
2162 )?;
2163 }
2164 } else if matches!(data_type, DataType::LargeBinary | DataType::LargeUtf8) {
2165 let (offsets, values) = get_byte_array_buffers::<i64>(array_data);
2166 for buffer in [offsets, values] {
2167 offset = encode_sink_buffer(
2168 buffer,
2169 meta,
2170 sink,
2171 offset,
2172 compression_codec,
2173 ipc_write_context,
2174 write_options.alignment,
2175 )?;
2176 }
2177 } else if DataType::is_numeric(data_type)
2178 || DataType::is_temporal(data_type)
2179 || matches!(
2180 array_data.data_type(),
2181 DataType::FixedSizeBinary(_) | DataType::Dictionary(_, _)
2182 )
2183 {
2184 assert_eq!(array_data.buffers().len(), 1);
2186
2187 let buffer = get_or_truncate_buffer(array_data);
2188 offset = encode_sink_buffer(
2189 buffer,
2190 meta,
2191 sink,
2192 offset,
2193 compression_codec,
2194 ipc_write_context,
2195 write_options.alignment,
2196 )?;
2197 } else if matches!(data_type, DataType::Boolean) {
2198 assert_eq!(array_data.buffers().len(), 1);
2201
2202 let buffer = &array_data.buffers()[0];
2203 let buffer = buffer.bit_slice(array_data.offset(), array_data.len());
2204 offset = encode_sink_buffer(
2205 buffer,
2206 meta,
2207 sink,
2208 offset,
2209 compression_codec,
2210 ipc_write_context,
2211 write_options.alignment,
2212 )?;
2213 } else if matches!(
2214 data_type,
2215 DataType::List(_) | DataType::LargeList(_) | DataType::Map(_, _)
2216 ) {
2217 assert_eq!(array_data.buffers().len(), 1);
2218 assert_eq!(array_data.child_data().len(), 1);
2219
2220 let (offsets, sliced_child_data) = match data_type {
2222 DataType::List(_) => get_list_array_buffers::<i32>(array_data),
2223 DataType::Map(_, _) => get_list_array_buffers::<i32>(array_data),
2224 DataType::LargeList(_) => get_list_array_buffers::<i64>(array_data),
2225 _ => unreachable!(),
2226 };
2227 offset = encode_sink_buffer(
2228 offsets,
2229 meta,
2230 sink,
2231 offset,
2232 compression_codec,
2233 ipc_write_context,
2234 write_options.alignment,
2235 )?;
2236 offset = write_array_data(
2237 &sliced_child_data,
2238 meta,
2239 sink,
2240 offset,
2241 compression_codec,
2242 ipc_write_context,
2243 write_options,
2244 )?;
2245 return Ok(offset);
2246 } else if matches!(
2247 data_type,
2248 DataType::ListView(_) | DataType::LargeListView(_)
2249 ) {
2250 assert_eq!(array_data.buffers().len(), 2); assert_eq!(array_data.child_data().len(), 1);
2252
2253 let (offsets, sizes, child_data) = match data_type {
2254 DataType::ListView(_) => get_list_view_array_buffers::<i32>(array_data),
2255 DataType::LargeListView(_) => get_list_view_array_buffers::<i64>(array_data),
2256 _ => unreachable!(),
2257 };
2258
2259 offset = encode_sink_buffer(
2260 offsets,
2261 meta,
2262 sink,
2263 offset,
2264 compression_codec,
2265 ipc_write_context,
2266 write_options.alignment,
2267 )?;
2268 offset = encode_sink_buffer(
2269 sizes,
2270 meta,
2271 sink,
2272 offset,
2273 compression_codec,
2274 ipc_write_context,
2275 write_options.alignment,
2276 )?;
2277
2278 offset = write_array_data(
2279 &child_data,
2280 meta,
2281 sink,
2282 offset,
2283 compression_codec,
2284 ipc_write_context,
2285 write_options,
2286 )?;
2287 return Ok(offset);
2288 } else if let DataType::FixedSizeList(_, fixed_size) = data_type {
2289 assert_eq!(array_data.child_data().len(), 1);
2290 let fixed_size = *fixed_size as usize;
2291
2292 let child_offset = array_data.offset() * fixed_size;
2293 let child_length = array_data.len() * fixed_size;
2294 let child_data = array_data.child_data()[0].slice(child_offset, child_length);
2295
2296 offset = write_array_data(
2297 &child_data,
2298 meta,
2299 sink,
2300 offset,
2301 compression_codec,
2302 ipc_write_context,
2303 write_options,
2304 )?;
2305 return Ok(offset);
2306 } else {
2307 for buffer in array_data.buffers() {
2308 offset = encode_sink_buffer(
2309 buffer.clone(),
2310 meta,
2311 sink,
2312 offset,
2313 compression_codec,
2314 ipc_write_context,
2315 write_options.alignment,
2316 )?;
2317 }
2318 }
2319
2320 match array_data.data_type() {
2321 DataType::Dictionary(_, _) => {}
2322 DataType::RunEndEncoded(_, _) => {
2323 let arr = unslice_run_array(array_data.clone())?;
2325 for data_ref in arr.child_data() {
2327 offset = write_array_data(
2329 data_ref,
2330 meta,
2331 sink,
2332 offset,
2333 compression_codec,
2334 ipc_write_context,
2335 write_options,
2336 )?;
2337 }
2338 }
2339 _ => {
2340 for data_ref in array_data.child_data() {
2342 offset = write_array_data(
2344 data_ref,
2345 meta,
2346 sink,
2347 offset,
2348 compression_codec,
2349 ipc_write_context,
2350 write_options,
2351 )?;
2352 }
2353 }
2354 }
2355 Ok(offset)
2356}
2357
2358fn encode_sink_buffer(
2372 buffer: Buffer,
2373 ipc_meta_data: &mut IpcMetadataBuilder,
2374 sink: &mut IpcBodySink<'_>,
2375 offset: i64,
2376 compression_codec: Option<CompressionCodec>,
2377 ipc_write_context: &mut IpcWriteContext,
2378 alignment: u8,
2379) -> Result<i64, ArrowError> {
2380 let (encoded, len) = match compression_codec {
2381 None => {
2382 let len = buffer.len() as i64;
2383 (EncodedBuffer::Raw(buffer), len)
2384 }
2385 Some(codec) => {
2386 let mut scratch = Vec::new();
2387 let written =
2388 codec.compress_to_vec(buffer.as_slice(), &mut scratch, ipc_write_context)?;
2389 let len = i64::try_from(written)
2390 .map_err(|e| ArrowError::InvalidArgumentError(format!("{e}")))?;
2391 (EncodedBuffer::Compressed(scratch), len)
2392 }
2393 };
2394
2395 let pad_len = pad_to_alignment(alignment, len as usize);
2396 sink.write(pad_len, encoded);
2397 ipc_meta_data.buffers.push(crate::Buffer::new(offset, len));
2398 Ok(offset + len + pad_len as i64)
2399}
2400
2401const PADDING: [u8; 64] = [0; 64];
2402
2403#[inline]
2409fn estimate_encoded_buffer_count(dt: &DataType) -> usize {
2410 match dt {
2411 DataType::Null => 0,
2412
2413 DataType::Binary | DataType::Utf8 | DataType::LargeBinary | DataType::LargeUtf8 => 3,
2414
2415 DataType::BinaryView | DataType::Utf8View => 3,
2416
2417 DataType::List(f) | DataType::LargeList(f) | DataType::Map(f, _) => {
2418 2 + estimate_encoded_buffer_count(f.data_type())
2419 }
2420
2421 DataType::ListView(f) | DataType::LargeListView(f) => {
2422 3 + estimate_encoded_buffer_count(f.data_type())
2423 }
2424
2425 DataType::FixedSizeList(f, _) => 1 + estimate_encoded_buffer_count(f.data_type()),
2426
2427 DataType::Struct(fields) => {
2428 1 + fields
2429 .iter()
2430 .map(|f| estimate_encoded_buffer_count(f.data_type()))
2431 .sum::<usize>()
2432 }
2433
2434 DataType::Dictionary(_, _) => 2,
2436
2437 DataType::Union(fields, UnionMode::Sparse) => {
2438 1 + fields
2439 .iter()
2440 .map(|(_, f)| estimate_encoded_buffer_count(f.data_type()))
2441 .sum::<usize>()
2442 }
2443 DataType::Union(fields, UnionMode::Dense) => {
2444 2 + fields
2445 .iter()
2446 .map(|(_, f)| estimate_encoded_buffer_count(f.data_type()))
2447 .sum::<usize>()
2448 }
2449
2450 DataType::RunEndEncoded(run_ends, values) => {
2451 estimate_encoded_buffer_count(run_ends.data_type())
2452 + estimate_encoded_buffer_count(values.data_type())
2453 }
2454 _ => 2,
2456 }
2457}
2458
2459#[inline]
2461fn pad_to_alignment(alignment: u8, len: usize) -> usize {
2462 let a = usize::from(alignment - 1);
2463 ((len + a) & !a) - len
2464}
2465
2466#[cfg(test)]
2467mod tests {
2468 use std::hash::Hasher;
2469 use std::io::Cursor;
2470 use std::io::Seek;
2471
2472 use arrow_array::builder::FixedSizeListBuilder;
2473 use arrow_array::builder::Float32Builder;
2474 use arrow_array::builder::Int64Builder;
2475 use arrow_array::builder::MapBuilder;
2476 use arrow_array::builder::StringViewBuilder;
2477 use arrow_array::builder::UnionBuilder;
2478 use arrow_array::builder::{
2479 GenericListBuilder, GenericListViewBuilder, ListBuilder, StringBuilder,
2480 };
2481 use arrow_array::builder::{PrimitiveRunBuilder, UInt32Builder};
2482 use arrow_array::types::*;
2483 use arrow_buffer::ScalarBuffer;
2484
2485 use crate::MetadataVersion;
2486 use crate::convert::fb_to_schema;
2487 use crate::reader::*;
2488 use crate::root_as_footer;
2489
2490 use super::*;
2491
2492 fn serialize_file(rb: &RecordBatch) -> Vec<u8> {
2493 let mut writer = FileWriter::try_new(vec![], rb.schema_ref()).unwrap();
2494 writer.write(rb).unwrap();
2495 writer.finish().unwrap();
2496 writer.into_inner().unwrap()
2497 }
2498
2499 fn deserialize_file(bytes: Vec<u8>) -> RecordBatch {
2500 let mut reader = FileReader::try_new(Cursor::new(bytes), None).unwrap();
2501 reader.next().unwrap().unwrap()
2502 }
2503
2504 fn serialize_stream(record: &RecordBatch) -> Vec<u8> {
2505 const IPC_ALIGNMENT: usize = 8;
2509
2510 let mut stream_writer = StreamWriter::try_new_with_options(
2511 vec![],
2512 record.schema_ref(),
2513 IpcWriteOptions::try_new(IPC_ALIGNMENT, false, MetadataVersion::V5).unwrap(),
2514 )
2515 .unwrap();
2516 stream_writer.write(record).unwrap();
2517 stream_writer.finish().unwrap();
2518 stream_writer.into_inner().unwrap()
2519 }
2520
2521 fn deserialize_stream(bytes: Vec<u8>) -> RecordBatch {
2522 let mut stream_reader = StreamReader::try_new(Cursor::new(bytes), None).unwrap();
2523 stream_reader.next().unwrap().unwrap()
2524 }
2525
2526 #[test]
2527 #[cfg(feature = "lz4")]
2528 fn test_write_empty_record_batch_lz4_compression() {
2529 let schema = Schema::new(vec![Field::new("field1", DataType::Int32, true)]);
2530 let values: Vec<Option<i32>> = vec![];
2531 let array = Int32Array::from(values);
2532 let record_batch =
2533 RecordBatch::try_new(Arc::new(schema.clone()), vec![Arc::new(array)]).unwrap();
2534
2535 let mut file = tempfile::tempfile().unwrap();
2536
2537 {
2538 let write_option = IpcWriteOptions::try_new(8, false, crate::MetadataVersion::V5)
2539 .unwrap()
2540 .try_with_compression(Some(crate::CompressionType::LZ4_FRAME))
2541 .unwrap();
2542
2543 let mut writer =
2544 FileWriter::try_new_with_options(&mut file, &schema, write_option).unwrap();
2545 writer.write(&record_batch).unwrap();
2546 writer.finish().unwrap();
2547 }
2548 file.rewind().unwrap();
2549 {
2550 let reader = FileReader::try_new(file, None).unwrap();
2552 for read_batch in reader {
2553 read_batch
2554 .unwrap()
2555 .columns()
2556 .iter()
2557 .zip(record_batch.columns())
2558 .for_each(|(a, b)| {
2559 assert_eq!(a.data_type(), b.data_type());
2560 assert_eq!(a.len(), b.len());
2561 assert_eq!(a.null_count(), b.null_count());
2562 });
2563 }
2564 }
2565 }
2566
2567 #[test]
2568 #[cfg(feature = "lz4")]
2569 fn test_write_file_with_lz4_compression() {
2570 let schema = Schema::new(vec![Field::new("field1", DataType::Int32, true)]);
2571 let values: Vec<Option<i32>> = vec![Some(12), Some(1)];
2572 let array = Int32Array::from(values);
2573 let record_batch =
2574 RecordBatch::try_new(Arc::new(schema.clone()), vec![Arc::new(array)]).unwrap();
2575
2576 let mut file = tempfile::tempfile().unwrap();
2577 {
2578 let write_option = IpcWriteOptions::try_new(8, false, crate::MetadataVersion::V5)
2579 .unwrap()
2580 .try_with_compression(Some(crate::CompressionType::LZ4_FRAME))
2581 .unwrap();
2582
2583 let mut writer =
2584 FileWriter::try_new_with_options(&mut file, &schema, write_option).unwrap();
2585 writer.write(&record_batch).unwrap();
2586 writer.finish().unwrap();
2587 }
2588 file.rewind().unwrap();
2589 {
2590 let reader = FileReader::try_new(file, None).unwrap();
2592 for read_batch in reader {
2593 read_batch
2594 .unwrap()
2595 .columns()
2596 .iter()
2597 .zip(record_batch.columns())
2598 .for_each(|(a, b)| {
2599 assert_eq!(a.data_type(), b.data_type());
2600 assert_eq!(a.len(), b.len());
2601 assert_eq!(a.null_count(), b.null_count());
2602 });
2603 }
2604 }
2605 }
2606
2607 #[test]
2608 #[cfg(feature = "zstd")]
2609 fn test_write_file_with_zstd_compression() {
2610 let schema = Schema::new(vec![Field::new("field1", DataType::Int32, true)]);
2611 let values: Vec<Option<i32>> = vec![Some(12), Some(1)];
2612 let array = Int32Array::from(values);
2613 let record_batch =
2614 RecordBatch::try_new(Arc::new(schema.clone()), vec![Arc::new(array)]).unwrap();
2615 let mut file = tempfile::tempfile().unwrap();
2616 {
2617 let write_option = IpcWriteOptions::try_new(8, false, crate::MetadataVersion::V5)
2618 .unwrap()
2619 .try_with_compression(Some(crate::CompressionType::ZSTD))
2620 .unwrap()
2621 .try_with_compression_level(Some(1))
2622 .unwrap();
2623
2624 let mut writer =
2625 FileWriter::try_new_with_options(&mut file, &schema, write_option).unwrap();
2626 writer.write(&record_batch).unwrap();
2627 writer.finish().unwrap();
2628 }
2629 file.rewind().unwrap();
2630 {
2631 let reader = FileReader::try_new(file, None).unwrap();
2633 for read_batch in reader {
2634 read_batch
2635 .unwrap()
2636 .columns()
2637 .iter()
2638 .zip(record_batch.columns())
2639 .for_each(|(a, b)| {
2640 assert_eq!(a.data_type(), b.data_type());
2641 assert_eq!(a.len(), b.len());
2642 assert_eq!(a.null_count(), b.null_count());
2643 });
2644 }
2645 }
2646 }
2647
2648 #[test]
2649 fn test_write_file() {
2650 let schema = Schema::new(vec![Field::new("field1", DataType::UInt32, true)]);
2651 let values: Vec<Option<u32>> = vec![
2652 Some(999),
2653 None,
2654 Some(235),
2655 Some(123),
2656 None,
2657 None,
2658 None,
2659 None,
2660 None,
2661 ];
2662 let array1 = UInt32Array::from(values);
2663 let batch =
2664 RecordBatch::try_new(Arc::new(schema.clone()), vec![Arc::new(array1) as ArrayRef])
2665 .unwrap();
2666 let mut file = tempfile::tempfile().unwrap();
2667 {
2668 let mut writer = FileWriter::try_new(&mut file, &schema).unwrap();
2669
2670 writer.write(&batch).unwrap();
2671 writer.finish().unwrap();
2672 }
2673 file.rewind().unwrap();
2674
2675 {
2676 let mut reader = FileReader::try_new(file, None).unwrap();
2677 while let Some(Ok(read_batch)) = reader.next() {
2678 read_batch
2679 .columns()
2680 .iter()
2681 .zip(batch.columns())
2682 .for_each(|(a, b)| {
2683 assert_eq!(a.data_type(), b.data_type());
2684 assert_eq!(a.len(), b.len());
2685 assert_eq!(a.null_count(), b.null_count());
2686 });
2687 }
2688 }
2689 }
2690
2691 #[test]
2692 fn test_empty_utf8_ipc_writes_nonempty_offsets_buffer() {
2693 let name = StringArray::from(Vec::<String>::new());
2694 let (offsets, values) = get_byte_array_buffers::<i32>(&name.to_data());
2695
2696 assert_eq!(name.len(), 0);
2697 assert_eq!(
2698 offsets.len(),
2699 std::mem::size_of::<i32>(),
2700 "offsets buffer should contain one zero i32 offset"
2701 );
2702 assert_eq!(values.len(), 0, "values buffer should remain empty");
2703 }
2704
2705 #[test]
2706 fn test_empty_large_utf8_ipc_writes_nonempty_offsets_buffer() {
2707 let name = LargeStringArray::from(Vec::<String>::new());
2708 let (offsets, values) = get_byte_array_buffers::<i64>(&name.to_data());
2709
2710 assert_eq!(name.len(), 0);
2711 assert_eq!(
2712 offsets.len(),
2713 std::mem::size_of::<i64>(),
2714 "offsets buffer should contain one zero i64 offset"
2715 );
2716 assert_eq!(values.len(), 0, "values buffer should remain empty");
2717 }
2718
2719 #[test]
2720 fn test_empty_list_ipc_writes_nonempty_offsets_buffer() {
2721 let list = GenericListBuilder::<i32, _>::new(UInt32Builder::new()).finish();
2722 let (offsets, child_data) = get_list_array_buffers::<i32>(&list.to_data());
2723
2724 assert_eq!(list.len(), 0);
2725 assert_eq!(
2726 offsets.len(),
2727 std::mem::size_of::<i32>(),
2728 "offsets buffer should contain one zero i32 offset"
2729 );
2730 assert_eq!(child_data.len(), 0, "child data should remain empty");
2731 }
2732
2733 #[test]
2734 fn test_empty_large_list_ipc_writes_nonempty_offsets_buffer() {
2735 let list = GenericListBuilder::<i64, _>::new(UInt32Builder::new()).finish();
2736 let (offsets, child_data) = get_list_array_buffers::<i64>(&list.to_data());
2737
2738 assert_eq!(list.len(), 0);
2739 assert_eq!(
2740 offsets.len(),
2741 std::mem::size_of::<i64>(),
2742 "offsets buffer should contain one zero i64 offset"
2743 );
2744 assert_eq!(child_data.len(), 0, "child data should remain empty");
2745 }
2746
2747 fn write_null_file(options: IpcWriteOptions) {
2748 let schema = Schema::new(vec![
2749 Field::new("nulls", DataType::Null, true),
2750 Field::new("int32s", DataType::Int32, false),
2751 Field::new("nulls2", DataType::Null, true),
2752 Field::new("f64s", DataType::Float64, false),
2753 ]);
2754 let array1 = NullArray::new(32);
2755 let array2 = Int32Array::from(vec![1; 32]);
2756 let array3 = NullArray::new(32);
2757 let array4 = Float64Array::from(vec![f64::NAN; 32]);
2758 let batch = RecordBatch::try_new(
2759 Arc::new(schema.clone()),
2760 vec![
2761 Arc::new(array1) as ArrayRef,
2762 Arc::new(array2) as ArrayRef,
2763 Arc::new(array3) as ArrayRef,
2764 Arc::new(array4) as ArrayRef,
2765 ],
2766 )
2767 .unwrap();
2768 let mut file = tempfile::tempfile().unwrap();
2769 {
2770 let mut writer = FileWriter::try_new_with_options(&mut file, &schema, options).unwrap();
2771
2772 writer.write(&batch).unwrap();
2773 writer.finish().unwrap();
2774 }
2775
2776 file.rewind().unwrap();
2777
2778 {
2779 let reader = FileReader::try_new(file, None).unwrap();
2780 reader.for_each(|maybe_batch| {
2781 maybe_batch
2782 .unwrap()
2783 .columns()
2784 .iter()
2785 .zip(batch.columns())
2786 .for_each(|(a, b)| {
2787 assert_eq!(a.data_type(), b.data_type());
2788 assert_eq!(a.len(), b.len());
2789 assert_eq!(a.null_count(), b.null_count());
2790 });
2791 });
2792 }
2793 }
2794 #[test]
2795 fn test_write_null_file_v4() {
2796 write_null_file(IpcWriteOptions::try_new(8, false, MetadataVersion::V4).unwrap());
2797 write_null_file(IpcWriteOptions::try_new(8, true, MetadataVersion::V4).unwrap());
2798 write_null_file(IpcWriteOptions::try_new(64, false, MetadataVersion::V4).unwrap());
2799 write_null_file(IpcWriteOptions::try_new(64, true, MetadataVersion::V4).unwrap());
2800 }
2801
2802 #[test]
2803 fn test_write_null_file_v5() {
2804 write_null_file(IpcWriteOptions::try_new(8, false, MetadataVersion::V5).unwrap());
2805 write_null_file(IpcWriteOptions::try_new(64, false, MetadataVersion::V5).unwrap());
2806 }
2807
2808 #[test]
2809 fn track_union_nested_dict() {
2810 let inner: DictionaryArray<Int32Type> = vec!["a", "b", "a"].into_iter().collect();
2811
2812 let array = Arc::new(inner) as ArrayRef;
2813
2814 #[allow(deprecated)]
2816 let dctfield = Field::new_dict("dict", array.data_type().clone(), false, 0, false);
2817 let union_fields = [(0, Arc::new(dctfield))].into_iter().collect();
2818
2819 let types = [0, 0, 0].into_iter().collect::<ScalarBuffer<i8>>();
2820 let offsets = [0, 1, 2].into_iter().collect::<ScalarBuffer<i32>>();
2821
2822 let union = UnionArray::try_new(union_fields, types, Some(offsets), vec![array]).unwrap();
2823
2824 let schema = Arc::new(Schema::new(vec![Field::new(
2825 "union",
2826 union.data_type().clone(),
2827 false,
2828 )]));
2829
2830 let r#gen = IpcDataGenerator::default();
2831 let mut dict_tracker = DictionaryTracker::new(false);
2832 r#gen.schema_to_bytes_with_dictionary_tracker(
2833 &schema,
2834 &mut dict_tracker,
2835 &IpcWriteOptions::default(),
2836 );
2837
2838 let batch = RecordBatch::try_new(schema, vec![Arc::new(union)]).unwrap();
2839
2840 r#gen
2841 .encode(
2842 &batch,
2843 &mut dict_tracker,
2844 &Default::default(),
2845 &mut Default::default(),
2846 )
2847 .unwrap();
2848
2849 assert!(dict_tracker.written.contains_key(&0));
2852 }
2853
2854 #[test]
2855 fn track_struct_nested_dict() {
2856 let inner: DictionaryArray<Int32Type> = vec!["a", "b", "a"].into_iter().collect();
2857
2858 let array = Arc::new(inner) as ArrayRef;
2859
2860 #[allow(deprecated)]
2862 let dctfield = Arc::new(Field::new_dict(
2863 "dict",
2864 array.data_type().clone(),
2865 false,
2866 2,
2867 false,
2868 ));
2869
2870 let s = StructArray::from(vec![(dctfield, array)]);
2871 let struct_array = Arc::new(s) as ArrayRef;
2872
2873 let schema = Arc::new(Schema::new(vec![Field::new(
2874 "struct",
2875 struct_array.data_type().clone(),
2876 false,
2877 )]));
2878
2879 let r#gen = IpcDataGenerator::default();
2880 let mut dict_tracker = DictionaryTracker::new(false);
2881 r#gen.schema_to_bytes_with_dictionary_tracker(
2882 &schema,
2883 &mut dict_tracker,
2884 &IpcWriteOptions::default(),
2885 );
2886
2887 let batch = RecordBatch::try_new(schema, vec![struct_array]).unwrap();
2888
2889 r#gen
2890 .encode(
2891 &batch,
2892 &mut dict_tracker,
2893 &Default::default(),
2894 &mut Default::default(),
2895 )
2896 .unwrap();
2897
2898 assert!(dict_tracker.written.contains_key(&0));
2899 }
2900
2901 fn write_union_file(options: IpcWriteOptions) {
2902 let schema = Schema::new(vec![Field::new_union(
2903 "union",
2904 vec![0, 1],
2905 vec![
2906 Field::new("a", DataType::Int32, false),
2907 Field::new("c", DataType::Float64, false),
2908 ],
2909 UnionMode::Sparse,
2910 )]);
2911 let mut builder = UnionBuilder::with_capacity_sparse(5);
2912 builder.append::<Int32Type>("a", 1).unwrap();
2913 builder.append_null::<Int32Type>("a").unwrap();
2914 builder.append::<Float64Type>("c", 3.0).unwrap();
2915 builder.append_null::<Float64Type>("c").unwrap();
2916 builder.append::<Int32Type>("a", 4).unwrap();
2917 let union = builder.build().unwrap();
2918
2919 let batch =
2920 RecordBatch::try_new(Arc::new(schema.clone()), vec![Arc::new(union) as ArrayRef])
2921 .unwrap();
2922
2923 let mut file = tempfile::tempfile().unwrap();
2924 {
2925 let mut writer = FileWriter::try_new_with_options(&mut file, &schema, options).unwrap();
2926
2927 writer.write(&batch).unwrap();
2928 writer.finish().unwrap();
2929 }
2930 file.rewind().unwrap();
2931
2932 {
2933 let reader = FileReader::try_new(file, None).unwrap();
2934 reader.for_each(|maybe_batch| {
2935 maybe_batch
2936 .unwrap()
2937 .columns()
2938 .iter()
2939 .zip(batch.columns())
2940 .for_each(|(a, b)| {
2941 assert_eq!(a.data_type(), b.data_type());
2942 assert_eq!(a.len(), b.len());
2943 assert_eq!(a.null_count(), b.null_count());
2944 });
2945 });
2946 }
2947 }
2948
2949 #[test]
2950 fn test_write_union_file_v4_v5() {
2951 write_union_file(IpcWriteOptions::try_new(8, false, MetadataVersion::V4).unwrap());
2952 write_union_file(IpcWriteOptions::try_new(8, false, MetadataVersion::V5).unwrap());
2953 }
2954
2955 #[test]
2956 fn test_write_view_types() {
2957 const LONG_TEST_STRING: &str =
2958 "This is a long string to make sure binary view array handles it";
2959 let schema = Schema::new(vec![
2960 Field::new("field1", DataType::BinaryView, true),
2961 Field::new("field2", DataType::Utf8View, true),
2962 ]);
2963 let values: Vec<Option<&[u8]>> = vec![
2964 Some(b"foo"),
2965 Some(b"bar"),
2966 Some(LONG_TEST_STRING.as_bytes()),
2967 ];
2968 let binary_array = BinaryViewArray::from_iter(values);
2969 let utf8_array =
2970 StringViewArray::from_iter(vec![Some("foo"), Some("bar"), Some(LONG_TEST_STRING)]);
2971 let record_batch = RecordBatch::try_new(
2972 Arc::new(schema.clone()),
2973 vec![Arc::new(binary_array), Arc::new(utf8_array)],
2974 )
2975 .unwrap();
2976
2977 let mut file = tempfile::tempfile().unwrap();
2978 {
2979 let mut writer = FileWriter::try_new(&mut file, &schema).unwrap();
2980 writer.write(&record_batch).unwrap();
2981 writer.finish().unwrap();
2982 }
2983 file.rewind().unwrap();
2984 {
2985 let mut reader = FileReader::try_new(&file, None).unwrap();
2986 let read_batch = reader.next().unwrap().unwrap();
2987 read_batch
2988 .columns()
2989 .iter()
2990 .zip(record_batch.columns())
2991 .for_each(|(a, b)| {
2992 assert_eq!(a, b);
2993 });
2994 }
2995 file.rewind().unwrap();
2996 {
2997 let mut reader = FileReader::try_new(&file, Some(vec![0])).unwrap();
2998 let read_batch = reader.next().unwrap().unwrap();
2999 assert_eq!(read_batch.num_columns(), 1);
3000 let read_array = read_batch.column(0);
3001 let write_array = record_batch.column(0);
3002 assert_eq!(read_array, write_array);
3003 }
3004 }
3005
3006 #[test]
3007 fn truncate_ipc_record_batch() {
3008 fn create_batch(rows: usize) -> RecordBatch {
3009 let schema = Schema::new(vec![
3010 Field::new("a", DataType::Int32, false),
3011 Field::new("b", DataType::Utf8, false),
3012 ]);
3013
3014 let a = Int32Array::from_iter_values(0..rows as i32);
3015 let b = StringArray::from_iter_values((0..rows).map(|i| i.to_string()));
3016
3017 RecordBatch::try_new(Arc::new(schema), vec![Arc::new(a), Arc::new(b)]).unwrap()
3018 }
3019
3020 let big_record_batch = create_batch(65536);
3021
3022 let length = 5;
3023 let small_record_batch = create_batch(length);
3024
3025 let offset = 2;
3026 let record_batch_slice = big_record_batch.slice(offset, length);
3027 assert!(
3028 serialize_stream(&big_record_batch).len() > serialize_stream(&small_record_batch).len()
3029 );
3030 assert_eq!(
3031 serialize_stream(&small_record_batch).len(),
3032 serialize_stream(&record_batch_slice).len()
3033 );
3034
3035 assert_eq!(
3036 deserialize_stream(serialize_stream(&record_batch_slice)),
3037 record_batch_slice
3038 );
3039 }
3040
3041 #[test]
3042 fn truncate_ipc_record_batch_with_nulls() {
3043 fn create_batch() -> RecordBatch {
3044 let schema = Schema::new(vec![
3045 Field::new("a", DataType::Int32, true),
3046 Field::new("b", DataType::Utf8, true),
3047 ]);
3048
3049 let a = Int32Array::from(vec![Some(1), None, Some(1), None, Some(1)]);
3050 let b = StringArray::from(vec![None, Some("a"), Some("a"), None, Some("a")]);
3051
3052 RecordBatch::try_new(Arc::new(schema), vec![Arc::new(a), Arc::new(b)]).unwrap()
3053 }
3054
3055 let record_batch = create_batch();
3056 let record_batch_slice = record_batch.slice(1, 2);
3057 let deserialized_batch = deserialize_stream(serialize_stream(&record_batch_slice));
3058
3059 assert!(
3060 serialize_stream(&record_batch).len() > serialize_stream(&record_batch_slice).len()
3061 );
3062
3063 assert!(deserialized_batch.column(0).is_null(0));
3064 assert!(deserialized_batch.column(0).is_valid(1));
3065 assert!(deserialized_batch.column(1).is_valid(0));
3066 assert!(deserialized_batch.column(1).is_valid(1));
3067
3068 assert_eq!(record_batch_slice, deserialized_batch);
3069 }
3070
3071 #[test]
3072 fn truncate_ipc_dictionary_array() {
3073 fn create_batch() -> RecordBatch {
3074 let values: StringArray = [Some("foo"), Some("bar"), Some("baz")]
3075 .into_iter()
3076 .collect();
3077 let keys: Int32Array = [Some(0), Some(2), None, Some(1)].into_iter().collect();
3078
3079 let array = DictionaryArray::new(keys, Arc::new(values));
3080
3081 let schema = Schema::new(vec![Field::new("dict", array.data_type().clone(), true)]);
3082
3083 RecordBatch::try_new(Arc::new(schema), vec![Arc::new(array)]).unwrap()
3084 }
3085
3086 let record_batch = create_batch();
3087 let record_batch_slice = record_batch.slice(1, 2);
3088 let deserialized_batch = deserialize_stream(serialize_stream(&record_batch_slice));
3089
3090 assert!(
3091 serialize_stream(&record_batch).len() > serialize_stream(&record_batch_slice).len()
3092 );
3093
3094 assert!(deserialized_batch.column(0).is_valid(0));
3095 assert!(deserialized_batch.column(0).is_null(1));
3096
3097 assert_eq!(record_batch_slice, deserialized_batch);
3098 }
3099
3100 #[test]
3101 fn truncate_ipc_struct_array() {
3102 fn create_batch() -> RecordBatch {
3103 let strings: StringArray = [Some("foo"), None, Some("bar"), Some("baz")]
3104 .into_iter()
3105 .collect();
3106 let ints: Int32Array = [Some(0), Some(2), None, Some(1)].into_iter().collect();
3107
3108 let struct_array = StructArray::from(vec![
3109 (
3110 Arc::new(Field::new("s", DataType::Utf8, true)),
3111 Arc::new(strings) as ArrayRef,
3112 ),
3113 (
3114 Arc::new(Field::new("c", DataType::Int32, true)),
3115 Arc::new(ints) as ArrayRef,
3116 ),
3117 ]);
3118
3119 let schema = Schema::new(vec![Field::new(
3120 "struct_array",
3121 struct_array.data_type().clone(),
3122 true,
3123 )]);
3124
3125 RecordBatch::try_new(Arc::new(schema), vec![Arc::new(struct_array)]).unwrap()
3126 }
3127
3128 let record_batch = create_batch();
3129 let record_batch_slice = record_batch.slice(1, 2);
3130 let deserialized_batch = deserialize_stream(serialize_stream(&record_batch_slice));
3131
3132 assert!(
3133 serialize_stream(&record_batch).len() > serialize_stream(&record_batch_slice).len()
3134 );
3135
3136 let structs = deserialized_batch
3137 .column(0)
3138 .as_any()
3139 .downcast_ref::<StructArray>()
3140 .unwrap();
3141
3142 assert!(structs.column(0).is_null(0));
3143 assert!(structs.column(0).is_valid(1));
3144 assert!(structs.column(1).is_valid(0));
3145 assert!(structs.column(1).is_null(1));
3146 assert_eq!(record_batch_slice, deserialized_batch);
3147 }
3148
3149 #[test]
3150 fn truncate_ipc_string_array_with_all_empty_string() {
3151 fn create_batch() -> RecordBatch {
3152 let schema = Schema::new(vec![Field::new("a", DataType::Utf8, true)]);
3153 let a = StringArray::from(vec![Some(""), Some(""), Some(""), Some(""), Some("")]);
3154 RecordBatch::try_new(Arc::new(schema), vec![Arc::new(a)]).unwrap()
3155 }
3156
3157 let record_batch = create_batch();
3158 let record_batch_slice = record_batch.slice(0, 1);
3159 let deserialized_batch = deserialize_stream(serialize_stream(&record_batch_slice));
3160
3161 assert!(
3162 serialize_stream(&record_batch).len() > serialize_stream(&record_batch_slice).len()
3163 );
3164 assert_eq!(record_batch_slice, deserialized_batch);
3165 }
3166
3167 #[test]
3168 fn test_stream_writer_writes_array_slice() {
3169 let array = UInt32Array::from(vec![Some(1), Some(2), Some(3)]);
3170 assert_eq!(
3171 vec![Some(1), Some(2), Some(3)],
3172 array.iter().collect::<Vec<_>>()
3173 );
3174
3175 let sliced = array.slice(1, 2);
3176 assert_eq!(vec![Some(2), Some(3)], sliced.iter().collect::<Vec<_>>());
3177
3178 let batch = RecordBatch::try_new(
3179 Arc::new(Schema::new(vec![Field::new("a", DataType::UInt32, true)])),
3180 vec![Arc::new(sliced)],
3181 )
3182 .expect("new batch");
3183
3184 let mut writer = StreamWriter::try_new(vec![], batch.schema_ref()).expect("new writer");
3185 writer.write(&batch).expect("write");
3186 let outbuf = writer.into_inner().expect("inner");
3187
3188 let mut reader = StreamReader::try_new(&outbuf[..], None).expect("new reader");
3189 let read_batch = reader.next().unwrap().expect("read batch");
3190
3191 let read_array: &UInt32Array = read_batch.column(0).as_primitive();
3192 assert_eq!(
3193 vec![Some(2), Some(3)],
3194 read_array.iter().collect::<Vec<_>>()
3195 );
3196 }
3197
3198 #[test]
3199 fn test_large_slice_uint32() {
3200 ensure_roundtrip(Arc::new(UInt32Array::from_iter(
3201 (0..8000).map(|i| if i % 2 == 0 { Some(i) } else { None }),
3202 )));
3203 }
3204
3205 #[test]
3206 fn test_large_slice_string() {
3207 let strings: Vec<_> = (0..8000)
3208 .map(|i| {
3209 if i % 2 == 0 {
3210 Some(format!("value{i}"))
3211 } else {
3212 None
3213 }
3214 })
3215 .collect();
3216
3217 ensure_roundtrip(Arc::new(StringArray::from(strings)));
3218 }
3219
3220 #[test]
3221 fn test_large_slice_string_list() {
3222 let mut ls = ListBuilder::new(StringBuilder::new());
3223
3224 let mut s = String::new();
3225 for row_number in 0..8000 {
3226 if row_number % 2 == 0 {
3227 for list_element in 0..1000 {
3228 s.clear();
3229 use std::fmt::Write;
3230 write!(&mut s, "value{row_number}-{list_element}").unwrap();
3231 ls.values().append_value(&s);
3232 }
3233 ls.append(true)
3234 } else {
3235 ls.append(false); }
3237 }
3238
3239 ensure_roundtrip(Arc::new(ls.finish()));
3240 }
3241
3242 #[test]
3243 fn test_large_slice_string_list_of_lists() {
3244 let mut ls = ListBuilder::new(ListBuilder::new(StringBuilder::new()));
3248
3249 for _ in 0..4000 {
3250 ls.values().append(true);
3251 ls.append(true)
3252 }
3253
3254 let mut s = String::new();
3255 for row_number in 0..4000 {
3256 if row_number % 2 == 0 {
3257 for list_element in 0..1000 {
3258 s.clear();
3259 use std::fmt::Write;
3260 write!(&mut s, "value{row_number}-{list_element}").unwrap();
3261 ls.values().values().append_value(&s);
3262 }
3263 ls.values().append(true);
3264 ls.append(true)
3265 } else {
3266 ls.append(false); }
3268 }
3269
3270 ensure_roundtrip(Arc::new(ls.finish()));
3271 }
3272
3273 fn ensure_roundtrip(array: ArrayRef) {
3275 let num_rows = array.len();
3276 let orig_batch = RecordBatch::try_from_iter(vec![("a", array)]).unwrap();
3277 let sliced_batch = orig_batch.slice(1, num_rows - 1);
3279
3280 let schema = orig_batch.schema();
3281 let stream_data = {
3282 let mut writer = StreamWriter::try_new(vec![], &schema).unwrap();
3283 writer.write(&sliced_batch).unwrap();
3284 writer.into_inner().unwrap()
3285 };
3286 let read_batch = {
3287 let projection = None;
3288 let mut reader = StreamReader::try_new(Cursor::new(stream_data), projection).unwrap();
3289 reader
3290 .next()
3291 .expect("expect no errors reading batch")
3292 .expect("expect batch")
3293 };
3294 assert_eq!(sliced_batch, read_batch);
3295
3296 let file_data = {
3297 let mut writer = FileWriter::try_new_buffered(vec![], &schema).unwrap();
3298 writer.write(&sliced_batch).unwrap();
3299 writer.into_inner().unwrap().into_inner().unwrap()
3300 };
3301 let read_batch = {
3302 let projection = None;
3303 let mut reader = FileReader::try_new(Cursor::new(file_data), projection).unwrap();
3304 reader
3305 .next()
3306 .expect("expect no errors reading batch")
3307 .expect("expect batch")
3308 };
3309 assert_eq!(sliced_batch, read_batch);
3310
3311 }
3313
3314 #[test]
3315 fn encode_bools_slice() {
3316 assert_bool_roundtrip([true, false], 1, 1);
3318
3319 assert_bool_roundtrip(
3321 [
3322 true, false, true, true, false, false, true, true, true, false, false, false, true,
3323 true, true, true, false, false, false, false, true, true, true, true, true, false,
3324 false, false, false, false,
3325 ],
3326 13,
3327 17,
3328 );
3329
3330 assert_bool_roundtrip(
3332 [
3333 true, false, true, true, false, false, true, true, true, false, false, false,
3334 ],
3335 8,
3336 2,
3337 );
3338
3339 assert_bool_roundtrip(
3341 [
3342 true, false, true, true, false, false, true, true, true, false, false, false, true,
3343 true, true, true, true, false, false, false, false, false,
3344 ],
3345 8,
3346 8,
3347 );
3348 }
3349
3350 fn assert_bool_roundtrip<const N: usize>(bools: [bool; N], offset: usize, length: usize) {
3351 let val_bool_field = Field::new("val", DataType::Boolean, false);
3352
3353 let schema = Arc::new(Schema::new(vec![val_bool_field]));
3354
3355 let bools = BooleanArray::from(bools.to_vec());
3356
3357 let batch = RecordBatch::try_new(Arc::clone(&schema), vec![Arc::new(bools)]).unwrap();
3358 let batch = batch.slice(offset, length);
3359
3360 let data = serialize_stream(&batch);
3361 let batch2 = deserialize_stream(data);
3362 assert_eq!(batch, batch2);
3363 }
3364
3365 #[test]
3366 fn test_run_array_unslice() {
3367 let total_len = 80;
3368 let vals: Vec<Option<i32>> = vec![Some(1), None, Some(2), Some(3), Some(4), None, Some(5)];
3369 let repeats: Vec<usize> = vec![3, 4, 1, 2];
3370 let mut input_array: Vec<Option<i32>> = Vec::with_capacity(total_len);
3371 for ix in 0_usize..32 {
3372 let repeat: usize = repeats[ix % repeats.len()];
3373 let val: Option<i32> = vals[ix % vals.len()];
3374 input_array.resize(input_array.len() + repeat, val);
3375 }
3376
3377 let mut builder =
3379 PrimitiveRunBuilder::<Int16Type, Int32Type>::with_capacity(input_array.len());
3380 builder.extend(input_array.iter().copied());
3381 let run_array = builder.finish();
3382
3383 for slice_len in 1..=total_len {
3385 let sliced_run_array: RunArray<Int16Type> =
3387 run_array.slice(0, slice_len).into_data().into();
3388
3389 let unsliced_run_array = into_zero_offset_run_array(sliced_run_array).unwrap();
3391 let typed = unsliced_run_array
3392 .downcast::<PrimitiveArray<Int32Type>>()
3393 .unwrap();
3394 let expected: Vec<Option<i32>> = input_array.iter().take(slice_len).copied().collect();
3395 let actual: Vec<Option<i32>> = typed.into_iter().collect();
3396 assert_eq!(expected, actual);
3397
3398 let sliced_run_array: RunArray<Int16Type> = run_array
3400 .slice(total_len - slice_len, slice_len)
3401 .into_data()
3402 .into();
3403
3404 let unsliced_run_array = into_zero_offset_run_array(sliced_run_array).unwrap();
3406 let typed = unsliced_run_array
3407 .downcast::<PrimitiveArray<Int32Type>>()
3408 .unwrap();
3409 let expected: Vec<Option<i32>> = input_array
3410 .iter()
3411 .skip(total_len - slice_len)
3412 .copied()
3413 .collect();
3414 let actual: Vec<Option<i32>> = typed.into_iter().collect();
3415 assert_eq!(expected, actual);
3416 }
3417 }
3418
3419 fn generate_list_data<O: OffsetSizeTrait>() -> GenericListArray<O> {
3420 let mut ls = GenericListBuilder::<O, _>::new(UInt32Builder::new());
3421
3422 for i in 0..100_000 {
3423 for value in [i, i, i] {
3424 ls.values().append_value(value);
3425 }
3426 ls.append(true)
3427 }
3428
3429 ls.finish()
3430 }
3431
3432 fn generate_utf8view_list_data<O: OffsetSizeTrait>() -> GenericListArray<O> {
3433 let mut ls = GenericListBuilder::<O, _>::new(StringViewBuilder::new());
3434
3435 for i in 0..100_000 {
3436 for value in [
3437 format!("value{}", i),
3438 format!("value{}", i),
3439 format!("value{}", i),
3440 ] {
3441 ls.values().append_value(&value);
3442 }
3443 ls.append(true)
3444 }
3445
3446 ls.finish()
3447 }
3448
3449 fn generate_string_list_data<O: OffsetSizeTrait>() -> GenericListArray<O> {
3450 let mut ls = GenericListBuilder::<O, _>::new(StringBuilder::new());
3451
3452 for i in 0..100_000 {
3453 for value in [
3454 format!("value{}", i),
3455 format!("value{}", i),
3456 format!("value{}", i),
3457 ] {
3458 ls.values().append_value(&value);
3459 }
3460 ls.append(true)
3461 }
3462
3463 ls.finish()
3464 }
3465
3466 fn generate_nested_list_data<O: OffsetSizeTrait>() -> GenericListArray<O> {
3467 let mut ls =
3468 GenericListBuilder::<O, _>::new(GenericListBuilder::<O, _>::new(UInt32Builder::new()));
3469
3470 for _i in 0..10_000 {
3471 for j in 0..10 {
3472 for value in [j, j, j, j] {
3473 ls.values().values().append_value(value);
3474 }
3475 ls.values().append(true)
3476 }
3477 ls.append(true);
3478 }
3479
3480 ls.finish()
3481 }
3482
3483 fn generate_nested_list_data_starting_at_zero<O: OffsetSizeTrait>() -> GenericListArray<O> {
3484 let mut ls =
3485 GenericListBuilder::<O, _>::new(GenericListBuilder::<O, _>::new(UInt32Builder::new()));
3486
3487 for _i in 0..999 {
3488 ls.values().append(true);
3489 ls.append(true);
3490 }
3491
3492 for j in 0..10 {
3493 for value in [j, j, j, j] {
3494 ls.values().values().append_value(value);
3495 }
3496 ls.values().append(true)
3497 }
3498 ls.append(true);
3499
3500 for i in 0..9_000 {
3501 for j in 0..10 {
3502 for value in [i + j, i + j, i + j, i + j] {
3503 ls.values().values().append_value(value);
3504 }
3505 ls.values().append(true)
3506 }
3507 ls.append(true);
3508 }
3509
3510 ls.finish()
3511 }
3512
3513 fn generate_map_array_data() -> MapArray {
3514 let keys_builder = UInt32Builder::new();
3515 let values_builder = UInt32Builder::new();
3516
3517 let mut builder = MapBuilder::new(None, keys_builder, values_builder);
3518
3519 for i in 0..100_000 {
3520 for _j in 0..3 {
3521 builder.keys().append_value(i);
3522 builder.values().append_value(i * 2);
3523 }
3524 builder.append(true).unwrap();
3525 }
3526
3527 builder.finish()
3528 }
3529
3530 #[test]
3531 fn reencode_offsets_when_first_offset_is_not_zero() {
3532 let original_list = generate_list_data::<i32>();
3533 let original_data = original_list.into_data();
3534 let slice_data = original_data.slice(75, 7);
3535 let (new_offsets, original_start, length) =
3536 reencode_offsets::<i32>(&slice_data.buffers()[0], &slice_data);
3537 assert_eq!(
3538 vec![0, 3, 6, 9, 12, 15, 18, 21],
3539 new_offsets.typed_data::<i32>()
3540 );
3541 assert_eq!(225, original_start);
3542 assert_eq!(21, length);
3543 }
3544
3545 #[test]
3546 fn reencode_offsets_when_first_offset_is_zero() {
3547 let mut ls = GenericListBuilder::<i32, _>::new(UInt32Builder::new());
3548 ls.append(true);
3550 ls.values().append_value(35);
3551 ls.values().append_value(42);
3552 ls.append(true);
3553 let original_list = ls.finish();
3554 let original_data = original_list.into_data();
3555
3556 let slice_data = original_data.slice(1, 1);
3557 let (new_offsets, original_start, length) =
3558 reencode_offsets::<i32>(&slice_data.buffers()[0], &slice_data);
3559 assert_eq!(vec![0, 2], new_offsets.typed_data::<i32>());
3560 assert_eq!(0, original_start);
3561 assert_eq!(2, length);
3562 }
3563
3564 fn roundtrip_ensure_sliced_smaller(in_batch: RecordBatch, expected_size_factor: usize) {
3567 let in_sliced = in_batch.slice(999, 1);
3569
3570 let bytes_batch = serialize_file(&in_batch);
3571 let bytes_sliced = serialize_file(&in_sliced);
3572
3573 assert!(bytes_sliced.len() < (bytes_batch.len() / expected_size_factor));
3575
3576 let out_batch = deserialize_file(bytes_batch);
3578 assert_eq!(in_batch, out_batch);
3579
3580 let out_sliced = deserialize_file(bytes_sliced);
3581 assert_eq!(in_sliced, out_sliced);
3582 }
3583
3584 #[test]
3585 fn encode_lists() {
3586 let val_inner = Field::new_list_field(DataType::UInt32, true);
3587 let val_list_field = Field::new("val", DataType::List(Arc::new(val_inner)), false);
3588 let schema = Arc::new(Schema::new(vec![val_list_field]));
3589
3590 let values = Arc::new(generate_list_data::<i32>());
3591
3592 let in_batch = RecordBatch::try_new(schema, vec![values]).unwrap();
3593 roundtrip_ensure_sliced_smaller(in_batch, 1000);
3594 }
3595
3596 #[test]
3597 fn encode_empty_list() {
3598 let val_inner = Field::new_list_field(DataType::UInt32, true);
3599 let val_list_field = Field::new("val", DataType::List(Arc::new(val_inner)), false);
3600 let schema = Arc::new(Schema::new(vec![val_list_field]));
3601
3602 let values = Arc::new(generate_list_data::<i32>());
3603
3604 let in_batch = RecordBatch::try_new(schema, vec![values])
3605 .unwrap()
3606 .slice(999, 0);
3607 let out_batch = deserialize_file(serialize_file(&in_batch));
3608 assert_eq!(in_batch, out_batch);
3609 }
3610
3611 #[test]
3612 fn encode_large_lists() {
3613 let val_inner = Field::new_list_field(DataType::UInt32, true);
3614 let val_list_field = Field::new("val", DataType::LargeList(Arc::new(val_inner)), false);
3615 let schema = Arc::new(Schema::new(vec![val_list_field]));
3616
3617 let values = Arc::new(generate_list_data::<i64>());
3618
3619 let in_batch = RecordBatch::try_new(schema, vec![values]).unwrap();
3622 roundtrip_ensure_sliced_smaller(in_batch, 1000);
3623 }
3624
3625 #[test]
3626 fn encode_large_lists_non_zero_offset() {
3627 let val_inner = Field::new_list_field(DataType::UInt32, true);
3628 let val_list_field = Field::new("val", DataType::LargeList(Arc::new(val_inner)), false);
3629 let schema = Arc::new(Schema::new(vec![val_list_field]));
3630
3631 let values = Arc::new(generate_list_data::<i64>());
3632
3633 check_sliced_list_array(schema, values);
3634 }
3635
3636 #[test]
3637 fn encode_large_lists_string_non_zero_offset() {
3638 let val_inner = Field::new_list_field(DataType::Utf8, true);
3639 let val_list_field = Field::new("val", DataType::LargeList(Arc::new(val_inner)), false);
3640 let schema = Arc::new(Schema::new(vec![val_list_field]));
3641
3642 let values = Arc::new(generate_string_list_data::<i64>());
3643
3644 check_sliced_list_array(schema, values);
3645 }
3646
3647 #[test]
3648 fn encode_large_list_string_view_non_zero_offset() {
3649 let val_inner = Field::new_list_field(DataType::Utf8View, true);
3650 let val_list_field = Field::new("val", DataType::LargeList(Arc::new(val_inner)), false);
3651 let schema = Arc::new(Schema::new(vec![val_list_field]));
3652
3653 let values = Arc::new(generate_utf8view_list_data::<i64>());
3654
3655 check_sliced_list_array(schema, values);
3656 }
3657
3658 fn check_sliced_list_array(schema: Arc<Schema>, values: Arc<GenericListArray<i64>>) {
3659 for (offset, len) in [(999, 1), (0, 13), (47, 12), (values.len() - 13, 13)] {
3660 let in_batch = RecordBatch::try_new(schema.clone(), vec![values.clone()])
3661 .unwrap()
3662 .slice(offset, len);
3663 let out_batch = deserialize_file(serialize_file(&in_batch));
3664 assert_eq!(in_batch, out_batch);
3665 }
3666 }
3667
3668 #[test]
3669 fn encode_nested_lists() {
3670 let inner_int = Arc::new(Field::new_list_field(DataType::UInt32, true));
3671 let inner_list_field = Arc::new(Field::new_list_field(DataType::List(inner_int), true));
3672 let list_field = Field::new("val", DataType::List(inner_list_field), true);
3673 let schema = Arc::new(Schema::new(vec![list_field]));
3674
3675 let values = Arc::new(generate_nested_list_data::<i32>());
3676
3677 let in_batch = RecordBatch::try_new(schema, vec![values]).unwrap();
3678 roundtrip_ensure_sliced_smaller(in_batch, 1000);
3679 }
3680
3681 #[test]
3682 fn encode_nested_lists_starting_at_zero() {
3683 let inner_int = Arc::new(Field::new("item", DataType::UInt32, true));
3684 let inner_list_field = Arc::new(Field::new("item", DataType::List(inner_int), true));
3685 let list_field = Field::new("val", DataType::List(inner_list_field), true);
3686 let schema = Arc::new(Schema::new(vec![list_field]));
3687
3688 let values = Arc::new(generate_nested_list_data_starting_at_zero::<i32>());
3689
3690 let in_batch = RecordBatch::try_new(schema, vec![values]).unwrap();
3691 roundtrip_ensure_sliced_smaller(in_batch, 1);
3692 }
3693
3694 #[test]
3695 fn encode_map_array() {
3696 let keys = Arc::new(Field::new(
3697 Field::MAP_KEY_FIELD_DEFAULT_NAME,
3698 DataType::UInt32,
3699 false,
3700 ));
3701 let values = Arc::new(Field::new(
3702 Field::MAP_VALUE_FIELD_DEFAULT_NAME,
3703 DataType::UInt32,
3704 true,
3705 ));
3706 let map_field = Field::new_map(
3707 "map",
3708 Field::MAP_ENTRIES_FIELD_DEFAULT_NAME,
3709 keys,
3710 values,
3711 false,
3712 true,
3713 );
3714 let schema = Arc::new(Schema::new(vec![map_field]));
3715
3716 let values = Arc::new(generate_map_array_data());
3717
3718 let in_batch = RecordBatch::try_new(schema, vec![values]).unwrap();
3719 roundtrip_ensure_sliced_smaller(in_batch, 1000);
3720 }
3721
3722 fn generate_list_view_data<O: OffsetSizeTrait>() -> GenericListViewArray<O> {
3723 let mut builder = GenericListViewBuilder::<O, _>::new(UInt32Builder::new());
3724
3725 for i in 0u32..100_000 {
3726 if i.is_multiple_of(10_000) {
3727 builder.append(false);
3728 continue;
3729 }
3730 for value in [i, i, i] {
3731 builder.values().append_value(value);
3732 }
3733 builder.append(true);
3734 }
3735
3736 builder.finish()
3737 }
3738
3739 #[test]
3740 fn encode_list_view_arrays() {
3741 let val_inner = Field::new_list_field(DataType::UInt32, true);
3742 let val_field = Field::new("val", DataType::ListView(Arc::new(val_inner)), true);
3743 let schema = Arc::new(Schema::new(vec![val_field]));
3744
3745 let values = Arc::new(generate_list_view_data::<i32>());
3746
3747 let in_batch = RecordBatch::try_new(schema, vec![values]).unwrap();
3748 let out_batch = deserialize_file(serialize_file(&in_batch));
3749 assert_eq!(in_batch, out_batch);
3750 }
3751
3752 #[test]
3753 fn encode_large_list_view_arrays() {
3754 let val_inner = Field::new_list_field(DataType::UInt32, true);
3755 let val_field = Field::new("val", DataType::LargeListView(Arc::new(val_inner)), true);
3756 let schema = Arc::new(Schema::new(vec![val_field]));
3757
3758 let values = Arc::new(generate_list_view_data::<i64>());
3759
3760 let in_batch = RecordBatch::try_new(schema, vec![values]).unwrap();
3761 let out_batch = deserialize_file(serialize_file(&in_batch));
3762 assert_eq!(in_batch, out_batch);
3763 }
3764
3765 #[test]
3766 fn check_sliced_list_view_array() {
3767 let inner = Field::new_list_field(DataType::UInt32, true);
3768 let field = Field::new("val", DataType::ListView(Arc::new(inner)), true);
3769 let schema = Arc::new(Schema::new(vec![field]));
3770 let values = Arc::new(generate_list_view_data::<i32>());
3771
3772 for (offset, len) in [(999, 1), (0, 13), (47, 12), (values.len() - 13, 13)] {
3773 let in_batch = RecordBatch::try_new(schema.clone(), vec![values.clone()])
3774 .unwrap()
3775 .slice(offset, len);
3776 let out_batch = deserialize_file(serialize_file(&in_batch));
3777 assert_eq!(in_batch, out_batch);
3778 }
3779 }
3780
3781 #[test]
3782 fn check_sliced_large_list_view_array() {
3783 let inner = Field::new_list_field(DataType::UInt32, true);
3784 let field = Field::new("val", DataType::LargeListView(Arc::new(inner)), true);
3785 let schema = Arc::new(Schema::new(vec![field]));
3786 let values = Arc::new(generate_list_view_data::<i64>());
3787
3788 for (offset, len) in [(999, 1), (0, 13), (47, 12), (values.len() - 13, 13)] {
3789 let in_batch = RecordBatch::try_new(schema.clone(), vec![values.clone()])
3790 .unwrap()
3791 .slice(offset, len);
3792 let out_batch = deserialize_file(serialize_file(&in_batch));
3793 assert_eq!(in_batch, out_batch);
3794 }
3795 }
3796
3797 fn generate_nested_list_view_data<O: OffsetSizeTrait>() -> GenericListViewArray<O> {
3798 let inner_builder = UInt32Builder::new();
3799 let middle_builder = GenericListViewBuilder::<O, _>::new(inner_builder);
3800 let mut outer_builder = GenericListViewBuilder::<O, _>::new(middle_builder);
3801
3802 for i in 0u32..10_000 {
3803 if i.is_multiple_of(1_000) {
3804 outer_builder.append(false);
3805 continue;
3806 }
3807
3808 for _ in 0..3 {
3809 for value in [i, i + 1, i + 2] {
3810 outer_builder.values().values().append_value(value);
3811 }
3812 outer_builder.values().append(true);
3813 }
3814 outer_builder.append(true);
3815 }
3816
3817 outer_builder.finish()
3818 }
3819
3820 #[test]
3821 fn encode_nested_list_views() {
3822 let inner_int = Arc::new(Field::new_list_field(DataType::UInt32, true));
3823 let inner_list_field = Arc::new(Field::new_list_field(DataType::ListView(inner_int), true));
3824 let list_field = Field::new("val", DataType::ListView(inner_list_field), true);
3825 let schema = Arc::new(Schema::new(vec![list_field]));
3826
3827 let values = Arc::new(generate_nested_list_view_data::<i32>());
3828
3829 let in_batch = RecordBatch::try_new(schema, vec![values]).unwrap();
3830 let out_batch = deserialize_file(serialize_file(&in_batch));
3831 assert_eq!(in_batch, out_batch);
3832 }
3833
3834 fn test_roundtrip_list_view_of_dict_impl<OffsetSize: OffsetSizeTrait, U: ArrowNativeType>(
3835 list_data_type: DataType,
3836 offsets: &[U; 5],
3837 sizes: &[U; 4],
3838 ) {
3839 let values = StringArray::from(vec![Some("alpha"), None, Some("beta"), Some("gamma")]);
3840 let keys = Int32Array::from_iter_values([0, 0, 1, 2, 3, 0, 2]);
3841 let dict_array = DictionaryArray::new(keys, Arc::new(values));
3842 let dict_data = dict_array.to_data();
3843
3844 let value_offsets = Buffer::from_slice_ref(offsets);
3845 let value_sizes = Buffer::from_slice_ref(sizes);
3846
3847 let list_data = ArrayData::builder(list_data_type)
3848 .len(4)
3849 .add_buffer(value_offsets)
3850 .add_buffer(value_sizes)
3851 .add_child_data(dict_data)
3852 .build()
3853 .unwrap();
3854 let list_view_array = GenericListViewArray::<OffsetSize>::from(list_data);
3855
3856 let schema = Arc::new(Schema::new(vec![Field::new(
3857 "f1",
3858 list_view_array.data_type().clone(),
3859 false,
3860 )]));
3861 let input_batch = RecordBatch::try_new(schema, vec![Arc::new(list_view_array)]).unwrap();
3862
3863 let output_batch = deserialize_file(serialize_file(&input_batch));
3864 assert_eq!(input_batch, output_batch);
3865
3866 let output_batch = deserialize_stream(serialize_stream(&input_batch));
3867 assert_eq!(input_batch, output_batch);
3868 }
3869
3870 #[test]
3871 fn test_roundtrip_list_view_of_dict() {
3872 #[allow(deprecated)]
3873 let list_data_type = DataType::ListView(Arc::new(Field::new_dict(
3874 "item",
3875 DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
3876 true,
3877 1,
3878 false,
3879 )));
3880 let offsets: &[i32; 5] = &[0, 2, 4, 4, 7];
3881 let sizes: &[i32; 4] = &[2, 2, 0, 3];
3882 test_roundtrip_list_view_of_dict_impl::<i32, i32>(list_data_type, offsets, sizes);
3883 }
3884
3885 #[test]
3886 fn test_roundtrip_large_list_view_of_dict() {
3887 #[allow(deprecated)]
3888 let list_data_type = DataType::LargeListView(Arc::new(Field::new_dict(
3889 "item",
3890 DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
3891 true,
3892 2,
3893 false,
3894 )));
3895 let offsets: &[i64; 5] = &[0, 2, 4, 4, 7];
3896 let sizes: &[i64; 4] = &[2, 2, 0, 3];
3897 test_roundtrip_list_view_of_dict_impl::<i64, i64>(list_data_type, offsets, sizes);
3898 }
3899
3900 #[test]
3901 fn test_roundtrip_sliced_list_view_of_dict() {
3902 #[allow(deprecated)]
3903 let list_data_type = DataType::ListView(Arc::new(Field::new_dict(
3904 "item",
3905 DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
3906 true,
3907 3,
3908 false,
3909 )));
3910
3911 let values = StringArray::from(vec![Some("alpha"), None, Some("beta"), Some("gamma")]);
3912 let keys = Int32Array::from_iter_values([0, 0, 1, 2, 3, 0, 2, 1, 0, 3, 2, 1]);
3913 let dict_array = DictionaryArray::new(keys, Arc::new(values));
3914 let dict_data = dict_array.to_data();
3915
3916 let offsets: &[i32; 7] = &[0, 2, 4, 4, 7, 9, 12];
3917 let sizes: &[i32; 6] = &[2, 2, 0, 3, 2, 3];
3918 let value_offsets = Buffer::from_slice_ref(offsets);
3919 let value_sizes = Buffer::from_slice_ref(sizes);
3920
3921 let list_data = ArrayData::builder(list_data_type)
3922 .len(6)
3923 .add_buffer(value_offsets)
3924 .add_buffer(value_sizes)
3925 .add_child_data(dict_data)
3926 .build()
3927 .unwrap();
3928 let list_view_array = GenericListViewArray::<i32>::from(list_data);
3929
3930 let schema = Arc::new(Schema::new(vec![Field::new(
3931 "f1",
3932 list_view_array.data_type().clone(),
3933 false,
3934 )]));
3935 let input_batch = RecordBatch::try_new(schema, vec![Arc::new(list_view_array)]).unwrap();
3936
3937 let sliced_batch = input_batch.slice(1, 4);
3938
3939 let output_batch = deserialize_file(serialize_file(&sliced_batch));
3940 assert_eq!(sliced_batch, output_batch);
3941
3942 let output_batch = deserialize_stream(serialize_stream(&sliced_batch));
3943 assert_eq!(sliced_batch, output_batch);
3944 }
3945
3946 #[test]
3947 fn test_roundtrip_dense_union_of_dict() {
3948 let values = StringArray::from(vec![Some("alpha"), None, Some("beta"), Some("gamma")]);
3949 let keys = Int32Array::from_iter_values([0, 0, 1, 2, 3, 0, 2]);
3950 let dict_array = DictionaryArray::new(keys, Arc::new(values));
3951
3952 #[allow(deprecated)]
3953 let dict_field = Arc::new(Field::new_dict(
3954 "dict",
3955 DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
3956 true,
3957 1,
3958 false,
3959 ));
3960 let int_field = Arc::new(Field::new("int", DataType::Int32, false));
3961 let union_fields = UnionFields::try_new(vec![0, 1], vec![dict_field, int_field]).unwrap();
3962
3963 let types = ScalarBuffer::from(vec![0i8, 0, 1, 0, 1, 0, 0]);
3964 let offsets = ScalarBuffer::from(vec![0i32, 1, 0, 2, 1, 3, 4]);
3965
3966 let int_array = Int32Array::from(vec![100, 200]);
3967
3968 let union = UnionArray::try_new(
3969 union_fields.clone(),
3970 types,
3971 Some(offsets),
3972 vec![Arc::new(dict_array), Arc::new(int_array)],
3973 )
3974 .unwrap();
3975
3976 let schema = Arc::new(Schema::new(vec![Field::new(
3977 "union",
3978 DataType::Union(union_fields, UnionMode::Dense),
3979 false,
3980 )]));
3981 let input_batch = RecordBatch::try_new(schema, vec![Arc::new(union)]).unwrap();
3982
3983 let output_batch = deserialize_file(serialize_file(&input_batch));
3984 assert_eq!(input_batch, output_batch);
3985
3986 let output_batch = deserialize_stream(serialize_stream(&input_batch));
3987 assert_eq!(input_batch, output_batch);
3988 }
3989
3990 #[test]
3991 fn test_roundtrip_sparse_union_of_dict() {
3992 let values = StringArray::from(vec![Some("alpha"), None, Some("beta"), Some("gamma")]);
3993 let keys = Int32Array::from_iter_values([0, 0, 1, 2, 3, 0, 2]);
3994 let dict_array = DictionaryArray::new(keys, Arc::new(values));
3995
3996 #[allow(deprecated)]
3997 let dict_field = Arc::new(Field::new_dict(
3998 "dict",
3999 DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
4000 true,
4001 2,
4002 false,
4003 ));
4004 let int_field = Arc::new(Field::new("int", DataType::Int32, false));
4005 let union_fields = UnionFields::try_new(vec![0, 1], vec![dict_field, int_field]).unwrap();
4006
4007 let types = ScalarBuffer::from(vec![0i8, 0, 1, 0, 1, 0, 0]);
4008
4009 let int_array = Int32Array::from(vec![0, 0, 100, 0, 200, 0, 0]);
4010
4011 let union = UnionArray::try_new(
4012 union_fields.clone(),
4013 types,
4014 None,
4015 vec![Arc::new(dict_array), Arc::new(int_array)],
4016 )
4017 .unwrap();
4018
4019 let schema = Arc::new(Schema::new(vec![Field::new(
4020 "union",
4021 DataType::Union(union_fields, UnionMode::Sparse),
4022 false,
4023 )]));
4024 let input_batch = RecordBatch::try_new(schema, vec![Arc::new(union)]).unwrap();
4025
4026 let output_batch = deserialize_file(serialize_file(&input_batch));
4027 assert_eq!(input_batch, output_batch);
4028
4029 let output_batch = deserialize_stream(serialize_stream(&input_batch));
4030 assert_eq!(input_batch, output_batch);
4031 }
4032
4033 #[test]
4034 fn test_roundtrip_map_with_dict_keys() {
4035 let key_values = StringArray::from(vec!["key_a", "key_b", "key_c"]);
4038 let keys = Int32Array::from_iter_values([0, 1, 2, 0, 1, 0]);
4039 let dict_keys = DictionaryArray::new(keys, Arc::new(key_values));
4040
4041 let values = Int32Array::from(vec![1, 2, 3, 4, 5, 6]);
4042
4043 #[allow(deprecated)]
4044 let entries_field = Arc::new(Field::new(
4045 Field::MAP_ENTRIES_FIELD_DEFAULT_NAME,
4046 DataType::Struct(
4047 vec![
4048 Field::new_dict(
4049 Field::MAP_KEY_FIELD_DEFAULT_NAME,
4050 DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
4051 false,
4052 1,
4053 false,
4054 ),
4055 Field::new(Field::MAP_VALUE_FIELD_DEFAULT_NAME, DataType::Int32, true),
4056 ]
4057 .into(),
4058 ),
4059 false,
4060 ));
4061
4062 let entries = StructArray::from(vec![
4063 (
4064 Arc::new(Field::new(
4065 Field::MAP_KEY_FIELD_DEFAULT_NAME,
4066 DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
4067 false,
4068 )),
4069 Arc::new(dict_keys) as ArrayRef,
4070 ),
4071 (
4072 Arc::new(Field::new(
4073 Field::MAP_VALUE_FIELD_DEFAULT_NAME,
4074 DataType::Int32,
4075 true,
4076 )),
4077 Arc::new(values) as ArrayRef,
4078 ),
4079 ]);
4080
4081 let offsets = Buffer::from_slice_ref([0i32, 2, 4, 6]);
4082
4083 let map_data = ArrayData::builder(DataType::Map(entries_field, false))
4084 .len(3)
4085 .add_buffer(offsets)
4086 .add_child_data(entries.into_data())
4087 .build()
4088 .unwrap();
4089 let map_array = MapArray::from(map_data);
4090
4091 let schema = Arc::new(Schema::new(vec![Field::new(
4092 "map",
4093 map_array.data_type().clone(),
4094 false,
4095 )]));
4096 let input_batch = RecordBatch::try_new(schema, vec![Arc::new(map_array)]).unwrap();
4097
4098 let output_batch = deserialize_file(serialize_file(&input_batch));
4099 assert_eq!(input_batch, output_batch);
4100
4101 let output_batch = deserialize_stream(serialize_stream(&input_batch));
4102 assert_eq!(input_batch, output_batch);
4103 }
4104
4105 #[test]
4106 fn test_roundtrip_map_with_dict_values() {
4107 let keys = StringArray::from(vec!["a", "b", "c", "d", "e", "f"]);
4110
4111 let value_values = StringArray::from(vec!["val_x", "val_y", "val_z"]);
4112 let value_keys = Int32Array::from_iter_values([0, 1, 2, 0, 1, 0]);
4113 let dict_values = DictionaryArray::new(value_keys, Arc::new(value_values));
4114
4115 #[allow(deprecated)]
4116 let entries_field = Arc::new(Field::new(
4117 Field::MAP_ENTRIES_FIELD_DEFAULT_NAME,
4118 DataType::Struct(
4119 vec![
4120 Field::new(Field::MAP_KEY_FIELD_DEFAULT_NAME, DataType::Utf8, false),
4121 Field::new_dict(
4122 Field::MAP_VALUE_FIELD_DEFAULT_NAME,
4123 DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
4124 true,
4125 2,
4126 false,
4127 ),
4128 ]
4129 .into(),
4130 ),
4131 false,
4132 ));
4133
4134 let entries = StructArray::from(vec![
4135 (
4136 Arc::new(Field::new(
4137 Field::MAP_KEY_FIELD_DEFAULT_NAME,
4138 DataType::Utf8,
4139 false,
4140 )),
4141 Arc::new(keys) as ArrayRef,
4142 ),
4143 (
4144 Arc::new(Field::new(
4145 Field::MAP_VALUE_FIELD_DEFAULT_NAME,
4146 DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
4147 true,
4148 )),
4149 Arc::new(dict_values) as ArrayRef,
4150 ),
4151 ]);
4152
4153 let offsets = Buffer::from_slice_ref([0i32, 2, 4, 6]);
4154
4155 let map_data = ArrayData::builder(DataType::Map(entries_field, false))
4156 .len(3)
4157 .add_buffer(offsets)
4158 .add_child_data(entries.into_data())
4159 .build()
4160 .unwrap();
4161 let map_array = MapArray::from(map_data);
4162
4163 let schema = Arc::new(Schema::new(vec![Field::new(
4164 "map",
4165 map_array.data_type().clone(),
4166 false,
4167 )]));
4168 let input_batch = RecordBatch::try_new(schema, vec![Arc::new(map_array)]).unwrap();
4169
4170 let output_batch = deserialize_file(serialize_file(&input_batch));
4171 assert_eq!(input_batch, output_batch);
4172
4173 let output_batch = deserialize_stream(serialize_stream(&input_batch));
4174 assert_eq!(input_batch, output_batch);
4175 }
4176
4177 #[test]
4178 fn test_decimal128_alignment16_is_sufficient() {
4179 const IPC_ALIGNMENT: usize = 16;
4180
4181 for num_cols in [1, 2, 3, 17, 50, 73, 99] {
4186 let num_rows = (num_cols * 7 + 11) % 100; let mut fields = Vec::new();
4189 let mut arrays = Vec::new();
4190 for i in 0..num_cols {
4191 let field = Field::new(format!("col_{i}"), DataType::Decimal128(38, 10), true);
4192 let array = Decimal128Array::from(vec![num_cols as i128; num_rows]);
4193 fields.push(field);
4194 arrays.push(Arc::new(array) as Arc<dyn Array>);
4195 }
4196 let schema = Schema::new(fields);
4197 let batch = RecordBatch::try_new(Arc::new(schema), arrays).unwrap();
4198
4199 let mut writer = FileWriter::try_new_with_options(
4200 Vec::new(),
4201 batch.schema_ref(),
4202 IpcWriteOptions::try_new(IPC_ALIGNMENT, false, MetadataVersion::V5).unwrap(),
4203 )
4204 .unwrap();
4205 writer.write(&batch).unwrap();
4206 writer.finish().unwrap();
4207
4208 let out: Vec<u8> = writer.into_inner().unwrap();
4209
4210 let buffer = Buffer::from_vec(out);
4211 let trailer_start = buffer.len() - 10;
4212 let footer_len =
4213 read_footer_length(buffer[trailer_start..].try_into().unwrap()).unwrap();
4214 let footer =
4215 root_as_footer(&buffer[trailer_start - footer_len..trailer_start]).unwrap();
4216
4217 let schema = fb_to_schema(footer.schema().unwrap());
4218
4219 let decoder =
4222 FileDecoder::new(Arc::new(schema), footer.version()).with_require_alignment(true);
4223
4224 let batches = footer.recordBatches().unwrap();
4225
4226 let block = batches.get(0);
4227 let block_len = block.bodyLength() as usize + block.metaDataLength() as usize;
4228 let data = buffer.slice_with_length(block.offset() as _, block_len);
4229
4230 let batch2 = decoder.read_record_batch(block, &data).unwrap().unwrap();
4231
4232 assert_eq!(batch, batch2);
4233 }
4234 }
4235
4236 #[test]
4237 fn test_decimal128_alignment8_is_unaligned() {
4238 const IPC_ALIGNMENT: usize = 8;
4239
4240 let num_cols = 2;
4241 let num_rows = 1;
4242
4243 let mut fields = Vec::new();
4244 let mut arrays = Vec::new();
4245 for i in 0..num_cols {
4246 let field = Field::new(format!("col_{i}"), DataType::Decimal128(38, 10), true);
4247 let array = Decimal128Array::from(vec![num_cols as i128; num_rows]);
4248 fields.push(field);
4249 arrays.push(Arc::new(array) as Arc<dyn Array>);
4250 }
4251 let schema = Schema::new(fields);
4252 let batch = RecordBatch::try_new(Arc::new(schema), arrays).unwrap();
4253
4254 let mut writer = FileWriter::try_new_with_options(
4255 Vec::new(),
4256 batch.schema_ref(),
4257 IpcWriteOptions::try_new(IPC_ALIGNMENT, false, MetadataVersion::V5).unwrap(),
4258 )
4259 .unwrap();
4260 writer.write(&batch).unwrap();
4261 writer.finish().unwrap();
4262
4263 let out: Vec<u8> = writer.into_inner().unwrap();
4264
4265 let buffer = Buffer::from_vec(out);
4266 let trailer_start = buffer.len() - 10;
4267 let footer_len = read_footer_length(buffer[trailer_start..].try_into().unwrap()).unwrap();
4268 let footer = root_as_footer(&buffer[trailer_start - footer_len..trailer_start]).unwrap();
4269 let schema = fb_to_schema(footer.schema().unwrap());
4270
4271 let decoder =
4274 FileDecoder::new(Arc::new(schema), footer.version()).with_require_alignment(true);
4275
4276 let batches = footer.recordBatches().unwrap();
4277
4278 let block = batches.get(0);
4279 let block_len = block.bodyLength() as usize + block.metaDataLength() as usize;
4280 let data = buffer.slice_with_length(block.offset() as _, block_len);
4281
4282 let result = decoder.read_record_batch(block, &data);
4283
4284 let error = result.unwrap_err();
4285 assert_eq!(
4286 error.to_string(),
4287 "Invalid argument error: Misaligned buffers[0] in array of type Decimal128(38, 10), \
4288 offset from expected alignment of 16 by 8"
4289 );
4290 }
4291
4292 #[test]
4293 fn test_flush() {
4294 let num_cols = 2;
4297 let mut fields = Vec::new();
4298 let options = IpcWriteOptions::try_new(8, false, MetadataVersion::V5).unwrap();
4299 for i in 0..num_cols {
4300 let field = Field::new(format!("col_{i}"), DataType::Decimal128(38, 10), true);
4301 fields.push(field);
4302 }
4303 let schema = Schema::new(fields);
4304 let inner_stream_writer = BufWriter::with_capacity(1024, Vec::new());
4305 let inner_file_writer = BufWriter::with_capacity(1024, Vec::new());
4306 let mut stream_writer =
4307 StreamWriter::try_new_with_options(inner_stream_writer, &schema, options.clone())
4308 .unwrap();
4309 let mut file_writer =
4310 FileWriter::try_new_with_options(inner_file_writer, &schema, options).unwrap();
4311
4312 let stream_bytes_written_on_new = stream_writer.get_ref().get_ref().len();
4313 let file_bytes_written_on_new = file_writer.get_ref().get_ref().len();
4314 stream_writer.flush().unwrap();
4315 file_writer.flush().unwrap();
4316 let stream_bytes_written_on_flush = stream_writer.get_ref().get_ref().len();
4317 let file_bytes_written_on_flush = file_writer.get_ref().get_ref().len();
4318 let stream_out = stream_writer.into_inner().unwrap().into_inner().unwrap();
4319 let expected_stream_flushed_bytes = stream_out.len() - 8;
4323 let expected_file_flushed_bytes = expected_stream_flushed_bytes + 8;
4326
4327 assert!(
4328 stream_bytes_written_on_new < stream_bytes_written_on_flush,
4329 "this test makes no sense if flush is not actually required"
4330 );
4331 assert!(
4332 file_bytes_written_on_new < file_bytes_written_on_flush,
4333 "this test makes no sense if flush is not actually required"
4334 );
4335 assert_eq!(stream_bytes_written_on_flush, expected_stream_flushed_bytes);
4336 assert_eq!(file_bytes_written_on_flush, expected_file_flushed_bytes);
4337 }
4338
4339 #[test]
4340 fn test_roundtrip_list_of_fixed_list() -> Result<(), ArrowError> {
4341 let l1_type =
4342 DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Float32, false)), 3);
4343 let l2_type = DataType::List(Arc::new(Field::new("item", l1_type.clone(), false)));
4344
4345 let l0_builder = Float32Builder::new();
4346 let l1_builder = FixedSizeListBuilder::new(l0_builder, 3).with_field(Arc::new(Field::new(
4347 "item",
4348 DataType::Float32,
4349 false,
4350 )));
4351 let mut l2_builder =
4352 ListBuilder::new(l1_builder).with_field(Arc::new(Field::new("item", l1_type, false)));
4353
4354 for point in [[1.0, 2.0, 3.0], [4.0, 5.0, 6.0], [7.0, 8.0, 9.0]] {
4355 l2_builder.values().values().append_value(point[0]);
4356 l2_builder.values().values().append_value(point[1]);
4357 l2_builder.values().values().append_value(point[2]);
4358
4359 l2_builder.values().append(true);
4360 }
4361 l2_builder.append(true);
4362
4363 let point = [10., 11., 12.];
4364 l2_builder.values().values().append_value(point[0]);
4365 l2_builder.values().values().append_value(point[1]);
4366 l2_builder.values().values().append_value(point[2]);
4367
4368 l2_builder.values().append(true);
4369 l2_builder.append(true);
4370
4371 let array = Arc::new(l2_builder.finish()) as ArrayRef;
4372
4373 let schema = Arc::new(Schema::new_with_metadata(
4374 vec![Field::new("points", l2_type, false)],
4375 HashMap::default(),
4376 ));
4377
4378 test_slices(&array, &schema, 0, 1)?;
4381 test_slices(&array, &schema, 0, 2)?;
4382 test_slices(&array, &schema, 1, 1)?;
4383
4384 Ok(())
4385 }
4386
4387 #[test]
4388 fn test_roundtrip_list_of_fixed_list_w_nulls() -> Result<(), ArrowError> {
4389 let l0_builder = Float32Builder::new();
4390 let l1_builder = FixedSizeListBuilder::new(l0_builder, 3);
4391 let mut l2_builder = ListBuilder::new(l1_builder);
4392
4393 for point in [
4394 [Some(1.0), Some(2.0), None],
4395 [Some(4.0), Some(5.0), Some(6.0)],
4396 [None, Some(8.0), Some(9.0)],
4397 ] {
4398 for p in point {
4399 match p {
4400 Some(p) => l2_builder.values().values().append_value(p),
4401 None => l2_builder.values().values().append_null(),
4402 }
4403 }
4404
4405 l2_builder.values().append(true);
4406 }
4407 l2_builder.append(true);
4408
4409 let point = [Some(10.), None, None];
4410 for p in point {
4411 match p {
4412 Some(p) => l2_builder.values().values().append_value(p),
4413 None => l2_builder.values().values().append_null(),
4414 }
4415 }
4416
4417 l2_builder.values().append(true);
4418 l2_builder.append(true);
4419
4420 let array = Arc::new(l2_builder.finish()) as ArrayRef;
4421
4422 let schema = Arc::new(Schema::new_with_metadata(
4423 vec![Field::new(
4424 "points",
4425 DataType::List(Arc::new(Field::new(
4426 "item",
4427 DataType::FixedSizeList(
4428 Arc::new(Field::new("item", DataType::Float32, true)),
4429 3,
4430 ),
4431 true,
4432 ))),
4433 true,
4434 )],
4435 HashMap::default(),
4436 ));
4437
4438 test_slices(&array, &schema, 0, 1)?;
4441 test_slices(&array, &schema, 0, 2)?;
4442 test_slices(&array, &schema, 1, 1)?;
4443
4444 Ok(())
4445 }
4446
4447 fn test_slices(
4448 parent_array: &ArrayRef,
4449 schema: &SchemaRef,
4450 offset: usize,
4451 length: usize,
4452 ) -> Result<(), ArrowError> {
4453 let subarray = parent_array.slice(offset, length);
4454 let original_batch = RecordBatch::try_new(schema.clone(), vec![subarray])?;
4455
4456 let mut bytes = Vec::new();
4457 let mut writer = StreamWriter::try_new(&mut bytes, schema)?;
4458 writer.write(&original_batch)?;
4459 writer.finish()?;
4460
4461 let mut cursor = std::io::Cursor::new(bytes);
4462 let mut reader = StreamReader::try_new(&mut cursor, None)?;
4463 let returned_batch = reader.next().unwrap()?;
4464
4465 assert_eq!(original_batch, returned_batch);
4466
4467 Ok(())
4468 }
4469
4470 #[test]
4471 fn test_roundtrip_fixed_list() -> Result<(), ArrowError> {
4472 let int_builder = Int64Builder::new();
4473 let mut fixed_list_builder = FixedSizeListBuilder::new(int_builder, 3)
4474 .with_field(Arc::new(Field::new("item", DataType::Int64, false)));
4475
4476 for point in [[1, 2, 3], [4, 5, 6], [7, 8, 9], [10, 11, 12]] {
4477 fixed_list_builder.values().append_value(point[0]);
4478 fixed_list_builder.values().append_value(point[1]);
4479 fixed_list_builder.values().append_value(point[2]);
4480
4481 fixed_list_builder.append(true);
4482 }
4483
4484 let array = Arc::new(fixed_list_builder.finish()) as ArrayRef;
4485
4486 let schema = Arc::new(Schema::new_with_metadata(
4487 vec![Field::new(
4488 "points",
4489 DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Int64, false)), 3),
4490 false,
4491 )],
4492 HashMap::default(),
4493 ));
4494
4495 test_slices(&array, &schema, 0, 4)?;
4498 test_slices(&array, &schema, 0, 2)?;
4499 test_slices(&array, &schema, 1, 3)?;
4500 test_slices(&array, &schema, 2, 1)?;
4501
4502 Ok(())
4503 }
4504
4505 #[test]
4506 fn test_roundtrip_fixed_list_w_nulls() -> Result<(), ArrowError> {
4507 let int_builder = Int64Builder::new();
4508 let mut fixed_list_builder = FixedSizeListBuilder::new(int_builder, 3);
4509
4510 for point in [
4511 [Some(1), Some(2), None],
4512 [Some(4), Some(5), Some(6)],
4513 [None, Some(8), Some(9)],
4514 [Some(10), None, None],
4515 ] {
4516 for p in point {
4517 match p {
4518 Some(p) => fixed_list_builder.values().append_value(p),
4519 None => fixed_list_builder.values().append_null(),
4520 }
4521 }
4522
4523 fixed_list_builder.append(true);
4524 }
4525
4526 let array = Arc::new(fixed_list_builder.finish()) as ArrayRef;
4527
4528 let schema = Arc::new(Schema::new_with_metadata(
4529 vec![Field::new(
4530 "points",
4531 DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Int64, true)), 3),
4532 true,
4533 )],
4534 HashMap::default(),
4535 ));
4536
4537 test_slices(&array, &schema, 0, 4)?;
4540 test_slices(&array, &schema, 0, 2)?;
4541 test_slices(&array, &schema, 1, 3)?;
4542 test_slices(&array, &schema, 2, 1)?;
4543
4544 Ok(())
4545 }
4546
4547 #[test]
4548 fn test_metadata_encoding_ordering() {
4549 fn create_hash() -> u64 {
4550 let metadata: HashMap<String, String> = [
4551 ("a", "1"), ("b", "2"), ("c", "3"), ("d", "4"), ("e", "5"), ]
4557 .into_iter()
4558 .map(|(k, v)| (k.to_owned(), v.to_owned()))
4559 .collect();
4560
4561 let schema = Arc::new(
4563 Schema::new(vec![
4564 Field::new("a", DataType::Int64, true).with_metadata(metadata.clone()),
4565 ])
4566 .with_metadata(metadata)
4567 .clone(),
4568 );
4569 let batch = RecordBatch::new_empty(schema.clone());
4570
4571 let mut bytes = Vec::new();
4572 let mut w = StreamWriter::try_new(&mut bytes, batch.schema_ref()).unwrap();
4573 w.write(&batch).unwrap();
4574 w.finish().unwrap();
4575
4576 let mut h = std::hash::DefaultHasher::new();
4577 h.write(&bytes);
4578 h.finish()
4579 }
4580
4581 let expected = create_hash();
4582
4583 let all_passed = (0..20).all(|_| create_hash() == expected);
4588 assert!(all_passed);
4589 }
4590
4591 #[test]
4592 fn test_dictionary_tracker_reset() {
4593 let data_gen = IpcDataGenerator::default();
4594 let mut dictionary_tracker = DictionaryTracker::new(false);
4595 let writer_options = IpcWriteOptions::default();
4596 let mut compression_ctx = IpcWriteContext::default();
4597
4598 let schema = Arc::new(Schema::new(vec![Field::new(
4599 "a",
4600 DataType::Dictionary(Box::new(DataType::UInt8), Box::new(DataType::Utf8)),
4601 false,
4602 )]));
4603
4604 let mut write_single_batch_stream =
4605 |batch: RecordBatch, dict_tracker: &mut DictionaryTracker| -> Vec<u8> {
4606 let mut buffer = Vec::new();
4607
4608 let stream_header = data_gen.schema_to_bytes_with_dictionary_tracker(
4610 &schema,
4611 dict_tracker,
4612 &writer_options,
4613 );
4614 _ = write_message(&mut buffer, stream_header, &writer_options).unwrap();
4615
4616 let (encoded_dicts, encoded_batch) = data_gen
4617 .encode(&batch, dict_tracker, &writer_options, &mut compression_ctx)
4618 .unwrap();
4619 for encoded_dict in encoded_dicts {
4620 _ = write_message(&mut buffer, encoded_dict, &writer_options).unwrap();
4621 }
4622 _ = write_message(&mut buffer, encoded_batch, &writer_options).unwrap();
4623
4624 buffer
4625 };
4626
4627 let batch1 = RecordBatch::try_new(
4628 schema.clone(),
4629 vec![Arc::new(DictionaryArray::new(
4630 UInt8Array::from_iter_values([0]),
4631 Arc::new(StringArray::from_iter_values(["a"])),
4632 ))],
4633 )
4634 .unwrap();
4635 let buffer = write_single_batch_stream(batch1.clone(), &mut dictionary_tracker);
4636
4637 let mut reader = StreamReader::try_new(Cursor::new(buffer), None).unwrap();
4639 let read_batch = reader.next().unwrap().unwrap();
4640 assert_eq!(read_batch, batch1);
4641
4642 dictionary_tracker.clear();
4644
4645 let batch2 = RecordBatch::try_new(
4647 schema.clone(),
4648 vec![Arc::new(DictionaryArray::new(
4649 UInt8Array::from_iter_values([0]),
4650 Arc::new(StringArray::from_iter_values(["a"])),
4651 ))],
4652 )
4653 .unwrap();
4654 let buffer = write_single_batch_stream(batch2.clone(), &mut dictionary_tracker);
4655 let mut reader = StreamReader::try_new(Cursor::new(buffer), None).unwrap();
4656 let read_batch = reader.next().unwrap().unwrap();
4657 assert_eq!(read_batch, batch2);
4658 }
4659}