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 let validity_buffer = BooleanBuffer::new(null_buffer, 0, len);
373 let null_buffer = if skip_validation {
374 unsafe { NullBuffer::new_unchecked(validity_buffer, null_count) }
376 } else {
377 let null_buffer = NullBuffer::new(validity_buffer);
378
379 if null_buffer.null_count() != null_count {
380 return Err(ArrowError::InvalidArgumentError(format!(
381 "null_count value ({}) doesn't match actual number of nulls in array ({})",
382 null_count,
383 null_buffer.null_count()
384 )));
385 }
386
387 null_buffer
388 };
389
390 Some(null_buffer)
391 } else {
392 None
393 };
394 if struct_arrays.is_empty() {
395 return Ok(Arc::new(StructArray::new_empty_fields(len, nulls)));
398 }
399
400 let struct_array = if skip_validation {
401 unsafe { StructArray::new_unchecked(struct_fields.clone(), struct_arrays, nulls) }
403 } else {
404 StructArray::try_new(struct_fields.clone(), struct_arrays, nulls)?
405 };
406
407 Ok(Arc::new(struct_array))
408 }
409
410 fn create_dictionary_array(
413 &self,
414 field_node: &FieldNode,
415 data_type: &DataType,
416 buffers: &[Buffer],
417 value_array: ArrayRef,
418 ) -> Result<ArrayRef, ArrowError> {
419 if let Dictionary(_, _) = *data_type {
420 let null_buffer = (field_node.null_count() > 0).then_some(buffers[0].clone());
421 let builder = ArrayData::builder(data_type.clone())
422 .len(field_node.length() as usize)
423 .add_buffer(buffers[1].clone())
424 .add_child_data(value_array.into_data())
425 .null_bit_buffer(null_buffer)
426 .null_count(field_node.null_count() as usize);
427 self.create_array_from_builder(builder)
428 } else {
429 unreachable!("Cannot create dictionary array from {:?}", data_type)
430 }
431 }
432}
433
434pub struct RecordBatchDecoder<'a> {
440 batch: crate::RecordBatch<'a>,
442 schema: SchemaRef,
444 dictionaries_by_id: &'a HashMap<i64, ArrayRef>,
446 compression: Option<CompressionCodec>,
448 decompression_context: DecompressionContext,
450 version: MetadataVersion,
452 data: &'a Buffer,
454 nodes: VectorIter<'a, FieldNode>,
456 buffers: VectorIter<'a, crate::Buffer>,
458 projection: Option<&'a [usize]>,
461 require_alignment: bool,
464 skip_validation: UnsafeFlag,
468}
469
470impl<'a> RecordBatchDecoder<'a> {
471 pub fn try_new(
473 buf: &'a Buffer,
474 batch: crate::RecordBatch<'a>,
475 schema: SchemaRef,
476 dictionaries_by_id: &'a HashMap<i64, ArrayRef>,
477 metadata: &'a MetadataVersion,
478 ) -> Result<Self, ArrowError> {
479 let buffers = batch.buffers().ok_or_else(|| {
480 ArrowError::IpcError("Unable to get buffers from IPC RecordBatch".to_string())
481 })?;
482 let field_nodes = batch.nodes().ok_or_else(|| {
483 ArrowError::IpcError("Unable to get field nodes from IPC RecordBatch".to_string())
484 })?;
485
486 let batch_compression = batch.compression();
487 let compression = batch_compression
488 .map(|batch_compression| batch_compression.codec().try_into())
489 .transpose()?;
490
491 Ok(Self {
492 batch,
493 schema,
494 dictionaries_by_id,
495 compression,
496 decompression_context: DecompressionContext::new(),
497 version: *metadata,
498 data: buf,
499 nodes: field_nodes.iter(),
500 buffers: buffers.iter(),
501 projection: None,
502 require_alignment: false,
503 skip_validation: UnsafeFlag::new(),
504 })
505 }
506
507 pub fn with_projection(mut self, projection: Option<&'a [usize]>) -> Self {
512 self.projection = projection;
513 self
514 }
515
516 pub fn with_require_alignment(mut self, require_alignment: bool) -> Self {
522 self.require_alignment = require_alignment;
523 self
524 }
525
526 pub fn with_skip_validation(mut self, skip_validation: UnsafeFlag) -> Self {
539 self.skip_validation = skip_validation;
540 self
541 }
542
543 pub fn read_record_batch(mut self) -> Result<RecordBatch, ArrowError> {
550 let mut variadic_counts: VecDeque<i64> = self
551 .batch
552 .variadicBufferCounts()
553 .into_iter()
554 .flatten()
555 .collect();
556
557 let options = RecordBatchOptions::new().with_row_count(Some(self.batch.length() as usize));
558
559 let schema = Arc::clone(&self.schema);
560 if let Some(projection) = self.projection {
561 let mut arrays = Vec::with_capacity(projection.len());
562 for (idx, field) in schema.fields().iter().enumerate() {
564 let mut decoded = None;
566 for (proj_idx, projected_idx) in projection.iter().enumerate() {
567 if *projected_idx == idx {
568 let child = match decoded.clone() {
570 Some(child) => child,
571 None => {
572 let child = self.create_array(field, &mut variadic_counts)?;
573 decoded = Some(Arc::clone(&child));
574 child
575 }
576 };
577 arrays.push((proj_idx, child));
578 }
579 }
580
581 if decoded.is_none() {
582 self.skip_field(field, &mut variadic_counts)?;
583 }
584 }
585
586 arrays.sort_by_key(|t| t.0);
587
588 let schema = Arc::new(schema.project(projection)?);
589 let columns = arrays.into_iter().map(|t| t.1).collect::<Vec<_>>();
590
591 if self.skip_validation.get() {
592 unsafe {
594 Ok(RecordBatch::new_unchecked(
595 schema,
596 columns,
597 self.batch.length() as usize,
598 ))
599 }
600 } else {
601 check_variadic_counts_consumed(&variadic_counts)?;
602 RecordBatch::try_new_with_options(schema, columns, &options)
603 }
604 } else {
605 let mut children = Vec::with_capacity(schema.fields().len());
606 for field in schema.fields() {
608 let child = self.create_array(field, &mut variadic_counts)?;
609 children.push(child);
610 }
611
612 if self.skip_validation.get() {
613 unsafe {
615 Ok(RecordBatch::new_unchecked(
616 schema,
617 children,
618 self.batch.length() as usize,
619 ))
620 }
621 } else {
622 check_variadic_counts_consumed(&variadic_counts)?;
623 RecordBatch::try_new_with_options(schema, children, &options)
624 }
625 }
626 }
627
628 fn next_buffer(&mut self) -> Result<Buffer, ArrowError> {
629 let buffer = self.buffers.next().ok_or_else(|| {
630 ArrowError::IpcError("Buffer count mismatched with metadata".to_string())
631 })?;
632 read_buffer(
633 buffer,
634 self.data,
635 self.compression,
636 &mut self.decompression_context,
637 )
638 }
639
640 fn skip_buffer(&mut self) -> Result<(), ArrowError> {
641 self.buffers.next().ok_or_else(|| {
642 ArrowError::IpcError("Buffer count mismatched with metadata".to_string())
643 })?;
644 Ok(())
645 }
646
647 fn next_node(&mut self, field: &Field) -> Result<&'a FieldNode, ArrowError> {
648 self.nodes.next().ok_or_else(|| {
649 ArrowError::SchemaError(format!(
650 "Invalid data for schema. {field} refers to node not found in schema",
651 ))
652 })
653 }
654
655 fn skip_field(
656 &mut self,
657 field: &Field,
658 variadic_count: &mut VecDeque<i64>,
659 ) -> Result<(), ArrowError> {
660 self.next_node(field)?;
661
662 match field.data_type() {
663 Utf8 | Binary | LargeBinary | LargeUtf8 => {
664 for _ in 0..3 {
665 self.skip_buffer()?;
666 }
667 }
668 Utf8View | BinaryView => {
669 let count = self.next_variadic_buffer_count(variadic_count, field.data_type())?;
670 for _ in 0..count {
671 self.skip_buffer()?;
672 }
673 }
674 FixedSizeBinary(_) => {
675 self.skip_buffer()?;
676 self.skip_buffer()?;
677 }
678 List(list_field) | LargeList(list_field) | Map(list_field, _) => {
679 self.skip_buffer()?;
680 self.skip_buffer()?;
681 self.skip_field(list_field, variadic_count)?;
682 }
683 ListView(list_field) | LargeListView(list_field) => {
684 self.skip_buffer()?; self.skip_buffer()?; self.skip_buffer()?; self.skip_field(list_field, variadic_count)?;
688 }
689 FixedSizeList(list_field, _) => {
690 self.skip_buffer()?;
691 self.skip_field(list_field, variadic_count)?;
692 }
693 Struct(struct_fields) => {
694 self.skip_buffer()?;
695
696 for struct_field in struct_fields {
698 self.skip_field(struct_field, variadic_count)?
699 }
700 }
701 RunEndEncoded(run_ends_field, values_field) => {
702 self.skip_field(run_ends_field, variadic_count)?;
703 self.skip_field(values_field, variadic_count)?;
704 }
705 Dictionary(_, _) => {
706 self.skip_buffer()?; self.skip_buffer()?; }
709 Union(fields, mode) => {
710 if self.version < MetadataVersion::V5 {
711 self.skip_buffer()?; }
713 self.skip_buffer()?; match mode {
716 UnionMode::Dense => self.skip_buffer()?, UnionMode::Sparse => {}
718 }
719
720 for (_, field) in fields.iter() {
721 self.skip_field(field, variadic_count)?
722 }
723 }
724 Null => {}
726
727 Boolean
729 | Int8
730 | Int16
731 | Int32
732 | Int64
733 | UInt8
734 | UInt16
735 | UInt32
736 | UInt64
737 | Float16
738 | Float32
739 | Float64
740 | Timestamp(_, _)
741 | Date32
742 | Date64
743 | Time32(_)
744 | Time64(_)
745 | Duration(_)
746 | Interval(_)
747 | Decimal32(_, _)
748 | Decimal64(_, _)
749 | Decimal128(_, _)
750 | Decimal256(_, _) => {
751 self.skip_buffer()?;
752 self.skip_buffer()?;
753 }
754 }
755 Ok(())
756 }
757}
758
759impl RecordBatchDecoder<'_> {
760 fn next_variadic_buffer_count(
766 &self,
767 variadic_counts: &mut VecDeque<i64>,
768 data_type: &DataType,
769 ) -> Result<usize, ArrowError> {
770 let count = variadic_counts.pop_front().ok_or_else(|| {
771 ArrowError::IpcError(format!("Missing variadic count for {data_type} column"))
772 })?;
773
774 let remaining = self.buffers.len();
775
776 usize::try_from(count)
778 .ok()
779 .and_then(|count| count.checked_add(2))
780 .filter(|total| *total <= remaining)
781 .ok_or_else(|| {
782 ArrowError::IpcError(format!(
783 "Invalid variadic count {count} for {data_type} column, \
784 with {remaining} buffer(s) left in the message"
785 ))
786 })
787 }
788}
789
790pub fn read_record_batch(
801 buf: &Buffer,
802 batch: crate::RecordBatch,
803 schema: SchemaRef,
804 dictionaries_by_id: &HashMap<i64, ArrayRef>,
805 projection: Option<&[usize]>,
806 metadata: &MetadataVersion,
807) -> Result<RecordBatch, ArrowError> {
808 RecordBatchDecoder::try_new(buf, batch, schema, dictionaries_by_id, metadata)?
809 .with_projection(projection)
810 .with_require_alignment(false)
811 .read_record_batch()
812}
813
814pub fn read_dictionary(
817 buf: &Buffer,
818 batch: crate::DictionaryBatch,
819 schema: &Schema,
820 dictionaries_by_id: &mut HashMap<i64, ArrayRef>,
821 metadata: &MetadataVersion,
822) -> Result<(), ArrowError> {
823 read_dictionary_impl(
824 buf,
825 batch,
826 schema,
827 dictionaries_by_id,
828 metadata,
829 false,
830 UnsafeFlag::new(),
831 )
832}
833
834pub fn read_dictionary_impl(
836 buf: &Buffer,
837 batch: crate::DictionaryBatch,
838 schema: &Schema,
839 dictionaries_by_id: &mut HashMap<i64, ArrayRef>,
840 metadata: &MetadataVersion,
841 require_alignment: bool,
842 skip_validation: UnsafeFlag,
843) -> Result<(), ArrowError> {
844 let id = batch.id();
845
846 let dictionary_values = get_dictionary_values(
847 buf,
848 batch,
849 schema,
850 dictionaries_by_id,
851 metadata,
852 require_alignment,
853 skip_validation,
854 )?;
855
856 update_dictionaries(dictionaries_by_id, batch.isDelta(), id, dictionary_values)?;
857
858 Ok(())
859}
860
861fn update_dictionaries(
870 dictionaries_by_id: &mut HashMap<i64, ArrayRef>,
871 is_delta: bool,
872 dict_id: i64,
873 dict_values: ArrayRef,
874) -> Result<(), ArrowError> {
875 if !is_delta {
876 dictionaries_by_id.insert(dict_id, dict_values.clone());
880 return Ok(());
881 }
882
883 let existing = dictionaries_by_id.get(&dict_id).ok_or_else(|| {
884 ArrowError::InvalidArgumentError(format!(
885 "No existing dictionary for delta dictionary with id '{dict_id}'"
886 ))
887 })?;
888
889 let combined = concat::concat(&[existing, &dict_values]).map_err(|e| {
890 ArrowError::InvalidArgumentError(format!("Failed to concat delta dictionary: {e}"))
891 })?;
892
893 dictionaries_by_id.insert(dict_id, combined);
894
895 Ok(())
896}
897
898fn get_dictionary_values(
902 buf: &Buffer,
903 batch: crate::DictionaryBatch,
904 schema: &Schema,
905 dictionaries_by_id: &HashMap<i64, ArrayRef>,
906 metadata: &MetadataVersion,
907 require_alignment: bool,
908 skip_validation: UnsafeFlag,
909) -> Result<ArrayRef, ArrowError> {
910 let id = batch.id();
911 #[expect(deprecated)]
912 let fields_using_this_dictionary = schema.fields_with_dict_id(id);
913 let first_field = fields_using_this_dictionary.first().ok_or_else(|| {
914 ArrowError::InvalidArgumentError(format!("dictionary id {id} not found in schema"))
915 })?;
916
917 let dictionary_values: ArrayRef = match first_field.data_type() {
921 DataType::Dictionary(_, value_type) => {
922 let value = value_type.as_ref().clone();
924 let schema = Schema::new(vec![Field::new("", value, true)]);
925 let Some(data) = batch.data() else {
927 return Err(ArrowError::ParseError(
928 "Dictionary batch is missing its data".to_string(),
929 ));
930 };
931 let record_batch = RecordBatchDecoder::try_new(
932 buf,
933 data,
934 Arc::new(schema),
935 dictionaries_by_id,
936 metadata,
937 )?
938 .with_require_alignment(require_alignment)
939 .with_skip_validation(skip_validation)
940 .read_record_batch()?;
941
942 Some(record_batch.column(0).clone())
943 }
944 _ => None,
945 }
946 .ok_or_else(|| {
947 ArrowError::InvalidArgumentError(format!("dictionary id {id} not found in schema"))
948 })?;
949
950 Ok(dictionary_values)
951}
952
953fn read_block<R: Read + Seek>(mut reader: R, block: &Block) -> Result<Buffer, ArrowError> {
955 reader.seek(SeekFrom::Start(block.offset() as u64))?;
956 let body_len = block.bodyLength().to_usize().unwrap();
957 let metadata_len = block.metaDataLength().to_usize().unwrap();
958 let total_len = body_len.checked_add(metadata_len).unwrap();
959
960 let mut buf = MutableBuffer::try_from_len_zeroed(total_len)
961 .map_err(|e| ArrowError::MemoryError(e.to_string()))?;
962 reader.read_exact(&mut buf)?;
963 Ok(buf.into())
964}
965
966fn check_variadic_counts_consumed(variadic_counts: &VecDeque<i64>) -> Result<(), ArrowError> {
971 if variadic_counts.is_empty() {
972 Ok(())
973 } else {
974 Err(ArrowError::IpcError(format!(
975 "Mismatch between schema and data: the IPC message declares {} more variadic \
976 buffer count(s) than the schema has BinaryView or Utf8View columns",
977 variadic_counts.len()
978 )))
979 }
980}
981
982fn parse_message(buf: &[u8]) -> Result<Message::Message<'_>, ArrowError> {
986 let buf = match buf[..4] == CONTINUATION_MARKER {
987 true => &buf[8..],
988 false => &buf[4..],
989 };
990 crate::root_as_message(buf)
991 .map_err(|err| ArrowError::ParseError(format!("Unable to get root as message: {err:?}")))
992}
993
994pub fn read_footer_length(buf: [u8; 10]) -> Result<usize, ArrowError> {
998 if buf[4..] != super::ARROW_MAGIC {
999 return Err(ArrowError::ParseError(
1000 "Arrow file does not contain correct footer".to_string(),
1001 ));
1002 }
1003
1004 let footer_len = i32::from_le_bytes([buf[0], buf[1], buf[2], buf[3]]);
1006 footer_len
1007 .try_into()
1008 .map_err(|_| ArrowError::ParseError(format!("Invalid footer length: {footer_len}")))
1009}
1010
1011#[derive(Debug)]
1076pub struct FileDecoder {
1077 schema: SchemaRef,
1078 dictionaries: HashMap<i64, ArrayRef>,
1079 version: MetadataVersion,
1080 projection: Option<Vec<usize>>,
1081 require_alignment: bool,
1082 skip_validation: UnsafeFlag,
1083}
1084
1085impl FileDecoder {
1086 pub fn new(schema: SchemaRef, version: MetadataVersion) -> Self {
1088 Self {
1089 schema,
1090 version,
1091 dictionaries: Default::default(),
1092 projection: None,
1093 require_alignment: false,
1094 skip_validation: UnsafeFlag::new(),
1095 }
1096 }
1097
1098 pub fn with_projection(mut self, projection: Vec<usize>) -> Self {
1100 self.projection = Some(projection);
1101 self
1102 }
1103
1104 pub fn with_require_alignment(mut self, require_alignment: bool) -> Self {
1117 self.require_alignment = require_alignment;
1118 self
1119 }
1120
1121 pub unsafe fn with_skip_validation(mut self, skip_validation: bool) -> Self {
1132 unsafe { self.skip_validation.set(skip_validation) };
1133 self
1134 }
1135
1136 fn read_message<'a>(&self, buf: &'a [u8]) -> Result<Message::Message<'a>, ArrowError> {
1137 let message = parse_message(buf)?;
1138
1139 if self.version != MetadataVersion::V1 && message.version() != self.version {
1141 return Err(ArrowError::IpcError(
1142 "Could not read IPC message as metadata versions mismatch".to_string(),
1143 ));
1144 }
1145 Ok(message)
1146 }
1147
1148 pub fn read_dictionary(&mut self, block: &Block, buf: &Buffer) -> Result<(), ArrowError> {
1150 let message = self.read_message(buf)?;
1151 match message.header_type() {
1152 crate::MessageHeader::DictionaryBatch => {
1153 let batch = message.header_as_dictionary_batch().unwrap();
1154 read_dictionary_impl(
1155 &buf.slice(block.metaDataLength() as _),
1156 batch,
1157 &self.schema,
1158 &mut self.dictionaries,
1159 &message.version(),
1160 self.require_alignment,
1161 self.skip_validation.clone(),
1162 )
1163 }
1164 t => Err(ArrowError::ParseError(format!(
1165 "Expecting DictionaryBatch in dictionary blocks, found {t:?}."
1166 ))),
1167 }
1168 }
1169
1170 pub fn read_record_batch(
1172 &self,
1173 block: &Block,
1174 buf: &Buffer,
1175 ) -> Result<Option<RecordBatch>, ArrowError> {
1176 let message = self.read_message(buf)?;
1177 match message.header_type() {
1178 crate::MessageHeader::Schema => Err(ArrowError::IpcError(
1179 "Not expecting a schema when messages are read".to_string(),
1180 )),
1181 crate::MessageHeader::RecordBatch => {
1182 let batch = message.header_as_record_batch().ok_or_else(|| {
1183 ArrowError::IpcError("Unable to read IPC message as record batch".to_string())
1184 })?;
1185 RecordBatchDecoder::try_new(
1187 &buf.slice(block.metaDataLength() as _),
1188 batch,
1189 self.schema.clone(),
1190 &self.dictionaries,
1191 &message.version(),
1192 )?
1193 .with_projection(self.projection.as_deref())
1194 .with_require_alignment(self.require_alignment)
1195 .with_skip_validation(self.skip_validation.clone())
1196 .read_record_batch()
1197 .map(Some)
1198 }
1199 crate::MessageHeader::NONE => Ok(None),
1200 t => Err(ArrowError::InvalidArgumentError(format!(
1201 "Reading types other than record batches not yet supported, unable to read {t:?}"
1202 ))),
1203 }
1204 }
1205}
1206
1207#[derive(Debug)]
1209pub struct FileReaderBuilder {
1210 projection: Option<Vec<usize>>,
1212 max_footer_fb_tables: usize,
1214 max_footer_fb_depth: usize,
1216}
1217
1218impl Default for FileReaderBuilder {
1219 fn default() -> Self {
1220 let verifier_options = VerifierOptions::default();
1221 Self {
1222 max_footer_fb_tables: verifier_options.max_tables,
1223 max_footer_fb_depth: verifier_options.max_depth,
1224 projection: None,
1225 }
1226 }
1227}
1228
1229impl FileReaderBuilder {
1230 pub fn new() -> Self {
1234 Self::default()
1235 }
1236
1237 pub fn with_projection(mut self, projection: Vec<usize>) -> Self {
1239 self.projection = Some(projection);
1240 self
1241 }
1242
1243 pub fn with_max_footer_fb_tables(mut self, max_footer_fb_tables: usize) -> Self {
1256 self.max_footer_fb_tables = max_footer_fb_tables;
1257 self
1258 }
1259
1260 pub fn with_max_footer_fb_depth(mut self, max_footer_fb_depth: usize) -> Self {
1273 self.max_footer_fb_depth = max_footer_fb_depth;
1274 self
1275 }
1276
1277 pub fn build<R: Read + Seek>(self, mut reader: R) -> Result<FileReader<R>, ArrowError> {
1279 let mut buffer = [0; 10];
1281 reader.seek(SeekFrom::End(-10))?;
1282 reader.read_exact(&mut buffer)?;
1283
1284 let footer_len = read_footer_length(buffer)?;
1285
1286 let mut footer_data = vec![0; footer_len];
1288 reader.seek(SeekFrom::End(-10 - footer_len as i64))?;
1289 reader.read_exact(&mut footer_data)?;
1290
1291 let verifier_options = VerifierOptions {
1292 max_tables: self.max_footer_fb_tables,
1293 max_depth: self.max_footer_fb_depth,
1294 ..Default::default()
1295 };
1296 let footer = crate::root_as_footer_with_opts(&verifier_options, &footer_data[..]).map_err(
1297 |err| ArrowError::ParseError(format!("Unable to get root as footer: {err:?}")),
1298 )?;
1299
1300 let blocks = footer.recordBatches().ok_or_else(|| {
1301 ArrowError::ParseError("Unable to get record batches from IPC Footer".to_string())
1302 })?;
1303
1304 let total_blocks = blocks.len();
1305
1306 let ipc_schema = footer.schema().ok_or_else(|| {
1307 ArrowError::ParseError("Unable to get schema from IPC Footer".to_string())
1308 })?;
1309 if !ipc_schema.endianness().equals_to_target_endianness() {
1310 return Err(ArrowError::IpcError(
1311 "the endianness of the source system does not match the endianness of the target system.".to_owned()
1312 ));
1313 }
1314
1315 let schema = Arc::new(crate::convert::try_fb_to_schema(ipc_schema)?);
1316
1317 let projected_schema = match &self.projection {
1318 Some(projection) => Arc::new(schema.project(projection)?),
1319 None => schema.clone(),
1320 };
1321
1322 let mut custom_metadata = HashMap::new();
1323 if let Some(fb_custom_metadata) = footer.custom_metadata() {
1324 for kv in fb_custom_metadata {
1325 let (Some(key), Some(value)) = (kv.key(), kv.value()) else {
1326 return Err(ArrowError::ParseError(
1327 "Custom metadata in the IPC footer is missing a key or a value".to_string(),
1328 ));
1329 };
1330 custom_metadata.insert(key.to_string(), value.to_string());
1331 }
1332 }
1333
1334 let mut decoder = FileDecoder::new(schema, footer.version());
1335 if let Some(projection) = self.projection {
1336 decoder = decoder.with_projection(projection)
1337 }
1338
1339 if let Some(dictionaries) = footer.dictionaries() {
1341 for block in dictionaries {
1342 let buf = read_block(&mut reader, block)?;
1343 decoder.read_dictionary(block, &buf)?;
1344 }
1345 }
1346
1347 Ok(FileReader {
1348 reader,
1349 blocks: blocks.iter().copied().collect(),
1350 current_block: 0,
1351 total_blocks,
1352 decoder,
1353 schema: projected_schema,
1354 custom_metadata,
1355 })
1356 }
1357}
1358
1359pub struct FileReader<R> {
1404 reader: R,
1406
1407 decoder: FileDecoder,
1409
1410 schema: SchemaRef,
1412
1413 blocks: Vec<Block>,
1417
1418 current_block: usize,
1420
1421 total_blocks: usize,
1423
1424 custom_metadata: HashMap<String, String>,
1426}
1427
1428impl<R> fmt::Debug for FileReader<R> {
1429 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> Result<(), fmt::Error> {
1430 f.debug_struct("FileReader<R>")
1431 .field("decoder", &self.decoder)
1432 .field("blocks", &self.blocks)
1433 .field("current_block", &self.current_block)
1434 .field("total_blocks", &self.total_blocks)
1435 .finish_non_exhaustive()
1436 }
1437}
1438
1439impl<R: Read + Seek> FileReader<BufReader<R>> {
1440 pub fn try_new_buffered(reader: R, projection: Option<Vec<usize>>) -> Result<Self, ArrowError> {
1444 Self::try_new(BufReader::new(reader), projection)
1445 }
1446}
1447
1448impl<R: Read + Seek> FileReader<R> {
1449 pub fn try_new(reader: R, projection: Option<Vec<usize>>) -> Result<Self, ArrowError> {
1461 let builder = FileReaderBuilder {
1462 projection,
1463 ..Default::default()
1464 };
1465 builder.build(reader)
1466 }
1467
1468 pub fn custom_metadata(&self) -> &HashMap<String, String> {
1470 &self.custom_metadata
1471 }
1472
1473 pub fn num_batches(&self) -> usize {
1475 self.total_blocks
1476 }
1477
1478 pub fn schema(&self) -> SchemaRef {
1480 self.schema.clone()
1481 }
1482
1483 pub fn set_index(&mut self, index: usize) -> Result<(), ArrowError> {
1487 if index >= self.total_blocks {
1488 Err(ArrowError::InvalidArgumentError(format!(
1489 "Cannot set batch to index {} from {} total batches",
1490 index, self.total_blocks
1491 )))
1492 } else {
1493 self.current_block = index;
1494 Ok(())
1495 }
1496 }
1497
1498 fn maybe_next(&mut self) -> Result<Option<RecordBatch>, ArrowError> {
1499 let block = &self.blocks[self.current_block];
1500 self.current_block += 1;
1501
1502 let buffer = read_block(&mut self.reader, block)?;
1504 self.decoder.read_record_batch(block, &buffer)
1505 }
1506
1507 pub fn get_ref(&self) -> &R {
1511 &self.reader
1512 }
1513
1514 pub fn get_mut(&mut self) -> &mut R {
1518 &mut self.reader
1519 }
1520
1521 pub unsafe fn with_skip_validation(mut self, skip_validation: bool) -> Self {
1527 self.decoder = unsafe { self.decoder.with_skip_validation(skip_validation) };
1528 self
1529 }
1530}
1531
1532impl<R: Read + Seek> Iterator for FileReader<R> {
1533 type Item = Result<RecordBatch, ArrowError>;
1534
1535 fn next(&mut self) -> Option<Self::Item> {
1536 if self.current_block < self.total_blocks {
1538 self.maybe_next().transpose()
1539 } else {
1540 None
1541 }
1542 }
1543}
1544
1545impl<R: Read + Seek> RecordBatchReader for FileReader<R> {
1546 fn schema(&self) -> SchemaRef {
1547 self.schema()
1548 }
1549}
1550
1551pub struct StreamReader<R> {
1585 reader: MessageReader<R>,
1587
1588 schema: SchemaRef,
1590
1591 dictionaries_by_id: HashMap<i64, ArrayRef>,
1595
1596 finished: bool,
1600
1601 projection: Option<(Vec<usize>, SchemaRef)>,
1603
1604 skip_validation: UnsafeFlag,
1608}
1609
1610impl<R> fmt::Debug for StreamReader<R> {
1611 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> std::result::Result<(), fmt::Error> {
1612 f.debug_struct("StreamReader<R>")
1613 .field("reader", &"R")
1614 .field("schema", &self.schema)
1615 .field("dictionaries_by_id", &self.dictionaries_by_id)
1616 .field("finished", &self.finished)
1617 .field("projection", &self.projection)
1618 .finish()
1619 }
1620}
1621
1622impl<R: Read> StreamReader<BufReader<R>> {
1623 pub fn try_new_buffered(reader: R, projection: Option<Vec<usize>>) -> Result<Self, ArrowError> {
1627 Self::try_new(BufReader::new(reader), projection)
1628 }
1629}
1630
1631impl<R: Read> StreamReader<R> {
1632 pub fn try_new(
1644 reader: R,
1645 projection: Option<Vec<usize>>,
1646 ) -> Result<StreamReader<R>, ArrowError> {
1647 let mut msg_reader = MessageReader::new(reader);
1648 let message = msg_reader.maybe_next()?;
1649 let Some((message, _)) = message else {
1650 return Err(ArrowError::IpcError(
1651 "Expected schema message, found empty stream.".to_string(),
1652 ));
1653 };
1654
1655 if message.header_type() != Message::MessageHeader::Schema {
1656 return Err(ArrowError::IpcError(format!(
1657 "Expected a schema as the first message in the stream, got: {:?}",
1658 message.header_type()
1659 )));
1660 }
1661
1662 let schema = message.header_as_schema().ok_or_else(|| {
1663 ArrowError::ParseError("Failed to parse schema from message header".to_string())
1664 })?;
1665 let schema = crate::convert::try_fb_to_schema(schema)?;
1666
1667 let dictionaries_by_id = HashMap::new();
1669
1670 let projection = match projection {
1671 Some(projection_indices) => {
1672 let schema = Arc::new(schema.project(&projection_indices)?);
1673 Some((projection_indices, schema))
1674 }
1675 _ => None,
1676 };
1677
1678 Ok(Self {
1679 reader: msg_reader,
1680 schema: Arc::new(schema),
1681 finished: false,
1682 dictionaries_by_id,
1683 projection,
1684 skip_validation: UnsafeFlag::new(),
1685 })
1686 }
1687
1688 pub fn schema(&self) -> SchemaRef {
1690 match &self.projection {
1691 Some((_, projected_schema)) => projected_schema.clone(),
1692 None => self.schema.clone(),
1693 }
1694 }
1695
1696 pub fn is_finished(&self) -> bool {
1698 self.finished
1699 }
1700
1701 fn maybe_next(&mut self) -> Result<Option<RecordBatch>, ArrowError> {
1702 if self.finished {
1703 return Ok(None);
1704 }
1705
1706 loop {
1708 let message = self.next_ipc_message()?;
1709 let Some(message) = message else {
1710 self.finished = true;
1712 return Ok(None);
1713 };
1714
1715 match message {
1716 IpcMessage::Schema(_) => {
1717 return Err(ArrowError::IpcError(
1718 "Expected a record batch, but found a schema".to_string(),
1719 ));
1720 }
1721 IpcMessage::RecordBatch(record_batch) => {
1722 return Ok(Some(record_batch));
1723 }
1724 IpcMessage::DictionaryBatch { .. } => {}
1725 }
1726 }
1727 }
1728
1729 pub(crate) fn next_ipc_message(&mut self) -> Result<Option<IpcMessage>, ArrowError> {
1737 let message = self.reader.maybe_next()?;
1738 let Some((message, body)) = message else {
1739 return Ok(None);
1741 };
1742
1743 let ipc_message = match message.header_type() {
1744 Message::MessageHeader::Schema => {
1745 let schema = message.header_as_schema().ok_or_else(|| {
1746 ArrowError::ParseError("Failed to parse schema from message header".to_string())
1747 })?;
1748 let arrow_schema = crate::convert::try_fb_to_schema(schema)?;
1749 IpcMessage::Schema(arrow_schema)
1750 }
1751 Message::MessageHeader::RecordBatch => {
1752 let batch = message.header_as_record_batch().ok_or_else(|| {
1753 ArrowError::IpcError("Unable to read IPC message as record batch".to_string())
1754 })?;
1755
1756 let version = message.version();
1757 let schema = self.schema.clone();
1758 let record_batch = RecordBatchDecoder::try_new(
1759 &body.into(),
1760 batch,
1761 schema,
1762 &self.dictionaries_by_id,
1763 &version,
1764 )?
1765 .with_projection(self.projection.as_ref().map(|x| x.0.as_ref()))
1766 .with_require_alignment(false)
1767 .with_skip_validation(self.skip_validation.clone())
1768 .read_record_batch()?;
1769 IpcMessage::RecordBatch(record_batch)
1770 }
1771 Message::MessageHeader::DictionaryBatch => {
1772 let dict = message.header_as_dictionary_batch().ok_or_else(|| {
1773 ArrowError::ParseError(
1774 "Failed to parse dictionary batch from message header".to_string(),
1775 )
1776 })?;
1777
1778 let version = message.version();
1779 let dict_values = get_dictionary_values(
1780 &body.into(),
1781 dict,
1782 &self.schema,
1783 &self.dictionaries_by_id,
1784 &version,
1785 false,
1786 self.skip_validation.clone(),
1787 )?;
1788
1789 update_dictionaries(
1790 &mut self.dictionaries_by_id,
1791 dict.isDelta(),
1792 dict.id(),
1793 dict_values.clone(),
1794 )?;
1795
1796 IpcMessage::DictionaryBatch {
1797 id: dict.id(),
1798 is_delta: (dict.isDelta()),
1799 values: (dict_values),
1800 }
1801 }
1802 x => {
1803 return Err(ArrowError::ParseError(format!(
1804 "Unsupported message header type in IPC stream: '{x:?}'"
1805 )));
1806 }
1807 };
1808
1809 Ok(Some(ipc_message))
1810 }
1811
1812 pub fn get_ref(&self) -> &R {
1816 self.reader.inner()
1817 }
1818
1819 pub fn get_mut(&mut self) -> &mut R {
1823 self.reader.inner_mut()
1824 }
1825
1826 pub unsafe fn with_skip_validation(mut self, skip_validation: bool) -> Self {
1832 unsafe { self.skip_validation.set(skip_validation) };
1833 self
1834 }
1835}
1836
1837impl<R: Read> Iterator for StreamReader<R> {
1838 type Item = Result<RecordBatch, ArrowError>;
1839
1840 fn next(&mut self) -> Option<Self::Item> {
1841 self.maybe_next().transpose()
1842 }
1843}
1844
1845impl<R: Read> RecordBatchReader for StreamReader<R> {
1846 fn schema(&self) -> SchemaRef {
1847 self.schema()
1848 }
1849}
1850
1851#[derive(Debug)]
1857#[expect(dead_code)]
1858pub(crate) enum IpcMessage {
1859 Schema(arrow_schema::Schema),
1860 RecordBatch(RecordBatch),
1861 DictionaryBatch {
1862 id: i64,
1863 is_delta: bool,
1864 values: ArrayRef,
1865 },
1866}
1867
1868const MAX_PREALLOC_BYTES: usize = 64 * 1024 * 1024;
1884
1885fn read_body_bounded<R: Read>(reader: &mut R, len: usize) -> Result<MutableBuffer, ArrowError> {
1887 let mut buf = MutableBuffer::try_from_len_zeroed(len.min(MAX_PREALLOC_BYTES))
1888 .map_err(|e| ArrowError::MemoryError(e.to_string()))?;
1889 let mut filled = 0;
1890 while filled < len {
1891 let target = buf.len();
1892 reader.read_exact(&mut buf.as_slice_mut()[filled..target])?;
1893 filled = target;
1894 if filled < len {
1895 buf.try_resize(len.min(target.saturating_mul(2)), 0)
1896 .map_err(|e| ArrowError::MemoryError(e.to_string()))?;
1897 }
1898 }
1899 Ok(buf)
1900}
1901
1902struct MessageReader<R> {
1905 reader: R,
1906 buf: Vec<u8>,
1907}
1908
1909impl<R: Read> MessageReader<R> {
1910 fn new(reader: R) -> Self {
1911 Self {
1912 reader,
1913 buf: Vec::new(),
1914 }
1915 }
1916
1917 fn maybe_next(&mut self) -> Result<Option<(Message::Message<'_>, MutableBuffer)>, ArrowError> {
1928 let meta_len = self.read_meta_len()?;
1929 let Some(meta_len) = meta_len else {
1930 return Ok(None);
1931 };
1932
1933 self.buf.clear();
1936 let read = (&mut self.reader)
1937 .take(meta_len as u64)
1938 .read_to_end(&mut self.buf)?;
1939 if read != meta_len {
1940 return Err(ArrowError::ParseError(format!(
1941 "Unexpected end of stream: expected {meta_len} metadata bytes, got {read}"
1942 )));
1943 }
1944
1945 let message = crate::root_as_message(self.buf.as_slice()).map_err(|err| {
1946 ArrowError::ParseError(format!("Unable to get root as message: {err:?}"))
1947 })?;
1948
1949 let body_len = usize::try_from(message.bodyLength()).map_err(|_| {
1950 ArrowError::ParseError(format!(
1951 "Invalid IPC message body length: {}",
1952 message.bodyLength()
1953 ))
1954 })?;
1955 let buf = read_body_bounded(&mut self.reader, body_len)?;
1956
1957 Ok(Some((message, buf)))
1958 }
1959
1960 fn inner_mut(&mut self) -> &mut R {
1962 &mut self.reader
1963 }
1964
1965 fn inner(&self) -> &R {
1967 &self.reader
1968 }
1969
1970 pub fn read_meta_len(&mut self) -> Result<Option<usize>, ArrowError> {
1979 let mut meta_len: [u8; 4] = [0; 4];
1980 match self.reader.read_exact(&mut meta_len) {
1981 Ok(()) => {}
1982 Err(e) => {
1983 return if e.kind() == std::io::ErrorKind::UnexpectedEof {
1984 Ok(None)
1988 } else {
1989 Err(ArrowError::from(e))
1990 };
1991 }
1992 }
1993
1994 let meta_len = {
1995 if meta_len == CONTINUATION_MARKER {
1998 self.reader.read_exact(&mut meta_len)?;
1999 }
2000
2001 i32::from_le_bytes(meta_len)
2002 };
2003
2004 if meta_len == 0 {
2005 return Ok(None);
2006 }
2007
2008 let meta_len = usize::try_from(meta_len)
2009 .map_err(|_| ArrowError::ParseError(format!("Invalid metadata length: {meta_len}")))?;
2010
2011 Ok(Some(meta_len))
2012 }
2013}
2014
2015#[cfg(test)]
2016mod tests {
2017 use std::io::Cursor;
2018
2019 use crate::convert::try_fb_to_schema;
2020 use crate::writer::{
2021 DictionaryTracker, IpcDataGenerator, IpcWriteOptions, unslice_run_array, write_message,
2022 };
2023
2024 use super::*;
2025
2026 use crate::{root_as_footer, root_as_message, size_prefixed_root_as_message};
2027 use arrow_array::builder::{PrimitiveRunBuilder, UnionBuilder};
2028 use arrow_array::types::*;
2029 use arrow_buffer::{NullBuffer, OffsetBuffer};
2030 use arrow_data::ArrayDataBuilder;
2031
2032 fn create_test_projection_schema() -> Schema {
2033 let list_data_type = DataType::List(Arc::new(Field::new_list_field(DataType::Int32, true)));
2035
2036 let fixed_size_list_data_type =
2037 DataType::FixedSizeList(Arc::new(Field::new_list_field(DataType::Int32, false)), 3);
2038
2039 let union_fields = UnionFields::from_fields(vec![
2040 Field::new("a", DataType::Int32, false),
2041 Field::new("b", DataType::Float64, false),
2042 ]);
2043
2044 let union_data_type = DataType::Union(union_fields, UnionMode::Dense);
2045
2046 let struct_fields = Fields::from(vec![
2047 Field::new("id", DataType::Int32, false),
2048 Field::new_list("list", Field::new_list_field(DataType::Int8, true), false),
2049 ]);
2050 let struct_data_type = DataType::Struct(struct_fields);
2051
2052 let run_encoded_data_type = DataType::RunEndEncoded(
2053 Arc::new(Field::new("run_ends", DataType::Int16, false)),
2054 Arc::new(Field::new("values", DataType::Int32, true)),
2055 );
2056
2057 Schema::new(vec![
2059 Field::new("f0", DataType::UInt32, false),
2060 Field::new("f1", DataType::Utf8, false),
2061 Field::new("f2", DataType::Boolean, false),
2062 Field::new("f3", union_data_type, true),
2063 Field::new("f4", DataType::Null, true),
2064 Field::new("f5", DataType::Float64, true),
2065 Field::new("f6", list_data_type, false),
2066 Field::new("f7", DataType::FixedSizeBinary(3), true),
2067 Field::new("f8", fixed_size_list_data_type, false),
2068 Field::new("f9", struct_data_type, false),
2069 Field::new("f10", run_encoded_data_type, false),
2070 Field::new("f11", DataType::Boolean, false),
2071 Field::new_dictionary("f12", DataType::Int8, DataType::Utf8, false),
2072 Field::new("f13", DataType::Utf8, false),
2073 ])
2074 }
2075
2076 fn create_test_projection_batch_data(schema: &Schema) -> RecordBatch {
2077 let array0 = UInt32Array::from(vec![1, 2, 3]);
2079 let array1 = StringArray::from(vec!["foo", "bar", "baz"]);
2080 let array2 = BooleanArray::from(vec![true, false, true]);
2081
2082 let mut union_builder = UnionBuilder::new_dense();
2083 union_builder.append::<Int32Type>("a", 1).unwrap();
2084 union_builder.append::<Float64Type>("b", 10.1).unwrap();
2085 union_builder.append_null::<Float64Type>("b").unwrap();
2086 let array3 = union_builder.build().unwrap();
2087
2088 let array4 = NullArray::new(3);
2089 let array5 = Float64Array::from(vec![Some(1.1), None, Some(3.3)]);
2090 let array6_values = vec![
2091 Some(vec![Some(10), Some(10), Some(10)]),
2092 Some(vec![Some(20), Some(20), Some(20)]),
2093 Some(vec![Some(30), Some(30)]),
2094 ];
2095 let array6 = ListArray::from_iter_primitive::<Int32Type, _, _>(array6_values);
2096 let array7_values = vec![vec![11, 12, 13], vec![22, 23, 24], vec![33, 34, 35]];
2097 let array7 = FixedSizeBinaryArray::try_from_iter(array7_values.into_iter()).unwrap();
2098
2099 let array8_values = ArrayData::builder(DataType::Int32)
2100 .len(9)
2101 .add_buffer(Buffer::from_slice_ref([40, 41, 42, 43, 44, 45, 46, 47, 48]))
2102 .build()
2103 .unwrap();
2104 let array8_data = ArrayData::builder(schema.field(8).data_type().clone())
2105 .len(3)
2106 .add_child_data(array8_values)
2107 .build()
2108 .unwrap();
2109 let array8 = FixedSizeListArray::from(array8_data);
2110
2111 let array9_id: ArrayRef = Arc::new(Int32Array::from(vec![1001, 1002, 1003]));
2112 let array9_list: ArrayRef =
2113 Arc::new(ListArray::from_iter_primitive::<Int8Type, _, _>(vec![
2114 Some(vec![Some(-10)]),
2115 Some(vec![Some(-20), Some(-20), Some(-20)]),
2116 Some(vec![Some(-30)]),
2117 ]));
2118 let array9 = ArrayDataBuilder::new(schema.field(9).data_type().clone())
2119 .add_child_data(array9_id.into_data())
2120 .add_child_data(array9_list.into_data())
2121 .len(3)
2122 .build()
2123 .unwrap();
2124 let array9 = StructArray::from(array9);
2125
2126 let array10_input = vec![Some(1_i32), None, None];
2127 let mut array10_builder = PrimitiveRunBuilder::<Int16Type, Int32Type>::new();
2128 array10_builder.extend(array10_input);
2129 let array10 = array10_builder.finish();
2130
2131 let array11 = BooleanArray::from(vec![false, false, true]);
2132
2133 let array12_values = StringArray::from(vec!["x", "yy", "zzz"]);
2134 let array12_keys = Int8Array::from_iter_values([1, 1, 2]);
2135 let array12 = DictionaryArray::new(array12_keys, Arc::new(array12_values));
2136
2137 let array13 = StringArray::from(vec!["a", "bb", "ccc"]);
2138
2139 RecordBatch::try_new(
2141 Arc::new(schema.clone()),
2142 vec![
2143 Arc::new(array0),
2144 Arc::new(array1),
2145 Arc::new(array2),
2146 Arc::new(array3),
2147 Arc::new(array4),
2148 Arc::new(array5),
2149 Arc::new(array6),
2150 Arc::new(array7),
2151 Arc::new(array8),
2152 Arc::new(array9),
2153 Arc::new(array10),
2154 Arc::new(array11),
2155 Arc::new(array12),
2156 Arc::new(array13),
2157 ],
2158 )
2159 .unwrap()
2160 }
2161
2162 #[test]
2163 fn test_negative_meta_len_start_stream() {
2164 let bytes = i32::to_le_bytes(-1);
2165 let mut buf = vec![];
2166 buf.extend(CONTINUATION_MARKER);
2167 buf.extend(bytes);
2168
2169 let reader_err = StreamReader::try_new(Cursor::new(buf), None).err();
2170 assert!(reader_err.is_some());
2171 assert_eq!(
2172 reader_err.unwrap().to_string(),
2173 "Parser error: Invalid metadata length: -1"
2174 );
2175 }
2176
2177 #[test]
2178 fn test_negative_meta_len_mid_stream() {
2179 let schema = Schema::new(vec![Field::new("a", DataType::Int32, false)]);
2180 let mut buf = Vec::new();
2181 {
2182 let mut writer = crate::writer::StreamWriter::try_new(&mut buf, &schema).unwrap();
2183 let batch =
2184 RecordBatch::try_new(Arc::new(schema), vec![Arc::new(Int32Array::from(vec![1]))])
2185 .unwrap();
2186 writer.write(&batch).unwrap();
2187 }
2188
2189 let bytes = i32::to_le_bytes(-1);
2190 buf.extend(CONTINUATION_MARKER);
2191 buf.extend(bytes);
2192
2193 let mut reader = StreamReader::try_new(Cursor::new(buf), None).unwrap();
2194 assert!(reader.maybe_next().is_ok());
2196 let batch_err = reader.maybe_next().err();
2198 assert!(batch_err.is_some());
2199 assert_eq!(
2200 batch_err.unwrap().to_string(),
2201 "Parser error: Invalid metadata length: -1"
2202 );
2203 }
2204
2205 #[test]
2206 fn test_invalid_dictionary_batch_without_data() {
2207 use crate::r#gen::Message::*;
2208 use flatbuffers::FlatBufferBuilder;
2209
2210 let schema = Schema::new(vec![Field::new(
2211 "col",
2212 DataType::Dictionary(Box::new(DataType::Int8), Box::new(DataType::Utf8)),
2213 true,
2214 )]);
2215
2216 let mut fbb = FlatBufferBuilder::new();
2219 let batch_offset = DictionaryBatch::create(
2220 &mut fbb,
2221 &DictionaryBatchArgs {
2222 id: 0,
2223 data: None,
2224 isDelta: false,
2225 },
2226 );
2227 fbb.finish_minimal(batch_offset);
2228 let batch_bytes = fbb.finished_data().to_vec();
2229 let batch = flatbuffers::root::<DictionaryBatch>(&batch_bytes).unwrap();
2230
2231 let data_buffer = Buffer::from(vec![0u8; 0]);
2232 let mut dictionaries: HashMap<i64, ArrayRef> = HashMap::new();
2233
2234 let err = read_dictionary(
2235 &data_buffer,
2236 batch,
2237 &schema,
2238 &mut dictionaries,
2239 &MetadataVersion::V5,
2240 )
2241 .unwrap_err();
2242
2243 assert_eq!(
2244 err.to_string(),
2245 "Parser error: Dictionary batch is missing its data"
2246 );
2247 }
2248
2249 #[test]
2250 fn test_missing_buffer_metadata_error() {
2251 use crate::r#gen::Message::*;
2252 use flatbuffers::FlatBufferBuilder;
2253
2254 let schema = Arc::new(Schema::new(vec![Field::new("col", DataType::Int32, true)]));
2255
2256 let mut fbb = FlatBufferBuilder::new();
2259 let nodes = fbb.create_vector(&[FieldNode::new(2, 0)]);
2260 let buffers = fbb.create_vector(&[crate::Buffer::new(0, 8)]);
2261 let batch_offset = RecordBatch::create(
2262 &mut fbb,
2263 &RecordBatchArgs {
2264 length: 2,
2265 nodes: Some(nodes),
2266 buffers: Some(buffers),
2267 compression: None,
2268 variadicBufferCounts: None,
2269 },
2270 );
2271 fbb.finish_minimal(batch_offset);
2272 let batch_bytes = fbb.finished_data().to_vec();
2273 let batch = flatbuffers::root::<RecordBatch>(&batch_bytes).unwrap();
2274
2275 let data_buffer = Buffer::from(vec![0u8; 8]);
2276 let dictionaries: HashMap<i64, ArrayRef> = HashMap::new();
2277 let metadata = MetadataVersion::V5;
2278
2279 let decoder = RecordBatchDecoder::try_new(
2280 &data_buffer,
2281 batch,
2282 schema.clone(),
2283 &dictionaries,
2284 &metadata,
2285 )
2286 .unwrap();
2287
2288 let result = decoder.read_record_batch();
2289
2290 match result {
2291 Err(ArrowError::IpcError(msg)) => {
2292 assert_eq!(msg, "Buffer count mismatched with metadata");
2293 }
2294 other => panic!("unexpected error: {other:?}"),
2295 }
2296 }
2297
2298 fn read_batch_with_variadic_count(count: i64, projection: Option<&[usize]>) -> ArrowError {
2301 use crate::r#gen::Message::*;
2302 use flatbuffers::FlatBufferBuilder;
2303
2304 let schema = Arc::new(Schema::new(vec![Field::new(
2305 "col",
2306 DataType::Utf8View,
2307 true,
2308 )]));
2309
2310 let mut fbb = FlatBufferBuilder::new();
2311 let nodes = fbb.create_vector(&[FieldNode::new(1, 0)]);
2312 let buffers = fbb.create_vector(&[crate::Buffer::new(0, 8), crate::Buffer::new(8, 8)]);
2313 let variadic_buffer_counts = fbb.create_vector(&[count]);
2314 let batch_offset = RecordBatch::create(
2315 &mut fbb,
2316 &RecordBatchArgs {
2317 length: 1,
2318 nodes: Some(nodes),
2319 buffers: Some(buffers),
2320 compression: None,
2321 variadicBufferCounts: Some(variadic_buffer_counts),
2322 },
2323 );
2324 fbb.finish_minimal(batch_offset);
2325 let batch_bytes = fbb.finished_data().to_vec();
2326 let batch = flatbuffers::root::<RecordBatch>(&batch_bytes).unwrap();
2327
2328 let data_buffer = Buffer::from(vec![0u8; 16]);
2329 let dictionaries: HashMap<i64, ArrayRef> = HashMap::new();
2330
2331 RecordBatchDecoder::try_new(
2332 &data_buffer,
2333 batch,
2334 schema,
2335 &dictionaries,
2336 &MetadataVersion::V5,
2337 )
2338 .unwrap()
2339 .with_projection(projection)
2340 .read_record_batch()
2341 .expect_err("should get error")
2342 }
2343
2344 #[test]
2347 fn test_invalid_variadic_buffer_count_error() {
2348 for count in [-2, -1, 1, i64::MAX] {
2351 for projection in [None, Some([].as_slice())] {
2352 let err = read_batch_with_variadic_count(count, projection);
2353 assert_eq!(
2354 err.to_string(),
2355 format!(
2356 "Ipc error: Invalid variadic count {count} for Utf8View column, \
2357 with 2 buffer(s) left in the message"
2358 ),
2359 "count {count}, projection {projection:?}"
2360 );
2361 }
2362 }
2363 }
2364
2365 #[test]
2367 fn test_valid_variadic_buffer_count_is_accepted() {
2368 let err = read_batch_with_variadic_count(0, None);
2369 assert!(!err.to_string().contains("Invalid variadic count"), "{err}");
2370 }
2371
2372 #[test]
2373 fn test_missing_footer_schema_error() {
2374 use crate::r#gen::File::{Footer, FooterArgs};
2375 use flatbuffers::FlatBufferBuilder;
2376
2377 let mut fbb = FlatBufferBuilder::new();
2380 let record_batches = fbb.create_vector::<Block>(&[]);
2381 let footer = Footer::create(
2382 &mut fbb,
2383 &FooterArgs {
2384 version: MetadataVersion::V5,
2385 schema: None,
2386 dictionaries: None,
2387 recordBatches: Some(record_batches),
2388 custom_metadata: None,
2389 },
2390 );
2391 fbb.finish(footer, None);
2392 let footer_data = fbb.finished_data();
2393
2394 let mut buf = Vec::new();
2396 buf.extend_from_slice(&crate::ARROW_MAGIC);
2397 buf.extend_from_slice(footer_data);
2398 buf.extend_from_slice(&(footer_data.len() as i32).to_le_bytes());
2399 buf.extend_from_slice(&crate::ARROW_MAGIC);
2400
2401 let err = FileReader::try_new(Cursor::new(buf), None)
2402 .expect_err("expected an error, not a panic");
2403 assert!(
2404 matches!(err, ArrowError::ParseError(_)),
2405 "expected ParseError, got {err:?}"
2406 );
2407 }
2408
2409 #[test]
2411 fn test_read_legacy_empty_list_without_offsets_buffer() {
2412 use crate::r#gen::Message::*;
2413 use flatbuffers::FlatBufferBuilder;
2414
2415 let schema = Arc::new(Schema::new(vec![Field::new_list(
2416 "items",
2417 Field::new_list_field(DataType::Int32, true),
2418 true,
2419 )]));
2420
2421 let mut fbb = FlatBufferBuilder::new();
2424 let nodes = fbb.create_vector(&[
2425 FieldNode::new(0, 0), FieldNode::new(0, 0), ]);
2428 let buffers = fbb.create_vector(&[
2429 crate::Buffer::new(0, 0), crate::Buffer::new(0, 0), crate::Buffer::new(0, 0), crate::Buffer::new(0, 0), ]);
2434 let batch_offset = RecordBatch::create(
2435 &mut fbb,
2436 &RecordBatchArgs {
2437 length: 0,
2438 nodes: Some(nodes),
2439 buffers: Some(buffers),
2440 compression: None,
2441 variadicBufferCounts: None,
2442 },
2443 );
2444 fbb.finish_minimal(batch_offset);
2445 let batch_bytes = fbb.finished_data().to_vec();
2446 let batch = flatbuffers::root::<RecordBatch>(&batch_bytes).unwrap();
2447
2448 let body = Buffer::from(Vec::<u8>::new());
2449 let dictionaries: HashMap<i64, ArrayRef> = HashMap::new();
2450 let metadata = MetadataVersion::V5;
2451
2452 let decoder =
2453 RecordBatchDecoder::try_new(&body, batch, schema.clone(), &dictionaries, &metadata)
2454 .unwrap();
2455
2456 let read_batch = decoder.read_record_batch().unwrap();
2457 assert_eq!(read_batch.num_rows(), 0);
2458
2459 let list = read_batch
2460 .column(0)
2461 .as_any()
2462 .downcast_ref::<ListArray>()
2463 .unwrap();
2464 assert_eq!(list.len(), 0);
2465 assert_eq!(list.values().len(), 0);
2466 }
2467
2468 #[test]
2470 fn test_read_legacy_empty_utf8_and_binary_without_offsets_buffer() {
2471 use crate::r#gen::Message::*;
2472 use flatbuffers::FlatBufferBuilder;
2473
2474 let schema = Arc::new(Schema::new(vec![
2475 Field::new("name", DataType::Utf8, true),
2476 Field::new("payload", DataType::Binary, true),
2477 ]));
2478
2479 let mut fbb = FlatBufferBuilder::new();
2482 let nodes = fbb.create_vector(&[
2483 FieldNode::new(0, 0), FieldNode::new(0, 0), ]);
2486 let buffers = fbb.create_vector(&[
2487 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), ]);
2494 let batch_offset = RecordBatch::create(
2495 &mut fbb,
2496 &RecordBatchArgs {
2497 length: 0,
2498 nodes: Some(nodes),
2499 buffers: Some(buffers),
2500 compression: None,
2501 variadicBufferCounts: None,
2502 },
2503 );
2504 fbb.finish_minimal(batch_offset);
2505 let batch_bytes = fbb.finished_data().to_vec();
2506 let batch = flatbuffers::root::<RecordBatch>(&batch_bytes).unwrap();
2507
2508 let body = Buffer::from(Vec::<u8>::new());
2509 let dictionaries: HashMap<i64, ArrayRef> = HashMap::new();
2510 let metadata = MetadataVersion::V5;
2511
2512 let decoder =
2513 RecordBatchDecoder::try_new(&body, batch, schema.clone(), &dictionaries, &metadata)
2514 .unwrap();
2515
2516 let read_batch = decoder.read_record_batch().unwrap();
2517 assert_eq!(read_batch.num_rows(), 0);
2518
2519 let utf8 = read_batch
2520 .column(0)
2521 .as_any()
2522 .downcast_ref::<StringArray>()
2523 .unwrap();
2524 assert_eq!(utf8.len(), 0);
2525 assert_eq!(utf8.value_offsets(), [0]);
2526
2527 let binary = read_batch
2528 .column(1)
2529 .as_any()
2530 .downcast_ref::<BinaryArray>()
2531 .unwrap();
2532 assert_eq!(binary.len(), 0);
2533 assert_eq!(binary.value_offsets(), [0]);
2534 }
2535
2536 #[test]
2537 fn test_projection_array_values() {
2538 let schema = create_test_projection_schema();
2540
2541 let batch = create_test_projection_batch_data(&schema);
2543
2544 let mut buf = Vec::new();
2546 {
2547 let mut writer = crate::writer::FileWriter::try_new(&mut buf, &schema).unwrap();
2548 writer.write(&batch).unwrap();
2549 writer.finish().unwrap();
2550 }
2551
2552 for index in 0..12 {
2554 let projection = vec![index];
2555 let reader = FileReader::try_new(std::io::Cursor::new(buf.clone()), Some(projection));
2556 let read_batch = reader.unwrap().next().unwrap().unwrap();
2557 let projected_column = read_batch.column(0);
2558 let expected_column = batch.column(index);
2559
2560 assert_eq!(projected_column.as_ref(), expected_column.as_ref());
2562 }
2563
2564 {
2565 let reader =
2567 FileReader::try_new(std::io::Cursor::new(buf.clone()), Some(vec![3, 2, 1]));
2568 let read_batch = reader.unwrap().next().unwrap().unwrap();
2569 let expected_batch = batch.project(&[3, 2, 1]).unwrap();
2570 assert_eq!(read_batch, expected_batch);
2571 }
2572 }
2573
2574 #[test]
2575 fn test_file_reader_projected_schema_matches_batch_schema() {
2576 let schema = create_test_projection_schema();
2577 let batch = create_test_projection_batch_data(&schema);
2578
2579 let mut buf = Vec::new();
2580 {
2581 let mut writer = crate::writer::FileWriter::try_new(&mut buf, &schema).unwrap();
2582 writer.write(&batch).unwrap();
2583 writer.finish().unwrap();
2584 }
2585
2586 let projection = vec![3, 2, 1];
2587 let mut reader = FileReader::try_new(Cursor::new(buf), Some(projection)).unwrap();
2588 let reader_schema = RecordBatchReader::schema(&reader);
2589 let read_batch = reader.next().unwrap().unwrap();
2590
2591 assert_eq!(reader_schema, read_batch.schema());
2592 }
2593
2594 #[test]
2595 fn test_stream_reader_projected_schema_matches_batch_schema() {
2596 let schema = create_test_projection_schema();
2597 let batch = create_test_projection_batch_data(&schema);
2598
2599 let mut buf = Vec::new();
2600 {
2601 let mut writer = crate::writer::StreamWriter::try_new(&mut buf, &schema).unwrap();
2602 writer.write(&batch).unwrap();
2603 writer.finish().unwrap();
2604 }
2605
2606 let projection = vec![3, 2, 1];
2607 let mut reader = StreamReader::try_new(Cursor::new(buf), Some(projection)).unwrap();
2608 let reader_schema = RecordBatchReader::schema(&reader);
2609 let read_batch = reader.next().unwrap().unwrap();
2610
2611 assert_eq!(reader_schema, read_batch.schema());
2612 }
2613
2614 #[test]
2615 fn test_file_reader_rejects_invalid_projection() {
2616 let schema = create_test_projection_schema();
2617 let batch = create_test_projection_batch_data(&schema);
2618
2619 let mut buf = Vec::new();
2620 {
2621 let mut writer = crate::writer::FileWriter::try_new(&mut buf, &schema).unwrap();
2622 writer.write(&batch).unwrap();
2623 writer.finish().unwrap();
2624 }
2625
2626 let result = FileReader::try_new(Cursor::new(buf), Some(vec![schema.fields().len()]));
2627
2628 assert!(matches!(result, Err(ArrowError::SchemaError(_))));
2629 }
2630
2631 #[test]
2632 fn test_projection_duplicate_indices() {
2633 let schema = create_test_projection_schema();
2634 let batch = create_test_projection_batch_data(&schema);
2635
2636 let mut buf = Vec::new();
2638 {
2639 let mut writer = crate::writer::FileWriter::try_new(&mut buf, &schema).unwrap();
2640 writer.write(&batch).unwrap();
2641 writer.finish().unwrap();
2642 }
2643
2644 for projection in [vec![1, 1], vec![2, 0, 2]] {
2646 let reader =
2647 FileReader::try_new(std::io::Cursor::new(buf.clone()), Some(projection.clone()));
2648 let read_batch = reader.unwrap().next().unwrap().unwrap();
2649
2650 let expected_batch = batch.project(&projection).unwrap();
2651 assert_eq!(read_batch, expected_batch);
2652 }
2653 }
2654
2655 #[test]
2656 fn test_arrow_single_float_row() {
2657 let schema = Schema::new(vec![
2658 Field::new("a", DataType::Float32, false),
2659 Field::new("b", DataType::Float32, false),
2660 Field::new("c", DataType::Int32, false),
2661 Field::new("d", DataType::Int32, false),
2662 ]);
2663 let arrays = vec![
2664 Arc::new(Float32Array::from(vec![1.23])) as ArrayRef,
2665 Arc::new(Float32Array::from(vec![-6.50])) as ArrayRef,
2666 Arc::new(Int32Array::from(vec![2])) as ArrayRef,
2667 Arc::new(Int32Array::from(vec![1])) as ArrayRef,
2668 ];
2669 let batch = RecordBatch::try_new(Arc::new(schema.clone()), arrays).unwrap();
2670 let mut file = tempfile::tempfile().unwrap();
2672 let mut stream_writer = crate::writer::StreamWriter::try_new(&mut file, &schema).unwrap();
2673 stream_writer.write(&batch).unwrap();
2674 stream_writer.finish().unwrap();
2675
2676 drop(stream_writer);
2677
2678 file.rewind().unwrap();
2679
2680 let reader = StreamReader::try_new(&mut file, None).unwrap();
2682
2683 reader.for_each(|batch| {
2684 let batch = batch.unwrap();
2685 assert!(
2686 batch
2687 .column(0)
2688 .as_any()
2689 .downcast_ref::<Float32Array>()
2690 .unwrap()
2691 .value(0)
2692 != 0.0
2693 );
2694 assert!(
2695 batch
2696 .column(1)
2697 .as_any()
2698 .downcast_ref::<Float32Array>()
2699 .unwrap()
2700 .value(0)
2701 != 0.0
2702 );
2703 });
2704
2705 file.rewind().unwrap();
2706
2707 let reader = StreamReader::try_new(file, Some(vec![0, 3])).unwrap();
2709
2710 reader.for_each(|batch| {
2711 let batch = batch.unwrap();
2712 assert_eq!(batch.schema().fields().len(), 2);
2713 assert_eq!(batch.schema().fields()[0].data_type(), &DataType::Float32);
2714 assert_eq!(batch.schema().fields()[1].data_type(), &DataType::Int32);
2715 });
2716 }
2717
2718 fn write_ipc(rb: &RecordBatch) -> Vec<u8> {
2720 let mut buf = Vec::new();
2721 let mut writer = crate::writer::FileWriter::try_new(&mut buf, rb.schema_ref()).unwrap();
2722 writer.write(rb).unwrap();
2723 writer.finish().unwrap();
2724 buf
2725 }
2726
2727 fn read_ipc(buf: &[u8]) -> Result<RecordBatch, ArrowError> {
2729 let mut reader = FileReader::try_new(std::io::Cursor::new(buf), None)?;
2730 reader.next().unwrap()
2731 }
2732
2733 fn read_ipc_skip_validation(buf: &[u8]) -> Result<RecordBatch, ArrowError> {
2736 let mut reader = unsafe {
2737 FileReader::try_new(std::io::Cursor::new(buf), None)?.with_skip_validation(true)
2738 };
2739 reader.next().unwrap()
2740 }
2741
2742 fn roundtrip_ipc(rb: &RecordBatch) -> RecordBatch {
2743 let buf = write_ipc(rb);
2744 read_ipc(&buf).unwrap()
2745 }
2746
2747 fn read_ipc_with_decoder(buf: Vec<u8>) -> Result<RecordBatch, ArrowError> {
2750 read_ipc_with_decoder_inner(buf, false)
2751 }
2752
2753 fn read_ipc_with_decoder_skip_validation(buf: Vec<u8>) -> Result<RecordBatch, ArrowError> {
2756 read_ipc_with_decoder_inner(buf, true)
2757 }
2758
2759 fn read_ipc_with_decoder_inner(
2760 buf: Vec<u8>,
2761 skip_validation: bool,
2762 ) -> Result<RecordBatch, ArrowError> {
2763 let buffer = Buffer::from_vec(buf);
2764 let trailer_start = buffer.len() - 10;
2765 let footer_len = read_footer_length(buffer[trailer_start..].try_into().unwrap())?;
2766 let footer = root_as_footer(&buffer[trailer_start - footer_len..trailer_start])
2767 .map_err(|e| ArrowError::InvalidArgumentError(format!("Invalid footer: {e}")))?;
2768
2769 let schema = try_fb_to_schema(footer.schema().unwrap()).unwrap();
2770
2771 let mut decoder = unsafe {
2772 FileDecoder::new(Arc::new(schema), footer.version())
2773 .with_skip_validation(skip_validation)
2774 };
2775 for block in footer.dictionaries().iter().flatten() {
2777 let block_len = block.bodyLength() as usize + block.metaDataLength() as usize;
2778 let data = buffer.slice_with_length(block.offset() as _, block_len);
2779 decoder.read_dictionary(block, &data)?
2780 }
2781
2782 let batches = footer.recordBatches().unwrap();
2784 assert_eq!(batches.len(), 1); let block = batches.get(0);
2787 let block_len = block.bodyLength() as usize + block.metaDataLength() as usize;
2788 let data = buffer.slice_with_length(block.offset() as _, block_len);
2789 Ok(decoder.read_record_batch(block, &data)?.unwrap())
2790 }
2791
2792 fn write_stream(rb: &RecordBatch) -> Vec<u8> {
2794 let mut buf = Vec::new();
2795 let mut writer = crate::writer::StreamWriter::try_new(&mut buf, rb.schema_ref()).unwrap();
2796 writer.write(rb).unwrap();
2797 writer.finish().unwrap();
2798 buf
2799 }
2800
2801 fn read_stream(buf: &[u8]) -> Result<RecordBatch, ArrowError> {
2803 let mut reader = StreamReader::try_new(std::io::Cursor::new(buf), None)?;
2804 reader.next().unwrap()
2805 }
2806
2807 fn read_stream_skip_validation(buf: &[u8]) -> Result<RecordBatch, ArrowError> {
2810 let mut reader = unsafe {
2811 StreamReader::try_new(std::io::Cursor::new(buf), None)?.with_skip_validation(true)
2812 };
2813 reader.next().unwrap()
2814 }
2815
2816 fn roundtrip_ipc_stream(rb: &RecordBatch) -> RecordBatch {
2817 let buf = write_stream(rb);
2818 read_stream(&buf).unwrap()
2819 }
2820
2821 #[test]
2822 fn test_roundtrip_with_custom_metadata() {
2823 let schema = Schema::new(vec![Field::new("dummy", DataType::Float64, false)]);
2824 let mut buf = Vec::new();
2825 let mut writer = crate::writer::FileWriter::try_new(&mut buf, &schema).unwrap();
2826 let mut test_metadata = HashMap::new();
2827 test_metadata.insert("abc".to_string(), "abc".to_string());
2828 test_metadata.insert("def".to_string(), "def".to_string());
2829 for (k, v) in &test_metadata {
2830 writer.write_metadata(k, v);
2831 }
2832 writer.finish().unwrap();
2833 drop(writer);
2834
2835 let reader = crate::reader::FileReader::try_new(std::io::Cursor::new(buf), None).unwrap();
2836 assert_eq!(reader.custom_metadata(), &test_metadata);
2837 }
2838
2839 #[test]
2840 fn test_roundtrip_nested_dict() {
2841 let inner: DictionaryArray<Int32Type> = vec!["a", "b", "a"].into_iter().collect();
2842
2843 let array = Arc::new(inner) as ArrayRef;
2844
2845 let dctfield = Arc::new(Field::new("dict", array.data_type().clone(), false));
2846
2847 let s = StructArray::from(vec![(dctfield, array)]);
2848 let struct_array = Arc::new(s) as ArrayRef;
2849
2850 let schema = Arc::new(Schema::new(vec![Field::new(
2851 "struct",
2852 struct_array.data_type().clone(),
2853 false,
2854 )]));
2855
2856 let batch = RecordBatch::try_new(schema, vec![struct_array]).unwrap();
2857
2858 assert_eq!(batch, roundtrip_ipc(&batch));
2859 }
2860
2861 #[test]
2862 fn test_roundtrip_nested_dict_no_preserve_dict_id() {
2863 let inner: DictionaryArray<Int32Type> = vec!["a", "b", "a"].into_iter().collect();
2864
2865 let array = Arc::new(inner) as ArrayRef;
2866
2867 let dctfield = Arc::new(Field::new("dict", array.data_type().clone(), false));
2868
2869 let s = StructArray::from(vec![(dctfield, array)]);
2870 let struct_array = Arc::new(s) as ArrayRef;
2871
2872 let schema = Arc::new(Schema::new(vec![Field::new(
2873 "struct",
2874 struct_array.data_type().clone(),
2875 false,
2876 )]));
2877
2878 let batch = RecordBatch::try_new(schema, vec![struct_array]).unwrap();
2879
2880 let mut buf = Vec::new();
2881 let mut writer = crate::writer::FileWriter::try_new_with_options(
2882 &mut buf,
2883 batch.schema_ref(),
2884 IpcWriteOptions::default(),
2885 )
2886 .unwrap();
2887 writer.write(&batch).unwrap();
2888 writer.finish().unwrap();
2889 drop(writer);
2890
2891 let mut reader = FileReader::try_new(std::io::Cursor::new(buf), None).unwrap();
2892
2893 assert_eq!(batch, reader.next().unwrap().unwrap());
2894 }
2895
2896 fn check_union_with_builder(mut builder: UnionBuilder) {
2897 builder.append::<Int32Type>("a", 1).unwrap();
2898 builder.append_null::<Int32Type>("a").unwrap();
2899 builder.append::<Float64Type>("c", 3.0).unwrap();
2900 builder.append::<Int32Type>("a", 4).unwrap();
2901 builder.append::<Int64Type>("d", 11).unwrap();
2902 let union = builder.build().unwrap();
2903
2904 let schema = Arc::new(Schema::new(vec![Field::new(
2905 "union",
2906 union.data_type().clone(),
2907 false,
2908 )]));
2909
2910 let union_array = Arc::new(union) as ArrayRef;
2911
2912 let rb = RecordBatch::try_new(schema, vec![union_array]).unwrap();
2913 let rb2 = roundtrip_ipc(&rb);
2914 assert_eq!(rb.schema(), rb2.schema());
2917 assert_eq!(rb.num_columns(), rb2.num_columns());
2918 assert_eq!(rb.num_rows(), rb2.num_rows());
2919 let union1 = rb.column(0);
2920 let union2 = rb2.column(0);
2921
2922 assert_eq!(union1, union2);
2923 }
2924
2925 #[test]
2926 fn test_roundtrip_dense_union() {
2927 check_union_with_builder(UnionBuilder::new_dense());
2928 }
2929
2930 #[test]
2931 fn test_roundtrip_sparse_union() {
2932 check_union_with_builder(UnionBuilder::new_sparse());
2933 }
2934
2935 #[test]
2936 fn test_roundtrip_struct_empty_fields() {
2937 let nulls = NullBuffer::from(&[true, true, false]);
2938 let rb = RecordBatch::try_from_iter([(
2939 "",
2940 Arc::new(StructArray::new_empty_fields(nulls.len(), Some(nulls))) as _,
2941 )])
2942 .unwrap();
2943 let rb2 = roundtrip_ipc(&rb);
2944 assert_eq!(rb, rb2);
2945 }
2946
2947 #[test]
2948 fn test_roundtrip_stream_run_array_sliced() {
2949 let run_array_1: Int32RunArray = vec!["a", "a", "a", "b", "b", "c", "c", "c"]
2950 .into_iter()
2951 .collect();
2952 let run_array_1_sliced = run_array_1.slice(2, 5);
2953
2954 let run_array_2_inupt = vec![Some(1_i32), None, None, Some(2), Some(2)];
2955 let mut run_array_2_builder = PrimitiveRunBuilder::<Int16Type, Int32Type>::new();
2956 run_array_2_builder.extend(run_array_2_inupt);
2957 let run_array_2 = run_array_2_builder.finish();
2958
2959 let schema = Arc::new(Schema::new(vec![
2960 Field::new(
2961 "run_array_1_sliced",
2962 run_array_1_sliced.data_type().clone(),
2963 false,
2964 ),
2965 Field::new("run_array_2", run_array_2.data_type().clone(), false),
2966 ]));
2967 let input_batch = RecordBatch::try_new(
2968 schema,
2969 vec![Arc::new(run_array_1_sliced.clone()), Arc::new(run_array_2)],
2970 )
2971 .unwrap();
2972 let output_batch = roundtrip_ipc_stream(&input_batch);
2973
2974 assert_eq!(input_batch.column(1), output_batch.column(1));
2978
2979 let run_array_1_unsliced = unslice_run_array(run_array_1_sliced.into_data()).unwrap();
2980 assert_eq!(run_array_1_unsliced, output_batch.column(0).into_data());
2981 }
2982
2983 #[test]
2984 fn test_roundtrip_stream_nested_dict() {
2985 let xs = vec!["AA", "BB", "AA", "CC", "BB"];
2986 let dict = Arc::new(
2987 xs.clone()
2988 .into_iter()
2989 .collect::<DictionaryArray<Int8Type>>(),
2990 );
2991 let string_array: ArrayRef = Arc::new(StringArray::from(xs.clone()));
2992 let struct_array = StructArray::from(vec![
2993 (
2994 Arc::new(Field::new("f2.1", DataType::Utf8, false)),
2995 string_array,
2996 ),
2997 (
2998 Arc::new(Field::new("f2.2_struct", dict.data_type().clone(), false)),
2999 dict.clone() as ArrayRef,
3000 ),
3001 ]);
3002 let schema = Arc::new(Schema::new(vec![
3003 Field::new("f1_string", DataType::Utf8, false),
3004 Field::new("f2_struct", struct_array.data_type().clone(), false),
3005 ]));
3006 let input_batch = RecordBatch::try_new(
3007 schema,
3008 vec![
3009 Arc::new(StringArray::from(xs.clone())),
3010 Arc::new(struct_array),
3011 ],
3012 )
3013 .unwrap();
3014 let output_batch = roundtrip_ipc_stream(&input_batch);
3015 assert_eq!(input_batch, output_batch);
3016 }
3017
3018 #[test]
3019 fn test_ipc_writers_reject_dictionary_of_dictionary_schema() {
3020 let values = Arc::new(StringArray::from(vec![Some("a"), Some("b")])) as ArrayRef;
3021 let inner = Arc::new(DictionaryArray::new(
3022 UInt32Array::from_iter_values([0, 1]),
3023 values,
3024 )) as ArrayRef;
3025 let outer = Arc::new(DictionaryArray::new(
3026 UInt32Array::from_iter_values([0, 1, 0]),
3027 inner,
3028 )) as ArrayRef;
3029
3030 let schema = Arc::new(Schema::new(vec![Field::new(
3031 "f1",
3032 outer.data_type().clone(),
3033 false,
3034 )]));
3035 let batch = RecordBatch::try_new(schema, vec![outer]).unwrap();
3036
3037 let mut stream = Vec::new();
3038 let Err(err) = crate::writer::StreamWriter::try_new(&mut stream, batch.schema_ref()) else {
3039 panic!("IPC stream writer should reject dictionary-of-dictionary schemas");
3040 };
3041 assert!(stream.is_empty());
3042
3043 assert!(
3044 err.to_string().contains("dictionary-of-dictionary values"),
3045 "unexpected error: {err}"
3046 );
3047
3048 let mut file = Vec::new();
3049 let Err(err) = crate::writer::FileWriter::try_new(&mut file, batch.schema_ref()) else {
3050 panic!("IPC file writer should reject dictionary-of-dictionary schemas");
3051 };
3052 assert!(file.is_empty());
3053
3054 assert!(
3055 err.to_string().contains("dictionary-of-dictionary values"),
3056 "unexpected error: {err}"
3057 );
3058 }
3059
3060 #[test]
3061 fn test_roundtrip_stream_nested_dict_of_map_of_dict() {
3062 let values = StringArray::from(vec![Some("a"), None, Some("b"), Some("c")]);
3063 let values = Arc::new(values) as ArrayRef;
3064 let value_dict_keys = Int8Array::from_iter_values([0, 1, 1, 2, 3, 1]);
3065 let value_dict_array = DictionaryArray::new(value_dict_keys, values.clone());
3066
3067 let key_dict_keys = Int8Array::from_iter_values([0, 0, 2, 2, 2, 3]);
3068 let key_dict_array = DictionaryArray::new(key_dict_keys, values);
3069
3070 #[expect(deprecated)]
3071 let keys_field = Arc::new(Field::new_dict(
3072 Field::MAP_KEY_FIELD_DEFAULT_NAME,
3073 DataType::Dictionary(Box::new(DataType::Int8), Box::new(DataType::Utf8)),
3074 false,
3075 1,
3076 false,
3077 ));
3078 #[expect(deprecated)]
3079 let values_field = Arc::new(Field::new_dict(
3080 Field::MAP_VALUE_FIELD_DEFAULT_NAME,
3081 DataType::Dictionary(Box::new(DataType::Int8), Box::new(DataType::Utf8)),
3082 true,
3083 2,
3084 false,
3085 ));
3086 let entry_struct = StructArray::from(vec![
3087 (keys_field, make_array(key_dict_array.into_data())),
3088 (values_field, make_array(value_dict_array.into_data())),
3089 ]);
3090 let map_data_type = DataType::Map(
3091 Arc::new(Field::new(
3092 Field::MAP_ENTRIES_FIELD_DEFAULT_NAME,
3093 entry_struct.data_type().clone(),
3094 false,
3095 )),
3096 false,
3097 );
3098
3099 let entry_offsets = Buffer::from_slice_ref([0, 2, 4, 6]);
3100 let map_data = ArrayData::builder(map_data_type)
3101 .len(3)
3102 .add_buffer(entry_offsets)
3103 .add_child_data(entry_struct.into_data())
3104 .build()
3105 .unwrap();
3106 let map_array = MapArray::from(map_data);
3107
3108 let dict_keys = Int8Array::from_iter_values([0, 1, 1, 2, 2, 1]);
3109 let dict_dict_array = DictionaryArray::new(dict_keys, Arc::new(map_array));
3110
3111 let schema = Arc::new(Schema::new(vec![Field::new(
3112 "f1",
3113 dict_dict_array.data_type().clone(),
3114 false,
3115 )]));
3116 let input_batch = RecordBatch::try_new(schema, vec![Arc::new(dict_dict_array)]).unwrap();
3117 let output_batch = roundtrip_ipc_stream(&input_batch);
3118 assert_eq!(input_batch, output_batch);
3119 }
3120
3121 fn test_roundtrip_stream_dict_of_list_of_dict_impl<
3122 OffsetSize: OffsetSizeTrait,
3123 U: ArrowNativeType,
3124 >(
3125 list_data_type: DataType,
3126 offsets: &[U; 5],
3127 ) {
3128 let values = StringArray::from(vec![Some("a"), None, Some("c"), None]);
3129 let keys = Int8Array::from_iter_values([0, 0, 1, 2, 0, 1, 3]);
3130 let dict_array = DictionaryArray::new(keys, Arc::new(values));
3131 let dict_data = dict_array.to_data();
3132
3133 let value_offsets = Buffer::from_slice_ref(offsets);
3134
3135 let list_data = ArrayData::builder(list_data_type)
3136 .len(4)
3137 .add_buffer(value_offsets)
3138 .add_child_data(dict_data)
3139 .build()
3140 .unwrap();
3141 let list_array = GenericListArray::<OffsetSize>::from(list_data);
3142
3143 let keys_for_dict = Int8Array::from_iter_values([0, 3, 0, 1, 1, 2, 0, 1, 3]);
3144 let dict_dict_array = DictionaryArray::new(keys_for_dict, Arc::new(list_array));
3145
3146 let schema = Arc::new(Schema::new(vec![Field::new(
3147 "f1",
3148 dict_dict_array.data_type().clone(),
3149 false,
3150 )]));
3151 let input_batch = RecordBatch::try_new(schema, vec![Arc::new(dict_dict_array)]).unwrap();
3152 let output_batch = roundtrip_ipc_stream(&input_batch);
3153 assert_eq!(input_batch, output_batch);
3154 }
3155
3156 #[test]
3157 fn test_roundtrip_stream_dict_of_list_of_dict() {
3158 #[expect(deprecated)]
3160 let list_data_type = DataType::List(Arc::new(Field::new_dict(
3161 "item",
3162 DataType::Dictionary(Box::new(DataType::Int8), Box::new(DataType::Utf8)),
3163 true,
3164 1,
3165 false,
3166 )));
3167 let offsets: &[i32; 5] = &[0, 2, 4, 4, 6];
3168 test_roundtrip_stream_dict_of_list_of_dict_impl::<i32, i32>(list_data_type, offsets);
3169
3170 #[expect(deprecated)]
3172 let list_data_type = DataType::LargeList(Arc::new(Field::new_dict(
3173 "item",
3174 DataType::Dictionary(Box::new(DataType::Int8), Box::new(DataType::Utf8)),
3175 true,
3176 1,
3177 false,
3178 )));
3179 let offsets: &[i64; 5] = &[0, 2, 4, 4, 7];
3180 test_roundtrip_stream_dict_of_list_of_dict_impl::<i64, i64>(list_data_type, offsets);
3181 }
3182
3183 #[test]
3184 fn test_roundtrip_stream_dict_of_fixed_size_list_of_dict() {
3185 let values = StringArray::from(vec![Some("a"), None, Some("c"), None]);
3186 let keys = Int8Array::from_iter_values([0, 0, 1, 2, 0, 1, 3, 1, 2]);
3187 let dict_array = DictionaryArray::new(keys, Arc::new(values));
3188 let dict_data = dict_array.into_data();
3189
3190 #[expect(deprecated)]
3191 let list_data_type = DataType::FixedSizeList(
3192 Arc::new(Field::new_dict(
3193 "item",
3194 DataType::Dictionary(Box::new(DataType::Int8), Box::new(DataType::Utf8)),
3195 true,
3196 1,
3197 false,
3198 )),
3199 3,
3200 );
3201 let list_data = ArrayData::builder(list_data_type)
3202 .len(3)
3203 .add_child_data(dict_data)
3204 .build()
3205 .unwrap();
3206 let list_array = FixedSizeListArray::from(list_data);
3207
3208 let keys_for_dict = Int8Array::from_iter_values([0, 1, 0, 1, 1, 2, 0, 1, 2]);
3209 let dict_dict_array = DictionaryArray::new(keys_for_dict, Arc::new(list_array));
3210
3211 let schema = Arc::new(Schema::new(vec![Field::new(
3212 "f1",
3213 dict_dict_array.data_type().clone(),
3214 false,
3215 )]));
3216 let input_batch = RecordBatch::try_new(schema, vec![Arc::new(dict_dict_array)]).unwrap();
3217 let output_batch = roundtrip_ipc_stream(&input_batch);
3218 assert_eq!(input_batch, output_batch);
3219 }
3220
3221 const LONG_TEST_STRING: &str =
3222 "This is a long string to make sure binary view array handles it";
3223
3224 #[test]
3225 fn test_roundtrip_view_types() {
3226 let schema = Schema::new(vec![
3227 Field::new("field_1", DataType::BinaryView, true),
3228 Field::new("field_2", DataType::Utf8, true),
3229 Field::new("field_3", DataType::Utf8View, true),
3230 ]);
3231 let bin_values: Vec<Option<&[u8]>> = vec![
3232 Some(b"foo"),
3233 None,
3234 Some(b"bar"),
3235 Some(LONG_TEST_STRING.as_bytes()),
3236 ];
3237 let utf8_values: Vec<Option<&str>> =
3238 vec![Some("foo"), None, Some("bar"), Some(LONG_TEST_STRING)];
3239 let bin_view_array = BinaryViewArray::from_iter(bin_values);
3240 let utf8_array = StringArray::from_iter(utf8_values.iter());
3241 let utf8_view_array = StringViewArray::from_iter(utf8_values);
3242 let record_batch = RecordBatch::try_new(
3243 Arc::new(schema.clone()),
3244 vec![
3245 Arc::new(bin_view_array),
3246 Arc::new(utf8_array),
3247 Arc::new(utf8_view_array),
3248 ],
3249 )
3250 .unwrap();
3251
3252 assert_eq!(record_batch, roundtrip_ipc(&record_batch));
3253 assert_eq!(record_batch, roundtrip_ipc_stream(&record_batch));
3254
3255 let sliced_batch = record_batch.slice(1, 2);
3256 assert_eq!(sliced_batch, roundtrip_ipc(&sliced_batch));
3257 assert_eq!(sliced_batch, roundtrip_ipc_stream(&sliced_batch));
3258 }
3259
3260 #[test]
3261 fn test_roundtrip_view_types_nested_dict() {
3262 let bin_values: Vec<Option<&[u8]>> = vec![
3263 Some(b"foo"),
3264 None,
3265 Some(b"bar"),
3266 Some(LONG_TEST_STRING.as_bytes()),
3267 Some(b"field"),
3268 ];
3269 let utf8_values: Vec<Option<&str>> = vec![
3270 Some("foo"),
3271 None,
3272 Some("bar"),
3273 Some(LONG_TEST_STRING),
3274 Some("field"),
3275 ];
3276 let bin_view_array = Arc::new(BinaryViewArray::from_iter(bin_values));
3277 let utf8_view_array = Arc::new(StringViewArray::from_iter(utf8_values));
3278
3279 let key_dict_keys = Int8Array::from_iter_values([0, 0, 2, 2, 0, 2, 3]);
3280 let key_dict_array = DictionaryArray::new(key_dict_keys, utf8_view_array.clone());
3281 #[expect(deprecated)]
3282 let keys_field = Arc::new(Field::new_dict(
3283 Field::MAP_KEY_FIELD_DEFAULT_NAME,
3284 DataType::Dictionary(Box::new(DataType::Int8), Box::new(DataType::Utf8View)),
3285 false,
3286 1,
3287 false,
3288 ));
3289
3290 let value_dict_keys = Int8Array::from_iter_values([0, 3, 0, 1, 2, 0, 1]);
3291 let value_dict_array = DictionaryArray::new(value_dict_keys, bin_view_array);
3292 #[expect(deprecated)]
3293 let values_field = Arc::new(Field::new_dict(
3294 Field::MAP_VALUE_FIELD_DEFAULT_NAME,
3295 DataType::Dictionary(Box::new(DataType::Int8), Box::new(DataType::BinaryView)),
3296 true,
3297 2,
3298 false,
3299 ));
3300 let entry_struct = StructArray::from(vec![
3301 (keys_field, make_array(key_dict_array.into_data())),
3302 (values_field, make_array(value_dict_array.into_data())),
3303 ]);
3304
3305 let map_data_type = DataType::Map(
3306 Arc::new(Field::new(
3307 Field::MAP_ENTRIES_FIELD_DEFAULT_NAME,
3308 entry_struct.data_type().clone(),
3309 false,
3310 )),
3311 false,
3312 );
3313 let entry_offsets = Buffer::from_slice_ref([0, 2, 4, 7]);
3314 let map_data = ArrayData::builder(map_data_type)
3315 .len(3)
3316 .add_buffer(entry_offsets)
3317 .add_child_data(entry_struct.into_data())
3318 .build()
3319 .unwrap();
3320 let map_array = MapArray::from(map_data);
3321
3322 let dict_keys = Int8Array::from_iter_values([0, 1, 0, 1, 1, 2, 0, 1, 2]);
3323 let dict_dict_array = DictionaryArray::new(dict_keys, Arc::new(map_array));
3324 let schema = Arc::new(Schema::new(vec![Field::new(
3325 "f1",
3326 dict_dict_array.data_type().clone(),
3327 false,
3328 )]));
3329 let batch = RecordBatch::try_new(schema, vec![Arc::new(dict_dict_array)]).unwrap();
3330 assert_eq!(batch, roundtrip_ipc(&batch));
3331 assert_eq!(batch, roundtrip_ipc_stream(&batch));
3332
3333 let sliced_batch = batch.slice(1, 2);
3334 assert_eq!(sliced_batch, roundtrip_ipc(&sliced_batch));
3335 assert_eq!(sliced_batch, roundtrip_ipc_stream(&sliced_batch));
3336 }
3337
3338 #[test]
3339 fn test_no_columns_batch() {
3340 let schema = Arc::new(Schema::empty());
3341 let options = RecordBatchOptions::new()
3342 .with_match_field_names(true)
3343 .with_row_count(Some(10));
3344 let input_batch = RecordBatch::try_new_with_options(schema, vec![], &options).unwrap();
3345 let output_batch = roundtrip_ipc_stream(&input_batch);
3346 assert_eq!(input_batch, output_batch);
3347 }
3348
3349 #[test]
3350 fn test_unaligned() {
3351 let batch = RecordBatch::try_from_iter(vec![(
3352 "i32",
3353 Arc::new(Int32Array::from(vec![1, 2, 3, 4])) as _,
3354 )])
3355 .unwrap();
3356
3357 let r#gen = IpcDataGenerator {};
3358 let mut dict_tracker = DictionaryTracker::new(false);
3359 let (_, encoded) = r#gen
3360 .encode(
3361 &batch,
3362 &mut dict_tracker,
3363 &Default::default(),
3364 &mut Default::default(),
3365 )
3366 .unwrap();
3367
3368 let message = root_as_message(&encoded.ipc_message).unwrap();
3369
3370 let mut buffer = MutableBuffer::with_capacity(encoded.arrow_data.len() + 1);
3372 buffer.push(0_u8);
3373 buffer.extend_from_slice(&encoded.arrow_data);
3374 let b = Buffer::from(buffer).slice(1);
3375 assert_ne!(b.as_ptr().align_offset(8), 0);
3376
3377 let ipc_batch = message.header_as_record_batch().unwrap();
3378 let roundtrip = RecordBatchDecoder::try_new(
3379 &b,
3380 ipc_batch,
3381 batch.schema(),
3382 &Default::default(),
3383 &message.version(),
3384 )
3385 .unwrap()
3386 .with_require_alignment(false)
3387 .read_record_batch()
3388 .unwrap();
3389 assert_eq!(batch, roundtrip);
3390 }
3391
3392 #[test]
3393 fn test_unaligned_throws_error_with_require_alignment() {
3394 let batch = RecordBatch::try_from_iter(vec![(
3395 "i32",
3396 Arc::new(Int32Array::from(vec![1, 2, 3, 4])) as _,
3397 )])
3398 .unwrap();
3399
3400 let r#gen = IpcDataGenerator {};
3401 let mut dict_tracker = DictionaryTracker::new(false);
3402 let (_, encoded) = r#gen
3403 .encode(
3404 &batch,
3405 &mut dict_tracker,
3406 &Default::default(),
3407 &mut Default::default(),
3408 )
3409 .unwrap();
3410
3411 let message = root_as_message(&encoded.ipc_message).unwrap();
3412
3413 let mut buffer = MutableBuffer::with_capacity(encoded.arrow_data.len() + 1);
3415 buffer.push(0_u8);
3416 buffer.extend_from_slice(&encoded.arrow_data);
3417 let b = Buffer::from(buffer).slice(1);
3418 assert_ne!(b.as_ptr().align_offset(8), 0);
3419
3420 let ipc_batch = message.header_as_record_batch().unwrap();
3421 let result = RecordBatchDecoder::try_new(
3422 &b,
3423 ipc_batch,
3424 batch.schema(),
3425 &Default::default(),
3426 &message.version(),
3427 )
3428 .unwrap()
3429 .with_require_alignment(true)
3430 .read_record_batch();
3431
3432 let error = result.unwrap_err();
3433 assert_eq!(
3434 error.to_string(),
3435 "Invalid argument error: Misaligned buffers[0] in array of type Int32, \
3436 offset from expected alignment of 4 by 1"
3437 );
3438 }
3439
3440 #[test]
3449 fn test_misaligned_buffers_error() {
3450 const OFFSET_RANGE: usize = 128;
3451
3452 let cases: Vec<(ArrayRef, usize)> = vec![
3456 (Arc::new(Int32Array::from_iter(0i32..100)) as _, 4),
3457 (Arc::new(Int64Array::from_iter(0i64..100)) as _, 8),
3458 (
3459 Arc::new(StringArray::from_iter_values(
3460 (0..100).map(|i| i.to_string()),
3461 )) as _,
3462 4,
3463 ),
3464 (
3465 Arc::new(LargeStringArray::from_iter_values(
3466 (0..100).map(|i| i.to_string()),
3467 )) as _,
3468 8,
3469 ),
3470 ];
3471
3472 for (array, alignment) in cases {
3473 let batch = RecordBatch::try_from_iter(vec![("col", Arc::clone(&array))]).unwrap();
3474 let encoder = IpcDataGenerator {};
3475 let mut dict_tracker = DictionaryTracker::new(false);
3476 let (_, encoded) = encoder
3477 .encode(
3478 &batch,
3479 &mut dict_tracker,
3480 &Default::default(),
3481 &mut Default::default(),
3482 )
3483 .unwrap();
3484 let message = root_as_message(&encoded.ipc_message).unwrap();
3485 let ipc_batch = message.header_as_record_batch().unwrap();
3486
3487 for offset in 0..OFFSET_RANGE {
3488 let mut storage = MutableBuffer::with_capacity(encoded.arrow_data.len() + offset);
3494 for _ in 0..offset {
3495 storage.push(0_u8);
3496 }
3497 storage.extend_from_slice(&encoded.arrow_data);
3498 let buf = Buffer::from(storage).slice(offset);
3499
3500 let result = RecordBatchDecoder::try_new(
3501 &buf,
3502 ipc_batch,
3503 batch.schema(),
3504 &Default::default(),
3505 &message.version(),
3506 )
3507 .unwrap()
3508 .with_require_alignment(true)
3509 .read_record_batch();
3510
3511 if offset % alignment == 0 {
3512 assert!(
3513 result.is_ok(),
3514 "type={} offset={offset}: expected Ok but got {:?}",
3515 array.data_type(),
3516 result.unwrap_err(),
3517 );
3518 } else {
3519 let err = result
3520 .expect_err(&format!(
3521 "type={} offset={offset}: expected Err for misaligned buffer",
3522 array.data_type()
3523 ))
3524 .to_string();
3525 assert!(
3526 err.contains("Misaligned buffers"),
3527 "type={} offset={offset}: unexpected error: {err}",
3528 array.data_type(),
3529 );
3530 }
3531 }
3532 }
3533 }
3534
3535 #[test]
3536 #[cfg_attr(miri, ignore)] fn test_file_with_massive_column_count() {
3538 let limit = 600_000;
3540
3541 let fields = (0..limit)
3542 .map(|i| Field::new(format!("{i}"), DataType::Boolean, false))
3543 .collect::<Vec<_>>();
3544 let schema = Arc::new(Schema::new(fields));
3545 let batch = RecordBatch::new_empty(schema);
3546
3547 let mut buf = Vec::new();
3548 let mut writer = crate::writer::FileWriter::try_new(&mut buf, batch.schema_ref()).unwrap();
3549 writer.write(&batch).unwrap();
3550 writer.finish().unwrap();
3551 drop(writer);
3552
3553 let mut reader = FileReaderBuilder::new()
3554 .with_max_footer_fb_tables(1_500_000)
3555 .build(std::io::Cursor::new(buf))
3556 .unwrap();
3557 let roundtrip_batch = reader.next().unwrap().unwrap();
3558
3559 assert_eq!(batch, roundtrip_batch);
3560 }
3561
3562 #[test]
3563 fn test_file_with_deeply_nested_columns() {
3564 let limit = 61;
3566
3567 let fields = (0..limit).fold(
3568 vec![Field::new("leaf", DataType::Boolean, false)],
3569 |field, index| vec![Field::new_struct(format!("{index}"), field, false)],
3570 );
3571 let schema = Arc::new(Schema::new(fields));
3572 let batch = RecordBatch::new_empty(schema);
3573
3574 let mut buf = Vec::new();
3575 let mut writer = crate::writer::FileWriter::try_new(&mut buf, batch.schema_ref()).unwrap();
3576 writer.write(&batch).unwrap();
3577 writer.finish().unwrap();
3578 drop(writer);
3579
3580 let mut reader = FileReaderBuilder::new()
3581 .with_max_footer_fb_depth(65)
3582 .build(std::io::Cursor::new(buf))
3583 .unwrap();
3584 let roundtrip_batch = reader.next().unwrap().unwrap();
3585
3586 assert_eq!(batch, roundtrip_batch);
3587 }
3588
3589 #[test]
3590 fn test_invalid_struct_array_ipc_read_errors() {
3591 let a_field = Field::new("a", DataType::Int32, false);
3592 let b_field = Field::new("b", DataType::Int32, false);
3593 let struct_fields = Fields::from(vec![a_field.clone(), b_field.clone()]);
3594
3595 let a_array_data = ArrayData::builder(a_field.data_type().clone())
3596 .len(4)
3597 .add_buffer(Buffer::from_slice_ref([1, 2, 3, 4]))
3598 .build()
3599 .unwrap();
3600 let b_array_data = ArrayData::builder(b_field.data_type().clone())
3601 .len(3)
3602 .add_buffer(Buffer::from_slice_ref([5, 6, 7]))
3603 .build()
3604 .unwrap();
3605
3606 let invalid_struct_arr = unsafe {
3607 StructArray::new_unchecked(
3608 struct_fields,
3609 vec![make_array(a_array_data), make_array(b_array_data)],
3610 None,
3611 )
3612 };
3613
3614 expect_ipc_validation_error(
3615 Arc::new(invalid_struct_arr),
3616 "Invalid argument error: Incorrect array length for StructArray field \"b\", expected 4 got 3",
3617 );
3618 }
3619
3620 #[test]
3621 fn test_invalid_nested_array_ipc_read_errors() {
3622 let a_field = Field::new("a", DataType::Int32, false);
3624 let b_field = Field::new("b", DataType::Utf8, false);
3625
3626 let schema = Arc::new(Schema::new(vec![Field::new_struct(
3627 "s",
3628 vec![a_field.clone(), b_field.clone()],
3629 false,
3630 )]));
3631
3632 let a_array_data = ArrayData::builder(a_field.data_type().clone())
3633 .len(4)
3634 .add_buffer(Buffer::from_slice_ref([1, 2, 3, 4]))
3635 .build()
3636 .unwrap();
3637 let b_array_data = {
3639 let valid: &[u8] = b" ";
3640 let mut invalid = vec![];
3641 invalid.extend_from_slice(b"ValidString");
3642 invalid.extend_from_slice(INVALID_UTF8_FIRST_CHAR);
3643 let binary_array =
3644 BinaryArray::from_iter(vec![None, Some(valid), None, Some(&invalid)]);
3645 let array = unsafe {
3646 StringArray::new_unchecked(
3647 binary_array.offsets().clone(),
3648 binary_array.values().clone(),
3649 binary_array.nulls().cloned(),
3650 )
3651 };
3652 array.into_data()
3653 };
3654 let struct_data_type = schema.field(0).data_type();
3655
3656 let invalid_struct_arr = unsafe {
3657 make_array(
3658 ArrayData::builder(struct_data_type.clone())
3659 .len(4)
3660 .add_child_data(a_array_data)
3661 .add_child_data(b_array_data)
3662 .build_unchecked(),
3663 )
3664 };
3665 expect_ipc_validation_error(
3666 invalid_struct_arr,
3667 "Invalid argument error: Invalid UTF8 sequence at string index 3 (3..18): invalid utf-8 sequence of 1 bytes from index 11",
3668 );
3669 }
3670
3671 #[test]
3672 fn test_same_dict_id_without_preserve() {
3673 let batch = RecordBatch::try_new(
3674 Arc::new(Schema::new(
3675 ["a", "b"]
3676 .iter()
3677 .map(|name| {
3678 #[expect(deprecated)]
3679 Field::new_dict(
3680 name.to_string(),
3681 DataType::Dictionary(
3682 Box::new(DataType::Int32),
3683 Box::new(DataType::Utf8),
3684 ),
3685 true,
3686 0,
3687 false,
3688 )
3689 })
3690 .collect::<Vec<Field>>(),
3691 )),
3692 vec![
3693 Arc::new(
3694 vec![Some("c"), Some("d")]
3695 .into_iter()
3696 .collect::<DictionaryArray<Int32Type>>(),
3697 ) as ArrayRef,
3698 Arc::new(
3699 vec![Some("e"), Some("f")]
3700 .into_iter()
3701 .collect::<DictionaryArray<Int32Type>>(),
3702 ) as ArrayRef,
3703 ],
3704 )
3705 .expect("Failed to create RecordBatch");
3706
3707 let mut buf = vec![];
3709 {
3710 let mut writer = crate::writer::StreamWriter::try_new_with_options(
3711 &mut buf,
3712 batch.schema().as_ref(),
3713 crate::writer::IpcWriteOptions::default(),
3714 )
3715 .expect("Failed to create StreamWriter");
3716 writer.write(&batch).expect("Failed to write RecordBatch");
3717 writer.finish().expect("Failed to finish StreamWriter");
3718 }
3719
3720 StreamReader::try_new(std::io::Cursor::new(buf), None)
3721 .expect("Failed to create StreamReader")
3722 .for_each(|decoded_batch| {
3723 assert_eq!(decoded_batch.expect("Failed to read RecordBatch"), batch);
3724 });
3725 }
3726
3727 #[test]
3728 fn test_validation_of_invalid_list_array() {
3729 let array = unsafe {
3731 let values = Int32Array::from(vec![1, 2, 3]);
3732 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);
3735 let nulls = None;
3736 ListArray::new(Arc::new(field), offsets, Arc::new(values), nulls)
3737 };
3738
3739 expect_ipc_validation_error(
3740 Arc::new(array),
3741 "Invalid argument error: Offset invariant failure: offset at position 2 out of bounds: 4 > 2",
3742 );
3743 }
3744
3745 #[test]
3746 fn test_validation_of_invalid_string_array() {
3747 let valid: &[u8] = b" ";
3748 let mut invalid = vec![];
3749 invalid.extend_from_slice(b"ThisStringIsCertainlyLongerThan12Bytes");
3750 invalid.extend_from_slice(INVALID_UTF8_FIRST_CHAR);
3751 let binary_array = BinaryArray::from_iter(vec![None, Some(valid), None, Some(&invalid)]);
3752 let array = unsafe {
3755 StringArray::new_unchecked(
3756 binary_array.offsets().clone(),
3757 binary_array.values().clone(),
3758 binary_array.nulls().cloned(),
3759 )
3760 };
3761 expect_ipc_validation_error(
3762 Arc::new(array),
3763 "Invalid argument error: Invalid UTF8 sequence at string index 3 (3..45): invalid utf-8 sequence of 1 bytes from index 38",
3764 );
3765 }
3766
3767 #[test]
3768 fn test_validation_of_invalid_string_view_array() {
3769 let valid: &[u8] = b" ";
3770 let mut invalid = vec![];
3771 invalid.extend_from_slice(b"ThisStringIsCertainlyLongerThan12Bytes");
3772 invalid.extend_from_slice(INVALID_UTF8_FIRST_CHAR);
3773 let binary_view_array =
3774 BinaryViewArray::from_iter(vec![None, Some(valid), None, Some(&invalid)]);
3775 let array = unsafe {
3778 StringViewArray::new_unchecked(
3779 binary_view_array.views().clone(),
3780 Arc::clone(binary_view_array.data_buffers()),
3781 binary_view_array.nulls().cloned(),
3782 )
3783 };
3784 expect_ipc_validation_error(
3785 Arc::new(array),
3786 "Invalid argument error: Encountered non-UTF-8 data at index 3: invalid utf-8 sequence of 1 bytes from index 38",
3787 );
3788 }
3789
3790 #[test]
3793 fn test_validation_of_invalid_dictionary_array() {
3794 let array = unsafe {
3795 let values = StringArray::from_iter_values(["a", "b", "c"]);
3796 let keys = Int32Array::from(vec![1, 200]); DictionaryArray::new_unchecked(keys, Arc::new(values))
3798 };
3799
3800 expect_ipc_validation_error(
3801 Arc::new(array),
3802 "Invalid argument error: Value at position 1 out of bounds: 200 (should be in [0, 2])",
3803 );
3804 }
3805
3806 #[test]
3807 fn test_validation_of_invalid_union_array() {
3808 let array = unsafe {
3809 let fields = UnionFields::try_new(
3810 vec![1, 3], vec![
3812 Field::new("a", DataType::Int32, false),
3813 Field::new("b", DataType::Utf8, false),
3814 ],
3815 )
3816 .unwrap();
3817 let type_ids = ScalarBuffer::from(vec![1i8, 2, 3]); let offsets = None;
3819 let children: Vec<ArrayRef> = vec![
3820 Arc::new(Int32Array::from(vec![10, 20, 30])),
3821 Arc::new(StringArray::from(vec![Some("a"), Some("b"), Some("c")])),
3822 ];
3823
3824 UnionArray::new_unchecked(fields, type_ids, offsets, children)
3825 };
3826
3827 expect_ipc_validation_error(
3828 Arc::new(array),
3829 "Invalid argument error: Type Ids values must match one of the field type ids",
3830 );
3831 }
3832
3833 const INVALID_UTF8_FIRST_CHAR: &[u8] = &[0xa0, 0xa1, 0x20, 0x20];
3836
3837 fn expect_ipc_validation_error(array: ArrayRef, expected_err: &str) {
3839 let rb = RecordBatch::try_from_iter([("a", array)]).unwrap();
3840
3841 let buf = write_stream(&rb); read_stream_skip_validation(&buf).unwrap();
3844 let err = read_stream(&buf).unwrap_err();
3845 assert_eq!(err.to_string(), expected_err);
3846
3847 let buf = write_ipc(&rb); read_ipc_skip_validation(&buf).unwrap();
3850 let err = read_ipc(&buf).unwrap_err();
3851 assert_eq!(err.to_string(), expected_err);
3852
3853 read_ipc_with_decoder_skip_validation(buf.clone()).unwrap();
3855 let err = read_ipc_with_decoder(buf).unwrap_err();
3856 assert_eq!(err.to_string(), expected_err);
3857 }
3858
3859 #[test]
3860 fn test_roundtrip_schema() {
3861 let schema = Schema::new(vec![
3862 Field::new(
3863 "a",
3864 DataType::Dictionary(Box::new(DataType::UInt16), Box::new(DataType::Utf8)),
3865 false,
3866 ),
3867 Field::new(
3868 "b",
3869 DataType::Dictionary(Box::new(DataType::UInt16), Box::new(DataType::Utf8)),
3870 false,
3871 ),
3872 ]);
3873
3874 let options = IpcWriteOptions::default();
3875 let data_gen = IpcDataGenerator::default();
3876 let mut dict_tracker = DictionaryTracker::new(false);
3877 let encoded_data =
3878 data_gen.schema_to_bytes_with_dictionary_tracker(&schema, &mut dict_tracker, &options);
3879 let mut schema_bytes = vec![];
3880 write_message(&mut schema_bytes, encoded_data, &options).expect("write_message");
3881
3882 let begin_offset: usize = if schema_bytes[0..4].eq(&CONTINUATION_MARKER) {
3883 4
3884 } else {
3885 0
3886 };
3887
3888 size_prefixed_root_as_message(&schema_bytes[begin_offset..])
3889 .expect_err("size_prefixed_root_as_message");
3890
3891 let msg = parse_message(&schema_bytes).expect("parse_message");
3892 let ipc_schema = msg.header_as_schema().expect("header_as_schema");
3893 let new_schema = try_fb_to_schema(ipc_schema).unwrap();
3894
3895 assert_eq!(schema, new_schema);
3896 }
3897
3898 #[test]
3899 fn test_negative_meta_len() {
3900 let bytes = i32::to_le_bytes(-1);
3901 let mut buf = vec![];
3902 buf.extend(CONTINUATION_MARKER);
3903 buf.extend(bytes);
3904
3905 let reader = StreamReader::try_new(Cursor::new(buf), None);
3906 assert!(reader.is_err());
3907 }
3908
3909 #[test]
3915 fn test_read_null_dict_without_dictionary_batch() {
3916 let keys = Int32Array::new_null(4);
3918 let values: ArrayRef = new_empty_array(&DataType::Utf8);
3919 let dict_array = DictionaryArray::new(keys, values);
3920
3921 let schema = Arc::new(Schema::new(vec![Field::new(
3922 "d",
3923 dict_array.data_type().clone(),
3924 true,
3925 )]));
3926 let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(dict_array)]).unwrap();
3927
3928 let full_stream = write_stream(&batch);
3930
3931 let mut stripped = Vec::new();
3935 let mut cursor = Cursor::new(&full_stream);
3936 loop {
3937 let mut header = [0u8; 4];
3940 if cursor.read_exact(&mut header).is_err() {
3941 break;
3942 }
3943 if header == CONTINUATION_MARKER && cursor.read_exact(&mut header).is_err() {
3944 break;
3945 }
3946 let meta_len = u32::from_le_bytes(header) as usize;
3947 if meta_len == 0 {
3948 stripped.extend_from_slice(&CONTINUATION_MARKER);
3950 stripped.extend_from_slice(&0u32.to_le_bytes());
3951 break;
3952 }
3953 let mut meta_buf = vec![0u8; meta_len];
3954 cursor.read_exact(&mut meta_buf).unwrap();
3955
3956 let message = root_as_message(&meta_buf).unwrap();
3957 let body_len = message.bodyLength() as usize;
3958 let mut body_buf = vec![0u8; body_len];
3959 cursor.read_exact(&mut body_buf).unwrap();
3960
3961 if message.header_type() == crate::MessageHeader::DictionaryBatch {
3962 continue;
3965 }
3966 stripped.extend_from_slice(&CONTINUATION_MARKER);
3967 stripped.extend_from_slice(&(meta_len as u32).to_le_bytes());
3968 stripped.extend_from_slice(&meta_buf);
3969 stripped.extend_from_slice(&body_buf);
3970 }
3971
3972 let result = read_stream(&stripped).unwrap();
3974 assert_eq!(result.num_rows(), 4);
3975 assert_eq!(result.num_columns(), 1);
3976
3977 let col = result.column(0);
3978 assert_eq!(col.null_count(), 4);
3979 assert_eq!(col.len(), 4);
3980 assert!(matches!(col.data_type(), DataType::Dictionary(_, _)));
3982 }
3983
3984 #[test]
3987 fn test_projection_skip_list_view() {
3988 use crate::reader::FileReader;
3989 use crate::writer::FileWriter;
3990 use arrow_array::{
3991 GenericListViewArray, Int32Array, RecordBatch,
3992 builder::{GenericListViewBuilder, UInt32Builder},
3993 };
3994 use arrow_schema::{DataType, Field, Schema};
3995 use std::sync::Arc;
3996
3997 let mut builder = GenericListViewBuilder::<i32, _>::new(UInt32Builder::new());
3999
4000 builder.values().append_value(1);
4001 builder.values().append_value(2);
4002 builder.append(true);
4003
4004 builder.append(false);
4005
4006 builder.values().append_value(3);
4007 builder.values().append_value(4);
4008 builder.append(true);
4009
4010 let list_view: GenericListViewArray<i32> = builder.finish();
4011
4012 let values = Int32Array::from(vec![10, 20, 30]);
4014
4015 let schema = Arc::new(Schema::new(vec![
4017 Field::new("a", list_view.data_type().clone(), true),
4018 Field::new("b", DataType::Int32, false),
4019 ]));
4020 let batch =
4022 RecordBatch::try_new(schema, vec![Arc::new(list_view), Arc::new(values.clone())])
4023 .unwrap();
4024
4025 let mut buf = Vec::new();
4027 {
4028 let mut writer = FileWriter::try_new(&mut buf, &batch.schema()).unwrap();
4029 writer.write(&batch).unwrap();
4030 writer.finish().unwrap();
4031 }
4032
4033 let mut reader = FileReader::try_new(std::io::Cursor::new(buf), Some(vec![1])).unwrap();
4035 let read_batch = reader.next().unwrap().unwrap();
4036
4037 assert_eq!(read_batch.num_columns(), 1);
4039 assert_eq!(read_batch.column(0).as_ref(), &values);
4040 }
4041
4042 #[test]
4045 fn test_projection_skip_union_v4() {
4046 use crate::MetadataVersion;
4047 use crate::reader::FileReader;
4048 use crate::writer::{FileWriter, IpcWriteOptions};
4049 use arrow_array::{
4050 ArrayRef, Int32Array, RecordBatch, builder::UnionBuilder, types::Int32Type,
4051 };
4052 use arrow_schema::{DataType, Field, Schema};
4053 use std::sync::Arc;
4054
4055 let mut builder = UnionBuilder::new_dense();
4057 builder.append::<Int32Type>("a", 1).unwrap();
4058 builder.append::<Int32Type>("a", 2).unwrap();
4059 builder.append::<Int32Type>("a", 3).unwrap();
4060 let union = builder.build().unwrap();
4061
4062 let values = Int32Array::from(vec![10, 20, 30]);
4064
4065 let schema = Arc::new(Schema::new(vec![
4067 Field::new("union", union.data_type().clone(), false),
4068 Field::new("values", DataType::Int32, false),
4069 ]));
4070
4071 let batch = RecordBatch::try_new(
4073 schema,
4074 vec![Arc::new(union) as ArrayRef, Arc::new(values.clone())],
4075 )
4076 .unwrap();
4077
4078 let mut buf = Vec::new();
4080 {
4081 let options = IpcWriteOptions::try_new(8, false, MetadataVersion::V4).unwrap();
4082 let mut writer =
4083 FileWriter::try_new_with_options(&mut buf, &batch.schema(), options).unwrap();
4084 writer.write(&batch).unwrap();
4085 writer.finish().unwrap();
4086 }
4087 let mut reader = FileReader::try_new(std::io::Cursor::new(buf), Some(vec![1])).unwrap();
4089 let read_batch = reader.next().unwrap().unwrap();
4090
4091 assert_eq!(read_batch.num_columns(), 1);
4093 assert_eq!(read_batch.column(0).as_ref(), &values);
4094 }
4095
4096 #[test]
4100 fn test_projection_skip_fixed_width_types() {
4101 use std::sync::Arc;
4102
4103 use arrow_array::{ArrayRef, BooleanArray, Int32Array, RecordBatch, make_array};
4104 use arrow_buffer::{Buffer, MutableBuffer};
4105 use arrow_data::ArrayData;
4106 use arrow_schema::{DataType, Field, IntervalUnit, Schema, TimeUnit};
4107
4108 use crate::reader::FileReader;
4109 use crate::writer::FileWriter;
4110
4111 fn make_array_for_type(data_type: DataType) -> ArrayRef {
4113 let len = 3;
4114
4115 if matches!(data_type, DataType::Boolean) {
4116 return Arc::new(BooleanArray::from(vec![true, false, true]));
4117 }
4118
4119 let width = data_type.primitive_width().unwrap();
4120 let data = ArrayData::builder(data_type)
4121 .len(len)
4122 .add_buffer(Buffer::from(MutableBuffer::from_len_zeroed(len * width)))
4123 .build()
4124 .unwrap();
4125
4126 make_array(data)
4127 }
4128
4129 let data_types = vec![
4131 DataType::Boolean,
4132 DataType::Int8,
4133 DataType::Int16,
4134 DataType::Int32,
4135 DataType::Int64,
4136 DataType::UInt8,
4137 DataType::UInt16,
4138 DataType::UInt32,
4139 DataType::UInt64,
4140 DataType::Float16,
4141 DataType::Float32,
4142 DataType::Float64,
4143 DataType::Timestamp(TimeUnit::Second, None),
4144 DataType::Date32,
4145 DataType::Date64,
4146 DataType::Time32(TimeUnit::Second),
4147 DataType::Time64(TimeUnit::Microsecond),
4148 DataType::Duration(TimeUnit::Second),
4149 DataType::Interval(IntervalUnit::YearMonth),
4150 DataType::Interval(IntervalUnit::DayTime),
4151 DataType::Interval(IntervalUnit::MonthDayNano),
4152 DataType::Decimal32(9, 2),
4153 DataType::Decimal64(18, 2),
4154 DataType::Decimal128(38, 2),
4155 DataType::Decimal256(76, 2),
4156 ];
4157
4158 for data_type in data_types {
4163 let skipped = make_array_for_type(data_type.clone());
4164 let values = Int32Array::from(vec![10, 20, 30]);
4165
4166 let schema = Arc::new(Schema::new(vec![
4167 Field::new("skipped", data_type, false),
4168 Field::new("values", DataType::Int32, false),
4169 ]));
4170
4171 let batch =
4172 RecordBatch::try_new(schema, vec![skipped, Arc::new(values.clone())]).unwrap();
4173
4174 let mut buf = Vec::new();
4176 {
4177 let mut writer = FileWriter::try_new(&mut buf, &batch.schema()).unwrap();
4178 writer.write(&batch).unwrap();
4179 writer.finish().unwrap();
4180 }
4181
4182 let mut reader = FileReader::try_new(std::io::Cursor::new(buf), Some(vec![1])).unwrap();
4184 let read_batch = reader.next().unwrap().unwrap();
4185
4186 assert_eq!(read_batch.num_columns(), 1);
4188 assert_eq!(read_batch.column(0).as_ref(), &values);
4189 }
4190 }
4191
4192 fn stream_with_declared_body(body_length: i64, body_bytes: usize) -> Vec<u8> {
4196 let mut fbb = flatbuffers::FlatBufferBuilder::new();
4197 let mut message = crate::MessageBuilder::new(&mut fbb);
4198 message.add_version(crate::MetadataVersion::V5);
4199 message.add_header_type(crate::MessageHeader::NONE);
4200 message.add_bodyLength(body_length);
4201 let root = message.finish();
4202 fbb.finish(root, None);
4203 let metadata = fbb.finished_data();
4204
4205 let mut stream = Vec::new();
4206 stream.extend_from_slice(&CONTINUATION_MARKER);
4207 stream.extend_from_slice(&(metadata.len() as i32).to_le_bytes());
4208 stream.extend_from_slice(metadata);
4209 stream.resize(stream.len() + body_bytes, 0xAB);
4210 stream
4211 }
4212
4213 #[test]
4218 fn test_stream_reader_rejects_implausible_body_length() {
4219 for body_length in [i64::MAX, 1 << 50, -1] {
4220 let stream = stream_with_declared_body(body_length, 0);
4221 let err = StreamReader::try_new(std::io::Cursor::new(stream), None).expect_err(
4222 &format!("a message declaring {body_length} body bytes must not be accepted"),
4223 );
4224 let err = err.to_string();
4225 assert!(
4226 err.contains("Invalid IPC message body length")
4227 || err.contains("Unexpected end of stream")
4228 || err.contains("failed to fill whole buffer"),
4229 "unexpected error for body_length {body_length}: {err}"
4230 );
4231 }
4232 }
4233
4234 #[test]
4238 #[cfg_attr(miri, ignore)] fn test_stream_reader_rejects_truncated_body() {
4240 let over_prealloc = MAX_PREALLOC_BYTES as i64 + 1;
4241 for (body_length, body_bytes) in [(1024, 10), (over_prealloc, MAX_PREALLOC_BYTES)] {
4242 let stream = stream_with_declared_body(body_length, body_bytes);
4243 let err = StreamReader::try_new(std::io::Cursor::new(stream), None)
4244 .expect_err("a body backed by fewer bytes than declared must not be accepted");
4245 assert!(
4246 err.to_string().contains("failed to fill whole buffer"),
4247 "unexpected error for body_length {body_length}: {err}"
4248 );
4249 }
4250 }
4251
4252 #[test]
4254 fn test_stream_reader_rejects_unbacked_metadata_length() {
4255 let mut stream = Vec::new();
4256 stream.extend_from_slice(&CONTINUATION_MARKER);
4257 stream.extend_from_slice(&i32::MAX.to_le_bytes());
4258 stream.extend_from_slice(b"not this many bytes");
4259
4260 let err = StreamReader::try_new(std::io::Cursor::new(stream), None)
4261 .expect_err("metadata length must be backed by the stream");
4262 assert!(
4263 err.to_string().contains("Unexpected end of stream"),
4264 "unexpected error: {err}"
4265 );
4266 }
4267}