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::{self};
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 {
543 self.skip_validation = skip_validation;
544 self
545 }
546
547 pub fn read_record_batch(mut self) -> Result<RecordBatch, ArrowError> {
549 let mut variadic_counts: VecDeque<i64> = self
550 .batch
551 .variadicBufferCounts()
552 .into_iter()
553 .flatten()
554 .collect();
555
556 let options = RecordBatchOptions::new().with_row_count(Some(self.batch.length() as usize));
557
558 let schema = Arc::clone(&self.schema);
559 if let Some(projection) = self.projection {
560 let mut arrays = Vec::with_capacity(projection.len());
561 for (idx, field) in schema.fields().iter().enumerate() {
563 let mut child = None;
565 for (proj_idx, projected_idx) in projection.iter().enumerate() {
566 if *projected_idx == idx {
567 if child.is_none() {
568 child = Some(self.create_array(field, &mut variadic_counts)?);
569 }
570
571 arrays.push((proj_idx, child.as_ref().unwrap().clone()));
573 }
574 }
575
576 if child.is_none() {
577 self.skip_field(field, &mut variadic_counts)?;
578 }
579 }
580
581 arrays.sort_by_key(|t| t.0);
582
583 let schema = Arc::new(schema.project(projection)?);
584 let columns = arrays.into_iter().map(|t| t.1).collect::<Vec<_>>();
585
586 if self.skip_validation.get() {
587 unsafe {
589 Ok(RecordBatch::new_unchecked(
590 schema,
591 columns,
592 self.batch.length() as usize,
593 ))
594 }
595 } else {
596 assert!(variadic_counts.is_empty());
597 RecordBatch::try_new_with_options(schema, columns, &options)
598 }
599 } else {
600 let mut children = Vec::with_capacity(schema.fields().len());
601 for field in schema.fields() {
603 let child = self.create_array(field, &mut variadic_counts)?;
604 children.push(child);
605 }
606
607 if self.skip_validation.get() {
608 unsafe {
610 Ok(RecordBatch::new_unchecked(
611 schema,
612 children,
613 self.batch.length() as usize,
614 ))
615 }
616 } else {
617 assert!(variadic_counts.is_empty());
618 RecordBatch::try_new_with_options(schema, children, &options)
619 }
620 }
621 }
622
623 fn next_buffer(&mut self) -> Result<Buffer, ArrowError> {
624 let buffer = self.buffers.next().ok_or_else(|| {
625 ArrowError::IpcError("Buffer count mismatched with metadata".to_string())
626 })?;
627 read_buffer(
628 buffer,
629 self.data,
630 self.compression,
631 &mut self.decompression_context,
632 )
633 }
634
635 fn skip_buffer(&mut self) {
636 self.buffers.next().unwrap();
637 }
638
639 fn next_node(&mut self, field: &Field) -> Result<&'a FieldNode, ArrowError> {
640 self.nodes.next().ok_or_else(|| {
641 ArrowError::SchemaError(format!(
642 "Invalid data for schema. {field} refers to node not found in schema",
643 ))
644 })
645 }
646
647 fn skip_field(
648 &mut self,
649 field: &Field,
650 variadic_count: &mut VecDeque<i64>,
651 ) -> Result<(), ArrowError> {
652 self.next_node(field)?;
653
654 match field.data_type() {
655 Utf8 | Binary | LargeBinary | LargeUtf8 => {
656 for _ in 0..3 {
657 self.skip_buffer()
658 }
659 }
660 Utf8View | BinaryView => {
661 let count = variadic_count
662 .pop_front()
663 .ok_or(ArrowError::IpcError(format!(
664 "Missing variadic count for {} column",
665 field.data_type()
666 )))?;
667 let count = count + 2; for _i in 0..count {
669 self.skip_buffer()
670 }
671 }
672 FixedSizeBinary(_) => {
673 self.skip_buffer();
674 self.skip_buffer();
675 }
676 List(list_field) | LargeList(list_field) | Map(list_field, _) => {
677 self.skip_buffer();
678 self.skip_buffer();
679 self.skip_field(list_field, variadic_count)?;
680 }
681 ListView(list_field) | LargeListView(list_field) => {
682 self.skip_buffer(); self.skip_buffer(); self.skip_buffer(); self.skip_field(list_field, variadic_count)?;
686 }
687 FixedSizeList(list_field, _) => {
688 self.skip_buffer();
689 self.skip_field(list_field, variadic_count)?;
690 }
691 Struct(struct_fields) => {
692 self.skip_buffer();
693
694 for struct_field in struct_fields {
696 self.skip_field(struct_field, variadic_count)?
697 }
698 }
699 RunEndEncoded(run_ends_field, values_field) => {
700 self.skip_field(run_ends_field, variadic_count)?;
701 self.skip_field(values_field, variadic_count)?;
702 }
703 Dictionary(_, _) => {
704 self.skip_buffer(); self.skip_buffer(); }
707 Union(fields, mode) => {
708 if self.version < MetadataVersion::V5 {
709 self.skip_buffer(); }
711 self.skip_buffer(); match mode {
714 UnionMode::Dense => self.skip_buffer(), UnionMode::Sparse => {}
716 };
717
718 for (_, field) in fields.iter() {
719 self.skip_field(field, variadic_count)?
720 }
721 }
722 Null => {}
724
725 Boolean
727 | Int8
728 | Int16
729 | Int32
730 | Int64
731 | UInt8
732 | UInt16
733 | UInt32
734 | UInt64
735 | Float16
736 | Float32
737 | Float64
738 | Timestamp(_, _)
739 | Date32
740 | Date64
741 | Time32(_)
742 | Time64(_)
743 | Duration(_)
744 | Interval(_)
745 | Decimal32(_, _)
746 | Decimal64(_, _)
747 | Decimal128(_, _)
748 | Decimal256(_, _) => {
749 self.skip_buffer();
750 self.skip_buffer();
751 }
752 };
753 Ok(())
754 }
755}
756
757pub fn read_record_batch(
768 buf: &Buffer,
769 batch: crate::RecordBatch,
770 schema: SchemaRef,
771 dictionaries_by_id: &HashMap<i64, ArrayRef>,
772 projection: Option<&[usize]>,
773 metadata: &MetadataVersion,
774) -> Result<RecordBatch, ArrowError> {
775 RecordBatchDecoder::try_new(buf, batch, schema, dictionaries_by_id, metadata)?
776 .with_projection(projection)
777 .with_require_alignment(false)
778 .read_record_batch()
779}
780
781pub fn read_dictionary(
784 buf: &Buffer,
785 batch: crate::DictionaryBatch,
786 schema: &Schema,
787 dictionaries_by_id: &mut HashMap<i64, ArrayRef>,
788 metadata: &MetadataVersion,
789) -> Result<(), ArrowError> {
790 read_dictionary_impl(
791 buf,
792 batch,
793 schema,
794 dictionaries_by_id,
795 metadata,
796 false,
797 UnsafeFlag::new(),
798 )
799}
800
801pub fn read_dictionary_impl(
803 buf: &Buffer,
804 batch: crate::DictionaryBatch,
805 schema: &Schema,
806 dictionaries_by_id: &mut HashMap<i64, ArrayRef>,
807 metadata: &MetadataVersion,
808 require_alignment: bool,
809 skip_validation: UnsafeFlag,
810) -> Result<(), ArrowError> {
811 let id = batch.id();
812
813 let dictionary_values = get_dictionary_values(
814 buf,
815 batch,
816 schema,
817 dictionaries_by_id,
818 metadata,
819 require_alignment,
820 skip_validation,
821 )?;
822
823 update_dictionaries(dictionaries_by_id, batch.isDelta(), id, dictionary_values)?;
824
825 Ok(())
826}
827
828fn update_dictionaries(
837 dictionaries_by_id: &mut HashMap<i64, ArrayRef>,
838 is_delta: bool,
839 dict_id: i64,
840 dict_values: ArrayRef,
841) -> Result<(), ArrowError> {
842 if !is_delta {
843 dictionaries_by_id.insert(dict_id, dict_values.clone());
847 return Ok(());
848 }
849
850 let existing = dictionaries_by_id.get(&dict_id).ok_or_else(|| {
851 ArrowError::InvalidArgumentError(format!(
852 "No existing dictionary for delta dictionary with id '{dict_id}'"
853 ))
854 })?;
855
856 let combined = concat::concat(&[existing, &dict_values]).map_err(|e| {
857 ArrowError::InvalidArgumentError(format!("Failed to concat delta dictionary: {e}"))
858 })?;
859
860 dictionaries_by_id.insert(dict_id, combined);
861
862 Ok(())
863}
864
865fn get_dictionary_values(
869 buf: &Buffer,
870 batch: crate::DictionaryBatch,
871 schema: &Schema,
872 dictionaries_by_id: &mut HashMap<i64, ArrayRef>,
873 metadata: &MetadataVersion,
874 require_alignment: bool,
875 skip_validation: UnsafeFlag,
876) -> Result<ArrayRef, ArrowError> {
877 let id = batch.id();
878 #[allow(deprecated)]
879 let fields_using_this_dictionary = schema.fields_with_dict_id(id);
880 let first_field = fields_using_this_dictionary.first().ok_or_else(|| {
881 ArrowError::InvalidArgumentError(format!("dictionary id {id} not found in schema"))
882 })?;
883
884 let dictionary_values: ArrayRef = match first_field.data_type() {
888 DataType::Dictionary(_, value_type) => {
889 let value = value_type.as_ref().clone();
891 let schema = Schema::new(vec![Field::new("", value, true)]);
892 let record_batch = RecordBatchDecoder::try_new(
894 buf,
895 batch.data().unwrap(),
896 Arc::new(schema),
897 dictionaries_by_id,
898 metadata,
899 )?
900 .with_require_alignment(require_alignment)
901 .with_skip_validation(skip_validation)
902 .read_record_batch()?;
903
904 Some(record_batch.column(0).clone())
905 }
906 _ => None,
907 }
908 .ok_or_else(|| {
909 ArrowError::InvalidArgumentError(format!("dictionary id {id} not found in schema"))
910 })?;
911
912 Ok(dictionary_values)
913}
914
915fn read_block<R: Read + Seek>(mut reader: R, block: &Block) -> Result<Buffer, ArrowError> {
917 reader.seek(SeekFrom::Start(block.offset() as u64))?;
918 let body_len = block.bodyLength().to_usize().unwrap();
919 let metadata_len = block.metaDataLength().to_usize().unwrap();
920 let total_len = body_len.checked_add(metadata_len).unwrap();
921
922 let mut buf = MutableBuffer::from_len_zeroed(total_len);
923 reader.read_exact(&mut buf)?;
924 Ok(buf.into())
925}
926
927fn parse_message(buf: &[u8]) -> Result<Message::Message<'_>, ArrowError> {
931 let buf = match buf[..4] == CONTINUATION_MARKER {
932 true => &buf[8..],
933 false => &buf[4..],
934 };
935 crate::root_as_message(buf)
936 .map_err(|err| ArrowError::ParseError(format!("Unable to get root as message: {err:?}")))
937}
938
939pub fn read_footer_length(buf: [u8; 10]) -> Result<usize, ArrowError> {
943 if buf[4..] != super::ARROW_MAGIC {
944 return Err(ArrowError::ParseError(
945 "Arrow file does not contain correct footer".to_string(),
946 ));
947 }
948
949 let footer_len = i32::from_le_bytes(buf[..4].try_into().unwrap());
951 footer_len
952 .try_into()
953 .map_err(|_| ArrowError::ParseError(format!("Invalid footer length: {footer_len}")))
954}
955
956#[derive(Debug)]
1021pub struct FileDecoder {
1022 schema: SchemaRef,
1023 dictionaries: HashMap<i64, ArrayRef>,
1024 version: MetadataVersion,
1025 projection: Option<Vec<usize>>,
1026 require_alignment: bool,
1027 skip_validation: UnsafeFlag,
1028}
1029
1030impl FileDecoder {
1031 pub fn new(schema: SchemaRef, version: MetadataVersion) -> Self {
1033 Self {
1034 schema,
1035 version,
1036 dictionaries: Default::default(),
1037 projection: None,
1038 require_alignment: false,
1039 skip_validation: UnsafeFlag::new(),
1040 }
1041 }
1042
1043 pub fn with_projection(mut self, projection: Vec<usize>) -> Self {
1045 self.projection = Some(projection);
1046 self
1047 }
1048
1049 pub fn with_require_alignment(mut self, require_alignment: bool) -> Self {
1062 self.require_alignment = require_alignment;
1063 self
1064 }
1065
1066 pub unsafe fn with_skip_validation(mut self, skip_validation: bool) -> Self {
1077 unsafe { self.skip_validation.set(skip_validation) };
1078 self
1079 }
1080
1081 fn read_message<'a>(&self, buf: &'a [u8]) -> Result<Message::Message<'a>, ArrowError> {
1082 let message = parse_message(buf)?;
1083
1084 if self.version != MetadataVersion::V1 && message.version() != self.version {
1086 return Err(ArrowError::IpcError(
1087 "Could not read IPC message as metadata versions mismatch".to_string(),
1088 ));
1089 }
1090 Ok(message)
1091 }
1092
1093 pub fn read_dictionary(&mut self, block: &Block, buf: &Buffer) -> Result<(), ArrowError> {
1095 let message = self.read_message(buf)?;
1096 match message.header_type() {
1097 crate::MessageHeader::DictionaryBatch => {
1098 let batch = message.header_as_dictionary_batch().unwrap();
1099 read_dictionary_impl(
1100 &buf.slice(block.metaDataLength() as _),
1101 batch,
1102 &self.schema,
1103 &mut self.dictionaries,
1104 &message.version(),
1105 self.require_alignment,
1106 self.skip_validation.clone(),
1107 )
1108 }
1109 t => Err(ArrowError::ParseError(format!(
1110 "Expecting DictionaryBatch in dictionary blocks, found {t:?}."
1111 ))),
1112 }
1113 }
1114
1115 pub fn read_record_batch(
1117 &self,
1118 block: &Block,
1119 buf: &Buffer,
1120 ) -> Result<Option<RecordBatch>, ArrowError> {
1121 let message = self.read_message(buf)?;
1122 match message.header_type() {
1123 crate::MessageHeader::Schema => Err(ArrowError::IpcError(
1124 "Not expecting a schema when messages are read".to_string(),
1125 )),
1126 crate::MessageHeader::RecordBatch => {
1127 let batch = message.header_as_record_batch().ok_or_else(|| {
1128 ArrowError::IpcError("Unable to read IPC message as record batch".to_string())
1129 })?;
1130 RecordBatchDecoder::try_new(
1132 &buf.slice(block.metaDataLength() as _),
1133 batch,
1134 self.schema.clone(),
1135 &self.dictionaries,
1136 &message.version(),
1137 )?
1138 .with_projection(self.projection.as_deref())
1139 .with_require_alignment(self.require_alignment)
1140 .with_skip_validation(self.skip_validation.clone())
1141 .read_record_batch()
1142 .map(Some)
1143 }
1144 crate::MessageHeader::NONE => Ok(None),
1145 t => Err(ArrowError::InvalidArgumentError(format!(
1146 "Reading types other than record batches not yet supported, unable to read {t:?}"
1147 ))),
1148 }
1149 }
1150}
1151
1152#[derive(Debug)]
1154pub struct FileReaderBuilder {
1155 projection: Option<Vec<usize>>,
1157 max_footer_fb_tables: usize,
1159 max_footer_fb_depth: usize,
1161}
1162
1163impl Default for FileReaderBuilder {
1164 fn default() -> Self {
1165 let verifier_options = VerifierOptions::default();
1166 Self {
1167 max_footer_fb_tables: verifier_options.max_tables,
1168 max_footer_fb_depth: verifier_options.max_depth,
1169 projection: None,
1170 }
1171 }
1172}
1173
1174impl FileReaderBuilder {
1175 pub fn new() -> Self {
1179 Self::default()
1180 }
1181
1182 pub fn with_projection(mut self, projection: Vec<usize>) -> Self {
1184 self.projection = Some(projection);
1185 self
1186 }
1187
1188 pub fn with_max_footer_fb_tables(mut self, max_footer_fb_tables: usize) -> Self {
1201 self.max_footer_fb_tables = max_footer_fb_tables;
1202 self
1203 }
1204
1205 pub fn with_max_footer_fb_depth(mut self, max_footer_fb_depth: usize) -> Self {
1218 self.max_footer_fb_depth = max_footer_fb_depth;
1219 self
1220 }
1221
1222 pub fn build<R: Read + Seek>(self, mut reader: R) -> Result<FileReader<R>, ArrowError> {
1224 let mut buffer = [0; 10];
1226 reader.seek(SeekFrom::End(-10))?;
1227 reader.read_exact(&mut buffer)?;
1228
1229 let footer_len = read_footer_length(buffer)?;
1230
1231 let mut footer_data = vec![0; footer_len];
1233 reader.seek(SeekFrom::End(-10 - footer_len as i64))?;
1234 reader.read_exact(&mut footer_data)?;
1235
1236 let verifier_options = VerifierOptions {
1237 max_tables: self.max_footer_fb_tables,
1238 max_depth: self.max_footer_fb_depth,
1239 ..Default::default()
1240 };
1241 let footer = crate::root_as_footer_with_opts(&verifier_options, &footer_data[..]).map_err(
1242 |err| ArrowError::ParseError(format!("Unable to get root as footer: {err:?}")),
1243 )?;
1244
1245 let blocks = footer.recordBatches().ok_or_else(|| {
1246 ArrowError::ParseError("Unable to get record batches from IPC Footer".to_string())
1247 })?;
1248
1249 let total_blocks = blocks.len();
1250
1251 let ipc_schema = footer.schema().unwrap();
1252 if !ipc_schema.endianness().equals_to_target_endianness() {
1253 return Err(ArrowError::IpcError(
1254 "the endianness of the source system does not match the endianness of the target system.".to_owned()
1255 ));
1256 }
1257
1258 let schema = crate::convert::fb_to_schema(ipc_schema);
1259
1260 let mut custom_metadata = HashMap::new();
1261 if let Some(fb_custom_metadata) = footer.custom_metadata() {
1262 for kv in fb_custom_metadata.into_iter() {
1263 custom_metadata.insert(
1264 kv.key().unwrap().to_string(),
1265 kv.value().unwrap().to_string(),
1266 );
1267 }
1268 }
1269
1270 let mut decoder = FileDecoder::new(Arc::new(schema), footer.version());
1271 if let Some(projection) = self.projection {
1272 decoder = decoder.with_projection(projection)
1273 }
1274
1275 if let Some(dictionaries) = footer.dictionaries() {
1277 for block in dictionaries {
1278 let buf = read_block(&mut reader, block)?;
1279 decoder.read_dictionary(block, &buf)?;
1280 }
1281 }
1282
1283 Ok(FileReader {
1284 reader,
1285 blocks: blocks.iter().copied().collect(),
1286 current_block: 0,
1287 total_blocks,
1288 decoder,
1289 custom_metadata,
1290 })
1291 }
1292}
1293
1294pub struct FileReader<R> {
1339 reader: R,
1341
1342 decoder: FileDecoder,
1344
1345 blocks: Vec<Block>,
1349
1350 current_block: usize,
1352
1353 total_blocks: usize,
1355
1356 custom_metadata: HashMap<String, String>,
1358}
1359
1360impl<R> fmt::Debug for FileReader<R> {
1361 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> Result<(), fmt::Error> {
1362 f.debug_struct("FileReader<R>")
1363 .field("decoder", &self.decoder)
1364 .field("blocks", &self.blocks)
1365 .field("current_block", &self.current_block)
1366 .field("total_blocks", &self.total_blocks)
1367 .finish_non_exhaustive()
1368 }
1369}
1370
1371impl<R: Read + Seek> FileReader<BufReader<R>> {
1372 pub fn try_new_buffered(reader: R, projection: Option<Vec<usize>>) -> Result<Self, ArrowError> {
1376 Self::try_new(BufReader::new(reader), projection)
1377 }
1378}
1379
1380impl<R: Read + Seek> FileReader<R> {
1381 pub fn try_new(reader: R, projection: Option<Vec<usize>>) -> Result<Self, ArrowError> {
1392 let builder = FileReaderBuilder {
1393 projection,
1394 ..Default::default()
1395 };
1396 builder.build(reader)
1397 }
1398
1399 pub fn custom_metadata(&self) -> &HashMap<String, String> {
1401 &self.custom_metadata
1402 }
1403
1404 pub fn num_batches(&self) -> usize {
1406 self.total_blocks
1407 }
1408
1409 pub fn schema(&self) -> SchemaRef {
1411 self.decoder.schema.clone()
1412 }
1413
1414 pub fn set_index(&mut self, index: usize) -> Result<(), ArrowError> {
1418 if index >= self.total_blocks {
1419 Err(ArrowError::InvalidArgumentError(format!(
1420 "Cannot set batch to index {} from {} total batches",
1421 index, self.total_blocks
1422 )))
1423 } else {
1424 self.current_block = index;
1425 Ok(())
1426 }
1427 }
1428
1429 fn maybe_next(&mut self) -> Result<Option<RecordBatch>, ArrowError> {
1430 let block = &self.blocks[self.current_block];
1431 self.current_block += 1;
1432
1433 let buffer = read_block(&mut self.reader, block)?;
1435 self.decoder.read_record_batch(block, &buffer)
1436 }
1437
1438 pub fn get_ref(&self) -> &R {
1442 &self.reader
1443 }
1444
1445 pub fn get_mut(&mut self) -> &mut R {
1449 &mut self.reader
1450 }
1451
1452 pub unsafe fn with_skip_validation(mut self, skip_validation: bool) -> Self {
1458 self.decoder = unsafe { self.decoder.with_skip_validation(skip_validation) };
1459 self
1460 }
1461}
1462
1463impl<R: Read + Seek> Iterator for FileReader<R> {
1464 type Item = Result<RecordBatch, ArrowError>;
1465
1466 fn next(&mut self) -> Option<Self::Item> {
1467 if self.current_block < self.total_blocks {
1469 self.maybe_next().transpose()
1470 } else {
1471 None
1472 }
1473 }
1474}
1475
1476impl<R: Read + Seek> RecordBatchReader for FileReader<R> {
1477 fn schema(&self) -> SchemaRef {
1478 self.schema()
1479 }
1480}
1481
1482pub struct StreamReader<R> {
1516 reader: MessageReader<R>,
1518
1519 schema: SchemaRef,
1521
1522 dictionaries_by_id: HashMap<i64, ArrayRef>,
1526
1527 finished: bool,
1531
1532 projection: Option<(Vec<usize>, Schema)>,
1534
1535 skip_validation: UnsafeFlag,
1539}
1540
1541impl<R> fmt::Debug for StreamReader<R> {
1542 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> std::result::Result<(), fmt::Error> {
1543 f.debug_struct("StreamReader<R>")
1544 .field("reader", &"R")
1545 .field("schema", &self.schema)
1546 .field("dictionaries_by_id", &self.dictionaries_by_id)
1547 .field("finished", &self.finished)
1548 .field("projection", &self.projection)
1549 .finish()
1550 }
1551}
1552
1553impl<R: Read> StreamReader<BufReader<R>> {
1554 pub fn try_new_buffered(reader: R, projection: Option<Vec<usize>>) -> Result<Self, ArrowError> {
1558 Self::try_new(BufReader::new(reader), projection)
1559 }
1560}
1561
1562impl<R: Read> StreamReader<R> {
1563 pub fn try_new(
1575 reader: R,
1576 projection: Option<Vec<usize>>,
1577 ) -> Result<StreamReader<R>, ArrowError> {
1578 let mut msg_reader = MessageReader::new(reader);
1579 let message = msg_reader.maybe_next()?;
1580 let Some((message, _)) = message else {
1581 return Err(ArrowError::IpcError(
1582 "Expected schema message, found empty stream.".to_string(),
1583 ));
1584 };
1585
1586 if message.header_type() != Message::MessageHeader::Schema {
1587 return Err(ArrowError::IpcError(format!(
1588 "Expected a schema as the first message in the stream, got: {:?}",
1589 message.header_type()
1590 )));
1591 }
1592
1593 let schema = message.header_as_schema().ok_or_else(|| {
1594 ArrowError::ParseError("Failed to parse schema from message header".to_string())
1595 })?;
1596 let schema = crate::convert::fb_to_schema(schema);
1597
1598 let dictionaries_by_id = HashMap::new();
1600
1601 let projection = match projection {
1602 Some(projection_indices) => {
1603 let schema = schema.project(&projection_indices)?;
1604 Some((projection_indices, schema))
1605 }
1606 _ => None,
1607 };
1608
1609 Ok(Self {
1610 reader: msg_reader,
1611 schema: Arc::new(schema),
1612 finished: false,
1613 dictionaries_by_id,
1614 projection,
1615 skip_validation: UnsafeFlag::new(),
1616 })
1617 }
1618
1619 #[deprecated(since = "53.0.0", note = "use `try_new` instead")]
1621 pub fn try_new_unbuffered(
1622 reader: R,
1623 projection: Option<Vec<usize>>,
1624 ) -> Result<Self, ArrowError> {
1625 Self::try_new(reader, projection)
1626 }
1627
1628 pub fn schema(&self) -> SchemaRef {
1630 self.schema.clone()
1631 }
1632
1633 pub fn is_finished(&self) -> bool {
1635 self.finished
1636 }
1637
1638 fn maybe_next(&mut self) -> Result<Option<RecordBatch>, ArrowError> {
1639 if self.finished {
1640 return Ok(None);
1641 }
1642
1643 loop {
1645 let message = self.next_ipc_message()?;
1646 let Some(message) = message else {
1647 self.finished = true;
1649 return Ok(None);
1650 };
1651
1652 match message {
1653 IpcMessage::Schema(_) => {
1654 return Err(ArrowError::IpcError(
1655 "Expected a record batch, but found a schema".to_string(),
1656 ));
1657 }
1658 IpcMessage::RecordBatch(record_batch) => {
1659 return Ok(Some(record_batch));
1660 }
1661 IpcMessage::DictionaryBatch { .. } => {
1662 continue;
1663 }
1664 };
1665 }
1666 }
1667
1668 pub(crate) fn next_ipc_message(&mut self) -> Result<Option<IpcMessage>, ArrowError> {
1676 let message = self.reader.maybe_next()?;
1677 let Some((message, body)) = message else {
1678 return Ok(None);
1680 };
1681
1682 let ipc_message = match message.header_type() {
1683 Message::MessageHeader::Schema => {
1684 let schema = message.header_as_schema().ok_or_else(|| {
1685 ArrowError::ParseError("Failed to parse schema from message header".to_string())
1686 })?;
1687 let arrow_schema = crate::convert::fb_to_schema(schema);
1688 IpcMessage::Schema(arrow_schema)
1689 }
1690 Message::MessageHeader::RecordBatch => {
1691 let batch = message.header_as_record_batch().ok_or_else(|| {
1692 ArrowError::IpcError("Unable to read IPC message as record batch".to_string())
1693 })?;
1694
1695 let version = message.version();
1696 let schema = self.schema.clone();
1697 let record_batch = RecordBatchDecoder::try_new(
1698 &body.into(),
1699 batch,
1700 schema,
1701 &self.dictionaries_by_id,
1702 &version,
1703 )?
1704 .with_projection(self.projection.as_ref().map(|x| x.0.as_ref()))
1705 .with_require_alignment(false)
1706 .with_skip_validation(self.skip_validation.clone())
1707 .read_record_batch()?;
1708 IpcMessage::RecordBatch(record_batch)
1709 }
1710 Message::MessageHeader::DictionaryBatch => {
1711 let dict = message.header_as_dictionary_batch().ok_or_else(|| {
1712 ArrowError::ParseError(
1713 "Failed to parse dictionary batch from message header".to_string(),
1714 )
1715 })?;
1716
1717 let version = message.version();
1718 let dict_values = get_dictionary_values(
1719 &body.into(),
1720 dict,
1721 &self.schema,
1722 &mut self.dictionaries_by_id,
1723 &version,
1724 false,
1725 self.skip_validation.clone(),
1726 )?;
1727
1728 update_dictionaries(
1729 &mut self.dictionaries_by_id,
1730 dict.isDelta(),
1731 dict.id(),
1732 dict_values.clone(),
1733 )?;
1734
1735 IpcMessage::DictionaryBatch {
1736 id: dict.id(),
1737 is_delta: (dict.isDelta()),
1738 values: (dict_values),
1739 }
1740 }
1741 x => {
1742 return Err(ArrowError::ParseError(format!(
1743 "Unsupported message header type in IPC stream: '{x:?}'"
1744 )));
1745 }
1746 };
1747
1748 Ok(Some(ipc_message))
1749 }
1750
1751 pub fn get_ref(&self) -> &R {
1755 self.reader.inner()
1756 }
1757
1758 pub fn get_mut(&mut self) -> &mut R {
1762 self.reader.inner_mut()
1763 }
1764
1765 pub unsafe fn with_skip_validation(mut self, skip_validation: bool) -> Self {
1771 unsafe { self.skip_validation.set(skip_validation) };
1772 self
1773 }
1774}
1775
1776impl<R: Read> Iterator for StreamReader<R> {
1777 type Item = Result<RecordBatch, ArrowError>;
1778
1779 fn next(&mut self) -> Option<Self::Item> {
1780 self.maybe_next().transpose()
1781 }
1782}
1783
1784impl<R: Read> RecordBatchReader for StreamReader<R> {
1785 fn schema(&self) -> SchemaRef {
1786 self.schema.clone()
1787 }
1788}
1789
1790#[derive(Debug)]
1796#[allow(dead_code)]
1797pub(crate) enum IpcMessage {
1798 Schema(arrow_schema::Schema),
1799 RecordBatch(RecordBatch),
1800 DictionaryBatch {
1801 id: i64,
1802 is_delta: bool,
1803 values: ArrayRef,
1804 },
1805}
1806
1807struct MessageReader<R> {
1810 reader: R,
1811 buf: Vec<u8>,
1812}
1813
1814impl<R: Read> MessageReader<R> {
1815 fn new(reader: R) -> Self {
1816 Self {
1817 reader,
1818 buf: Vec::new(),
1819 }
1820 }
1821
1822 fn maybe_next(&mut self) -> Result<Option<(Message::Message<'_>, MutableBuffer)>, ArrowError> {
1833 let meta_len = self.read_meta_len()?;
1834 let Some(meta_len) = meta_len else {
1835 return Ok(None);
1836 };
1837
1838 self.buf.resize(meta_len, 0);
1839 self.reader.read_exact(&mut self.buf)?;
1840
1841 let message = crate::root_as_message(self.buf.as_slice()).map_err(|err| {
1842 ArrowError::ParseError(format!("Unable to get root as message: {err:?}"))
1843 })?;
1844
1845 let mut buf = MutableBuffer::from_len_zeroed(message.bodyLength() as usize);
1846 self.reader.read_exact(&mut buf)?;
1847
1848 Ok(Some((message, buf)))
1849 }
1850
1851 fn inner_mut(&mut self) -> &mut R {
1853 &mut self.reader
1854 }
1855
1856 fn inner(&self) -> &R {
1858 &self.reader
1859 }
1860
1861 pub fn read_meta_len(&mut self) -> Result<Option<usize>, ArrowError> {
1870 let mut meta_len: [u8; 4] = [0; 4];
1871 match self.reader.read_exact(&mut meta_len) {
1872 Ok(_) => {}
1873 Err(e) => {
1874 return if e.kind() == std::io::ErrorKind::UnexpectedEof {
1875 Ok(None)
1879 } else {
1880 Err(ArrowError::from(e))
1881 };
1882 }
1883 };
1884
1885 let meta_len = {
1886 if meta_len == CONTINUATION_MARKER {
1889 self.reader.read_exact(&mut meta_len)?;
1890 }
1891
1892 i32::from_le_bytes(meta_len)
1893 };
1894
1895 if meta_len == 0 {
1896 return Ok(None);
1897 }
1898
1899 let meta_len = usize::try_from(meta_len)
1900 .map_err(|_| ArrowError::ParseError(format!("Invalid metadata length: {meta_len}")))?;
1901
1902 Ok(Some(meta_len))
1903 }
1904}
1905
1906#[cfg(test)]
1907mod tests {
1908 use std::io::Cursor;
1909
1910 use crate::convert::fb_to_schema;
1911 use crate::writer::{
1912 DictionaryTracker, IpcDataGenerator, IpcWriteOptions, unslice_run_array, write_message,
1913 };
1914
1915 use super::*;
1916
1917 use crate::{root_as_footer, root_as_message, size_prefixed_root_as_message};
1918 use arrow_array::builder::{PrimitiveRunBuilder, UnionBuilder};
1919 use arrow_array::types::*;
1920 use arrow_buffer::{NullBuffer, OffsetBuffer};
1921 use arrow_data::ArrayDataBuilder;
1922
1923 fn create_test_projection_schema() -> Schema {
1924 let list_data_type = DataType::List(Arc::new(Field::new_list_field(DataType::Int32, true)));
1926
1927 let fixed_size_list_data_type =
1928 DataType::FixedSizeList(Arc::new(Field::new_list_field(DataType::Int32, false)), 3);
1929
1930 let union_fields = UnionFields::from_fields(vec![
1931 Field::new("a", DataType::Int32, false),
1932 Field::new("b", DataType::Float64, false),
1933 ]);
1934
1935 let union_data_type = DataType::Union(union_fields, UnionMode::Dense);
1936
1937 let struct_fields = Fields::from(vec![
1938 Field::new("id", DataType::Int32, false),
1939 Field::new_list("list", Field::new_list_field(DataType::Int8, true), false),
1940 ]);
1941 let struct_data_type = DataType::Struct(struct_fields);
1942
1943 let run_encoded_data_type = DataType::RunEndEncoded(
1944 Arc::new(Field::new("run_ends", DataType::Int16, false)),
1945 Arc::new(Field::new("values", DataType::Int32, true)),
1946 );
1947
1948 Schema::new(vec![
1950 Field::new("f0", DataType::UInt32, false),
1951 Field::new("f1", DataType::Utf8, false),
1952 Field::new("f2", DataType::Boolean, false),
1953 Field::new("f3", union_data_type, true),
1954 Field::new("f4", DataType::Null, true),
1955 Field::new("f5", DataType::Float64, true),
1956 Field::new("f6", list_data_type, false),
1957 Field::new("f7", DataType::FixedSizeBinary(3), true),
1958 Field::new("f8", fixed_size_list_data_type, false),
1959 Field::new("f9", struct_data_type, false),
1960 Field::new("f10", run_encoded_data_type, false),
1961 Field::new("f11", DataType::Boolean, false),
1962 Field::new_dictionary("f12", DataType::Int8, DataType::Utf8, false),
1963 Field::new("f13", DataType::Utf8, false),
1964 ])
1965 }
1966
1967 fn create_test_projection_batch_data(schema: &Schema) -> RecordBatch {
1968 let array0 = UInt32Array::from(vec![1, 2, 3]);
1970 let array1 = StringArray::from(vec!["foo", "bar", "baz"]);
1971 let array2 = BooleanArray::from(vec![true, false, true]);
1972
1973 let mut union_builder = UnionBuilder::new_dense();
1974 union_builder.append::<Int32Type>("a", 1).unwrap();
1975 union_builder.append::<Float64Type>("b", 10.1).unwrap();
1976 union_builder.append_null::<Float64Type>("b").unwrap();
1977 let array3 = union_builder.build().unwrap();
1978
1979 let array4 = NullArray::new(3);
1980 let array5 = Float64Array::from(vec![Some(1.1), None, Some(3.3)]);
1981 let array6_values = vec![
1982 Some(vec![Some(10), Some(10), Some(10)]),
1983 Some(vec![Some(20), Some(20), Some(20)]),
1984 Some(vec![Some(30), Some(30)]),
1985 ];
1986 let array6 = ListArray::from_iter_primitive::<Int32Type, _, _>(array6_values);
1987 let array7_values = vec![vec![11, 12, 13], vec![22, 23, 24], vec![33, 34, 35]];
1988 let array7 = FixedSizeBinaryArray::try_from_iter(array7_values.into_iter()).unwrap();
1989
1990 let array8_values = ArrayData::builder(DataType::Int32)
1991 .len(9)
1992 .add_buffer(Buffer::from_slice_ref([40, 41, 42, 43, 44, 45, 46, 47, 48]))
1993 .build()
1994 .unwrap();
1995 let array8_data = ArrayData::builder(schema.field(8).data_type().clone())
1996 .len(3)
1997 .add_child_data(array8_values)
1998 .build()
1999 .unwrap();
2000 let array8 = FixedSizeListArray::from(array8_data);
2001
2002 let array9_id: ArrayRef = Arc::new(Int32Array::from(vec![1001, 1002, 1003]));
2003 let array9_list: ArrayRef =
2004 Arc::new(ListArray::from_iter_primitive::<Int8Type, _, _>(vec![
2005 Some(vec![Some(-10)]),
2006 Some(vec![Some(-20), Some(-20), Some(-20)]),
2007 Some(vec![Some(-30)]),
2008 ]));
2009 let array9 = ArrayDataBuilder::new(schema.field(9).data_type().clone())
2010 .add_child_data(array9_id.into_data())
2011 .add_child_data(array9_list.into_data())
2012 .len(3)
2013 .build()
2014 .unwrap();
2015 let array9 = StructArray::from(array9);
2016
2017 let array10_input = vec![Some(1_i32), None, None];
2018 let mut array10_builder = PrimitiveRunBuilder::<Int16Type, Int32Type>::new();
2019 array10_builder.extend(array10_input);
2020 let array10 = array10_builder.finish();
2021
2022 let array11 = BooleanArray::from(vec![false, false, true]);
2023
2024 let array12_values = StringArray::from(vec!["x", "yy", "zzz"]);
2025 let array12_keys = Int8Array::from_iter_values([1, 1, 2]);
2026 let array12 = DictionaryArray::new(array12_keys, Arc::new(array12_values));
2027
2028 let array13 = StringArray::from(vec!["a", "bb", "ccc"]);
2029
2030 RecordBatch::try_new(
2032 Arc::new(schema.clone()),
2033 vec![
2034 Arc::new(array0),
2035 Arc::new(array1),
2036 Arc::new(array2),
2037 Arc::new(array3),
2038 Arc::new(array4),
2039 Arc::new(array5),
2040 Arc::new(array6),
2041 Arc::new(array7),
2042 Arc::new(array8),
2043 Arc::new(array9),
2044 Arc::new(array10),
2045 Arc::new(array11),
2046 Arc::new(array12),
2047 Arc::new(array13),
2048 ],
2049 )
2050 .unwrap()
2051 }
2052
2053 #[test]
2054 fn test_negative_meta_len_start_stream() {
2055 let bytes = i32::to_le_bytes(-1);
2056 let mut buf = vec![];
2057 buf.extend(CONTINUATION_MARKER);
2058 buf.extend(bytes);
2059
2060 let reader_err = StreamReader::try_new(Cursor::new(buf), None).err();
2061 assert!(reader_err.is_some());
2062 assert_eq!(
2063 reader_err.unwrap().to_string(),
2064 "Parser error: Invalid metadata length: -1"
2065 );
2066 }
2067
2068 #[test]
2069 fn test_negative_meta_len_mid_stream() {
2070 let schema = Schema::new(vec![Field::new("a", DataType::Int32, false)]);
2071 let mut buf = Vec::new();
2072 {
2073 let mut writer = crate::writer::StreamWriter::try_new(&mut buf, &schema).unwrap();
2074 let batch =
2075 RecordBatch::try_new(Arc::new(schema), vec![Arc::new(Int32Array::from(vec![1]))])
2076 .unwrap();
2077 writer.write(&batch).unwrap();
2078 }
2079
2080 let bytes = i32::to_le_bytes(-1);
2081 buf.extend(CONTINUATION_MARKER);
2082 buf.extend(bytes);
2083
2084 let mut reader = StreamReader::try_new(Cursor::new(buf), None).unwrap();
2085 assert!(reader.maybe_next().is_ok());
2087 let batch_err = reader.maybe_next().err();
2089 assert!(batch_err.is_some());
2090 assert_eq!(
2091 batch_err.unwrap().to_string(),
2092 "Parser error: Invalid metadata length: -1"
2093 );
2094 }
2095
2096 #[test]
2097 fn test_missing_buffer_metadata_error() {
2098 use crate::r#gen::Message::*;
2099 use flatbuffers::FlatBufferBuilder;
2100
2101 let schema = Arc::new(Schema::new(vec![Field::new("col", DataType::Int32, true)]));
2102
2103 let mut fbb = FlatBufferBuilder::new();
2106 let nodes = fbb.create_vector(&[FieldNode::new(2, 0)]);
2107 let buffers = fbb.create_vector(&[crate::Buffer::new(0, 8)]);
2108 let batch_offset = RecordBatch::create(
2109 &mut fbb,
2110 &RecordBatchArgs {
2111 length: 2,
2112 nodes: Some(nodes),
2113 buffers: Some(buffers),
2114 compression: None,
2115 variadicBufferCounts: None,
2116 },
2117 );
2118 fbb.finish_minimal(batch_offset);
2119 let batch_bytes = fbb.finished_data().to_vec();
2120 let batch = flatbuffers::root::<RecordBatch>(&batch_bytes).unwrap();
2121
2122 let data_buffer = Buffer::from(vec![0u8; 8]);
2123 let dictionaries: HashMap<i64, ArrayRef> = HashMap::new();
2124 let metadata = MetadataVersion::V5;
2125
2126 let decoder = RecordBatchDecoder::try_new(
2127 &data_buffer,
2128 batch,
2129 schema.clone(),
2130 &dictionaries,
2131 &metadata,
2132 )
2133 .unwrap();
2134
2135 let result = decoder.read_record_batch();
2136
2137 match result {
2138 Err(ArrowError::IpcError(msg)) => {
2139 assert_eq!(msg, "Buffer count mismatched with metadata");
2140 }
2141 other => panic!("unexpected error: {other:?}"),
2142 }
2143 }
2144
2145 #[test]
2147 fn test_read_legacy_empty_list_without_offsets_buffer() {
2148 use crate::r#gen::Message::*;
2149 use flatbuffers::FlatBufferBuilder;
2150
2151 let schema = Arc::new(Schema::new(vec![Field::new_list(
2152 "items",
2153 Field::new_list_field(DataType::Int32, true),
2154 true,
2155 )]));
2156
2157 let mut fbb = FlatBufferBuilder::new();
2160 let nodes = fbb.create_vector(&[
2161 FieldNode::new(0, 0), FieldNode::new(0, 0), ]);
2164 let buffers = fbb.create_vector(&[
2165 crate::Buffer::new(0, 0), crate::Buffer::new(0, 0), crate::Buffer::new(0, 0), crate::Buffer::new(0, 0), ]);
2170 let batch_offset = RecordBatch::create(
2171 &mut fbb,
2172 &RecordBatchArgs {
2173 length: 0,
2174 nodes: Some(nodes),
2175 buffers: Some(buffers),
2176 compression: None,
2177 variadicBufferCounts: None,
2178 },
2179 );
2180 fbb.finish_minimal(batch_offset);
2181 let batch_bytes = fbb.finished_data().to_vec();
2182 let batch = flatbuffers::root::<RecordBatch>(&batch_bytes).unwrap();
2183
2184 let body = Buffer::from(Vec::<u8>::new());
2185 let dictionaries: HashMap<i64, ArrayRef> = HashMap::new();
2186 let metadata = MetadataVersion::V5;
2187
2188 let decoder =
2189 RecordBatchDecoder::try_new(&body, batch, schema.clone(), &dictionaries, &metadata)
2190 .unwrap();
2191
2192 let read_batch = decoder.read_record_batch().unwrap();
2193 assert_eq!(read_batch.num_rows(), 0);
2194
2195 let list = read_batch
2196 .column(0)
2197 .as_any()
2198 .downcast_ref::<ListArray>()
2199 .unwrap();
2200 assert_eq!(list.len(), 0);
2201 assert_eq!(list.values().len(), 0);
2202 }
2203
2204 #[test]
2206 fn test_read_legacy_empty_utf8_and_binary_without_offsets_buffer() {
2207 use crate::r#gen::Message::*;
2208 use flatbuffers::FlatBufferBuilder;
2209
2210 let schema = Arc::new(Schema::new(vec![
2211 Field::new("name", DataType::Utf8, true),
2212 Field::new("payload", DataType::Binary, true),
2213 ]));
2214
2215 let mut fbb = FlatBufferBuilder::new();
2218 let nodes = fbb.create_vector(&[
2219 FieldNode::new(0, 0), FieldNode::new(0, 0), ]);
2222 let buffers = fbb.create_vector(&[
2223 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), ]);
2230 let batch_offset = RecordBatch::create(
2231 &mut fbb,
2232 &RecordBatchArgs {
2233 length: 0,
2234 nodes: Some(nodes),
2235 buffers: Some(buffers),
2236 compression: None,
2237 variadicBufferCounts: None,
2238 },
2239 );
2240 fbb.finish_minimal(batch_offset);
2241 let batch_bytes = fbb.finished_data().to_vec();
2242 let batch = flatbuffers::root::<RecordBatch>(&batch_bytes).unwrap();
2243
2244 let body = Buffer::from(Vec::<u8>::new());
2245 let dictionaries: HashMap<i64, ArrayRef> = HashMap::new();
2246 let metadata = MetadataVersion::V5;
2247
2248 let decoder =
2249 RecordBatchDecoder::try_new(&body, batch, schema.clone(), &dictionaries, &metadata)
2250 .unwrap();
2251
2252 let read_batch = decoder.read_record_batch().unwrap();
2253 assert_eq!(read_batch.num_rows(), 0);
2254
2255 let utf8 = read_batch
2256 .column(0)
2257 .as_any()
2258 .downcast_ref::<StringArray>()
2259 .unwrap();
2260 assert_eq!(utf8.len(), 0);
2261 assert_eq!(utf8.value_offsets(), [0]);
2262
2263 let binary = read_batch
2264 .column(1)
2265 .as_any()
2266 .downcast_ref::<BinaryArray>()
2267 .unwrap();
2268 assert_eq!(binary.len(), 0);
2269 assert_eq!(binary.value_offsets(), [0]);
2270 }
2271
2272 #[test]
2273 fn test_projection_array_values() {
2274 let schema = create_test_projection_schema();
2276
2277 let batch = create_test_projection_batch_data(&schema);
2279
2280 let mut buf = Vec::new();
2282 {
2283 let mut writer = crate::writer::FileWriter::try_new(&mut buf, &schema).unwrap();
2284 writer.write(&batch).unwrap();
2285 writer.finish().unwrap();
2286 }
2287
2288 for index in 0..12 {
2290 let projection = vec![index];
2291 let reader = FileReader::try_new(std::io::Cursor::new(buf.clone()), Some(projection));
2292 let read_batch = reader.unwrap().next().unwrap().unwrap();
2293 let projected_column = read_batch.column(0);
2294 let expected_column = batch.column(index);
2295
2296 assert_eq!(projected_column.as_ref(), expected_column.as_ref());
2298 }
2299
2300 {
2301 let reader =
2303 FileReader::try_new(std::io::Cursor::new(buf.clone()), Some(vec![3, 2, 1]));
2304 let read_batch = reader.unwrap().next().unwrap().unwrap();
2305 let expected_batch = batch.project(&[3, 2, 1]).unwrap();
2306 assert_eq!(read_batch, expected_batch);
2307 }
2308 }
2309
2310 #[test]
2311 fn test_projection_duplicate_indices() {
2312 let schema = create_test_projection_schema();
2313 let batch = create_test_projection_batch_data(&schema);
2314
2315 let mut buf = Vec::new();
2317 {
2318 let mut writer = crate::writer::FileWriter::try_new(&mut buf, &schema).unwrap();
2319 writer.write(&batch).unwrap();
2320 writer.finish().unwrap();
2321 }
2322
2323 for projection in [vec![1, 1], vec![2, 0, 2]] {
2325 let reader =
2326 FileReader::try_new(std::io::Cursor::new(buf.clone()), Some(projection.clone()));
2327 let read_batch = reader.unwrap().next().unwrap().unwrap();
2328
2329 let expected_batch = batch.project(&projection).unwrap();
2330 assert_eq!(read_batch, expected_batch);
2331 }
2332 }
2333
2334 #[test]
2335 fn test_arrow_single_float_row() {
2336 let schema = Schema::new(vec![
2337 Field::new("a", DataType::Float32, false),
2338 Field::new("b", DataType::Float32, false),
2339 Field::new("c", DataType::Int32, false),
2340 Field::new("d", DataType::Int32, false),
2341 ]);
2342 let arrays = vec![
2343 Arc::new(Float32Array::from(vec![1.23])) as ArrayRef,
2344 Arc::new(Float32Array::from(vec![-6.50])) as ArrayRef,
2345 Arc::new(Int32Array::from(vec![2])) as ArrayRef,
2346 Arc::new(Int32Array::from(vec![1])) as ArrayRef,
2347 ];
2348 let batch = RecordBatch::try_new(Arc::new(schema.clone()), arrays).unwrap();
2349 let mut file = tempfile::tempfile().unwrap();
2351 let mut stream_writer = crate::writer::StreamWriter::try_new(&mut file, &schema).unwrap();
2352 stream_writer.write(&batch).unwrap();
2353 stream_writer.finish().unwrap();
2354
2355 drop(stream_writer);
2356
2357 file.rewind().unwrap();
2358
2359 let reader = StreamReader::try_new(&mut file, None).unwrap();
2361
2362 reader.for_each(|batch| {
2363 let batch = batch.unwrap();
2364 assert!(
2365 batch
2366 .column(0)
2367 .as_any()
2368 .downcast_ref::<Float32Array>()
2369 .unwrap()
2370 .value(0)
2371 != 0.0
2372 );
2373 assert!(
2374 batch
2375 .column(1)
2376 .as_any()
2377 .downcast_ref::<Float32Array>()
2378 .unwrap()
2379 .value(0)
2380 != 0.0
2381 );
2382 });
2383
2384 file.rewind().unwrap();
2385
2386 let reader = StreamReader::try_new(file, Some(vec![0, 3])).unwrap();
2388
2389 reader.for_each(|batch| {
2390 let batch = batch.unwrap();
2391 assert_eq!(batch.schema().fields().len(), 2);
2392 assert_eq!(batch.schema().fields()[0].data_type(), &DataType::Float32);
2393 assert_eq!(batch.schema().fields()[1].data_type(), &DataType::Int32);
2394 });
2395 }
2396
2397 fn write_ipc(rb: &RecordBatch) -> Vec<u8> {
2399 let mut buf = Vec::new();
2400 let mut writer = crate::writer::FileWriter::try_new(&mut buf, rb.schema_ref()).unwrap();
2401 writer.write(rb).unwrap();
2402 writer.finish().unwrap();
2403 buf
2404 }
2405
2406 fn read_ipc(buf: &[u8]) -> Result<RecordBatch, ArrowError> {
2408 let mut reader = FileReader::try_new(std::io::Cursor::new(buf), None)?;
2409 reader.next().unwrap()
2410 }
2411
2412 fn read_ipc_skip_validation(buf: &[u8]) -> Result<RecordBatch, ArrowError> {
2415 let mut reader = unsafe {
2416 FileReader::try_new(std::io::Cursor::new(buf), None)?.with_skip_validation(true)
2417 };
2418 reader.next().unwrap()
2419 }
2420
2421 fn roundtrip_ipc(rb: &RecordBatch) -> RecordBatch {
2422 let buf = write_ipc(rb);
2423 read_ipc(&buf).unwrap()
2424 }
2425
2426 fn read_ipc_with_decoder(buf: Vec<u8>) -> Result<RecordBatch, ArrowError> {
2429 read_ipc_with_decoder_inner(buf, false)
2430 }
2431
2432 fn read_ipc_with_decoder_skip_validation(buf: Vec<u8>) -> Result<RecordBatch, ArrowError> {
2435 read_ipc_with_decoder_inner(buf, true)
2436 }
2437
2438 fn read_ipc_with_decoder_inner(
2439 buf: Vec<u8>,
2440 skip_validation: bool,
2441 ) -> Result<RecordBatch, ArrowError> {
2442 let buffer = Buffer::from_vec(buf);
2443 let trailer_start = buffer.len() - 10;
2444 let footer_len = read_footer_length(buffer[trailer_start..].try_into().unwrap())?;
2445 let footer = root_as_footer(&buffer[trailer_start - footer_len..trailer_start])
2446 .map_err(|e| ArrowError::InvalidArgumentError(format!("Invalid footer: {e}")))?;
2447
2448 let schema = fb_to_schema(footer.schema().unwrap());
2449
2450 let mut decoder = unsafe {
2451 FileDecoder::new(Arc::new(schema), footer.version())
2452 .with_skip_validation(skip_validation)
2453 };
2454 for block in footer.dictionaries().iter().flatten() {
2456 let block_len = block.bodyLength() as usize + block.metaDataLength() as usize;
2457 let data = buffer.slice_with_length(block.offset() as _, block_len);
2458 decoder.read_dictionary(block, &data)?
2459 }
2460
2461 let batches = footer.recordBatches().unwrap();
2463 assert_eq!(batches.len(), 1); let block = batches.get(0);
2466 let block_len = block.bodyLength() as usize + block.metaDataLength() as usize;
2467 let data = buffer.slice_with_length(block.offset() as _, block_len);
2468 Ok(decoder.read_record_batch(block, &data)?.unwrap())
2469 }
2470
2471 fn write_stream(rb: &RecordBatch) -> Vec<u8> {
2473 let mut buf = Vec::new();
2474 let mut writer = crate::writer::StreamWriter::try_new(&mut buf, rb.schema_ref()).unwrap();
2475 writer.write(rb).unwrap();
2476 writer.finish().unwrap();
2477 buf
2478 }
2479
2480 fn read_stream(buf: &[u8]) -> Result<RecordBatch, ArrowError> {
2482 let mut reader = StreamReader::try_new(std::io::Cursor::new(buf), None)?;
2483 reader.next().unwrap()
2484 }
2485
2486 fn read_stream_skip_validation(buf: &[u8]) -> Result<RecordBatch, ArrowError> {
2489 let mut reader = unsafe {
2490 StreamReader::try_new(std::io::Cursor::new(buf), None)?.with_skip_validation(true)
2491 };
2492 reader.next().unwrap()
2493 }
2494
2495 fn roundtrip_ipc_stream(rb: &RecordBatch) -> RecordBatch {
2496 let buf = write_stream(rb);
2497 read_stream(&buf).unwrap()
2498 }
2499
2500 #[test]
2501 fn test_roundtrip_with_custom_metadata() {
2502 let schema = Schema::new(vec![Field::new("dummy", DataType::Float64, false)]);
2503 let mut buf = Vec::new();
2504 let mut writer = crate::writer::FileWriter::try_new(&mut buf, &schema).unwrap();
2505 let mut test_metadata = HashMap::new();
2506 test_metadata.insert("abc".to_string(), "abc".to_string());
2507 test_metadata.insert("def".to_string(), "def".to_string());
2508 for (k, v) in &test_metadata {
2509 writer.write_metadata(k, v);
2510 }
2511 writer.finish().unwrap();
2512 drop(writer);
2513
2514 let reader = crate::reader::FileReader::try_new(std::io::Cursor::new(buf), None).unwrap();
2515 assert_eq!(reader.custom_metadata(), &test_metadata);
2516 }
2517
2518 #[test]
2519 fn test_roundtrip_nested_dict() {
2520 let inner: DictionaryArray<Int32Type> = vec!["a", "b", "a"].into_iter().collect();
2521
2522 let array = Arc::new(inner) as ArrayRef;
2523
2524 let dctfield = Arc::new(Field::new("dict", array.data_type().clone(), false));
2525
2526 let s = StructArray::from(vec![(dctfield, array)]);
2527 let struct_array = Arc::new(s) as ArrayRef;
2528
2529 let schema = Arc::new(Schema::new(vec![Field::new(
2530 "struct",
2531 struct_array.data_type().clone(),
2532 false,
2533 )]));
2534
2535 let batch = RecordBatch::try_new(schema, vec![struct_array]).unwrap();
2536
2537 assert_eq!(batch, roundtrip_ipc(&batch));
2538 }
2539
2540 #[test]
2541 fn test_roundtrip_nested_dict_no_preserve_dict_id() {
2542 let inner: DictionaryArray<Int32Type> = vec!["a", "b", "a"].into_iter().collect();
2543
2544 let array = Arc::new(inner) as ArrayRef;
2545
2546 let dctfield = Arc::new(Field::new("dict", array.data_type().clone(), false));
2547
2548 let s = StructArray::from(vec![(dctfield, array)]);
2549 let struct_array = Arc::new(s) as ArrayRef;
2550
2551 let schema = Arc::new(Schema::new(vec![Field::new(
2552 "struct",
2553 struct_array.data_type().clone(),
2554 false,
2555 )]));
2556
2557 let batch = RecordBatch::try_new(schema, vec![struct_array]).unwrap();
2558
2559 let mut buf = Vec::new();
2560 let mut writer = crate::writer::FileWriter::try_new_with_options(
2561 &mut buf,
2562 batch.schema_ref(),
2563 IpcWriteOptions::default(),
2564 )
2565 .unwrap();
2566 writer.write(&batch).unwrap();
2567 writer.finish().unwrap();
2568 drop(writer);
2569
2570 let mut reader = FileReader::try_new(std::io::Cursor::new(buf), None).unwrap();
2571
2572 assert_eq!(batch, reader.next().unwrap().unwrap());
2573 }
2574
2575 fn check_union_with_builder(mut builder: UnionBuilder) {
2576 builder.append::<Int32Type>("a", 1).unwrap();
2577 builder.append_null::<Int32Type>("a").unwrap();
2578 builder.append::<Float64Type>("c", 3.0).unwrap();
2579 builder.append::<Int32Type>("a", 4).unwrap();
2580 builder.append::<Int64Type>("d", 11).unwrap();
2581 let union = builder.build().unwrap();
2582
2583 let schema = Arc::new(Schema::new(vec![Field::new(
2584 "union",
2585 union.data_type().clone(),
2586 false,
2587 )]));
2588
2589 let union_array = Arc::new(union) as ArrayRef;
2590
2591 let rb = RecordBatch::try_new(schema, vec![union_array]).unwrap();
2592 let rb2 = roundtrip_ipc(&rb);
2593 assert_eq!(rb.schema(), rb2.schema());
2596 assert_eq!(rb.num_columns(), rb2.num_columns());
2597 assert_eq!(rb.num_rows(), rb2.num_rows());
2598 let union1 = rb.column(0);
2599 let union2 = rb2.column(0);
2600
2601 assert_eq!(union1, union2);
2602 }
2603
2604 #[test]
2605 fn test_roundtrip_dense_union() {
2606 check_union_with_builder(UnionBuilder::new_dense());
2607 }
2608
2609 #[test]
2610 fn test_roundtrip_sparse_union() {
2611 check_union_with_builder(UnionBuilder::new_sparse());
2612 }
2613
2614 #[test]
2615 fn test_roundtrip_struct_empty_fields() {
2616 let nulls = NullBuffer::from(&[true, true, false]);
2617 let rb = RecordBatch::try_from_iter([(
2618 "",
2619 Arc::new(StructArray::new_empty_fields(nulls.len(), Some(nulls))) as _,
2620 )])
2621 .unwrap();
2622 let rb2 = roundtrip_ipc(&rb);
2623 assert_eq!(rb, rb2);
2624 }
2625
2626 #[test]
2627 fn test_roundtrip_stream_run_array_sliced() {
2628 let run_array_1: Int32RunArray = vec!["a", "a", "a", "b", "b", "c", "c", "c"]
2629 .into_iter()
2630 .collect();
2631 let run_array_1_sliced = run_array_1.slice(2, 5);
2632
2633 let run_array_2_inupt = vec![Some(1_i32), None, None, Some(2), Some(2)];
2634 let mut run_array_2_builder = PrimitiveRunBuilder::<Int16Type, Int32Type>::new();
2635 run_array_2_builder.extend(run_array_2_inupt);
2636 let run_array_2 = run_array_2_builder.finish();
2637
2638 let schema = Arc::new(Schema::new(vec![
2639 Field::new(
2640 "run_array_1_sliced",
2641 run_array_1_sliced.data_type().clone(),
2642 false,
2643 ),
2644 Field::new("run_array_2", run_array_2.data_type().clone(), false),
2645 ]));
2646 let input_batch = RecordBatch::try_new(
2647 schema,
2648 vec![Arc::new(run_array_1_sliced.clone()), Arc::new(run_array_2)],
2649 )
2650 .unwrap();
2651 let output_batch = roundtrip_ipc_stream(&input_batch);
2652
2653 assert_eq!(input_batch.column(1), output_batch.column(1));
2657
2658 let run_array_1_unsliced = unslice_run_array(run_array_1_sliced.into_data()).unwrap();
2659 assert_eq!(run_array_1_unsliced, output_batch.column(0).into_data());
2660 }
2661
2662 #[test]
2663 fn test_roundtrip_stream_nested_dict() {
2664 let xs = vec!["AA", "BB", "AA", "CC", "BB"];
2665 let dict = Arc::new(
2666 xs.clone()
2667 .into_iter()
2668 .collect::<DictionaryArray<Int8Type>>(),
2669 );
2670 let string_array: ArrayRef = Arc::new(StringArray::from(xs.clone()));
2671 let struct_array = StructArray::from(vec![
2672 (
2673 Arc::new(Field::new("f2.1", DataType::Utf8, false)),
2674 string_array,
2675 ),
2676 (
2677 Arc::new(Field::new("f2.2_struct", dict.data_type().clone(), false)),
2678 dict.clone() as ArrayRef,
2679 ),
2680 ]);
2681 let schema = Arc::new(Schema::new(vec![
2682 Field::new("f1_string", DataType::Utf8, false),
2683 Field::new("f2_struct", struct_array.data_type().clone(), false),
2684 ]));
2685 let input_batch = RecordBatch::try_new(
2686 schema,
2687 vec![
2688 Arc::new(StringArray::from(xs.clone())),
2689 Arc::new(struct_array),
2690 ],
2691 )
2692 .unwrap();
2693 let output_batch = roundtrip_ipc_stream(&input_batch);
2694 assert_eq!(input_batch, output_batch);
2695 }
2696
2697 #[test]
2698 fn test_ipc_writers_reject_dictionary_of_dictionary_schema() {
2699 let values = Arc::new(StringArray::from(vec![Some("a"), Some("b")])) as ArrayRef;
2700 let inner = Arc::new(DictionaryArray::new(
2701 UInt32Array::from_iter_values([0, 1]),
2702 values,
2703 )) as ArrayRef;
2704 let outer = Arc::new(DictionaryArray::new(
2705 UInt32Array::from_iter_values([0, 1, 0]),
2706 inner,
2707 )) as ArrayRef;
2708
2709 let schema = Arc::new(Schema::new(vec![Field::new(
2710 "f1",
2711 outer.data_type().clone(),
2712 false,
2713 )]));
2714 let batch = RecordBatch::try_new(schema, vec![outer]).unwrap();
2715
2716 let mut stream = Vec::new();
2717 let Err(err) = crate::writer::StreamWriter::try_new(&mut stream, batch.schema_ref()) else {
2718 panic!("IPC stream writer should reject dictionary-of-dictionary schemas");
2719 };
2720 assert!(stream.is_empty());
2721
2722 assert!(
2723 err.to_string().contains("dictionary-of-dictionary values"),
2724 "unexpected error: {err}"
2725 );
2726
2727 let mut file = Vec::new();
2728 let Err(err) = crate::writer::FileWriter::try_new(&mut file, batch.schema_ref()) else {
2729 panic!("IPC file writer should reject dictionary-of-dictionary schemas");
2730 };
2731 assert!(file.is_empty());
2732
2733 assert!(
2734 err.to_string().contains("dictionary-of-dictionary values"),
2735 "unexpected error: {err}"
2736 );
2737 }
2738
2739 #[test]
2740 fn test_roundtrip_stream_nested_dict_of_map_of_dict() {
2741 let values = StringArray::from(vec![Some("a"), None, Some("b"), Some("c")]);
2742 let values = Arc::new(values) as ArrayRef;
2743 let value_dict_keys = Int8Array::from_iter_values([0, 1, 1, 2, 3, 1]);
2744 let value_dict_array = DictionaryArray::new(value_dict_keys, values.clone());
2745
2746 let key_dict_keys = Int8Array::from_iter_values([0, 0, 2, 2, 2, 3]);
2747 let key_dict_array = DictionaryArray::new(key_dict_keys, values);
2748
2749 #[allow(deprecated)]
2750 let keys_field = Arc::new(Field::new_dict(
2751 "keys",
2752 DataType::Dictionary(Box::new(DataType::Int8), Box::new(DataType::Utf8)),
2753 false,
2754 1,
2755 false,
2756 ));
2757 #[allow(deprecated)]
2758 let values_field = Arc::new(Field::new_dict(
2759 "values",
2760 DataType::Dictionary(Box::new(DataType::Int8), Box::new(DataType::Utf8)),
2761 true,
2762 2,
2763 false,
2764 ));
2765 let entry_struct = StructArray::from(vec![
2766 (keys_field, make_array(key_dict_array.into_data())),
2767 (values_field, make_array(value_dict_array.into_data())),
2768 ]);
2769 let map_data_type = DataType::Map(
2770 Arc::new(Field::new(
2771 "entries",
2772 entry_struct.data_type().clone(),
2773 false,
2774 )),
2775 false,
2776 );
2777
2778 let entry_offsets = Buffer::from_slice_ref([0, 2, 4, 6]);
2779 let map_data = ArrayData::builder(map_data_type)
2780 .len(3)
2781 .add_buffer(entry_offsets)
2782 .add_child_data(entry_struct.into_data())
2783 .build()
2784 .unwrap();
2785 let map_array = MapArray::from(map_data);
2786
2787 let dict_keys = Int8Array::from_iter_values([0, 1, 1, 2, 2, 1]);
2788 let dict_dict_array = DictionaryArray::new(dict_keys, Arc::new(map_array));
2789
2790 let schema = Arc::new(Schema::new(vec![Field::new(
2791 "f1",
2792 dict_dict_array.data_type().clone(),
2793 false,
2794 )]));
2795 let input_batch = RecordBatch::try_new(schema, vec![Arc::new(dict_dict_array)]).unwrap();
2796 let output_batch = roundtrip_ipc_stream(&input_batch);
2797 assert_eq!(input_batch, output_batch);
2798 }
2799
2800 fn test_roundtrip_stream_dict_of_list_of_dict_impl<
2801 OffsetSize: OffsetSizeTrait,
2802 U: ArrowNativeType,
2803 >(
2804 list_data_type: DataType,
2805 offsets: &[U; 5],
2806 ) {
2807 let values = StringArray::from(vec![Some("a"), None, Some("c"), None]);
2808 let keys = Int8Array::from_iter_values([0, 0, 1, 2, 0, 1, 3]);
2809 let dict_array = DictionaryArray::new(keys, Arc::new(values));
2810 let dict_data = dict_array.to_data();
2811
2812 let value_offsets = Buffer::from_slice_ref(offsets);
2813
2814 let list_data = ArrayData::builder(list_data_type)
2815 .len(4)
2816 .add_buffer(value_offsets)
2817 .add_child_data(dict_data)
2818 .build()
2819 .unwrap();
2820 let list_array = GenericListArray::<OffsetSize>::from(list_data);
2821
2822 let keys_for_dict = Int8Array::from_iter_values([0, 3, 0, 1, 1, 2, 0, 1, 3]);
2823 let dict_dict_array = DictionaryArray::new(keys_for_dict, Arc::new(list_array));
2824
2825 let schema = Arc::new(Schema::new(vec![Field::new(
2826 "f1",
2827 dict_dict_array.data_type().clone(),
2828 false,
2829 )]));
2830 let input_batch = RecordBatch::try_new(schema, vec![Arc::new(dict_dict_array)]).unwrap();
2831 let output_batch = roundtrip_ipc_stream(&input_batch);
2832 assert_eq!(input_batch, output_batch);
2833 }
2834
2835 #[test]
2836 fn test_roundtrip_stream_dict_of_list_of_dict() {
2837 #[allow(deprecated)]
2839 let list_data_type = DataType::List(Arc::new(Field::new_dict(
2840 "item",
2841 DataType::Dictionary(Box::new(DataType::Int8), Box::new(DataType::Utf8)),
2842 true,
2843 1,
2844 false,
2845 )));
2846 let offsets: &[i32; 5] = &[0, 2, 4, 4, 6];
2847 test_roundtrip_stream_dict_of_list_of_dict_impl::<i32, i32>(list_data_type, offsets);
2848
2849 #[allow(deprecated)]
2851 let list_data_type = DataType::LargeList(Arc::new(Field::new_dict(
2852 "item",
2853 DataType::Dictionary(Box::new(DataType::Int8), Box::new(DataType::Utf8)),
2854 true,
2855 1,
2856 false,
2857 )));
2858 let offsets: &[i64; 5] = &[0, 2, 4, 4, 7];
2859 test_roundtrip_stream_dict_of_list_of_dict_impl::<i64, i64>(list_data_type, offsets);
2860 }
2861
2862 #[test]
2863 fn test_roundtrip_stream_dict_of_fixed_size_list_of_dict() {
2864 let values = StringArray::from(vec![Some("a"), None, Some("c"), None]);
2865 let keys = Int8Array::from_iter_values([0, 0, 1, 2, 0, 1, 3, 1, 2]);
2866 let dict_array = DictionaryArray::new(keys, Arc::new(values));
2867 let dict_data = dict_array.into_data();
2868
2869 #[allow(deprecated)]
2870 let list_data_type = DataType::FixedSizeList(
2871 Arc::new(Field::new_dict(
2872 "item",
2873 DataType::Dictionary(Box::new(DataType::Int8), Box::new(DataType::Utf8)),
2874 true,
2875 1,
2876 false,
2877 )),
2878 3,
2879 );
2880 let list_data = ArrayData::builder(list_data_type)
2881 .len(3)
2882 .add_child_data(dict_data)
2883 .build()
2884 .unwrap();
2885 let list_array = FixedSizeListArray::from(list_data);
2886
2887 let keys_for_dict = Int8Array::from_iter_values([0, 1, 0, 1, 1, 2, 0, 1, 2]);
2888 let dict_dict_array = DictionaryArray::new(keys_for_dict, Arc::new(list_array));
2889
2890 let schema = Arc::new(Schema::new(vec![Field::new(
2891 "f1",
2892 dict_dict_array.data_type().clone(),
2893 false,
2894 )]));
2895 let input_batch = RecordBatch::try_new(schema, vec![Arc::new(dict_dict_array)]).unwrap();
2896 let output_batch = roundtrip_ipc_stream(&input_batch);
2897 assert_eq!(input_batch, output_batch);
2898 }
2899
2900 const LONG_TEST_STRING: &str =
2901 "This is a long string to make sure binary view array handles it";
2902
2903 #[test]
2904 fn test_roundtrip_view_types() {
2905 let schema = Schema::new(vec![
2906 Field::new("field_1", DataType::BinaryView, true),
2907 Field::new("field_2", DataType::Utf8, true),
2908 Field::new("field_3", DataType::Utf8View, true),
2909 ]);
2910 let bin_values: Vec<Option<&[u8]>> = vec![
2911 Some(b"foo"),
2912 None,
2913 Some(b"bar"),
2914 Some(LONG_TEST_STRING.as_bytes()),
2915 ];
2916 let utf8_values: Vec<Option<&str>> =
2917 vec![Some("foo"), None, Some("bar"), Some(LONG_TEST_STRING)];
2918 let bin_view_array = BinaryViewArray::from_iter(bin_values);
2919 let utf8_array = StringArray::from_iter(utf8_values.iter());
2920 let utf8_view_array = StringViewArray::from_iter(utf8_values);
2921 let record_batch = RecordBatch::try_new(
2922 Arc::new(schema.clone()),
2923 vec![
2924 Arc::new(bin_view_array),
2925 Arc::new(utf8_array),
2926 Arc::new(utf8_view_array),
2927 ],
2928 )
2929 .unwrap();
2930
2931 assert_eq!(record_batch, roundtrip_ipc(&record_batch));
2932 assert_eq!(record_batch, roundtrip_ipc_stream(&record_batch));
2933
2934 let sliced_batch = record_batch.slice(1, 2);
2935 assert_eq!(sliced_batch, roundtrip_ipc(&sliced_batch));
2936 assert_eq!(sliced_batch, roundtrip_ipc_stream(&sliced_batch));
2937 }
2938
2939 #[test]
2940 fn test_roundtrip_view_types_nested_dict() {
2941 let bin_values: Vec<Option<&[u8]>> = vec![
2942 Some(b"foo"),
2943 None,
2944 Some(b"bar"),
2945 Some(LONG_TEST_STRING.as_bytes()),
2946 Some(b"field"),
2947 ];
2948 let utf8_values: Vec<Option<&str>> = vec![
2949 Some("foo"),
2950 None,
2951 Some("bar"),
2952 Some(LONG_TEST_STRING),
2953 Some("field"),
2954 ];
2955 let bin_view_array = Arc::new(BinaryViewArray::from_iter(bin_values));
2956 let utf8_view_array = Arc::new(StringViewArray::from_iter(utf8_values));
2957
2958 let key_dict_keys = Int8Array::from_iter_values([0, 0, 2, 2, 0, 2, 3]);
2959 let key_dict_array = DictionaryArray::new(key_dict_keys, utf8_view_array.clone());
2960 #[allow(deprecated)]
2961 let keys_field = Arc::new(Field::new_dict(
2962 "keys",
2963 DataType::Dictionary(Box::new(DataType::Int8), Box::new(DataType::Utf8View)),
2964 false,
2965 1,
2966 false,
2967 ));
2968
2969 let value_dict_keys = Int8Array::from_iter_values([0, 3, 0, 1, 2, 0, 1]);
2970 let value_dict_array = DictionaryArray::new(value_dict_keys, bin_view_array);
2971 #[allow(deprecated)]
2972 let values_field = Arc::new(Field::new_dict(
2973 "values",
2974 DataType::Dictionary(Box::new(DataType::Int8), Box::new(DataType::BinaryView)),
2975 true,
2976 2,
2977 false,
2978 ));
2979 let entry_struct = StructArray::from(vec![
2980 (keys_field, make_array(key_dict_array.into_data())),
2981 (values_field, make_array(value_dict_array.into_data())),
2982 ]);
2983
2984 let map_data_type = DataType::Map(
2985 Arc::new(Field::new(
2986 "entries",
2987 entry_struct.data_type().clone(),
2988 false,
2989 )),
2990 false,
2991 );
2992 let entry_offsets = Buffer::from_slice_ref([0, 2, 4, 7]);
2993 let map_data = ArrayData::builder(map_data_type)
2994 .len(3)
2995 .add_buffer(entry_offsets)
2996 .add_child_data(entry_struct.into_data())
2997 .build()
2998 .unwrap();
2999 let map_array = MapArray::from(map_data);
3000
3001 let dict_keys = Int8Array::from_iter_values([0, 1, 0, 1, 1, 2, 0, 1, 2]);
3002 let dict_dict_array = DictionaryArray::new(dict_keys, Arc::new(map_array));
3003 let schema = Arc::new(Schema::new(vec![Field::new(
3004 "f1",
3005 dict_dict_array.data_type().clone(),
3006 false,
3007 )]));
3008 let batch = RecordBatch::try_new(schema, vec![Arc::new(dict_dict_array)]).unwrap();
3009 assert_eq!(batch, roundtrip_ipc(&batch));
3010 assert_eq!(batch, roundtrip_ipc_stream(&batch));
3011
3012 let sliced_batch = batch.slice(1, 2);
3013 assert_eq!(sliced_batch, roundtrip_ipc(&sliced_batch));
3014 assert_eq!(sliced_batch, roundtrip_ipc_stream(&sliced_batch));
3015 }
3016
3017 #[test]
3018 fn test_no_columns_batch() {
3019 let schema = Arc::new(Schema::empty());
3020 let options = RecordBatchOptions::new()
3021 .with_match_field_names(true)
3022 .with_row_count(Some(10));
3023 let input_batch = RecordBatch::try_new_with_options(schema, vec![], &options).unwrap();
3024 let output_batch = roundtrip_ipc_stream(&input_batch);
3025 assert_eq!(input_batch, output_batch);
3026 }
3027
3028 #[test]
3029 fn test_unaligned() {
3030 let batch = RecordBatch::try_from_iter(vec![(
3031 "i32",
3032 Arc::new(Int32Array::from(vec![1, 2, 3, 4])) as _,
3033 )])
3034 .unwrap();
3035
3036 let r#gen = IpcDataGenerator {};
3037 let mut dict_tracker = DictionaryTracker::new(false);
3038 let (_, encoded) = r#gen
3039 .encode(
3040 &batch,
3041 &mut dict_tracker,
3042 &Default::default(),
3043 &mut Default::default(),
3044 )
3045 .unwrap();
3046
3047 let message = root_as_message(&encoded.ipc_message).unwrap();
3048
3049 let mut buffer = MutableBuffer::with_capacity(encoded.arrow_data.len() + 1);
3051 buffer.push(0_u8);
3052 buffer.extend_from_slice(&encoded.arrow_data);
3053 let b = Buffer::from(buffer).slice(1);
3054 assert_ne!(b.as_ptr().align_offset(8), 0);
3055
3056 let ipc_batch = message.header_as_record_batch().unwrap();
3057 let roundtrip = RecordBatchDecoder::try_new(
3058 &b,
3059 ipc_batch,
3060 batch.schema(),
3061 &Default::default(),
3062 &message.version(),
3063 )
3064 .unwrap()
3065 .with_require_alignment(false)
3066 .read_record_batch()
3067 .unwrap();
3068 assert_eq!(batch, roundtrip);
3069 }
3070
3071 #[test]
3072 fn test_unaligned_throws_error_with_require_alignment() {
3073 let batch = RecordBatch::try_from_iter(vec![(
3074 "i32",
3075 Arc::new(Int32Array::from(vec![1, 2, 3, 4])) as _,
3076 )])
3077 .unwrap();
3078
3079 let r#gen = IpcDataGenerator {};
3080 let mut dict_tracker = DictionaryTracker::new(false);
3081 let (_, encoded) = r#gen
3082 .encode(
3083 &batch,
3084 &mut dict_tracker,
3085 &Default::default(),
3086 &mut Default::default(),
3087 )
3088 .unwrap();
3089
3090 let message = root_as_message(&encoded.ipc_message).unwrap();
3091
3092 let mut buffer = MutableBuffer::with_capacity(encoded.arrow_data.len() + 1);
3094 buffer.push(0_u8);
3095 buffer.extend_from_slice(&encoded.arrow_data);
3096 let b = Buffer::from(buffer).slice(1);
3097 assert_ne!(b.as_ptr().align_offset(8), 0);
3098
3099 let ipc_batch = message.header_as_record_batch().unwrap();
3100 let result = RecordBatchDecoder::try_new(
3101 &b,
3102 ipc_batch,
3103 batch.schema(),
3104 &Default::default(),
3105 &message.version(),
3106 )
3107 .unwrap()
3108 .with_require_alignment(true)
3109 .read_record_batch();
3110
3111 let error = result.unwrap_err();
3112 assert_eq!(
3113 error.to_string(),
3114 "Invalid argument error: Misaligned buffers[0] in array of type Int32, \
3115 offset from expected alignment of 4 by 1"
3116 );
3117 }
3118
3119 #[test]
3120 fn test_file_with_massive_column_count() {
3121 let limit = 600_000;
3123
3124 let fields = (0..limit)
3125 .map(|i| Field::new(format!("{i}"), DataType::Boolean, false))
3126 .collect::<Vec<_>>();
3127 let schema = Arc::new(Schema::new(fields));
3128 let batch = RecordBatch::new_empty(schema);
3129
3130 let mut buf = Vec::new();
3131 let mut writer = crate::writer::FileWriter::try_new(&mut buf, batch.schema_ref()).unwrap();
3132 writer.write(&batch).unwrap();
3133 writer.finish().unwrap();
3134 drop(writer);
3135
3136 let mut reader = FileReaderBuilder::new()
3137 .with_max_footer_fb_tables(1_500_000)
3138 .build(std::io::Cursor::new(buf))
3139 .unwrap();
3140 let roundtrip_batch = reader.next().unwrap().unwrap();
3141
3142 assert_eq!(batch, roundtrip_batch);
3143 }
3144
3145 #[test]
3146 fn test_file_with_deeply_nested_columns() {
3147 let limit = 61;
3149
3150 let fields = (0..limit).fold(
3151 vec![Field::new("leaf", DataType::Boolean, false)],
3152 |field, index| vec![Field::new_struct(format!("{index}"), field, false)],
3153 );
3154 let schema = Arc::new(Schema::new(fields));
3155 let batch = RecordBatch::new_empty(schema);
3156
3157 let mut buf = Vec::new();
3158 let mut writer = crate::writer::FileWriter::try_new(&mut buf, batch.schema_ref()).unwrap();
3159 writer.write(&batch).unwrap();
3160 writer.finish().unwrap();
3161 drop(writer);
3162
3163 let mut reader = FileReaderBuilder::new()
3164 .with_max_footer_fb_depth(65)
3165 .build(std::io::Cursor::new(buf))
3166 .unwrap();
3167 let roundtrip_batch = reader.next().unwrap().unwrap();
3168
3169 assert_eq!(batch, roundtrip_batch);
3170 }
3171
3172 #[test]
3173 fn test_invalid_struct_array_ipc_read_errors() {
3174 let a_field = Field::new("a", DataType::Int32, false);
3175 let b_field = Field::new("b", DataType::Int32, false);
3176 let struct_fields = Fields::from(vec![a_field.clone(), b_field.clone()]);
3177
3178 let a_array_data = ArrayData::builder(a_field.data_type().clone())
3179 .len(4)
3180 .add_buffer(Buffer::from_slice_ref([1, 2, 3, 4]))
3181 .build()
3182 .unwrap();
3183 let b_array_data = ArrayData::builder(b_field.data_type().clone())
3184 .len(3)
3185 .add_buffer(Buffer::from_slice_ref([5, 6, 7]))
3186 .build()
3187 .unwrap();
3188
3189 let invalid_struct_arr = unsafe {
3190 StructArray::new_unchecked(
3191 struct_fields,
3192 vec![make_array(a_array_data), make_array(b_array_data)],
3193 None,
3194 )
3195 };
3196
3197 expect_ipc_validation_error(
3198 Arc::new(invalid_struct_arr),
3199 "Invalid argument error: Incorrect array length for StructArray field \"b\", expected 4 got 3",
3200 );
3201 }
3202
3203 #[test]
3204 fn test_invalid_nested_array_ipc_read_errors() {
3205 let a_field = Field::new("a", DataType::Int32, false);
3207 let b_field = Field::new("b", DataType::Utf8, false);
3208
3209 let schema = Arc::new(Schema::new(vec![Field::new_struct(
3210 "s",
3211 vec![a_field.clone(), b_field.clone()],
3212 false,
3213 )]));
3214
3215 let a_array_data = ArrayData::builder(a_field.data_type().clone())
3216 .len(4)
3217 .add_buffer(Buffer::from_slice_ref([1, 2, 3, 4]))
3218 .build()
3219 .unwrap();
3220 let b_array_data = {
3222 let valid: &[u8] = b" ";
3223 let mut invalid = vec![];
3224 invalid.extend_from_slice(b"ValidString");
3225 invalid.extend_from_slice(INVALID_UTF8_FIRST_CHAR);
3226 let binary_array =
3227 BinaryArray::from_iter(vec![None, Some(valid), None, Some(&invalid)]);
3228 let array = unsafe {
3229 StringArray::new_unchecked(
3230 binary_array.offsets().clone(),
3231 binary_array.values().clone(),
3232 binary_array.nulls().cloned(),
3233 )
3234 };
3235 array.into_data()
3236 };
3237 let struct_data_type = schema.field(0).data_type();
3238
3239 let invalid_struct_arr = unsafe {
3240 make_array(
3241 ArrayData::builder(struct_data_type.clone())
3242 .len(4)
3243 .add_child_data(a_array_data)
3244 .add_child_data(b_array_data)
3245 .build_unchecked(),
3246 )
3247 };
3248 expect_ipc_validation_error(
3249 invalid_struct_arr,
3250 "Invalid argument error: Invalid UTF8 sequence at string index 3 (3..18): invalid utf-8 sequence of 1 bytes from index 11",
3251 );
3252 }
3253
3254 #[test]
3255 fn test_same_dict_id_without_preserve() {
3256 let batch = RecordBatch::try_new(
3257 Arc::new(Schema::new(
3258 ["a", "b"]
3259 .iter()
3260 .map(|name| {
3261 #[allow(deprecated)]
3262 Field::new_dict(
3263 name.to_string(),
3264 DataType::Dictionary(
3265 Box::new(DataType::Int32),
3266 Box::new(DataType::Utf8),
3267 ),
3268 true,
3269 0,
3270 false,
3271 )
3272 })
3273 .collect::<Vec<Field>>(),
3274 )),
3275 vec![
3276 Arc::new(
3277 vec![Some("c"), Some("d")]
3278 .into_iter()
3279 .collect::<DictionaryArray<Int32Type>>(),
3280 ) as ArrayRef,
3281 Arc::new(
3282 vec![Some("e"), Some("f")]
3283 .into_iter()
3284 .collect::<DictionaryArray<Int32Type>>(),
3285 ) as ArrayRef,
3286 ],
3287 )
3288 .expect("Failed to create RecordBatch");
3289
3290 let mut buf = vec![];
3292 {
3293 let mut writer = crate::writer::StreamWriter::try_new_with_options(
3294 &mut buf,
3295 batch.schema().as_ref(),
3296 crate::writer::IpcWriteOptions::default(),
3297 )
3298 .expect("Failed to create StreamWriter");
3299 writer.write(&batch).expect("Failed to write RecordBatch");
3300 writer.finish().expect("Failed to finish StreamWriter");
3301 }
3302
3303 StreamReader::try_new(std::io::Cursor::new(buf), None)
3304 .expect("Failed to create StreamReader")
3305 .for_each(|decoded_batch| {
3306 assert_eq!(decoded_batch.expect("Failed to read RecordBatch"), batch);
3307 });
3308 }
3309
3310 #[test]
3311 fn test_validation_of_invalid_list_array() {
3312 let array = unsafe {
3314 let values = Int32Array::from(vec![1, 2, 3]);
3315 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);
3318 let nulls = None;
3319 ListArray::new(Arc::new(field), offsets, Arc::new(values), nulls)
3320 };
3321
3322 expect_ipc_validation_error(
3323 Arc::new(array),
3324 "Invalid argument error: Offset invariant failure: offset at position 2 out of bounds: 4 > 2",
3325 );
3326 }
3327
3328 #[test]
3329 fn test_validation_of_invalid_string_array() {
3330 let valid: &[u8] = b" ";
3331 let mut invalid = vec![];
3332 invalid.extend_from_slice(b"ThisStringIsCertainlyLongerThan12Bytes");
3333 invalid.extend_from_slice(INVALID_UTF8_FIRST_CHAR);
3334 let binary_array = BinaryArray::from_iter(vec![None, Some(valid), None, Some(&invalid)]);
3335 let array = unsafe {
3338 StringArray::new_unchecked(
3339 binary_array.offsets().clone(),
3340 binary_array.values().clone(),
3341 binary_array.nulls().cloned(),
3342 )
3343 };
3344 expect_ipc_validation_error(
3345 Arc::new(array),
3346 "Invalid argument error: Invalid UTF8 sequence at string index 3 (3..45): invalid utf-8 sequence of 1 bytes from index 38",
3347 );
3348 }
3349
3350 #[test]
3351 fn test_validation_of_invalid_string_view_array() {
3352 let valid: &[u8] = b" ";
3353 let mut invalid = vec![];
3354 invalid.extend_from_slice(b"ThisStringIsCertainlyLongerThan12Bytes");
3355 invalid.extend_from_slice(INVALID_UTF8_FIRST_CHAR);
3356 let binary_view_array =
3357 BinaryViewArray::from_iter(vec![None, Some(valid), None, Some(&invalid)]);
3358 let array = unsafe {
3361 StringViewArray::new_unchecked(
3362 binary_view_array.views().clone(),
3363 binary_view_array.data_buffers().to_vec(),
3364 binary_view_array.nulls().cloned(),
3365 )
3366 };
3367 expect_ipc_validation_error(
3368 Arc::new(array),
3369 "Invalid argument error: Encountered non-UTF-8 data at index 3: invalid utf-8 sequence of 1 bytes from index 38",
3370 );
3371 }
3372
3373 #[test]
3376 fn test_validation_of_invalid_dictionary_array() {
3377 let array = unsafe {
3378 let values = StringArray::from_iter_values(["a", "b", "c"]);
3379 let keys = Int32Array::from(vec![1, 200]); DictionaryArray::new_unchecked(keys, Arc::new(values))
3381 };
3382
3383 expect_ipc_validation_error(
3384 Arc::new(array),
3385 "Invalid argument error: Value at position 1 out of bounds: 200 (should be in [0, 2])",
3386 );
3387 }
3388
3389 #[test]
3390 fn test_validation_of_invalid_union_array() {
3391 let array = unsafe {
3392 let fields = UnionFields::try_new(
3393 vec![1, 3], vec![
3395 Field::new("a", DataType::Int32, false),
3396 Field::new("b", DataType::Utf8, false),
3397 ],
3398 )
3399 .unwrap();
3400 let type_ids = ScalarBuffer::from(vec![1i8, 2, 3]); let offsets = None;
3402 let children: Vec<ArrayRef> = vec![
3403 Arc::new(Int32Array::from(vec![10, 20, 30])),
3404 Arc::new(StringArray::from(vec![Some("a"), Some("b"), Some("c")])),
3405 ];
3406
3407 UnionArray::new_unchecked(fields, type_ids, offsets, children)
3408 };
3409
3410 expect_ipc_validation_error(
3411 Arc::new(array),
3412 "Invalid argument error: Type Ids values must match one of the field type ids",
3413 );
3414 }
3415
3416 const INVALID_UTF8_FIRST_CHAR: &[u8] = &[0xa0, 0xa1, 0x20, 0x20];
3419
3420 fn expect_ipc_validation_error(array: ArrayRef, expected_err: &str) {
3422 let rb = RecordBatch::try_from_iter([("a", array)]).unwrap();
3423
3424 let buf = write_stream(&rb); read_stream_skip_validation(&buf).unwrap();
3427 let err = read_stream(&buf).unwrap_err();
3428 assert_eq!(err.to_string(), expected_err);
3429
3430 let buf = write_ipc(&rb); read_ipc_skip_validation(&buf).unwrap();
3433 let err = read_ipc(&buf).unwrap_err();
3434 assert_eq!(err.to_string(), expected_err);
3435
3436 read_ipc_with_decoder_skip_validation(buf.clone()).unwrap();
3438 let err = read_ipc_with_decoder(buf).unwrap_err();
3439 assert_eq!(err.to_string(), expected_err);
3440 }
3441
3442 #[test]
3443 fn test_roundtrip_schema() {
3444 let schema = Schema::new(vec![
3445 Field::new(
3446 "a",
3447 DataType::Dictionary(Box::new(DataType::UInt16), Box::new(DataType::Utf8)),
3448 false,
3449 ),
3450 Field::new(
3451 "b",
3452 DataType::Dictionary(Box::new(DataType::UInt16), Box::new(DataType::Utf8)),
3453 false,
3454 ),
3455 ]);
3456
3457 let options = IpcWriteOptions::default();
3458 let data_gen = IpcDataGenerator::default();
3459 let mut dict_tracker = DictionaryTracker::new(false);
3460 let encoded_data =
3461 data_gen.schema_to_bytes_with_dictionary_tracker(&schema, &mut dict_tracker, &options);
3462 let mut schema_bytes = vec![];
3463 write_message(&mut schema_bytes, encoded_data, &options).expect("write_message");
3464
3465 let begin_offset: usize = if schema_bytes[0..4].eq(&CONTINUATION_MARKER) {
3466 4
3467 } else {
3468 0
3469 };
3470
3471 size_prefixed_root_as_message(&schema_bytes[begin_offset..])
3472 .expect_err("size_prefixed_root_as_message");
3473
3474 let msg = parse_message(&schema_bytes).expect("parse_message");
3475 let ipc_schema = msg.header_as_schema().expect("header_as_schema");
3476 let new_schema = fb_to_schema(ipc_schema);
3477
3478 assert_eq!(schema, new_schema);
3479 }
3480
3481 #[test]
3482 fn test_negative_meta_len() {
3483 let bytes = i32::to_le_bytes(-1);
3484 let mut buf = vec![];
3485 buf.extend(CONTINUATION_MARKER);
3486 buf.extend(bytes);
3487
3488 let reader = StreamReader::try_new(Cursor::new(buf), None);
3489 assert!(reader.is_err());
3490 }
3491
3492 #[test]
3498 fn test_read_null_dict_without_dictionary_batch() {
3499 let keys = Int32Array::new_null(4);
3501 let values: ArrayRef = new_empty_array(&DataType::Utf8);
3502 let dict_array = DictionaryArray::new(keys, values);
3503
3504 let schema = Arc::new(Schema::new(vec![Field::new(
3505 "d",
3506 dict_array.data_type().clone(),
3507 true,
3508 )]));
3509 let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(dict_array)]).unwrap();
3510
3511 let full_stream = write_stream(&batch);
3513
3514 let mut stripped = Vec::new();
3518 let mut cursor = Cursor::new(&full_stream);
3519 loop {
3520 let mut header = [0u8; 4];
3523 if cursor.read_exact(&mut header).is_err() {
3524 break;
3525 }
3526 if header == CONTINUATION_MARKER && cursor.read_exact(&mut header).is_err() {
3527 break;
3528 }
3529 let meta_len = u32::from_le_bytes(header) as usize;
3530 if meta_len == 0 {
3531 stripped.extend_from_slice(&CONTINUATION_MARKER);
3533 stripped.extend_from_slice(&0u32.to_le_bytes());
3534 break;
3535 }
3536 let mut meta_buf = vec![0u8; meta_len];
3537 cursor.read_exact(&mut meta_buf).unwrap();
3538
3539 let message = root_as_message(&meta_buf).unwrap();
3540 let body_len = message.bodyLength() as usize;
3541 let mut body_buf = vec![0u8; body_len];
3542 cursor.read_exact(&mut body_buf).unwrap();
3543
3544 if message.header_type() == crate::MessageHeader::DictionaryBatch {
3545 continue;
3548 }
3549 stripped.extend_from_slice(&CONTINUATION_MARKER);
3550 stripped.extend_from_slice(&(meta_len as u32).to_le_bytes());
3551 stripped.extend_from_slice(&meta_buf);
3552 stripped.extend_from_slice(&body_buf);
3553 }
3554
3555 let result = read_stream(&stripped).unwrap();
3557 assert_eq!(result.num_rows(), 4);
3558 assert_eq!(result.num_columns(), 1);
3559
3560 let col = result.column(0);
3561 assert_eq!(col.null_count(), 4);
3562 assert_eq!(col.len(), 4);
3563 assert!(matches!(col.data_type(), DataType::Dictionary(_, _)));
3565 }
3566
3567 #[test]
3570 fn test_projection_skip_list_view() {
3571 use crate::reader::FileReader;
3572 use crate::writer::FileWriter;
3573 use arrow_array::{
3574 GenericListViewArray, Int32Array, RecordBatch,
3575 builder::{GenericListViewBuilder, UInt32Builder},
3576 };
3577 use arrow_schema::{DataType, Field, Schema};
3578 use std::sync::Arc;
3579
3580 let mut builder = GenericListViewBuilder::<i32, _>::new(UInt32Builder::new());
3582
3583 builder.values().append_value(1);
3584 builder.values().append_value(2);
3585 builder.append(true);
3586
3587 builder.append(false);
3588
3589 builder.values().append_value(3);
3590 builder.values().append_value(4);
3591 builder.append(true);
3592
3593 let list_view: GenericListViewArray<i32> = builder.finish();
3594
3595 let values = Int32Array::from(vec![10, 20, 30]);
3597
3598 let schema = Arc::new(Schema::new(vec![
3600 Field::new("a", list_view.data_type().clone(), true),
3601 Field::new("b", DataType::Int32, false),
3602 ]));
3603 let batch =
3605 RecordBatch::try_new(schema, vec![Arc::new(list_view), Arc::new(values.clone())])
3606 .unwrap();
3607
3608 let mut buf = Vec::new();
3610 {
3611 let mut writer = FileWriter::try_new(&mut buf, &batch.schema()).unwrap();
3612 writer.write(&batch).unwrap();
3613 writer.finish().unwrap();
3614 }
3615
3616 let mut reader = FileReader::try_new(std::io::Cursor::new(buf), Some(vec![1])).unwrap();
3618 let read_batch = reader.next().unwrap().unwrap();
3619
3620 assert_eq!(read_batch.num_columns(), 1);
3622 assert_eq!(read_batch.column(0).as_ref(), &values);
3623 }
3624
3625 #[test]
3628 fn test_projection_skip_union_v4() {
3629 use crate::MetadataVersion;
3630 use crate::reader::FileReader;
3631 use crate::writer::{FileWriter, IpcWriteOptions};
3632 use arrow_array::{
3633 ArrayRef, Int32Array, RecordBatch, builder::UnionBuilder, types::Int32Type,
3634 };
3635 use arrow_schema::{DataType, Field, Schema};
3636 use std::sync::Arc;
3637
3638 let mut builder = UnionBuilder::new_dense();
3640 builder.append::<Int32Type>("a", 1).unwrap();
3641 builder.append::<Int32Type>("a", 2).unwrap();
3642 builder.append::<Int32Type>("a", 3).unwrap();
3643 let union = builder.build().unwrap();
3644
3645 let values = Int32Array::from(vec![10, 20, 30]);
3647
3648 let schema = Arc::new(Schema::new(vec![
3650 Field::new("union", union.data_type().clone(), false),
3651 Field::new("values", DataType::Int32, false),
3652 ]));
3653
3654 let batch = RecordBatch::try_new(
3656 schema,
3657 vec![Arc::new(union) as ArrayRef, Arc::new(values.clone())],
3658 )
3659 .unwrap();
3660
3661 let mut buf = Vec::new();
3663 {
3664 let options = IpcWriteOptions::try_new(8, false, MetadataVersion::V4).unwrap();
3665 let mut writer =
3666 FileWriter::try_new_with_options(&mut buf, &batch.schema(), options).unwrap();
3667 writer.write(&batch).unwrap();
3668 writer.finish().unwrap();
3669 }
3670 let mut reader = FileReader::try_new(std::io::Cursor::new(buf), Some(vec![1])).unwrap();
3672 let read_batch = reader.next().unwrap().unwrap();
3673
3674 assert_eq!(read_batch.num_columns(), 1);
3676 assert_eq!(read_batch.column(0).as_ref(), &values);
3677 }
3678
3679 #[test]
3683 fn test_projection_skip_fixed_width_types() {
3684 use std::sync::Arc;
3685
3686 use arrow_array::{ArrayRef, BooleanArray, Int32Array, RecordBatch, make_array};
3687 use arrow_buffer::Buffer;
3688 use arrow_data::ArrayData;
3689 use arrow_schema::{DataType, Field, IntervalUnit, Schema, TimeUnit};
3690
3691 use crate::reader::FileReader;
3692 use crate::writer::FileWriter;
3693
3694 fn make_array_for_type(data_type: DataType) -> ArrayRef {
3696 let len = 3;
3697
3698 if matches!(data_type, DataType::Boolean) {
3699 return Arc::new(BooleanArray::from(vec![true, false, true]));
3700 }
3701
3702 let width = data_type.primitive_width().unwrap();
3703 let data = ArrayData::builder(data_type)
3704 .len(len)
3705 .add_buffer(Buffer::from(vec![0_u8; len * width]))
3706 .build()
3707 .unwrap();
3708
3709 make_array(data)
3710 }
3711
3712 let data_types = vec![
3714 DataType::Boolean,
3715 DataType::Int8,
3716 DataType::Int16,
3717 DataType::Int32,
3718 DataType::Int64,
3719 DataType::UInt8,
3720 DataType::UInt16,
3721 DataType::UInt32,
3722 DataType::UInt64,
3723 DataType::Float16,
3724 DataType::Float32,
3725 DataType::Float64,
3726 DataType::Timestamp(TimeUnit::Second, None),
3727 DataType::Date32,
3728 DataType::Date64,
3729 DataType::Time32(TimeUnit::Second),
3730 DataType::Time64(TimeUnit::Microsecond),
3731 DataType::Duration(TimeUnit::Second),
3732 DataType::Interval(IntervalUnit::YearMonth),
3733 DataType::Interval(IntervalUnit::DayTime),
3734 DataType::Interval(IntervalUnit::MonthDayNano),
3735 DataType::Decimal32(9, 2),
3736 DataType::Decimal64(18, 2),
3737 DataType::Decimal128(38, 2),
3738 DataType::Decimal256(76, 2),
3739 ];
3740
3741 for data_type in data_types {
3746 let skipped = make_array_for_type(data_type.clone());
3747 let values = Int32Array::from(vec![10, 20, 30]);
3748
3749 let schema = Arc::new(Schema::new(vec![
3750 Field::new("skipped", data_type, false),
3751 Field::new("values", DataType::Int32, false),
3752 ]));
3753
3754 let batch =
3755 RecordBatch::try_new(schema, vec![skipped, Arc::new(values.clone())]).unwrap();
3756
3757 let mut buf = Vec::new();
3759 {
3760 let mut writer = FileWriter::try_new(&mut buf, &batch.schema()).unwrap();
3761 writer.write(&batch).unwrap();
3762 writer.finish().unwrap();
3763 }
3764
3765 let mut reader = FileReader::try_new(std::io::Cursor::new(buf), Some(vec![1])).unwrap();
3767 let read_batch = reader.next().unwrap().unwrap();
3768
3769 assert_eq!(read_batch.num_columns(), 1);
3771 assert_eq!(read_batch.column(0).as_ref(), &values);
3772 }
3773 }
3774}