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(
2054 Field::REE_RUN_ENDS_FIELD_DEFAULT_NAME,
2055 DataType::Int16,
2056 false,
2057 )),
2058 Arc::new(Field::new(
2059 Field::REE_VALUES_FIELD_DEFAULT_NAME,
2060 DataType::Int32,
2061 true,
2062 )),
2063 );
2064
2065 Schema::new(vec![
2067 Field::new("f0", DataType::UInt32, false),
2068 Field::new("f1", DataType::Utf8, false),
2069 Field::new("f2", DataType::Boolean, false),
2070 Field::new("f3", union_data_type, true),
2071 Field::new("f4", DataType::Null, true),
2072 Field::new("f5", DataType::Float64, true),
2073 Field::new("f6", list_data_type, false),
2074 Field::new("f7", DataType::FixedSizeBinary(3), true),
2075 Field::new("f8", fixed_size_list_data_type, false),
2076 Field::new("f9", struct_data_type, false),
2077 Field::new("f10", run_encoded_data_type, false),
2078 Field::new("f11", DataType::Boolean, false),
2079 Field::new_dictionary("f12", DataType::Int8, DataType::Utf8, false),
2080 Field::new("f13", DataType::Utf8, false),
2081 ])
2082 }
2083
2084 fn create_test_projection_batch_data(schema: &Schema) -> RecordBatch {
2085 let array0 = UInt32Array::from(vec![1, 2, 3]);
2087 let array1 = StringArray::from(vec!["foo", "bar", "baz"]);
2088 let array2 = BooleanArray::from(vec![true, false, true]);
2089
2090 let mut union_builder = UnionBuilder::new_dense();
2091 union_builder.append::<Int32Type>("a", 1).unwrap();
2092 union_builder.append::<Float64Type>("b", 10.1).unwrap();
2093 union_builder.append_null::<Float64Type>("b").unwrap();
2094 let array3 = union_builder.build().unwrap();
2095
2096 let array4 = NullArray::new(3);
2097 let array5 = Float64Array::from(vec![Some(1.1), None, Some(3.3)]);
2098 let array6_values = vec![
2099 Some(vec![Some(10), Some(10), Some(10)]),
2100 Some(vec![Some(20), Some(20), Some(20)]),
2101 Some(vec![Some(30), Some(30)]),
2102 ];
2103 let array6 = ListArray::from_iter_primitive::<Int32Type, _, _>(array6_values);
2104 let array7_values = vec![vec![11, 12, 13], vec![22, 23, 24], vec![33, 34, 35]];
2105 let array7 = FixedSizeBinaryArray::try_from_iter(array7_values.into_iter()).unwrap();
2106
2107 let array8_values = ArrayData::builder(DataType::Int32)
2108 .len(9)
2109 .add_buffer(Buffer::from_slice_ref([40, 41, 42, 43, 44, 45, 46, 47, 48]))
2110 .build()
2111 .unwrap();
2112 let array8_data = ArrayData::builder(schema.field(8).data_type().clone())
2113 .len(3)
2114 .add_child_data(array8_values)
2115 .build()
2116 .unwrap();
2117 let array8 = FixedSizeListArray::from(array8_data);
2118
2119 let array9_id: ArrayRef = Arc::new(Int32Array::from(vec![1001, 1002, 1003]));
2120 let array9_list: ArrayRef =
2121 Arc::new(ListArray::from_iter_primitive::<Int8Type, _, _>(vec![
2122 Some(vec![Some(-10)]),
2123 Some(vec![Some(-20), Some(-20), Some(-20)]),
2124 Some(vec![Some(-30)]),
2125 ]));
2126 let array9 = ArrayDataBuilder::new(schema.field(9).data_type().clone())
2127 .add_child_data(array9_id.into_data())
2128 .add_child_data(array9_list.into_data())
2129 .len(3)
2130 .build()
2131 .unwrap();
2132 let array9 = StructArray::from(array9);
2133
2134 let array10_input = vec![Some(1_i32), None, None];
2135 let mut array10_builder = PrimitiveRunBuilder::<Int16Type, Int32Type>::new();
2136 array10_builder.extend(array10_input);
2137 let array10 = array10_builder.finish();
2138
2139 let array11 = BooleanArray::from(vec![false, false, true]);
2140
2141 let array12_values = StringArray::from(vec!["x", "yy", "zzz"]);
2142 let array12_keys = Int8Array::from_iter_values([1, 1, 2]);
2143 let array12 = DictionaryArray::new(array12_keys, Arc::new(array12_values));
2144
2145 let array13 = StringArray::from(vec!["a", "bb", "ccc"]);
2146
2147 RecordBatch::try_new(
2149 Arc::new(schema.clone()),
2150 vec![
2151 Arc::new(array0),
2152 Arc::new(array1),
2153 Arc::new(array2),
2154 Arc::new(array3),
2155 Arc::new(array4),
2156 Arc::new(array5),
2157 Arc::new(array6),
2158 Arc::new(array7),
2159 Arc::new(array8),
2160 Arc::new(array9),
2161 Arc::new(array10),
2162 Arc::new(array11),
2163 Arc::new(array12),
2164 Arc::new(array13),
2165 ],
2166 )
2167 .unwrap()
2168 }
2169
2170 #[test]
2171 fn test_negative_meta_len_start_stream() {
2172 let bytes = i32::to_le_bytes(-1);
2173 let mut buf = vec![];
2174 buf.extend(CONTINUATION_MARKER);
2175 buf.extend(bytes);
2176
2177 let reader_err = StreamReader::try_new(Cursor::new(buf), None).err();
2178 assert!(reader_err.is_some());
2179 assert_eq!(
2180 reader_err.unwrap().to_string(),
2181 "Parser error: Invalid metadata length: -1"
2182 );
2183 }
2184
2185 #[test]
2186 fn test_negative_meta_len_mid_stream() {
2187 let schema = Schema::new(vec![Field::new("a", DataType::Int32, false)]);
2188 let mut buf = Vec::new();
2189 {
2190 let mut writer = crate::writer::StreamWriter::try_new(&mut buf, &schema).unwrap();
2191 let batch =
2192 RecordBatch::try_new(Arc::new(schema), vec![Arc::new(Int32Array::from(vec![1]))])
2193 .unwrap();
2194 writer.write(&batch).unwrap();
2195 }
2196
2197 let bytes = i32::to_le_bytes(-1);
2198 buf.extend(CONTINUATION_MARKER);
2199 buf.extend(bytes);
2200
2201 let mut reader = StreamReader::try_new(Cursor::new(buf), None).unwrap();
2202 assert!(reader.maybe_next().is_ok());
2204 let batch_err = reader.maybe_next().err();
2206 assert!(batch_err.is_some());
2207 assert_eq!(
2208 batch_err.unwrap().to_string(),
2209 "Parser error: Invalid metadata length: -1"
2210 );
2211 }
2212
2213 #[test]
2214 fn test_invalid_dictionary_batch_without_data() {
2215 use crate::r#gen::Message::*;
2216 use flatbuffers::FlatBufferBuilder;
2217
2218 let schema = Schema::new(vec![Field::new(
2219 "col",
2220 DataType::Dictionary(Box::new(DataType::Int8), Box::new(DataType::Utf8)),
2221 true,
2222 )]);
2223
2224 let mut fbb = FlatBufferBuilder::new();
2227 let batch_offset = DictionaryBatch::create(
2228 &mut fbb,
2229 &DictionaryBatchArgs {
2230 id: 0,
2231 data: None,
2232 isDelta: false,
2233 },
2234 );
2235 fbb.finish_minimal(batch_offset);
2236 let batch_bytes = fbb.finished_data().to_vec();
2237 let batch = flatbuffers::root::<DictionaryBatch>(&batch_bytes).unwrap();
2238
2239 let data_buffer = Buffer::from(vec![0u8; 0]);
2240 let mut dictionaries: HashMap<i64, ArrayRef> = HashMap::new();
2241
2242 let err = read_dictionary(
2243 &data_buffer,
2244 batch,
2245 &schema,
2246 &mut dictionaries,
2247 &MetadataVersion::V5,
2248 )
2249 .unwrap_err();
2250
2251 assert_eq!(
2252 err.to_string(),
2253 "Parser error: Dictionary batch is missing its data"
2254 );
2255 }
2256
2257 #[test]
2258 fn test_missing_buffer_metadata_error() {
2259 use crate::r#gen::Message::*;
2260 use flatbuffers::FlatBufferBuilder;
2261
2262 let schema = Arc::new(Schema::new(vec![Field::new("col", DataType::Int32, true)]));
2263
2264 let mut fbb = FlatBufferBuilder::new();
2267 let nodes = fbb.create_vector(&[FieldNode::new(2, 0)]);
2268 let buffers = fbb.create_vector(&[crate::Buffer::new(0, 8)]);
2269 let batch_offset = RecordBatch::create(
2270 &mut fbb,
2271 &RecordBatchArgs {
2272 length: 2,
2273 nodes: Some(nodes),
2274 buffers: Some(buffers),
2275 compression: None,
2276 variadicBufferCounts: None,
2277 },
2278 );
2279 fbb.finish_minimal(batch_offset);
2280 let batch_bytes = fbb.finished_data().to_vec();
2281 let batch = flatbuffers::root::<RecordBatch>(&batch_bytes).unwrap();
2282
2283 let data_buffer = Buffer::from(vec![0u8; 8]);
2284 let dictionaries: HashMap<i64, ArrayRef> = HashMap::new();
2285 let metadata = MetadataVersion::V5;
2286
2287 let decoder = RecordBatchDecoder::try_new(
2288 &data_buffer,
2289 batch,
2290 schema.clone(),
2291 &dictionaries,
2292 &metadata,
2293 )
2294 .unwrap();
2295
2296 let result = decoder.read_record_batch();
2297
2298 match result {
2299 Err(ArrowError::IpcError(msg)) => {
2300 assert_eq!(msg, "Buffer count mismatched with metadata");
2301 }
2302 other => panic!("unexpected error: {other:?}"),
2303 }
2304 }
2305
2306 fn read_batch_with_variadic_count(count: i64, projection: Option<&[usize]>) -> ArrowError {
2309 use crate::r#gen::Message::*;
2310 use flatbuffers::FlatBufferBuilder;
2311
2312 let schema = Arc::new(Schema::new(vec![Field::new(
2313 "col",
2314 DataType::Utf8View,
2315 true,
2316 )]));
2317
2318 let mut fbb = FlatBufferBuilder::new();
2319 let nodes = fbb.create_vector(&[FieldNode::new(1, 0)]);
2320 let buffers = fbb.create_vector(&[crate::Buffer::new(0, 8), crate::Buffer::new(8, 8)]);
2321 let variadic_buffer_counts = fbb.create_vector(&[count]);
2322 let batch_offset = RecordBatch::create(
2323 &mut fbb,
2324 &RecordBatchArgs {
2325 length: 1,
2326 nodes: Some(nodes),
2327 buffers: Some(buffers),
2328 compression: None,
2329 variadicBufferCounts: Some(variadic_buffer_counts),
2330 },
2331 );
2332 fbb.finish_minimal(batch_offset);
2333 let batch_bytes = fbb.finished_data().to_vec();
2334 let batch = flatbuffers::root::<RecordBatch>(&batch_bytes).unwrap();
2335
2336 let data_buffer = Buffer::from(vec![0u8; 16]);
2337 let dictionaries: HashMap<i64, ArrayRef> = HashMap::new();
2338
2339 RecordBatchDecoder::try_new(
2340 &data_buffer,
2341 batch,
2342 schema,
2343 &dictionaries,
2344 &MetadataVersion::V5,
2345 )
2346 .unwrap()
2347 .with_projection(projection)
2348 .read_record_batch()
2349 .expect_err("should get error")
2350 }
2351
2352 #[test]
2355 fn test_invalid_variadic_buffer_count_error() {
2356 for count in [-2, -1, 1, i64::MAX] {
2359 for projection in [None, Some([].as_slice())] {
2360 let err = read_batch_with_variadic_count(count, projection);
2361 assert_eq!(
2362 err.to_string(),
2363 format!(
2364 "Ipc error: Invalid variadic count {count} for Utf8View column, \
2365 with 2 buffer(s) left in the message"
2366 ),
2367 "count {count}, projection {projection:?}"
2368 );
2369 }
2370 }
2371 }
2372
2373 #[test]
2375 fn test_valid_variadic_buffer_count_is_accepted() {
2376 let err = read_batch_with_variadic_count(0, None);
2377 assert!(!err.to_string().contains("Invalid variadic count"), "{err}");
2378 }
2379
2380 #[test]
2381 fn test_missing_footer_schema_error() {
2382 use crate::r#gen::File::{Footer, FooterArgs};
2383 use flatbuffers::FlatBufferBuilder;
2384
2385 let mut fbb = FlatBufferBuilder::new();
2388 let record_batches = fbb.create_vector::<Block>(&[]);
2389 let footer = Footer::create(
2390 &mut fbb,
2391 &FooterArgs {
2392 version: MetadataVersion::V5,
2393 schema: None,
2394 dictionaries: None,
2395 recordBatches: Some(record_batches),
2396 custom_metadata: None,
2397 },
2398 );
2399 fbb.finish(footer, None);
2400 let footer_data = fbb.finished_data();
2401
2402 let mut buf = Vec::new();
2404 buf.extend_from_slice(&crate::ARROW_MAGIC);
2405 buf.extend_from_slice(footer_data);
2406 buf.extend_from_slice(&(footer_data.len() as i32).to_le_bytes());
2407 buf.extend_from_slice(&crate::ARROW_MAGIC);
2408
2409 let err = FileReader::try_new(Cursor::new(buf), None)
2410 .expect_err("expected an error, not a panic");
2411 assert!(
2412 matches!(err, ArrowError::ParseError(_)),
2413 "expected ParseError, got {err:?}"
2414 );
2415 }
2416
2417 #[test]
2419 fn test_read_legacy_empty_list_without_offsets_buffer() {
2420 use crate::r#gen::Message::*;
2421 use flatbuffers::FlatBufferBuilder;
2422
2423 let schema = Arc::new(Schema::new(vec![Field::new_list(
2424 "items",
2425 Field::new_list_field(DataType::Int32, true),
2426 true,
2427 )]));
2428
2429 let mut fbb = FlatBufferBuilder::new();
2432 let nodes = fbb.create_vector(&[
2433 FieldNode::new(0, 0), FieldNode::new(0, 0), ]);
2436 let buffers = fbb.create_vector(&[
2437 crate::Buffer::new(0, 0), crate::Buffer::new(0, 0), crate::Buffer::new(0, 0), crate::Buffer::new(0, 0), ]);
2442 let batch_offset = RecordBatch::create(
2443 &mut fbb,
2444 &RecordBatchArgs {
2445 length: 0,
2446 nodes: Some(nodes),
2447 buffers: Some(buffers),
2448 compression: None,
2449 variadicBufferCounts: None,
2450 },
2451 );
2452 fbb.finish_minimal(batch_offset);
2453 let batch_bytes = fbb.finished_data().to_vec();
2454 let batch = flatbuffers::root::<RecordBatch>(&batch_bytes).unwrap();
2455
2456 let body = Buffer::from(Vec::<u8>::new());
2457 let dictionaries: HashMap<i64, ArrayRef> = HashMap::new();
2458 let metadata = MetadataVersion::V5;
2459
2460 let decoder =
2461 RecordBatchDecoder::try_new(&body, batch, schema.clone(), &dictionaries, &metadata)
2462 .unwrap();
2463
2464 let read_batch = decoder.read_record_batch().unwrap();
2465 assert_eq!(read_batch.num_rows(), 0);
2466
2467 let list = read_batch
2468 .column(0)
2469 .as_any()
2470 .downcast_ref::<ListArray>()
2471 .unwrap();
2472 assert_eq!(list.len(), 0);
2473 assert_eq!(list.values().len(), 0);
2474 }
2475
2476 #[test]
2478 fn test_read_legacy_empty_utf8_and_binary_without_offsets_buffer() {
2479 use crate::r#gen::Message::*;
2480 use flatbuffers::FlatBufferBuilder;
2481
2482 let schema = Arc::new(Schema::new(vec![
2483 Field::new("name", DataType::Utf8, true),
2484 Field::new("payload", DataType::Binary, true),
2485 ]));
2486
2487 let mut fbb = FlatBufferBuilder::new();
2490 let nodes = fbb.create_vector(&[
2491 FieldNode::new(0, 0), FieldNode::new(0, 0), ]);
2494 let buffers = fbb.create_vector(&[
2495 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), ]);
2502 let batch_offset = RecordBatch::create(
2503 &mut fbb,
2504 &RecordBatchArgs {
2505 length: 0,
2506 nodes: Some(nodes),
2507 buffers: Some(buffers),
2508 compression: None,
2509 variadicBufferCounts: None,
2510 },
2511 );
2512 fbb.finish_minimal(batch_offset);
2513 let batch_bytes = fbb.finished_data().to_vec();
2514 let batch = flatbuffers::root::<RecordBatch>(&batch_bytes).unwrap();
2515
2516 let body = Buffer::from(Vec::<u8>::new());
2517 let dictionaries: HashMap<i64, ArrayRef> = HashMap::new();
2518 let metadata = MetadataVersion::V5;
2519
2520 let decoder =
2521 RecordBatchDecoder::try_new(&body, batch, schema.clone(), &dictionaries, &metadata)
2522 .unwrap();
2523
2524 let read_batch = decoder.read_record_batch().unwrap();
2525 assert_eq!(read_batch.num_rows(), 0);
2526
2527 let utf8 = read_batch
2528 .column(0)
2529 .as_any()
2530 .downcast_ref::<StringArray>()
2531 .unwrap();
2532 assert_eq!(utf8.len(), 0);
2533 assert_eq!(utf8.value_offsets(), [0]);
2534
2535 let binary = read_batch
2536 .column(1)
2537 .as_any()
2538 .downcast_ref::<BinaryArray>()
2539 .unwrap();
2540 assert_eq!(binary.len(), 0);
2541 assert_eq!(binary.value_offsets(), [0]);
2542 }
2543
2544 #[test]
2545 fn test_projection_array_values() {
2546 let schema = create_test_projection_schema();
2548
2549 let batch = create_test_projection_batch_data(&schema);
2551
2552 let mut buf = Vec::new();
2554 {
2555 let mut writer = crate::writer::FileWriter::try_new(&mut buf, &schema).unwrap();
2556 writer.write(&batch).unwrap();
2557 writer.finish().unwrap();
2558 }
2559
2560 for index in 0..12 {
2562 let projection = vec![index];
2563 let reader = FileReader::try_new(std::io::Cursor::new(buf.clone()), Some(projection));
2564 let read_batch = reader.unwrap().next().unwrap().unwrap();
2565 let projected_column = read_batch.column(0);
2566 let expected_column = batch.column(index);
2567
2568 assert_eq!(projected_column.as_ref(), expected_column.as_ref());
2570 }
2571
2572 {
2573 let reader =
2575 FileReader::try_new(std::io::Cursor::new(buf.clone()), Some(vec![3, 2, 1]));
2576 let read_batch = reader.unwrap().next().unwrap().unwrap();
2577 let expected_batch = batch.project(&[3, 2, 1]).unwrap();
2578 assert_eq!(read_batch, expected_batch);
2579 }
2580 }
2581
2582 #[test]
2583 fn test_file_reader_projected_schema_matches_batch_schema() {
2584 let schema = create_test_projection_schema();
2585 let batch = create_test_projection_batch_data(&schema);
2586
2587 let mut buf = Vec::new();
2588 {
2589 let mut writer = crate::writer::FileWriter::try_new(&mut buf, &schema).unwrap();
2590 writer.write(&batch).unwrap();
2591 writer.finish().unwrap();
2592 }
2593
2594 let projection = vec![3, 2, 1];
2595 let mut reader = FileReader::try_new(Cursor::new(buf), Some(projection)).unwrap();
2596 let reader_schema = RecordBatchReader::schema(&reader);
2597 let read_batch = reader.next().unwrap().unwrap();
2598
2599 assert_eq!(reader_schema, read_batch.schema());
2600 }
2601
2602 #[test]
2603 fn test_stream_reader_projected_schema_matches_batch_schema() {
2604 let schema = create_test_projection_schema();
2605 let batch = create_test_projection_batch_data(&schema);
2606
2607 let mut buf = Vec::new();
2608 {
2609 let mut writer = crate::writer::StreamWriter::try_new(&mut buf, &schema).unwrap();
2610 writer.write(&batch).unwrap();
2611 writer.finish().unwrap();
2612 }
2613
2614 let projection = vec![3, 2, 1];
2615 let mut reader = StreamReader::try_new(Cursor::new(buf), Some(projection)).unwrap();
2616 let reader_schema = RecordBatchReader::schema(&reader);
2617 let read_batch = reader.next().unwrap().unwrap();
2618
2619 assert_eq!(reader_schema, read_batch.schema());
2620 }
2621
2622 #[test]
2623 fn test_file_reader_rejects_invalid_projection() {
2624 let schema = create_test_projection_schema();
2625 let batch = create_test_projection_batch_data(&schema);
2626
2627 let mut buf = Vec::new();
2628 {
2629 let mut writer = crate::writer::FileWriter::try_new(&mut buf, &schema).unwrap();
2630 writer.write(&batch).unwrap();
2631 writer.finish().unwrap();
2632 }
2633
2634 let result = FileReader::try_new(Cursor::new(buf), Some(vec![schema.fields().len()]));
2635
2636 assert!(matches!(result, Err(ArrowError::SchemaError(_))));
2637 }
2638
2639 #[test]
2640 fn test_projection_duplicate_indices() {
2641 let schema = create_test_projection_schema();
2642 let batch = create_test_projection_batch_data(&schema);
2643
2644 let mut buf = Vec::new();
2646 {
2647 let mut writer = crate::writer::FileWriter::try_new(&mut buf, &schema).unwrap();
2648 writer.write(&batch).unwrap();
2649 writer.finish().unwrap();
2650 }
2651
2652 for projection in [vec![1, 1], vec![2, 0, 2]] {
2654 let reader =
2655 FileReader::try_new(std::io::Cursor::new(buf.clone()), Some(projection.clone()));
2656 let read_batch = reader.unwrap().next().unwrap().unwrap();
2657
2658 let expected_batch = batch.project(&projection).unwrap();
2659 assert_eq!(read_batch, expected_batch);
2660 }
2661 }
2662
2663 #[test]
2664 fn test_arrow_single_float_row() {
2665 let schema = Schema::new(vec![
2666 Field::new("a", DataType::Float32, false),
2667 Field::new("b", DataType::Float32, false),
2668 Field::new("c", DataType::Int32, false),
2669 Field::new("d", DataType::Int32, false),
2670 ]);
2671 let arrays = vec![
2672 Arc::new(Float32Array::from(vec![1.23])) as ArrayRef,
2673 Arc::new(Float32Array::from(vec![-6.50])) as ArrayRef,
2674 Arc::new(Int32Array::from(vec![2])) as ArrayRef,
2675 Arc::new(Int32Array::from(vec![1])) as ArrayRef,
2676 ];
2677 let batch = RecordBatch::try_new(Arc::new(schema.clone()), arrays).unwrap();
2678 let mut file = tempfile::tempfile().unwrap();
2680 let mut stream_writer = crate::writer::StreamWriter::try_new(&mut file, &schema).unwrap();
2681 stream_writer.write(&batch).unwrap();
2682 stream_writer.finish().unwrap();
2683
2684 drop(stream_writer);
2685
2686 file.rewind().unwrap();
2687
2688 let reader = StreamReader::try_new(&mut file, None).unwrap();
2690
2691 reader.for_each(|batch| {
2692 let batch = batch.unwrap();
2693 assert!(
2694 batch
2695 .column(0)
2696 .as_any()
2697 .downcast_ref::<Float32Array>()
2698 .unwrap()
2699 .value(0)
2700 != 0.0
2701 );
2702 assert!(
2703 batch
2704 .column(1)
2705 .as_any()
2706 .downcast_ref::<Float32Array>()
2707 .unwrap()
2708 .value(0)
2709 != 0.0
2710 );
2711 });
2712
2713 file.rewind().unwrap();
2714
2715 let reader = StreamReader::try_new(file, Some(vec![0, 3])).unwrap();
2717
2718 reader.for_each(|batch| {
2719 let batch = batch.unwrap();
2720 assert_eq!(batch.schema().fields().len(), 2);
2721 assert_eq!(batch.schema().fields()[0].data_type(), &DataType::Float32);
2722 assert_eq!(batch.schema().fields()[1].data_type(), &DataType::Int32);
2723 });
2724 }
2725
2726 fn write_ipc(rb: &RecordBatch) -> Vec<u8> {
2728 let mut buf = Vec::new();
2729 let mut writer = crate::writer::FileWriter::try_new(&mut buf, rb.schema_ref()).unwrap();
2730 writer.write(rb).unwrap();
2731 writer.finish().unwrap();
2732 buf
2733 }
2734
2735 fn read_ipc(buf: &[u8]) -> Result<RecordBatch, ArrowError> {
2737 let mut reader = FileReader::try_new(std::io::Cursor::new(buf), None)?;
2738 reader.next().unwrap()
2739 }
2740
2741 fn read_ipc_skip_validation(buf: &[u8]) -> Result<RecordBatch, ArrowError> {
2744 let mut reader = unsafe {
2745 FileReader::try_new(std::io::Cursor::new(buf), None)?.with_skip_validation(true)
2746 };
2747 reader.next().unwrap()
2748 }
2749
2750 fn roundtrip_ipc(rb: &RecordBatch) -> RecordBatch {
2751 let buf = write_ipc(rb);
2752 read_ipc(&buf).unwrap()
2753 }
2754
2755 fn read_ipc_with_decoder(buf: Vec<u8>) -> Result<RecordBatch, ArrowError> {
2758 read_ipc_with_decoder_inner(buf, false)
2759 }
2760
2761 fn read_ipc_with_decoder_skip_validation(buf: Vec<u8>) -> Result<RecordBatch, ArrowError> {
2764 read_ipc_with_decoder_inner(buf, true)
2765 }
2766
2767 fn read_ipc_with_decoder_inner(
2768 buf: Vec<u8>,
2769 skip_validation: bool,
2770 ) -> Result<RecordBatch, ArrowError> {
2771 let buffer = Buffer::from_vec(buf);
2772 let trailer_start = buffer.len() - 10;
2773 let footer_len = read_footer_length(buffer[trailer_start..].try_into().unwrap())?;
2774 let footer = root_as_footer(&buffer[trailer_start - footer_len..trailer_start])
2775 .map_err(|e| ArrowError::InvalidArgumentError(format!("Invalid footer: {e}")))?;
2776
2777 let schema = try_fb_to_schema(footer.schema().unwrap()).unwrap();
2778
2779 let mut decoder = unsafe {
2780 FileDecoder::new(Arc::new(schema), footer.version())
2781 .with_skip_validation(skip_validation)
2782 };
2783 for block in footer.dictionaries().iter().flatten() {
2785 let block_len = block.bodyLength() as usize + block.metaDataLength() as usize;
2786 let data = buffer.slice_with_length(block.offset() as _, block_len);
2787 decoder.read_dictionary(block, &data)?
2788 }
2789
2790 let batches = footer.recordBatches().unwrap();
2792 assert_eq!(batches.len(), 1); let block = batches.get(0);
2795 let block_len = block.bodyLength() as usize + block.metaDataLength() as usize;
2796 let data = buffer.slice_with_length(block.offset() as _, block_len);
2797 Ok(decoder.read_record_batch(block, &data)?.unwrap())
2798 }
2799
2800 fn write_stream(rb: &RecordBatch) -> Vec<u8> {
2802 let mut buf = Vec::new();
2803 let mut writer = crate::writer::StreamWriter::try_new(&mut buf, rb.schema_ref()).unwrap();
2804 writer.write(rb).unwrap();
2805 writer.finish().unwrap();
2806 buf
2807 }
2808
2809 fn read_stream(buf: &[u8]) -> Result<RecordBatch, ArrowError> {
2811 let mut reader = StreamReader::try_new(std::io::Cursor::new(buf), None)?;
2812 reader.next().unwrap()
2813 }
2814
2815 fn read_stream_skip_validation(buf: &[u8]) -> Result<RecordBatch, ArrowError> {
2818 let mut reader = unsafe {
2819 StreamReader::try_new(std::io::Cursor::new(buf), None)?.with_skip_validation(true)
2820 };
2821 reader.next().unwrap()
2822 }
2823
2824 fn roundtrip_ipc_stream(rb: &RecordBatch) -> RecordBatch {
2825 let buf = write_stream(rb);
2826 read_stream(&buf).unwrap()
2827 }
2828
2829 #[test]
2830 fn test_roundtrip_with_custom_metadata() {
2831 let schema = Schema::new(vec![Field::new("dummy", DataType::Float64, false)]);
2832 let mut buf = Vec::new();
2833 let mut writer = crate::writer::FileWriter::try_new(&mut buf, &schema).unwrap();
2834 let mut test_metadata = HashMap::new();
2835 test_metadata.insert("abc".to_string(), "abc".to_string());
2836 test_metadata.insert("def".to_string(), "def".to_string());
2837 for (k, v) in &test_metadata {
2838 writer.write_metadata(k, v);
2839 }
2840 writer.finish().unwrap();
2841 drop(writer);
2842
2843 let reader = crate::reader::FileReader::try_new(std::io::Cursor::new(buf), None).unwrap();
2844 assert_eq!(reader.custom_metadata(), &test_metadata);
2845 }
2846
2847 #[test]
2848 fn test_roundtrip_nested_dict() {
2849 let inner: DictionaryArray<Int32Type> = vec!["a", "b", "a"].into_iter().collect();
2850
2851 let array = Arc::new(inner) as ArrayRef;
2852
2853 let dctfield = Arc::new(Field::new("dict", array.data_type().clone(), false));
2854
2855 let s = StructArray::from(vec![(dctfield, array)]);
2856 let struct_array = Arc::new(s) as ArrayRef;
2857
2858 let schema = Arc::new(Schema::new(vec![Field::new(
2859 "struct",
2860 struct_array.data_type().clone(),
2861 false,
2862 )]));
2863
2864 let batch = RecordBatch::try_new(schema, vec![struct_array]).unwrap();
2865
2866 assert_eq!(batch, roundtrip_ipc(&batch));
2867 }
2868
2869 #[test]
2870 fn test_roundtrip_nested_dict_no_preserve_dict_id() {
2871 let inner: DictionaryArray<Int32Type> = vec!["a", "b", "a"].into_iter().collect();
2872
2873 let array = Arc::new(inner) as ArrayRef;
2874
2875 let dctfield = Arc::new(Field::new("dict", array.data_type().clone(), false));
2876
2877 let s = StructArray::from(vec![(dctfield, array)]);
2878 let struct_array = Arc::new(s) as ArrayRef;
2879
2880 let schema = Arc::new(Schema::new(vec![Field::new(
2881 "struct",
2882 struct_array.data_type().clone(),
2883 false,
2884 )]));
2885
2886 let batch = RecordBatch::try_new(schema, vec![struct_array]).unwrap();
2887
2888 let mut buf = Vec::new();
2889 let mut writer = crate::writer::FileWriter::try_new_with_options(
2890 &mut buf,
2891 batch.schema_ref(),
2892 IpcWriteOptions::default(),
2893 )
2894 .unwrap();
2895 writer.write(&batch).unwrap();
2896 writer.finish().unwrap();
2897 drop(writer);
2898
2899 let mut reader = FileReader::try_new(std::io::Cursor::new(buf), None).unwrap();
2900
2901 assert_eq!(batch, reader.next().unwrap().unwrap());
2902 }
2903
2904 fn check_union_with_builder(mut builder: UnionBuilder) {
2905 builder.append::<Int32Type>("a", 1).unwrap();
2906 builder.append_null::<Int32Type>("a").unwrap();
2907 builder.append::<Float64Type>("c", 3.0).unwrap();
2908 builder.append::<Int32Type>("a", 4).unwrap();
2909 builder.append::<Int64Type>("d", 11).unwrap();
2910 let union = builder.build().unwrap();
2911
2912 let schema = Arc::new(Schema::new(vec![Field::new(
2913 "union",
2914 union.data_type().clone(),
2915 false,
2916 )]));
2917
2918 let union_array = Arc::new(union) as ArrayRef;
2919
2920 let rb = RecordBatch::try_new(schema, vec![union_array]).unwrap();
2921 let rb2 = roundtrip_ipc(&rb);
2922 assert_eq!(rb.schema(), rb2.schema());
2925 assert_eq!(rb.num_columns(), rb2.num_columns());
2926 assert_eq!(rb.num_rows(), rb2.num_rows());
2927 let union1 = rb.column(0);
2928 let union2 = rb2.column(0);
2929
2930 assert_eq!(union1, union2);
2931 }
2932
2933 #[test]
2934 fn test_roundtrip_dense_union() {
2935 check_union_with_builder(UnionBuilder::new_dense());
2936 }
2937
2938 #[test]
2939 fn test_roundtrip_sparse_union() {
2940 check_union_with_builder(UnionBuilder::new_sparse());
2941 }
2942
2943 #[test]
2944 fn test_roundtrip_struct_empty_fields() {
2945 let nulls = NullBuffer::from(&[true, true, false]);
2946 let rb = RecordBatch::try_from_iter([(
2947 "",
2948 Arc::new(StructArray::new_empty_fields(nulls.len(), Some(nulls))) as _,
2949 )])
2950 .unwrap();
2951 let rb2 = roundtrip_ipc(&rb);
2952 assert_eq!(rb, rb2);
2953 }
2954
2955 #[test]
2956 fn test_roundtrip_stream_run_array_sliced() {
2957 let run_array_1: Int32RunArray = vec!["a", "a", "a", "b", "b", "c", "c", "c"]
2958 .into_iter()
2959 .collect();
2960 let run_array_1_sliced = run_array_1.slice(2, 5);
2961
2962 let run_array_2_inupt = vec![Some(1_i32), None, None, Some(2), Some(2)];
2963 let mut run_array_2_builder = PrimitiveRunBuilder::<Int16Type, Int32Type>::new();
2964 run_array_2_builder.extend(run_array_2_inupt);
2965 let run_array_2 = run_array_2_builder.finish();
2966
2967 let schema = Arc::new(Schema::new(vec![
2968 Field::new(
2969 "run_array_1_sliced",
2970 run_array_1_sliced.data_type().clone(),
2971 false,
2972 ),
2973 Field::new("run_array_2", run_array_2.data_type().clone(), false),
2974 ]));
2975 let input_batch = RecordBatch::try_new(
2976 schema,
2977 vec![Arc::new(run_array_1_sliced.clone()), Arc::new(run_array_2)],
2978 )
2979 .unwrap();
2980 let output_batch = roundtrip_ipc_stream(&input_batch);
2981
2982 assert_eq!(input_batch.column(1), output_batch.column(1));
2986
2987 let run_array_1_unsliced = unslice_run_array(run_array_1_sliced.into_data()).unwrap();
2988 assert_eq!(run_array_1_unsliced, output_batch.column(0).into_data());
2989 }
2990
2991 #[test]
2992 fn test_roundtrip_stream_nested_dict() {
2993 let xs = vec!["AA", "BB", "AA", "CC", "BB"];
2994 let dict = Arc::new(
2995 xs.clone()
2996 .into_iter()
2997 .collect::<DictionaryArray<Int8Type>>(),
2998 );
2999 let string_array: ArrayRef = Arc::new(StringArray::from(xs.clone()));
3000 let struct_array = StructArray::from(vec![
3001 (
3002 Arc::new(Field::new("f2.1", DataType::Utf8, false)),
3003 string_array,
3004 ),
3005 (
3006 Arc::new(Field::new("f2.2_struct", dict.data_type().clone(), false)),
3007 dict.clone() as ArrayRef,
3008 ),
3009 ]);
3010 let schema = Arc::new(Schema::new(vec![
3011 Field::new("f1_string", DataType::Utf8, false),
3012 Field::new("f2_struct", struct_array.data_type().clone(), false),
3013 ]));
3014 let input_batch = RecordBatch::try_new(
3015 schema,
3016 vec![
3017 Arc::new(StringArray::from(xs.clone())),
3018 Arc::new(struct_array),
3019 ],
3020 )
3021 .unwrap();
3022 let output_batch = roundtrip_ipc_stream(&input_batch);
3023 assert_eq!(input_batch, output_batch);
3024 }
3025
3026 #[test]
3027 fn test_ipc_writers_reject_dictionary_of_dictionary_schema() {
3028 let values = Arc::new(StringArray::from(vec![Some("a"), Some("b")])) as ArrayRef;
3029 let inner = Arc::new(DictionaryArray::new(
3030 UInt32Array::from_iter_values([0, 1]),
3031 values,
3032 )) as ArrayRef;
3033 let outer = Arc::new(DictionaryArray::new(
3034 UInt32Array::from_iter_values([0, 1, 0]),
3035 inner,
3036 )) as ArrayRef;
3037
3038 let schema = Arc::new(Schema::new(vec![Field::new(
3039 "f1",
3040 outer.data_type().clone(),
3041 false,
3042 )]));
3043 let batch = RecordBatch::try_new(schema, vec![outer]).unwrap();
3044
3045 let mut stream = Vec::new();
3046 let Err(err) = crate::writer::StreamWriter::try_new(&mut stream, batch.schema_ref()) else {
3047 panic!("IPC stream writer should reject dictionary-of-dictionary schemas");
3048 };
3049 assert!(stream.is_empty());
3050
3051 assert!(
3052 err.to_string().contains("dictionary-of-dictionary values"),
3053 "unexpected error: {err}"
3054 );
3055
3056 let mut file = Vec::new();
3057 let Err(err) = crate::writer::FileWriter::try_new(&mut file, batch.schema_ref()) else {
3058 panic!("IPC file writer should reject dictionary-of-dictionary schemas");
3059 };
3060 assert!(file.is_empty());
3061
3062 assert!(
3063 err.to_string().contains("dictionary-of-dictionary values"),
3064 "unexpected error: {err}"
3065 );
3066 }
3067
3068 #[test]
3069 fn test_roundtrip_stream_nested_dict_of_map_of_dict() {
3070 let values = StringArray::from(vec![Some("a"), None, Some("b"), Some("c")]);
3071 let values = Arc::new(values) as ArrayRef;
3072 let value_dict_keys = Int8Array::from_iter_values([0, 1, 1, 2, 3, 1]);
3073 let value_dict_array = DictionaryArray::new(value_dict_keys, values.clone());
3074
3075 let key_dict_keys = Int8Array::from_iter_values([0, 0, 2, 2, 2, 3]);
3076 let key_dict_array = DictionaryArray::new(key_dict_keys, values);
3077
3078 #[expect(deprecated)]
3079 let keys_field = Arc::new(Field::new_dict(
3080 Field::MAP_KEY_FIELD_DEFAULT_NAME,
3081 DataType::Dictionary(Box::new(DataType::Int8), Box::new(DataType::Utf8)),
3082 false,
3083 1,
3084 false,
3085 ));
3086 #[expect(deprecated)]
3087 let values_field = Arc::new(Field::new_dict(
3088 Field::MAP_VALUE_FIELD_DEFAULT_NAME,
3089 DataType::Dictionary(Box::new(DataType::Int8), Box::new(DataType::Utf8)),
3090 true,
3091 2,
3092 false,
3093 ));
3094 let entry_struct = StructArray::from(vec![
3095 (keys_field, make_array(key_dict_array.into_data())),
3096 (values_field, make_array(value_dict_array.into_data())),
3097 ]);
3098 let map_data_type = DataType::Map(
3099 Arc::new(Field::new(
3100 Field::MAP_ENTRIES_FIELD_DEFAULT_NAME,
3101 entry_struct.data_type().clone(),
3102 false,
3103 )),
3104 false,
3105 );
3106
3107 let entry_offsets = Buffer::from_slice_ref([0, 2, 4, 6]);
3108 let map_data = ArrayData::builder(map_data_type)
3109 .len(3)
3110 .add_buffer(entry_offsets)
3111 .add_child_data(entry_struct.into_data())
3112 .build()
3113 .unwrap();
3114 let map_array = MapArray::from(map_data);
3115
3116 let dict_keys = Int8Array::from_iter_values([0, 1, 1, 2, 2, 1]);
3117 let dict_dict_array = DictionaryArray::new(dict_keys, Arc::new(map_array));
3118
3119 let schema = Arc::new(Schema::new(vec![Field::new(
3120 "f1",
3121 dict_dict_array.data_type().clone(),
3122 false,
3123 )]));
3124 let input_batch = RecordBatch::try_new(schema, vec![Arc::new(dict_dict_array)]).unwrap();
3125 let output_batch = roundtrip_ipc_stream(&input_batch);
3126 assert_eq!(input_batch, output_batch);
3127 }
3128
3129 fn test_roundtrip_stream_dict_of_list_of_dict_impl<
3130 OffsetSize: OffsetSizeTrait,
3131 U: ArrowNativeType,
3132 >(
3133 list_data_type: DataType,
3134 offsets: &[U; 5],
3135 ) {
3136 let values = StringArray::from(vec![Some("a"), None, Some("c"), None]);
3137 let keys = Int8Array::from_iter_values([0, 0, 1, 2, 0, 1, 3]);
3138 let dict_array = DictionaryArray::new(keys, Arc::new(values));
3139 let dict_data = dict_array.to_data();
3140
3141 let value_offsets = Buffer::from_slice_ref(offsets);
3142
3143 let list_data = ArrayData::builder(list_data_type)
3144 .len(4)
3145 .add_buffer(value_offsets)
3146 .add_child_data(dict_data)
3147 .build()
3148 .unwrap();
3149 let list_array = GenericListArray::<OffsetSize>::from(list_data);
3150
3151 let keys_for_dict = Int8Array::from_iter_values([0, 3, 0, 1, 1, 2, 0, 1, 3]);
3152 let dict_dict_array = DictionaryArray::new(keys_for_dict, Arc::new(list_array));
3153
3154 let schema = Arc::new(Schema::new(vec![Field::new(
3155 "f1",
3156 dict_dict_array.data_type().clone(),
3157 false,
3158 )]));
3159 let input_batch = RecordBatch::try_new(schema, vec![Arc::new(dict_dict_array)]).unwrap();
3160 let output_batch = roundtrip_ipc_stream(&input_batch);
3161 assert_eq!(input_batch, output_batch);
3162 }
3163
3164 #[test]
3165 fn test_roundtrip_stream_dict_of_list_of_dict() {
3166 #[expect(deprecated)]
3168 let list_data_type = DataType::List(Arc::new(Field::new_dict(
3169 "item",
3170 DataType::Dictionary(Box::new(DataType::Int8), Box::new(DataType::Utf8)),
3171 true,
3172 1,
3173 false,
3174 )));
3175 let offsets: &[i32; 5] = &[0, 2, 4, 4, 6];
3176 test_roundtrip_stream_dict_of_list_of_dict_impl::<i32, i32>(list_data_type, offsets);
3177
3178 #[expect(deprecated)]
3180 let list_data_type = DataType::LargeList(Arc::new(Field::new_dict(
3181 "item",
3182 DataType::Dictionary(Box::new(DataType::Int8), Box::new(DataType::Utf8)),
3183 true,
3184 1,
3185 false,
3186 )));
3187 let offsets: &[i64; 5] = &[0, 2, 4, 4, 7];
3188 test_roundtrip_stream_dict_of_list_of_dict_impl::<i64, i64>(list_data_type, offsets);
3189 }
3190
3191 #[test]
3192 fn test_roundtrip_stream_dict_of_fixed_size_list_of_dict() {
3193 let values = StringArray::from(vec![Some("a"), None, Some("c"), None]);
3194 let keys = Int8Array::from_iter_values([0, 0, 1, 2, 0, 1, 3, 1, 2]);
3195 let dict_array = DictionaryArray::new(keys, Arc::new(values));
3196 let dict_data = dict_array.into_data();
3197
3198 #[expect(deprecated)]
3199 let list_data_type = DataType::FixedSizeList(
3200 Arc::new(Field::new_dict(
3201 "item",
3202 DataType::Dictionary(Box::new(DataType::Int8), Box::new(DataType::Utf8)),
3203 true,
3204 1,
3205 false,
3206 )),
3207 3,
3208 );
3209 let list_data = ArrayData::builder(list_data_type)
3210 .len(3)
3211 .add_child_data(dict_data)
3212 .build()
3213 .unwrap();
3214 let list_array = FixedSizeListArray::from(list_data);
3215
3216 let keys_for_dict = Int8Array::from_iter_values([0, 1, 0, 1, 1, 2, 0, 1, 2]);
3217 let dict_dict_array = DictionaryArray::new(keys_for_dict, Arc::new(list_array));
3218
3219 let schema = Arc::new(Schema::new(vec![Field::new(
3220 "f1",
3221 dict_dict_array.data_type().clone(),
3222 false,
3223 )]));
3224 let input_batch = RecordBatch::try_new(schema, vec![Arc::new(dict_dict_array)]).unwrap();
3225 let output_batch = roundtrip_ipc_stream(&input_batch);
3226 assert_eq!(input_batch, output_batch);
3227 }
3228
3229 const LONG_TEST_STRING: &str =
3230 "This is a long string to make sure binary view array handles it";
3231
3232 #[test]
3233 fn test_roundtrip_view_types() {
3234 let schema = Schema::new(vec![
3235 Field::new("field_1", DataType::BinaryView, true),
3236 Field::new("field_2", DataType::Utf8, true),
3237 Field::new("field_3", DataType::Utf8View, true),
3238 ]);
3239 let bin_values: Vec<Option<&[u8]>> = vec![
3240 Some(b"foo"),
3241 None,
3242 Some(b"bar"),
3243 Some(LONG_TEST_STRING.as_bytes()),
3244 ];
3245 let utf8_values: Vec<Option<&str>> =
3246 vec![Some("foo"), None, Some("bar"), Some(LONG_TEST_STRING)];
3247 let bin_view_array = BinaryViewArray::from_iter(bin_values);
3248 let utf8_array = StringArray::from_iter(utf8_values.iter());
3249 let utf8_view_array = StringViewArray::from_iter(utf8_values);
3250 let record_batch = RecordBatch::try_new(
3251 Arc::new(schema.clone()),
3252 vec![
3253 Arc::new(bin_view_array),
3254 Arc::new(utf8_array),
3255 Arc::new(utf8_view_array),
3256 ],
3257 )
3258 .unwrap();
3259
3260 assert_eq!(record_batch, roundtrip_ipc(&record_batch));
3261 assert_eq!(record_batch, roundtrip_ipc_stream(&record_batch));
3262
3263 let sliced_batch = record_batch.slice(1, 2);
3264 assert_eq!(sliced_batch, roundtrip_ipc(&sliced_batch));
3265 assert_eq!(sliced_batch, roundtrip_ipc_stream(&sliced_batch));
3266 }
3267
3268 #[test]
3269 fn test_roundtrip_view_types_nested_dict() {
3270 let bin_values: Vec<Option<&[u8]>> = vec![
3271 Some(b"foo"),
3272 None,
3273 Some(b"bar"),
3274 Some(LONG_TEST_STRING.as_bytes()),
3275 Some(b"field"),
3276 ];
3277 let utf8_values: Vec<Option<&str>> = vec![
3278 Some("foo"),
3279 None,
3280 Some("bar"),
3281 Some(LONG_TEST_STRING),
3282 Some("field"),
3283 ];
3284 let bin_view_array = Arc::new(BinaryViewArray::from_iter(bin_values));
3285 let utf8_view_array = Arc::new(StringViewArray::from_iter(utf8_values));
3286
3287 let key_dict_keys = Int8Array::from_iter_values([0, 0, 2, 2, 0, 2, 3]);
3288 let key_dict_array = DictionaryArray::new(key_dict_keys, utf8_view_array.clone());
3289 #[expect(deprecated)]
3290 let keys_field = Arc::new(Field::new_dict(
3291 Field::MAP_KEY_FIELD_DEFAULT_NAME,
3292 DataType::Dictionary(Box::new(DataType::Int8), Box::new(DataType::Utf8View)),
3293 false,
3294 1,
3295 false,
3296 ));
3297
3298 let value_dict_keys = Int8Array::from_iter_values([0, 3, 0, 1, 2, 0, 1]);
3299 let value_dict_array = DictionaryArray::new(value_dict_keys, bin_view_array);
3300 #[expect(deprecated)]
3301 let values_field = Arc::new(Field::new_dict(
3302 Field::MAP_VALUE_FIELD_DEFAULT_NAME,
3303 DataType::Dictionary(Box::new(DataType::Int8), Box::new(DataType::BinaryView)),
3304 true,
3305 2,
3306 false,
3307 ));
3308 let entry_struct = StructArray::from(vec![
3309 (keys_field, make_array(key_dict_array.into_data())),
3310 (values_field, make_array(value_dict_array.into_data())),
3311 ]);
3312
3313 let map_data_type = DataType::Map(
3314 Arc::new(Field::new(
3315 Field::MAP_ENTRIES_FIELD_DEFAULT_NAME,
3316 entry_struct.data_type().clone(),
3317 false,
3318 )),
3319 false,
3320 );
3321 let entry_offsets = Buffer::from_slice_ref([0, 2, 4, 7]);
3322 let map_data = ArrayData::builder(map_data_type)
3323 .len(3)
3324 .add_buffer(entry_offsets)
3325 .add_child_data(entry_struct.into_data())
3326 .build()
3327 .unwrap();
3328 let map_array = MapArray::from(map_data);
3329
3330 let dict_keys = Int8Array::from_iter_values([0, 1, 0, 1, 1, 2, 0, 1, 2]);
3331 let dict_dict_array = DictionaryArray::new(dict_keys, Arc::new(map_array));
3332 let schema = Arc::new(Schema::new(vec![Field::new(
3333 "f1",
3334 dict_dict_array.data_type().clone(),
3335 false,
3336 )]));
3337 let batch = RecordBatch::try_new(schema, vec![Arc::new(dict_dict_array)]).unwrap();
3338 assert_eq!(batch, roundtrip_ipc(&batch));
3339 assert_eq!(batch, roundtrip_ipc_stream(&batch));
3340
3341 let sliced_batch = batch.slice(1, 2);
3342 assert_eq!(sliced_batch, roundtrip_ipc(&sliced_batch));
3343 assert_eq!(sliced_batch, roundtrip_ipc_stream(&sliced_batch));
3344 }
3345
3346 #[test]
3347 fn test_no_columns_batch() {
3348 let schema = Arc::new(Schema::empty());
3349 let options = RecordBatchOptions::new()
3350 .with_match_field_names(true)
3351 .with_row_count(Some(10));
3352 let input_batch = RecordBatch::try_new_with_options(schema, vec![], &options).unwrap();
3353 let output_batch = roundtrip_ipc_stream(&input_batch);
3354 assert_eq!(input_batch, output_batch);
3355 }
3356
3357 #[test]
3358 fn test_unaligned() {
3359 let batch = RecordBatch::try_from_iter(vec![(
3360 "i32",
3361 Arc::new(Int32Array::from(vec![1, 2, 3, 4])) as _,
3362 )])
3363 .unwrap();
3364
3365 let r#gen = IpcDataGenerator {};
3366 let mut dict_tracker = DictionaryTracker::new(false);
3367 let (_, encoded) = r#gen
3368 .encode(
3369 &batch,
3370 &mut dict_tracker,
3371 &Default::default(),
3372 &mut Default::default(),
3373 )
3374 .unwrap();
3375
3376 let message = root_as_message(&encoded.ipc_message).unwrap();
3377
3378 let mut buffer = MutableBuffer::with_capacity(encoded.arrow_data.len() + 1);
3380 buffer.push(0_u8);
3381 buffer.extend_from_slice(&encoded.arrow_data);
3382 let b = Buffer::from(buffer).slice(1);
3383 assert_ne!(b.as_ptr().align_offset(8), 0);
3384
3385 let ipc_batch = message.header_as_record_batch().unwrap();
3386 let roundtrip = RecordBatchDecoder::try_new(
3387 &b,
3388 ipc_batch,
3389 batch.schema(),
3390 &Default::default(),
3391 &message.version(),
3392 )
3393 .unwrap()
3394 .with_require_alignment(false)
3395 .read_record_batch()
3396 .unwrap();
3397 assert_eq!(batch, roundtrip);
3398 }
3399
3400 #[test]
3401 fn test_unaligned_throws_error_with_require_alignment() {
3402 let batch = RecordBatch::try_from_iter(vec![(
3403 "i32",
3404 Arc::new(Int32Array::from(vec![1, 2, 3, 4])) as _,
3405 )])
3406 .unwrap();
3407
3408 let r#gen = IpcDataGenerator {};
3409 let mut dict_tracker = DictionaryTracker::new(false);
3410 let (_, encoded) = r#gen
3411 .encode(
3412 &batch,
3413 &mut dict_tracker,
3414 &Default::default(),
3415 &mut Default::default(),
3416 )
3417 .unwrap();
3418
3419 let message = root_as_message(&encoded.ipc_message).unwrap();
3420
3421 let mut buffer = MutableBuffer::with_capacity(encoded.arrow_data.len() + 1);
3423 buffer.push(0_u8);
3424 buffer.extend_from_slice(&encoded.arrow_data);
3425 let b = Buffer::from(buffer).slice(1);
3426 assert_ne!(b.as_ptr().align_offset(8), 0);
3427
3428 let ipc_batch = message.header_as_record_batch().unwrap();
3429 let result = RecordBatchDecoder::try_new(
3430 &b,
3431 ipc_batch,
3432 batch.schema(),
3433 &Default::default(),
3434 &message.version(),
3435 )
3436 .unwrap()
3437 .with_require_alignment(true)
3438 .read_record_batch();
3439
3440 let error = result.unwrap_err();
3441 assert_eq!(
3442 error.to_string(),
3443 "Invalid argument error: Misaligned buffers[0] in array of type Int32, \
3444 offset from expected alignment of 4 by 1"
3445 );
3446 }
3447
3448 #[test]
3457 fn test_misaligned_buffers_error() {
3458 const OFFSET_RANGE: usize = 128;
3459
3460 let cases: Vec<(ArrayRef, usize)> = vec![
3464 (Arc::new(Int32Array::from_iter(0i32..100)) as _, 4),
3465 (Arc::new(Int64Array::from_iter(0i64..100)) as _, 8),
3466 (
3467 Arc::new(StringArray::from_iter_values(
3468 (0..100).map(|i| i.to_string()),
3469 )) as _,
3470 4,
3471 ),
3472 (
3473 Arc::new(LargeStringArray::from_iter_values(
3474 (0..100).map(|i| i.to_string()),
3475 )) as _,
3476 8,
3477 ),
3478 ];
3479
3480 for (array, alignment) in cases {
3481 let batch = RecordBatch::try_from_iter(vec![("col", Arc::clone(&array))]).unwrap();
3482 let encoder = IpcDataGenerator {};
3483 let mut dict_tracker = DictionaryTracker::new(false);
3484 let (_, encoded) = encoder
3485 .encode(
3486 &batch,
3487 &mut dict_tracker,
3488 &Default::default(),
3489 &mut Default::default(),
3490 )
3491 .unwrap();
3492 let message = root_as_message(&encoded.ipc_message).unwrap();
3493 let ipc_batch = message.header_as_record_batch().unwrap();
3494
3495 for offset in 0..OFFSET_RANGE {
3496 let mut storage = MutableBuffer::with_capacity(encoded.arrow_data.len() + offset);
3502 for _ in 0..offset {
3503 storage.push(0_u8);
3504 }
3505 storage.extend_from_slice(&encoded.arrow_data);
3506 let buf = Buffer::from(storage).slice(offset);
3507
3508 let result = RecordBatchDecoder::try_new(
3509 &buf,
3510 ipc_batch,
3511 batch.schema(),
3512 &Default::default(),
3513 &message.version(),
3514 )
3515 .unwrap()
3516 .with_require_alignment(true)
3517 .read_record_batch();
3518
3519 if offset % alignment == 0 {
3520 assert!(
3521 result.is_ok(),
3522 "type={} offset={offset}: expected Ok but got {:?}",
3523 array.data_type(),
3524 result.unwrap_err(),
3525 );
3526 } else {
3527 let err = result
3528 .expect_err(&format!(
3529 "type={} offset={offset}: expected Err for misaligned buffer",
3530 array.data_type()
3531 ))
3532 .to_string();
3533 assert!(
3534 err.contains("Misaligned buffers"),
3535 "type={} offset={offset}: unexpected error: {err}",
3536 array.data_type(),
3537 );
3538 }
3539 }
3540 }
3541 }
3542
3543 #[test]
3544 #[cfg_attr(miri, ignore)] fn test_file_with_massive_column_count() {
3546 let limit = 600_000;
3548
3549 let fields = (0..limit)
3550 .map(|i| Field::new(format!("{i}"), DataType::Boolean, false))
3551 .collect::<Vec<_>>();
3552 let schema = Arc::new(Schema::new(fields));
3553 let batch = RecordBatch::new_empty(schema);
3554
3555 let mut buf = Vec::new();
3556 let mut writer = crate::writer::FileWriter::try_new(&mut buf, batch.schema_ref()).unwrap();
3557 writer.write(&batch).unwrap();
3558 writer.finish().unwrap();
3559 drop(writer);
3560
3561 let mut reader = FileReaderBuilder::new()
3562 .with_max_footer_fb_tables(1_500_000)
3563 .build(std::io::Cursor::new(buf))
3564 .unwrap();
3565 let roundtrip_batch = reader.next().unwrap().unwrap();
3566
3567 assert_eq!(batch, roundtrip_batch);
3568 }
3569
3570 #[test]
3571 fn test_file_with_deeply_nested_columns() {
3572 let limit = 61;
3574
3575 let fields = (0..limit).fold(
3576 vec![Field::new("leaf", DataType::Boolean, false)],
3577 |field, index| vec![Field::new_struct(format!("{index}"), field, false)],
3578 );
3579 let schema = Arc::new(Schema::new(fields));
3580 let batch = RecordBatch::new_empty(schema);
3581
3582 let mut buf = Vec::new();
3583 let mut writer = crate::writer::FileWriter::try_new(&mut buf, batch.schema_ref()).unwrap();
3584 writer.write(&batch).unwrap();
3585 writer.finish().unwrap();
3586 drop(writer);
3587
3588 let mut reader = FileReaderBuilder::new()
3589 .with_max_footer_fb_depth(65)
3590 .build(std::io::Cursor::new(buf))
3591 .unwrap();
3592 let roundtrip_batch = reader.next().unwrap().unwrap();
3593
3594 assert_eq!(batch, roundtrip_batch);
3595 }
3596
3597 #[test]
3598 fn test_stream_reader_rejects_short_validity_buffer() {
3599 let data = ArrayDataBuilder::new(DataType::Int32)
3601 .len(8000)
3602 .add_buffer(ScalarBuffer::<i32>::from_iter(0..8000).into())
3603 .nulls(Some(NullBuffer::from(&[true, false, true, false])));
3604 let array: ArrayRef = unsafe { Arc::new(Int32Array::from(data.build_unchecked())) };
3605 let batch = RecordBatch::try_from_iter([("a", array)]).unwrap();
3606
3607 let mut stream = Vec::new();
3608 let mut writer =
3609 crate::writer::StreamWriter::try_new(&mut stream, &batch.schema()).unwrap();
3610 writer.write(&batch).unwrap();
3611 writer.finish().unwrap();
3612
3613 let mut reader = StreamReader::try_new(Cursor::new(stream), None).unwrap();
3614 let err = reader.next().unwrap().unwrap_err();
3615 assert_eq!(
3616 err.to_string(),
3617 "Invalid argument error: null_bit_buffer size too small. got 1 needed 1000"
3618 );
3619 }
3620
3621 #[test]
3622 fn test_invalid_struct_array_ipc_read_errors() {
3623 let a_field = Field::new("a", DataType::Int32, false);
3624 let b_field = Field::new("b", DataType::Int32, false);
3625 let struct_fields = Fields::from(vec![a_field.clone(), b_field.clone()]);
3626
3627 let a_array_data = ArrayData::builder(a_field.data_type().clone())
3628 .len(4)
3629 .add_buffer(Buffer::from_slice_ref([1, 2, 3, 4]))
3630 .build()
3631 .unwrap();
3632 let b_array_data = ArrayData::builder(b_field.data_type().clone())
3633 .len(3)
3634 .add_buffer(Buffer::from_slice_ref([5, 6, 7]))
3635 .build()
3636 .unwrap();
3637
3638 let invalid_struct_arr = unsafe {
3639 StructArray::new_unchecked(
3640 struct_fields,
3641 vec![make_array(a_array_data), make_array(b_array_data)],
3642 None,
3643 )
3644 };
3645
3646 expect_ipc_validation_error(
3647 Arc::new(invalid_struct_arr),
3648 "Invalid argument error: Incorrect array length for StructArray field \"b\", expected 4 got 3",
3649 );
3650 }
3651
3652 #[test]
3653 fn test_invalid_nested_array_ipc_read_errors() {
3654 let a_field = Field::new("a", DataType::Int32, false);
3656 let b_field = Field::new("b", DataType::Utf8, false);
3657
3658 let schema = Arc::new(Schema::new(vec![Field::new_struct(
3659 "s",
3660 vec![a_field.clone(), b_field.clone()],
3661 false,
3662 )]));
3663
3664 let a_array_data = ArrayData::builder(a_field.data_type().clone())
3665 .len(4)
3666 .add_buffer(Buffer::from_slice_ref([1, 2, 3, 4]))
3667 .build()
3668 .unwrap();
3669 let b_array_data = {
3671 let valid: &[u8] = b" ";
3672 let mut invalid = vec![];
3673 invalid.extend_from_slice(b"ValidString");
3674 invalid.extend_from_slice(INVALID_UTF8_FIRST_CHAR);
3675 let binary_array =
3676 BinaryArray::from_iter(vec![None, Some(valid), None, Some(&invalid)]);
3677 let array = unsafe {
3678 StringArray::new_unchecked(
3679 binary_array.offsets().clone(),
3680 binary_array.values().clone(),
3681 binary_array.nulls().cloned(),
3682 )
3683 };
3684 array.into_data()
3685 };
3686 let struct_data_type = schema.field(0).data_type();
3687
3688 let invalid_struct_arr = unsafe {
3689 make_array(
3690 ArrayData::builder(struct_data_type.clone())
3691 .len(4)
3692 .add_child_data(a_array_data)
3693 .add_child_data(b_array_data)
3694 .build_unchecked(),
3695 )
3696 };
3697 expect_ipc_validation_error(
3698 invalid_struct_arr,
3699 "Invalid argument error: Invalid UTF8 sequence at string index 3 (3..18): invalid utf-8 sequence of 1 bytes from index 11",
3700 );
3701 }
3702
3703 #[test]
3704 fn test_same_dict_id_without_preserve() {
3705 let batch = RecordBatch::try_new(
3706 Arc::new(Schema::new(
3707 ["a", "b"]
3708 .iter()
3709 .map(|name| {
3710 #[expect(deprecated)]
3711 Field::new_dict(
3712 name.to_string(),
3713 DataType::Dictionary(
3714 Box::new(DataType::Int32),
3715 Box::new(DataType::Utf8),
3716 ),
3717 true,
3718 0,
3719 false,
3720 )
3721 })
3722 .collect::<Vec<Field>>(),
3723 )),
3724 vec![
3725 Arc::new(
3726 vec![Some("c"), Some("d")]
3727 .into_iter()
3728 .collect::<DictionaryArray<Int32Type>>(),
3729 ) as ArrayRef,
3730 Arc::new(
3731 vec![Some("e"), Some("f")]
3732 .into_iter()
3733 .collect::<DictionaryArray<Int32Type>>(),
3734 ) as ArrayRef,
3735 ],
3736 )
3737 .expect("Failed to create RecordBatch");
3738
3739 let mut buf = vec![];
3741 {
3742 let mut writer = crate::writer::StreamWriter::try_new_with_options(
3743 &mut buf,
3744 batch.schema().as_ref(),
3745 crate::writer::IpcWriteOptions::default(),
3746 )
3747 .expect("Failed to create StreamWriter");
3748 writer.write(&batch).expect("Failed to write RecordBatch");
3749 writer.finish().expect("Failed to finish StreamWriter");
3750 }
3751
3752 StreamReader::try_new(std::io::Cursor::new(buf), None)
3753 .expect("Failed to create StreamReader")
3754 .for_each(|decoded_batch| {
3755 assert_eq!(decoded_batch.expect("Failed to read RecordBatch"), batch);
3756 });
3757 }
3758
3759 #[test]
3760 fn test_validation_of_invalid_list_array() {
3761 let array = unsafe {
3763 let values = Int32Array::from(vec![1, 2, 3]);
3764 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);
3767 let nulls = None;
3768 ListArray::new(Arc::new(field), offsets, Arc::new(values), nulls)
3769 };
3770
3771 expect_ipc_validation_error(
3772 Arc::new(array),
3773 "Invalid argument error: Offset invariant failure: offset at position 2 out of bounds: 4 > 2",
3774 );
3775 }
3776
3777 #[test]
3778 fn test_validation_of_invalid_string_array() {
3779 let valid: &[u8] = b" ";
3780 let mut invalid = vec![];
3781 invalid.extend_from_slice(b"ThisStringIsCertainlyLongerThan12Bytes");
3782 invalid.extend_from_slice(INVALID_UTF8_FIRST_CHAR);
3783 let binary_array = BinaryArray::from_iter(vec![None, Some(valid), None, Some(&invalid)]);
3784 let array = unsafe {
3787 StringArray::new_unchecked(
3788 binary_array.offsets().clone(),
3789 binary_array.values().clone(),
3790 binary_array.nulls().cloned(),
3791 )
3792 };
3793 expect_ipc_validation_error(
3794 Arc::new(array),
3795 "Invalid argument error: Invalid UTF8 sequence at string index 3 (3..45): invalid utf-8 sequence of 1 bytes from index 38",
3796 );
3797 }
3798
3799 #[test]
3800 fn test_validation_of_invalid_string_view_array() {
3801 let valid: &[u8] = b" ";
3802 let mut invalid = vec![];
3803 invalid.extend_from_slice(b"ThisStringIsCertainlyLongerThan12Bytes");
3804 invalid.extend_from_slice(INVALID_UTF8_FIRST_CHAR);
3805 let binary_view_array =
3806 BinaryViewArray::from_iter(vec![None, Some(valid), None, Some(&invalid)]);
3807 let array = unsafe {
3810 StringViewArray::new_unchecked(
3811 binary_view_array.views().clone(),
3812 Arc::clone(binary_view_array.data_buffers()),
3813 binary_view_array.nulls().cloned(),
3814 )
3815 };
3816 expect_ipc_validation_error(
3817 Arc::new(array),
3818 "Invalid argument error: Encountered non-UTF-8 data at index 3: invalid utf-8 sequence of 1 bytes from index 38",
3819 );
3820 }
3821
3822 #[test]
3825 fn test_validation_of_invalid_dictionary_array() {
3826 let array = unsafe {
3827 let values = StringArray::from_iter_values(["a", "b", "c"]);
3828 let keys = Int32Array::from(vec![1, 200]); DictionaryArray::new_unchecked(keys, Arc::new(values))
3830 };
3831
3832 expect_ipc_validation_error(
3833 Arc::new(array),
3834 "Invalid argument error: Value at position 1 out of bounds: 200 (should be in [0, 2])",
3835 );
3836 }
3837
3838 #[test]
3839 fn test_validation_of_invalid_union_array() {
3840 let array = unsafe {
3841 let fields = UnionFields::try_new(
3842 vec![1, 3], vec![
3844 Field::new("a", DataType::Int32, false),
3845 Field::new("b", DataType::Utf8, false),
3846 ],
3847 )
3848 .unwrap();
3849 let type_ids = ScalarBuffer::from(vec![1i8, 2, 3]); let offsets = None;
3851 let children: Vec<ArrayRef> = vec![
3852 Arc::new(Int32Array::from(vec![10, 20, 30])),
3853 Arc::new(StringArray::from(vec![Some("a"), Some("b"), Some("c")])),
3854 ];
3855
3856 UnionArray::new_unchecked(fields, type_ids, offsets, children)
3857 };
3858
3859 expect_ipc_validation_error(
3860 Arc::new(array),
3861 "Invalid argument error: Type Ids values must match one of the field type ids",
3862 );
3863 }
3864
3865 const INVALID_UTF8_FIRST_CHAR: &[u8] = &[0xa0, 0xa1, 0x20, 0x20];
3868
3869 fn expect_ipc_validation_error(array: ArrayRef, expected_err: &str) {
3871 let rb = RecordBatch::try_from_iter([("a", array)]).unwrap();
3872
3873 let buf = write_stream(&rb); read_stream_skip_validation(&buf).unwrap();
3876 let err = read_stream(&buf).unwrap_err();
3877 assert_eq!(err.to_string(), expected_err);
3878
3879 let buf = write_ipc(&rb); read_ipc_skip_validation(&buf).unwrap();
3882 let err = read_ipc(&buf).unwrap_err();
3883 assert_eq!(err.to_string(), expected_err);
3884
3885 read_ipc_with_decoder_skip_validation(buf.clone()).unwrap();
3887 let err = read_ipc_with_decoder(buf).unwrap_err();
3888 assert_eq!(err.to_string(), expected_err);
3889 }
3890
3891 #[test]
3892 fn test_roundtrip_schema() {
3893 let schema = Schema::new(vec![
3894 Field::new(
3895 "a",
3896 DataType::Dictionary(Box::new(DataType::UInt16), Box::new(DataType::Utf8)),
3897 false,
3898 ),
3899 Field::new(
3900 "b",
3901 DataType::Dictionary(Box::new(DataType::UInt16), Box::new(DataType::Utf8)),
3902 false,
3903 ),
3904 ]);
3905
3906 let options = IpcWriteOptions::default();
3907 let data_gen = IpcDataGenerator::default();
3908 let mut dict_tracker = DictionaryTracker::new(false);
3909 let encoded_data =
3910 data_gen.schema_to_bytes_with_dictionary_tracker(&schema, &mut dict_tracker, &options);
3911 let mut schema_bytes = vec![];
3912 write_message(&mut schema_bytes, encoded_data, &options).expect("write_message");
3913
3914 let begin_offset: usize = if schema_bytes[0..4].eq(&CONTINUATION_MARKER) {
3915 4
3916 } else {
3917 0
3918 };
3919
3920 size_prefixed_root_as_message(&schema_bytes[begin_offset..])
3921 .expect_err("size_prefixed_root_as_message");
3922
3923 let msg = parse_message(&schema_bytes).expect("parse_message");
3924 let ipc_schema = msg.header_as_schema().expect("header_as_schema");
3925 let new_schema = try_fb_to_schema(ipc_schema).unwrap();
3926
3927 assert_eq!(schema, new_schema);
3928 }
3929
3930 #[test]
3931 fn test_negative_meta_len() {
3932 let bytes = i32::to_le_bytes(-1);
3933 let mut buf = vec![];
3934 buf.extend(CONTINUATION_MARKER);
3935 buf.extend(bytes);
3936
3937 let reader = StreamReader::try_new(Cursor::new(buf), None);
3938 assert!(reader.is_err());
3939 }
3940
3941 #[test]
3947 fn test_read_null_dict_without_dictionary_batch() {
3948 let keys = Int32Array::new_null(4);
3950 let values: ArrayRef = new_empty_array(&DataType::Utf8);
3951 let dict_array = DictionaryArray::new(keys, values);
3952
3953 let schema = Arc::new(Schema::new(vec![Field::new(
3954 "d",
3955 dict_array.data_type().clone(),
3956 true,
3957 )]));
3958 let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(dict_array)]).unwrap();
3959
3960 let full_stream = write_stream(&batch);
3962
3963 let mut stripped = Vec::new();
3967 let mut cursor = Cursor::new(&full_stream);
3968 loop {
3969 let mut header = [0u8; 4];
3972 if cursor.read_exact(&mut header).is_err() {
3973 break;
3974 }
3975 if header == CONTINUATION_MARKER && cursor.read_exact(&mut header).is_err() {
3976 break;
3977 }
3978 let meta_len = u32::from_le_bytes(header) as usize;
3979 if meta_len == 0 {
3980 stripped.extend_from_slice(&CONTINUATION_MARKER);
3982 stripped.extend_from_slice(&0u32.to_le_bytes());
3983 break;
3984 }
3985 let mut meta_buf = vec![0u8; meta_len];
3986 cursor.read_exact(&mut meta_buf).unwrap();
3987
3988 let message = root_as_message(&meta_buf).unwrap();
3989 let body_len = message.bodyLength() as usize;
3990 let mut body_buf = vec![0u8; body_len];
3991 cursor.read_exact(&mut body_buf).unwrap();
3992
3993 if message.header_type() == crate::MessageHeader::DictionaryBatch {
3994 continue;
3997 }
3998 stripped.extend_from_slice(&CONTINUATION_MARKER);
3999 stripped.extend_from_slice(&(meta_len as u32).to_le_bytes());
4000 stripped.extend_from_slice(&meta_buf);
4001 stripped.extend_from_slice(&body_buf);
4002 }
4003
4004 let result = read_stream(&stripped).unwrap();
4006 assert_eq!(result.num_rows(), 4);
4007 assert_eq!(result.num_columns(), 1);
4008
4009 let col = result.column(0);
4010 assert_eq!(col.null_count(), 4);
4011 assert_eq!(col.len(), 4);
4012 assert!(matches!(col.data_type(), DataType::Dictionary(_, _)));
4014 }
4015
4016 #[test]
4019 fn test_projection_skip_list_view() {
4020 use crate::reader::FileReader;
4021 use crate::writer::FileWriter;
4022 use arrow_array::{
4023 GenericListViewArray, Int32Array, RecordBatch,
4024 builder::{GenericListViewBuilder, UInt32Builder},
4025 };
4026 use arrow_schema::{DataType, Field, Schema};
4027 use std::sync::Arc;
4028
4029 let mut builder = GenericListViewBuilder::<i32, _>::new(UInt32Builder::new());
4031
4032 builder.values().append_value(1);
4033 builder.values().append_value(2);
4034 builder.append(true);
4035
4036 builder.append(false);
4037
4038 builder.values().append_value(3);
4039 builder.values().append_value(4);
4040 builder.append(true);
4041
4042 let list_view: GenericListViewArray<i32> = builder.finish();
4043
4044 let values = Int32Array::from(vec![10, 20, 30]);
4046
4047 let schema = Arc::new(Schema::new(vec![
4049 Field::new("a", list_view.data_type().clone(), true),
4050 Field::new("b", DataType::Int32, false),
4051 ]));
4052 let batch =
4054 RecordBatch::try_new(schema, vec![Arc::new(list_view), Arc::new(values.clone())])
4055 .unwrap();
4056
4057 let mut buf = Vec::new();
4059 {
4060 let mut writer = FileWriter::try_new(&mut buf, &batch.schema()).unwrap();
4061 writer.write(&batch).unwrap();
4062 writer.finish().unwrap();
4063 }
4064
4065 let mut reader = FileReader::try_new(std::io::Cursor::new(buf), Some(vec![1])).unwrap();
4067 let read_batch = reader.next().unwrap().unwrap();
4068
4069 assert_eq!(read_batch.num_columns(), 1);
4071 assert_eq!(read_batch.column(0).as_ref(), &values);
4072 }
4073
4074 #[test]
4077 fn test_projection_skip_union_v4() {
4078 use crate::MetadataVersion;
4079 use crate::reader::FileReader;
4080 use crate::writer::{FileWriter, IpcWriteOptions};
4081 use arrow_array::{
4082 ArrayRef, Int32Array, RecordBatch, builder::UnionBuilder, types::Int32Type,
4083 };
4084 use arrow_schema::{DataType, Field, Schema};
4085 use std::sync::Arc;
4086
4087 let mut builder = UnionBuilder::new_dense();
4089 builder.append::<Int32Type>("a", 1).unwrap();
4090 builder.append::<Int32Type>("a", 2).unwrap();
4091 builder.append::<Int32Type>("a", 3).unwrap();
4092 let union = builder.build().unwrap();
4093
4094 let values = Int32Array::from(vec![10, 20, 30]);
4096
4097 let schema = Arc::new(Schema::new(vec![
4099 Field::new("union", union.data_type().clone(), false),
4100 Field::new("values", DataType::Int32, false),
4101 ]));
4102
4103 let batch = RecordBatch::try_new(
4105 schema,
4106 vec![Arc::new(union) as ArrayRef, Arc::new(values.clone())],
4107 )
4108 .unwrap();
4109
4110 let mut buf = Vec::new();
4112 {
4113 let options = IpcWriteOptions::try_new(8, false, MetadataVersion::V4).unwrap();
4114 let mut writer =
4115 FileWriter::try_new_with_options(&mut buf, &batch.schema(), options).unwrap();
4116 writer.write(&batch).unwrap();
4117 writer.finish().unwrap();
4118 }
4119 let mut reader = FileReader::try_new(std::io::Cursor::new(buf), Some(vec![1])).unwrap();
4121 let read_batch = reader.next().unwrap().unwrap();
4122
4123 assert_eq!(read_batch.num_columns(), 1);
4125 assert_eq!(read_batch.column(0).as_ref(), &values);
4126 }
4127
4128 #[test]
4132 fn test_projection_skip_fixed_width_types() {
4133 use std::sync::Arc;
4134
4135 use arrow_array::{ArrayRef, BooleanArray, Int32Array, RecordBatch, make_array};
4136 use arrow_buffer::{Buffer, MutableBuffer};
4137 use arrow_data::ArrayData;
4138 use arrow_schema::{DataType, Field, IntervalUnit, Schema, TimeUnit};
4139
4140 use crate::reader::FileReader;
4141 use crate::writer::FileWriter;
4142
4143 fn make_array_for_type(data_type: DataType) -> ArrayRef {
4145 let len = 3;
4146
4147 if matches!(data_type, DataType::Boolean) {
4148 return Arc::new(BooleanArray::from(vec![true, false, true]));
4149 }
4150
4151 let width = data_type.primitive_width().unwrap();
4152 let data = ArrayData::builder(data_type)
4153 .len(len)
4154 .add_buffer(Buffer::from(MutableBuffer::from_len_zeroed(len * width)))
4155 .build()
4156 .unwrap();
4157
4158 make_array(data)
4159 }
4160
4161 let data_types = vec![
4163 DataType::Boolean,
4164 DataType::Int8,
4165 DataType::Int16,
4166 DataType::Int32,
4167 DataType::Int64,
4168 DataType::UInt8,
4169 DataType::UInt16,
4170 DataType::UInt32,
4171 DataType::UInt64,
4172 DataType::Float16,
4173 DataType::Float32,
4174 DataType::Float64,
4175 DataType::Timestamp(TimeUnit::Second, None),
4176 DataType::Date32,
4177 DataType::Date64,
4178 DataType::Time32(TimeUnit::Second),
4179 DataType::Time64(TimeUnit::Microsecond),
4180 DataType::Duration(TimeUnit::Second),
4181 DataType::Interval(IntervalUnit::YearMonth),
4182 DataType::Interval(IntervalUnit::DayTime),
4183 DataType::Interval(IntervalUnit::MonthDayNano),
4184 DataType::Decimal32(9, 2),
4185 DataType::Decimal64(18, 2),
4186 DataType::Decimal128(38, 2),
4187 DataType::Decimal256(76, 2),
4188 ];
4189
4190 for data_type in data_types {
4195 let skipped = make_array_for_type(data_type.clone());
4196 let values = Int32Array::from(vec![10, 20, 30]);
4197
4198 let schema = Arc::new(Schema::new(vec![
4199 Field::new("skipped", data_type, false),
4200 Field::new("values", DataType::Int32, false),
4201 ]));
4202
4203 let batch =
4204 RecordBatch::try_new(schema, vec![skipped, Arc::new(values.clone())]).unwrap();
4205
4206 let mut buf = Vec::new();
4208 {
4209 let mut writer = FileWriter::try_new(&mut buf, &batch.schema()).unwrap();
4210 writer.write(&batch).unwrap();
4211 writer.finish().unwrap();
4212 }
4213
4214 let mut reader = FileReader::try_new(std::io::Cursor::new(buf), Some(vec![1])).unwrap();
4216 let read_batch = reader.next().unwrap().unwrap();
4217
4218 assert_eq!(read_batch.num_columns(), 1);
4220 assert_eq!(read_batch.column(0).as_ref(), &values);
4221 }
4222 }
4223
4224 fn stream_with_declared_body(body_length: i64, body_bytes: usize) -> Vec<u8> {
4228 let mut fbb = flatbuffers::FlatBufferBuilder::new();
4229 let mut message = crate::MessageBuilder::new(&mut fbb);
4230 message.add_version(crate::MetadataVersion::V5);
4231 message.add_header_type(crate::MessageHeader::NONE);
4232 message.add_bodyLength(body_length);
4233 let root = message.finish();
4234 fbb.finish(root, None);
4235 let metadata = fbb.finished_data();
4236
4237 let mut stream = Vec::new();
4238 stream.extend_from_slice(&CONTINUATION_MARKER);
4239 stream.extend_from_slice(&(metadata.len() as i32).to_le_bytes());
4240 stream.extend_from_slice(metadata);
4241 stream.resize(stream.len() + body_bytes, 0xAB);
4242 stream
4243 }
4244
4245 #[test]
4250 fn test_stream_reader_rejects_implausible_body_length() {
4251 for body_length in [i64::MAX, 1 << 50, -1] {
4252 let stream = stream_with_declared_body(body_length, 0);
4253 let err = StreamReader::try_new(std::io::Cursor::new(stream), None).expect_err(
4254 &format!("a message declaring {body_length} body bytes must not be accepted"),
4255 );
4256 let err = err.to_string();
4257 assert!(
4258 err.contains("Invalid IPC message body length")
4259 || err.contains("Unexpected end of stream")
4260 || err.contains("failed to fill whole buffer"),
4261 "unexpected error for body_length {body_length}: {err}"
4262 );
4263 }
4264 }
4265
4266 #[test]
4270 #[cfg_attr(miri, ignore)] fn test_stream_reader_rejects_truncated_body() {
4272 let over_prealloc = MAX_PREALLOC_BYTES as i64 + 1;
4273 for (body_length, body_bytes) in [(1024, 10), (over_prealloc, MAX_PREALLOC_BYTES)] {
4274 let stream = stream_with_declared_body(body_length, body_bytes);
4275 let err = StreamReader::try_new(std::io::Cursor::new(stream), None)
4276 .expect_err("a body backed by fewer bytes than declared must not be accepted");
4277 assert!(
4278 err.to_string().contains("failed to fill whole buffer"),
4279 "unexpected error for body_length {body_length}: {err}"
4280 );
4281 }
4282 }
4283
4284 #[test]
4286 fn test_stream_reader_rejects_unbacked_metadata_length() {
4287 let mut stream = Vec::new();
4288 stream.extend_from_slice(&CONTINUATION_MARKER);
4289 stream.extend_from_slice(&i32::MAX.to_le_bytes());
4290 stream.extend_from_slice(b"not this many bytes");
4291
4292 let err = StreamReader::try_new(std::io::Cursor::new(stream), None)
4293 .expect_err("metadata length must be backed by the stream");
4294 assert!(
4295 err.to_string().contains("Unexpected end of stream"),
4296 "unexpected error: {err}"
4297 );
4298 }
4299}