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