1mod stream;
28pub use stream::*;
29
30use arrow_select::concat;
31
32use flatbuffers::{VectorIter, VerifierOptions};
33use std::collections::{HashMap, VecDeque};
34use std::fmt;
35use std::io::{BufReader, Read, Seek, SeekFrom};
36use std::sync::Arc;
37
38use arrow_array::*;
39use arrow_buffer::{
40 ArrowNativeType, BooleanBuffer, Buffer, MutableBuffer, NullBuffer, ScalarBuffer,
41};
42use arrow_data::{ArrayData, ArrayDataBuilder, UnsafeFlag};
43use arrow_schema::*;
44
45use crate::compression::{CompressionCodec, DecompressionContext};
46use crate::r#gen::Message;
47use crate::{Block, CONTINUATION_MARKER, FieldNode, MetadataVersion};
48use DataType::*;
49
50fn read_buffer(
60 buf: &crate::Buffer,
61 a_data: &Buffer,
62 compression_codec: Option<CompressionCodec>,
63 decompression_context: &mut DecompressionContext,
64) -> Result<Buffer, ArrowError> {
65 let start_offset = buf.offset() as usize;
66 let buf_data = a_data.slice_with_length(start_offset, buf.length() as usize);
67 match (buf_data.is_empty(), compression_codec) {
69 (true, _) | (_, None) => Ok(buf_data),
70 (false, Some(decompressor)) => {
71 decompressor.decompress_to_buffer(&buf_data, decompression_context)
72 }
73 }
74}
75impl RecordBatchDecoder<'_> {
76 fn create_array(
89 &mut self,
90 field: &Field,
91 variadic_counts: &mut VecDeque<i64>,
92 ) -> Result<ArrayRef, ArrowError> {
93 let data_type = field.data_type();
94 match data_type {
95 Utf8 | Binary | LargeBinary | LargeUtf8 => {
96 let field_node = self.next_node(field)?;
97 let buffers = [
98 self.next_buffer()?,
99 self.next_buffer()?,
100 self.next_buffer()?,
101 ];
102 self.create_primitive_array(field_node, data_type, &buffers)
103 }
104 BinaryView | Utf8View => {
105 let count = variadic_counts
106 .pop_front()
107 .ok_or(ArrowError::IpcError(format!(
108 "Missing variadic count for {data_type} column"
109 )))?;
110 let count = count + 2; let buffers = (0..count)
112 .map(|_| self.next_buffer())
113 .collect::<Result<Vec<_>, _>>()?;
114 let field_node = self.next_node(field)?;
115 self.create_primitive_array(field_node, data_type, &buffers)
116 }
117 FixedSizeBinary(_) => {
118 let field_node = self.next_node(field)?;
119 let buffers = [self.next_buffer()?, self.next_buffer()?];
120 self.create_primitive_array(field_node, data_type, &buffers)
121 }
122 List(list_field) | LargeList(list_field) | Map(list_field, _) => {
123 let list_node = self.next_node(field)?;
124 let list_buffers = [self.next_buffer()?, self.next_buffer()?];
125 let values = self.create_array(list_field, variadic_counts)?;
126 self.create_list_array(list_node, data_type, &list_buffers, values)
127 }
128 ListView(list_field) | LargeListView(list_field) => {
129 let list_node = self.next_node(field)?;
130 let list_buffers = [
131 self.next_buffer()?, self.next_buffer()?, self.next_buffer()?, ];
135 let values = self.create_array(list_field, variadic_counts)?;
136 self.create_list_view_array(list_node, data_type, &list_buffers, values)
137 }
138 FixedSizeList(list_field, _) => {
139 let list_node = self.next_node(field)?;
140 let list_buffers = [self.next_buffer()?];
141 let values = self.create_array(list_field, variadic_counts)?;
142 self.create_list_array(list_node, data_type, &list_buffers, values)
143 }
144 Struct(struct_fields) => {
145 let struct_node = self.next_node(field)?;
146 let null_buffer = self.next_buffer()?;
147
148 let mut struct_arrays = Vec::with_capacity(struct_fields.len());
150 for struct_field in struct_fields {
153 let child = self.create_array(struct_field, variadic_counts)?;
154 struct_arrays.push(child);
155 }
156 self.create_struct_array(struct_node, null_buffer, struct_fields, struct_arrays)
157 }
158 RunEndEncoded(run_ends_field, values_field) => {
159 let run_node = self.next_node(field)?;
160 let run_ends = self.create_array(run_ends_field, variadic_counts)?;
161 let values = self.create_array(values_field, variadic_counts)?;
162
163 let run_array_length = run_node.length() as usize;
164 let builder = ArrayData::builder(data_type.clone())
165 .len(run_array_length)
166 .offset(0)
167 .add_child_data(run_ends.into_data())
168 .add_child_data(values.into_data())
169 .null_count(run_node.null_count() as usize);
170
171 self.create_array_from_builder(builder)
172 }
173 Dictionary(_, _) => {
175 let index_node = self.next_node(field)?;
176 let index_buffers = [self.next_buffer()?, self.next_buffer()?];
177
178 #[allow(deprecated)]
179 let dict_id = field.dict_id().ok_or_else(|| {
180 ArrowError::ParseError(format!("Field {field} does not have dict id"))
181 })?;
182
183 let value_array = match self.dictionaries_by_id.get(&dict_id) {
184 Some(array) => array.clone(),
185 None => {
186 if let Dictionary(_, value_type) = data_type {
190 arrow_array::new_empty_array(value_type.as_ref())
191 } else {
192 unreachable!()
193 }
194 }
195 };
196
197 self.create_dictionary_array(index_node, data_type, &index_buffers, value_array)
198 }
199 Union(fields, mode) => {
200 let union_node = self.next_node(field)?;
201 let len = union_node.length() as usize;
202
203 if self.version < MetadataVersion::V5 {
206 self.next_buffer()?;
207 }
208
209 let type_ids: ScalarBuffer<i8> =
210 self.next_buffer()?.slice_with_length(0, len).into();
211
212 let value_offsets = match mode {
213 UnionMode::Dense => {
214 let offsets: ScalarBuffer<i32> =
215 self.next_buffer()?.slice_with_length(0, len * 4).into();
216 Some(offsets)
217 }
218 UnionMode::Sparse => None,
219 };
220
221 let mut children = Vec::with_capacity(fields.len());
222
223 for (_id, field) in fields.iter() {
224 let child = self.create_array(field, variadic_counts)?;
225 children.push(child);
226 }
227
228 let array = if self.skip_validation.get() {
229 unsafe {
231 UnionArray::new_unchecked(fields.clone(), type_ids, value_offsets, children)
232 }
233 } else {
234 UnionArray::try_new(fields.clone(), type_ids, value_offsets, children)?
235 };
236 Ok(Arc::new(array))
237 }
238 Null => {
239 let node = self.next_node(field)?;
240 let length = node.length();
241 let null_count = node.null_count();
242
243 if length != null_count {
244 return Err(ArrowError::SchemaError(format!(
245 "Field {field} of NullArray has unequal null_count {null_count} and len {length}"
246 )));
247 }
248
249 let builder = ArrayData::builder(data_type.clone())
250 .len(length as usize)
251 .offset(0);
252 self.create_array_from_builder(builder)
253 }
254 _ => {
255 let field_node = self.next_node(field)?;
256 let buffers = [self.next_buffer()?, self.next_buffer()?];
257 self.create_primitive_array(field_node, data_type, &buffers)
258 }
259 }
260 }
261
262 fn create_primitive_array(
265 &self,
266 field_node: &FieldNode,
267 data_type: &DataType,
268 buffers: &[Buffer],
269 ) -> Result<ArrayRef, ArrowError> {
270 let length = field_node.length() as usize;
271 let null_buffer = (field_node.null_count() > 0).then_some(buffers[0].clone());
272 let mut builder = match data_type {
273 Utf8 | Binary | LargeBinary | LargeUtf8 => {
274 ArrayData::builder(data_type.clone())
276 .len(length)
277 .buffers(buffers[1..3].to_vec())
278 .null_bit_buffer(null_buffer)
279 }
280 BinaryView | Utf8View => ArrayData::builder(data_type.clone())
281 .len(length)
282 .buffers(buffers[1..].to_vec())
283 .null_bit_buffer(null_buffer),
284 _ if data_type.is_primitive() || matches!(data_type, Boolean | FixedSizeBinary(_)) => {
285 ArrayData::builder(data_type.clone())
287 .len(length)
288 .add_buffer(buffers[1].clone())
289 .null_bit_buffer(null_buffer)
290 }
291 t => unreachable!("Data type {:?} either unsupported or not primitive", t),
292 };
293
294 builder = builder.null_count(field_node.null_count() as usize);
295
296 self.create_array_from_builder(builder)
297 }
298
299 fn create_array_from_builder(&self, builder: ArrayDataBuilder) -> Result<ArrayRef, ArrowError> {
301 let mut builder = builder.align_buffers(!self.require_alignment);
302 if self.skip_validation.get() {
303 unsafe { builder = builder.skip_validation(true) }
305 };
306 Ok(make_array(builder.build()?))
307 }
308
309 fn create_list_array(
312 &self,
313 field_node: &FieldNode,
314 data_type: &DataType,
315 buffers: &[Buffer],
316 child_array: ArrayRef,
317 ) -> Result<ArrayRef, ArrowError> {
318 let null_buffer = (field_node.null_count() > 0).then_some(buffers[0].clone());
319 let length = field_node.length() as usize;
320 let child_data = child_array.into_data();
321 let mut builder = match data_type {
322 List(_) | LargeList(_) | Map(_, _) => ArrayData::builder(data_type.clone())
323 .len(length)
324 .add_buffer(buffers[1].clone())
325 .add_child_data(child_data)
326 .null_bit_buffer(null_buffer),
327
328 FixedSizeList(_, _) => ArrayData::builder(data_type.clone())
329 .len(length)
330 .add_child_data(child_data)
331 .null_bit_buffer(null_buffer),
332
333 _ => unreachable!("Cannot create list or map array from {:?}", data_type),
334 };
335
336 builder = builder.null_count(field_node.null_count() as usize);
337
338 self.create_array_from_builder(builder)
339 }
340
341 fn create_list_view_array(
342 &self,
343 field_node: &FieldNode,
344 data_type: &DataType,
345 buffers: &[Buffer],
346 child_array: ArrayRef,
347 ) -> Result<ArrayRef, ArrowError> {
348 assert!(matches!(data_type, ListView(_) | LargeListView(_)));
349
350 let null_buffer = (field_node.null_count() > 0).then_some(buffers[0].clone());
351 let length = field_node.length() as usize;
352 let child_data = child_array.into_data();
353
354 self.create_array_from_builder(
355 ArrayData::builder(data_type.clone())
356 .len(length)
357 .add_buffer(buffers[1].clone()) .add_buffer(buffers[2].clone()) .add_child_data(child_data)
360 .null_bit_buffer(null_buffer)
361 .null_count(field_node.null_count() as usize),
362 )
363 }
364
365 fn create_struct_array(
366 &self,
367 struct_node: &FieldNode,
368 null_buffer: Buffer,
369 struct_fields: &Fields,
370 struct_arrays: Vec<ArrayRef>,
371 ) -> Result<ArrayRef, ArrowError> {
372 let null_count = struct_node.null_count() as usize;
373 let len = struct_node.length() as usize;
374 let skip_validation = self.skip_validation.get();
375
376 let nulls = if null_count > 0 {
377 let validity_buffer = BooleanBuffer::new(null_buffer, 0, len);
378 let null_buffer = if skip_validation {
379 unsafe { NullBuffer::new_unchecked(validity_buffer, null_count) }
381 } else {
382 let null_buffer = NullBuffer::new(validity_buffer);
383
384 if null_buffer.null_count() != null_count {
385 return Err(ArrowError::InvalidArgumentError(format!(
386 "null_count value ({}) doesn't match actual number of nulls in array ({})",
387 null_count,
388 null_buffer.null_count()
389 )));
390 }
391
392 null_buffer
393 };
394
395 Some(null_buffer)
396 } else {
397 None
398 };
399 if struct_arrays.is_empty() {
400 return Ok(Arc::new(StructArray::new_empty_fields(len, nulls)));
403 }
404
405 let struct_array = if skip_validation {
406 unsafe { StructArray::new_unchecked(struct_fields.clone(), struct_arrays, nulls) }
408 } else {
409 StructArray::try_new(struct_fields.clone(), struct_arrays, nulls)?
410 };
411
412 Ok(Arc::new(struct_array))
413 }
414
415 fn create_dictionary_array(
418 &self,
419 field_node: &FieldNode,
420 data_type: &DataType,
421 buffers: &[Buffer],
422 value_array: ArrayRef,
423 ) -> Result<ArrayRef, ArrowError> {
424 if let Dictionary(_, _) = *data_type {
425 let null_buffer = (field_node.null_count() > 0).then_some(buffers[0].clone());
426 let builder = ArrayData::builder(data_type.clone())
427 .len(field_node.length() as usize)
428 .add_buffer(buffers[1].clone())
429 .add_child_data(value_array.into_data())
430 .null_bit_buffer(null_buffer)
431 .null_count(field_node.null_count() as usize);
432 self.create_array_from_builder(builder)
433 } else {
434 unreachable!("Cannot create dictionary array from {:?}", data_type)
435 }
436 }
437}
438
439pub struct RecordBatchDecoder<'a> {
445 batch: crate::RecordBatch<'a>,
447 schema: SchemaRef,
449 dictionaries_by_id: &'a HashMap<i64, ArrayRef>,
451 compression: Option<CompressionCodec>,
453 decompression_context: DecompressionContext,
455 version: MetadataVersion,
457 data: &'a Buffer,
459 nodes: VectorIter<'a, FieldNode>,
461 buffers: VectorIter<'a, crate::Buffer>,
463 projection: Option<&'a [usize]>,
466 require_alignment: bool,
469 skip_validation: UnsafeFlag,
473}
474
475impl<'a> RecordBatchDecoder<'a> {
476 pub fn try_new(
478 buf: &'a Buffer,
479 batch: crate::RecordBatch<'a>,
480 schema: SchemaRef,
481 dictionaries_by_id: &'a HashMap<i64, ArrayRef>,
482 metadata: &'a MetadataVersion,
483 ) -> Result<Self, ArrowError> {
484 let buffers = batch.buffers().ok_or_else(|| {
485 ArrowError::IpcError("Unable to get buffers from IPC RecordBatch".to_string())
486 })?;
487 let field_nodes = batch.nodes().ok_or_else(|| {
488 ArrowError::IpcError("Unable to get field nodes from IPC RecordBatch".to_string())
489 })?;
490
491 let batch_compression = batch.compression();
492 let compression = batch_compression
493 .map(|batch_compression| batch_compression.codec().try_into())
494 .transpose()?;
495
496 Ok(Self {
497 batch,
498 schema,
499 dictionaries_by_id,
500 compression,
501 decompression_context: DecompressionContext::new(),
502 version: *metadata,
503 data: buf,
504 nodes: field_nodes.iter(),
505 buffers: buffers.iter(),
506 projection: None,
507 require_alignment: false,
508 skip_validation: UnsafeFlag::new(),
509 })
510 }
511
512 pub fn with_projection(mut self, projection: Option<&'a [usize]>) -> Self {
517 self.projection = projection;
518 self
519 }
520
521 pub fn with_require_alignment(mut self, require_alignment: bool) -> Self {
527 self.require_alignment = require_alignment;
528 self
529 }
530
531 pub fn with_skip_validation(mut self, skip_validation: UnsafeFlag) -> Self {
544 self.skip_validation = skip_validation;
545 self
546 }
547
548 pub fn read_record_batch(mut self) -> Result<RecordBatch, ArrowError> {
550 let mut variadic_counts: VecDeque<i64> = self
551 .batch
552 .variadicBufferCounts()
553 .into_iter()
554 .flatten()
555 .collect();
556
557 let options = RecordBatchOptions::new().with_row_count(Some(self.batch.length() as usize));
558
559 let schema = Arc::clone(&self.schema);
560 if let Some(projection) = self.projection {
561 let mut arrays = Vec::with_capacity(projection.len());
562 for (idx, field) in schema.fields().iter().enumerate() {
564 let mut child = None;
566 for (proj_idx, projected_idx) in projection.iter().enumerate() {
567 if *projected_idx == idx {
568 if child.is_none() {
569 child = Some(self.create_array(field, &mut variadic_counts)?);
570 }
571
572 arrays.push((proj_idx, child.as_ref().unwrap().clone()));
574 }
575 }
576
577 if child.is_none() {
578 self.skip_field(field, &mut variadic_counts)?;
579 }
580 }
581
582 arrays.sort_by_key(|t| t.0);
583
584 let schema = Arc::new(schema.project(projection)?);
585 let columns = arrays.into_iter().map(|t| t.1).collect::<Vec<_>>();
586
587 if self.skip_validation.get() {
588 unsafe {
590 Ok(RecordBatch::new_unchecked(
591 schema,
592 columns,
593 self.batch.length() as usize,
594 ))
595 }
596 } else {
597 assert!(variadic_counts.is_empty());
598 RecordBatch::try_new_with_options(schema, columns, &options)
599 }
600 } else {
601 let mut children = Vec::with_capacity(schema.fields().len());
602 for field in schema.fields() {
604 let child = self.create_array(field, &mut variadic_counts)?;
605 children.push(child);
606 }
607
608 if self.skip_validation.get() {
609 unsafe {
611 Ok(RecordBatch::new_unchecked(
612 schema,
613 children,
614 self.batch.length() as usize,
615 ))
616 }
617 } else {
618 assert!(variadic_counts.is_empty());
619 RecordBatch::try_new_with_options(schema, children, &options)
620 }
621 }
622 }
623
624 fn next_buffer(&mut self) -> Result<Buffer, ArrowError> {
625 let buffer = self.buffers.next().ok_or_else(|| {
626 ArrowError::IpcError("Buffer count mismatched with metadata".to_string())
627 })?;
628 read_buffer(
629 buffer,
630 self.data,
631 self.compression,
632 &mut self.decompression_context,
633 )
634 }
635
636 fn skip_buffer(&mut self) {
637 self.buffers.next().unwrap();
638 }
639
640 fn next_node(&mut self, field: &Field) -> Result<&'a FieldNode, ArrowError> {
641 self.nodes.next().ok_or_else(|| {
642 ArrowError::SchemaError(format!(
643 "Invalid data for schema. {field} refers to node not found in schema",
644 ))
645 })
646 }
647
648 fn skip_field(
649 &mut self,
650 field: &Field,
651 variadic_count: &mut VecDeque<i64>,
652 ) -> Result<(), ArrowError> {
653 self.next_node(field)?;
654
655 match field.data_type() {
656 Utf8 | Binary | LargeBinary | LargeUtf8 => {
657 for _ in 0..3 {
658 self.skip_buffer()
659 }
660 }
661 Utf8View | BinaryView => {
662 let count = variadic_count
663 .pop_front()
664 .ok_or(ArrowError::IpcError(format!(
665 "Missing variadic count for {} column",
666 field.data_type()
667 )))?;
668 let count = count + 2; for _i in 0..count {
670 self.skip_buffer()
671 }
672 }
673 FixedSizeBinary(_) => {
674 self.skip_buffer();
675 self.skip_buffer();
676 }
677 List(list_field) | LargeList(list_field) | Map(list_field, _) => {
678 self.skip_buffer();
679 self.skip_buffer();
680 self.skip_field(list_field, variadic_count)?;
681 }
682 ListView(list_field) | LargeListView(list_field) => {
683 self.skip_buffer(); self.skip_buffer(); self.skip_buffer(); self.skip_field(list_field, variadic_count)?;
687 }
688 FixedSizeList(list_field, _) => {
689 self.skip_buffer();
690 self.skip_field(list_field, variadic_count)?;
691 }
692 Struct(struct_fields) => {
693 self.skip_buffer();
694
695 for struct_field in struct_fields {
697 self.skip_field(struct_field, variadic_count)?
698 }
699 }
700 RunEndEncoded(run_ends_field, values_field) => {
701 self.skip_field(run_ends_field, variadic_count)?;
702 self.skip_field(values_field, variadic_count)?;
703 }
704 Dictionary(_, _) => {
705 self.skip_buffer(); self.skip_buffer(); }
708 Union(fields, mode) => {
709 if self.version < MetadataVersion::V5 {
710 self.skip_buffer(); }
712 self.skip_buffer(); match mode {
715 UnionMode::Dense => self.skip_buffer(), UnionMode::Sparse => {}
717 };
718
719 for (_, field) in fields.iter() {
720 self.skip_field(field, variadic_count)?
721 }
722 }
723 Null => {}
725
726 Boolean
728 | Int8
729 | Int16
730 | Int32
731 | Int64
732 | UInt8
733 | UInt16
734 | UInt32
735 | UInt64
736 | Float16
737 | Float32
738 | Float64
739 | Timestamp(_, _)
740 | Date32
741 | Date64
742 | Time32(_)
743 | Time64(_)
744 | Duration(_)
745 | Interval(_)
746 | Decimal32(_, _)
747 | Decimal64(_, _)
748 | Decimal128(_, _)
749 | Decimal256(_, _) => {
750 self.skip_buffer();
751 self.skip_buffer();
752 }
753 };
754 Ok(())
755 }
756}
757
758pub fn read_record_batch(
769 buf: &Buffer,
770 batch: crate::RecordBatch,
771 schema: SchemaRef,
772 dictionaries_by_id: &HashMap<i64, ArrayRef>,
773 projection: Option<&[usize]>,
774 metadata: &MetadataVersion,
775) -> Result<RecordBatch, ArrowError> {
776 RecordBatchDecoder::try_new(buf, batch, schema, dictionaries_by_id, metadata)?
777 .with_projection(projection)
778 .with_require_alignment(false)
779 .read_record_batch()
780}
781
782pub fn read_dictionary(
785 buf: &Buffer,
786 batch: crate::DictionaryBatch,
787 schema: &Schema,
788 dictionaries_by_id: &mut HashMap<i64, ArrayRef>,
789 metadata: &MetadataVersion,
790) -> Result<(), ArrowError> {
791 read_dictionary_impl(
792 buf,
793 batch,
794 schema,
795 dictionaries_by_id,
796 metadata,
797 false,
798 UnsafeFlag::new(),
799 )
800}
801
802pub fn read_dictionary_impl(
804 buf: &Buffer,
805 batch: crate::DictionaryBatch,
806 schema: &Schema,
807 dictionaries_by_id: &mut HashMap<i64, ArrayRef>,
808 metadata: &MetadataVersion,
809 require_alignment: bool,
810 skip_validation: UnsafeFlag,
811) -> Result<(), ArrowError> {
812 let id = batch.id();
813
814 let dictionary_values = get_dictionary_values(
815 buf,
816 batch,
817 schema,
818 dictionaries_by_id,
819 metadata,
820 require_alignment,
821 skip_validation,
822 )?;
823
824 update_dictionaries(dictionaries_by_id, batch.isDelta(), id, dictionary_values)?;
825
826 Ok(())
827}
828
829fn update_dictionaries(
838 dictionaries_by_id: &mut HashMap<i64, ArrayRef>,
839 is_delta: bool,
840 dict_id: i64,
841 dict_values: ArrayRef,
842) -> Result<(), ArrowError> {
843 if !is_delta {
844 dictionaries_by_id.insert(dict_id, dict_values.clone());
848 return Ok(());
849 }
850
851 let existing = dictionaries_by_id.get(&dict_id).ok_or_else(|| {
852 ArrowError::InvalidArgumentError(format!(
853 "No existing dictionary for delta dictionary with id '{dict_id}'"
854 ))
855 })?;
856
857 let combined = concat::concat(&[existing, &dict_values]).map_err(|e| {
858 ArrowError::InvalidArgumentError(format!("Failed to concat delta dictionary: {e}"))
859 })?;
860
861 dictionaries_by_id.insert(dict_id, combined);
862
863 Ok(())
864}
865
866fn get_dictionary_values(
870 buf: &Buffer,
871 batch: crate::DictionaryBatch,
872 schema: &Schema,
873 dictionaries_by_id: &mut HashMap<i64, ArrayRef>,
874 metadata: &MetadataVersion,
875 require_alignment: bool,
876 skip_validation: UnsafeFlag,
877) -> Result<ArrayRef, ArrowError> {
878 let id = batch.id();
879 #[allow(deprecated)]
880 let fields_using_this_dictionary = schema.fields_with_dict_id(id);
881 let first_field = fields_using_this_dictionary.first().ok_or_else(|| {
882 ArrowError::InvalidArgumentError(format!("dictionary id {id} not found in schema"))
883 })?;
884
885 let dictionary_values: ArrayRef = match first_field.data_type() {
889 DataType::Dictionary(_, value_type) => {
890 let value = value_type.as_ref().clone();
892 let schema = Schema::new(vec![Field::new("", value, true)]);
893 let record_batch = RecordBatchDecoder::try_new(
895 buf,
896 batch.data().unwrap(),
897 Arc::new(schema),
898 dictionaries_by_id,
899 metadata,
900 )?
901 .with_require_alignment(require_alignment)
902 .with_skip_validation(skip_validation)
903 .read_record_batch()?;
904
905 Some(record_batch.column(0).clone())
906 }
907 _ => None,
908 }
909 .ok_or_else(|| {
910 ArrowError::InvalidArgumentError(format!("dictionary id {id} not found in schema"))
911 })?;
912
913 Ok(dictionary_values)
914}
915
916fn read_block<R: Read + Seek>(mut reader: R, block: &Block) -> Result<Buffer, ArrowError> {
918 reader.seek(SeekFrom::Start(block.offset() as u64))?;
919 let body_len = block.bodyLength().to_usize().unwrap();
920 let metadata_len = block.metaDataLength().to_usize().unwrap();
921 let total_len = body_len.checked_add(metadata_len).unwrap();
922
923 let mut buf = MutableBuffer::from_len_zeroed(total_len);
924 reader.read_exact(&mut buf)?;
925 Ok(buf.into())
926}
927
928fn parse_message(buf: &[u8]) -> Result<Message::Message<'_>, ArrowError> {
932 let buf = match buf[..4] == CONTINUATION_MARKER {
933 true => &buf[8..],
934 false => &buf[4..],
935 };
936 crate::root_as_message(buf)
937 .map_err(|err| ArrowError::ParseError(format!("Unable to get root as message: {err:?}")))
938}
939
940pub fn read_footer_length(buf: [u8; 10]) -> Result<usize, ArrowError> {
944 if buf[4..] != super::ARROW_MAGIC {
945 return Err(ArrowError::ParseError(
946 "Arrow file does not contain correct footer".to_string(),
947 ));
948 }
949
950 let footer_len = i32::from_le_bytes(buf[..4].try_into().unwrap());
952 footer_len
953 .try_into()
954 .map_err(|_| ArrowError::ParseError(format!("Invalid footer length: {footer_len}")))
955}
956
957#[derive(Debug)]
1022pub struct FileDecoder {
1023 schema: SchemaRef,
1024 dictionaries: HashMap<i64, ArrayRef>,
1025 version: MetadataVersion,
1026 projection: Option<Vec<usize>>,
1027 require_alignment: bool,
1028 skip_validation: UnsafeFlag,
1029}
1030
1031impl FileDecoder {
1032 pub fn new(schema: SchemaRef, version: MetadataVersion) -> Self {
1034 Self {
1035 schema,
1036 version,
1037 dictionaries: Default::default(),
1038 projection: None,
1039 require_alignment: false,
1040 skip_validation: UnsafeFlag::new(),
1041 }
1042 }
1043
1044 pub fn with_projection(mut self, projection: Vec<usize>) -> Self {
1046 self.projection = Some(projection);
1047 self
1048 }
1049
1050 pub fn with_require_alignment(mut self, require_alignment: bool) -> Self {
1063 self.require_alignment = require_alignment;
1064 self
1065 }
1066
1067 pub unsafe fn with_skip_validation(mut self, skip_validation: bool) -> Self {
1078 unsafe { self.skip_validation.set(skip_validation) };
1079 self
1080 }
1081
1082 fn read_message<'a>(&self, buf: &'a [u8]) -> Result<Message::Message<'a>, ArrowError> {
1083 let message = parse_message(buf)?;
1084
1085 if self.version != MetadataVersion::V1 && message.version() != self.version {
1087 return Err(ArrowError::IpcError(
1088 "Could not read IPC message as metadata versions mismatch".to_string(),
1089 ));
1090 }
1091 Ok(message)
1092 }
1093
1094 pub fn read_dictionary(&mut self, block: &Block, buf: &Buffer) -> Result<(), ArrowError> {
1096 let message = self.read_message(buf)?;
1097 match message.header_type() {
1098 crate::MessageHeader::DictionaryBatch => {
1099 let batch = message.header_as_dictionary_batch().unwrap();
1100 read_dictionary_impl(
1101 &buf.slice(block.metaDataLength() as _),
1102 batch,
1103 &self.schema,
1104 &mut self.dictionaries,
1105 &message.version(),
1106 self.require_alignment,
1107 self.skip_validation.clone(),
1108 )
1109 }
1110 t => Err(ArrowError::ParseError(format!(
1111 "Expecting DictionaryBatch in dictionary blocks, found {t:?}."
1112 ))),
1113 }
1114 }
1115
1116 pub fn read_record_batch(
1118 &self,
1119 block: &Block,
1120 buf: &Buffer,
1121 ) -> Result<Option<RecordBatch>, ArrowError> {
1122 let message = self.read_message(buf)?;
1123 match message.header_type() {
1124 crate::MessageHeader::Schema => Err(ArrowError::IpcError(
1125 "Not expecting a schema when messages are read".to_string(),
1126 )),
1127 crate::MessageHeader::RecordBatch => {
1128 let batch = message.header_as_record_batch().ok_or_else(|| {
1129 ArrowError::IpcError("Unable to read IPC message as record batch".to_string())
1130 })?;
1131 RecordBatchDecoder::try_new(
1133 &buf.slice(block.metaDataLength() as _),
1134 batch,
1135 self.schema.clone(),
1136 &self.dictionaries,
1137 &message.version(),
1138 )?
1139 .with_projection(self.projection.as_deref())
1140 .with_require_alignment(self.require_alignment)
1141 .with_skip_validation(self.skip_validation.clone())
1142 .read_record_batch()
1143 .map(Some)
1144 }
1145 crate::MessageHeader::NONE => Ok(None),
1146 t => Err(ArrowError::InvalidArgumentError(format!(
1147 "Reading types other than record batches not yet supported, unable to read {t:?}"
1148 ))),
1149 }
1150 }
1151}
1152
1153#[derive(Debug)]
1155pub struct FileReaderBuilder {
1156 projection: Option<Vec<usize>>,
1158 max_footer_fb_tables: usize,
1160 max_footer_fb_depth: usize,
1162}
1163
1164impl Default for FileReaderBuilder {
1165 fn default() -> Self {
1166 let verifier_options = VerifierOptions::default();
1167 Self {
1168 max_footer_fb_tables: verifier_options.max_tables,
1169 max_footer_fb_depth: verifier_options.max_depth,
1170 projection: None,
1171 }
1172 }
1173}
1174
1175impl FileReaderBuilder {
1176 pub fn new() -> Self {
1180 Self::default()
1181 }
1182
1183 pub fn with_projection(mut self, projection: Vec<usize>) -> Self {
1185 self.projection = Some(projection);
1186 self
1187 }
1188
1189 pub fn with_max_footer_fb_tables(mut self, max_footer_fb_tables: usize) -> Self {
1202 self.max_footer_fb_tables = max_footer_fb_tables;
1203 self
1204 }
1205
1206 pub fn with_max_footer_fb_depth(mut self, max_footer_fb_depth: usize) -> Self {
1219 self.max_footer_fb_depth = max_footer_fb_depth;
1220 self
1221 }
1222
1223 pub fn build<R: Read + Seek>(self, mut reader: R) -> Result<FileReader<R>, ArrowError> {
1225 let mut buffer = [0; 10];
1227 reader.seek(SeekFrom::End(-10))?;
1228 reader.read_exact(&mut buffer)?;
1229
1230 let footer_len = read_footer_length(buffer)?;
1231
1232 let mut footer_data = vec![0; footer_len];
1234 reader.seek(SeekFrom::End(-10 - footer_len as i64))?;
1235 reader.read_exact(&mut footer_data)?;
1236
1237 let verifier_options = VerifierOptions {
1238 max_tables: self.max_footer_fb_tables,
1239 max_depth: self.max_footer_fb_depth,
1240 ..Default::default()
1241 };
1242 let footer = crate::root_as_footer_with_opts(&verifier_options, &footer_data[..]).map_err(
1243 |err| ArrowError::ParseError(format!("Unable to get root as footer: {err:?}")),
1244 )?;
1245
1246 let blocks = footer.recordBatches().ok_or_else(|| {
1247 ArrowError::ParseError("Unable to get record batches from IPC Footer".to_string())
1248 })?;
1249
1250 let total_blocks = blocks.len();
1251
1252 let ipc_schema = footer.schema().unwrap();
1253 if !ipc_schema.endianness().equals_to_target_endianness() {
1254 return Err(ArrowError::IpcError(
1255 "the endianness of the source system does not match the endianness of the target system.".to_owned()
1256 ));
1257 }
1258
1259 let schema = crate::convert::fb_to_schema(ipc_schema);
1260
1261 let mut custom_metadata = HashMap::new();
1262 if let Some(fb_custom_metadata) = footer.custom_metadata() {
1263 for kv in fb_custom_metadata.into_iter() {
1264 custom_metadata.insert(
1265 kv.key().unwrap().to_string(),
1266 kv.value().unwrap().to_string(),
1267 );
1268 }
1269 }
1270
1271 let mut decoder = FileDecoder::new(Arc::new(schema), footer.version());
1272 if let Some(projection) = self.projection {
1273 decoder = decoder.with_projection(projection)
1274 }
1275
1276 if let Some(dictionaries) = footer.dictionaries() {
1278 for block in dictionaries {
1279 let buf = read_block(&mut reader, block)?;
1280 decoder.read_dictionary(block, &buf)?;
1281 }
1282 }
1283
1284 Ok(FileReader {
1285 reader,
1286 blocks: blocks.iter().copied().collect(),
1287 current_block: 0,
1288 total_blocks,
1289 decoder,
1290 custom_metadata,
1291 })
1292 }
1293}
1294
1295pub struct FileReader<R> {
1340 reader: R,
1342
1343 decoder: FileDecoder,
1345
1346 blocks: Vec<Block>,
1350
1351 current_block: usize,
1353
1354 total_blocks: usize,
1356
1357 custom_metadata: HashMap<String, String>,
1359}
1360
1361impl<R> fmt::Debug for FileReader<R> {
1362 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> Result<(), fmt::Error> {
1363 f.debug_struct("FileReader<R>")
1364 .field("decoder", &self.decoder)
1365 .field("blocks", &self.blocks)
1366 .field("current_block", &self.current_block)
1367 .field("total_blocks", &self.total_blocks)
1368 .finish_non_exhaustive()
1369 }
1370}
1371
1372impl<R: Read + Seek> FileReader<BufReader<R>> {
1373 pub fn try_new_buffered(reader: R, projection: Option<Vec<usize>>) -> Result<Self, ArrowError> {
1377 Self::try_new(BufReader::new(reader), projection)
1378 }
1379}
1380
1381impl<R: Read + Seek> FileReader<R> {
1382 pub fn try_new(reader: R, projection: Option<Vec<usize>>) -> Result<Self, ArrowError> {
1393 let builder = FileReaderBuilder {
1394 projection,
1395 ..Default::default()
1396 };
1397 builder.build(reader)
1398 }
1399
1400 pub fn custom_metadata(&self) -> &HashMap<String, String> {
1402 &self.custom_metadata
1403 }
1404
1405 pub fn num_batches(&self) -> usize {
1407 self.total_blocks
1408 }
1409
1410 pub fn schema(&self) -> SchemaRef {
1412 self.decoder.schema.clone()
1413 }
1414
1415 pub fn set_index(&mut self, index: usize) -> Result<(), ArrowError> {
1419 if index >= self.total_blocks {
1420 Err(ArrowError::InvalidArgumentError(format!(
1421 "Cannot set batch to index {} from {} total batches",
1422 index, self.total_blocks
1423 )))
1424 } else {
1425 self.current_block = index;
1426 Ok(())
1427 }
1428 }
1429
1430 fn maybe_next(&mut self) -> Result<Option<RecordBatch>, ArrowError> {
1431 let block = &self.blocks[self.current_block];
1432 self.current_block += 1;
1433
1434 let buffer = read_block(&mut self.reader, block)?;
1436 self.decoder.read_record_batch(block, &buffer)
1437 }
1438
1439 pub fn get_ref(&self) -> &R {
1443 &self.reader
1444 }
1445
1446 pub fn get_mut(&mut self) -> &mut R {
1450 &mut self.reader
1451 }
1452
1453 pub unsafe fn with_skip_validation(mut self, skip_validation: bool) -> Self {
1459 self.decoder = unsafe { self.decoder.with_skip_validation(skip_validation) };
1460 self
1461 }
1462}
1463
1464impl<R: Read + Seek> Iterator for FileReader<R> {
1465 type Item = Result<RecordBatch, ArrowError>;
1466
1467 fn next(&mut self) -> Option<Self::Item> {
1468 if self.current_block < self.total_blocks {
1470 self.maybe_next().transpose()
1471 } else {
1472 None
1473 }
1474 }
1475}
1476
1477impl<R: Read + Seek> RecordBatchReader for FileReader<R> {
1478 fn schema(&self) -> SchemaRef {
1479 self.schema()
1480 }
1481}
1482
1483pub struct StreamReader<R> {
1517 reader: MessageReader<R>,
1519
1520 schema: SchemaRef,
1522
1523 dictionaries_by_id: HashMap<i64, ArrayRef>,
1527
1528 finished: bool,
1532
1533 projection: Option<(Vec<usize>, Schema)>,
1535
1536 skip_validation: UnsafeFlag,
1540}
1541
1542impl<R> fmt::Debug for StreamReader<R> {
1543 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> std::result::Result<(), fmt::Error> {
1544 f.debug_struct("StreamReader<R>")
1545 .field("reader", &"R")
1546 .field("schema", &self.schema)
1547 .field("dictionaries_by_id", &self.dictionaries_by_id)
1548 .field("finished", &self.finished)
1549 .field("projection", &self.projection)
1550 .finish()
1551 }
1552}
1553
1554impl<R: Read> StreamReader<BufReader<R>> {
1555 pub fn try_new_buffered(reader: R, projection: Option<Vec<usize>>) -> Result<Self, ArrowError> {
1559 Self::try_new(BufReader::new(reader), projection)
1560 }
1561}
1562
1563impl<R: Read> StreamReader<R> {
1564 pub fn try_new(
1576 reader: R,
1577 projection: Option<Vec<usize>>,
1578 ) -> Result<StreamReader<R>, ArrowError> {
1579 let mut msg_reader = MessageReader::new(reader);
1580 let message = msg_reader.maybe_next()?;
1581 let Some((message, _)) = message else {
1582 return Err(ArrowError::IpcError(
1583 "Expected schema message, found empty stream.".to_string(),
1584 ));
1585 };
1586
1587 if message.header_type() != Message::MessageHeader::Schema {
1588 return Err(ArrowError::IpcError(format!(
1589 "Expected a schema as the first message in the stream, got: {:?}",
1590 message.header_type()
1591 )));
1592 }
1593
1594 let schema = message.header_as_schema().ok_or_else(|| {
1595 ArrowError::ParseError("Failed to parse schema from message header".to_string())
1596 })?;
1597 let schema = crate::convert::fb_to_schema(schema);
1598
1599 let dictionaries_by_id = HashMap::new();
1601
1602 let projection = match projection {
1603 Some(projection_indices) => {
1604 let schema = schema.project(&projection_indices)?;
1605 Some((projection_indices, schema))
1606 }
1607 _ => None,
1608 };
1609
1610 Ok(Self {
1611 reader: msg_reader,
1612 schema: Arc::new(schema),
1613 finished: false,
1614 dictionaries_by_id,
1615 projection,
1616 skip_validation: UnsafeFlag::new(),
1617 })
1618 }
1619
1620 pub fn schema(&self) -> SchemaRef {
1622 self.schema.clone()
1623 }
1624
1625 pub fn is_finished(&self) -> bool {
1627 self.finished
1628 }
1629
1630 fn maybe_next(&mut self) -> Result<Option<RecordBatch>, ArrowError> {
1631 if self.finished {
1632 return Ok(None);
1633 }
1634
1635 loop {
1637 let message = self.next_ipc_message()?;
1638 let Some(message) = message else {
1639 self.finished = true;
1641 return Ok(None);
1642 };
1643
1644 match message {
1645 IpcMessage::Schema(_) => {
1646 return Err(ArrowError::IpcError(
1647 "Expected a record batch, but found a schema".to_string(),
1648 ));
1649 }
1650 IpcMessage::RecordBatch(record_batch) => {
1651 return Ok(Some(record_batch));
1652 }
1653 IpcMessage::DictionaryBatch { .. } => {
1654 continue;
1655 }
1656 };
1657 }
1658 }
1659
1660 pub(crate) fn next_ipc_message(&mut self) -> Result<Option<IpcMessage>, ArrowError> {
1668 let message = self.reader.maybe_next()?;
1669 let Some((message, body)) = message else {
1670 return Ok(None);
1672 };
1673
1674 let ipc_message = match message.header_type() {
1675 Message::MessageHeader::Schema => {
1676 let schema = message.header_as_schema().ok_or_else(|| {
1677 ArrowError::ParseError("Failed to parse schema from message header".to_string())
1678 })?;
1679 let arrow_schema = crate::convert::fb_to_schema(schema);
1680 IpcMessage::Schema(arrow_schema)
1681 }
1682 Message::MessageHeader::RecordBatch => {
1683 let batch = message.header_as_record_batch().ok_or_else(|| {
1684 ArrowError::IpcError("Unable to read IPC message as record batch".to_string())
1685 })?;
1686
1687 let version = message.version();
1688 let schema = self.schema.clone();
1689 let record_batch = RecordBatchDecoder::try_new(
1690 &body.into(),
1691 batch,
1692 schema,
1693 &self.dictionaries_by_id,
1694 &version,
1695 )?
1696 .with_projection(self.projection.as_ref().map(|x| x.0.as_ref()))
1697 .with_require_alignment(false)
1698 .with_skip_validation(self.skip_validation.clone())
1699 .read_record_batch()?;
1700 IpcMessage::RecordBatch(record_batch)
1701 }
1702 Message::MessageHeader::DictionaryBatch => {
1703 let dict = message.header_as_dictionary_batch().ok_or_else(|| {
1704 ArrowError::ParseError(
1705 "Failed to parse dictionary batch from message header".to_string(),
1706 )
1707 })?;
1708
1709 let version = message.version();
1710 let dict_values = get_dictionary_values(
1711 &body.into(),
1712 dict,
1713 &self.schema,
1714 &mut self.dictionaries_by_id,
1715 &version,
1716 false,
1717 self.skip_validation.clone(),
1718 )?;
1719
1720 update_dictionaries(
1721 &mut self.dictionaries_by_id,
1722 dict.isDelta(),
1723 dict.id(),
1724 dict_values.clone(),
1725 )?;
1726
1727 IpcMessage::DictionaryBatch {
1728 id: dict.id(),
1729 is_delta: (dict.isDelta()),
1730 values: (dict_values),
1731 }
1732 }
1733 x => {
1734 return Err(ArrowError::ParseError(format!(
1735 "Unsupported message header type in IPC stream: '{x:?}'"
1736 )));
1737 }
1738 };
1739
1740 Ok(Some(ipc_message))
1741 }
1742
1743 pub fn get_ref(&self) -> &R {
1747 self.reader.inner()
1748 }
1749
1750 pub fn get_mut(&mut self) -> &mut R {
1754 self.reader.inner_mut()
1755 }
1756
1757 pub unsafe fn with_skip_validation(mut self, skip_validation: bool) -> Self {
1763 unsafe { self.skip_validation.set(skip_validation) };
1764 self
1765 }
1766}
1767
1768impl<R: Read> Iterator for StreamReader<R> {
1769 type Item = Result<RecordBatch, ArrowError>;
1770
1771 fn next(&mut self) -> Option<Self::Item> {
1772 self.maybe_next().transpose()
1773 }
1774}
1775
1776impl<R: Read> RecordBatchReader for StreamReader<R> {
1777 fn schema(&self) -> SchemaRef {
1778 self.schema.clone()
1779 }
1780}
1781
1782#[derive(Debug)]
1788#[allow(dead_code)]
1789pub(crate) enum IpcMessage {
1790 Schema(arrow_schema::Schema),
1791 RecordBatch(RecordBatch),
1792 DictionaryBatch {
1793 id: i64,
1794 is_delta: bool,
1795 values: ArrayRef,
1796 },
1797}
1798
1799struct MessageReader<R> {
1802 reader: R,
1803 buf: Vec<u8>,
1804}
1805
1806impl<R: Read> MessageReader<R> {
1807 fn new(reader: R) -> Self {
1808 Self {
1809 reader,
1810 buf: Vec::new(),
1811 }
1812 }
1813
1814 fn maybe_next(&mut self) -> Result<Option<(Message::Message<'_>, MutableBuffer)>, ArrowError> {
1825 let meta_len = self.read_meta_len()?;
1826 let Some(meta_len) = meta_len else {
1827 return Ok(None);
1828 };
1829
1830 self.buf.resize(meta_len, 0);
1831 self.reader.read_exact(&mut self.buf)?;
1832
1833 let message = crate::root_as_message(self.buf.as_slice()).map_err(|err| {
1834 ArrowError::ParseError(format!("Unable to get root as message: {err:?}"))
1835 })?;
1836
1837 let mut buf = MutableBuffer::from_len_zeroed(message.bodyLength() as usize);
1838 self.reader.read_exact(&mut buf)?;
1839
1840 Ok(Some((message, buf)))
1841 }
1842
1843 fn inner_mut(&mut self) -> &mut R {
1845 &mut self.reader
1846 }
1847
1848 fn inner(&self) -> &R {
1850 &self.reader
1851 }
1852
1853 pub fn read_meta_len(&mut self) -> Result<Option<usize>, ArrowError> {
1862 let mut meta_len: [u8; 4] = [0; 4];
1863 match self.reader.read_exact(&mut meta_len) {
1864 Ok(_) => {}
1865 Err(e) => {
1866 return if e.kind() == std::io::ErrorKind::UnexpectedEof {
1867 Ok(None)
1871 } else {
1872 Err(ArrowError::from(e))
1873 };
1874 }
1875 };
1876
1877 let meta_len = {
1878 if meta_len == CONTINUATION_MARKER {
1881 self.reader.read_exact(&mut meta_len)?;
1882 }
1883
1884 i32::from_le_bytes(meta_len)
1885 };
1886
1887 if meta_len == 0 {
1888 return Ok(None);
1889 }
1890
1891 let meta_len = usize::try_from(meta_len)
1892 .map_err(|_| ArrowError::ParseError(format!("Invalid metadata length: {meta_len}")))?;
1893
1894 Ok(Some(meta_len))
1895 }
1896}
1897
1898#[cfg(test)]
1899mod tests {
1900 use std::io::Cursor;
1901
1902 use crate::convert::fb_to_schema;
1903 use crate::writer::{
1904 DictionaryTracker, IpcDataGenerator, IpcWriteOptions, unslice_run_array, write_message,
1905 };
1906
1907 use super::*;
1908
1909 use crate::{root_as_footer, root_as_message, size_prefixed_root_as_message};
1910 use arrow_array::builder::{PrimitiveRunBuilder, UnionBuilder};
1911 use arrow_array::types::*;
1912 use arrow_buffer::{NullBuffer, OffsetBuffer};
1913 use arrow_data::ArrayDataBuilder;
1914
1915 fn create_test_projection_schema() -> Schema {
1916 let list_data_type = DataType::List(Arc::new(Field::new_list_field(DataType::Int32, true)));
1918
1919 let fixed_size_list_data_type =
1920 DataType::FixedSizeList(Arc::new(Field::new_list_field(DataType::Int32, false)), 3);
1921
1922 let union_fields = UnionFields::from_fields(vec![
1923 Field::new("a", DataType::Int32, false),
1924 Field::new("b", DataType::Float64, false),
1925 ]);
1926
1927 let union_data_type = DataType::Union(union_fields, UnionMode::Dense);
1928
1929 let struct_fields = Fields::from(vec![
1930 Field::new("id", DataType::Int32, false),
1931 Field::new_list("list", Field::new_list_field(DataType::Int8, true), false),
1932 ]);
1933 let struct_data_type = DataType::Struct(struct_fields);
1934
1935 let run_encoded_data_type = DataType::RunEndEncoded(
1936 Arc::new(Field::new("run_ends", DataType::Int16, false)),
1937 Arc::new(Field::new("values", DataType::Int32, true)),
1938 );
1939
1940 Schema::new(vec![
1942 Field::new("f0", DataType::UInt32, false),
1943 Field::new("f1", DataType::Utf8, false),
1944 Field::new("f2", DataType::Boolean, false),
1945 Field::new("f3", union_data_type, true),
1946 Field::new("f4", DataType::Null, true),
1947 Field::new("f5", DataType::Float64, true),
1948 Field::new("f6", list_data_type, false),
1949 Field::new("f7", DataType::FixedSizeBinary(3), true),
1950 Field::new("f8", fixed_size_list_data_type, false),
1951 Field::new("f9", struct_data_type, false),
1952 Field::new("f10", run_encoded_data_type, false),
1953 Field::new("f11", DataType::Boolean, false),
1954 Field::new_dictionary("f12", DataType::Int8, DataType::Utf8, false),
1955 Field::new("f13", DataType::Utf8, false),
1956 ])
1957 }
1958
1959 fn create_test_projection_batch_data(schema: &Schema) -> RecordBatch {
1960 let array0 = UInt32Array::from(vec![1, 2, 3]);
1962 let array1 = StringArray::from(vec!["foo", "bar", "baz"]);
1963 let array2 = BooleanArray::from(vec![true, false, true]);
1964
1965 let mut union_builder = UnionBuilder::new_dense();
1966 union_builder.append::<Int32Type>("a", 1).unwrap();
1967 union_builder.append::<Float64Type>("b", 10.1).unwrap();
1968 union_builder.append_null::<Float64Type>("b").unwrap();
1969 let array3 = union_builder.build().unwrap();
1970
1971 let array4 = NullArray::new(3);
1972 let array5 = Float64Array::from(vec![Some(1.1), None, Some(3.3)]);
1973 let array6_values = vec![
1974 Some(vec![Some(10), Some(10), Some(10)]),
1975 Some(vec![Some(20), Some(20), Some(20)]),
1976 Some(vec![Some(30), Some(30)]),
1977 ];
1978 let array6 = ListArray::from_iter_primitive::<Int32Type, _, _>(array6_values);
1979 let array7_values = vec![vec![11, 12, 13], vec![22, 23, 24], vec![33, 34, 35]];
1980 let array7 = FixedSizeBinaryArray::try_from_iter(array7_values.into_iter()).unwrap();
1981
1982 let array8_values = ArrayData::builder(DataType::Int32)
1983 .len(9)
1984 .add_buffer(Buffer::from_slice_ref([40, 41, 42, 43, 44, 45, 46, 47, 48]))
1985 .build()
1986 .unwrap();
1987 let array8_data = ArrayData::builder(schema.field(8).data_type().clone())
1988 .len(3)
1989 .add_child_data(array8_values)
1990 .build()
1991 .unwrap();
1992 let array8 = FixedSizeListArray::from(array8_data);
1993
1994 let array9_id: ArrayRef = Arc::new(Int32Array::from(vec![1001, 1002, 1003]));
1995 let array9_list: ArrayRef =
1996 Arc::new(ListArray::from_iter_primitive::<Int8Type, _, _>(vec![
1997 Some(vec![Some(-10)]),
1998 Some(vec![Some(-20), Some(-20), Some(-20)]),
1999 Some(vec![Some(-30)]),
2000 ]));
2001 let array9 = ArrayDataBuilder::new(schema.field(9).data_type().clone())
2002 .add_child_data(array9_id.into_data())
2003 .add_child_data(array9_list.into_data())
2004 .len(3)
2005 .build()
2006 .unwrap();
2007 let array9 = StructArray::from(array9);
2008
2009 let array10_input = vec![Some(1_i32), None, None];
2010 let mut array10_builder = PrimitiveRunBuilder::<Int16Type, Int32Type>::new();
2011 array10_builder.extend(array10_input);
2012 let array10 = array10_builder.finish();
2013
2014 let array11 = BooleanArray::from(vec![false, false, true]);
2015
2016 let array12_values = StringArray::from(vec!["x", "yy", "zzz"]);
2017 let array12_keys = Int8Array::from_iter_values([1, 1, 2]);
2018 let array12 = DictionaryArray::new(array12_keys, Arc::new(array12_values));
2019
2020 let array13 = StringArray::from(vec!["a", "bb", "ccc"]);
2021
2022 RecordBatch::try_new(
2024 Arc::new(schema.clone()),
2025 vec![
2026 Arc::new(array0),
2027 Arc::new(array1),
2028 Arc::new(array2),
2029 Arc::new(array3),
2030 Arc::new(array4),
2031 Arc::new(array5),
2032 Arc::new(array6),
2033 Arc::new(array7),
2034 Arc::new(array8),
2035 Arc::new(array9),
2036 Arc::new(array10),
2037 Arc::new(array11),
2038 Arc::new(array12),
2039 Arc::new(array13),
2040 ],
2041 )
2042 .unwrap()
2043 }
2044
2045 #[test]
2046 fn test_negative_meta_len_start_stream() {
2047 let bytes = i32::to_le_bytes(-1);
2048 let mut buf = vec![];
2049 buf.extend(CONTINUATION_MARKER);
2050 buf.extend(bytes);
2051
2052 let reader_err = StreamReader::try_new(Cursor::new(buf), None).err();
2053 assert!(reader_err.is_some());
2054 assert_eq!(
2055 reader_err.unwrap().to_string(),
2056 "Parser error: Invalid metadata length: -1"
2057 );
2058 }
2059
2060 #[test]
2061 fn test_negative_meta_len_mid_stream() {
2062 let schema = Schema::new(vec![Field::new("a", DataType::Int32, false)]);
2063 let mut buf = Vec::new();
2064 {
2065 let mut writer = crate::writer::StreamWriter::try_new(&mut buf, &schema).unwrap();
2066 let batch =
2067 RecordBatch::try_new(Arc::new(schema), vec![Arc::new(Int32Array::from(vec![1]))])
2068 .unwrap();
2069 writer.write(&batch).unwrap();
2070 }
2071
2072 let bytes = i32::to_le_bytes(-1);
2073 buf.extend(CONTINUATION_MARKER);
2074 buf.extend(bytes);
2075
2076 let mut reader = StreamReader::try_new(Cursor::new(buf), None).unwrap();
2077 assert!(reader.maybe_next().is_ok());
2079 let batch_err = reader.maybe_next().err();
2081 assert!(batch_err.is_some());
2082 assert_eq!(
2083 batch_err.unwrap().to_string(),
2084 "Parser error: Invalid metadata length: -1"
2085 );
2086 }
2087
2088 #[test]
2089 fn test_missing_buffer_metadata_error() {
2090 use crate::r#gen::Message::*;
2091 use flatbuffers::FlatBufferBuilder;
2092
2093 let schema = Arc::new(Schema::new(vec![Field::new("col", DataType::Int32, true)]));
2094
2095 let mut fbb = FlatBufferBuilder::new();
2098 let nodes = fbb.create_vector(&[FieldNode::new(2, 0)]);
2099 let buffers = fbb.create_vector(&[crate::Buffer::new(0, 8)]);
2100 let batch_offset = RecordBatch::create(
2101 &mut fbb,
2102 &RecordBatchArgs {
2103 length: 2,
2104 nodes: Some(nodes),
2105 buffers: Some(buffers),
2106 compression: None,
2107 variadicBufferCounts: None,
2108 },
2109 );
2110 fbb.finish_minimal(batch_offset);
2111 let batch_bytes = fbb.finished_data().to_vec();
2112 let batch = flatbuffers::root::<RecordBatch>(&batch_bytes).unwrap();
2113
2114 let data_buffer = Buffer::from(vec![0u8; 8]);
2115 let dictionaries: HashMap<i64, ArrayRef> = HashMap::new();
2116 let metadata = MetadataVersion::V5;
2117
2118 let decoder = RecordBatchDecoder::try_new(
2119 &data_buffer,
2120 batch,
2121 schema.clone(),
2122 &dictionaries,
2123 &metadata,
2124 )
2125 .unwrap();
2126
2127 let result = decoder.read_record_batch();
2128
2129 match result {
2130 Err(ArrowError::IpcError(msg)) => {
2131 assert_eq!(msg, "Buffer count mismatched with metadata");
2132 }
2133 other => panic!("unexpected error: {other:?}"),
2134 }
2135 }
2136
2137 #[test]
2139 fn test_read_legacy_empty_list_without_offsets_buffer() {
2140 use crate::r#gen::Message::*;
2141 use flatbuffers::FlatBufferBuilder;
2142
2143 let schema = Arc::new(Schema::new(vec![Field::new_list(
2144 "items",
2145 Field::new_list_field(DataType::Int32, true),
2146 true,
2147 )]));
2148
2149 let mut fbb = FlatBufferBuilder::new();
2152 let nodes = fbb.create_vector(&[
2153 FieldNode::new(0, 0), FieldNode::new(0, 0), ]);
2156 let buffers = fbb.create_vector(&[
2157 crate::Buffer::new(0, 0), crate::Buffer::new(0, 0), crate::Buffer::new(0, 0), crate::Buffer::new(0, 0), ]);
2162 let batch_offset = RecordBatch::create(
2163 &mut fbb,
2164 &RecordBatchArgs {
2165 length: 0,
2166 nodes: Some(nodes),
2167 buffers: Some(buffers),
2168 compression: None,
2169 variadicBufferCounts: None,
2170 },
2171 );
2172 fbb.finish_minimal(batch_offset);
2173 let batch_bytes = fbb.finished_data().to_vec();
2174 let batch = flatbuffers::root::<RecordBatch>(&batch_bytes).unwrap();
2175
2176 let body = Buffer::from(Vec::<u8>::new());
2177 let dictionaries: HashMap<i64, ArrayRef> = HashMap::new();
2178 let metadata = MetadataVersion::V5;
2179
2180 let decoder =
2181 RecordBatchDecoder::try_new(&body, batch, schema.clone(), &dictionaries, &metadata)
2182 .unwrap();
2183
2184 let read_batch = decoder.read_record_batch().unwrap();
2185 assert_eq!(read_batch.num_rows(), 0);
2186
2187 let list = read_batch
2188 .column(0)
2189 .as_any()
2190 .downcast_ref::<ListArray>()
2191 .unwrap();
2192 assert_eq!(list.len(), 0);
2193 assert_eq!(list.values().len(), 0);
2194 }
2195
2196 #[test]
2198 fn test_read_legacy_empty_utf8_and_binary_without_offsets_buffer() {
2199 use crate::r#gen::Message::*;
2200 use flatbuffers::FlatBufferBuilder;
2201
2202 let schema = Arc::new(Schema::new(vec![
2203 Field::new("name", DataType::Utf8, true),
2204 Field::new("payload", DataType::Binary, true),
2205 ]));
2206
2207 let mut fbb = FlatBufferBuilder::new();
2210 let nodes = fbb.create_vector(&[
2211 FieldNode::new(0, 0), FieldNode::new(0, 0), ]);
2214 let buffers = fbb.create_vector(&[
2215 crate::Buffer::new(0, 0), crate::Buffer::new(0, 0), crate::Buffer::new(0, 0), crate::Buffer::new(0, 0), crate::Buffer::new(0, 0), crate::Buffer::new(0, 0), ]);
2222 let batch_offset = RecordBatch::create(
2223 &mut fbb,
2224 &RecordBatchArgs {
2225 length: 0,
2226 nodes: Some(nodes),
2227 buffers: Some(buffers),
2228 compression: None,
2229 variadicBufferCounts: None,
2230 },
2231 );
2232 fbb.finish_minimal(batch_offset);
2233 let batch_bytes = fbb.finished_data().to_vec();
2234 let batch = flatbuffers::root::<RecordBatch>(&batch_bytes).unwrap();
2235
2236 let body = Buffer::from(Vec::<u8>::new());
2237 let dictionaries: HashMap<i64, ArrayRef> = HashMap::new();
2238 let metadata = MetadataVersion::V5;
2239
2240 let decoder =
2241 RecordBatchDecoder::try_new(&body, batch, schema.clone(), &dictionaries, &metadata)
2242 .unwrap();
2243
2244 let read_batch = decoder.read_record_batch().unwrap();
2245 assert_eq!(read_batch.num_rows(), 0);
2246
2247 let utf8 = read_batch
2248 .column(0)
2249 .as_any()
2250 .downcast_ref::<StringArray>()
2251 .unwrap();
2252 assert_eq!(utf8.len(), 0);
2253 assert_eq!(utf8.value_offsets(), [0]);
2254
2255 let binary = read_batch
2256 .column(1)
2257 .as_any()
2258 .downcast_ref::<BinaryArray>()
2259 .unwrap();
2260 assert_eq!(binary.len(), 0);
2261 assert_eq!(binary.value_offsets(), [0]);
2262 }
2263
2264 #[test]
2265 fn test_projection_array_values() {
2266 let schema = create_test_projection_schema();
2268
2269 let batch = create_test_projection_batch_data(&schema);
2271
2272 let mut buf = Vec::new();
2274 {
2275 let mut writer = crate::writer::FileWriter::try_new(&mut buf, &schema).unwrap();
2276 writer.write(&batch).unwrap();
2277 writer.finish().unwrap();
2278 }
2279
2280 for index in 0..12 {
2282 let projection = vec![index];
2283 let reader = FileReader::try_new(std::io::Cursor::new(buf.clone()), Some(projection));
2284 let read_batch = reader.unwrap().next().unwrap().unwrap();
2285 let projected_column = read_batch.column(0);
2286 let expected_column = batch.column(index);
2287
2288 assert_eq!(projected_column.as_ref(), expected_column.as_ref());
2290 }
2291
2292 {
2293 let reader =
2295 FileReader::try_new(std::io::Cursor::new(buf.clone()), Some(vec![3, 2, 1]));
2296 let read_batch = reader.unwrap().next().unwrap().unwrap();
2297 let expected_batch = batch.project(&[3, 2, 1]).unwrap();
2298 assert_eq!(read_batch, expected_batch);
2299 }
2300 }
2301
2302 #[test]
2303 fn test_projection_duplicate_indices() {
2304 let schema = create_test_projection_schema();
2305 let batch = create_test_projection_batch_data(&schema);
2306
2307 let mut buf = Vec::new();
2309 {
2310 let mut writer = crate::writer::FileWriter::try_new(&mut buf, &schema).unwrap();
2311 writer.write(&batch).unwrap();
2312 writer.finish().unwrap();
2313 }
2314
2315 for projection in [vec![1, 1], vec![2, 0, 2]] {
2317 let reader =
2318 FileReader::try_new(std::io::Cursor::new(buf.clone()), Some(projection.clone()));
2319 let read_batch = reader.unwrap().next().unwrap().unwrap();
2320
2321 let expected_batch = batch.project(&projection).unwrap();
2322 assert_eq!(read_batch, expected_batch);
2323 }
2324 }
2325
2326 #[test]
2327 fn test_arrow_single_float_row() {
2328 let schema = Schema::new(vec![
2329 Field::new("a", DataType::Float32, false),
2330 Field::new("b", DataType::Float32, false),
2331 Field::new("c", DataType::Int32, false),
2332 Field::new("d", DataType::Int32, false),
2333 ]);
2334 let arrays = vec![
2335 Arc::new(Float32Array::from(vec![1.23])) as ArrayRef,
2336 Arc::new(Float32Array::from(vec![-6.50])) as ArrayRef,
2337 Arc::new(Int32Array::from(vec![2])) as ArrayRef,
2338 Arc::new(Int32Array::from(vec![1])) as ArrayRef,
2339 ];
2340 let batch = RecordBatch::try_new(Arc::new(schema.clone()), arrays).unwrap();
2341 let mut file = tempfile::tempfile().unwrap();
2343 let mut stream_writer = crate::writer::StreamWriter::try_new(&mut file, &schema).unwrap();
2344 stream_writer.write(&batch).unwrap();
2345 stream_writer.finish().unwrap();
2346
2347 drop(stream_writer);
2348
2349 file.rewind().unwrap();
2350
2351 let reader = StreamReader::try_new(&mut file, None).unwrap();
2353
2354 reader.for_each(|batch| {
2355 let batch = batch.unwrap();
2356 assert!(
2357 batch
2358 .column(0)
2359 .as_any()
2360 .downcast_ref::<Float32Array>()
2361 .unwrap()
2362 .value(0)
2363 != 0.0
2364 );
2365 assert!(
2366 batch
2367 .column(1)
2368 .as_any()
2369 .downcast_ref::<Float32Array>()
2370 .unwrap()
2371 .value(0)
2372 != 0.0
2373 );
2374 });
2375
2376 file.rewind().unwrap();
2377
2378 let reader = StreamReader::try_new(file, Some(vec![0, 3])).unwrap();
2380
2381 reader.for_each(|batch| {
2382 let batch = batch.unwrap();
2383 assert_eq!(batch.schema().fields().len(), 2);
2384 assert_eq!(batch.schema().fields()[0].data_type(), &DataType::Float32);
2385 assert_eq!(batch.schema().fields()[1].data_type(), &DataType::Int32);
2386 });
2387 }
2388
2389 fn write_ipc(rb: &RecordBatch) -> Vec<u8> {
2391 let mut buf = Vec::new();
2392 let mut writer = crate::writer::FileWriter::try_new(&mut buf, rb.schema_ref()).unwrap();
2393 writer.write(rb).unwrap();
2394 writer.finish().unwrap();
2395 buf
2396 }
2397
2398 fn read_ipc(buf: &[u8]) -> Result<RecordBatch, ArrowError> {
2400 let mut reader = FileReader::try_new(std::io::Cursor::new(buf), None)?;
2401 reader.next().unwrap()
2402 }
2403
2404 fn read_ipc_skip_validation(buf: &[u8]) -> Result<RecordBatch, ArrowError> {
2407 let mut reader = unsafe {
2408 FileReader::try_new(std::io::Cursor::new(buf), None)?.with_skip_validation(true)
2409 };
2410 reader.next().unwrap()
2411 }
2412
2413 fn roundtrip_ipc(rb: &RecordBatch) -> RecordBatch {
2414 let buf = write_ipc(rb);
2415 read_ipc(&buf).unwrap()
2416 }
2417
2418 fn read_ipc_with_decoder(buf: Vec<u8>) -> Result<RecordBatch, ArrowError> {
2421 read_ipc_with_decoder_inner(buf, false)
2422 }
2423
2424 fn read_ipc_with_decoder_skip_validation(buf: Vec<u8>) -> Result<RecordBatch, ArrowError> {
2427 read_ipc_with_decoder_inner(buf, true)
2428 }
2429
2430 fn read_ipc_with_decoder_inner(
2431 buf: Vec<u8>,
2432 skip_validation: bool,
2433 ) -> Result<RecordBatch, ArrowError> {
2434 let buffer = Buffer::from_vec(buf);
2435 let trailer_start = buffer.len() - 10;
2436 let footer_len = read_footer_length(buffer[trailer_start..].try_into().unwrap())?;
2437 let footer = root_as_footer(&buffer[trailer_start - footer_len..trailer_start])
2438 .map_err(|e| ArrowError::InvalidArgumentError(format!("Invalid footer: {e}")))?;
2439
2440 let schema = fb_to_schema(footer.schema().unwrap());
2441
2442 let mut decoder = unsafe {
2443 FileDecoder::new(Arc::new(schema), footer.version())
2444 .with_skip_validation(skip_validation)
2445 };
2446 for block in footer.dictionaries().iter().flatten() {
2448 let block_len = block.bodyLength() as usize + block.metaDataLength() as usize;
2449 let data = buffer.slice_with_length(block.offset() as _, block_len);
2450 decoder.read_dictionary(block, &data)?
2451 }
2452
2453 let batches = footer.recordBatches().unwrap();
2455 assert_eq!(batches.len(), 1); let block = batches.get(0);
2458 let block_len = block.bodyLength() as usize + block.metaDataLength() as usize;
2459 let data = buffer.slice_with_length(block.offset() as _, block_len);
2460 Ok(decoder.read_record_batch(block, &data)?.unwrap())
2461 }
2462
2463 fn write_stream(rb: &RecordBatch) -> Vec<u8> {
2465 let mut buf = Vec::new();
2466 let mut writer = crate::writer::StreamWriter::try_new(&mut buf, rb.schema_ref()).unwrap();
2467 writer.write(rb).unwrap();
2468 writer.finish().unwrap();
2469 buf
2470 }
2471
2472 fn read_stream(buf: &[u8]) -> Result<RecordBatch, ArrowError> {
2474 let mut reader = StreamReader::try_new(std::io::Cursor::new(buf), None)?;
2475 reader.next().unwrap()
2476 }
2477
2478 fn read_stream_skip_validation(buf: &[u8]) -> Result<RecordBatch, ArrowError> {
2481 let mut reader = unsafe {
2482 StreamReader::try_new(std::io::Cursor::new(buf), None)?.with_skip_validation(true)
2483 };
2484 reader.next().unwrap()
2485 }
2486
2487 fn roundtrip_ipc_stream(rb: &RecordBatch) -> RecordBatch {
2488 let buf = write_stream(rb);
2489 read_stream(&buf).unwrap()
2490 }
2491
2492 #[test]
2493 fn test_roundtrip_with_custom_metadata() {
2494 let schema = Schema::new(vec![Field::new("dummy", DataType::Float64, false)]);
2495 let mut buf = Vec::new();
2496 let mut writer = crate::writer::FileWriter::try_new(&mut buf, &schema).unwrap();
2497 let mut test_metadata = HashMap::new();
2498 test_metadata.insert("abc".to_string(), "abc".to_string());
2499 test_metadata.insert("def".to_string(), "def".to_string());
2500 for (k, v) in &test_metadata {
2501 writer.write_metadata(k, v);
2502 }
2503 writer.finish().unwrap();
2504 drop(writer);
2505
2506 let reader = crate::reader::FileReader::try_new(std::io::Cursor::new(buf), None).unwrap();
2507 assert_eq!(reader.custom_metadata(), &test_metadata);
2508 }
2509
2510 #[test]
2511 fn test_roundtrip_nested_dict() {
2512 let inner: DictionaryArray<Int32Type> = vec!["a", "b", "a"].into_iter().collect();
2513
2514 let array = Arc::new(inner) as ArrayRef;
2515
2516 let dctfield = Arc::new(Field::new("dict", array.data_type().clone(), false));
2517
2518 let s = StructArray::from(vec![(dctfield, array)]);
2519 let struct_array = Arc::new(s) as ArrayRef;
2520
2521 let schema = Arc::new(Schema::new(vec![Field::new(
2522 "struct",
2523 struct_array.data_type().clone(),
2524 false,
2525 )]));
2526
2527 let batch = RecordBatch::try_new(schema, vec![struct_array]).unwrap();
2528
2529 assert_eq!(batch, roundtrip_ipc(&batch));
2530 }
2531
2532 #[test]
2533 fn test_roundtrip_nested_dict_no_preserve_dict_id() {
2534 let inner: DictionaryArray<Int32Type> = vec!["a", "b", "a"].into_iter().collect();
2535
2536 let array = Arc::new(inner) as ArrayRef;
2537
2538 let dctfield = Arc::new(Field::new("dict", array.data_type().clone(), false));
2539
2540 let s = StructArray::from(vec![(dctfield, array)]);
2541 let struct_array = Arc::new(s) as ArrayRef;
2542
2543 let schema = Arc::new(Schema::new(vec![Field::new(
2544 "struct",
2545 struct_array.data_type().clone(),
2546 false,
2547 )]));
2548
2549 let batch = RecordBatch::try_new(schema, vec![struct_array]).unwrap();
2550
2551 let mut buf = Vec::new();
2552 let mut writer = crate::writer::FileWriter::try_new_with_options(
2553 &mut buf,
2554 batch.schema_ref(),
2555 IpcWriteOptions::default(),
2556 )
2557 .unwrap();
2558 writer.write(&batch).unwrap();
2559 writer.finish().unwrap();
2560 drop(writer);
2561
2562 let mut reader = FileReader::try_new(std::io::Cursor::new(buf), None).unwrap();
2563
2564 assert_eq!(batch, reader.next().unwrap().unwrap());
2565 }
2566
2567 fn check_union_with_builder(mut builder: UnionBuilder) {
2568 builder.append::<Int32Type>("a", 1).unwrap();
2569 builder.append_null::<Int32Type>("a").unwrap();
2570 builder.append::<Float64Type>("c", 3.0).unwrap();
2571 builder.append::<Int32Type>("a", 4).unwrap();
2572 builder.append::<Int64Type>("d", 11).unwrap();
2573 let union = builder.build().unwrap();
2574
2575 let schema = Arc::new(Schema::new(vec![Field::new(
2576 "union",
2577 union.data_type().clone(),
2578 false,
2579 )]));
2580
2581 let union_array = Arc::new(union) as ArrayRef;
2582
2583 let rb = RecordBatch::try_new(schema, vec![union_array]).unwrap();
2584 let rb2 = roundtrip_ipc(&rb);
2585 assert_eq!(rb.schema(), rb2.schema());
2588 assert_eq!(rb.num_columns(), rb2.num_columns());
2589 assert_eq!(rb.num_rows(), rb2.num_rows());
2590 let union1 = rb.column(0);
2591 let union2 = rb2.column(0);
2592
2593 assert_eq!(union1, union2);
2594 }
2595
2596 #[test]
2597 fn test_roundtrip_dense_union() {
2598 check_union_with_builder(UnionBuilder::new_dense());
2599 }
2600
2601 #[test]
2602 fn test_roundtrip_sparse_union() {
2603 check_union_with_builder(UnionBuilder::new_sparse());
2604 }
2605
2606 #[test]
2607 fn test_roundtrip_struct_empty_fields() {
2608 let nulls = NullBuffer::from(&[true, true, false]);
2609 let rb = RecordBatch::try_from_iter([(
2610 "",
2611 Arc::new(StructArray::new_empty_fields(nulls.len(), Some(nulls))) as _,
2612 )])
2613 .unwrap();
2614 let rb2 = roundtrip_ipc(&rb);
2615 assert_eq!(rb, rb2);
2616 }
2617
2618 #[test]
2619 fn test_roundtrip_stream_run_array_sliced() {
2620 let run_array_1: Int32RunArray = vec!["a", "a", "a", "b", "b", "c", "c", "c"]
2621 .into_iter()
2622 .collect();
2623 let run_array_1_sliced = run_array_1.slice(2, 5);
2624
2625 let run_array_2_inupt = vec![Some(1_i32), None, None, Some(2), Some(2)];
2626 let mut run_array_2_builder = PrimitiveRunBuilder::<Int16Type, Int32Type>::new();
2627 run_array_2_builder.extend(run_array_2_inupt);
2628 let run_array_2 = run_array_2_builder.finish();
2629
2630 let schema = Arc::new(Schema::new(vec![
2631 Field::new(
2632 "run_array_1_sliced",
2633 run_array_1_sliced.data_type().clone(),
2634 false,
2635 ),
2636 Field::new("run_array_2", run_array_2.data_type().clone(), false),
2637 ]));
2638 let input_batch = RecordBatch::try_new(
2639 schema,
2640 vec![Arc::new(run_array_1_sliced.clone()), Arc::new(run_array_2)],
2641 )
2642 .unwrap();
2643 let output_batch = roundtrip_ipc_stream(&input_batch);
2644
2645 assert_eq!(input_batch.column(1), output_batch.column(1));
2649
2650 let run_array_1_unsliced = unslice_run_array(run_array_1_sliced.into_data()).unwrap();
2651 assert_eq!(run_array_1_unsliced, output_batch.column(0).into_data());
2652 }
2653
2654 #[test]
2655 fn test_roundtrip_stream_nested_dict() {
2656 let xs = vec!["AA", "BB", "AA", "CC", "BB"];
2657 let dict = Arc::new(
2658 xs.clone()
2659 .into_iter()
2660 .collect::<DictionaryArray<Int8Type>>(),
2661 );
2662 let string_array: ArrayRef = Arc::new(StringArray::from(xs.clone()));
2663 let struct_array = StructArray::from(vec![
2664 (
2665 Arc::new(Field::new("f2.1", DataType::Utf8, false)),
2666 string_array,
2667 ),
2668 (
2669 Arc::new(Field::new("f2.2_struct", dict.data_type().clone(), false)),
2670 dict.clone() as ArrayRef,
2671 ),
2672 ]);
2673 let schema = Arc::new(Schema::new(vec![
2674 Field::new("f1_string", DataType::Utf8, false),
2675 Field::new("f2_struct", struct_array.data_type().clone(), false),
2676 ]));
2677 let input_batch = RecordBatch::try_new(
2678 schema,
2679 vec![
2680 Arc::new(StringArray::from(xs.clone())),
2681 Arc::new(struct_array),
2682 ],
2683 )
2684 .unwrap();
2685 let output_batch = roundtrip_ipc_stream(&input_batch);
2686 assert_eq!(input_batch, output_batch);
2687 }
2688
2689 #[test]
2690 fn test_ipc_writers_reject_dictionary_of_dictionary_schema() {
2691 let values = Arc::new(StringArray::from(vec![Some("a"), Some("b")])) as ArrayRef;
2692 let inner = Arc::new(DictionaryArray::new(
2693 UInt32Array::from_iter_values([0, 1]),
2694 values,
2695 )) as ArrayRef;
2696 let outer = Arc::new(DictionaryArray::new(
2697 UInt32Array::from_iter_values([0, 1, 0]),
2698 inner,
2699 )) as ArrayRef;
2700
2701 let schema = Arc::new(Schema::new(vec![Field::new(
2702 "f1",
2703 outer.data_type().clone(),
2704 false,
2705 )]));
2706 let batch = RecordBatch::try_new(schema, vec![outer]).unwrap();
2707
2708 let mut stream = Vec::new();
2709 let Err(err) = crate::writer::StreamWriter::try_new(&mut stream, batch.schema_ref()) else {
2710 panic!("IPC stream writer should reject dictionary-of-dictionary schemas");
2711 };
2712 assert!(stream.is_empty());
2713
2714 assert!(
2715 err.to_string().contains("dictionary-of-dictionary values"),
2716 "unexpected error: {err}"
2717 );
2718
2719 let mut file = Vec::new();
2720 let Err(err) = crate::writer::FileWriter::try_new(&mut file, batch.schema_ref()) else {
2721 panic!("IPC file writer should reject dictionary-of-dictionary schemas");
2722 };
2723 assert!(file.is_empty());
2724
2725 assert!(
2726 err.to_string().contains("dictionary-of-dictionary values"),
2727 "unexpected error: {err}"
2728 );
2729 }
2730
2731 #[test]
2732 fn test_roundtrip_stream_nested_dict_of_map_of_dict() {
2733 let values = StringArray::from(vec![Some("a"), None, Some("b"), Some("c")]);
2734 let values = Arc::new(values) as ArrayRef;
2735 let value_dict_keys = Int8Array::from_iter_values([0, 1, 1, 2, 3, 1]);
2736 let value_dict_array = DictionaryArray::new(value_dict_keys, values.clone());
2737
2738 let key_dict_keys = Int8Array::from_iter_values([0, 0, 2, 2, 2, 3]);
2739 let key_dict_array = DictionaryArray::new(key_dict_keys, values);
2740
2741 #[allow(deprecated)]
2742 let keys_field = Arc::new(Field::new_dict(
2743 Field::MAP_KEY_FIELD_DEFAULT_NAME,
2744 DataType::Dictionary(Box::new(DataType::Int8), Box::new(DataType::Utf8)),
2745 false,
2746 1,
2747 false,
2748 ));
2749 #[allow(deprecated)]
2750 let values_field = Arc::new(Field::new_dict(
2751 Field::MAP_VALUE_FIELD_DEFAULT_NAME,
2752 DataType::Dictionary(Box::new(DataType::Int8), Box::new(DataType::Utf8)),
2753 true,
2754 2,
2755 false,
2756 ));
2757 let entry_struct = StructArray::from(vec![
2758 (keys_field, make_array(key_dict_array.into_data())),
2759 (values_field, make_array(value_dict_array.into_data())),
2760 ]);
2761 let map_data_type = DataType::Map(
2762 Arc::new(Field::new(
2763 Field::MAP_ENTRIES_FIELD_DEFAULT_NAME,
2764 entry_struct.data_type().clone(),
2765 false,
2766 )),
2767 false,
2768 );
2769
2770 let entry_offsets = Buffer::from_slice_ref([0, 2, 4, 6]);
2771 let map_data = ArrayData::builder(map_data_type)
2772 .len(3)
2773 .add_buffer(entry_offsets)
2774 .add_child_data(entry_struct.into_data())
2775 .build()
2776 .unwrap();
2777 let map_array = MapArray::from(map_data);
2778
2779 let dict_keys = Int8Array::from_iter_values([0, 1, 1, 2, 2, 1]);
2780 let dict_dict_array = DictionaryArray::new(dict_keys, Arc::new(map_array));
2781
2782 let schema = Arc::new(Schema::new(vec![Field::new(
2783 "f1",
2784 dict_dict_array.data_type().clone(),
2785 false,
2786 )]));
2787 let input_batch = RecordBatch::try_new(schema, vec![Arc::new(dict_dict_array)]).unwrap();
2788 let output_batch = roundtrip_ipc_stream(&input_batch);
2789 assert_eq!(input_batch, output_batch);
2790 }
2791
2792 fn test_roundtrip_stream_dict_of_list_of_dict_impl<
2793 OffsetSize: OffsetSizeTrait,
2794 U: ArrowNativeType,
2795 >(
2796 list_data_type: DataType,
2797 offsets: &[U; 5],
2798 ) {
2799 let values = StringArray::from(vec![Some("a"), None, Some("c"), None]);
2800 let keys = Int8Array::from_iter_values([0, 0, 1, 2, 0, 1, 3]);
2801 let dict_array = DictionaryArray::new(keys, Arc::new(values));
2802 let dict_data = dict_array.to_data();
2803
2804 let value_offsets = Buffer::from_slice_ref(offsets);
2805
2806 let list_data = ArrayData::builder(list_data_type)
2807 .len(4)
2808 .add_buffer(value_offsets)
2809 .add_child_data(dict_data)
2810 .build()
2811 .unwrap();
2812 let list_array = GenericListArray::<OffsetSize>::from(list_data);
2813
2814 let keys_for_dict = Int8Array::from_iter_values([0, 3, 0, 1, 1, 2, 0, 1, 3]);
2815 let dict_dict_array = DictionaryArray::new(keys_for_dict, Arc::new(list_array));
2816
2817 let schema = Arc::new(Schema::new(vec![Field::new(
2818 "f1",
2819 dict_dict_array.data_type().clone(),
2820 false,
2821 )]));
2822 let input_batch = RecordBatch::try_new(schema, vec![Arc::new(dict_dict_array)]).unwrap();
2823 let output_batch = roundtrip_ipc_stream(&input_batch);
2824 assert_eq!(input_batch, output_batch);
2825 }
2826
2827 #[test]
2828 fn test_roundtrip_stream_dict_of_list_of_dict() {
2829 #[allow(deprecated)]
2831 let list_data_type = DataType::List(Arc::new(Field::new_dict(
2832 "item",
2833 DataType::Dictionary(Box::new(DataType::Int8), Box::new(DataType::Utf8)),
2834 true,
2835 1,
2836 false,
2837 )));
2838 let offsets: &[i32; 5] = &[0, 2, 4, 4, 6];
2839 test_roundtrip_stream_dict_of_list_of_dict_impl::<i32, i32>(list_data_type, offsets);
2840
2841 #[allow(deprecated)]
2843 let list_data_type = DataType::LargeList(Arc::new(Field::new_dict(
2844 "item",
2845 DataType::Dictionary(Box::new(DataType::Int8), Box::new(DataType::Utf8)),
2846 true,
2847 1,
2848 false,
2849 )));
2850 let offsets: &[i64; 5] = &[0, 2, 4, 4, 7];
2851 test_roundtrip_stream_dict_of_list_of_dict_impl::<i64, i64>(list_data_type, offsets);
2852 }
2853
2854 #[test]
2855 fn test_roundtrip_stream_dict_of_fixed_size_list_of_dict() {
2856 let values = StringArray::from(vec![Some("a"), None, Some("c"), None]);
2857 let keys = Int8Array::from_iter_values([0, 0, 1, 2, 0, 1, 3, 1, 2]);
2858 let dict_array = DictionaryArray::new(keys, Arc::new(values));
2859 let dict_data = dict_array.into_data();
2860
2861 #[allow(deprecated)]
2862 let list_data_type = DataType::FixedSizeList(
2863 Arc::new(Field::new_dict(
2864 "item",
2865 DataType::Dictionary(Box::new(DataType::Int8), Box::new(DataType::Utf8)),
2866 true,
2867 1,
2868 false,
2869 )),
2870 3,
2871 );
2872 let list_data = ArrayData::builder(list_data_type)
2873 .len(3)
2874 .add_child_data(dict_data)
2875 .build()
2876 .unwrap();
2877 let list_array = FixedSizeListArray::from(list_data);
2878
2879 let keys_for_dict = Int8Array::from_iter_values([0, 1, 0, 1, 1, 2, 0, 1, 2]);
2880 let dict_dict_array = DictionaryArray::new(keys_for_dict, Arc::new(list_array));
2881
2882 let schema = Arc::new(Schema::new(vec![Field::new(
2883 "f1",
2884 dict_dict_array.data_type().clone(),
2885 false,
2886 )]));
2887 let input_batch = RecordBatch::try_new(schema, vec![Arc::new(dict_dict_array)]).unwrap();
2888 let output_batch = roundtrip_ipc_stream(&input_batch);
2889 assert_eq!(input_batch, output_batch);
2890 }
2891
2892 const LONG_TEST_STRING: &str =
2893 "This is a long string to make sure binary view array handles it";
2894
2895 #[test]
2896 fn test_roundtrip_view_types() {
2897 let schema = Schema::new(vec![
2898 Field::new("field_1", DataType::BinaryView, true),
2899 Field::new("field_2", DataType::Utf8, true),
2900 Field::new("field_3", DataType::Utf8View, true),
2901 ]);
2902 let bin_values: Vec<Option<&[u8]>> = vec![
2903 Some(b"foo"),
2904 None,
2905 Some(b"bar"),
2906 Some(LONG_TEST_STRING.as_bytes()),
2907 ];
2908 let utf8_values: Vec<Option<&str>> =
2909 vec![Some("foo"), None, Some("bar"), Some(LONG_TEST_STRING)];
2910 let bin_view_array = BinaryViewArray::from_iter(bin_values);
2911 let utf8_array = StringArray::from_iter(utf8_values.iter());
2912 let utf8_view_array = StringViewArray::from_iter(utf8_values);
2913 let record_batch = RecordBatch::try_new(
2914 Arc::new(schema.clone()),
2915 vec![
2916 Arc::new(bin_view_array),
2917 Arc::new(utf8_array),
2918 Arc::new(utf8_view_array),
2919 ],
2920 )
2921 .unwrap();
2922
2923 assert_eq!(record_batch, roundtrip_ipc(&record_batch));
2924 assert_eq!(record_batch, roundtrip_ipc_stream(&record_batch));
2925
2926 let sliced_batch = record_batch.slice(1, 2);
2927 assert_eq!(sliced_batch, roundtrip_ipc(&sliced_batch));
2928 assert_eq!(sliced_batch, roundtrip_ipc_stream(&sliced_batch));
2929 }
2930
2931 #[test]
2932 fn test_roundtrip_view_types_nested_dict() {
2933 let bin_values: Vec<Option<&[u8]>> = vec![
2934 Some(b"foo"),
2935 None,
2936 Some(b"bar"),
2937 Some(LONG_TEST_STRING.as_bytes()),
2938 Some(b"field"),
2939 ];
2940 let utf8_values: Vec<Option<&str>> = vec![
2941 Some("foo"),
2942 None,
2943 Some("bar"),
2944 Some(LONG_TEST_STRING),
2945 Some("field"),
2946 ];
2947 let bin_view_array = Arc::new(BinaryViewArray::from_iter(bin_values));
2948 let utf8_view_array = Arc::new(StringViewArray::from_iter(utf8_values));
2949
2950 let key_dict_keys = Int8Array::from_iter_values([0, 0, 2, 2, 0, 2, 3]);
2951 let key_dict_array = DictionaryArray::new(key_dict_keys, utf8_view_array.clone());
2952 #[allow(deprecated)]
2953 let keys_field = Arc::new(Field::new_dict(
2954 Field::MAP_KEY_FIELD_DEFAULT_NAME,
2955 DataType::Dictionary(Box::new(DataType::Int8), Box::new(DataType::Utf8View)),
2956 false,
2957 1,
2958 false,
2959 ));
2960
2961 let value_dict_keys = Int8Array::from_iter_values([0, 3, 0, 1, 2, 0, 1]);
2962 let value_dict_array = DictionaryArray::new(value_dict_keys, bin_view_array);
2963 #[allow(deprecated)]
2964 let values_field = Arc::new(Field::new_dict(
2965 Field::MAP_VALUE_FIELD_DEFAULT_NAME,
2966 DataType::Dictionary(Box::new(DataType::Int8), Box::new(DataType::BinaryView)),
2967 true,
2968 2,
2969 false,
2970 ));
2971 let entry_struct = StructArray::from(vec![
2972 (keys_field, make_array(key_dict_array.into_data())),
2973 (values_field, make_array(value_dict_array.into_data())),
2974 ]);
2975
2976 let map_data_type = DataType::Map(
2977 Arc::new(Field::new(
2978 Field::MAP_ENTRIES_FIELD_DEFAULT_NAME,
2979 entry_struct.data_type().clone(),
2980 false,
2981 )),
2982 false,
2983 );
2984 let entry_offsets = Buffer::from_slice_ref([0, 2, 4, 7]);
2985 let map_data = ArrayData::builder(map_data_type)
2986 .len(3)
2987 .add_buffer(entry_offsets)
2988 .add_child_data(entry_struct.into_data())
2989 .build()
2990 .unwrap();
2991 let map_array = MapArray::from(map_data);
2992
2993 let dict_keys = Int8Array::from_iter_values([0, 1, 0, 1, 1, 2, 0, 1, 2]);
2994 let dict_dict_array = DictionaryArray::new(dict_keys, Arc::new(map_array));
2995 let schema = Arc::new(Schema::new(vec![Field::new(
2996 "f1",
2997 dict_dict_array.data_type().clone(),
2998 false,
2999 )]));
3000 let batch = RecordBatch::try_new(schema, vec![Arc::new(dict_dict_array)]).unwrap();
3001 assert_eq!(batch, roundtrip_ipc(&batch));
3002 assert_eq!(batch, roundtrip_ipc_stream(&batch));
3003
3004 let sliced_batch = batch.slice(1, 2);
3005 assert_eq!(sliced_batch, roundtrip_ipc(&sliced_batch));
3006 assert_eq!(sliced_batch, roundtrip_ipc_stream(&sliced_batch));
3007 }
3008
3009 #[test]
3010 fn test_no_columns_batch() {
3011 let schema = Arc::new(Schema::empty());
3012 let options = RecordBatchOptions::new()
3013 .with_match_field_names(true)
3014 .with_row_count(Some(10));
3015 let input_batch = RecordBatch::try_new_with_options(schema, vec![], &options).unwrap();
3016 let output_batch = roundtrip_ipc_stream(&input_batch);
3017 assert_eq!(input_batch, output_batch);
3018 }
3019
3020 #[test]
3021 fn test_unaligned() {
3022 let batch = RecordBatch::try_from_iter(vec![(
3023 "i32",
3024 Arc::new(Int32Array::from(vec![1, 2, 3, 4])) as _,
3025 )])
3026 .unwrap();
3027
3028 let r#gen = IpcDataGenerator {};
3029 let mut dict_tracker = DictionaryTracker::new(false);
3030 let (_, encoded) = r#gen
3031 .encode(
3032 &batch,
3033 &mut dict_tracker,
3034 &Default::default(),
3035 &mut Default::default(),
3036 )
3037 .unwrap();
3038
3039 let message = root_as_message(&encoded.ipc_message).unwrap();
3040
3041 let mut buffer = MutableBuffer::with_capacity(encoded.arrow_data.len() + 1);
3043 buffer.push(0_u8);
3044 buffer.extend_from_slice(&encoded.arrow_data);
3045 let b = Buffer::from(buffer).slice(1);
3046 assert_ne!(b.as_ptr().align_offset(8), 0);
3047
3048 let ipc_batch = message.header_as_record_batch().unwrap();
3049 let roundtrip = RecordBatchDecoder::try_new(
3050 &b,
3051 ipc_batch,
3052 batch.schema(),
3053 &Default::default(),
3054 &message.version(),
3055 )
3056 .unwrap()
3057 .with_require_alignment(false)
3058 .read_record_batch()
3059 .unwrap();
3060 assert_eq!(batch, roundtrip);
3061 }
3062
3063 #[test]
3064 fn test_unaligned_throws_error_with_require_alignment() {
3065 let batch = RecordBatch::try_from_iter(vec![(
3066 "i32",
3067 Arc::new(Int32Array::from(vec![1, 2, 3, 4])) as _,
3068 )])
3069 .unwrap();
3070
3071 let r#gen = IpcDataGenerator {};
3072 let mut dict_tracker = DictionaryTracker::new(false);
3073 let (_, encoded) = r#gen
3074 .encode(
3075 &batch,
3076 &mut dict_tracker,
3077 &Default::default(),
3078 &mut Default::default(),
3079 )
3080 .unwrap();
3081
3082 let message = root_as_message(&encoded.ipc_message).unwrap();
3083
3084 let mut buffer = MutableBuffer::with_capacity(encoded.arrow_data.len() + 1);
3086 buffer.push(0_u8);
3087 buffer.extend_from_slice(&encoded.arrow_data);
3088 let b = Buffer::from(buffer).slice(1);
3089 assert_ne!(b.as_ptr().align_offset(8), 0);
3090
3091 let ipc_batch = message.header_as_record_batch().unwrap();
3092 let result = RecordBatchDecoder::try_new(
3093 &b,
3094 ipc_batch,
3095 batch.schema(),
3096 &Default::default(),
3097 &message.version(),
3098 )
3099 .unwrap()
3100 .with_require_alignment(true)
3101 .read_record_batch();
3102
3103 let error = result.unwrap_err();
3104 assert_eq!(
3105 error.to_string(),
3106 "Invalid argument error: Misaligned buffers[0] in array of type Int32, \
3107 offset from expected alignment of 4 by 1"
3108 );
3109 }
3110
3111 #[test]
3120 fn test_misaligned_buffers_error() {
3121 const OFFSET_RANGE: usize = 128;
3122
3123 let cases: Vec<(ArrayRef, usize)> = vec![
3127 (Arc::new(Int32Array::from_iter(0i32..100)) as _, 4),
3128 (Arc::new(Int64Array::from_iter(0i64..100)) as _, 8),
3129 (
3130 Arc::new(StringArray::from_iter_values(
3131 (0..100).map(|i| i.to_string()),
3132 )) as _,
3133 4,
3134 ),
3135 (
3136 Arc::new(LargeStringArray::from_iter_values(
3137 (0..100).map(|i| i.to_string()),
3138 )) as _,
3139 8,
3140 ),
3141 ];
3142
3143 for (array, alignment) in cases {
3144 let batch = RecordBatch::try_from_iter(vec![("col", Arc::clone(&array))]).unwrap();
3145 let encoder = IpcDataGenerator {};
3146 let mut dict_tracker = DictionaryTracker::new(false);
3147 let (_, encoded) = encoder
3148 .encode(
3149 &batch,
3150 &mut dict_tracker,
3151 &Default::default(),
3152 &mut Default::default(),
3153 )
3154 .unwrap();
3155 let message = root_as_message(&encoded.ipc_message).unwrap();
3156 let ipc_batch = message.header_as_record_batch().unwrap();
3157
3158 for offset in 0..OFFSET_RANGE {
3159 let mut storage = MutableBuffer::with_capacity(encoded.arrow_data.len() + offset);
3165 for _ in 0..offset {
3166 storage.push(0_u8);
3167 }
3168 storage.extend_from_slice(&encoded.arrow_data);
3169 let buf = Buffer::from(storage).slice(offset);
3170
3171 let result = RecordBatchDecoder::try_new(
3172 &buf,
3173 ipc_batch,
3174 batch.schema(),
3175 &Default::default(),
3176 &message.version(),
3177 )
3178 .unwrap()
3179 .with_require_alignment(true)
3180 .read_record_batch();
3181
3182 if offset % alignment == 0 {
3183 assert!(
3184 result.is_ok(),
3185 "type={} offset={offset}: expected Ok but got {:?}",
3186 array.data_type(),
3187 result.unwrap_err(),
3188 );
3189 } else {
3190 let err = result
3191 .expect_err(&format!(
3192 "type={} offset={offset}: expected Err for misaligned buffer",
3193 array.data_type()
3194 ))
3195 .to_string();
3196 assert!(
3197 err.contains("Misaligned buffers"),
3198 "type={} offset={offset}: unexpected error: {err}",
3199 array.data_type(),
3200 );
3201 }
3202 }
3203 }
3204 }
3205
3206 #[test]
3207 fn test_file_with_massive_column_count() {
3208 let limit = 600_000;
3210
3211 let fields = (0..limit)
3212 .map(|i| Field::new(format!("{i}"), DataType::Boolean, false))
3213 .collect::<Vec<_>>();
3214 let schema = Arc::new(Schema::new(fields));
3215 let batch = RecordBatch::new_empty(schema);
3216
3217 let mut buf = Vec::new();
3218 let mut writer = crate::writer::FileWriter::try_new(&mut buf, batch.schema_ref()).unwrap();
3219 writer.write(&batch).unwrap();
3220 writer.finish().unwrap();
3221 drop(writer);
3222
3223 let mut reader = FileReaderBuilder::new()
3224 .with_max_footer_fb_tables(1_500_000)
3225 .build(std::io::Cursor::new(buf))
3226 .unwrap();
3227 let roundtrip_batch = reader.next().unwrap().unwrap();
3228
3229 assert_eq!(batch, roundtrip_batch);
3230 }
3231
3232 #[test]
3233 fn test_file_with_deeply_nested_columns() {
3234 let limit = 61;
3236
3237 let fields = (0..limit).fold(
3238 vec![Field::new("leaf", DataType::Boolean, false)],
3239 |field, index| vec![Field::new_struct(format!("{index}"), field, false)],
3240 );
3241 let schema = Arc::new(Schema::new(fields));
3242 let batch = RecordBatch::new_empty(schema);
3243
3244 let mut buf = Vec::new();
3245 let mut writer = crate::writer::FileWriter::try_new(&mut buf, batch.schema_ref()).unwrap();
3246 writer.write(&batch).unwrap();
3247 writer.finish().unwrap();
3248 drop(writer);
3249
3250 let mut reader = FileReaderBuilder::new()
3251 .with_max_footer_fb_depth(65)
3252 .build(std::io::Cursor::new(buf))
3253 .unwrap();
3254 let roundtrip_batch = reader.next().unwrap().unwrap();
3255
3256 assert_eq!(batch, roundtrip_batch);
3257 }
3258
3259 #[test]
3260 fn test_invalid_struct_array_ipc_read_errors() {
3261 let a_field = Field::new("a", DataType::Int32, false);
3262 let b_field = Field::new("b", DataType::Int32, false);
3263 let struct_fields = Fields::from(vec![a_field.clone(), b_field.clone()]);
3264
3265 let a_array_data = ArrayData::builder(a_field.data_type().clone())
3266 .len(4)
3267 .add_buffer(Buffer::from_slice_ref([1, 2, 3, 4]))
3268 .build()
3269 .unwrap();
3270 let b_array_data = ArrayData::builder(b_field.data_type().clone())
3271 .len(3)
3272 .add_buffer(Buffer::from_slice_ref([5, 6, 7]))
3273 .build()
3274 .unwrap();
3275
3276 let invalid_struct_arr = unsafe {
3277 StructArray::new_unchecked(
3278 struct_fields,
3279 vec![make_array(a_array_data), make_array(b_array_data)],
3280 None,
3281 )
3282 };
3283
3284 expect_ipc_validation_error(
3285 Arc::new(invalid_struct_arr),
3286 "Invalid argument error: Incorrect array length for StructArray field \"b\", expected 4 got 3",
3287 );
3288 }
3289
3290 #[test]
3291 fn test_invalid_nested_array_ipc_read_errors() {
3292 let a_field = Field::new("a", DataType::Int32, false);
3294 let b_field = Field::new("b", DataType::Utf8, false);
3295
3296 let schema = Arc::new(Schema::new(vec![Field::new_struct(
3297 "s",
3298 vec![a_field.clone(), b_field.clone()],
3299 false,
3300 )]));
3301
3302 let a_array_data = ArrayData::builder(a_field.data_type().clone())
3303 .len(4)
3304 .add_buffer(Buffer::from_slice_ref([1, 2, 3, 4]))
3305 .build()
3306 .unwrap();
3307 let b_array_data = {
3309 let valid: &[u8] = b" ";
3310 let mut invalid = vec![];
3311 invalid.extend_from_slice(b"ValidString");
3312 invalid.extend_from_slice(INVALID_UTF8_FIRST_CHAR);
3313 let binary_array =
3314 BinaryArray::from_iter(vec![None, Some(valid), None, Some(&invalid)]);
3315 let array = unsafe {
3316 StringArray::new_unchecked(
3317 binary_array.offsets().clone(),
3318 binary_array.values().clone(),
3319 binary_array.nulls().cloned(),
3320 )
3321 };
3322 array.into_data()
3323 };
3324 let struct_data_type = schema.field(0).data_type();
3325
3326 let invalid_struct_arr = unsafe {
3327 make_array(
3328 ArrayData::builder(struct_data_type.clone())
3329 .len(4)
3330 .add_child_data(a_array_data)
3331 .add_child_data(b_array_data)
3332 .build_unchecked(),
3333 )
3334 };
3335 expect_ipc_validation_error(
3336 invalid_struct_arr,
3337 "Invalid argument error: Invalid UTF8 sequence at string index 3 (3..18): invalid utf-8 sequence of 1 bytes from index 11",
3338 );
3339 }
3340
3341 #[test]
3342 fn test_same_dict_id_without_preserve() {
3343 let batch = RecordBatch::try_new(
3344 Arc::new(Schema::new(
3345 ["a", "b"]
3346 .iter()
3347 .map(|name| {
3348 #[allow(deprecated)]
3349 Field::new_dict(
3350 name.to_string(),
3351 DataType::Dictionary(
3352 Box::new(DataType::Int32),
3353 Box::new(DataType::Utf8),
3354 ),
3355 true,
3356 0,
3357 false,
3358 )
3359 })
3360 .collect::<Vec<Field>>(),
3361 )),
3362 vec![
3363 Arc::new(
3364 vec![Some("c"), Some("d")]
3365 .into_iter()
3366 .collect::<DictionaryArray<Int32Type>>(),
3367 ) as ArrayRef,
3368 Arc::new(
3369 vec![Some("e"), Some("f")]
3370 .into_iter()
3371 .collect::<DictionaryArray<Int32Type>>(),
3372 ) as ArrayRef,
3373 ],
3374 )
3375 .expect("Failed to create RecordBatch");
3376
3377 let mut buf = vec![];
3379 {
3380 let mut writer = crate::writer::StreamWriter::try_new_with_options(
3381 &mut buf,
3382 batch.schema().as_ref(),
3383 crate::writer::IpcWriteOptions::default(),
3384 )
3385 .expect("Failed to create StreamWriter");
3386 writer.write(&batch).expect("Failed to write RecordBatch");
3387 writer.finish().expect("Failed to finish StreamWriter");
3388 }
3389
3390 StreamReader::try_new(std::io::Cursor::new(buf), None)
3391 .expect("Failed to create StreamReader")
3392 .for_each(|decoded_batch| {
3393 assert_eq!(decoded_batch.expect("Failed to read RecordBatch"), batch);
3394 });
3395 }
3396
3397 #[test]
3398 fn test_validation_of_invalid_list_array() {
3399 let array = unsafe {
3401 let values = Int32Array::from(vec![1, 2, 3]);
3402 let bad_offsets = ScalarBuffer::<i32>::from(vec![0, 2, 4, 2]); let offsets = OffsetBuffer::new_unchecked(bad_offsets); let field = Field::new_list_field(DataType::Int32, true);
3405 let nulls = None;
3406 ListArray::new(Arc::new(field), offsets, Arc::new(values), nulls)
3407 };
3408
3409 expect_ipc_validation_error(
3410 Arc::new(array),
3411 "Invalid argument error: Offset invariant failure: offset at position 2 out of bounds: 4 > 2",
3412 );
3413 }
3414
3415 #[test]
3416 fn test_validation_of_invalid_string_array() {
3417 let valid: &[u8] = b" ";
3418 let mut invalid = vec![];
3419 invalid.extend_from_slice(b"ThisStringIsCertainlyLongerThan12Bytes");
3420 invalid.extend_from_slice(INVALID_UTF8_FIRST_CHAR);
3421 let binary_array = BinaryArray::from_iter(vec![None, Some(valid), None, Some(&invalid)]);
3422 let array = unsafe {
3425 StringArray::new_unchecked(
3426 binary_array.offsets().clone(),
3427 binary_array.values().clone(),
3428 binary_array.nulls().cloned(),
3429 )
3430 };
3431 expect_ipc_validation_error(
3432 Arc::new(array),
3433 "Invalid argument error: Invalid UTF8 sequence at string index 3 (3..45): invalid utf-8 sequence of 1 bytes from index 38",
3434 );
3435 }
3436
3437 #[test]
3438 fn test_validation_of_invalid_string_view_array() {
3439 let valid: &[u8] = b" ";
3440 let mut invalid = vec![];
3441 invalid.extend_from_slice(b"ThisStringIsCertainlyLongerThan12Bytes");
3442 invalid.extend_from_slice(INVALID_UTF8_FIRST_CHAR);
3443 let binary_view_array =
3444 BinaryViewArray::from_iter(vec![None, Some(valid), None, Some(&invalid)]);
3445 let array = unsafe {
3448 StringViewArray::new_unchecked(
3449 binary_view_array.views().clone(),
3450 binary_view_array.data_buffers().to_vec(),
3451 binary_view_array.nulls().cloned(),
3452 )
3453 };
3454 expect_ipc_validation_error(
3455 Arc::new(array),
3456 "Invalid argument error: Encountered non-UTF-8 data at index 3: invalid utf-8 sequence of 1 bytes from index 38",
3457 );
3458 }
3459
3460 #[test]
3463 fn test_validation_of_invalid_dictionary_array() {
3464 let array = unsafe {
3465 let values = StringArray::from_iter_values(["a", "b", "c"]);
3466 let keys = Int32Array::from(vec![1, 200]); DictionaryArray::new_unchecked(keys, Arc::new(values))
3468 };
3469
3470 expect_ipc_validation_error(
3471 Arc::new(array),
3472 "Invalid argument error: Value at position 1 out of bounds: 200 (should be in [0, 2])",
3473 );
3474 }
3475
3476 #[test]
3477 fn test_validation_of_invalid_union_array() {
3478 let array = unsafe {
3479 let fields = UnionFields::try_new(
3480 vec![1, 3], vec![
3482 Field::new("a", DataType::Int32, false),
3483 Field::new("b", DataType::Utf8, false),
3484 ],
3485 )
3486 .unwrap();
3487 let type_ids = ScalarBuffer::from(vec![1i8, 2, 3]); let offsets = None;
3489 let children: Vec<ArrayRef> = vec![
3490 Arc::new(Int32Array::from(vec![10, 20, 30])),
3491 Arc::new(StringArray::from(vec![Some("a"), Some("b"), Some("c")])),
3492 ];
3493
3494 UnionArray::new_unchecked(fields, type_ids, offsets, children)
3495 };
3496
3497 expect_ipc_validation_error(
3498 Arc::new(array),
3499 "Invalid argument error: Type Ids values must match one of the field type ids",
3500 );
3501 }
3502
3503 const INVALID_UTF8_FIRST_CHAR: &[u8] = &[0xa0, 0xa1, 0x20, 0x20];
3506
3507 fn expect_ipc_validation_error(array: ArrayRef, expected_err: &str) {
3509 let rb = RecordBatch::try_from_iter([("a", array)]).unwrap();
3510
3511 let buf = write_stream(&rb); read_stream_skip_validation(&buf).unwrap();
3514 let err = read_stream(&buf).unwrap_err();
3515 assert_eq!(err.to_string(), expected_err);
3516
3517 let buf = write_ipc(&rb); read_ipc_skip_validation(&buf).unwrap();
3520 let err = read_ipc(&buf).unwrap_err();
3521 assert_eq!(err.to_string(), expected_err);
3522
3523 read_ipc_with_decoder_skip_validation(buf.clone()).unwrap();
3525 let err = read_ipc_with_decoder(buf).unwrap_err();
3526 assert_eq!(err.to_string(), expected_err);
3527 }
3528
3529 #[test]
3530 fn test_roundtrip_schema() {
3531 let schema = Schema::new(vec![
3532 Field::new(
3533 "a",
3534 DataType::Dictionary(Box::new(DataType::UInt16), Box::new(DataType::Utf8)),
3535 false,
3536 ),
3537 Field::new(
3538 "b",
3539 DataType::Dictionary(Box::new(DataType::UInt16), Box::new(DataType::Utf8)),
3540 false,
3541 ),
3542 ]);
3543
3544 let options = IpcWriteOptions::default();
3545 let data_gen = IpcDataGenerator::default();
3546 let mut dict_tracker = DictionaryTracker::new(false);
3547 let encoded_data =
3548 data_gen.schema_to_bytes_with_dictionary_tracker(&schema, &mut dict_tracker, &options);
3549 let mut schema_bytes = vec![];
3550 write_message(&mut schema_bytes, encoded_data, &options).expect("write_message");
3551
3552 let begin_offset: usize = if schema_bytes[0..4].eq(&CONTINUATION_MARKER) {
3553 4
3554 } else {
3555 0
3556 };
3557
3558 size_prefixed_root_as_message(&schema_bytes[begin_offset..])
3559 .expect_err("size_prefixed_root_as_message");
3560
3561 let msg = parse_message(&schema_bytes).expect("parse_message");
3562 let ipc_schema = msg.header_as_schema().expect("header_as_schema");
3563 let new_schema = fb_to_schema(ipc_schema);
3564
3565 assert_eq!(schema, new_schema);
3566 }
3567
3568 #[test]
3569 fn test_negative_meta_len() {
3570 let bytes = i32::to_le_bytes(-1);
3571 let mut buf = vec![];
3572 buf.extend(CONTINUATION_MARKER);
3573 buf.extend(bytes);
3574
3575 let reader = StreamReader::try_new(Cursor::new(buf), None);
3576 assert!(reader.is_err());
3577 }
3578
3579 #[test]
3585 fn test_read_null_dict_without_dictionary_batch() {
3586 let keys = Int32Array::new_null(4);
3588 let values: ArrayRef = new_empty_array(&DataType::Utf8);
3589 let dict_array = DictionaryArray::new(keys, values);
3590
3591 let schema = Arc::new(Schema::new(vec![Field::new(
3592 "d",
3593 dict_array.data_type().clone(),
3594 true,
3595 )]));
3596 let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(dict_array)]).unwrap();
3597
3598 let full_stream = write_stream(&batch);
3600
3601 let mut stripped = Vec::new();
3605 let mut cursor = Cursor::new(&full_stream);
3606 loop {
3607 let mut header = [0u8; 4];
3610 if cursor.read_exact(&mut header).is_err() {
3611 break;
3612 }
3613 if header == CONTINUATION_MARKER && cursor.read_exact(&mut header).is_err() {
3614 break;
3615 }
3616 let meta_len = u32::from_le_bytes(header) as usize;
3617 if meta_len == 0 {
3618 stripped.extend_from_slice(&CONTINUATION_MARKER);
3620 stripped.extend_from_slice(&0u32.to_le_bytes());
3621 break;
3622 }
3623 let mut meta_buf = vec![0u8; meta_len];
3624 cursor.read_exact(&mut meta_buf).unwrap();
3625
3626 let message = root_as_message(&meta_buf).unwrap();
3627 let body_len = message.bodyLength() as usize;
3628 let mut body_buf = vec![0u8; body_len];
3629 cursor.read_exact(&mut body_buf).unwrap();
3630
3631 if message.header_type() == crate::MessageHeader::DictionaryBatch {
3632 continue;
3635 }
3636 stripped.extend_from_slice(&CONTINUATION_MARKER);
3637 stripped.extend_from_slice(&(meta_len as u32).to_le_bytes());
3638 stripped.extend_from_slice(&meta_buf);
3639 stripped.extend_from_slice(&body_buf);
3640 }
3641
3642 let result = read_stream(&stripped).unwrap();
3644 assert_eq!(result.num_rows(), 4);
3645 assert_eq!(result.num_columns(), 1);
3646
3647 let col = result.column(0);
3648 assert_eq!(col.null_count(), 4);
3649 assert_eq!(col.len(), 4);
3650 assert!(matches!(col.data_type(), DataType::Dictionary(_, _)));
3652 }
3653
3654 #[test]
3657 fn test_projection_skip_list_view() {
3658 use crate::reader::FileReader;
3659 use crate::writer::FileWriter;
3660 use arrow_array::{
3661 GenericListViewArray, Int32Array, RecordBatch,
3662 builder::{GenericListViewBuilder, UInt32Builder},
3663 };
3664 use arrow_schema::{DataType, Field, Schema};
3665 use std::sync::Arc;
3666
3667 let mut builder = GenericListViewBuilder::<i32, _>::new(UInt32Builder::new());
3669
3670 builder.values().append_value(1);
3671 builder.values().append_value(2);
3672 builder.append(true);
3673
3674 builder.append(false);
3675
3676 builder.values().append_value(3);
3677 builder.values().append_value(4);
3678 builder.append(true);
3679
3680 let list_view: GenericListViewArray<i32> = builder.finish();
3681
3682 let values = Int32Array::from(vec![10, 20, 30]);
3684
3685 let schema = Arc::new(Schema::new(vec![
3687 Field::new("a", list_view.data_type().clone(), true),
3688 Field::new("b", DataType::Int32, false),
3689 ]));
3690 let batch =
3692 RecordBatch::try_new(schema, vec![Arc::new(list_view), Arc::new(values.clone())])
3693 .unwrap();
3694
3695 let mut buf = Vec::new();
3697 {
3698 let mut writer = FileWriter::try_new(&mut buf, &batch.schema()).unwrap();
3699 writer.write(&batch).unwrap();
3700 writer.finish().unwrap();
3701 }
3702
3703 let mut reader = FileReader::try_new(std::io::Cursor::new(buf), Some(vec![1])).unwrap();
3705 let read_batch = reader.next().unwrap().unwrap();
3706
3707 assert_eq!(read_batch.num_columns(), 1);
3709 assert_eq!(read_batch.column(0).as_ref(), &values);
3710 }
3711
3712 #[test]
3715 fn test_projection_skip_union_v4() {
3716 use crate::MetadataVersion;
3717 use crate::reader::FileReader;
3718 use crate::writer::{FileWriter, IpcWriteOptions};
3719 use arrow_array::{
3720 ArrayRef, Int32Array, RecordBatch, builder::UnionBuilder, types::Int32Type,
3721 };
3722 use arrow_schema::{DataType, Field, Schema};
3723 use std::sync::Arc;
3724
3725 let mut builder = UnionBuilder::new_dense();
3727 builder.append::<Int32Type>("a", 1).unwrap();
3728 builder.append::<Int32Type>("a", 2).unwrap();
3729 builder.append::<Int32Type>("a", 3).unwrap();
3730 let union = builder.build().unwrap();
3731
3732 let values = Int32Array::from(vec![10, 20, 30]);
3734
3735 let schema = Arc::new(Schema::new(vec![
3737 Field::new("union", union.data_type().clone(), false),
3738 Field::new("values", DataType::Int32, false),
3739 ]));
3740
3741 let batch = RecordBatch::try_new(
3743 schema,
3744 vec![Arc::new(union) as ArrayRef, Arc::new(values.clone())],
3745 )
3746 .unwrap();
3747
3748 let mut buf = Vec::new();
3750 {
3751 let options = IpcWriteOptions::try_new(8, false, MetadataVersion::V4).unwrap();
3752 let mut writer =
3753 FileWriter::try_new_with_options(&mut buf, &batch.schema(), options).unwrap();
3754 writer.write(&batch).unwrap();
3755 writer.finish().unwrap();
3756 }
3757 let mut reader = FileReader::try_new(std::io::Cursor::new(buf), Some(vec![1])).unwrap();
3759 let read_batch = reader.next().unwrap().unwrap();
3760
3761 assert_eq!(read_batch.num_columns(), 1);
3763 assert_eq!(read_batch.column(0).as_ref(), &values);
3764 }
3765
3766 #[test]
3770 fn test_projection_skip_fixed_width_types() {
3771 use std::sync::Arc;
3772
3773 use arrow_array::{ArrayRef, BooleanArray, Int32Array, RecordBatch, make_array};
3774 use arrow_buffer::Buffer;
3775 use arrow_data::ArrayData;
3776 use arrow_schema::{DataType, Field, IntervalUnit, Schema, TimeUnit};
3777
3778 use crate::reader::FileReader;
3779 use crate::writer::FileWriter;
3780
3781 fn make_array_for_type(data_type: DataType) -> ArrayRef {
3783 let len = 3;
3784
3785 if matches!(data_type, DataType::Boolean) {
3786 return Arc::new(BooleanArray::from(vec![true, false, true]));
3787 }
3788
3789 let width = data_type.primitive_width().unwrap();
3790 let data = ArrayData::builder(data_type)
3791 .len(len)
3792 .add_buffer(Buffer::from(vec![0_u8; len * width]))
3793 .build()
3794 .unwrap();
3795
3796 make_array(data)
3797 }
3798
3799 let data_types = vec![
3801 DataType::Boolean,
3802 DataType::Int8,
3803 DataType::Int16,
3804 DataType::Int32,
3805 DataType::Int64,
3806 DataType::UInt8,
3807 DataType::UInt16,
3808 DataType::UInt32,
3809 DataType::UInt64,
3810 DataType::Float16,
3811 DataType::Float32,
3812 DataType::Float64,
3813 DataType::Timestamp(TimeUnit::Second, None),
3814 DataType::Date32,
3815 DataType::Date64,
3816 DataType::Time32(TimeUnit::Second),
3817 DataType::Time64(TimeUnit::Microsecond),
3818 DataType::Duration(TimeUnit::Second),
3819 DataType::Interval(IntervalUnit::YearMonth),
3820 DataType::Interval(IntervalUnit::DayTime),
3821 DataType::Interval(IntervalUnit::MonthDayNano),
3822 DataType::Decimal32(9, 2),
3823 DataType::Decimal64(18, 2),
3824 DataType::Decimal128(38, 2),
3825 DataType::Decimal256(76, 2),
3826 ];
3827
3828 for data_type in data_types {
3833 let skipped = make_array_for_type(data_type.clone());
3834 let values = Int32Array::from(vec![10, 20, 30]);
3835
3836 let schema = Arc::new(Schema::new(vec![
3837 Field::new("skipped", data_type, false),
3838 Field::new("values", DataType::Int32, false),
3839 ]));
3840
3841 let batch =
3842 RecordBatch::try_new(schema, vec![skipped, Arc::new(values.clone())]).unwrap();
3843
3844 let mut buf = Vec::new();
3846 {
3847 let mut writer = FileWriter::try_new(&mut buf, &batch.schema()).unwrap();
3848 writer.write(&batch).unwrap();
3849 writer.finish().unwrap();
3850 }
3851
3852 let mut reader = FileReader::try_new(std::io::Cursor::new(buf), Some(vec![1])).unwrap();
3854 let read_batch = reader.next().unwrap().unwrap();
3855
3856 assert_eq!(read_batch.num_columns(), 1);
3858 assert_eq!(read_batch.column(0).as_ref(), &values);
3859 }
3860 }
3861}