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 = self.next_variadic_buffer_count(variadic_counts, data_type)?;
106 let buffers = (0..count)
107 .map(|_| self.next_buffer())
108 .collect::<Result<Vec<_>, _>>()?;
109 let field_node = self.next_node(field)?;
110 self.create_primitive_array(field_node, data_type, &buffers)
111 }
112 FixedSizeBinary(_) => {
113 let field_node = self.next_node(field)?;
114 let buffers = [self.next_buffer()?, self.next_buffer()?];
115 self.create_primitive_array(field_node, data_type, &buffers)
116 }
117 List(list_field) | LargeList(list_field) | Map(list_field, _) => {
118 let list_node = self.next_node(field)?;
119 let list_buffers = [self.next_buffer()?, self.next_buffer()?];
120 let values = self.create_array(list_field, variadic_counts)?;
121 self.create_list_array(list_node, data_type, &list_buffers, values)
122 }
123 ListView(list_field) | LargeListView(list_field) => {
124 let list_node = self.next_node(field)?;
125 let list_buffers = [
126 self.next_buffer()?, self.next_buffer()?, self.next_buffer()?, ];
130 let values = self.create_array(list_field, variadic_counts)?;
131 self.create_list_view_array(list_node, data_type, &list_buffers, values)
132 }
133 FixedSizeList(list_field, _) => {
134 let list_node = self.next_node(field)?;
135 let list_buffers = [self.next_buffer()?];
136 let values = self.create_array(list_field, variadic_counts)?;
137 self.create_list_array(list_node, data_type, &list_buffers, values)
138 }
139 Struct(struct_fields) => {
140 let struct_node = self.next_node(field)?;
141 let null_buffer = self.next_buffer()?;
142
143 let mut struct_arrays = Vec::with_capacity(struct_fields.len());
145 for struct_field in struct_fields {
148 let child = self.create_array(struct_field, variadic_counts)?;
149 struct_arrays.push(child);
150 }
151 self.create_struct_array(struct_node, null_buffer, struct_fields, struct_arrays)
152 }
153 RunEndEncoded(run_ends_field, values_field) => {
154 let run_node = self.next_node(field)?;
155 let run_ends = self.create_array(run_ends_field, variadic_counts)?;
156 let values = self.create_array(values_field, variadic_counts)?;
157
158 let run_array_length = run_node.length() as usize;
159 let builder = ArrayData::builder(data_type.clone())
160 .len(run_array_length)
161 .offset(0)
162 .add_child_data(run_ends.into_data())
163 .add_child_data(values.into_data())
164 .null_count(run_node.null_count() as usize);
165
166 self.create_array_from_builder(builder)
167 }
168 Dictionary(_, _) => {
170 let index_node = self.next_node(field)?;
171 let index_buffers = [self.next_buffer()?, self.next_buffer()?];
172
173 #[expect(deprecated)]
174 let dict_id = field.dict_id().ok_or_else(|| {
175 ArrowError::ParseError(format!("Field {field} does not have dict id"))
176 })?;
177
178 let value_array = match self.dictionaries_by_id.get(&dict_id) {
179 Some(array) => array.clone(),
180 None => {
181 if let Dictionary(_, value_type) = data_type {
185 arrow_array::new_empty_array(value_type.as_ref())
186 } else {
187 unreachable!()
188 }
189 }
190 };
191
192 self.create_dictionary_array(index_node, data_type, &index_buffers, value_array)
193 }
194 Union(fields, mode) => {
195 let union_node = self.next_node(field)?;
196 let len = union_node.length() as usize;
197
198 if self.version < MetadataVersion::V5 {
201 self.next_buffer()?;
202 }
203
204 let type_ids: ScalarBuffer<i8> =
205 self.next_buffer()?.slice_with_length(0, len).into();
206
207 let value_offsets = match mode {
208 UnionMode::Dense => {
209 let offsets: ScalarBuffer<i32> =
210 self.next_buffer()?.slice_with_length(0, len * 4).into();
211 Some(offsets)
212 }
213 UnionMode::Sparse => None,
214 };
215
216 let mut children = Vec::with_capacity(fields.len());
217
218 for (_id, field) in fields.iter() {
219 let child = self.create_array(field, variadic_counts)?;
220 children.push(child);
221 }
222
223 let array = if self.skip_validation.get() {
224 unsafe {
226 UnionArray::new_unchecked(fields.clone(), type_ids, value_offsets, children)
227 }
228 } else {
229 UnionArray::try_new(fields.clone(), type_ids, value_offsets, children)?
230 };
231 Ok(Arc::new(array))
232 }
233 Null => {
234 let node = self.next_node(field)?;
235 let length = node.length();
236 let null_count = node.null_count();
237
238 if length != null_count {
239 return Err(ArrowError::SchemaError(format!(
240 "Field {field} of NullArray has unequal null_count {null_count} and len {length}"
241 )));
242 }
243
244 let builder = ArrayData::builder(data_type.clone())
245 .len(length as usize)
246 .offset(0);
247 self.create_array_from_builder(builder)
248 }
249 _ => {
250 let field_node = self.next_node(field)?;
251 let buffers = [self.next_buffer()?, self.next_buffer()?];
252 self.create_primitive_array(field_node, data_type, &buffers)
253 }
254 }
255 }
256
257 fn create_primitive_array(
260 &self,
261 field_node: &FieldNode,
262 data_type: &DataType,
263 buffers: &[Buffer],
264 ) -> Result<ArrayRef, ArrowError> {
265 let length = field_node.length() as usize;
266 let null_buffer = (field_node.null_count() > 0).then_some(buffers[0].clone());
267 let mut builder = match data_type {
268 Utf8 | Binary | LargeBinary | LargeUtf8 => {
269 ArrayData::builder(data_type.clone())
271 .len(length)
272 .buffers(buffers[1..3].to_vec())
273 .null_bit_buffer(null_buffer)
274 }
275 BinaryView | Utf8View => ArrayData::builder(data_type.clone())
276 .len(length)
277 .buffers(buffers[1..].to_vec())
278 .null_bit_buffer(null_buffer),
279 _ if data_type.is_primitive() || matches!(data_type, Boolean | FixedSizeBinary(_)) => {
280 ArrayData::builder(data_type.clone())
282 .len(length)
283 .add_buffer(buffers[1].clone())
284 .null_bit_buffer(null_buffer)
285 }
286 t => unreachable!("Data type {:?} either unsupported or not primitive", t),
287 };
288
289 builder = builder.null_count(field_node.null_count() as usize);
290
291 self.create_array_from_builder(builder)
292 }
293
294 fn create_array_from_builder(&self, builder: ArrayDataBuilder) -> Result<ArrayRef, ArrowError> {
296 let mut builder = builder.align_buffers(!self.require_alignment);
297 if self.skip_validation.get() {
298 unsafe { builder = builder.skip_validation(true) }
300 }
301 Ok(make_array(builder.build()?))
302 }
303
304 fn create_list_array(
307 &self,
308 field_node: &FieldNode,
309 data_type: &DataType,
310 buffers: &[Buffer],
311 child_array: ArrayRef,
312 ) -> Result<ArrayRef, ArrowError> {
313 let null_buffer = (field_node.null_count() > 0).then_some(buffers[0].clone());
314 let length = field_node.length() as usize;
315 let child_data = child_array.into_data();
316 let mut builder = match data_type {
317 List(_) | LargeList(_) | Map(_, _) => ArrayData::builder(data_type.clone())
318 .len(length)
319 .add_buffer(buffers[1].clone())
320 .add_child_data(child_data)
321 .null_bit_buffer(null_buffer),
322
323 FixedSizeList(_, _) => ArrayData::builder(data_type.clone())
324 .len(length)
325 .add_child_data(child_data)
326 .null_bit_buffer(null_buffer),
327
328 _ => unreachable!("Cannot create list or map array from {:?}", data_type),
329 };
330
331 builder = builder.null_count(field_node.null_count() as usize);
332
333 self.create_array_from_builder(builder)
334 }
335
336 fn create_list_view_array(
337 &self,
338 field_node: &FieldNode,
339 data_type: &DataType,
340 buffers: &[Buffer],
341 child_array: ArrayRef,
342 ) -> Result<ArrayRef, ArrowError> {
343 assert!(matches!(data_type, ListView(_) | LargeListView(_)));
344
345 let null_buffer = (field_node.null_count() > 0).then_some(buffers[0].clone());
346 let length = field_node.length() as usize;
347 let child_data = child_array.into_data();
348
349 self.create_array_from_builder(
350 ArrayData::builder(data_type.clone())
351 .len(length)
352 .add_buffer(buffers[1].clone()) .add_buffer(buffers[2].clone()) .add_child_data(child_data)
355 .null_bit_buffer(null_buffer)
356 .null_count(field_node.null_count() as usize),
357 )
358 }
359
360 fn create_struct_array(
361 &self,
362 struct_node: &FieldNode,
363 null_buffer: Buffer,
364 struct_fields: &Fields,
365 struct_arrays: Vec<ArrayRef>,
366 ) -> Result<ArrayRef, ArrowError> {
367 let null_count = struct_node.null_count() as usize;
368 let len = struct_node.length() as usize;
369 let skip_validation = self.skip_validation.get();
370
371 let nulls = if null_count > 0 {
372 if !skip_validation {
373 let needed_len = len.div_ceil(8);
375 if null_buffer.len() < needed_len {
376 return Err(ArrowError::InvalidArgumentError(format!(
377 "null_bit_buffer size too small. got {} needed {}",
378 null_buffer.len(),
379 needed_len
380 )));
381 }
382 }
383 let validity_buffer = BooleanBuffer::new(null_buffer, 0, len);
384 let null_buffer = if skip_validation {
385 unsafe { NullBuffer::new_unchecked(validity_buffer, null_count) }
387 } else {
388 let null_buffer = NullBuffer::new(validity_buffer);
389
390 if null_buffer.null_count() != null_count {
391 return Err(ArrowError::InvalidArgumentError(format!(
392 "null_count value ({}) doesn't match actual number of nulls in array ({})",
393 null_count,
394 null_buffer.null_count()
395 )));
396 }
397
398 null_buffer
399 };
400
401 Some(null_buffer)
402 } else {
403 None
404 };
405 if struct_arrays.is_empty() {
406 return Ok(Arc::new(StructArray::new_empty_fields(len, nulls)));
409 }
410
411 let struct_array = if skip_validation {
412 unsafe { StructArray::new_unchecked(struct_fields.clone(), struct_arrays, nulls) }
414 } else {
415 StructArray::try_new(struct_fields.clone(), struct_arrays, nulls)?
416 };
417
418 Ok(Arc::new(struct_array))
419 }
420
421 fn create_dictionary_array(
424 &self,
425 field_node: &FieldNode,
426 data_type: &DataType,
427 buffers: &[Buffer],
428 value_array: ArrayRef,
429 ) -> Result<ArrayRef, ArrowError> {
430 if let Dictionary(_, _) = *data_type {
431 let null_buffer = (field_node.null_count() > 0).then_some(buffers[0].clone());
432 let builder = ArrayData::builder(data_type.clone())
433 .len(field_node.length() as usize)
434 .add_buffer(buffers[1].clone())
435 .add_child_data(value_array.into_data())
436 .null_bit_buffer(null_buffer)
437 .null_count(field_node.null_count() as usize);
438 self.create_array_from_builder(builder)
439 } else {
440 unreachable!("Cannot create dictionary array from {:?}", data_type)
441 }
442 }
443}
444
445pub struct RecordBatchDecoder<'a> {
451 batch: crate::RecordBatch<'a>,
453 schema: SchemaRef,
455 dictionaries_by_id: &'a HashMap<i64, ArrayRef>,
457 compression: Option<CompressionCodec>,
459 decompression_context: DecompressionContext,
461 version: MetadataVersion,
463 data: &'a Buffer,
465 nodes: VectorIter<'a, FieldNode>,
467 buffers: VectorIter<'a, crate::Buffer>,
469 projection: Option<&'a [usize]>,
472 require_alignment: bool,
475 skip_validation: UnsafeFlag,
479}
480
481impl<'a> RecordBatchDecoder<'a> {
482 pub fn try_new(
484 buf: &'a Buffer,
485 batch: crate::RecordBatch<'a>,
486 schema: SchemaRef,
487 dictionaries_by_id: &'a HashMap<i64, ArrayRef>,
488 metadata: &'a MetadataVersion,
489 ) -> Result<Self, ArrowError> {
490 let buffers = batch.buffers().ok_or_else(|| {
491 ArrowError::IpcError("Unable to get buffers from IPC RecordBatch".to_string())
492 })?;
493 let field_nodes = batch.nodes().ok_or_else(|| {
494 ArrowError::IpcError("Unable to get field nodes from IPC RecordBatch".to_string())
495 })?;
496
497 let batch_compression = batch.compression();
498 let compression = batch_compression
499 .map(|batch_compression| batch_compression.codec().try_into())
500 .transpose()?;
501
502 Ok(Self {
503 batch,
504 schema,
505 dictionaries_by_id,
506 compression,
507 decompression_context: DecompressionContext::new(),
508 version: *metadata,
509 data: buf,
510 nodes: field_nodes.iter(),
511 buffers: buffers.iter(),
512 projection: None,
513 require_alignment: false,
514 skip_validation: UnsafeFlag::new(),
515 })
516 }
517
518 pub fn with_projection(mut self, projection: Option<&'a [usize]>) -> Self {
523 self.projection = projection;
524 self
525 }
526
527 pub fn with_require_alignment(mut self, require_alignment: bool) -> Self {
533 self.require_alignment = require_alignment;
534 self
535 }
536
537 pub fn with_skip_validation(mut self, skip_validation: UnsafeFlag) -> Self {
550 self.skip_validation = skip_validation;
551 self
552 }
553
554 pub fn read_record_batch(mut self) -> Result<RecordBatch, ArrowError> {
561 let mut variadic_counts: VecDeque<i64> = self
562 .batch
563 .variadicBufferCounts()
564 .into_iter()
565 .flatten()
566 .collect();
567
568 let options = RecordBatchOptions::new().with_row_count(Some(self.batch.length() as usize));
569
570 let schema = Arc::clone(&self.schema);
571 if let Some(projection) = self.projection {
572 let mut arrays = Vec::with_capacity(projection.len());
573 for (idx, field) in schema.fields().iter().enumerate() {
575 let mut decoded = None;
577 for (proj_idx, projected_idx) in projection.iter().enumerate() {
578 if *projected_idx == idx {
579 let child = match decoded.clone() {
581 Some(child) => child,
582 None => {
583 let child = self.create_array(field, &mut variadic_counts)?;
584 decoded = Some(Arc::clone(&child));
585 child
586 }
587 };
588 arrays.push((proj_idx, child));
589 }
590 }
591
592 if decoded.is_none() {
593 self.skip_field(field, &mut variadic_counts)?;
594 }
595 }
596
597 arrays.sort_by_key(|t| t.0);
598
599 let schema = Arc::new(schema.project(projection)?);
600 let columns = arrays.into_iter().map(|t| t.1).collect::<Vec<_>>();
601
602 if self.skip_validation.get() {
603 unsafe {
605 Ok(RecordBatch::new_unchecked(
606 schema,
607 columns,
608 self.batch.length() as usize,
609 ))
610 }
611 } else {
612 check_variadic_counts_consumed(&variadic_counts)?;
613 RecordBatch::try_new_with_options(schema, columns, &options)
614 }
615 } else {
616 let mut children = Vec::with_capacity(schema.fields().len());
617 for field in schema.fields() {
619 let child = self.create_array(field, &mut variadic_counts)?;
620 children.push(child);
621 }
622
623 if self.skip_validation.get() {
624 unsafe {
626 Ok(RecordBatch::new_unchecked(
627 schema,
628 children,
629 self.batch.length() as usize,
630 ))
631 }
632 } else {
633 check_variadic_counts_consumed(&variadic_counts)?;
634 RecordBatch::try_new_with_options(schema, children, &options)
635 }
636 }
637 }
638
639 fn next_buffer(&mut self) -> Result<Buffer, ArrowError> {
640 let buffer = self.buffers.next().ok_or_else(|| {
641 ArrowError::IpcError("Buffer count mismatched with metadata".to_string())
642 })?;
643 read_buffer(
644 buffer,
645 self.data,
646 self.compression,
647 &mut self.decompression_context,
648 )
649 }
650
651 fn skip_buffer(&mut self) -> Result<(), ArrowError> {
652 self.buffers.next().ok_or_else(|| {
653 ArrowError::IpcError("Buffer count mismatched with metadata".to_string())
654 })?;
655 Ok(())
656 }
657
658 fn next_node(&mut self, field: &Field) -> Result<&'a FieldNode, ArrowError> {
659 self.nodes.next().ok_or_else(|| {
660 ArrowError::SchemaError(format!(
661 "Invalid data for schema. {field} refers to node not found in schema",
662 ))
663 })
664 }
665
666 fn skip_field(
667 &mut self,
668 field: &Field,
669 variadic_count: &mut VecDeque<i64>,
670 ) -> Result<(), ArrowError> {
671 self.next_node(field)?;
672
673 match field.data_type() {
674 Utf8 | Binary | LargeBinary | LargeUtf8 => {
675 for _ in 0..3 {
676 self.skip_buffer()?;
677 }
678 }
679 Utf8View | BinaryView => {
680 let count = self.next_variadic_buffer_count(variadic_count, field.data_type())?;
681 for _ in 0..count {
682 self.skip_buffer()?;
683 }
684 }
685 FixedSizeBinary(_) => {
686 self.skip_buffer()?;
687 self.skip_buffer()?;
688 }
689 List(list_field) | LargeList(list_field) | Map(list_field, _) => {
690 self.skip_buffer()?;
691 self.skip_buffer()?;
692 self.skip_field(list_field, variadic_count)?;
693 }
694 ListView(list_field) | LargeListView(list_field) => {
695 self.skip_buffer()?; self.skip_buffer()?; self.skip_buffer()?; self.skip_field(list_field, variadic_count)?;
699 }
700 FixedSizeList(list_field, _) => {
701 self.skip_buffer()?;
702 self.skip_field(list_field, variadic_count)?;
703 }
704 Struct(struct_fields) => {
705 self.skip_buffer()?;
706
707 for struct_field in struct_fields {
709 self.skip_field(struct_field, variadic_count)?
710 }
711 }
712 RunEndEncoded(run_ends_field, values_field) => {
713 self.skip_field(run_ends_field, variadic_count)?;
714 self.skip_field(values_field, variadic_count)?;
715 }
716 Dictionary(_, _) => {
717 self.skip_buffer()?; self.skip_buffer()?; }
720 Union(fields, mode) => {
721 if self.version < MetadataVersion::V5 {
722 self.skip_buffer()?; }
724 self.skip_buffer()?; match mode {
727 UnionMode::Dense => self.skip_buffer()?, UnionMode::Sparse => {}
729 }
730
731 for (_, field) in fields.iter() {
732 self.skip_field(field, variadic_count)?
733 }
734 }
735 Null => {}
737
738 Boolean
740 | Int8
741 | Int16
742 | Int32
743 | Int64
744 | UInt8
745 | UInt16
746 | UInt32
747 | UInt64
748 | Float16
749 | Float32
750 | Float64
751 | Timestamp(_, _)
752 | Date32
753 | Date64
754 | Time32(_)
755 | Time64(_)
756 | Duration(_)
757 | Interval(_)
758 | Decimal32(_, _)
759 | Decimal64(_, _)
760 | Decimal128(_, _)
761 | Decimal256(_, _) => {
762 self.skip_buffer()?;
763 self.skip_buffer()?;
764 }
765 }
766 Ok(())
767 }
768}
769
770impl RecordBatchDecoder<'_> {
771 fn next_variadic_buffer_count(
777 &self,
778 variadic_counts: &mut VecDeque<i64>,
779 data_type: &DataType,
780 ) -> Result<usize, ArrowError> {
781 let count = variadic_counts.pop_front().ok_or_else(|| {
782 ArrowError::IpcError(format!("Missing variadic count for {data_type} column"))
783 })?;
784
785 let remaining = self.buffers.len();
786
787 usize::try_from(count)
789 .ok()
790 .and_then(|count| count.checked_add(2))
791 .filter(|total| *total <= remaining)
792 .ok_or_else(|| {
793 ArrowError::IpcError(format!(
794 "Invalid variadic count {count} for {data_type} column, \
795 with {remaining} buffer(s) left in the message"
796 ))
797 })
798 }
799}
800
801pub fn read_record_batch(
812 buf: &Buffer,
813 batch: crate::RecordBatch,
814 schema: SchemaRef,
815 dictionaries_by_id: &HashMap<i64, ArrayRef>,
816 projection: Option<&[usize]>,
817 metadata: &MetadataVersion,
818) -> Result<RecordBatch, ArrowError> {
819 RecordBatchDecoder::try_new(buf, batch, schema, dictionaries_by_id, metadata)?
820 .with_projection(projection)
821 .with_require_alignment(false)
822 .read_record_batch()
823}
824
825pub fn read_dictionary(
828 buf: &Buffer,
829 batch: crate::DictionaryBatch,
830 schema: &Schema,
831 dictionaries_by_id: &mut HashMap<i64, ArrayRef>,
832 metadata: &MetadataVersion,
833) -> Result<(), ArrowError> {
834 read_dictionary_impl(
835 buf,
836 batch,
837 schema,
838 dictionaries_by_id,
839 metadata,
840 false,
841 UnsafeFlag::new(),
842 )
843}
844
845pub fn read_dictionary_impl(
847 buf: &Buffer,
848 batch: crate::DictionaryBatch,
849 schema: &Schema,
850 dictionaries_by_id: &mut HashMap<i64, ArrayRef>,
851 metadata: &MetadataVersion,
852 require_alignment: bool,
853 skip_validation: UnsafeFlag,
854) -> Result<(), ArrowError> {
855 let id = batch.id();
856
857 let dictionary_values = get_dictionary_values(
858 buf,
859 batch,
860 schema,
861 dictionaries_by_id,
862 metadata,
863 require_alignment,
864 skip_validation,
865 )?;
866
867 update_dictionaries(dictionaries_by_id, batch.isDelta(), id, dictionary_values)?;
868
869 Ok(())
870}
871
872fn update_dictionaries(
881 dictionaries_by_id: &mut HashMap<i64, ArrayRef>,
882 is_delta: bool,
883 dict_id: i64,
884 dict_values: ArrayRef,
885) -> Result<(), ArrowError> {
886 if !is_delta {
887 dictionaries_by_id.insert(dict_id, dict_values.clone());
891 return Ok(());
892 }
893
894 let existing = dictionaries_by_id.get(&dict_id).ok_or_else(|| {
895 ArrowError::InvalidArgumentError(format!(
896 "No existing dictionary for delta dictionary with id '{dict_id}'"
897 ))
898 })?;
899
900 let combined = concat::concat(&[existing, &dict_values]).map_err(|e| {
901 ArrowError::InvalidArgumentError(format!("Failed to concat delta dictionary: {e}"))
902 })?;
903
904 dictionaries_by_id.insert(dict_id, combined);
905
906 Ok(())
907}
908
909fn get_dictionary_values(
913 buf: &Buffer,
914 batch: crate::DictionaryBatch,
915 schema: &Schema,
916 dictionaries_by_id: &HashMap<i64, ArrayRef>,
917 metadata: &MetadataVersion,
918 require_alignment: bool,
919 skip_validation: UnsafeFlag,
920) -> Result<ArrayRef, ArrowError> {
921 let id = batch.id();
922 #[expect(deprecated)]
923 let fields_using_this_dictionary = schema.fields_with_dict_id(id);
924 let first_field = fields_using_this_dictionary.first().ok_or_else(|| {
925 ArrowError::InvalidArgumentError(format!("dictionary id {id} not found in schema"))
926 })?;
927
928 let dictionary_values: ArrayRef = match first_field.data_type() {
932 DataType::Dictionary(_, value_type) => {
933 let value = value_type.as_ref().clone();
935 let schema = Schema::new(vec![Field::new("", value, true)]);
936 let Some(data) = batch.data() else {
938 return Err(ArrowError::ParseError(
939 "Dictionary batch is missing its data".to_string(),
940 ));
941 };
942 let record_batch = RecordBatchDecoder::try_new(
943 buf,
944 data,
945 Arc::new(schema),
946 dictionaries_by_id,
947 metadata,
948 )?
949 .with_require_alignment(require_alignment)
950 .with_skip_validation(skip_validation)
951 .read_record_batch()?;
952
953 Some(record_batch.column(0).clone())
954 }
955 _ => None,
956 }
957 .ok_or_else(|| {
958 ArrowError::InvalidArgumentError(format!("dictionary id {id} not found in schema"))
959 })?;
960
961 Ok(dictionary_values)
962}
963
964fn read_block<R: Read + Seek>(mut reader: R, block: &Block) -> Result<Buffer, ArrowError> {
966 reader.seek(SeekFrom::Start(block.offset() as u64))?;
967 let body_len = block.bodyLength().to_usize().unwrap();
968 let metadata_len = block.metaDataLength().to_usize().unwrap();
969 let total_len = body_len.checked_add(metadata_len).unwrap();
970
971 let mut buf = MutableBuffer::try_from_len_zeroed(total_len)
972 .map_err(|e| ArrowError::MemoryError(e.to_string()))?;
973 reader.read_exact(&mut buf)?;
974 Ok(buf.into())
975}
976
977fn check_variadic_counts_consumed(variadic_counts: &VecDeque<i64>) -> Result<(), ArrowError> {
982 if variadic_counts.is_empty() {
983 Ok(())
984 } else {
985 Err(ArrowError::IpcError(format!(
986 "Mismatch between schema and data: the IPC message declares {} more variadic \
987 buffer count(s) than the schema has BinaryView or Utf8View columns",
988 variadic_counts.len()
989 )))
990 }
991}
992
993fn parse_message(buf: &[u8]) -> Result<Message::Message<'_>, ArrowError> {
997 let buf = match buf[..4] == CONTINUATION_MARKER {
998 true => &buf[8..],
999 false => &buf[4..],
1000 };
1001 crate::root_as_message(buf)
1002 .map_err(|err| ArrowError::ParseError(format!("Unable to get root as message: {err:?}")))
1003}
1004
1005pub fn read_footer_length(buf: [u8; 10]) -> Result<usize, ArrowError> {
1009 if buf[4..] != super::ARROW_MAGIC {
1010 return Err(ArrowError::ParseError(
1011 "Arrow file does not contain correct footer".to_string(),
1012 ));
1013 }
1014
1015 let footer_len = i32::from_le_bytes([buf[0], buf[1], buf[2], buf[3]]);
1017 footer_len
1018 .try_into()
1019 .map_err(|_| ArrowError::ParseError(format!("Invalid footer length: {footer_len}")))
1020}
1021
1022#[derive(Debug)]
1087pub struct FileDecoder {
1088 schema: SchemaRef,
1089 dictionaries: HashMap<i64, ArrayRef>,
1090 version: MetadataVersion,
1091 projection: Option<Vec<usize>>,
1092 require_alignment: bool,
1093 skip_validation: UnsafeFlag,
1094}
1095
1096impl FileDecoder {
1097 pub fn new(schema: SchemaRef, version: MetadataVersion) -> Self {
1099 Self {
1100 schema,
1101 version,
1102 dictionaries: Default::default(),
1103 projection: None,
1104 require_alignment: false,
1105 skip_validation: UnsafeFlag::new(),
1106 }
1107 }
1108
1109 pub fn with_projection(mut self, projection: Vec<usize>) -> Self {
1111 self.projection = Some(projection);
1112 self
1113 }
1114
1115 pub fn with_require_alignment(mut self, require_alignment: bool) -> Self {
1128 self.require_alignment = require_alignment;
1129 self
1130 }
1131
1132 pub unsafe fn with_skip_validation(mut self, skip_validation: bool) -> Self {
1143 unsafe { self.skip_validation.set(skip_validation) };
1144 self
1145 }
1146
1147 fn read_message<'a>(&self, buf: &'a [u8]) -> Result<Message::Message<'a>, ArrowError> {
1148 let message = parse_message(buf)?;
1149
1150 if self.version != MetadataVersion::V1 && message.version() != self.version {
1152 return Err(ArrowError::IpcError(
1153 "Could not read IPC message as metadata versions mismatch".to_string(),
1154 ));
1155 }
1156 Ok(message)
1157 }
1158
1159 pub fn read_dictionary(&mut self, block: &Block, buf: &Buffer) -> Result<(), ArrowError> {
1161 let message = self.read_message(buf)?;
1162 match message.header_type() {
1163 crate::MessageHeader::DictionaryBatch => {
1164 let batch = message.header_as_dictionary_batch().unwrap();
1165 read_dictionary_impl(
1166 &buf.slice(block.metaDataLength() as _),
1167 batch,
1168 &self.schema,
1169 &mut self.dictionaries,
1170 &message.version(),
1171 self.require_alignment,
1172 self.skip_validation.clone(),
1173 )
1174 }
1175 t => Err(ArrowError::ParseError(format!(
1176 "Expecting DictionaryBatch in dictionary blocks, found {t:?}."
1177 ))),
1178 }
1179 }
1180
1181 pub fn read_record_batch(
1183 &self,
1184 block: &Block,
1185 buf: &Buffer,
1186 ) -> Result<Option<RecordBatch>, ArrowError> {
1187 let message = self.read_message(buf)?;
1188 match message.header_type() {
1189 crate::MessageHeader::Schema => Err(ArrowError::IpcError(
1190 "Not expecting a schema when messages are read".to_string(),
1191 )),
1192 crate::MessageHeader::RecordBatch => {
1193 let batch = message.header_as_record_batch().ok_or_else(|| {
1194 ArrowError::IpcError("Unable to read IPC message as record batch".to_string())
1195 })?;
1196 RecordBatchDecoder::try_new(
1198 &buf.slice(block.metaDataLength() as _),
1199 batch,
1200 self.schema.clone(),
1201 &self.dictionaries,
1202 &message.version(),
1203 )?
1204 .with_projection(self.projection.as_deref())
1205 .with_require_alignment(self.require_alignment)
1206 .with_skip_validation(self.skip_validation.clone())
1207 .read_record_batch()
1208 .map(Some)
1209 }
1210 crate::MessageHeader::NONE => Ok(None),
1211 t => Err(ArrowError::InvalidArgumentError(format!(
1212 "Reading types other than record batches not yet supported, unable to read {t:?}"
1213 ))),
1214 }
1215 }
1216}
1217
1218#[derive(Debug)]
1220pub struct FileReaderBuilder {
1221 projection: Option<Vec<usize>>,
1223 max_footer_fb_tables: usize,
1225 max_footer_fb_depth: usize,
1227}
1228
1229impl Default for FileReaderBuilder {
1230 fn default() -> Self {
1231 let verifier_options = VerifierOptions::default();
1232 Self {
1233 max_footer_fb_tables: verifier_options.max_tables,
1234 max_footer_fb_depth: verifier_options.max_depth,
1235 projection: None,
1236 }
1237 }
1238}
1239
1240impl FileReaderBuilder {
1241 pub fn new() -> Self {
1245 Self::default()
1246 }
1247
1248 pub fn with_projection(mut self, projection: Vec<usize>) -> Self {
1250 self.projection = Some(projection);
1251 self
1252 }
1253
1254 pub fn with_max_footer_fb_tables(mut self, max_footer_fb_tables: usize) -> Self {
1267 self.max_footer_fb_tables = max_footer_fb_tables;
1268 self
1269 }
1270
1271 pub fn with_max_footer_fb_depth(mut self, max_footer_fb_depth: usize) -> Self {
1284 self.max_footer_fb_depth = max_footer_fb_depth;
1285 self
1286 }
1287
1288 pub fn build<R: Read + Seek>(self, mut reader: R) -> Result<FileReader<R>, ArrowError> {
1290 let mut buffer = [0; 10];
1292 reader.seek(SeekFrom::End(-10))?;
1293 reader.read_exact(&mut buffer)?;
1294
1295 let footer_len = read_footer_length(buffer)?;
1296
1297 let mut footer_data = vec![0; footer_len];
1299 reader.seek(SeekFrom::End(-10 - footer_len as i64))?;
1300 reader.read_exact(&mut footer_data)?;
1301
1302 let verifier_options = VerifierOptions {
1303 max_tables: self.max_footer_fb_tables,
1304 max_depth: self.max_footer_fb_depth,
1305 ..Default::default()
1306 };
1307 let footer = crate::root_as_footer_with_opts(&verifier_options, &footer_data[..]).map_err(
1308 |err| ArrowError::ParseError(format!("Unable to get root as footer: {err:?}")),
1309 )?;
1310
1311 let blocks = footer.recordBatches().ok_or_else(|| {
1312 ArrowError::ParseError("Unable to get record batches from IPC Footer".to_string())
1313 })?;
1314
1315 let total_blocks = blocks.len();
1316
1317 let ipc_schema = footer.schema().ok_or_else(|| {
1318 ArrowError::ParseError("Unable to get schema from IPC Footer".to_string())
1319 })?;
1320 if !ipc_schema.endianness().equals_to_target_endianness() {
1321 return Err(ArrowError::IpcError(
1322 "the endianness of the source system does not match the endianness of the target system.".to_owned()
1323 ));
1324 }
1325
1326 let schema = Arc::new(crate::convert::try_fb_to_schema(ipc_schema)?);
1327
1328 let projected_schema = match &self.projection {
1329 Some(projection) => Arc::new(schema.project(projection)?),
1330 None => schema.clone(),
1331 };
1332
1333 let mut custom_metadata = HashMap::new();
1334 if let Some(fb_custom_metadata) = footer.custom_metadata() {
1335 for kv in fb_custom_metadata {
1336 let (Some(key), Some(value)) = (kv.key(), kv.value()) else {
1337 return Err(ArrowError::ParseError(
1338 "Custom metadata in the IPC footer is missing a key or a value".to_string(),
1339 ));
1340 };
1341 custom_metadata.insert(key.to_string(), value.to_string());
1342 }
1343 }
1344
1345 let mut decoder = FileDecoder::new(schema, footer.version());
1346 if let Some(projection) = self.projection {
1347 decoder = decoder.with_projection(projection)
1348 }
1349
1350 if let Some(dictionaries) = footer.dictionaries() {
1352 for block in dictionaries {
1353 let buf = read_block(&mut reader, block)?;
1354 decoder.read_dictionary(block, &buf)?;
1355 }
1356 }
1357
1358 Ok(FileReader {
1359 reader,
1360 blocks: blocks.iter().copied().collect(),
1361 current_block: 0,
1362 total_blocks,
1363 decoder,
1364 schema: projected_schema,
1365 custom_metadata,
1366 })
1367 }
1368}
1369
1370pub struct FileReader<R> {
1415 reader: R,
1417
1418 decoder: FileDecoder,
1420
1421 schema: SchemaRef,
1423
1424 blocks: Vec<Block>,
1428
1429 current_block: usize,
1431
1432 total_blocks: usize,
1434
1435 custom_metadata: HashMap<String, String>,
1437}
1438
1439impl<R> fmt::Debug for FileReader<R> {
1440 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> Result<(), fmt::Error> {
1441 f.debug_struct("FileReader<R>")
1442 .field("decoder", &self.decoder)
1443 .field("blocks", &self.blocks)
1444 .field("current_block", &self.current_block)
1445 .field("total_blocks", &self.total_blocks)
1446 .finish_non_exhaustive()
1447 }
1448}
1449
1450impl<R: Read + Seek> FileReader<BufReader<R>> {
1451 pub fn try_new_buffered(reader: R, projection: Option<Vec<usize>>) -> Result<Self, ArrowError> {
1455 Self::try_new(BufReader::new(reader), projection)
1456 }
1457}
1458
1459impl<R: Read + Seek> FileReader<R> {
1460 pub fn try_new(reader: R, projection: Option<Vec<usize>>) -> Result<Self, ArrowError> {
1472 let builder = FileReaderBuilder {
1473 projection,
1474 ..Default::default()
1475 };
1476 builder.build(reader)
1477 }
1478
1479 pub fn custom_metadata(&self) -> &HashMap<String, String> {
1481 &self.custom_metadata
1482 }
1483
1484 pub fn num_batches(&self) -> usize {
1486 self.total_blocks
1487 }
1488
1489 pub fn schema(&self) -> SchemaRef {
1491 self.schema.clone()
1492 }
1493
1494 pub fn set_index(&mut self, index: usize) -> Result<(), ArrowError> {
1498 if index >= self.total_blocks {
1499 Err(ArrowError::InvalidArgumentError(format!(
1500 "Cannot set batch to index {} from {} total batches",
1501 index, self.total_blocks
1502 )))
1503 } else {
1504 self.current_block = index;
1505 Ok(())
1506 }
1507 }
1508
1509 fn maybe_next(&mut self) -> Result<Option<RecordBatch>, ArrowError> {
1510 let block = &self.blocks[self.current_block];
1511 self.current_block += 1;
1512
1513 let buffer = read_block(&mut self.reader, block)?;
1515 self.decoder.read_record_batch(block, &buffer)
1516 }
1517
1518 pub fn get_ref(&self) -> &R {
1522 &self.reader
1523 }
1524
1525 pub fn get_mut(&mut self) -> &mut R {
1529 &mut self.reader
1530 }
1531
1532 pub unsafe fn with_skip_validation(mut self, skip_validation: bool) -> Self {
1538 self.decoder = unsafe { self.decoder.with_skip_validation(skip_validation) };
1539 self
1540 }
1541}
1542
1543impl<R: Read + Seek> Iterator for FileReader<R> {
1544 type Item = Result<RecordBatch, ArrowError>;
1545
1546 fn next(&mut self) -> Option<Self::Item> {
1547 if self.current_block < self.total_blocks {
1549 self.maybe_next().transpose()
1550 } else {
1551 None
1552 }
1553 }
1554}
1555
1556impl<R: Read + Seek> RecordBatchReader for FileReader<R> {
1557 fn schema(&self) -> SchemaRef {
1558 self.schema()
1559 }
1560}
1561
1562pub struct StreamReader<R> {
1596 reader: MessageReader<R>,
1598
1599 schema: SchemaRef,
1601
1602 dictionaries_by_id: HashMap<i64, ArrayRef>,
1606
1607 finished: bool,
1611
1612 projection: Option<(Vec<usize>, SchemaRef)>,
1614
1615 skip_validation: UnsafeFlag,
1619}
1620
1621impl<R> fmt::Debug for StreamReader<R> {
1622 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> std::result::Result<(), fmt::Error> {
1623 f.debug_struct("StreamReader<R>")
1624 .field("reader", &"R")
1625 .field("schema", &self.schema)
1626 .field("dictionaries_by_id", &self.dictionaries_by_id)
1627 .field("finished", &self.finished)
1628 .field("projection", &self.projection)
1629 .finish()
1630 }
1631}
1632
1633impl<R: Read> StreamReader<BufReader<R>> {
1634 pub fn try_new_buffered(reader: R, projection: Option<Vec<usize>>) -> Result<Self, ArrowError> {
1638 Self::try_new(BufReader::new(reader), projection)
1639 }
1640}
1641
1642impl<R: Read> StreamReader<R> {
1643 pub fn try_new(
1655 reader: R,
1656 projection: Option<Vec<usize>>,
1657 ) -> Result<StreamReader<R>, ArrowError> {
1658 let mut msg_reader = MessageReader::new(reader);
1659 let message = msg_reader.maybe_next()?;
1660 let Some((message, _)) = message else {
1661 return Err(ArrowError::IpcError(
1662 "Expected schema message, found empty stream.".to_string(),
1663 ));
1664 };
1665
1666 if message.header_type() != Message::MessageHeader::Schema {
1667 return Err(ArrowError::IpcError(format!(
1668 "Expected a schema as the first message in the stream, got: {:?}",
1669 message.header_type()
1670 )));
1671 }
1672
1673 let schema = message.header_as_schema().ok_or_else(|| {
1674 ArrowError::ParseError("Failed to parse schema from message header".to_string())
1675 })?;
1676 let schema = crate::convert::try_fb_to_schema(schema)?;
1677
1678 let dictionaries_by_id = HashMap::new();
1680
1681 let projection = match projection {
1682 Some(projection_indices) => {
1683 let schema = Arc::new(schema.project(&projection_indices)?);
1684 Some((projection_indices, schema))
1685 }
1686 _ => None,
1687 };
1688
1689 Ok(Self {
1690 reader: msg_reader,
1691 schema: Arc::new(schema),
1692 finished: false,
1693 dictionaries_by_id,
1694 projection,
1695 skip_validation: UnsafeFlag::new(),
1696 })
1697 }
1698
1699 pub fn schema(&self) -> SchemaRef {
1701 match &self.projection {
1702 Some((_, projected_schema)) => projected_schema.clone(),
1703 None => self.schema.clone(),
1704 }
1705 }
1706
1707 pub fn is_finished(&self) -> bool {
1709 self.finished
1710 }
1711
1712 fn maybe_next(&mut self) -> Result<Option<RecordBatch>, ArrowError> {
1713 if self.finished {
1714 return Ok(None);
1715 }
1716
1717 loop {
1719 let message = self.next_ipc_message()?;
1720 let Some(message) = message else {
1721 self.finished = true;
1723 return Ok(None);
1724 };
1725
1726 match message {
1727 IpcMessage::Schema(_) => {
1728 return Err(ArrowError::IpcError(
1729 "Expected a record batch, but found a schema".to_string(),
1730 ));
1731 }
1732 IpcMessage::RecordBatch(record_batch) => {
1733 return Ok(Some(record_batch));
1734 }
1735 IpcMessage::DictionaryBatch { .. } => {}
1736 }
1737 }
1738 }
1739
1740 pub(crate) fn next_ipc_message(&mut self) -> Result<Option<IpcMessage>, ArrowError> {
1748 let message = self.reader.maybe_next()?;
1749 let Some((message, body)) = message else {
1750 return Ok(None);
1752 };
1753
1754 let ipc_message = match message.header_type() {
1755 Message::MessageHeader::Schema => {
1756 let schema = message.header_as_schema().ok_or_else(|| {
1757 ArrowError::ParseError("Failed to parse schema from message header".to_string())
1758 })?;
1759 let arrow_schema = crate::convert::try_fb_to_schema(schema)?;
1760 IpcMessage::Schema(arrow_schema)
1761 }
1762 Message::MessageHeader::RecordBatch => {
1763 let batch = message.header_as_record_batch().ok_or_else(|| {
1764 ArrowError::IpcError("Unable to read IPC message as record batch".to_string())
1765 })?;
1766
1767 let version = message.version();
1768 let schema = self.schema.clone();
1769 let record_batch = RecordBatchDecoder::try_new(
1770 &body.into(),
1771 batch,
1772 schema,
1773 &self.dictionaries_by_id,
1774 &version,
1775 )?
1776 .with_projection(self.projection.as_ref().map(|x| x.0.as_ref()))
1777 .with_require_alignment(false)
1778 .with_skip_validation(self.skip_validation.clone())
1779 .read_record_batch()?;
1780 IpcMessage::RecordBatch(record_batch)
1781 }
1782 Message::MessageHeader::DictionaryBatch => {
1783 let dict = message.header_as_dictionary_batch().ok_or_else(|| {
1784 ArrowError::ParseError(
1785 "Failed to parse dictionary batch from message header".to_string(),
1786 )
1787 })?;
1788
1789 let version = message.version();
1790 let dict_values = get_dictionary_values(
1791 &body.into(),
1792 dict,
1793 &self.schema,
1794 &self.dictionaries_by_id,
1795 &version,
1796 false,
1797 self.skip_validation.clone(),
1798 )?;
1799
1800 update_dictionaries(
1801 &mut self.dictionaries_by_id,
1802 dict.isDelta(),
1803 dict.id(),
1804 dict_values.clone(),
1805 )?;
1806
1807 IpcMessage::DictionaryBatch {
1808 id: dict.id(),
1809 is_delta: (dict.isDelta()),
1810 values: (dict_values),
1811 }
1812 }
1813 x => {
1814 return Err(ArrowError::ParseError(format!(
1815 "Unsupported message header type in IPC stream: '{x:?}'"
1816 )));
1817 }
1818 };
1819
1820 Ok(Some(ipc_message))
1821 }
1822
1823 pub fn get_ref(&self) -> &R {
1827 self.reader.inner()
1828 }
1829
1830 pub fn get_mut(&mut self) -> &mut R {
1834 self.reader.inner_mut()
1835 }
1836
1837 pub unsafe fn with_skip_validation(mut self, skip_validation: bool) -> Self {
1843 unsafe { self.skip_validation.set(skip_validation) };
1844 self
1845 }
1846}
1847
1848impl<R: Read> Iterator for StreamReader<R> {
1849 type Item = Result<RecordBatch, ArrowError>;
1850
1851 fn next(&mut self) -> Option<Self::Item> {
1852 self.maybe_next().transpose()
1853 }
1854}
1855
1856impl<R: Read> RecordBatchReader for StreamReader<R> {
1857 fn schema(&self) -> SchemaRef {
1858 self.schema()
1859 }
1860}
1861
1862#[derive(Debug)]
1868#[expect(dead_code)]
1869pub(crate) enum IpcMessage {
1870 Schema(arrow_schema::Schema),
1871 RecordBatch(RecordBatch),
1872 DictionaryBatch {
1873 id: i64,
1874 is_delta: bool,
1875 values: ArrayRef,
1876 },
1877}
1878
1879const MAX_PREALLOC_BYTES: usize = 64 * 1024 * 1024;
1895
1896fn read_body_bounded<R: Read>(reader: &mut R, len: usize) -> Result<MutableBuffer, ArrowError> {
1898 let mut buf = MutableBuffer::try_from_len_zeroed(len.min(MAX_PREALLOC_BYTES))
1899 .map_err(|e| ArrowError::MemoryError(e.to_string()))?;
1900 let mut filled = 0;
1901 while filled < len {
1902 let target = buf.len();
1903 reader.read_exact(&mut buf.as_slice_mut()[filled..target])?;
1904 filled = target;
1905 if filled < len {
1906 buf.try_resize(len.min(target.saturating_mul(2)), 0)
1907 .map_err(|e| ArrowError::MemoryError(e.to_string()))?;
1908 }
1909 }
1910 Ok(buf)
1911}
1912
1913struct MessageReader<R> {
1916 reader: R,
1917 buf: Vec<u8>,
1918}
1919
1920impl<R: Read> MessageReader<R> {
1921 fn new(reader: R) -> Self {
1922 Self {
1923 reader,
1924 buf: Vec::new(),
1925 }
1926 }
1927
1928 fn maybe_next(&mut self) -> Result<Option<(Message::Message<'_>, MutableBuffer)>, ArrowError> {
1939 let meta_len = self.read_meta_len()?;
1940 let Some(meta_len) = meta_len else {
1941 return Ok(None);
1942 };
1943
1944 self.buf.clear();
1947 let read = (&mut self.reader)
1948 .take(meta_len as u64)
1949 .read_to_end(&mut self.buf)?;
1950 if read != meta_len {
1951 return Err(ArrowError::ParseError(format!(
1952 "Unexpected end of stream: expected {meta_len} metadata bytes, got {read}"
1953 )));
1954 }
1955
1956 let message = crate::root_as_message(self.buf.as_slice()).map_err(|err| {
1957 ArrowError::ParseError(format!("Unable to get root as message: {err:?}"))
1958 })?;
1959
1960 let body_len = usize::try_from(message.bodyLength()).map_err(|_| {
1961 ArrowError::ParseError(format!(
1962 "Invalid IPC message body length: {}",
1963 message.bodyLength()
1964 ))
1965 })?;
1966 let buf = read_body_bounded(&mut self.reader, body_len)?;
1967
1968 Ok(Some((message, buf)))
1969 }
1970
1971 fn inner_mut(&mut self) -> &mut R {
1973 &mut self.reader
1974 }
1975
1976 fn inner(&self) -> &R {
1978 &self.reader
1979 }
1980
1981 pub fn read_meta_len(&mut self) -> Result<Option<usize>, ArrowError> {
1990 let mut meta_len: [u8; 4] = [0; 4];
1991 match self.reader.read_exact(&mut meta_len) {
1992 Ok(()) => {}
1993 Err(e) => {
1994 return if e.kind() == std::io::ErrorKind::UnexpectedEof {
1995 Ok(None)
1999 } else {
2000 Err(ArrowError::from(e))
2001 };
2002 }
2003 }
2004
2005 let meta_len = {
2006 if meta_len == CONTINUATION_MARKER {
2009 self.reader.read_exact(&mut meta_len)?;
2010 }
2011
2012 i32::from_le_bytes(meta_len)
2013 };
2014
2015 if meta_len == 0 {
2016 return Ok(None);
2017 }
2018
2019 let meta_len = usize::try_from(meta_len)
2020 .map_err(|_| ArrowError::ParseError(format!("Invalid metadata length: {meta_len}")))?;
2021
2022 Ok(Some(meta_len))
2023 }
2024}
2025
2026#[cfg(test)]
2027mod tests {
2028 use std::io::Cursor;
2029
2030 use crate::convert::try_fb_to_schema;
2031 use crate::writer::{
2032 DictionaryTracker, IpcDataGenerator, IpcWriteOptions, unslice_run_array, write_message,
2033 };
2034
2035 use super::*;
2036
2037 use crate::{root_as_footer, root_as_message, size_prefixed_root_as_message};
2038 use arrow_array::builder::{PrimitiveRunBuilder, UnionBuilder};
2039 use arrow_array::types::*;
2040 use arrow_buffer::{NullBuffer, OffsetBuffer};
2041 use arrow_data::ArrayDataBuilder;
2042
2043 fn create_test_projection_schema() -> Schema {
2044 let list_data_type = DataType::List(Arc::new(Field::new_list_field(DataType::Int32, true)));
2046
2047 let fixed_size_list_data_type =
2048 DataType::FixedSizeList(Arc::new(Field::new_list_field(DataType::Int32, false)), 3);
2049
2050 let union_fields = UnionFields::from_fields(vec![
2051 Field::new("a", DataType::Int32, false),
2052 Field::new("b", DataType::Float64, false),
2053 ]);
2054
2055 let union_data_type = DataType::Union(union_fields, UnionMode::Dense);
2056
2057 let struct_fields = Fields::from(vec![
2058 Field::new("id", DataType::Int32, false),
2059 Field::new_list("list", Field::new_list_field(DataType::Int8, true), false),
2060 ]);
2061 let struct_data_type = DataType::Struct(struct_fields);
2062
2063 let run_encoded_data_type = DataType::RunEndEncoded(
2064 Arc::new(Field::new(
2065 Field::REE_RUN_ENDS_FIELD_DEFAULT_NAME,
2066 DataType::Int16,
2067 false,
2068 )),
2069 Arc::new(Field::new(
2070 Field::REE_VALUES_FIELD_DEFAULT_NAME,
2071 DataType::Int32,
2072 true,
2073 )),
2074 );
2075
2076 Schema::new(vec![
2078 Field::new("f0", DataType::UInt32, false),
2079 Field::new("f1", DataType::Utf8, false),
2080 Field::new("f2", DataType::Boolean, false),
2081 Field::new("f3", union_data_type, true),
2082 Field::new("f4", DataType::Null, true),
2083 Field::new("f5", DataType::Float64, true),
2084 Field::new("f6", list_data_type, false),
2085 Field::new("f7", DataType::FixedSizeBinary(3), true),
2086 Field::new("f8", fixed_size_list_data_type, false),
2087 Field::new("f9", struct_data_type, false),
2088 Field::new("f10", run_encoded_data_type, false),
2089 Field::new("f11", DataType::Boolean, false),
2090 Field::new_dictionary("f12", DataType::Int8, DataType::Utf8, false),
2091 Field::new("f13", DataType::Utf8, false),
2092 ])
2093 }
2094
2095 fn create_test_projection_batch_data(schema: &Schema) -> RecordBatch {
2096 let array0 = UInt32Array::from(vec![1, 2, 3]);
2098 let array1 = StringArray::from(vec!["foo", "bar", "baz"]);
2099 let array2 = BooleanArray::from(vec![true, false, true]);
2100
2101 let mut union_builder = UnionBuilder::new_dense();
2102 union_builder.append::<Int32Type>("a", 1).unwrap();
2103 union_builder.append::<Float64Type>("b", 10.1).unwrap();
2104 union_builder.append_null::<Float64Type>("b").unwrap();
2105 let array3 = union_builder.build().unwrap();
2106
2107 let array4 = NullArray::new(3);
2108 let array5 = Float64Array::from(vec![Some(1.1), None, Some(3.3)]);
2109 let array6_values = vec![
2110 Some(vec![Some(10), Some(10), Some(10)]),
2111 Some(vec![Some(20), Some(20), Some(20)]),
2112 Some(vec![Some(30), Some(30)]),
2113 ];
2114 let array6 = ListArray::from_iter_primitive::<Int32Type, _, _>(array6_values);
2115 let array7_values = vec![vec![11, 12, 13], vec![22, 23, 24], vec![33, 34, 35]];
2116 let array7 = FixedSizeBinaryArray::try_from_iter(array7_values.into_iter()).unwrap();
2117
2118 let array8_values = ArrayData::builder(DataType::Int32)
2119 .len(9)
2120 .add_buffer(Buffer::from_slice_ref([40, 41, 42, 43, 44, 45, 46, 47, 48]))
2121 .build()
2122 .unwrap();
2123 let array8_data = ArrayData::builder(schema.field(8).data_type().clone())
2124 .len(3)
2125 .add_child_data(array8_values)
2126 .build()
2127 .unwrap();
2128 let array8 = FixedSizeListArray::from(array8_data);
2129
2130 let array9_id: ArrayRef = Arc::new(Int32Array::from(vec![1001, 1002, 1003]));
2131 let array9_list: ArrayRef =
2132 Arc::new(ListArray::from_iter_primitive::<Int8Type, _, _>(vec![
2133 Some(vec![Some(-10)]),
2134 Some(vec![Some(-20), Some(-20), Some(-20)]),
2135 Some(vec![Some(-30)]),
2136 ]));
2137 let array9 = ArrayDataBuilder::new(schema.field(9).data_type().clone())
2138 .add_child_data(array9_id.into_data())
2139 .add_child_data(array9_list.into_data())
2140 .len(3)
2141 .build()
2142 .unwrap();
2143 let array9 = StructArray::from(array9);
2144
2145 let array10_input = vec![Some(1_i32), None, None];
2146 let mut array10_builder = PrimitiveRunBuilder::<Int16Type, Int32Type>::new();
2147 array10_builder.extend(array10_input);
2148 let array10 = array10_builder.finish();
2149
2150 let array11 = BooleanArray::from(vec![false, false, true]);
2151
2152 let array12_values = StringArray::from(vec!["x", "yy", "zzz"]);
2153 let array12_keys = Int8Array::from_iter_values([1, 1, 2]);
2154 let array12 = DictionaryArray::new(array12_keys, Arc::new(array12_values));
2155
2156 let array13 = StringArray::from(vec!["a", "bb", "ccc"]);
2157
2158 RecordBatch::try_new(
2160 Arc::new(schema.clone()),
2161 vec![
2162 Arc::new(array0),
2163 Arc::new(array1),
2164 Arc::new(array2),
2165 Arc::new(array3),
2166 Arc::new(array4),
2167 Arc::new(array5),
2168 Arc::new(array6),
2169 Arc::new(array7),
2170 Arc::new(array8),
2171 Arc::new(array9),
2172 Arc::new(array10),
2173 Arc::new(array11),
2174 Arc::new(array12),
2175 Arc::new(array13),
2176 ],
2177 )
2178 .unwrap()
2179 }
2180
2181 #[test]
2182 fn test_negative_meta_len_start_stream() {
2183 let bytes = i32::to_le_bytes(-1);
2184 let mut buf = vec![];
2185 buf.extend(CONTINUATION_MARKER);
2186 buf.extend(bytes);
2187
2188 let reader_err = StreamReader::try_new(Cursor::new(buf), None).err();
2189 assert!(reader_err.is_some());
2190 assert_eq!(
2191 reader_err.unwrap().to_string(),
2192 "Parser error: Invalid metadata length: -1"
2193 );
2194 }
2195
2196 #[test]
2197 fn test_negative_meta_len_mid_stream() {
2198 let schema = Schema::new(vec![Field::new("a", DataType::Int32, false)]);
2199 let mut buf = Vec::new();
2200 {
2201 let mut writer = crate::writer::StreamWriter::try_new(&mut buf, &schema).unwrap();
2202 let batch =
2203 RecordBatch::try_new(Arc::new(schema), vec![Arc::new(Int32Array::from(vec![1]))])
2204 .unwrap();
2205 writer.write(&batch).unwrap();
2206 }
2207
2208 let bytes = i32::to_le_bytes(-1);
2209 buf.extend(CONTINUATION_MARKER);
2210 buf.extend(bytes);
2211
2212 let mut reader = StreamReader::try_new(Cursor::new(buf), None).unwrap();
2213 assert!(reader.maybe_next().is_ok());
2215 let batch_err = reader.maybe_next().err();
2217 assert!(batch_err.is_some());
2218 assert_eq!(
2219 batch_err.unwrap().to_string(),
2220 "Parser error: Invalid metadata length: -1"
2221 );
2222 }
2223
2224 #[test]
2225 fn test_invalid_dictionary_batch_without_data() {
2226 use crate::r#gen::Message::*;
2227 use flatbuffers::FlatBufferBuilder;
2228
2229 let schema = Schema::new(vec![Field::new(
2230 "col",
2231 DataType::Dictionary(Box::new(DataType::Int8), Box::new(DataType::Utf8)),
2232 true,
2233 )]);
2234
2235 let mut fbb = FlatBufferBuilder::new();
2238 let batch_offset = DictionaryBatch::create(
2239 &mut fbb,
2240 &DictionaryBatchArgs {
2241 id: 0,
2242 data: None,
2243 isDelta: false,
2244 },
2245 );
2246 fbb.finish_minimal(batch_offset);
2247 let batch_bytes = fbb.finished_data().to_vec();
2248 let batch = flatbuffers::root::<DictionaryBatch>(&batch_bytes).unwrap();
2249
2250 let data_buffer = Buffer::from(vec![0u8; 0]);
2251 let mut dictionaries: HashMap<i64, ArrayRef> = HashMap::new();
2252
2253 let err = read_dictionary(
2254 &data_buffer,
2255 batch,
2256 &schema,
2257 &mut dictionaries,
2258 &MetadataVersion::V5,
2259 )
2260 .unwrap_err();
2261
2262 assert_eq!(
2263 err.to_string(),
2264 "Parser error: Dictionary batch is missing its data"
2265 );
2266 }
2267
2268 #[test]
2269 fn test_missing_buffer_metadata_error() {
2270 use crate::r#gen::Message::*;
2271 use flatbuffers::FlatBufferBuilder;
2272
2273 let schema = Arc::new(Schema::new(vec![Field::new("col", DataType::Int32, true)]));
2274
2275 let mut fbb = FlatBufferBuilder::new();
2278 let nodes = fbb.create_vector(&[FieldNode::new(2, 0)]);
2279 let buffers = fbb.create_vector(&[crate::Buffer::new(0, 8)]);
2280 let batch_offset = RecordBatch::create(
2281 &mut fbb,
2282 &RecordBatchArgs {
2283 length: 2,
2284 nodes: Some(nodes),
2285 buffers: Some(buffers),
2286 compression: None,
2287 variadicBufferCounts: None,
2288 },
2289 );
2290 fbb.finish_minimal(batch_offset);
2291 let batch_bytes = fbb.finished_data().to_vec();
2292 let batch = flatbuffers::root::<RecordBatch>(&batch_bytes).unwrap();
2293
2294 let data_buffer = Buffer::from(vec![0u8; 8]);
2295 let dictionaries: HashMap<i64, ArrayRef> = HashMap::new();
2296 let metadata = MetadataVersion::V5;
2297
2298 let decoder = RecordBatchDecoder::try_new(
2299 &data_buffer,
2300 batch,
2301 schema.clone(),
2302 &dictionaries,
2303 &metadata,
2304 )
2305 .unwrap();
2306
2307 let result = decoder.read_record_batch();
2308
2309 match result {
2310 Err(ArrowError::IpcError(msg)) => {
2311 assert_eq!(msg, "Buffer count mismatched with metadata");
2312 }
2313 other => panic!("unexpected error: {other:?}"),
2314 }
2315 }
2316
2317 fn read_batch_with_variadic_count(count: i64, projection: Option<&[usize]>) -> ArrowError {
2320 use crate::r#gen::Message::*;
2321 use flatbuffers::FlatBufferBuilder;
2322
2323 let schema = Arc::new(Schema::new(vec![Field::new(
2324 "col",
2325 DataType::Utf8View,
2326 true,
2327 )]));
2328
2329 let mut fbb = FlatBufferBuilder::new();
2330 let nodes = fbb.create_vector(&[FieldNode::new(1, 0)]);
2331 let buffers = fbb.create_vector(&[crate::Buffer::new(0, 8), crate::Buffer::new(8, 8)]);
2332 let variadic_buffer_counts = fbb.create_vector(&[count]);
2333 let batch_offset = RecordBatch::create(
2334 &mut fbb,
2335 &RecordBatchArgs {
2336 length: 1,
2337 nodes: Some(nodes),
2338 buffers: Some(buffers),
2339 compression: None,
2340 variadicBufferCounts: Some(variadic_buffer_counts),
2341 },
2342 );
2343 fbb.finish_minimal(batch_offset);
2344 let batch_bytes = fbb.finished_data().to_vec();
2345 let batch = flatbuffers::root::<RecordBatch>(&batch_bytes).unwrap();
2346
2347 let data_buffer = Buffer::from(vec![0u8; 16]);
2348 let dictionaries: HashMap<i64, ArrayRef> = HashMap::new();
2349
2350 RecordBatchDecoder::try_new(
2351 &data_buffer,
2352 batch,
2353 schema,
2354 &dictionaries,
2355 &MetadataVersion::V5,
2356 )
2357 .unwrap()
2358 .with_projection(projection)
2359 .read_record_batch()
2360 .expect_err("should get error")
2361 }
2362
2363 #[test]
2366 fn test_invalid_variadic_buffer_count_error() {
2367 for count in [-2, -1, 1, i64::MAX] {
2370 for projection in [None, Some([].as_slice())] {
2371 let err = read_batch_with_variadic_count(count, projection);
2372 assert_eq!(
2373 err.to_string(),
2374 format!(
2375 "Ipc error: Invalid variadic count {count} for Utf8View column, \
2376 with 2 buffer(s) left in the message"
2377 ),
2378 "count {count}, projection {projection:?}"
2379 );
2380 }
2381 }
2382 }
2383
2384 #[test]
2386 fn test_valid_variadic_buffer_count_is_accepted() {
2387 let err = read_batch_with_variadic_count(0, None);
2388 assert!(!err.to_string().contains("Invalid variadic count"), "{err}");
2389 }
2390
2391 #[test]
2392 fn test_missing_footer_schema_error() {
2393 use crate::r#gen::File::{Footer, FooterArgs};
2394 use flatbuffers::FlatBufferBuilder;
2395
2396 let mut fbb = FlatBufferBuilder::new();
2399 let record_batches = fbb.create_vector::<Block>(&[]);
2400 let footer = Footer::create(
2401 &mut fbb,
2402 &FooterArgs {
2403 version: MetadataVersion::V5,
2404 schema: None,
2405 dictionaries: None,
2406 recordBatches: Some(record_batches),
2407 custom_metadata: None,
2408 },
2409 );
2410 fbb.finish(footer, None);
2411 let footer_data = fbb.finished_data();
2412
2413 let mut buf = Vec::new();
2415 buf.extend_from_slice(&crate::ARROW_MAGIC);
2416 buf.extend_from_slice(footer_data);
2417 buf.extend_from_slice(&(footer_data.len() as i32).to_le_bytes());
2418 buf.extend_from_slice(&crate::ARROW_MAGIC);
2419
2420 let err = FileReader::try_new(Cursor::new(buf), None)
2421 .expect_err("expected an error, not a panic");
2422 assert!(
2423 matches!(err, ArrowError::ParseError(_)),
2424 "expected ParseError, got {err:?}"
2425 );
2426 }
2427
2428 #[test]
2430 fn test_read_legacy_empty_list_without_offsets_buffer() {
2431 use crate::r#gen::Message::*;
2432 use flatbuffers::FlatBufferBuilder;
2433
2434 let schema = Arc::new(Schema::new(vec![Field::new_list(
2435 "items",
2436 Field::new_list_field(DataType::Int32, true),
2437 true,
2438 )]));
2439
2440 let mut fbb = FlatBufferBuilder::new();
2443 let nodes = fbb.create_vector(&[
2444 FieldNode::new(0, 0), FieldNode::new(0, 0), ]);
2447 let buffers = fbb.create_vector(&[
2448 crate::Buffer::new(0, 0), crate::Buffer::new(0, 0), crate::Buffer::new(0, 0), crate::Buffer::new(0, 0), ]);
2453 let batch_offset = RecordBatch::create(
2454 &mut fbb,
2455 &RecordBatchArgs {
2456 length: 0,
2457 nodes: Some(nodes),
2458 buffers: Some(buffers),
2459 compression: None,
2460 variadicBufferCounts: None,
2461 },
2462 );
2463 fbb.finish_minimal(batch_offset);
2464 let batch_bytes = fbb.finished_data().to_vec();
2465 let batch = flatbuffers::root::<RecordBatch>(&batch_bytes).unwrap();
2466
2467 let body = Buffer::from(Vec::<u8>::new());
2468 let dictionaries: HashMap<i64, ArrayRef> = HashMap::new();
2469 let metadata = MetadataVersion::V5;
2470
2471 let decoder =
2472 RecordBatchDecoder::try_new(&body, batch, schema.clone(), &dictionaries, &metadata)
2473 .unwrap();
2474
2475 let read_batch = decoder.read_record_batch().unwrap();
2476 assert_eq!(read_batch.num_rows(), 0);
2477
2478 let list = read_batch
2479 .column(0)
2480 .as_any()
2481 .downcast_ref::<ListArray>()
2482 .unwrap();
2483 assert_eq!(list.len(), 0);
2484 assert_eq!(list.values().len(), 0);
2485 }
2486
2487 #[test]
2489 fn test_read_legacy_empty_utf8_and_binary_without_offsets_buffer() {
2490 use crate::r#gen::Message::*;
2491 use flatbuffers::FlatBufferBuilder;
2492
2493 let schema = Arc::new(Schema::new(vec![
2494 Field::new("name", DataType::Utf8, true),
2495 Field::new("payload", DataType::Binary, true),
2496 ]));
2497
2498 let mut fbb = FlatBufferBuilder::new();
2501 let nodes = fbb.create_vector(&[
2502 FieldNode::new(0, 0), FieldNode::new(0, 0), ]);
2505 let buffers = fbb.create_vector(&[
2506 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), ]);
2513 let batch_offset = RecordBatch::create(
2514 &mut fbb,
2515 &RecordBatchArgs {
2516 length: 0,
2517 nodes: Some(nodes),
2518 buffers: Some(buffers),
2519 compression: None,
2520 variadicBufferCounts: None,
2521 },
2522 );
2523 fbb.finish_minimal(batch_offset);
2524 let batch_bytes = fbb.finished_data().to_vec();
2525 let batch = flatbuffers::root::<RecordBatch>(&batch_bytes).unwrap();
2526
2527 let body = Buffer::from(Vec::<u8>::new());
2528 let dictionaries: HashMap<i64, ArrayRef> = HashMap::new();
2529 let metadata = MetadataVersion::V5;
2530
2531 let decoder =
2532 RecordBatchDecoder::try_new(&body, batch, schema.clone(), &dictionaries, &metadata)
2533 .unwrap();
2534
2535 let read_batch = decoder.read_record_batch().unwrap();
2536 assert_eq!(read_batch.num_rows(), 0);
2537
2538 let utf8 = read_batch
2539 .column(0)
2540 .as_any()
2541 .downcast_ref::<StringArray>()
2542 .unwrap();
2543 assert_eq!(utf8.len(), 0);
2544 assert_eq!(utf8.value_offsets(), [0]);
2545
2546 let binary = read_batch
2547 .column(1)
2548 .as_any()
2549 .downcast_ref::<BinaryArray>()
2550 .unwrap();
2551 assert_eq!(binary.len(), 0);
2552 assert_eq!(binary.value_offsets(), [0]);
2553 }
2554
2555 #[test]
2556 fn test_projection_array_values() {
2557 let schema = create_test_projection_schema();
2559
2560 let batch = create_test_projection_batch_data(&schema);
2562
2563 let mut buf = Vec::new();
2565 {
2566 let mut writer = crate::writer::FileWriter::try_new(&mut buf, &schema).unwrap();
2567 writer.write(&batch).unwrap();
2568 writer.finish().unwrap();
2569 }
2570
2571 for index in 0..12 {
2573 let projection = vec![index];
2574 let reader = FileReader::try_new(std::io::Cursor::new(buf.clone()), Some(projection));
2575 let read_batch = reader.unwrap().next().unwrap().unwrap();
2576 let projected_column = read_batch.column(0);
2577 let expected_column = batch.column(index);
2578
2579 assert_eq!(projected_column.as_ref(), expected_column.as_ref());
2581 }
2582
2583 {
2584 let reader =
2586 FileReader::try_new(std::io::Cursor::new(buf.clone()), Some(vec![3, 2, 1]));
2587 let read_batch = reader.unwrap().next().unwrap().unwrap();
2588 let expected_batch = batch.project(&[3, 2, 1]).unwrap();
2589 assert_eq!(read_batch, expected_batch);
2590 }
2591 }
2592
2593 #[test]
2594 fn test_file_reader_projected_schema_matches_batch_schema() {
2595 let schema = create_test_projection_schema();
2596 let batch = create_test_projection_batch_data(&schema);
2597
2598 let mut buf = Vec::new();
2599 {
2600 let mut writer = crate::writer::FileWriter::try_new(&mut buf, &schema).unwrap();
2601 writer.write(&batch).unwrap();
2602 writer.finish().unwrap();
2603 }
2604
2605 let projection = vec![3, 2, 1];
2606 let mut reader = FileReader::try_new(Cursor::new(buf), Some(projection)).unwrap();
2607 let reader_schema = RecordBatchReader::schema(&reader);
2608 let read_batch = reader.next().unwrap().unwrap();
2609
2610 assert_eq!(reader_schema, read_batch.schema());
2611 }
2612
2613 #[test]
2614 fn test_stream_reader_projected_schema_matches_batch_schema() {
2615 let schema = create_test_projection_schema();
2616 let batch = create_test_projection_batch_data(&schema);
2617
2618 let mut buf = Vec::new();
2619 {
2620 let mut writer = crate::writer::StreamWriter::try_new(&mut buf, &schema).unwrap();
2621 writer.write(&batch).unwrap();
2622 writer.finish().unwrap();
2623 }
2624
2625 let projection = vec![3, 2, 1];
2626 let mut reader = StreamReader::try_new(Cursor::new(buf), Some(projection)).unwrap();
2627 let reader_schema = RecordBatchReader::schema(&reader);
2628 let read_batch = reader.next().unwrap().unwrap();
2629
2630 assert_eq!(reader_schema, read_batch.schema());
2631 }
2632
2633 #[test]
2634 fn test_file_reader_rejects_invalid_projection() {
2635 let schema = create_test_projection_schema();
2636 let batch = create_test_projection_batch_data(&schema);
2637
2638 let mut buf = Vec::new();
2639 {
2640 let mut writer = crate::writer::FileWriter::try_new(&mut buf, &schema).unwrap();
2641 writer.write(&batch).unwrap();
2642 writer.finish().unwrap();
2643 }
2644
2645 let result = FileReader::try_new(Cursor::new(buf), Some(vec![schema.fields().len()]));
2646
2647 assert!(matches!(result, Err(ArrowError::SchemaError(_))));
2648 }
2649
2650 #[test]
2651 fn test_projection_duplicate_indices() {
2652 let schema = create_test_projection_schema();
2653 let batch = create_test_projection_batch_data(&schema);
2654
2655 let mut buf = Vec::new();
2657 {
2658 let mut writer = crate::writer::FileWriter::try_new(&mut buf, &schema).unwrap();
2659 writer.write(&batch).unwrap();
2660 writer.finish().unwrap();
2661 }
2662
2663 for projection in [vec![1, 1], vec![2, 0, 2]] {
2665 let reader =
2666 FileReader::try_new(std::io::Cursor::new(buf.clone()), Some(projection.clone()));
2667 let read_batch = reader.unwrap().next().unwrap().unwrap();
2668
2669 let expected_batch = batch.project(&projection).unwrap();
2670 assert_eq!(read_batch, expected_batch);
2671 }
2672 }
2673
2674 #[test]
2675 fn test_arrow_single_float_row() {
2676 let schema = Schema::new(vec![
2677 Field::new("a", DataType::Float32, false),
2678 Field::new("b", DataType::Float32, false),
2679 Field::new("c", DataType::Int32, false),
2680 Field::new("d", DataType::Int32, false),
2681 ]);
2682 let arrays = vec![
2683 Arc::new(Float32Array::from(vec![1.23])) as ArrayRef,
2684 Arc::new(Float32Array::from(vec![-6.50])) as ArrayRef,
2685 Arc::new(Int32Array::from(vec![2])) as ArrayRef,
2686 Arc::new(Int32Array::from(vec![1])) as ArrayRef,
2687 ];
2688 let batch = RecordBatch::try_new(Arc::new(schema.clone()), arrays).unwrap();
2689 let mut file = tempfile::tempfile().unwrap();
2691 let mut stream_writer = crate::writer::StreamWriter::try_new(&mut file, &schema).unwrap();
2692 stream_writer.write(&batch).unwrap();
2693 stream_writer.finish().unwrap();
2694
2695 drop(stream_writer);
2696
2697 file.rewind().unwrap();
2698
2699 let reader = StreamReader::try_new(&mut file, None).unwrap();
2701
2702 reader.for_each(|batch| {
2703 let batch = batch.unwrap();
2704 assert!(
2705 batch
2706 .column(0)
2707 .as_any()
2708 .downcast_ref::<Float32Array>()
2709 .unwrap()
2710 .value(0)
2711 != 0.0
2712 );
2713 assert!(
2714 batch
2715 .column(1)
2716 .as_any()
2717 .downcast_ref::<Float32Array>()
2718 .unwrap()
2719 .value(0)
2720 != 0.0
2721 );
2722 });
2723
2724 file.rewind().unwrap();
2725
2726 let reader = StreamReader::try_new(file, Some(vec![0, 3])).unwrap();
2728
2729 reader.for_each(|batch| {
2730 let batch = batch.unwrap();
2731 assert_eq!(batch.schema().fields().len(), 2);
2732 assert_eq!(batch.schema().fields()[0].data_type(), &DataType::Float32);
2733 assert_eq!(batch.schema().fields()[1].data_type(), &DataType::Int32);
2734 });
2735 }
2736
2737 fn write_ipc(rb: &RecordBatch) -> Vec<u8> {
2739 let mut buf = Vec::new();
2740 let mut writer = crate::writer::FileWriter::try_new(&mut buf, rb.schema_ref()).unwrap();
2741 writer.write(rb).unwrap();
2742 writer.finish().unwrap();
2743 buf
2744 }
2745
2746 fn read_ipc(buf: &[u8]) -> Result<RecordBatch, ArrowError> {
2748 let mut reader = FileReader::try_new(std::io::Cursor::new(buf), None)?;
2749 reader.next().unwrap()
2750 }
2751
2752 fn read_ipc_skip_validation(buf: &[u8]) -> Result<RecordBatch, ArrowError> {
2755 let mut reader = unsafe {
2756 FileReader::try_new(std::io::Cursor::new(buf), None)?.with_skip_validation(true)
2757 };
2758 reader.next().unwrap()
2759 }
2760
2761 fn roundtrip_ipc(rb: &RecordBatch) -> RecordBatch {
2762 let buf = write_ipc(rb);
2763 read_ipc(&buf).unwrap()
2764 }
2765
2766 fn read_ipc_with_decoder(buf: Vec<u8>) -> Result<RecordBatch, ArrowError> {
2769 read_ipc_with_decoder_inner(buf, false)
2770 }
2771
2772 fn read_ipc_with_decoder_skip_validation(buf: Vec<u8>) -> Result<RecordBatch, ArrowError> {
2775 read_ipc_with_decoder_inner(buf, true)
2776 }
2777
2778 fn read_ipc_with_decoder_inner(
2779 buf: Vec<u8>,
2780 skip_validation: bool,
2781 ) -> Result<RecordBatch, ArrowError> {
2782 let buffer = Buffer::from_vec(buf);
2783 let trailer_start = buffer.len() - 10;
2784 let footer_len = read_footer_length(buffer[trailer_start..].try_into().unwrap())?;
2785 let footer = root_as_footer(&buffer[trailer_start - footer_len..trailer_start])
2786 .map_err(|e| ArrowError::InvalidArgumentError(format!("Invalid footer: {e}")))?;
2787
2788 let schema = try_fb_to_schema(footer.schema().unwrap()).unwrap();
2789
2790 let mut decoder = unsafe {
2791 FileDecoder::new(Arc::new(schema), footer.version())
2792 .with_skip_validation(skip_validation)
2793 };
2794 for block in footer.dictionaries().iter().flatten() {
2796 let block_len = block.bodyLength() as usize + block.metaDataLength() as usize;
2797 let data = buffer.slice_with_length(block.offset() as _, block_len);
2798 decoder.read_dictionary(block, &data)?
2799 }
2800
2801 let batches = footer.recordBatches().unwrap();
2803 assert_eq!(batches.len(), 1); let block = batches.get(0);
2806 let block_len = block.bodyLength() as usize + block.metaDataLength() as usize;
2807 let data = buffer.slice_with_length(block.offset() as _, block_len);
2808 Ok(decoder.read_record_batch(block, &data)?.unwrap())
2809 }
2810
2811 fn write_stream(rb: &RecordBatch) -> Vec<u8> {
2813 let mut buf = Vec::new();
2814 let mut writer = crate::writer::StreamWriter::try_new(&mut buf, rb.schema_ref()).unwrap();
2815 writer.write(rb).unwrap();
2816 writer.finish().unwrap();
2817 buf
2818 }
2819
2820 fn read_stream(buf: &[u8]) -> Result<RecordBatch, ArrowError> {
2822 let mut reader = StreamReader::try_new(std::io::Cursor::new(buf), None)?;
2823 reader.next().unwrap()
2824 }
2825
2826 fn read_stream_skip_validation(buf: &[u8]) -> Result<RecordBatch, ArrowError> {
2829 let mut reader = unsafe {
2830 StreamReader::try_new(std::io::Cursor::new(buf), None)?.with_skip_validation(true)
2831 };
2832 reader.next().unwrap()
2833 }
2834
2835 fn roundtrip_ipc_stream(rb: &RecordBatch) -> RecordBatch {
2836 let buf = write_stream(rb);
2837 read_stream(&buf).unwrap()
2838 }
2839
2840 #[test]
2841 fn test_roundtrip_with_custom_metadata() {
2842 let schema = Schema::new(vec![Field::new("dummy", DataType::Float64, false)]);
2843 let mut buf = Vec::new();
2844 let mut writer = crate::writer::FileWriter::try_new(&mut buf, &schema).unwrap();
2845 let mut test_metadata = HashMap::new();
2846 test_metadata.insert("abc".to_string(), "abc".to_string());
2847 test_metadata.insert("def".to_string(), "def".to_string());
2848 for (k, v) in &test_metadata {
2849 writer.write_metadata(k, v);
2850 }
2851 writer.finish().unwrap();
2852 drop(writer);
2853
2854 let reader = crate::reader::FileReader::try_new(std::io::Cursor::new(buf), None).unwrap();
2855 assert_eq!(reader.custom_metadata(), &test_metadata);
2856 }
2857
2858 #[test]
2859 fn test_roundtrip_nested_dict() {
2860 let inner: DictionaryArray<Int32Type> = vec!["a", "b", "a"].into_iter().collect();
2861
2862 let array = Arc::new(inner) as ArrayRef;
2863
2864 let dctfield = Arc::new(Field::new("dict", array.data_type().clone(), false));
2865
2866 let s = StructArray::from(vec![(dctfield, array)]);
2867 let struct_array = Arc::new(s) as ArrayRef;
2868
2869 let schema = Arc::new(Schema::new(vec![Field::new(
2870 "struct",
2871 struct_array.data_type().clone(),
2872 false,
2873 )]));
2874
2875 let batch = RecordBatch::try_new(schema, vec![struct_array]).unwrap();
2876
2877 assert_eq!(batch, roundtrip_ipc(&batch));
2878 }
2879
2880 #[test]
2881 fn test_roundtrip_nested_dict_no_preserve_dict_id() {
2882 let inner: DictionaryArray<Int32Type> = vec!["a", "b", "a"].into_iter().collect();
2883
2884 let array = Arc::new(inner) as ArrayRef;
2885
2886 let dctfield = Arc::new(Field::new("dict", array.data_type().clone(), false));
2887
2888 let s = StructArray::from(vec![(dctfield, array)]);
2889 let struct_array = Arc::new(s) as ArrayRef;
2890
2891 let schema = Arc::new(Schema::new(vec![Field::new(
2892 "struct",
2893 struct_array.data_type().clone(),
2894 false,
2895 )]));
2896
2897 let batch = RecordBatch::try_new(schema, vec![struct_array]).unwrap();
2898
2899 let mut buf = Vec::new();
2900 let mut writer = crate::writer::FileWriter::try_new_with_options(
2901 &mut buf,
2902 batch.schema_ref(),
2903 IpcWriteOptions::default(),
2904 )
2905 .unwrap();
2906 writer.write(&batch).unwrap();
2907 writer.finish().unwrap();
2908 drop(writer);
2909
2910 let mut reader = FileReader::try_new(std::io::Cursor::new(buf), None).unwrap();
2911
2912 assert_eq!(batch, reader.next().unwrap().unwrap());
2913 }
2914
2915 fn check_union_with_builder(mut builder: UnionBuilder) {
2916 builder.append::<Int32Type>("a", 1).unwrap();
2917 builder.append_null::<Int32Type>("a").unwrap();
2918 builder.append::<Float64Type>("c", 3.0).unwrap();
2919 builder.append::<Int32Type>("a", 4).unwrap();
2920 builder.append::<Int64Type>("d", 11).unwrap();
2921 let union = builder.build().unwrap();
2922
2923 let schema = Arc::new(Schema::new(vec![Field::new(
2924 "union",
2925 union.data_type().clone(),
2926 false,
2927 )]));
2928
2929 let union_array = Arc::new(union) as ArrayRef;
2930
2931 let rb = RecordBatch::try_new(schema, vec![union_array]).unwrap();
2932 let rb2 = roundtrip_ipc(&rb);
2933 assert_eq!(rb.schema(), rb2.schema());
2936 assert_eq!(rb.num_columns(), rb2.num_columns());
2937 assert_eq!(rb.num_rows(), rb2.num_rows());
2938 let union1 = rb.column(0);
2939 let union2 = rb2.column(0);
2940
2941 assert_eq!(union1, union2);
2942 }
2943
2944 #[test]
2945 fn test_roundtrip_dense_union() {
2946 check_union_with_builder(UnionBuilder::new_dense());
2947 }
2948
2949 #[test]
2950 fn test_roundtrip_sparse_union() {
2951 check_union_with_builder(UnionBuilder::new_sparse());
2952 }
2953
2954 #[test]
2955 fn test_roundtrip_struct_empty_fields() {
2956 let nulls = NullBuffer::from(&[true, true, false]);
2957 let rb = RecordBatch::try_from_iter([(
2958 "",
2959 Arc::new(StructArray::new_empty_fields(nulls.len(), Some(nulls))) as _,
2960 )])
2961 .unwrap();
2962 let rb2 = roundtrip_ipc(&rb);
2963 assert_eq!(rb, rb2);
2964 }
2965
2966 #[test]
2967 fn test_roundtrip_stream_run_array_sliced() {
2968 let run_array_1: Int32RunArray = vec!["a", "a", "a", "b", "b", "c", "c", "c"]
2969 .into_iter()
2970 .collect();
2971 let run_array_1_sliced = run_array_1.slice(2, 5);
2972
2973 let run_array_2_inupt = vec![Some(1_i32), None, None, Some(2), Some(2)];
2974 let mut run_array_2_builder = PrimitiveRunBuilder::<Int16Type, Int32Type>::new();
2975 run_array_2_builder.extend(run_array_2_inupt);
2976 let run_array_2 = run_array_2_builder.finish();
2977
2978 let schema = Arc::new(Schema::new(vec![
2979 Field::new(
2980 "run_array_1_sliced",
2981 run_array_1_sliced.data_type().clone(),
2982 false,
2983 ),
2984 Field::new("run_array_2", run_array_2.data_type().clone(), false),
2985 ]));
2986 let input_batch = RecordBatch::try_new(
2987 schema,
2988 vec![Arc::new(run_array_1_sliced.clone()), Arc::new(run_array_2)],
2989 )
2990 .unwrap();
2991 let output_batch = roundtrip_ipc_stream(&input_batch);
2992
2993 assert_eq!(input_batch.column(1), output_batch.column(1));
2997
2998 let run_array_1_unsliced = unslice_run_array(run_array_1_sliced.into_data()).unwrap();
2999 assert_eq!(run_array_1_unsliced, output_batch.column(0).into_data());
3000 }
3001
3002 #[test]
3003 fn test_roundtrip_stream_nested_dict() {
3004 let xs = vec!["AA", "BB", "AA", "CC", "BB"];
3005 let dict = Arc::new(
3006 xs.clone()
3007 .into_iter()
3008 .collect::<DictionaryArray<Int8Type>>(),
3009 );
3010 let string_array: ArrayRef = Arc::new(StringArray::from(xs.clone()));
3011 let struct_array = StructArray::from(vec![
3012 (
3013 Arc::new(Field::new("f2.1", DataType::Utf8, false)),
3014 string_array,
3015 ),
3016 (
3017 Arc::new(Field::new("f2.2_struct", dict.data_type().clone(), false)),
3018 dict.clone() as ArrayRef,
3019 ),
3020 ]);
3021 let schema = Arc::new(Schema::new(vec![
3022 Field::new("f1_string", DataType::Utf8, false),
3023 Field::new("f2_struct", struct_array.data_type().clone(), false),
3024 ]));
3025 let input_batch = RecordBatch::try_new(
3026 schema,
3027 vec![
3028 Arc::new(StringArray::from(xs.clone())),
3029 Arc::new(struct_array),
3030 ],
3031 )
3032 .unwrap();
3033 let output_batch = roundtrip_ipc_stream(&input_batch);
3034 assert_eq!(input_batch, output_batch);
3035 }
3036
3037 #[test]
3038 fn test_ipc_writers_reject_dictionary_of_dictionary_schema() {
3039 let values = Arc::new(StringArray::from(vec![Some("a"), Some("b")])) as ArrayRef;
3040 let inner = Arc::new(DictionaryArray::new(
3041 UInt32Array::from_iter_values([0, 1]),
3042 values,
3043 )) as ArrayRef;
3044 let outer = Arc::new(DictionaryArray::new(
3045 UInt32Array::from_iter_values([0, 1, 0]),
3046 inner,
3047 )) as ArrayRef;
3048
3049 let schema = Arc::new(Schema::new(vec![Field::new(
3050 "f1",
3051 outer.data_type().clone(),
3052 false,
3053 )]));
3054 let batch = RecordBatch::try_new(schema, vec![outer]).unwrap();
3055
3056 let mut stream = Vec::new();
3057 let Err(err) = crate::writer::StreamWriter::try_new(&mut stream, batch.schema_ref()) else {
3058 panic!("IPC stream writer should reject dictionary-of-dictionary schemas");
3059 };
3060 assert!(stream.is_empty());
3061
3062 assert!(
3063 err.to_string().contains("dictionary-of-dictionary values"),
3064 "unexpected error: {err}"
3065 );
3066
3067 let mut file = Vec::new();
3068 let Err(err) = crate::writer::FileWriter::try_new(&mut file, batch.schema_ref()) else {
3069 panic!("IPC file writer should reject dictionary-of-dictionary schemas");
3070 };
3071 assert!(file.is_empty());
3072
3073 assert!(
3074 err.to_string().contains("dictionary-of-dictionary values"),
3075 "unexpected error: {err}"
3076 );
3077 }
3078
3079 #[test]
3080 fn test_roundtrip_stream_nested_dict_of_map_of_dict() {
3081 let values = StringArray::from(vec![Some("a"), None, Some("b"), Some("c")]);
3082 let values = Arc::new(values) as ArrayRef;
3083 let value_dict_keys = Int8Array::from_iter_values([0, 1, 1, 2, 3, 1]);
3084 let value_dict_array = DictionaryArray::new(value_dict_keys, values.clone());
3085
3086 let key_dict_keys = Int8Array::from_iter_values([0, 0, 2, 2, 2, 3]);
3087 let key_dict_array = DictionaryArray::new(key_dict_keys, values);
3088
3089 #[expect(deprecated)]
3090 let keys_field = Arc::new(Field::new_dict(
3091 Field::MAP_KEY_FIELD_DEFAULT_NAME,
3092 DataType::Dictionary(Box::new(DataType::Int8), Box::new(DataType::Utf8)),
3093 false,
3094 1,
3095 false,
3096 ));
3097 #[expect(deprecated)]
3098 let values_field = Arc::new(Field::new_dict(
3099 Field::MAP_VALUE_FIELD_DEFAULT_NAME,
3100 DataType::Dictionary(Box::new(DataType::Int8), Box::new(DataType::Utf8)),
3101 true,
3102 2,
3103 false,
3104 ));
3105 let entry_struct = StructArray::from(vec![
3106 (keys_field, make_array(key_dict_array.into_data())),
3107 (values_field, make_array(value_dict_array.into_data())),
3108 ]);
3109 let map_data_type = DataType::Map(
3110 Arc::new(Field::new(
3111 Field::MAP_ENTRIES_FIELD_DEFAULT_NAME,
3112 entry_struct.data_type().clone(),
3113 false,
3114 )),
3115 false,
3116 );
3117
3118 let entry_offsets = Buffer::from_slice_ref([0, 2, 4, 6]);
3119 let map_data = ArrayData::builder(map_data_type)
3120 .len(3)
3121 .add_buffer(entry_offsets)
3122 .add_child_data(entry_struct.into_data())
3123 .build()
3124 .unwrap();
3125 let map_array = MapArray::from(map_data);
3126
3127 let dict_keys = Int8Array::from_iter_values([0, 1, 1, 2, 2, 1]);
3128 let dict_dict_array = DictionaryArray::new(dict_keys, Arc::new(map_array));
3129
3130 let schema = Arc::new(Schema::new(vec![Field::new(
3131 "f1",
3132 dict_dict_array.data_type().clone(),
3133 false,
3134 )]));
3135 let input_batch = RecordBatch::try_new(schema, vec![Arc::new(dict_dict_array)]).unwrap();
3136 let output_batch = roundtrip_ipc_stream(&input_batch);
3137 assert_eq!(input_batch, output_batch);
3138 }
3139
3140 fn test_roundtrip_stream_dict_of_list_of_dict_impl<
3141 OffsetSize: OffsetSizeTrait,
3142 U: ArrowNativeType,
3143 >(
3144 list_data_type: DataType,
3145 offsets: &[U; 5],
3146 ) {
3147 let values = StringArray::from(vec![Some("a"), None, Some("c"), None]);
3148 let keys = Int8Array::from_iter_values([0, 0, 1, 2, 0, 1, 3]);
3149 let dict_array = DictionaryArray::new(keys, Arc::new(values));
3150 let dict_data = dict_array.to_data();
3151
3152 let value_offsets = Buffer::from_slice_ref(offsets);
3153
3154 let list_data = ArrayData::builder(list_data_type)
3155 .len(4)
3156 .add_buffer(value_offsets)
3157 .add_child_data(dict_data)
3158 .build()
3159 .unwrap();
3160 let list_array = GenericListArray::<OffsetSize>::from(list_data);
3161
3162 let keys_for_dict = Int8Array::from_iter_values([0, 3, 0, 1, 1, 2, 0, 1, 3]);
3163 let dict_dict_array = DictionaryArray::new(keys_for_dict, Arc::new(list_array));
3164
3165 let schema = Arc::new(Schema::new(vec![Field::new(
3166 "f1",
3167 dict_dict_array.data_type().clone(),
3168 false,
3169 )]));
3170 let input_batch = RecordBatch::try_new(schema, vec![Arc::new(dict_dict_array)]).unwrap();
3171 let output_batch = roundtrip_ipc_stream(&input_batch);
3172 assert_eq!(input_batch, output_batch);
3173 }
3174
3175 #[test]
3176 fn test_roundtrip_stream_dict_of_list_of_dict() {
3177 #[expect(deprecated)]
3179 let list_data_type = DataType::List(Arc::new(Field::new_dict(
3180 "item",
3181 DataType::Dictionary(Box::new(DataType::Int8), Box::new(DataType::Utf8)),
3182 true,
3183 1,
3184 false,
3185 )));
3186 let offsets: &[i32; 5] = &[0, 2, 4, 4, 6];
3187 test_roundtrip_stream_dict_of_list_of_dict_impl::<i32, i32>(list_data_type, offsets);
3188
3189 #[expect(deprecated)]
3191 let list_data_type = DataType::LargeList(Arc::new(Field::new_dict(
3192 "item",
3193 DataType::Dictionary(Box::new(DataType::Int8), Box::new(DataType::Utf8)),
3194 true,
3195 1,
3196 false,
3197 )));
3198 let offsets: &[i64; 5] = &[0, 2, 4, 4, 7];
3199 test_roundtrip_stream_dict_of_list_of_dict_impl::<i64, i64>(list_data_type, offsets);
3200 }
3201
3202 #[test]
3203 fn test_roundtrip_stream_dict_of_fixed_size_list_of_dict() {
3204 let values = StringArray::from(vec![Some("a"), None, Some("c"), None]);
3205 let keys = Int8Array::from_iter_values([0, 0, 1, 2, 0, 1, 3, 1, 2]);
3206 let dict_array = DictionaryArray::new(keys, Arc::new(values));
3207 let dict_data = dict_array.into_data();
3208
3209 #[expect(deprecated)]
3210 let list_data_type = DataType::FixedSizeList(
3211 Arc::new(Field::new_dict(
3212 "item",
3213 DataType::Dictionary(Box::new(DataType::Int8), Box::new(DataType::Utf8)),
3214 true,
3215 1,
3216 false,
3217 )),
3218 3,
3219 );
3220 let list_data = ArrayData::builder(list_data_type)
3221 .len(3)
3222 .add_child_data(dict_data)
3223 .build()
3224 .unwrap();
3225 let list_array = FixedSizeListArray::from(list_data);
3226
3227 let keys_for_dict = Int8Array::from_iter_values([0, 1, 0, 1, 1, 2, 0, 1, 2]);
3228 let dict_dict_array = DictionaryArray::new(keys_for_dict, Arc::new(list_array));
3229
3230 let schema = Arc::new(Schema::new(vec![Field::new(
3231 "f1",
3232 dict_dict_array.data_type().clone(),
3233 false,
3234 )]));
3235 let input_batch = RecordBatch::try_new(schema, vec![Arc::new(dict_dict_array)]).unwrap();
3236 let output_batch = roundtrip_ipc_stream(&input_batch);
3237 assert_eq!(input_batch, output_batch);
3238 }
3239
3240 const LONG_TEST_STRING: &str =
3241 "This is a long string to make sure binary view array handles it";
3242
3243 #[test]
3244 fn test_roundtrip_view_types() {
3245 let schema = Schema::new(vec![
3246 Field::new("field_1", DataType::BinaryView, true),
3247 Field::new("field_2", DataType::Utf8, true),
3248 Field::new("field_3", DataType::Utf8View, true),
3249 ]);
3250 let bin_values: Vec<Option<&[u8]>> = vec![
3251 Some(b"foo"),
3252 None,
3253 Some(b"bar"),
3254 Some(LONG_TEST_STRING.as_bytes()),
3255 ];
3256 let utf8_values: Vec<Option<&str>> =
3257 vec![Some("foo"), None, Some("bar"), Some(LONG_TEST_STRING)];
3258 let bin_view_array = BinaryViewArray::from_iter(bin_values);
3259 let utf8_array = StringArray::from_iter(utf8_values.iter());
3260 let utf8_view_array = StringViewArray::from_iter(utf8_values);
3261 let record_batch = RecordBatch::try_new(
3262 Arc::new(schema.clone()),
3263 vec![
3264 Arc::new(bin_view_array),
3265 Arc::new(utf8_array),
3266 Arc::new(utf8_view_array),
3267 ],
3268 )
3269 .unwrap();
3270
3271 assert_eq!(record_batch, roundtrip_ipc(&record_batch));
3272 assert_eq!(record_batch, roundtrip_ipc_stream(&record_batch));
3273
3274 let sliced_batch = record_batch.slice(1, 2);
3275 assert_eq!(sliced_batch, roundtrip_ipc(&sliced_batch));
3276 assert_eq!(sliced_batch, roundtrip_ipc_stream(&sliced_batch));
3277 }
3278
3279 #[test]
3280 fn test_roundtrip_view_types_nested_dict() {
3281 let bin_values: Vec<Option<&[u8]>> = vec![
3282 Some(b"foo"),
3283 None,
3284 Some(b"bar"),
3285 Some(LONG_TEST_STRING.as_bytes()),
3286 Some(b"field"),
3287 ];
3288 let utf8_values: Vec<Option<&str>> = vec![
3289 Some("foo"),
3290 None,
3291 Some("bar"),
3292 Some(LONG_TEST_STRING),
3293 Some("field"),
3294 ];
3295 let bin_view_array = Arc::new(BinaryViewArray::from_iter(bin_values));
3296 let utf8_view_array = Arc::new(StringViewArray::from_iter(utf8_values));
3297
3298 let key_dict_keys = Int8Array::from_iter_values([0, 0, 2, 2, 0, 2, 3]);
3299 let key_dict_array = DictionaryArray::new(key_dict_keys, utf8_view_array.clone());
3300 #[expect(deprecated)]
3301 let keys_field = Arc::new(Field::new_dict(
3302 Field::MAP_KEY_FIELD_DEFAULT_NAME,
3303 DataType::Dictionary(Box::new(DataType::Int8), Box::new(DataType::Utf8View)),
3304 false,
3305 1,
3306 false,
3307 ));
3308
3309 let value_dict_keys = Int8Array::from_iter_values([0, 3, 0, 1, 2, 0, 1]);
3310 let value_dict_array = DictionaryArray::new(value_dict_keys, bin_view_array);
3311 #[expect(deprecated)]
3312 let values_field = Arc::new(Field::new_dict(
3313 Field::MAP_VALUE_FIELD_DEFAULT_NAME,
3314 DataType::Dictionary(Box::new(DataType::Int8), Box::new(DataType::BinaryView)),
3315 true,
3316 2,
3317 false,
3318 ));
3319 let entry_struct = StructArray::from(vec![
3320 (keys_field, make_array(key_dict_array.into_data())),
3321 (values_field, make_array(value_dict_array.into_data())),
3322 ]);
3323
3324 let map_data_type = DataType::Map(
3325 Arc::new(Field::new(
3326 Field::MAP_ENTRIES_FIELD_DEFAULT_NAME,
3327 entry_struct.data_type().clone(),
3328 false,
3329 )),
3330 false,
3331 );
3332 let entry_offsets = Buffer::from_slice_ref([0, 2, 4, 7]);
3333 let map_data = ArrayData::builder(map_data_type)
3334 .len(3)
3335 .add_buffer(entry_offsets)
3336 .add_child_data(entry_struct.into_data())
3337 .build()
3338 .unwrap();
3339 let map_array = MapArray::from(map_data);
3340
3341 let dict_keys = Int8Array::from_iter_values([0, 1, 0, 1, 1, 2, 0, 1, 2]);
3342 let dict_dict_array = DictionaryArray::new(dict_keys, Arc::new(map_array));
3343 let schema = Arc::new(Schema::new(vec![Field::new(
3344 "f1",
3345 dict_dict_array.data_type().clone(),
3346 false,
3347 )]));
3348 let batch = RecordBatch::try_new(schema, vec![Arc::new(dict_dict_array)]).unwrap();
3349 assert_eq!(batch, roundtrip_ipc(&batch));
3350 assert_eq!(batch, roundtrip_ipc_stream(&batch));
3351
3352 let sliced_batch = batch.slice(1, 2);
3353 assert_eq!(sliced_batch, roundtrip_ipc(&sliced_batch));
3354 assert_eq!(sliced_batch, roundtrip_ipc_stream(&sliced_batch));
3355 }
3356
3357 #[test]
3358 fn test_no_columns_batch() {
3359 let schema = Arc::new(Schema::empty());
3360 let options = RecordBatchOptions::new()
3361 .with_match_field_names(true)
3362 .with_row_count(Some(10));
3363 let input_batch = RecordBatch::try_new_with_options(schema, vec![], &options).unwrap();
3364 let output_batch = roundtrip_ipc_stream(&input_batch);
3365 assert_eq!(input_batch, output_batch);
3366 }
3367
3368 #[test]
3369 fn test_unaligned() {
3370 let batch = RecordBatch::try_from_iter(vec![(
3371 "i32",
3372 Arc::new(Int32Array::from(vec![1, 2, 3, 4])) as _,
3373 )])
3374 .unwrap();
3375
3376 let r#gen = IpcDataGenerator {};
3377 let mut dict_tracker = DictionaryTracker::new(false);
3378 let (_, encoded) = r#gen
3379 .encode(
3380 &batch,
3381 &mut dict_tracker,
3382 &Default::default(),
3383 &mut Default::default(),
3384 )
3385 .unwrap();
3386
3387 let message = root_as_message(&encoded.ipc_message).unwrap();
3388
3389 let mut buffer = MutableBuffer::with_capacity(encoded.arrow_data.len() + 1);
3391 buffer.push(0_u8);
3392 buffer.extend_from_slice(&encoded.arrow_data);
3393 let b = Buffer::from(buffer).slice(1);
3394 assert_ne!(b.as_ptr().align_offset(8), 0);
3395
3396 let ipc_batch = message.header_as_record_batch().unwrap();
3397 let roundtrip = RecordBatchDecoder::try_new(
3398 &b,
3399 ipc_batch,
3400 batch.schema(),
3401 &Default::default(),
3402 &message.version(),
3403 )
3404 .unwrap()
3405 .with_require_alignment(false)
3406 .read_record_batch()
3407 .unwrap();
3408 assert_eq!(batch, roundtrip);
3409 }
3410
3411 #[test]
3412 fn test_unaligned_throws_error_with_require_alignment() {
3413 let batch = RecordBatch::try_from_iter(vec![(
3414 "i32",
3415 Arc::new(Int32Array::from(vec![1, 2, 3, 4])) as _,
3416 )])
3417 .unwrap();
3418
3419 let r#gen = IpcDataGenerator {};
3420 let mut dict_tracker = DictionaryTracker::new(false);
3421 let (_, encoded) = r#gen
3422 .encode(
3423 &batch,
3424 &mut dict_tracker,
3425 &Default::default(),
3426 &mut Default::default(),
3427 )
3428 .unwrap();
3429
3430 let message = root_as_message(&encoded.ipc_message).unwrap();
3431
3432 let mut buffer = MutableBuffer::with_capacity(encoded.arrow_data.len() + 1);
3434 buffer.push(0_u8);
3435 buffer.extend_from_slice(&encoded.arrow_data);
3436 let b = Buffer::from(buffer).slice(1);
3437 assert_ne!(b.as_ptr().align_offset(8), 0);
3438
3439 let ipc_batch = message.header_as_record_batch().unwrap();
3440 let result = RecordBatchDecoder::try_new(
3441 &b,
3442 ipc_batch,
3443 batch.schema(),
3444 &Default::default(),
3445 &message.version(),
3446 )
3447 .unwrap()
3448 .with_require_alignment(true)
3449 .read_record_batch();
3450
3451 let error = result.unwrap_err();
3452 assert_eq!(
3453 error.to_string(),
3454 "Invalid argument error: Misaligned buffers[0] in array of type Int32, \
3455 offset from expected alignment of 4 by 1"
3456 );
3457 }
3458
3459 #[test]
3468 fn test_misaligned_buffers_error() {
3469 const OFFSET_RANGE: usize = 128;
3470
3471 let cases: Vec<(ArrayRef, usize)> = vec![
3475 (Arc::new(Int32Array::from_iter(0i32..100)) as _, 4),
3476 (Arc::new(Int64Array::from_iter(0i64..100)) as _, 8),
3477 (
3478 Arc::new(StringArray::from_iter_values(
3479 (0..100).map(|i| i.to_string()),
3480 )) as _,
3481 4,
3482 ),
3483 (
3484 Arc::new(LargeStringArray::from_iter_values(
3485 (0..100).map(|i| i.to_string()),
3486 )) as _,
3487 8,
3488 ),
3489 ];
3490
3491 for (array, alignment) in cases {
3492 let batch = RecordBatch::try_from_iter(vec![("col", Arc::clone(&array))]).unwrap();
3493 let encoder = IpcDataGenerator {};
3494 let mut dict_tracker = DictionaryTracker::new(false);
3495 let (_, encoded) = encoder
3496 .encode(
3497 &batch,
3498 &mut dict_tracker,
3499 &Default::default(),
3500 &mut Default::default(),
3501 )
3502 .unwrap();
3503 let message = root_as_message(&encoded.ipc_message).unwrap();
3504 let ipc_batch = message.header_as_record_batch().unwrap();
3505
3506 for offset in 0..OFFSET_RANGE {
3507 let mut storage = MutableBuffer::with_capacity(encoded.arrow_data.len() + offset);
3513 for _ in 0..offset {
3514 storage.push(0_u8);
3515 }
3516 storage.extend_from_slice(&encoded.arrow_data);
3517 let buf = Buffer::from(storage).slice(offset);
3518
3519 let result = RecordBatchDecoder::try_new(
3520 &buf,
3521 ipc_batch,
3522 batch.schema(),
3523 &Default::default(),
3524 &message.version(),
3525 )
3526 .unwrap()
3527 .with_require_alignment(true)
3528 .read_record_batch();
3529
3530 if offset % alignment == 0 {
3531 assert!(
3532 result.is_ok(),
3533 "type={} offset={offset}: expected Ok but got {:?}",
3534 array.data_type(),
3535 result.unwrap_err(),
3536 );
3537 } else {
3538 let err = result
3539 .expect_err(&format!(
3540 "type={} offset={offset}: expected Err for misaligned buffer",
3541 array.data_type()
3542 ))
3543 .to_string();
3544 assert!(
3545 err.contains("Misaligned buffers"),
3546 "type={} offset={offset}: unexpected error: {err}",
3547 array.data_type(),
3548 );
3549 }
3550 }
3551 }
3552 }
3553
3554 #[test]
3555 #[cfg_attr(miri, ignore)] fn test_file_with_massive_column_count() {
3557 let limit = 600_000;
3559
3560 let fields = (0..limit)
3561 .map(|i| Field::new(format!("{i}"), DataType::Boolean, false))
3562 .collect::<Vec<_>>();
3563 let schema = Arc::new(Schema::new(fields));
3564 let batch = RecordBatch::new_empty(schema);
3565
3566 let mut buf = Vec::new();
3567 let mut writer = crate::writer::FileWriter::try_new(&mut buf, batch.schema_ref()).unwrap();
3568 writer.write(&batch).unwrap();
3569 writer.finish().unwrap();
3570 drop(writer);
3571
3572 let mut reader = FileReaderBuilder::new()
3573 .with_max_footer_fb_tables(1_500_000)
3574 .build(std::io::Cursor::new(buf))
3575 .unwrap();
3576 let roundtrip_batch = reader.next().unwrap().unwrap();
3577
3578 assert_eq!(batch, roundtrip_batch);
3579 }
3580
3581 #[test]
3582 fn test_file_with_deeply_nested_columns() {
3583 let limit = 61;
3585
3586 let fields = (0..limit).fold(
3587 vec![Field::new("leaf", DataType::Boolean, false)],
3588 |field, index| vec![Field::new_struct(format!("{index}"), field, false)],
3589 );
3590 let schema = Arc::new(Schema::new(fields));
3591 let batch = RecordBatch::new_empty(schema);
3592
3593 let mut buf = Vec::new();
3594 let mut writer = crate::writer::FileWriter::try_new(&mut buf, batch.schema_ref()).unwrap();
3595 writer.write(&batch).unwrap();
3596 writer.finish().unwrap();
3597 drop(writer);
3598
3599 let mut reader = FileReaderBuilder::new()
3600 .with_max_footer_fb_depth(65)
3601 .build(std::io::Cursor::new(buf))
3602 .unwrap();
3603 let roundtrip_batch = reader.next().unwrap().unwrap();
3604
3605 assert_eq!(batch, roundtrip_batch);
3606 }
3607
3608 #[test]
3609 fn test_stream_reader_rejects_short_validity_buffer() {
3610 let data = ArrayDataBuilder::new(DataType::Int32)
3612 .len(8000)
3613 .add_buffer(ScalarBuffer::<i32>::from_iter(0..8000).into())
3614 .nulls(Some(NullBuffer::from(&[true, false, true, false])));
3615 let array: ArrayRef = unsafe { Arc::new(Int32Array::from(data.build_unchecked())) };
3616 let batch = RecordBatch::try_from_iter([("a", array)]).unwrap();
3617
3618 let mut stream = Vec::new();
3619 let mut writer =
3620 crate::writer::StreamWriter::try_new(&mut stream, &batch.schema()).unwrap();
3621 writer.write(&batch).unwrap();
3622 writer.finish().unwrap();
3623
3624 let mut reader = StreamReader::try_new(Cursor::new(stream), None).unwrap();
3625 let err = reader.next().unwrap().unwrap_err();
3626 assert_eq!(
3627 err.to_string(),
3628 "Invalid argument error: null_bit_buffer size too small. got 1 needed 1000"
3629 );
3630 }
3631
3632 #[test]
3633 fn test_stream_reader_rejects_short_struct_validity_buffer() {
3634 let fields = Fields::from(vec![Field::new("b", DataType::Int32, false)]);
3636 let child: ArrayRef = Arc::new(Int32Array::from_iter_values(0..8000));
3637 let nulls = NullBuffer::from(&[true, false, true, false]);
3638 let array: ArrayRef =
3639 unsafe { Arc::new(StructArray::new_unchecked(fields, vec![child], Some(nulls))) };
3640 let batch = RecordBatch::try_from_iter([("a", array)]).unwrap();
3641
3642 let mut stream = Vec::new();
3643 let mut writer =
3644 crate::writer::StreamWriter::try_new(&mut stream, &batch.schema()).unwrap();
3645 writer.write(&batch).unwrap();
3646 writer.finish().unwrap();
3647
3648 let mut reader = StreamReader::try_new(Cursor::new(stream), None).unwrap();
3649 let err = reader.next().unwrap().unwrap_err();
3650 assert_eq!(
3651 err.to_string(),
3652 "Invalid argument error: null_bit_buffer size too small. got 1 needed 1000"
3653 );
3654 }
3655
3656 #[test]
3657 fn test_invalid_struct_array_ipc_read_errors() {
3658 let a_field = Field::new("a", DataType::Int32, false);
3659 let b_field = Field::new("b", DataType::Int32, false);
3660 let struct_fields = Fields::from(vec![a_field.clone(), b_field.clone()]);
3661
3662 let a_array_data = ArrayData::builder(a_field.data_type().clone())
3663 .len(4)
3664 .add_buffer(Buffer::from_slice_ref([1, 2, 3, 4]))
3665 .build()
3666 .unwrap();
3667 let b_array_data = ArrayData::builder(b_field.data_type().clone())
3668 .len(3)
3669 .add_buffer(Buffer::from_slice_ref([5, 6, 7]))
3670 .build()
3671 .unwrap();
3672
3673 let invalid_struct_arr = unsafe {
3674 StructArray::new_unchecked(
3675 struct_fields,
3676 vec![make_array(a_array_data), make_array(b_array_data)],
3677 None,
3678 )
3679 };
3680
3681 expect_ipc_validation_error(
3682 Arc::new(invalid_struct_arr),
3683 "Invalid argument error: Incorrect array length for StructArray field \"b\", expected 4 got 3",
3684 );
3685 }
3686
3687 #[test]
3688 fn test_invalid_nested_array_ipc_read_errors() {
3689 let a_field = Field::new("a", DataType::Int32, false);
3691 let b_field = Field::new("b", DataType::Utf8, false);
3692
3693 let schema = Arc::new(Schema::new(vec![Field::new_struct(
3694 "s",
3695 vec![a_field.clone(), b_field.clone()],
3696 false,
3697 )]));
3698
3699 let a_array_data = ArrayData::builder(a_field.data_type().clone())
3700 .len(4)
3701 .add_buffer(Buffer::from_slice_ref([1, 2, 3, 4]))
3702 .build()
3703 .unwrap();
3704 let b_array_data = {
3706 let valid: &[u8] = b" ";
3707 let mut invalid = vec![];
3708 invalid.extend_from_slice(b"ValidString");
3709 invalid.extend_from_slice(INVALID_UTF8_FIRST_CHAR);
3710 let binary_array =
3711 BinaryArray::from_iter(vec![None, Some(valid), None, Some(&invalid)]);
3712 let array = unsafe {
3713 StringArray::new_unchecked(
3714 binary_array.offsets().clone(),
3715 binary_array.values().clone(),
3716 binary_array.nulls().cloned(),
3717 )
3718 };
3719 array.into_data()
3720 };
3721 let struct_data_type = schema.field(0).data_type();
3722
3723 let invalid_struct_arr = unsafe {
3724 make_array(
3725 ArrayData::builder(struct_data_type.clone())
3726 .len(4)
3727 .add_child_data(a_array_data)
3728 .add_child_data(b_array_data)
3729 .build_unchecked(),
3730 )
3731 };
3732 expect_ipc_validation_error(
3733 invalid_struct_arr,
3734 "Invalid argument error: Invalid UTF8 sequence at string index 3 (3..18): invalid utf-8 sequence of 1 bytes from index 11",
3735 );
3736 }
3737
3738 #[test]
3739 fn test_same_dict_id_without_preserve() {
3740 let batch = RecordBatch::try_new(
3741 Arc::new(Schema::new(
3742 ["a", "b"]
3743 .iter()
3744 .map(|name| {
3745 #[expect(deprecated)]
3746 Field::new_dict(
3747 name.to_string(),
3748 DataType::Dictionary(
3749 Box::new(DataType::Int32),
3750 Box::new(DataType::Utf8),
3751 ),
3752 true,
3753 0,
3754 false,
3755 )
3756 })
3757 .collect::<Vec<Field>>(),
3758 )),
3759 vec![
3760 Arc::new(
3761 vec![Some("c"), Some("d")]
3762 .into_iter()
3763 .collect::<DictionaryArray<Int32Type>>(),
3764 ) as ArrayRef,
3765 Arc::new(
3766 vec![Some("e"), Some("f")]
3767 .into_iter()
3768 .collect::<DictionaryArray<Int32Type>>(),
3769 ) as ArrayRef,
3770 ],
3771 )
3772 .expect("Failed to create RecordBatch");
3773
3774 let mut buf = vec![];
3776 {
3777 let mut writer = crate::writer::StreamWriter::try_new_with_options(
3778 &mut buf,
3779 batch.schema().as_ref(),
3780 crate::writer::IpcWriteOptions::default(),
3781 )
3782 .expect("Failed to create StreamWriter");
3783 writer.write(&batch).expect("Failed to write RecordBatch");
3784 writer.finish().expect("Failed to finish StreamWriter");
3785 }
3786
3787 StreamReader::try_new(std::io::Cursor::new(buf), None)
3788 .expect("Failed to create StreamReader")
3789 .for_each(|decoded_batch| {
3790 assert_eq!(decoded_batch.expect("Failed to read RecordBatch"), batch);
3791 });
3792 }
3793
3794 #[test]
3795 fn test_validation_of_invalid_list_array() {
3796 let array = unsafe {
3798 let values = Int32Array::from(vec![1, 2, 3]);
3799 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);
3802 let nulls = None;
3803 ListArray::new(Arc::new(field), offsets, Arc::new(values), nulls)
3804 };
3805
3806 expect_ipc_validation_error(
3807 Arc::new(array),
3808 "Invalid argument error: Offset invariant failure: offset at position 2 out of bounds: 4 > 2",
3809 );
3810 }
3811
3812 #[test]
3813 fn test_validation_of_invalid_string_array() {
3814 let valid: &[u8] = b" ";
3815 let mut invalid = vec![];
3816 invalid.extend_from_slice(b"ThisStringIsCertainlyLongerThan12Bytes");
3817 invalid.extend_from_slice(INVALID_UTF8_FIRST_CHAR);
3818 let binary_array = BinaryArray::from_iter(vec![None, Some(valid), None, Some(&invalid)]);
3819 let array = unsafe {
3822 StringArray::new_unchecked(
3823 binary_array.offsets().clone(),
3824 binary_array.values().clone(),
3825 binary_array.nulls().cloned(),
3826 )
3827 };
3828 expect_ipc_validation_error(
3829 Arc::new(array),
3830 "Invalid argument error: Invalid UTF8 sequence at string index 3 (3..45): invalid utf-8 sequence of 1 bytes from index 38",
3831 );
3832 }
3833
3834 #[test]
3835 fn test_validation_of_invalid_string_view_array() {
3836 let valid: &[u8] = b" ";
3837 let mut invalid = vec![];
3838 invalid.extend_from_slice(b"ThisStringIsCertainlyLongerThan12Bytes");
3839 invalid.extend_from_slice(INVALID_UTF8_FIRST_CHAR);
3840 let binary_view_array =
3841 BinaryViewArray::from_iter(vec![None, Some(valid), None, Some(&invalid)]);
3842 let array = unsafe {
3845 StringViewArray::new_unchecked(
3846 binary_view_array.views().clone(),
3847 Arc::clone(binary_view_array.data_buffers()),
3848 binary_view_array.nulls().cloned(),
3849 )
3850 };
3851 expect_ipc_validation_error(
3852 Arc::new(array),
3853 "Invalid argument error: Encountered non-UTF-8 data at index 3: invalid utf-8 sequence of 1 bytes from index 38",
3854 );
3855 }
3856
3857 #[test]
3860 fn test_validation_of_invalid_dictionary_array() {
3861 let array = unsafe {
3862 let values = StringArray::from_iter_values(["a", "b", "c"]);
3863 let keys = Int32Array::from(vec![1, 200]); DictionaryArray::new_unchecked(keys, Arc::new(values))
3865 };
3866
3867 expect_ipc_validation_error(
3868 Arc::new(array),
3869 "Invalid argument error: Value at position 1 out of bounds: 200 (should be in [0, 2])",
3870 );
3871 }
3872
3873 #[test]
3874 fn test_validation_of_invalid_union_array() {
3875 let array = unsafe {
3876 let fields = UnionFields::try_new(
3877 vec![1, 3], vec![
3879 Field::new("a", DataType::Int32, false),
3880 Field::new("b", DataType::Utf8, false),
3881 ],
3882 )
3883 .unwrap();
3884 let type_ids = ScalarBuffer::from(vec![1i8, 2, 3]); let offsets = None;
3886 let children: Vec<ArrayRef> = vec![
3887 Arc::new(Int32Array::from(vec![10, 20, 30])),
3888 Arc::new(StringArray::from(vec![Some("a"), Some("b"), Some("c")])),
3889 ];
3890
3891 UnionArray::new_unchecked(fields, type_ids, offsets, children)
3892 };
3893
3894 expect_ipc_validation_error(
3895 Arc::new(array),
3896 "Invalid argument error: Type Ids values must match one of the field type ids",
3897 );
3898 }
3899
3900 const INVALID_UTF8_FIRST_CHAR: &[u8] = &[0xa0, 0xa1, 0x20, 0x20];
3903
3904 fn expect_ipc_validation_error(array: ArrayRef, expected_err: &str) {
3906 let rb = RecordBatch::try_from_iter([("a", array)]).unwrap();
3907
3908 let buf = write_stream(&rb); read_stream_skip_validation(&buf).unwrap();
3911 let err = read_stream(&buf).unwrap_err();
3912 assert_eq!(err.to_string(), expected_err);
3913
3914 let buf = write_ipc(&rb); read_ipc_skip_validation(&buf).unwrap();
3917 let err = read_ipc(&buf).unwrap_err();
3918 assert_eq!(err.to_string(), expected_err);
3919
3920 read_ipc_with_decoder_skip_validation(buf.clone()).unwrap();
3922 let err = read_ipc_with_decoder(buf).unwrap_err();
3923 assert_eq!(err.to_string(), expected_err);
3924 }
3925
3926 #[test]
3927 fn test_roundtrip_schema() {
3928 let schema = Schema::new(vec![
3929 Field::new(
3930 "a",
3931 DataType::Dictionary(Box::new(DataType::UInt16), Box::new(DataType::Utf8)),
3932 false,
3933 ),
3934 Field::new(
3935 "b",
3936 DataType::Dictionary(Box::new(DataType::UInt16), Box::new(DataType::Utf8)),
3937 false,
3938 ),
3939 ]);
3940
3941 let options = IpcWriteOptions::default();
3942 let data_gen = IpcDataGenerator::default();
3943 let mut dict_tracker = DictionaryTracker::new(false);
3944 let encoded_data =
3945 data_gen.schema_to_bytes_with_dictionary_tracker(&schema, &mut dict_tracker, &options);
3946 let mut schema_bytes = vec![];
3947 write_message(&mut schema_bytes, encoded_data, &options).expect("write_message");
3948
3949 let begin_offset: usize = if schema_bytes[0..4].eq(&CONTINUATION_MARKER) {
3950 4
3951 } else {
3952 0
3953 };
3954
3955 size_prefixed_root_as_message(&schema_bytes[begin_offset..])
3956 .expect_err("size_prefixed_root_as_message");
3957
3958 let msg = parse_message(&schema_bytes).expect("parse_message");
3959 let ipc_schema = msg.header_as_schema().expect("header_as_schema");
3960 let new_schema = try_fb_to_schema(ipc_schema).unwrap();
3961
3962 assert_eq!(schema, new_schema);
3963 }
3964
3965 #[test]
3966 fn test_negative_meta_len() {
3967 let bytes = i32::to_le_bytes(-1);
3968 let mut buf = vec![];
3969 buf.extend(CONTINUATION_MARKER);
3970 buf.extend(bytes);
3971
3972 let reader = StreamReader::try_new(Cursor::new(buf), None);
3973 assert!(reader.is_err());
3974 }
3975
3976 #[test]
3982 fn test_read_null_dict_without_dictionary_batch() {
3983 let keys = Int32Array::new_null(4);
3985 let values: ArrayRef = new_empty_array(&DataType::Utf8);
3986 let dict_array = DictionaryArray::new(keys, values);
3987
3988 let schema = Arc::new(Schema::new(vec![Field::new(
3989 "d",
3990 dict_array.data_type().clone(),
3991 true,
3992 )]));
3993 let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(dict_array)]).unwrap();
3994
3995 let full_stream = write_stream(&batch);
3997
3998 let mut stripped = Vec::new();
4002 let mut cursor = Cursor::new(&full_stream);
4003 loop {
4004 let mut header = [0u8; 4];
4007 if cursor.read_exact(&mut header).is_err() {
4008 break;
4009 }
4010 if header == CONTINUATION_MARKER && cursor.read_exact(&mut header).is_err() {
4011 break;
4012 }
4013 let meta_len = u32::from_le_bytes(header) as usize;
4014 if meta_len == 0 {
4015 stripped.extend_from_slice(&CONTINUATION_MARKER);
4017 stripped.extend_from_slice(&0u32.to_le_bytes());
4018 break;
4019 }
4020 let mut meta_buf = vec![0u8; meta_len];
4021 cursor.read_exact(&mut meta_buf).unwrap();
4022
4023 let message = root_as_message(&meta_buf).unwrap();
4024 let body_len = message.bodyLength() as usize;
4025 let mut body_buf = vec![0u8; body_len];
4026 cursor.read_exact(&mut body_buf).unwrap();
4027
4028 if message.header_type() == crate::MessageHeader::DictionaryBatch {
4029 continue;
4032 }
4033 stripped.extend_from_slice(&CONTINUATION_MARKER);
4034 stripped.extend_from_slice(&(meta_len as u32).to_le_bytes());
4035 stripped.extend_from_slice(&meta_buf);
4036 stripped.extend_from_slice(&body_buf);
4037 }
4038
4039 let result = read_stream(&stripped).unwrap();
4041 assert_eq!(result.num_rows(), 4);
4042 assert_eq!(result.num_columns(), 1);
4043
4044 let col = result.column(0);
4045 assert_eq!(col.null_count(), 4);
4046 assert_eq!(col.len(), 4);
4047 assert!(matches!(col.data_type(), DataType::Dictionary(_, _)));
4049 }
4050
4051 #[test]
4054 fn test_projection_skip_list_view() {
4055 use crate::reader::FileReader;
4056 use crate::writer::FileWriter;
4057 use arrow_array::{
4058 GenericListViewArray, Int32Array, RecordBatch,
4059 builder::{GenericListViewBuilder, UInt32Builder},
4060 };
4061 use arrow_schema::{DataType, Field, Schema};
4062 use std::sync::Arc;
4063
4064 let mut builder = GenericListViewBuilder::<i32, _>::new(UInt32Builder::new());
4066
4067 builder.values().append_value(1);
4068 builder.values().append_value(2);
4069 builder.append(true);
4070
4071 builder.append(false);
4072
4073 builder.values().append_value(3);
4074 builder.values().append_value(4);
4075 builder.append(true);
4076
4077 let list_view: GenericListViewArray<i32> = builder.finish();
4078
4079 let values = Int32Array::from(vec![10, 20, 30]);
4081
4082 let schema = Arc::new(Schema::new(vec![
4084 Field::new("a", list_view.data_type().clone(), true),
4085 Field::new("b", DataType::Int32, false),
4086 ]));
4087 let batch =
4089 RecordBatch::try_new(schema, vec![Arc::new(list_view), Arc::new(values.clone())])
4090 .unwrap();
4091
4092 let mut buf = Vec::new();
4094 {
4095 let mut writer = FileWriter::try_new(&mut buf, &batch.schema()).unwrap();
4096 writer.write(&batch).unwrap();
4097 writer.finish().unwrap();
4098 }
4099
4100 let mut reader = FileReader::try_new(std::io::Cursor::new(buf), Some(vec![1])).unwrap();
4102 let read_batch = reader.next().unwrap().unwrap();
4103
4104 assert_eq!(read_batch.num_columns(), 1);
4106 assert_eq!(read_batch.column(0).as_ref(), &values);
4107 }
4108
4109 #[test]
4112 fn test_projection_skip_union_v4() {
4113 use crate::MetadataVersion;
4114 use crate::reader::FileReader;
4115 use crate::writer::{FileWriter, IpcWriteOptions};
4116 use arrow_array::{
4117 ArrayRef, Int32Array, RecordBatch, builder::UnionBuilder, types::Int32Type,
4118 };
4119 use arrow_schema::{DataType, Field, Schema};
4120 use std::sync::Arc;
4121
4122 let mut builder = UnionBuilder::new_dense();
4124 builder.append::<Int32Type>("a", 1).unwrap();
4125 builder.append::<Int32Type>("a", 2).unwrap();
4126 builder.append::<Int32Type>("a", 3).unwrap();
4127 let union = builder.build().unwrap();
4128
4129 let values = Int32Array::from(vec![10, 20, 30]);
4131
4132 let schema = Arc::new(Schema::new(vec![
4134 Field::new("union", union.data_type().clone(), false),
4135 Field::new("values", DataType::Int32, false),
4136 ]));
4137
4138 let batch = RecordBatch::try_new(
4140 schema,
4141 vec![Arc::new(union) as ArrayRef, Arc::new(values.clone())],
4142 )
4143 .unwrap();
4144
4145 let mut buf = Vec::new();
4147 {
4148 let options = IpcWriteOptions::try_new(8, false, MetadataVersion::V4).unwrap();
4149 let mut writer =
4150 FileWriter::try_new_with_options(&mut buf, &batch.schema(), options).unwrap();
4151 writer.write(&batch).unwrap();
4152 writer.finish().unwrap();
4153 }
4154 let mut reader = FileReader::try_new(std::io::Cursor::new(buf), Some(vec![1])).unwrap();
4156 let read_batch = reader.next().unwrap().unwrap();
4157
4158 assert_eq!(read_batch.num_columns(), 1);
4160 assert_eq!(read_batch.column(0).as_ref(), &values);
4161 }
4162
4163 #[test]
4167 fn test_projection_skip_fixed_width_types() {
4168 use std::sync::Arc;
4169
4170 use arrow_array::{ArrayRef, BooleanArray, Int32Array, RecordBatch, make_array};
4171 use arrow_buffer::{Buffer, MutableBuffer};
4172 use arrow_data::ArrayData;
4173 use arrow_schema::{DataType, Field, IntervalUnit, Schema, TimeUnit};
4174
4175 use crate::reader::FileReader;
4176 use crate::writer::FileWriter;
4177
4178 fn make_array_for_type(data_type: DataType) -> ArrayRef {
4180 let len = 3;
4181
4182 if matches!(data_type, DataType::Boolean) {
4183 return Arc::new(BooleanArray::from(vec![true, false, true]));
4184 }
4185
4186 let width = data_type.primitive_width().unwrap();
4187 let data = ArrayData::builder(data_type)
4188 .len(len)
4189 .add_buffer(Buffer::from(MutableBuffer::from_len_zeroed(len * width)))
4190 .build()
4191 .unwrap();
4192
4193 make_array(data)
4194 }
4195
4196 let data_types = vec![
4198 DataType::Boolean,
4199 DataType::Int8,
4200 DataType::Int16,
4201 DataType::Int32,
4202 DataType::Int64,
4203 DataType::UInt8,
4204 DataType::UInt16,
4205 DataType::UInt32,
4206 DataType::UInt64,
4207 DataType::Float16,
4208 DataType::Float32,
4209 DataType::Float64,
4210 DataType::Timestamp(TimeUnit::Second, None),
4211 DataType::Date32,
4212 DataType::Date64,
4213 DataType::Time32(TimeUnit::Second),
4214 DataType::Time64(TimeUnit::Microsecond),
4215 DataType::Duration(TimeUnit::Second),
4216 DataType::Interval(IntervalUnit::YearMonth),
4217 DataType::Interval(IntervalUnit::DayTime),
4218 DataType::Interval(IntervalUnit::MonthDayNano),
4219 DataType::Decimal32(9, 2),
4220 DataType::Decimal64(18, 2),
4221 DataType::Decimal128(38, 2),
4222 DataType::Decimal256(76, 2),
4223 ];
4224
4225 for data_type in data_types {
4230 let skipped = make_array_for_type(data_type.clone());
4231 let values = Int32Array::from(vec![10, 20, 30]);
4232
4233 let schema = Arc::new(Schema::new(vec![
4234 Field::new("skipped", data_type, false),
4235 Field::new("values", DataType::Int32, false),
4236 ]));
4237
4238 let batch =
4239 RecordBatch::try_new(schema, vec![skipped, Arc::new(values.clone())]).unwrap();
4240
4241 let mut buf = Vec::new();
4243 {
4244 let mut writer = FileWriter::try_new(&mut buf, &batch.schema()).unwrap();
4245 writer.write(&batch).unwrap();
4246 writer.finish().unwrap();
4247 }
4248
4249 let mut reader = FileReader::try_new(std::io::Cursor::new(buf), Some(vec![1])).unwrap();
4251 let read_batch = reader.next().unwrap().unwrap();
4252
4253 assert_eq!(read_batch.num_columns(), 1);
4255 assert_eq!(read_batch.column(0).as_ref(), &values);
4256 }
4257 }
4258
4259 fn stream_with_declared_body(body_length: i64, body_bytes: usize) -> Vec<u8> {
4263 let mut fbb = flatbuffers::FlatBufferBuilder::new();
4264 let mut message = crate::MessageBuilder::new(&mut fbb);
4265 message.add_version(crate::MetadataVersion::V5);
4266 message.add_header_type(crate::MessageHeader::NONE);
4267 message.add_bodyLength(body_length);
4268 let root = message.finish();
4269 fbb.finish(root, None);
4270 let metadata = fbb.finished_data();
4271
4272 let mut stream = Vec::new();
4273 stream.extend_from_slice(&CONTINUATION_MARKER);
4274 stream.extend_from_slice(&(metadata.len() as i32).to_le_bytes());
4275 stream.extend_from_slice(metadata);
4276 stream.resize(stream.len() + body_bytes, 0xAB);
4277 stream
4278 }
4279
4280 #[test]
4285 fn test_stream_reader_rejects_implausible_body_length() {
4286 for body_length in [i64::MAX, 1 << 50, -1] {
4287 let stream = stream_with_declared_body(body_length, 0);
4288 let err = StreamReader::try_new(std::io::Cursor::new(stream), None).expect_err(
4289 &format!("a message declaring {body_length} body bytes must not be accepted"),
4290 );
4291 let err = err.to_string();
4292 assert!(
4293 err.contains("Invalid IPC message body length")
4294 || err.contains("Unexpected end of stream")
4295 || err.contains("failed to fill whole buffer"),
4296 "unexpected error for body_length {body_length}: {err}"
4297 );
4298 }
4299 }
4300
4301 #[test]
4305 #[cfg_attr(miri, ignore)] fn test_stream_reader_rejects_truncated_body() {
4307 let over_prealloc = MAX_PREALLOC_BYTES as i64 + 1;
4308 for (body_length, body_bytes) in [(1024, 10), (over_prealloc, MAX_PREALLOC_BYTES)] {
4309 let stream = stream_with_declared_body(body_length, body_bytes);
4310 let err = StreamReader::try_new(std::io::Cursor::new(stream), None)
4311 .expect_err("a body backed by fewer bytes than declared must not be accepted");
4312 assert!(
4313 err.to_string().contains("failed to fill whole buffer"),
4314 "unexpected error for body_length {body_length}: {err}"
4315 );
4316 }
4317 }
4318
4319 #[test]
4321 fn test_stream_reader_rejects_unbacked_metadata_length() {
4322 let mut stream = Vec::new();
4323 stream.extend_from_slice(&CONTINUATION_MARKER);
4324 stream.extend_from_slice(&i32::MAX.to_le_bytes());
4325 stream.extend_from_slice(b"not this many bytes");
4326
4327 let err = StreamReader::try_new(std::io::Cursor::new(stream), None)
4328 .expect_err("metadata length must be backed by the stream");
4329 assert!(
4330 err.to_string().contains("Unexpected end of stream"),
4331 "unexpected error: {err}"
4332 );
4333 }
4334}