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 record_batch = RecordBatchDecoder::try_new(
927 buf,
928 batch.data().unwrap(),
929 Arc::new(schema),
930 dictionaries_by_id,
931 metadata,
932 )?
933 .with_require_alignment(require_alignment)
934 .with_skip_validation(skip_validation)
935 .read_record_batch()?;
936
937 Some(record_batch.column(0).clone())
938 }
939 _ => None,
940 }
941 .ok_or_else(|| {
942 ArrowError::InvalidArgumentError(format!("dictionary id {id} not found in schema"))
943 })?;
944
945 Ok(dictionary_values)
946}
947
948fn read_block<R: Read + Seek>(mut reader: R, block: &Block) -> Result<Buffer, ArrowError> {
950 reader.seek(SeekFrom::Start(block.offset() as u64))?;
951 let body_len = block.bodyLength().to_usize().unwrap();
952 let metadata_len = block.metaDataLength().to_usize().unwrap();
953 let total_len = body_len.checked_add(metadata_len).unwrap();
954
955 let mut buf = MutableBuffer::try_from_len_zeroed(total_len)
956 .map_err(|e| ArrowError::MemoryError(e.to_string()))?;
957 reader.read_exact(&mut buf)?;
958 Ok(buf.into())
959}
960
961fn check_variadic_counts_consumed(variadic_counts: &VecDeque<i64>) -> Result<(), ArrowError> {
966 if variadic_counts.is_empty() {
967 Ok(())
968 } else {
969 Err(ArrowError::IpcError(format!(
970 "Mismatch between schema and data: the IPC message declares {} more variadic \
971 buffer count(s) than the schema has BinaryView or Utf8View columns",
972 variadic_counts.len()
973 )))
974 }
975}
976
977fn parse_message(buf: &[u8]) -> Result<Message::Message<'_>, ArrowError> {
981 let buf = match buf[..4] == CONTINUATION_MARKER {
982 true => &buf[8..],
983 false => &buf[4..],
984 };
985 crate::root_as_message(buf)
986 .map_err(|err| ArrowError::ParseError(format!("Unable to get root as message: {err:?}")))
987}
988
989pub fn read_footer_length(buf: [u8; 10]) -> Result<usize, ArrowError> {
993 if buf[4..] != super::ARROW_MAGIC {
994 return Err(ArrowError::ParseError(
995 "Arrow file does not contain correct footer".to_string(),
996 ));
997 }
998
999 let footer_len = i32::from_le_bytes([buf[0], buf[1], buf[2], buf[3]]);
1001 footer_len
1002 .try_into()
1003 .map_err(|_| ArrowError::ParseError(format!("Invalid footer length: {footer_len}")))
1004}
1005
1006#[derive(Debug)]
1071pub struct FileDecoder {
1072 schema: SchemaRef,
1073 dictionaries: HashMap<i64, ArrayRef>,
1074 version: MetadataVersion,
1075 projection: Option<Vec<usize>>,
1076 require_alignment: bool,
1077 skip_validation: UnsafeFlag,
1078}
1079
1080impl FileDecoder {
1081 pub fn new(schema: SchemaRef, version: MetadataVersion) -> Self {
1083 Self {
1084 schema,
1085 version,
1086 dictionaries: Default::default(),
1087 projection: None,
1088 require_alignment: false,
1089 skip_validation: UnsafeFlag::new(),
1090 }
1091 }
1092
1093 pub fn with_projection(mut self, projection: Vec<usize>) -> Self {
1095 self.projection = Some(projection);
1096 self
1097 }
1098
1099 pub fn with_require_alignment(mut self, require_alignment: bool) -> Self {
1112 self.require_alignment = require_alignment;
1113 self
1114 }
1115
1116 pub unsafe fn with_skip_validation(mut self, skip_validation: bool) -> Self {
1127 unsafe { self.skip_validation.set(skip_validation) };
1128 self
1129 }
1130
1131 fn read_message<'a>(&self, buf: &'a [u8]) -> Result<Message::Message<'a>, ArrowError> {
1132 let message = parse_message(buf)?;
1133
1134 if self.version != MetadataVersion::V1 && message.version() != self.version {
1136 return Err(ArrowError::IpcError(
1137 "Could not read IPC message as metadata versions mismatch".to_string(),
1138 ));
1139 }
1140 Ok(message)
1141 }
1142
1143 pub fn read_dictionary(&mut self, block: &Block, buf: &Buffer) -> Result<(), ArrowError> {
1145 let message = self.read_message(buf)?;
1146 match message.header_type() {
1147 crate::MessageHeader::DictionaryBatch => {
1148 let batch = message.header_as_dictionary_batch().unwrap();
1149 read_dictionary_impl(
1150 &buf.slice(block.metaDataLength() as _),
1151 batch,
1152 &self.schema,
1153 &mut self.dictionaries,
1154 &message.version(),
1155 self.require_alignment,
1156 self.skip_validation.clone(),
1157 )
1158 }
1159 t => Err(ArrowError::ParseError(format!(
1160 "Expecting DictionaryBatch in dictionary blocks, found {t:?}."
1161 ))),
1162 }
1163 }
1164
1165 pub fn read_record_batch(
1167 &self,
1168 block: &Block,
1169 buf: &Buffer,
1170 ) -> Result<Option<RecordBatch>, ArrowError> {
1171 let message = self.read_message(buf)?;
1172 match message.header_type() {
1173 crate::MessageHeader::Schema => Err(ArrowError::IpcError(
1174 "Not expecting a schema when messages are read".to_string(),
1175 )),
1176 crate::MessageHeader::RecordBatch => {
1177 let batch = message.header_as_record_batch().ok_or_else(|| {
1178 ArrowError::IpcError("Unable to read IPC message as record batch".to_string())
1179 })?;
1180 RecordBatchDecoder::try_new(
1182 &buf.slice(block.metaDataLength() as _),
1183 batch,
1184 self.schema.clone(),
1185 &self.dictionaries,
1186 &message.version(),
1187 )?
1188 .with_projection(self.projection.as_deref())
1189 .with_require_alignment(self.require_alignment)
1190 .with_skip_validation(self.skip_validation.clone())
1191 .read_record_batch()
1192 .map(Some)
1193 }
1194 crate::MessageHeader::NONE => Ok(None),
1195 t => Err(ArrowError::InvalidArgumentError(format!(
1196 "Reading types other than record batches not yet supported, unable to read {t:?}"
1197 ))),
1198 }
1199 }
1200}
1201
1202#[derive(Debug)]
1204pub struct FileReaderBuilder {
1205 projection: Option<Vec<usize>>,
1207 max_footer_fb_tables: usize,
1209 max_footer_fb_depth: usize,
1211}
1212
1213impl Default for FileReaderBuilder {
1214 fn default() -> Self {
1215 let verifier_options = VerifierOptions::default();
1216 Self {
1217 max_footer_fb_tables: verifier_options.max_tables,
1218 max_footer_fb_depth: verifier_options.max_depth,
1219 projection: None,
1220 }
1221 }
1222}
1223
1224impl FileReaderBuilder {
1225 pub fn new() -> Self {
1229 Self::default()
1230 }
1231
1232 pub fn with_projection(mut self, projection: Vec<usize>) -> Self {
1234 self.projection = Some(projection);
1235 self
1236 }
1237
1238 pub fn with_max_footer_fb_tables(mut self, max_footer_fb_tables: usize) -> Self {
1251 self.max_footer_fb_tables = max_footer_fb_tables;
1252 self
1253 }
1254
1255 pub fn with_max_footer_fb_depth(mut self, max_footer_fb_depth: usize) -> Self {
1268 self.max_footer_fb_depth = max_footer_fb_depth;
1269 self
1270 }
1271
1272 pub fn build<R: Read + Seek>(self, mut reader: R) -> Result<FileReader<R>, ArrowError> {
1274 let mut buffer = [0; 10];
1276 reader.seek(SeekFrom::End(-10))?;
1277 reader.read_exact(&mut buffer)?;
1278
1279 let footer_len = read_footer_length(buffer)?;
1280
1281 let mut footer_data = vec![0; footer_len];
1283 reader.seek(SeekFrom::End(-10 - footer_len as i64))?;
1284 reader.read_exact(&mut footer_data)?;
1285
1286 let verifier_options = VerifierOptions {
1287 max_tables: self.max_footer_fb_tables,
1288 max_depth: self.max_footer_fb_depth,
1289 ..Default::default()
1290 };
1291 let footer = crate::root_as_footer_with_opts(&verifier_options, &footer_data[..]).map_err(
1292 |err| ArrowError::ParseError(format!("Unable to get root as footer: {err:?}")),
1293 )?;
1294
1295 let blocks = footer.recordBatches().ok_or_else(|| {
1296 ArrowError::ParseError("Unable to get record batches from IPC Footer".to_string())
1297 })?;
1298
1299 let total_blocks = blocks.len();
1300
1301 let ipc_schema = footer.schema().ok_or_else(|| {
1302 ArrowError::ParseError("Unable to get schema from IPC Footer".to_string())
1303 })?;
1304 if !ipc_schema.endianness().equals_to_target_endianness() {
1305 return Err(ArrowError::IpcError(
1306 "the endianness of the source system does not match the endianness of the target system.".to_owned()
1307 ));
1308 }
1309
1310 let schema = Arc::new(crate::convert::try_fb_to_schema(ipc_schema)?);
1311
1312 let projected_schema = match &self.projection {
1313 Some(projection) => Arc::new(schema.project(projection)?),
1314 None => schema.clone(),
1315 };
1316
1317 let mut custom_metadata = HashMap::new();
1318 if let Some(fb_custom_metadata) = footer.custom_metadata() {
1319 for kv in fb_custom_metadata {
1320 let (Some(key), Some(value)) = (kv.key(), kv.value()) else {
1321 return Err(ArrowError::ParseError(
1322 "Custom metadata in the IPC footer is missing a key or a value".to_string(),
1323 ));
1324 };
1325 custom_metadata.insert(key.to_string(), value.to_string());
1326 }
1327 }
1328
1329 let mut decoder = FileDecoder::new(schema, footer.version());
1330 if let Some(projection) = self.projection {
1331 decoder = decoder.with_projection(projection)
1332 }
1333
1334 if let Some(dictionaries) = footer.dictionaries() {
1336 for block in dictionaries {
1337 let buf = read_block(&mut reader, block)?;
1338 decoder.read_dictionary(block, &buf)?;
1339 }
1340 }
1341
1342 Ok(FileReader {
1343 reader,
1344 blocks: blocks.iter().copied().collect(),
1345 current_block: 0,
1346 total_blocks,
1347 decoder,
1348 schema: projected_schema,
1349 custom_metadata,
1350 })
1351 }
1352}
1353
1354pub struct FileReader<R> {
1399 reader: R,
1401
1402 decoder: FileDecoder,
1404
1405 schema: SchemaRef,
1407
1408 blocks: Vec<Block>,
1412
1413 current_block: usize,
1415
1416 total_blocks: usize,
1418
1419 custom_metadata: HashMap<String, String>,
1421}
1422
1423impl<R> fmt::Debug for FileReader<R> {
1424 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> Result<(), fmt::Error> {
1425 f.debug_struct("FileReader<R>")
1426 .field("decoder", &self.decoder)
1427 .field("blocks", &self.blocks)
1428 .field("current_block", &self.current_block)
1429 .field("total_blocks", &self.total_blocks)
1430 .finish_non_exhaustive()
1431 }
1432}
1433
1434impl<R: Read + Seek> FileReader<BufReader<R>> {
1435 pub fn try_new_buffered(reader: R, projection: Option<Vec<usize>>) -> Result<Self, ArrowError> {
1439 Self::try_new(BufReader::new(reader), projection)
1440 }
1441}
1442
1443impl<R: Read + Seek> FileReader<R> {
1444 pub fn try_new(reader: R, projection: Option<Vec<usize>>) -> Result<Self, ArrowError> {
1456 let builder = FileReaderBuilder {
1457 projection,
1458 ..Default::default()
1459 };
1460 builder.build(reader)
1461 }
1462
1463 pub fn custom_metadata(&self) -> &HashMap<String, String> {
1465 &self.custom_metadata
1466 }
1467
1468 pub fn num_batches(&self) -> usize {
1470 self.total_blocks
1471 }
1472
1473 pub fn schema(&self) -> SchemaRef {
1475 self.schema.clone()
1476 }
1477
1478 pub fn set_index(&mut self, index: usize) -> Result<(), ArrowError> {
1482 if index >= self.total_blocks {
1483 Err(ArrowError::InvalidArgumentError(format!(
1484 "Cannot set batch to index {} from {} total batches",
1485 index, self.total_blocks
1486 )))
1487 } else {
1488 self.current_block = index;
1489 Ok(())
1490 }
1491 }
1492
1493 fn maybe_next(&mut self) -> Result<Option<RecordBatch>, ArrowError> {
1494 let block = &self.blocks[self.current_block];
1495 self.current_block += 1;
1496
1497 let buffer = read_block(&mut self.reader, block)?;
1499 self.decoder.read_record_batch(block, &buffer)
1500 }
1501
1502 pub fn get_ref(&self) -> &R {
1506 &self.reader
1507 }
1508
1509 pub fn get_mut(&mut self) -> &mut R {
1513 &mut self.reader
1514 }
1515
1516 pub unsafe fn with_skip_validation(mut self, skip_validation: bool) -> Self {
1522 self.decoder = unsafe { self.decoder.with_skip_validation(skip_validation) };
1523 self
1524 }
1525}
1526
1527impl<R: Read + Seek> Iterator for FileReader<R> {
1528 type Item = Result<RecordBatch, ArrowError>;
1529
1530 fn next(&mut self) -> Option<Self::Item> {
1531 if self.current_block < self.total_blocks {
1533 self.maybe_next().transpose()
1534 } else {
1535 None
1536 }
1537 }
1538}
1539
1540impl<R: Read + Seek> RecordBatchReader for FileReader<R> {
1541 fn schema(&self) -> SchemaRef {
1542 self.schema()
1543 }
1544}
1545
1546pub struct StreamReader<R> {
1580 reader: MessageReader<R>,
1582
1583 schema: SchemaRef,
1585
1586 dictionaries_by_id: HashMap<i64, ArrayRef>,
1590
1591 finished: bool,
1595
1596 projection: Option<(Vec<usize>, SchemaRef)>,
1598
1599 skip_validation: UnsafeFlag,
1603}
1604
1605impl<R> fmt::Debug for StreamReader<R> {
1606 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> std::result::Result<(), fmt::Error> {
1607 f.debug_struct("StreamReader<R>")
1608 .field("reader", &"R")
1609 .field("schema", &self.schema)
1610 .field("dictionaries_by_id", &self.dictionaries_by_id)
1611 .field("finished", &self.finished)
1612 .field("projection", &self.projection)
1613 .finish()
1614 }
1615}
1616
1617impl<R: Read> StreamReader<BufReader<R>> {
1618 pub fn try_new_buffered(reader: R, projection: Option<Vec<usize>>) -> Result<Self, ArrowError> {
1622 Self::try_new(BufReader::new(reader), projection)
1623 }
1624}
1625
1626impl<R: Read> StreamReader<R> {
1627 pub fn try_new(
1639 reader: R,
1640 projection: Option<Vec<usize>>,
1641 ) -> Result<StreamReader<R>, ArrowError> {
1642 let mut msg_reader = MessageReader::new(reader);
1643 let message = msg_reader.maybe_next()?;
1644 let Some((message, _)) = message else {
1645 return Err(ArrowError::IpcError(
1646 "Expected schema message, found empty stream.".to_string(),
1647 ));
1648 };
1649
1650 if message.header_type() != Message::MessageHeader::Schema {
1651 return Err(ArrowError::IpcError(format!(
1652 "Expected a schema as the first message in the stream, got: {:?}",
1653 message.header_type()
1654 )));
1655 }
1656
1657 let schema = message.header_as_schema().ok_or_else(|| {
1658 ArrowError::ParseError("Failed to parse schema from message header".to_string())
1659 })?;
1660 let schema = crate::convert::try_fb_to_schema(schema)?;
1661
1662 let dictionaries_by_id = HashMap::new();
1664
1665 let projection = match projection {
1666 Some(projection_indices) => {
1667 let schema = Arc::new(schema.project(&projection_indices)?);
1668 Some((projection_indices, schema))
1669 }
1670 _ => None,
1671 };
1672
1673 Ok(Self {
1674 reader: msg_reader,
1675 schema: Arc::new(schema),
1676 finished: false,
1677 dictionaries_by_id,
1678 projection,
1679 skip_validation: UnsafeFlag::new(),
1680 })
1681 }
1682
1683 pub fn schema(&self) -> SchemaRef {
1685 match &self.projection {
1686 Some((_, projected_schema)) => projected_schema.clone(),
1687 None => self.schema.clone(),
1688 }
1689 }
1690
1691 pub fn is_finished(&self) -> bool {
1693 self.finished
1694 }
1695
1696 fn maybe_next(&mut self) -> Result<Option<RecordBatch>, ArrowError> {
1697 if self.finished {
1698 return Ok(None);
1699 }
1700
1701 loop {
1703 let message = self.next_ipc_message()?;
1704 let Some(message) = message else {
1705 self.finished = true;
1707 return Ok(None);
1708 };
1709
1710 match message {
1711 IpcMessage::Schema(_) => {
1712 return Err(ArrowError::IpcError(
1713 "Expected a record batch, but found a schema".to_string(),
1714 ));
1715 }
1716 IpcMessage::RecordBatch(record_batch) => {
1717 return Ok(Some(record_batch));
1718 }
1719 IpcMessage::DictionaryBatch { .. } => {}
1720 }
1721 }
1722 }
1723
1724 pub(crate) fn next_ipc_message(&mut self) -> Result<Option<IpcMessage>, ArrowError> {
1732 let message = self.reader.maybe_next()?;
1733 let Some((message, body)) = message else {
1734 return Ok(None);
1736 };
1737
1738 let ipc_message = match message.header_type() {
1739 Message::MessageHeader::Schema => {
1740 let schema = message.header_as_schema().ok_or_else(|| {
1741 ArrowError::ParseError("Failed to parse schema from message header".to_string())
1742 })?;
1743 let arrow_schema = crate::convert::try_fb_to_schema(schema)?;
1744 IpcMessage::Schema(arrow_schema)
1745 }
1746 Message::MessageHeader::RecordBatch => {
1747 let batch = message.header_as_record_batch().ok_or_else(|| {
1748 ArrowError::IpcError("Unable to read IPC message as record batch".to_string())
1749 })?;
1750
1751 let version = message.version();
1752 let schema = self.schema.clone();
1753 let record_batch = RecordBatchDecoder::try_new(
1754 &body.into(),
1755 batch,
1756 schema,
1757 &self.dictionaries_by_id,
1758 &version,
1759 )?
1760 .with_projection(self.projection.as_ref().map(|x| x.0.as_ref()))
1761 .with_require_alignment(false)
1762 .with_skip_validation(self.skip_validation.clone())
1763 .read_record_batch()?;
1764 IpcMessage::RecordBatch(record_batch)
1765 }
1766 Message::MessageHeader::DictionaryBatch => {
1767 let dict = message.header_as_dictionary_batch().ok_or_else(|| {
1768 ArrowError::ParseError(
1769 "Failed to parse dictionary batch from message header".to_string(),
1770 )
1771 })?;
1772
1773 let version = message.version();
1774 let dict_values = get_dictionary_values(
1775 &body.into(),
1776 dict,
1777 &self.schema,
1778 &self.dictionaries_by_id,
1779 &version,
1780 false,
1781 self.skip_validation.clone(),
1782 )?;
1783
1784 update_dictionaries(
1785 &mut self.dictionaries_by_id,
1786 dict.isDelta(),
1787 dict.id(),
1788 dict_values.clone(),
1789 )?;
1790
1791 IpcMessage::DictionaryBatch {
1792 id: dict.id(),
1793 is_delta: (dict.isDelta()),
1794 values: (dict_values),
1795 }
1796 }
1797 x => {
1798 return Err(ArrowError::ParseError(format!(
1799 "Unsupported message header type in IPC stream: '{x:?}'"
1800 )));
1801 }
1802 };
1803
1804 Ok(Some(ipc_message))
1805 }
1806
1807 pub fn get_ref(&self) -> &R {
1811 self.reader.inner()
1812 }
1813
1814 pub fn get_mut(&mut self) -> &mut R {
1818 self.reader.inner_mut()
1819 }
1820
1821 pub unsafe fn with_skip_validation(mut self, skip_validation: bool) -> Self {
1827 unsafe { self.skip_validation.set(skip_validation) };
1828 self
1829 }
1830}
1831
1832impl<R: Read> Iterator for StreamReader<R> {
1833 type Item = Result<RecordBatch, ArrowError>;
1834
1835 fn next(&mut self) -> Option<Self::Item> {
1836 self.maybe_next().transpose()
1837 }
1838}
1839
1840impl<R: Read> RecordBatchReader for StreamReader<R> {
1841 fn schema(&self) -> SchemaRef {
1842 self.schema()
1843 }
1844}
1845
1846#[derive(Debug)]
1852#[expect(dead_code)]
1853pub(crate) enum IpcMessage {
1854 Schema(arrow_schema::Schema),
1855 RecordBatch(RecordBatch),
1856 DictionaryBatch {
1857 id: i64,
1858 is_delta: bool,
1859 values: ArrayRef,
1860 },
1861}
1862
1863const MAX_PREALLOC_BYTES: usize = 64 * 1024 * 1024;
1879
1880fn read_body_bounded<R: Read>(reader: &mut R, len: usize) -> Result<MutableBuffer, ArrowError> {
1882 let mut buf = MutableBuffer::try_from_len_zeroed(len.min(MAX_PREALLOC_BYTES))
1883 .map_err(|e| ArrowError::MemoryError(e.to_string()))?;
1884 let mut filled = 0;
1885 while filled < len {
1886 let target = buf.len();
1887 reader.read_exact(&mut buf.as_slice_mut()[filled..target])?;
1888 filled = target;
1889 if filled < len {
1890 buf.try_resize(len.min(target.saturating_mul(2)), 0)
1891 .map_err(|e| ArrowError::MemoryError(e.to_string()))?;
1892 }
1893 }
1894 Ok(buf)
1895}
1896
1897struct MessageReader<R> {
1900 reader: R,
1901 buf: Vec<u8>,
1902}
1903
1904impl<R: Read> MessageReader<R> {
1905 fn new(reader: R) -> Self {
1906 Self {
1907 reader,
1908 buf: Vec::new(),
1909 }
1910 }
1911
1912 fn maybe_next(&mut self) -> Result<Option<(Message::Message<'_>, MutableBuffer)>, ArrowError> {
1923 let meta_len = self.read_meta_len()?;
1924 let Some(meta_len) = meta_len else {
1925 return Ok(None);
1926 };
1927
1928 self.buf.clear();
1931 let read = (&mut self.reader)
1932 .take(meta_len as u64)
1933 .read_to_end(&mut self.buf)?;
1934 if read != meta_len {
1935 return Err(ArrowError::ParseError(format!(
1936 "Unexpected end of stream: expected {meta_len} metadata bytes, got {read}"
1937 )));
1938 }
1939
1940 let message = crate::root_as_message(self.buf.as_slice()).map_err(|err| {
1941 ArrowError::ParseError(format!("Unable to get root as message: {err:?}"))
1942 })?;
1943
1944 let body_len = usize::try_from(message.bodyLength()).map_err(|_| {
1945 ArrowError::ParseError(format!(
1946 "Invalid IPC message body length: {}",
1947 message.bodyLength()
1948 ))
1949 })?;
1950 let buf = read_body_bounded(&mut self.reader, body_len)?;
1951
1952 Ok(Some((message, buf)))
1953 }
1954
1955 fn inner_mut(&mut self) -> &mut R {
1957 &mut self.reader
1958 }
1959
1960 fn inner(&self) -> &R {
1962 &self.reader
1963 }
1964
1965 pub fn read_meta_len(&mut self) -> Result<Option<usize>, ArrowError> {
1974 let mut meta_len: [u8; 4] = [0; 4];
1975 match self.reader.read_exact(&mut meta_len) {
1976 Ok(()) => {}
1977 Err(e) => {
1978 return if e.kind() == std::io::ErrorKind::UnexpectedEof {
1979 Ok(None)
1983 } else {
1984 Err(ArrowError::from(e))
1985 };
1986 }
1987 }
1988
1989 let meta_len = {
1990 if meta_len == CONTINUATION_MARKER {
1993 self.reader.read_exact(&mut meta_len)?;
1994 }
1995
1996 i32::from_le_bytes(meta_len)
1997 };
1998
1999 if meta_len == 0 {
2000 return Ok(None);
2001 }
2002
2003 let meta_len = usize::try_from(meta_len)
2004 .map_err(|_| ArrowError::ParseError(format!("Invalid metadata length: {meta_len}")))?;
2005
2006 Ok(Some(meta_len))
2007 }
2008}
2009
2010#[cfg(test)]
2011mod tests {
2012 use std::io::Cursor;
2013
2014 use crate::convert::try_fb_to_schema;
2015 use crate::writer::{
2016 DictionaryTracker, IpcDataGenerator, IpcWriteOptions, unslice_run_array, write_message,
2017 };
2018
2019 use super::*;
2020
2021 use crate::{root_as_footer, root_as_message, size_prefixed_root_as_message};
2022 use arrow_array::builder::{PrimitiveRunBuilder, UnionBuilder};
2023 use arrow_array::types::*;
2024 use arrow_buffer::{NullBuffer, OffsetBuffer};
2025 use arrow_data::ArrayDataBuilder;
2026
2027 fn create_test_projection_schema() -> Schema {
2028 let list_data_type = DataType::List(Arc::new(Field::new_list_field(DataType::Int32, true)));
2030
2031 let fixed_size_list_data_type =
2032 DataType::FixedSizeList(Arc::new(Field::new_list_field(DataType::Int32, false)), 3);
2033
2034 let union_fields = UnionFields::from_fields(vec![
2035 Field::new("a", DataType::Int32, false),
2036 Field::new("b", DataType::Float64, false),
2037 ]);
2038
2039 let union_data_type = DataType::Union(union_fields, UnionMode::Dense);
2040
2041 let struct_fields = Fields::from(vec![
2042 Field::new("id", DataType::Int32, false),
2043 Field::new_list("list", Field::new_list_field(DataType::Int8, true), false),
2044 ]);
2045 let struct_data_type = DataType::Struct(struct_fields);
2046
2047 let run_encoded_data_type = DataType::RunEndEncoded(
2048 Arc::new(Field::new("run_ends", DataType::Int16, false)),
2049 Arc::new(Field::new("values", DataType::Int32, true)),
2050 );
2051
2052 Schema::new(vec![
2054 Field::new("f0", DataType::UInt32, false),
2055 Field::new("f1", DataType::Utf8, false),
2056 Field::new("f2", DataType::Boolean, false),
2057 Field::new("f3", union_data_type, true),
2058 Field::new("f4", DataType::Null, true),
2059 Field::new("f5", DataType::Float64, true),
2060 Field::new("f6", list_data_type, false),
2061 Field::new("f7", DataType::FixedSizeBinary(3), true),
2062 Field::new("f8", fixed_size_list_data_type, false),
2063 Field::new("f9", struct_data_type, false),
2064 Field::new("f10", run_encoded_data_type, false),
2065 Field::new("f11", DataType::Boolean, false),
2066 Field::new_dictionary("f12", DataType::Int8, DataType::Utf8, false),
2067 Field::new("f13", DataType::Utf8, false),
2068 ])
2069 }
2070
2071 fn create_test_projection_batch_data(schema: &Schema) -> RecordBatch {
2072 let array0 = UInt32Array::from(vec![1, 2, 3]);
2074 let array1 = StringArray::from(vec!["foo", "bar", "baz"]);
2075 let array2 = BooleanArray::from(vec![true, false, true]);
2076
2077 let mut union_builder = UnionBuilder::new_dense();
2078 union_builder.append::<Int32Type>("a", 1).unwrap();
2079 union_builder.append::<Float64Type>("b", 10.1).unwrap();
2080 union_builder.append_null::<Float64Type>("b").unwrap();
2081 let array3 = union_builder.build().unwrap();
2082
2083 let array4 = NullArray::new(3);
2084 let array5 = Float64Array::from(vec![Some(1.1), None, Some(3.3)]);
2085 let array6_values = vec![
2086 Some(vec![Some(10), Some(10), Some(10)]),
2087 Some(vec![Some(20), Some(20), Some(20)]),
2088 Some(vec![Some(30), Some(30)]),
2089 ];
2090 let array6 = ListArray::from_iter_primitive::<Int32Type, _, _>(array6_values);
2091 let array7_values = vec![vec![11, 12, 13], vec![22, 23, 24], vec![33, 34, 35]];
2092 let array7 = FixedSizeBinaryArray::try_from_iter(array7_values.into_iter()).unwrap();
2093
2094 let array8_values = ArrayData::builder(DataType::Int32)
2095 .len(9)
2096 .add_buffer(Buffer::from_slice_ref([40, 41, 42, 43, 44, 45, 46, 47, 48]))
2097 .build()
2098 .unwrap();
2099 let array8_data = ArrayData::builder(schema.field(8).data_type().clone())
2100 .len(3)
2101 .add_child_data(array8_values)
2102 .build()
2103 .unwrap();
2104 let array8 = FixedSizeListArray::from(array8_data);
2105
2106 let array9_id: ArrayRef = Arc::new(Int32Array::from(vec![1001, 1002, 1003]));
2107 let array9_list: ArrayRef =
2108 Arc::new(ListArray::from_iter_primitive::<Int8Type, _, _>(vec![
2109 Some(vec![Some(-10)]),
2110 Some(vec![Some(-20), Some(-20), Some(-20)]),
2111 Some(vec![Some(-30)]),
2112 ]));
2113 let array9 = ArrayDataBuilder::new(schema.field(9).data_type().clone())
2114 .add_child_data(array9_id.into_data())
2115 .add_child_data(array9_list.into_data())
2116 .len(3)
2117 .build()
2118 .unwrap();
2119 let array9 = StructArray::from(array9);
2120
2121 let array10_input = vec![Some(1_i32), None, None];
2122 let mut array10_builder = PrimitiveRunBuilder::<Int16Type, Int32Type>::new();
2123 array10_builder.extend(array10_input);
2124 let array10 = array10_builder.finish();
2125
2126 let array11 = BooleanArray::from(vec![false, false, true]);
2127
2128 let array12_values = StringArray::from(vec!["x", "yy", "zzz"]);
2129 let array12_keys = Int8Array::from_iter_values([1, 1, 2]);
2130 let array12 = DictionaryArray::new(array12_keys, Arc::new(array12_values));
2131
2132 let array13 = StringArray::from(vec!["a", "bb", "ccc"]);
2133
2134 RecordBatch::try_new(
2136 Arc::new(schema.clone()),
2137 vec![
2138 Arc::new(array0),
2139 Arc::new(array1),
2140 Arc::new(array2),
2141 Arc::new(array3),
2142 Arc::new(array4),
2143 Arc::new(array5),
2144 Arc::new(array6),
2145 Arc::new(array7),
2146 Arc::new(array8),
2147 Arc::new(array9),
2148 Arc::new(array10),
2149 Arc::new(array11),
2150 Arc::new(array12),
2151 Arc::new(array13),
2152 ],
2153 )
2154 .unwrap()
2155 }
2156
2157 #[test]
2158 fn test_negative_meta_len_start_stream() {
2159 let bytes = i32::to_le_bytes(-1);
2160 let mut buf = vec![];
2161 buf.extend(CONTINUATION_MARKER);
2162 buf.extend(bytes);
2163
2164 let reader_err = StreamReader::try_new(Cursor::new(buf), None).err();
2165 assert!(reader_err.is_some());
2166 assert_eq!(
2167 reader_err.unwrap().to_string(),
2168 "Parser error: Invalid metadata length: -1"
2169 );
2170 }
2171
2172 #[test]
2173 fn test_negative_meta_len_mid_stream() {
2174 let schema = Schema::new(vec![Field::new("a", DataType::Int32, false)]);
2175 let mut buf = Vec::new();
2176 {
2177 let mut writer = crate::writer::StreamWriter::try_new(&mut buf, &schema).unwrap();
2178 let batch =
2179 RecordBatch::try_new(Arc::new(schema), vec![Arc::new(Int32Array::from(vec![1]))])
2180 .unwrap();
2181 writer.write(&batch).unwrap();
2182 }
2183
2184 let bytes = i32::to_le_bytes(-1);
2185 buf.extend(CONTINUATION_MARKER);
2186 buf.extend(bytes);
2187
2188 let mut reader = StreamReader::try_new(Cursor::new(buf), None).unwrap();
2189 assert!(reader.maybe_next().is_ok());
2191 let batch_err = reader.maybe_next().err();
2193 assert!(batch_err.is_some());
2194 assert_eq!(
2195 batch_err.unwrap().to_string(),
2196 "Parser error: Invalid metadata length: -1"
2197 );
2198 }
2199
2200 #[test]
2201 fn test_missing_buffer_metadata_error() {
2202 use crate::r#gen::Message::*;
2203 use flatbuffers::FlatBufferBuilder;
2204
2205 let schema = Arc::new(Schema::new(vec![Field::new("col", DataType::Int32, true)]));
2206
2207 let mut fbb = FlatBufferBuilder::new();
2210 let nodes = fbb.create_vector(&[FieldNode::new(2, 0)]);
2211 let buffers = fbb.create_vector(&[crate::Buffer::new(0, 8)]);
2212 let batch_offset = RecordBatch::create(
2213 &mut fbb,
2214 &RecordBatchArgs {
2215 length: 2,
2216 nodes: Some(nodes),
2217 buffers: Some(buffers),
2218 compression: None,
2219 variadicBufferCounts: None,
2220 },
2221 );
2222 fbb.finish_minimal(batch_offset);
2223 let batch_bytes = fbb.finished_data().to_vec();
2224 let batch = flatbuffers::root::<RecordBatch>(&batch_bytes).unwrap();
2225
2226 let data_buffer = Buffer::from(vec![0u8; 8]);
2227 let dictionaries: HashMap<i64, ArrayRef> = HashMap::new();
2228 let metadata = MetadataVersion::V5;
2229
2230 let decoder = RecordBatchDecoder::try_new(
2231 &data_buffer,
2232 batch,
2233 schema.clone(),
2234 &dictionaries,
2235 &metadata,
2236 )
2237 .unwrap();
2238
2239 let result = decoder.read_record_batch();
2240
2241 match result {
2242 Err(ArrowError::IpcError(msg)) => {
2243 assert_eq!(msg, "Buffer count mismatched with metadata");
2244 }
2245 other => panic!("unexpected error: {other:?}"),
2246 }
2247 }
2248
2249 fn read_batch_with_variadic_count(count: i64, projection: Option<&[usize]>) -> ArrowError {
2252 use crate::r#gen::Message::*;
2253 use flatbuffers::FlatBufferBuilder;
2254
2255 let schema = Arc::new(Schema::new(vec![Field::new(
2256 "col",
2257 DataType::Utf8View,
2258 true,
2259 )]));
2260
2261 let mut fbb = FlatBufferBuilder::new();
2262 let nodes = fbb.create_vector(&[FieldNode::new(1, 0)]);
2263 let buffers = fbb.create_vector(&[crate::Buffer::new(0, 8), crate::Buffer::new(8, 8)]);
2264 let variadic_buffer_counts = fbb.create_vector(&[count]);
2265 let batch_offset = RecordBatch::create(
2266 &mut fbb,
2267 &RecordBatchArgs {
2268 length: 1,
2269 nodes: Some(nodes),
2270 buffers: Some(buffers),
2271 compression: None,
2272 variadicBufferCounts: Some(variadic_buffer_counts),
2273 },
2274 );
2275 fbb.finish_minimal(batch_offset);
2276 let batch_bytes = fbb.finished_data().to_vec();
2277 let batch = flatbuffers::root::<RecordBatch>(&batch_bytes).unwrap();
2278
2279 let data_buffer = Buffer::from(vec![0u8; 16]);
2280 let dictionaries: HashMap<i64, ArrayRef> = HashMap::new();
2281
2282 RecordBatchDecoder::try_new(
2283 &data_buffer,
2284 batch,
2285 schema,
2286 &dictionaries,
2287 &MetadataVersion::V5,
2288 )
2289 .unwrap()
2290 .with_projection(projection)
2291 .read_record_batch()
2292 .expect_err("should get error")
2293 }
2294
2295 #[test]
2298 fn test_invalid_variadic_buffer_count_error() {
2299 for count in [-2, -1, 1, i64::MAX] {
2302 for projection in [None, Some([].as_slice())] {
2303 let err = read_batch_with_variadic_count(count, projection);
2304 assert_eq!(
2305 err.to_string(),
2306 format!(
2307 "Ipc error: Invalid variadic count {count} for Utf8View column, \
2308 with 2 buffer(s) left in the message"
2309 ),
2310 "count {count}, projection {projection:?}"
2311 );
2312 }
2313 }
2314 }
2315
2316 #[test]
2318 fn test_valid_variadic_buffer_count_is_accepted() {
2319 let err = read_batch_with_variadic_count(0, None);
2320 assert!(!err.to_string().contains("Invalid variadic count"), "{err}");
2321 }
2322
2323 #[test]
2324 fn test_missing_footer_schema_error() {
2325 use crate::r#gen::File::{Footer, FooterArgs};
2326 use flatbuffers::FlatBufferBuilder;
2327
2328 let mut fbb = FlatBufferBuilder::new();
2331 let record_batches = fbb.create_vector::<Block>(&[]);
2332 let footer = Footer::create(
2333 &mut fbb,
2334 &FooterArgs {
2335 version: MetadataVersion::V5,
2336 schema: None,
2337 dictionaries: None,
2338 recordBatches: Some(record_batches),
2339 custom_metadata: None,
2340 },
2341 );
2342 fbb.finish(footer, None);
2343 let footer_data = fbb.finished_data();
2344
2345 let mut buf = Vec::new();
2347 buf.extend_from_slice(&crate::ARROW_MAGIC);
2348 buf.extend_from_slice(footer_data);
2349 buf.extend_from_slice(&(footer_data.len() as i32).to_le_bytes());
2350 buf.extend_from_slice(&crate::ARROW_MAGIC);
2351
2352 let err = FileReader::try_new(Cursor::new(buf), None)
2353 .expect_err("expected an error, not a panic");
2354 assert!(
2355 matches!(err, ArrowError::ParseError(_)),
2356 "expected ParseError, got {err:?}"
2357 );
2358 }
2359
2360 #[test]
2362 fn test_read_legacy_empty_list_without_offsets_buffer() {
2363 use crate::r#gen::Message::*;
2364 use flatbuffers::FlatBufferBuilder;
2365
2366 let schema = Arc::new(Schema::new(vec![Field::new_list(
2367 "items",
2368 Field::new_list_field(DataType::Int32, true),
2369 true,
2370 )]));
2371
2372 let mut fbb = FlatBufferBuilder::new();
2375 let nodes = fbb.create_vector(&[
2376 FieldNode::new(0, 0), FieldNode::new(0, 0), ]);
2379 let buffers = fbb.create_vector(&[
2380 crate::Buffer::new(0, 0), crate::Buffer::new(0, 0), crate::Buffer::new(0, 0), crate::Buffer::new(0, 0), ]);
2385 let batch_offset = RecordBatch::create(
2386 &mut fbb,
2387 &RecordBatchArgs {
2388 length: 0,
2389 nodes: Some(nodes),
2390 buffers: Some(buffers),
2391 compression: None,
2392 variadicBufferCounts: None,
2393 },
2394 );
2395 fbb.finish_minimal(batch_offset);
2396 let batch_bytes = fbb.finished_data().to_vec();
2397 let batch = flatbuffers::root::<RecordBatch>(&batch_bytes).unwrap();
2398
2399 let body = Buffer::from(Vec::<u8>::new());
2400 let dictionaries: HashMap<i64, ArrayRef> = HashMap::new();
2401 let metadata = MetadataVersion::V5;
2402
2403 let decoder =
2404 RecordBatchDecoder::try_new(&body, batch, schema.clone(), &dictionaries, &metadata)
2405 .unwrap();
2406
2407 let read_batch = decoder.read_record_batch().unwrap();
2408 assert_eq!(read_batch.num_rows(), 0);
2409
2410 let list = read_batch
2411 .column(0)
2412 .as_any()
2413 .downcast_ref::<ListArray>()
2414 .unwrap();
2415 assert_eq!(list.len(), 0);
2416 assert_eq!(list.values().len(), 0);
2417 }
2418
2419 #[test]
2421 fn test_read_legacy_empty_utf8_and_binary_without_offsets_buffer() {
2422 use crate::r#gen::Message::*;
2423 use flatbuffers::FlatBufferBuilder;
2424
2425 let schema = Arc::new(Schema::new(vec![
2426 Field::new("name", DataType::Utf8, true),
2427 Field::new("payload", DataType::Binary, true),
2428 ]));
2429
2430 let mut fbb = FlatBufferBuilder::new();
2433 let nodes = fbb.create_vector(&[
2434 FieldNode::new(0, 0), FieldNode::new(0, 0), ]);
2437 let buffers = fbb.create_vector(&[
2438 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), ]);
2445 let batch_offset = RecordBatch::create(
2446 &mut fbb,
2447 &RecordBatchArgs {
2448 length: 0,
2449 nodes: Some(nodes),
2450 buffers: Some(buffers),
2451 compression: None,
2452 variadicBufferCounts: None,
2453 },
2454 );
2455 fbb.finish_minimal(batch_offset);
2456 let batch_bytes = fbb.finished_data().to_vec();
2457 let batch = flatbuffers::root::<RecordBatch>(&batch_bytes).unwrap();
2458
2459 let body = Buffer::from(Vec::<u8>::new());
2460 let dictionaries: HashMap<i64, ArrayRef> = HashMap::new();
2461 let metadata = MetadataVersion::V5;
2462
2463 let decoder =
2464 RecordBatchDecoder::try_new(&body, batch, schema.clone(), &dictionaries, &metadata)
2465 .unwrap();
2466
2467 let read_batch = decoder.read_record_batch().unwrap();
2468 assert_eq!(read_batch.num_rows(), 0);
2469
2470 let utf8 = read_batch
2471 .column(0)
2472 .as_any()
2473 .downcast_ref::<StringArray>()
2474 .unwrap();
2475 assert_eq!(utf8.len(), 0);
2476 assert_eq!(utf8.value_offsets(), [0]);
2477
2478 let binary = read_batch
2479 .column(1)
2480 .as_any()
2481 .downcast_ref::<BinaryArray>()
2482 .unwrap();
2483 assert_eq!(binary.len(), 0);
2484 assert_eq!(binary.value_offsets(), [0]);
2485 }
2486
2487 #[test]
2488 fn test_projection_array_values() {
2489 let schema = create_test_projection_schema();
2491
2492 let batch = create_test_projection_batch_data(&schema);
2494
2495 let mut buf = Vec::new();
2497 {
2498 let mut writer = crate::writer::FileWriter::try_new(&mut buf, &schema).unwrap();
2499 writer.write(&batch).unwrap();
2500 writer.finish().unwrap();
2501 }
2502
2503 for index in 0..12 {
2505 let projection = vec![index];
2506 let reader = FileReader::try_new(std::io::Cursor::new(buf.clone()), Some(projection));
2507 let read_batch = reader.unwrap().next().unwrap().unwrap();
2508 let projected_column = read_batch.column(0);
2509 let expected_column = batch.column(index);
2510
2511 assert_eq!(projected_column.as_ref(), expected_column.as_ref());
2513 }
2514
2515 {
2516 let reader =
2518 FileReader::try_new(std::io::Cursor::new(buf.clone()), Some(vec![3, 2, 1]));
2519 let read_batch = reader.unwrap().next().unwrap().unwrap();
2520 let expected_batch = batch.project(&[3, 2, 1]).unwrap();
2521 assert_eq!(read_batch, expected_batch);
2522 }
2523 }
2524
2525 #[test]
2526 fn test_file_reader_projected_schema_matches_batch_schema() {
2527 let schema = create_test_projection_schema();
2528 let batch = create_test_projection_batch_data(&schema);
2529
2530 let mut buf = Vec::new();
2531 {
2532 let mut writer = crate::writer::FileWriter::try_new(&mut buf, &schema).unwrap();
2533 writer.write(&batch).unwrap();
2534 writer.finish().unwrap();
2535 }
2536
2537 let projection = vec![3, 2, 1];
2538 let mut reader = FileReader::try_new(Cursor::new(buf), Some(projection)).unwrap();
2539 let reader_schema = RecordBatchReader::schema(&reader);
2540 let read_batch = reader.next().unwrap().unwrap();
2541
2542 assert_eq!(reader_schema, read_batch.schema());
2543 }
2544
2545 #[test]
2546 fn test_stream_reader_projected_schema_matches_batch_schema() {
2547 let schema = create_test_projection_schema();
2548 let batch = create_test_projection_batch_data(&schema);
2549
2550 let mut buf = Vec::new();
2551 {
2552 let mut writer = crate::writer::StreamWriter::try_new(&mut buf, &schema).unwrap();
2553 writer.write(&batch).unwrap();
2554 writer.finish().unwrap();
2555 }
2556
2557 let projection = vec![3, 2, 1];
2558 let mut reader = StreamReader::try_new(Cursor::new(buf), Some(projection)).unwrap();
2559 let reader_schema = RecordBatchReader::schema(&reader);
2560 let read_batch = reader.next().unwrap().unwrap();
2561
2562 assert_eq!(reader_schema, read_batch.schema());
2563 }
2564
2565 #[test]
2566 fn test_file_reader_rejects_invalid_projection() {
2567 let schema = create_test_projection_schema();
2568 let batch = create_test_projection_batch_data(&schema);
2569
2570 let mut buf = Vec::new();
2571 {
2572 let mut writer = crate::writer::FileWriter::try_new(&mut buf, &schema).unwrap();
2573 writer.write(&batch).unwrap();
2574 writer.finish().unwrap();
2575 }
2576
2577 let result = FileReader::try_new(Cursor::new(buf), Some(vec![schema.fields().len()]));
2578
2579 assert!(matches!(result, Err(ArrowError::SchemaError(_))));
2580 }
2581
2582 #[test]
2583 fn test_projection_duplicate_indices() {
2584 let schema = create_test_projection_schema();
2585 let batch = create_test_projection_batch_data(&schema);
2586
2587 let mut buf = Vec::new();
2589 {
2590 let mut writer = crate::writer::FileWriter::try_new(&mut buf, &schema).unwrap();
2591 writer.write(&batch).unwrap();
2592 writer.finish().unwrap();
2593 }
2594
2595 for projection in [vec![1, 1], vec![2, 0, 2]] {
2597 let reader =
2598 FileReader::try_new(std::io::Cursor::new(buf.clone()), Some(projection.clone()));
2599 let read_batch = reader.unwrap().next().unwrap().unwrap();
2600
2601 let expected_batch = batch.project(&projection).unwrap();
2602 assert_eq!(read_batch, expected_batch);
2603 }
2604 }
2605
2606 #[test]
2607 fn test_arrow_single_float_row() {
2608 let schema = Schema::new(vec![
2609 Field::new("a", DataType::Float32, false),
2610 Field::new("b", DataType::Float32, false),
2611 Field::new("c", DataType::Int32, false),
2612 Field::new("d", DataType::Int32, false),
2613 ]);
2614 let arrays = vec![
2615 Arc::new(Float32Array::from(vec![1.23])) as ArrayRef,
2616 Arc::new(Float32Array::from(vec![-6.50])) as ArrayRef,
2617 Arc::new(Int32Array::from(vec![2])) as ArrayRef,
2618 Arc::new(Int32Array::from(vec![1])) as ArrayRef,
2619 ];
2620 let batch = RecordBatch::try_new(Arc::new(schema.clone()), arrays).unwrap();
2621 let mut file = tempfile::tempfile().unwrap();
2623 let mut stream_writer = crate::writer::StreamWriter::try_new(&mut file, &schema).unwrap();
2624 stream_writer.write(&batch).unwrap();
2625 stream_writer.finish().unwrap();
2626
2627 drop(stream_writer);
2628
2629 file.rewind().unwrap();
2630
2631 let reader = StreamReader::try_new(&mut file, None).unwrap();
2633
2634 reader.for_each(|batch| {
2635 let batch = batch.unwrap();
2636 assert!(
2637 batch
2638 .column(0)
2639 .as_any()
2640 .downcast_ref::<Float32Array>()
2641 .unwrap()
2642 .value(0)
2643 != 0.0
2644 );
2645 assert!(
2646 batch
2647 .column(1)
2648 .as_any()
2649 .downcast_ref::<Float32Array>()
2650 .unwrap()
2651 .value(0)
2652 != 0.0
2653 );
2654 });
2655
2656 file.rewind().unwrap();
2657
2658 let reader = StreamReader::try_new(file, Some(vec![0, 3])).unwrap();
2660
2661 reader.for_each(|batch| {
2662 let batch = batch.unwrap();
2663 assert_eq!(batch.schema().fields().len(), 2);
2664 assert_eq!(batch.schema().fields()[0].data_type(), &DataType::Float32);
2665 assert_eq!(batch.schema().fields()[1].data_type(), &DataType::Int32);
2666 });
2667 }
2668
2669 fn write_ipc(rb: &RecordBatch) -> Vec<u8> {
2671 let mut buf = Vec::new();
2672 let mut writer = crate::writer::FileWriter::try_new(&mut buf, rb.schema_ref()).unwrap();
2673 writer.write(rb).unwrap();
2674 writer.finish().unwrap();
2675 buf
2676 }
2677
2678 fn read_ipc(buf: &[u8]) -> Result<RecordBatch, ArrowError> {
2680 let mut reader = FileReader::try_new(std::io::Cursor::new(buf), None)?;
2681 reader.next().unwrap()
2682 }
2683
2684 fn read_ipc_skip_validation(buf: &[u8]) -> Result<RecordBatch, ArrowError> {
2687 let mut reader = unsafe {
2688 FileReader::try_new(std::io::Cursor::new(buf), None)?.with_skip_validation(true)
2689 };
2690 reader.next().unwrap()
2691 }
2692
2693 fn roundtrip_ipc(rb: &RecordBatch) -> RecordBatch {
2694 let buf = write_ipc(rb);
2695 read_ipc(&buf).unwrap()
2696 }
2697
2698 fn read_ipc_with_decoder(buf: Vec<u8>) -> Result<RecordBatch, ArrowError> {
2701 read_ipc_with_decoder_inner(buf, false)
2702 }
2703
2704 fn read_ipc_with_decoder_skip_validation(buf: Vec<u8>) -> Result<RecordBatch, ArrowError> {
2707 read_ipc_with_decoder_inner(buf, true)
2708 }
2709
2710 fn read_ipc_with_decoder_inner(
2711 buf: Vec<u8>,
2712 skip_validation: bool,
2713 ) -> Result<RecordBatch, ArrowError> {
2714 let buffer = Buffer::from_vec(buf);
2715 let trailer_start = buffer.len() - 10;
2716 let footer_len = read_footer_length(buffer[trailer_start..].try_into().unwrap())?;
2717 let footer = root_as_footer(&buffer[trailer_start - footer_len..trailer_start])
2718 .map_err(|e| ArrowError::InvalidArgumentError(format!("Invalid footer: {e}")))?;
2719
2720 let schema = try_fb_to_schema(footer.schema().unwrap()).unwrap();
2721
2722 let mut decoder = unsafe {
2723 FileDecoder::new(Arc::new(schema), footer.version())
2724 .with_skip_validation(skip_validation)
2725 };
2726 for block in footer.dictionaries().iter().flatten() {
2728 let block_len = block.bodyLength() as usize + block.metaDataLength() as usize;
2729 let data = buffer.slice_with_length(block.offset() as _, block_len);
2730 decoder.read_dictionary(block, &data)?
2731 }
2732
2733 let batches = footer.recordBatches().unwrap();
2735 assert_eq!(batches.len(), 1); let block = batches.get(0);
2738 let block_len = block.bodyLength() as usize + block.metaDataLength() as usize;
2739 let data = buffer.slice_with_length(block.offset() as _, block_len);
2740 Ok(decoder.read_record_batch(block, &data)?.unwrap())
2741 }
2742
2743 fn write_stream(rb: &RecordBatch) -> Vec<u8> {
2745 let mut buf = Vec::new();
2746 let mut writer = crate::writer::StreamWriter::try_new(&mut buf, rb.schema_ref()).unwrap();
2747 writer.write(rb).unwrap();
2748 writer.finish().unwrap();
2749 buf
2750 }
2751
2752 fn read_stream(buf: &[u8]) -> Result<RecordBatch, ArrowError> {
2754 let mut reader = StreamReader::try_new(std::io::Cursor::new(buf), None)?;
2755 reader.next().unwrap()
2756 }
2757
2758 fn read_stream_skip_validation(buf: &[u8]) -> Result<RecordBatch, ArrowError> {
2761 let mut reader = unsafe {
2762 StreamReader::try_new(std::io::Cursor::new(buf), None)?.with_skip_validation(true)
2763 };
2764 reader.next().unwrap()
2765 }
2766
2767 fn roundtrip_ipc_stream(rb: &RecordBatch) -> RecordBatch {
2768 let buf = write_stream(rb);
2769 read_stream(&buf).unwrap()
2770 }
2771
2772 #[test]
2773 fn test_roundtrip_with_custom_metadata() {
2774 let schema = Schema::new(vec![Field::new("dummy", DataType::Float64, false)]);
2775 let mut buf = Vec::new();
2776 let mut writer = crate::writer::FileWriter::try_new(&mut buf, &schema).unwrap();
2777 let mut test_metadata = HashMap::new();
2778 test_metadata.insert("abc".to_string(), "abc".to_string());
2779 test_metadata.insert("def".to_string(), "def".to_string());
2780 for (k, v) in &test_metadata {
2781 writer.write_metadata(k, v);
2782 }
2783 writer.finish().unwrap();
2784 drop(writer);
2785
2786 let reader = crate::reader::FileReader::try_new(std::io::Cursor::new(buf), None).unwrap();
2787 assert_eq!(reader.custom_metadata(), &test_metadata);
2788 }
2789
2790 #[test]
2791 fn test_roundtrip_nested_dict() {
2792 let inner: DictionaryArray<Int32Type> = vec!["a", "b", "a"].into_iter().collect();
2793
2794 let array = Arc::new(inner) as ArrayRef;
2795
2796 let dctfield = Arc::new(Field::new("dict", array.data_type().clone(), false));
2797
2798 let s = StructArray::from(vec![(dctfield, array)]);
2799 let struct_array = Arc::new(s) as ArrayRef;
2800
2801 let schema = Arc::new(Schema::new(vec![Field::new(
2802 "struct",
2803 struct_array.data_type().clone(),
2804 false,
2805 )]));
2806
2807 let batch = RecordBatch::try_new(schema, vec![struct_array]).unwrap();
2808
2809 assert_eq!(batch, roundtrip_ipc(&batch));
2810 }
2811
2812 #[test]
2813 fn test_roundtrip_nested_dict_no_preserve_dict_id() {
2814 let inner: DictionaryArray<Int32Type> = vec!["a", "b", "a"].into_iter().collect();
2815
2816 let array = Arc::new(inner) as ArrayRef;
2817
2818 let dctfield = Arc::new(Field::new("dict", array.data_type().clone(), false));
2819
2820 let s = StructArray::from(vec![(dctfield, array)]);
2821 let struct_array = Arc::new(s) as ArrayRef;
2822
2823 let schema = Arc::new(Schema::new(vec![Field::new(
2824 "struct",
2825 struct_array.data_type().clone(),
2826 false,
2827 )]));
2828
2829 let batch = RecordBatch::try_new(schema, vec![struct_array]).unwrap();
2830
2831 let mut buf = Vec::new();
2832 let mut writer = crate::writer::FileWriter::try_new_with_options(
2833 &mut buf,
2834 batch.schema_ref(),
2835 IpcWriteOptions::default(),
2836 )
2837 .unwrap();
2838 writer.write(&batch).unwrap();
2839 writer.finish().unwrap();
2840 drop(writer);
2841
2842 let mut reader = FileReader::try_new(std::io::Cursor::new(buf), None).unwrap();
2843
2844 assert_eq!(batch, reader.next().unwrap().unwrap());
2845 }
2846
2847 fn check_union_with_builder(mut builder: UnionBuilder) {
2848 builder.append::<Int32Type>("a", 1).unwrap();
2849 builder.append_null::<Int32Type>("a").unwrap();
2850 builder.append::<Float64Type>("c", 3.0).unwrap();
2851 builder.append::<Int32Type>("a", 4).unwrap();
2852 builder.append::<Int64Type>("d", 11).unwrap();
2853 let union = builder.build().unwrap();
2854
2855 let schema = Arc::new(Schema::new(vec![Field::new(
2856 "union",
2857 union.data_type().clone(),
2858 false,
2859 )]));
2860
2861 let union_array = Arc::new(union) as ArrayRef;
2862
2863 let rb = RecordBatch::try_new(schema, vec![union_array]).unwrap();
2864 let rb2 = roundtrip_ipc(&rb);
2865 assert_eq!(rb.schema(), rb2.schema());
2868 assert_eq!(rb.num_columns(), rb2.num_columns());
2869 assert_eq!(rb.num_rows(), rb2.num_rows());
2870 let union1 = rb.column(0);
2871 let union2 = rb2.column(0);
2872
2873 assert_eq!(union1, union2);
2874 }
2875
2876 #[test]
2877 fn test_roundtrip_dense_union() {
2878 check_union_with_builder(UnionBuilder::new_dense());
2879 }
2880
2881 #[test]
2882 fn test_roundtrip_sparse_union() {
2883 check_union_with_builder(UnionBuilder::new_sparse());
2884 }
2885
2886 #[test]
2887 fn test_roundtrip_struct_empty_fields() {
2888 let nulls = NullBuffer::from(&[true, true, false]);
2889 let rb = RecordBatch::try_from_iter([(
2890 "",
2891 Arc::new(StructArray::new_empty_fields(nulls.len(), Some(nulls))) as _,
2892 )])
2893 .unwrap();
2894 let rb2 = roundtrip_ipc(&rb);
2895 assert_eq!(rb, rb2);
2896 }
2897
2898 #[test]
2899 fn test_roundtrip_stream_run_array_sliced() {
2900 let run_array_1: Int32RunArray = vec!["a", "a", "a", "b", "b", "c", "c", "c"]
2901 .into_iter()
2902 .collect();
2903 let run_array_1_sliced = run_array_1.slice(2, 5);
2904
2905 let run_array_2_inupt = vec![Some(1_i32), None, None, Some(2), Some(2)];
2906 let mut run_array_2_builder = PrimitiveRunBuilder::<Int16Type, Int32Type>::new();
2907 run_array_2_builder.extend(run_array_2_inupt);
2908 let run_array_2 = run_array_2_builder.finish();
2909
2910 let schema = Arc::new(Schema::new(vec![
2911 Field::new(
2912 "run_array_1_sliced",
2913 run_array_1_sliced.data_type().clone(),
2914 false,
2915 ),
2916 Field::new("run_array_2", run_array_2.data_type().clone(), false),
2917 ]));
2918 let input_batch = RecordBatch::try_new(
2919 schema,
2920 vec![Arc::new(run_array_1_sliced.clone()), Arc::new(run_array_2)],
2921 )
2922 .unwrap();
2923 let output_batch = roundtrip_ipc_stream(&input_batch);
2924
2925 assert_eq!(input_batch.column(1), output_batch.column(1));
2929
2930 let run_array_1_unsliced = unslice_run_array(run_array_1_sliced.into_data()).unwrap();
2931 assert_eq!(run_array_1_unsliced, output_batch.column(0).into_data());
2932 }
2933
2934 #[test]
2935 fn test_roundtrip_stream_nested_dict() {
2936 let xs = vec!["AA", "BB", "AA", "CC", "BB"];
2937 let dict = Arc::new(
2938 xs.clone()
2939 .into_iter()
2940 .collect::<DictionaryArray<Int8Type>>(),
2941 );
2942 let string_array: ArrayRef = Arc::new(StringArray::from(xs.clone()));
2943 let struct_array = StructArray::from(vec![
2944 (
2945 Arc::new(Field::new("f2.1", DataType::Utf8, false)),
2946 string_array,
2947 ),
2948 (
2949 Arc::new(Field::new("f2.2_struct", dict.data_type().clone(), false)),
2950 dict.clone() as ArrayRef,
2951 ),
2952 ]);
2953 let schema = Arc::new(Schema::new(vec![
2954 Field::new("f1_string", DataType::Utf8, false),
2955 Field::new("f2_struct", struct_array.data_type().clone(), false),
2956 ]));
2957 let input_batch = RecordBatch::try_new(
2958 schema,
2959 vec![
2960 Arc::new(StringArray::from(xs.clone())),
2961 Arc::new(struct_array),
2962 ],
2963 )
2964 .unwrap();
2965 let output_batch = roundtrip_ipc_stream(&input_batch);
2966 assert_eq!(input_batch, output_batch);
2967 }
2968
2969 #[test]
2970 fn test_ipc_writers_reject_dictionary_of_dictionary_schema() {
2971 let values = Arc::new(StringArray::from(vec![Some("a"), Some("b")])) as ArrayRef;
2972 let inner = Arc::new(DictionaryArray::new(
2973 UInt32Array::from_iter_values([0, 1]),
2974 values,
2975 )) as ArrayRef;
2976 let outer = Arc::new(DictionaryArray::new(
2977 UInt32Array::from_iter_values([0, 1, 0]),
2978 inner,
2979 )) as ArrayRef;
2980
2981 let schema = Arc::new(Schema::new(vec![Field::new(
2982 "f1",
2983 outer.data_type().clone(),
2984 false,
2985 )]));
2986 let batch = RecordBatch::try_new(schema, vec![outer]).unwrap();
2987
2988 let mut stream = Vec::new();
2989 let Err(err) = crate::writer::StreamWriter::try_new(&mut stream, batch.schema_ref()) else {
2990 panic!("IPC stream writer should reject dictionary-of-dictionary schemas");
2991 };
2992 assert!(stream.is_empty());
2993
2994 assert!(
2995 err.to_string().contains("dictionary-of-dictionary values"),
2996 "unexpected error: {err}"
2997 );
2998
2999 let mut file = Vec::new();
3000 let Err(err) = crate::writer::FileWriter::try_new(&mut file, batch.schema_ref()) else {
3001 panic!("IPC file writer should reject dictionary-of-dictionary schemas");
3002 };
3003 assert!(file.is_empty());
3004
3005 assert!(
3006 err.to_string().contains("dictionary-of-dictionary values"),
3007 "unexpected error: {err}"
3008 );
3009 }
3010
3011 #[test]
3012 fn test_roundtrip_stream_nested_dict_of_map_of_dict() {
3013 let values = StringArray::from(vec![Some("a"), None, Some("b"), Some("c")]);
3014 let values = Arc::new(values) as ArrayRef;
3015 let value_dict_keys = Int8Array::from_iter_values([0, 1, 1, 2, 3, 1]);
3016 let value_dict_array = DictionaryArray::new(value_dict_keys, values.clone());
3017
3018 let key_dict_keys = Int8Array::from_iter_values([0, 0, 2, 2, 2, 3]);
3019 let key_dict_array = DictionaryArray::new(key_dict_keys, values);
3020
3021 #[expect(deprecated)]
3022 let keys_field = Arc::new(Field::new_dict(
3023 Field::MAP_KEY_FIELD_DEFAULT_NAME,
3024 DataType::Dictionary(Box::new(DataType::Int8), Box::new(DataType::Utf8)),
3025 false,
3026 1,
3027 false,
3028 ));
3029 #[expect(deprecated)]
3030 let values_field = Arc::new(Field::new_dict(
3031 Field::MAP_VALUE_FIELD_DEFAULT_NAME,
3032 DataType::Dictionary(Box::new(DataType::Int8), Box::new(DataType::Utf8)),
3033 true,
3034 2,
3035 false,
3036 ));
3037 let entry_struct = StructArray::from(vec![
3038 (keys_field, make_array(key_dict_array.into_data())),
3039 (values_field, make_array(value_dict_array.into_data())),
3040 ]);
3041 let map_data_type = DataType::Map(
3042 Arc::new(Field::new(
3043 Field::MAP_ENTRIES_FIELD_DEFAULT_NAME,
3044 entry_struct.data_type().clone(),
3045 false,
3046 )),
3047 false,
3048 );
3049
3050 let entry_offsets = Buffer::from_slice_ref([0, 2, 4, 6]);
3051 let map_data = ArrayData::builder(map_data_type)
3052 .len(3)
3053 .add_buffer(entry_offsets)
3054 .add_child_data(entry_struct.into_data())
3055 .build()
3056 .unwrap();
3057 let map_array = MapArray::from(map_data);
3058
3059 let dict_keys = Int8Array::from_iter_values([0, 1, 1, 2, 2, 1]);
3060 let dict_dict_array = DictionaryArray::new(dict_keys, Arc::new(map_array));
3061
3062 let schema = Arc::new(Schema::new(vec![Field::new(
3063 "f1",
3064 dict_dict_array.data_type().clone(),
3065 false,
3066 )]));
3067 let input_batch = RecordBatch::try_new(schema, vec![Arc::new(dict_dict_array)]).unwrap();
3068 let output_batch = roundtrip_ipc_stream(&input_batch);
3069 assert_eq!(input_batch, output_batch);
3070 }
3071
3072 fn test_roundtrip_stream_dict_of_list_of_dict_impl<
3073 OffsetSize: OffsetSizeTrait,
3074 U: ArrowNativeType,
3075 >(
3076 list_data_type: DataType,
3077 offsets: &[U; 5],
3078 ) {
3079 let values = StringArray::from(vec![Some("a"), None, Some("c"), None]);
3080 let keys = Int8Array::from_iter_values([0, 0, 1, 2, 0, 1, 3]);
3081 let dict_array = DictionaryArray::new(keys, Arc::new(values));
3082 let dict_data = dict_array.to_data();
3083
3084 let value_offsets = Buffer::from_slice_ref(offsets);
3085
3086 let list_data = ArrayData::builder(list_data_type)
3087 .len(4)
3088 .add_buffer(value_offsets)
3089 .add_child_data(dict_data)
3090 .build()
3091 .unwrap();
3092 let list_array = GenericListArray::<OffsetSize>::from(list_data);
3093
3094 let keys_for_dict = Int8Array::from_iter_values([0, 3, 0, 1, 1, 2, 0, 1, 3]);
3095 let dict_dict_array = DictionaryArray::new(keys_for_dict, Arc::new(list_array));
3096
3097 let schema = Arc::new(Schema::new(vec![Field::new(
3098 "f1",
3099 dict_dict_array.data_type().clone(),
3100 false,
3101 )]));
3102 let input_batch = RecordBatch::try_new(schema, vec![Arc::new(dict_dict_array)]).unwrap();
3103 let output_batch = roundtrip_ipc_stream(&input_batch);
3104 assert_eq!(input_batch, output_batch);
3105 }
3106
3107 #[test]
3108 fn test_roundtrip_stream_dict_of_list_of_dict() {
3109 #[expect(deprecated)]
3111 let list_data_type = DataType::List(Arc::new(Field::new_dict(
3112 "item",
3113 DataType::Dictionary(Box::new(DataType::Int8), Box::new(DataType::Utf8)),
3114 true,
3115 1,
3116 false,
3117 )));
3118 let offsets: &[i32; 5] = &[0, 2, 4, 4, 6];
3119 test_roundtrip_stream_dict_of_list_of_dict_impl::<i32, i32>(list_data_type, offsets);
3120
3121 #[expect(deprecated)]
3123 let list_data_type = DataType::LargeList(Arc::new(Field::new_dict(
3124 "item",
3125 DataType::Dictionary(Box::new(DataType::Int8), Box::new(DataType::Utf8)),
3126 true,
3127 1,
3128 false,
3129 )));
3130 let offsets: &[i64; 5] = &[0, 2, 4, 4, 7];
3131 test_roundtrip_stream_dict_of_list_of_dict_impl::<i64, i64>(list_data_type, offsets);
3132 }
3133
3134 #[test]
3135 fn test_roundtrip_stream_dict_of_fixed_size_list_of_dict() {
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, 1, 2]);
3138 let dict_array = DictionaryArray::new(keys, Arc::new(values));
3139 let dict_data = dict_array.into_data();
3140
3141 #[expect(deprecated)]
3142 let list_data_type = DataType::FixedSizeList(
3143 Arc::new(Field::new_dict(
3144 "item",
3145 DataType::Dictionary(Box::new(DataType::Int8), Box::new(DataType::Utf8)),
3146 true,
3147 1,
3148 false,
3149 )),
3150 3,
3151 );
3152 let list_data = ArrayData::builder(list_data_type)
3153 .len(3)
3154 .add_child_data(dict_data)
3155 .build()
3156 .unwrap();
3157 let list_array = FixedSizeListArray::from(list_data);
3158
3159 let keys_for_dict = Int8Array::from_iter_values([0, 1, 0, 1, 1, 2, 0, 1, 2]);
3160 let dict_dict_array = DictionaryArray::new(keys_for_dict, Arc::new(list_array));
3161
3162 let schema = Arc::new(Schema::new(vec![Field::new(
3163 "f1",
3164 dict_dict_array.data_type().clone(),
3165 false,
3166 )]));
3167 let input_batch = RecordBatch::try_new(schema, vec![Arc::new(dict_dict_array)]).unwrap();
3168 let output_batch = roundtrip_ipc_stream(&input_batch);
3169 assert_eq!(input_batch, output_batch);
3170 }
3171
3172 const LONG_TEST_STRING: &str =
3173 "This is a long string to make sure binary view array handles it";
3174
3175 #[test]
3176 fn test_roundtrip_view_types() {
3177 let schema = Schema::new(vec![
3178 Field::new("field_1", DataType::BinaryView, true),
3179 Field::new("field_2", DataType::Utf8, true),
3180 Field::new("field_3", DataType::Utf8View, true),
3181 ]);
3182 let bin_values: Vec<Option<&[u8]>> = vec![
3183 Some(b"foo"),
3184 None,
3185 Some(b"bar"),
3186 Some(LONG_TEST_STRING.as_bytes()),
3187 ];
3188 let utf8_values: Vec<Option<&str>> =
3189 vec![Some("foo"), None, Some("bar"), Some(LONG_TEST_STRING)];
3190 let bin_view_array = BinaryViewArray::from_iter(bin_values);
3191 let utf8_array = StringArray::from_iter(utf8_values.iter());
3192 let utf8_view_array = StringViewArray::from_iter(utf8_values);
3193 let record_batch = RecordBatch::try_new(
3194 Arc::new(schema.clone()),
3195 vec![
3196 Arc::new(bin_view_array),
3197 Arc::new(utf8_array),
3198 Arc::new(utf8_view_array),
3199 ],
3200 )
3201 .unwrap();
3202
3203 assert_eq!(record_batch, roundtrip_ipc(&record_batch));
3204 assert_eq!(record_batch, roundtrip_ipc_stream(&record_batch));
3205
3206 let sliced_batch = record_batch.slice(1, 2);
3207 assert_eq!(sliced_batch, roundtrip_ipc(&sliced_batch));
3208 assert_eq!(sliced_batch, roundtrip_ipc_stream(&sliced_batch));
3209 }
3210
3211 #[test]
3212 fn test_roundtrip_view_types_nested_dict() {
3213 let bin_values: Vec<Option<&[u8]>> = vec![
3214 Some(b"foo"),
3215 None,
3216 Some(b"bar"),
3217 Some(LONG_TEST_STRING.as_bytes()),
3218 Some(b"field"),
3219 ];
3220 let utf8_values: Vec<Option<&str>> = vec![
3221 Some("foo"),
3222 None,
3223 Some("bar"),
3224 Some(LONG_TEST_STRING),
3225 Some("field"),
3226 ];
3227 let bin_view_array = Arc::new(BinaryViewArray::from_iter(bin_values));
3228 let utf8_view_array = Arc::new(StringViewArray::from_iter(utf8_values));
3229
3230 let key_dict_keys = Int8Array::from_iter_values([0, 0, 2, 2, 0, 2, 3]);
3231 let key_dict_array = DictionaryArray::new(key_dict_keys, utf8_view_array.clone());
3232 #[expect(deprecated)]
3233 let keys_field = Arc::new(Field::new_dict(
3234 Field::MAP_KEY_FIELD_DEFAULT_NAME,
3235 DataType::Dictionary(Box::new(DataType::Int8), Box::new(DataType::Utf8View)),
3236 false,
3237 1,
3238 false,
3239 ));
3240
3241 let value_dict_keys = Int8Array::from_iter_values([0, 3, 0, 1, 2, 0, 1]);
3242 let value_dict_array = DictionaryArray::new(value_dict_keys, bin_view_array);
3243 #[expect(deprecated)]
3244 let values_field = Arc::new(Field::new_dict(
3245 Field::MAP_VALUE_FIELD_DEFAULT_NAME,
3246 DataType::Dictionary(Box::new(DataType::Int8), Box::new(DataType::BinaryView)),
3247 true,
3248 2,
3249 false,
3250 ));
3251 let entry_struct = StructArray::from(vec![
3252 (keys_field, make_array(key_dict_array.into_data())),
3253 (values_field, make_array(value_dict_array.into_data())),
3254 ]);
3255
3256 let map_data_type = DataType::Map(
3257 Arc::new(Field::new(
3258 Field::MAP_ENTRIES_FIELD_DEFAULT_NAME,
3259 entry_struct.data_type().clone(),
3260 false,
3261 )),
3262 false,
3263 );
3264 let entry_offsets = Buffer::from_slice_ref([0, 2, 4, 7]);
3265 let map_data = ArrayData::builder(map_data_type)
3266 .len(3)
3267 .add_buffer(entry_offsets)
3268 .add_child_data(entry_struct.into_data())
3269 .build()
3270 .unwrap();
3271 let map_array = MapArray::from(map_data);
3272
3273 let dict_keys = Int8Array::from_iter_values([0, 1, 0, 1, 1, 2, 0, 1, 2]);
3274 let dict_dict_array = DictionaryArray::new(dict_keys, Arc::new(map_array));
3275 let schema = Arc::new(Schema::new(vec![Field::new(
3276 "f1",
3277 dict_dict_array.data_type().clone(),
3278 false,
3279 )]));
3280 let batch = RecordBatch::try_new(schema, vec![Arc::new(dict_dict_array)]).unwrap();
3281 assert_eq!(batch, roundtrip_ipc(&batch));
3282 assert_eq!(batch, roundtrip_ipc_stream(&batch));
3283
3284 let sliced_batch = batch.slice(1, 2);
3285 assert_eq!(sliced_batch, roundtrip_ipc(&sliced_batch));
3286 assert_eq!(sliced_batch, roundtrip_ipc_stream(&sliced_batch));
3287 }
3288
3289 #[test]
3290 fn test_no_columns_batch() {
3291 let schema = Arc::new(Schema::empty());
3292 let options = RecordBatchOptions::new()
3293 .with_match_field_names(true)
3294 .with_row_count(Some(10));
3295 let input_batch = RecordBatch::try_new_with_options(schema, vec![], &options).unwrap();
3296 let output_batch = roundtrip_ipc_stream(&input_batch);
3297 assert_eq!(input_batch, output_batch);
3298 }
3299
3300 #[test]
3301 fn test_unaligned() {
3302 let batch = RecordBatch::try_from_iter(vec![(
3303 "i32",
3304 Arc::new(Int32Array::from(vec![1, 2, 3, 4])) as _,
3305 )])
3306 .unwrap();
3307
3308 let r#gen = IpcDataGenerator {};
3309 let mut dict_tracker = DictionaryTracker::new(false);
3310 let (_, encoded) = r#gen
3311 .encode(
3312 &batch,
3313 &mut dict_tracker,
3314 &Default::default(),
3315 &mut Default::default(),
3316 )
3317 .unwrap();
3318
3319 let message = root_as_message(&encoded.ipc_message).unwrap();
3320
3321 let mut buffer = MutableBuffer::with_capacity(encoded.arrow_data.len() + 1);
3323 buffer.push(0_u8);
3324 buffer.extend_from_slice(&encoded.arrow_data);
3325 let b = Buffer::from(buffer).slice(1);
3326 assert_ne!(b.as_ptr().align_offset(8), 0);
3327
3328 let ipc_batch = message.header_as_record_batch().unwrap();
3329 let roundtrip = RecordBatchDecoder::try_new(
3330 &b,
3331 ipc_batch,
3332 batch.schema(),
3333 &Default::default(),
3334 &message.version(),
3335 )
3336 .unwrap()
3337 .with_require_alignment(false)
3338 .read_record_batch()
3339 .unwrap();
3340 assert_eq!(batch, roundtrip);
3341 }
3342
3343 #[test]
3344 fn test_unaligned_throws_error_with_require_alignment() {
3345 let batch = RecordBatch::try_from_iter(vec![(
3346 "i32",
3347 Arc::new(Int32Array::from(vec![1, 2, 3, 4])) as _,
3348 )])
3349 .unwrap();
3350
3351 let r#gen = IpcDataGenerator {};
3352 let mut dict_tracker = DictionaryTracker::new(false);
3353 let (_, encoded) = r#gen
3354 .encode(
3355 &batch,
3356 &mut dict_tracker,
3357 &Default::default(),
3358 &mut Default::default(),
3359 )
3360 .unwrap();
3361
3362 let message = root_as_message(&encoded.ipc_message).unwrap();
3363
3364 let mut buffer = MutableBuffer::with_capacity(encoded.arrow_data.len() + 1);
3366 buffer.push(0_u8);
3367 buffer.extend_from_slice(&encoded.arrow_data);
3368 let b = Buffer::from(buffer).slice(1);
3369 assert_ne!(b.as_ptr().align_offset(8), 0);
3370
3371 let ipc_batch = message.header_as_record_batch().unwrap();
3372 let result = RecordBatchDecoder::try_new(
3373 &b,
3374 ipc_batch,
3375 batch.schema(),
3376 &Default::default(),
3377 &message.version(),
3378 )
3379 .unwrap()
3380 .with_require_alignment(true)
3381 .read_record_batch();
3382
3383 let error = result.unwrap_err();
3384 assert_eq!(
3385 error.to_string(),
3386 "Invalid argument error: Misaligned buffers[0] in array of type Int32, \
3387 offset from expected alignment of 4 by 1"
3388 );
3389 }
3390
3391 #[test]
3400 fn test_misaligned_buffers_error() {
3401 const OFFSET_RANGE: usize = 128;
3402
3403 let cases: Vec<(ArrayRef, usize)> = vec![
3407 (Arc::new(Int32Array::from_iter(0i32..100)) as _, 4),
3408 (Arc::new(Int64Array::from_iter(0i64..100)) as _, 8),
3409 (
3410 Arc::new(StringArray::from_iter_values(
3411 (0..100).map(|i| i.to_string()),
3412 )) as _,
3413 4,
3414 ),
3415 (
3416 Arc::new(LargeStringArray::from_iter_values(
3417 (0..100).map(|i| i.to_string()),
3418 )) as _,
3419 8,
3420 ),
3421 ];
3422
3423 for (array, alignment) in cases {
3424 let batch = RecordBatch::try_from_iter(vec![("col", Arc::clone(&array))]).unwrap();
3425 let encoder = IpcDataGenerator {};
3426 let mut dict_tracker = DictionaryTracker::new(false);
3427 let (_, encoded) = encoder
3428 .encode(
3429 &batch,
3430 &mut dict_tracker,
3431 &Default::default(),
3432 &mut Default::default(),
3433 )
3434 .unwrap();
3435 let message = root_as_message(&encoded.ipc_message).unwrap();
3436 let ipc_batch = message.header_as_record_batch().unwrap();
3437
3438 for offset in 0..OFFSET_RANGE {
3439 let mut storage = MutableBuffer::with_capacity(encoded.arrow_data.len() + offset);
3445 for _ in 0..offset {
3446 storage.push(0_u8);
3447 }
3448 storage.extend_from_slice(&encoded.arrow_data);
3449 let buf = Buffer::from(storage).slice(offset);
3450
3451 let result = RecordBatchDecoder::try_new(
3452 &buf,
3453 ipc_batch,
3454 batch.schema(),
3455 &Default::default(),
3456 &message.version(),
3457 )
3458 .unwrap()
3459 .with_require_alignment(true)
3460 .read_record_batch();
3461
3462 if offset % alignment == 0 {
3463 assert!(
3464 result.is_ok(),
3465 "type={} offset={offset}: expected Ok but got {:?}",
3466 array.data_type(),
3467 result.unwrap_err(),
3468 );
3469 } else {
3470 let err = result
3471 .expect_err(&format!(
3472 "type={} offset={offset}: expected Err for misaligned buffer",
3473 array.data_type()
3474 ))
3475 .to_string();
3476 assert!(
3477 err.contains("Misaligned buffers"),
3478 "type={} offset={offset}: unexpected error: {err}",
3479 array.data_type(),
3480 );
3481 }
3482 }
3483 }
3484 }
3485
3486 #[test]
3487 #[cfg_attr(miri, ignore)] fn test_file_with_massive_column_count() {
3489 let limit = 600_000;
3491
3492 let fields = (0..limit)
3493 .map(|i| Field::new(format!("{i}"), DataType::Boolean, false))
3494 .collect::<Vec<_>>();
3495 let schema = Arc::new(Schema::new(fields));
3496 let batch = RecordBatch::new_empty(schema);
3497
3498 let mut buf = Vec::new();
3499 let mut writer = crate::writer::FileWriter::try_new(&mut buf, batch.schema_ref()).unwrap();
3500 writer.write(&batch).unwrap();
3501 writer.finish().unwrap();
3502 drop(writer);
3503
3504 let mut reader = FileReaderBuilder::new()
3505 .with_max_footer_fb_tables(1_500_000)
3506 .build(std::io::Cursor::new(buf))
3507 .unwrap();
3508 let roundtrip_batch = reader.next().unwrap().unwrap();
3509
3510 assert_eq!(batch, roundtrip_batch);
3511 }
3512
3513 #[test]
3514 fn test_file_with_deeply_nested_columns() {
3515 let limit = 61;
3517
3518 let fields = (0..limit).fold(
3519 vec![Field::new("leaf", DataType::Boolean, false)],
3520 |field, index| vec![Field::new_struct(format!("{index}"), field, false)],
3521 );
3522 let schema = Arc::new(Schema::new(fields));
3523 let batch = RecordBatch::new_empty(schema);
3524
3525 let mut buf = Vec::new();
3526 let mut writer = crate::writer::FileWriter::try_new(&mut buf, batch.schema_ref()).unwrap();
3527 writer.write(&batch).unwrap();
3528 writer.finish().unwrap();
3529 drop(writer);
3530
3531 let mut reader = FileReaderBuilder::new()
3532 .with_max_footer_fb_depth(65)
3533 .build(std::io::Cursor::new(buf))
3534 .unwrap();
3535 let roundtrip_batch = reader.next().unwrap().unwrap();
3536
3537 assert_eq!(batch, roundtrip_batch);
3538 }
3539
3540 #[test]
3541 fn test_invalid_struct_array_ipc_read_errors() {
3542 let a_field = Field::new("a", DataType::Int32, false);
3543 let b_field = Field::new("b", DataType::Int32, false);
3544 let struct_fields = Fields::from(vec![a_field.clone(), b_field.clone()]);
3545
3546 let a_array_data = ArrayData::builder(a_field.data_type().clone())
3547 .len(4)
3548 .add_buffer(Buffer::from_slice_ref([1, 2, 3, 4]))
3549 .build()
3550 .unwrap();
3551 let b_array_data = ArrayData::builder(b_field.data_type().clone())
3552 .len(3)
3553 .add_buffer(Buffer::from_slice_ref([5, 6, 7]))
3554 .build()
3555 .unwrap();
3556
3557 let invalid_struct_arr = unsafe {
3558 StructArray::new_unchecked(
3559 struct_fields,
3560 vec![make_array(a_array_data), make_array(b_array_data)],
3561 None,
3562 )
3563 };
3564
3565 expect_ipc_validation_error(
3566 Arc::new(invalid_struct_arr),
3567 "Invalid argument error: Incorrect array length for StructArray field \"b\", expected 4 got 3",
3568 );
3569 }
3570
3571 #[test]
3572 fn test_invalid_nested_array_ipc_read_errors() {
3573 let a_field = Field::new("a", DataType::Int32, false);
3575 let b_field = Field::new("b", DataType::Utf8, false);
3576
3577 let schema = Arc::new(Schema::new(vec![Field::new_struct(
3578 "s",
3579 vec![a_field.clone(), b_field.clone()],
3580 false,
3581 )]));
3582
3583 let a_array_data = ArrayData::builder(a_field.data_type().clone())
3584 .len(4)
3585 .add_buffer(Buffer::from_slice_ref([1, 2, 3, 4]))
3586 .build()
3587 .unwrap();
3588 let b_array_data = {
3590 let valid: &[u8] = b" ";
3591 let mut invalid = vec![];
3592 invalid.extend_from_slice(b"ValidString");
3593 invalid.extend_from_slice(INVALID_UTF8_FIRST_CHAR);
3594 let binary_array =
3595 BinaryArray::from_iter(vec![None, Some(valid), None, Some(&invalid)]);
3596 let array = unsafe {
3597 StringArray::new_unchecked(
3598 binary_array.offsets().clone(),
3599 binary_array.values().clone(),
3600 binary_array.nulls().cloned(),
3601 )
3602 };
3603 array.into_data()
3604 };
3605 let struct_data_type = schema.field(0).data_type();
3606
3607 let invalid_struct_arr = unsafe {
3608 make_array(
3609 ArrayData::builder(struct_data_type.clone())
3610 .len(4)
3611 .add_child_data(a_array_data)
3612 .add_child_data(b_array_data)
3613 .build_unchecked(),
3614 )
3615 };
3616 expect_ipc_validation_error(
3617 invalid_struct_arr,
3618 "Invalid argument error: Invalid UTF8 sequence at string index 3 (3..18): invalid utf-8 sequence of 1 bytes from index 11",
3619 );
3620 }
3621
3622 #[test]
3623 fn test_same_dict_id_without_preserve() {
3624 let batch = RecordBatch::try_new(
3625 Arc::new(Schema::new(
3626 ["a", "b"]
3627 .iter()
3628 .map(|name| {
3629 #[expect(deprecated)]
3630 Field::new_dict(
3631 name.to_string(),
3632 DataType::Dictionary(
3633 Box::new(DataType::Int32),
3634 Box::new(DataType::Utf8),
3635 ),
3636 true,
3637 0,
3638 false,
3639 )
3640 })
3641 .collect::<Vec<Field>>(),
3642 )),
3643 vec![
3644 Arc::new(
3645 vec![Some("c"), Some("d")]
3646 .into_iter()
3647 .collect::<DictionaryArray<Int32Type>>(),
3648 ) as ArrayRef,
3649 Arc::new(
3650 vec![Some("e"), Some("f")]
3651 .into_iter()
3652 .collect::<DictionaryArray<Int32Type>>(),
3653 ) as ArrayRef,
3654 ],
3655 )
3656 .expect("Failed to create RecordBatch");
3657
3658 let mut buf = vec![];
3660 {
3661 let mut writer = crate::writer::StreamWriter::try_new_with_options(
3662 &mut buf,
3663 batch.schema().as_ref(),
3664 crate::writer::IpcWriteOptions::default(),
3665 )
3666 .expect("Failed to create StreamWriter");
3667 writer.write(&batch).expect("Failed to write RecordBatch");
3668 writer.finish().expect("Failed to finish StreamWriter");
3669 }
3670
3671 StreamReader::try_new(std::io::Cursor::new(buf), None)
3672 .expect("Failed to create StreamReader")
3673 .for_each(|decoded_batch| {
3674 assert_eq!(decoded_batch.expect("Failed to read RecordBatch"), batch);
3675 });
3676 }
3677
3678 #[test]
3679 fn test_validation_of_invalid_list_array() {
3680 let array = unsafe {
3682 let values = Int32Array::from(vec![1, 2, 3]);
3683 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);
3686 let nulls = None;
3687 ListArray::new(Arc::new(field), offsets, Arc::new(values), nulls)
3688 };
3689
3690 expect_ipc_validation_error(
3691 Arc::new(array),
3692 "Invalid argument error: Offset invariant failure: offset at position 2 out of bounds: 4 > 2",
3693 );
3694 }
3695
3696 #[test]
3697 fn test_validation_of_invalid_string_array() {
3698 let valid: &[u8] = b" ";
3699 let mut invalid = vec![];
3700 invalid.extend_from_slice(b"ThisStringIsCertainlyLongerThan12Bytes");
3701 invalid.extend_from_slice(INVALID_UTF8_FIRST_CHAR);
3702 let binary_array = BinaryArray::from_iter(vec![None, Some(valid), None, Some(&invalid)]);
3703 let array = unsafe {
3706 StringArray::new_unchecked(
3707 binary_array.offsets().clone(),
3708 binary_array.values().clone(),
3709 binary_array.nulls().cloned(),
3710 )
3711 };
3712 expect_ipc_validation_error(
3713 Arc::new(array),
3714 "Invalid argument error: Invalid UTF8 sequence at string index 3 (3..45): invalid utf-8 sequence of 1 bytes from index 38",
3715 );
3716 }
3717
3718 #[test]
3719 fn test_validation_of_invalid_string_view_array() {
3720 let valid: &[u8] = b" ";
3721 let mut invalid = vec![];
3722 invalid.extend_from_slice(b"ThisStringIsCertainlyLongerThan12Bytes");
3723 invalid.extend_from_slice(INVALID_UTF8_FIRST_CHAR);
3724 let binary_view_array =
3725 BinaryViewArray::from_iter(vec![None, Some(valid), None, Some(&invalid)]);
3726 let array = unsafe {
3729 StringViewArray::new_unchecked(
3730 binary_view_array.views().clone(),
3731 Arc::clone(binary_view_array.data_buffers()),
3732 binary_view_array.nulls().cloned(),
3733 )
3734 };
3735 expect_ipc_validation_error(
3736 Arc::new(array),
3737 "Invalid argument error: Encountered non-UTF-8 data at index 3: invalid utf-8 sequence of 1 bytes from index 38",
3738 );
3739 }
3740
3741 #[test]
3744 fn test_validation_of_invalid_dictionary_array() {
3745 let array = unsafe {
3746 let values = StringArray::from_iter_values(["a", "b", "c"]);
3747 let keys = Int32Array::from(vec![1, 200]); DictionaryArray::new_unchecked(keys, Arc::new(values))
3749 };
3750
3751 expect_ipc_validation_error(
3752 Arc::new(array),
3753 "Invalid argument error: Value at position 1 out of bounds: 200 (should be in [0, 2])",
3754 );
3755 }
3756
3757 #[test]
3758 fn test_validation_of_invalid_union_array() {
3759 let array = unsafe {
3760 let fields = UnionFields::try_new(
3761 vec![1, 3], vec![
3763 Field::new("a", DataType::Int32, false),
3764 Field::new("b", DataType::Utf8, false),
3765 ],
3766 )
3767 .unwrap();
3768 let type_ids = ScalarBuffer::from(vec![1i8, 2, 3]); let offsets = None;
3770 let children: Vec<ArrayRef> = vec![
3771 Arc::new(Int32Array::from(vec![10, 20, 30])),
3772 Arc::new(StringArray::from(vec![Some("a"), Some("b"), Some("c")])),
3773 ];
3774
3775 UnionArray::new_unchecked(fields, type_ids, offsets, children)
3776 };
3777
3778 expect_ipc_validation_error(
3779 Arc::new(array),
3780 "Invalid argument error: Type Ids values must match one of the field type ids",
3781 );
3782 }
3783
3784 const INVALID_UTF8_FIRST_CHAR: &[u8] = &[0xa0, 0xa1, 0x20, 0x20];
3787
3788 fn expect_ipc_validation_error(array: ArrayRef, expected_err: &str) {
3790 let rb = RecordBatch::try_from_iter([("a", array)]).unwrap();
3791
3792 let buf = write_stream(&rb); read_stream_skip_validation(&buf).unwrap();
3795 let err = read_stream(&buf).unwrap_err();
3796 assert_eq!(err.to_string(), expected_err);
3797
3798 let buf = write_ipc(&rb); read_ipc_skip_validation(&buf).unwrap();
3801 let err = read_ipc(&buf).unwrap_err();
3802 assert_eq!(err.to_string(), expected_err);
3803
3804 read_ipc_with_decoder_skip_validation(buf.clone()).unwrap();
3806 let err = read_ipc_with_decoder(buf).unwrap_err();
3807 assert_eq!(err.to_string(), expected_err);
3808 }
3809
3810 #[test]
3811 fn test_roundtrip_schema() {
3812 let schema = Schema::new(vec![
3813 Field::new(
3814 "a",
3815 DataType::Dictionary(Box::new(DataType::UInt16), Box::new(DataType::Utf8)),
3816 false,
3817 ),
3818 Field::new(
3819 "b",
3820 DataType::Dictionary(Box::new(DataType::UInt16), Box::new(DataType::Utf8)),
3821 false,
3822 ),
3823 ]);
3824
3825 let options = IpcWriteOptions::default();
3826 let data_gen = IpcDataGenerator::default();
3827 let mut dict_tracker = DictionaryTracker::new(false);
3828 let encoded_data =
3829 data_gen.schema_to_bytes_with_dictionary_tracker(&schema, &mut dict_tracker, &options);
3830 let mut schema_bytes = vec![];
3831 write_message(&mut schema_bytes, encoded_data, &options).expect("write_message");
3832
3833 let begin_offset: usize = if schema_bytes[0..4].eq(&CONTINUATION_MARKER) {
3834 4
3835 } else {
3836 0
3837 };
3838
3839 size_prefixed_root_as_message(&schema_bytes[begin_offset..])
3840 .expect_err("size_prefixed_root_as_message");
3841
3842 let msg = parse_message(&schema_bytes).expect("parse_message");
3843 let ipc_schema = msg.header_as_schema().expect("header_as_schema");
3844 let new_schema = try_fb_to_schema(ipc_schema).unwrap();
3845
3846 assert_eq!(schema, new_schema);
3847 }
3848
3849 #[test]
3850 fn test_negative_meta_len() {
3851 let bytes = i32::to_le_bytes(-1);
3852 let mut buf = vec![];
3853 buf.extend(CONTINUATION_MARKER);
3854 buf.extend(bytes);
3855
3856 let reader = StreamReader::try_new(Cursor::new(buf), None);
3857 assert!(reader.is_err());
3858 }
3859
3860 #[test]
3866 fn test_read_null_dict_without_dictionary_batch() {
3867 let keys = Int32Array::new_null(4);
3869 let values: ArrayRef = new_empty_array(&DataType::Utf8);
3870 let dict_array = DictionaryArray::new(keys, values);
3871
3872 let schema = Arc::new(Schema::new(vec![Field::new(
3873 "d",
3874 dict_array.data_type().clone(),
3875 true,
3876 )]));
3877 let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(dict_array)]).unwrap();
3878
3879 let full_stream = write_stream(&batch);
3881
3882 let mut stripped = Vec::new();
3886 let mut cursor = Cursor::new(&full_stream);
3887 loop {
3888 let mut header = [0u8; 4];
3891 if cursor.read_exact(&mut header).is_err() {
3892 break;
3893 }
3894 if header == CONTINUATION_MARKER && cursor.read_exact(&mut header).is_err() {
3895 break;
3896 }
3897 let meta_len = u32::from_le_bytes(header) as usize;
3898 if meta_len == 0 {
3899 stripped.extend_from_slice(&CONTINUATION_MARKER);
3901 stripped.extend_from_slice(&0u32.to_le_bytes());
3902 break;
3903 }
3904 let mut meta_buf = vec![0u8; meta_len];
3905 cursor.read_exact(&mut meta_buf).unwrap();
3906
3907 let message = root_as_message(&meta_buf).unwrap();
3908 let body_len = message.bodyLength() as usize;
3909 let mut body_buf = vec![0u8; body_len];
3910 cursor.read_exact(&mut body_buf).unwrap();
3911
3912 if message.header_type() == crate::MessageHeader::DictionaryBatch {
3913 continue;
3916 }
3917 stripped.extend_from_slice(&CONTINUATION_MARKER);
3918 stripped.extend_from_slice(&(meta_len as u32).to_le_bytes());
3919 stripped.extend_from_slice(&meta_buf);
3920 stripped.extend_from_slice(&body_buf);
3921 }
3922
3923 let result = read_stream(&stripped).unwrap();
3925 assert_eq!(result.num_rows(), 4);
3926 assert_eq!(result.num_columns(), 1);
3927
3928 let col = result.column(0);
3929 assert_eq!(col.null_count(), 4);
3930 assert_eq!(col.len(), 4);
3931 assert!(matches!(col.data_type(), DataType::Dictionary(_, _)));
3933 }
3934
3935 #[test]
3938 fn test_projection_skip_list_view() {
3939 use crate::reader::FileReader;
3940 use crate::writer::FileWriter;
3941 use arrow_array::{
3942 GenericListViewArray, Int32Array, RecordBatch,
3943 builder::{GenericListViewBuilder, UInt32Builder},
3944 };
3945 use arrow_schema::{DataType, Field, Schema};
3946 use std::sync::Arc;
3947
3948 let mut builder = GenericListViewBuilder::<i32, _>::new(UInt32Builder::new());
3950
3951 builder.values().append_value(1);
3952 builder.values().append_value(2);
3953 builder.append(true);
3954
3955 builder.append(false);
3956
3957 builder.values().append_value(3);
3958 builder.values().append_value(4);
3959 builder.append(true);
3960
3961 let list_view: GenericListViewArray<i32> = builder.finish();
3962
3963 let values = Int32Array::from(vec![10, 20, 30]);
3965
3966 let schema = Arc::new(Schema::new(vec![
3968 Field::new("a", list_view.data_type().clone(), true),
3969 Field::new("b", DataType::Int32, false),
3970 ]));
3971 let batch =
3973 RecordBatch::try_new(schema, vec![Arc::new(list_view), Arc::new(values.clone())])
3974 .unwrap();
3975
3976 let mut buf = Vec::new();
3978 {
3979 let mut writer = FileWriter::try_new(&mut buf, &batch.schema()).unwrap();
3980 writer.write(&batch).unwrap();
3981 writer.finish().unwrap();
3982 }
3983
3984 let mut reader = FileReader::try_new(std::io::Cursor::new(buf), Some(vec![1])).unwrap();
3986 let read_batch = reader.next().unwrap().unwrap();
3987
3988 assert_eq!(read_batch.num_columns(), 1);
3990 assert_eq!(read_batch.column(0).as_ref(), &values);
3991 }
3992
3993 #[test]
3996 fn test_projection_skip_union_v4() {
3997 use crate::MetadataVersion;
3998 use crate::reader::FileReader;
3999 use crate::writer::{FileWriter, IpcWriteOptions};
4000 use arrow_array::{
4001 ArrayRef, Int32Array, RecordBatch, builder::UnionBuilder, types::Int32Type,
4002 };
4003 use arrow_schema::{DataType, Field, Schema};
4004 use std::sync::Arc;
4005
4006 let mut builder = UnionBuilder::new_dense();
4008 builder.append::<Int32Type>("a", 1).unwrap();
4009 builder.append::<Int32Type>("a", 2).unwrap();
4010 builder.append::<Int32Type>("a", 3).unwrap();
4011 let union = builder.build().unwrap();
4012
4013 let values = Int32Array::from(vec![10, 20, 30]);
4015
4016 let schema = Arc::new(Schema::new(vec![
4018 Field::new("union", union.data_type().clone(), false),
4019 Field::new("values", DataType::Int32, false),
4020 ]));
4021
4022 let batch = RecordBatch::try_new(
4024 schema,
4025 vec![Arc::new(union) as ArrayRef, Arc::new(values.clone())],
4026 )
4027 .unwrap();
4028
4029 let mut buf = Vec::new();
4031 {
4032 let options = IpcWriteOptions::try_new(8, false, MetadataVersion::V4).unwrap();
4033 let mut writer =
4034 FileWriter::try_new_with_options(&mut buf, &batch.schema(), options).unwrap();
4035 writer.write(&batch).unwrap();
4036 writer.finish().unwrap();
4037 }
4038 let mut reader = FileReader::try_new(std::io::Cursor::new(buf), Some(vec![1])).unwrap();
4040 let read_batch = reader.next().unwrap().unwrap();
4041
4042 assert_eq!(read_batch.num_columns(), 1);
4044 assert_eq!(read_batch.column(0).as_ref(), &values);
4045 }
4046
4047 #[test]
4051 fn test_projection_skip_fixed_width_types() {
4052 use std::sync::Arc;
4053
4054 use arrow_array::{ArrayRef, BooleanArray, Int32Array, RecordBatch, make_array};
4055 use arrow_buffer::{Buffer, MutableBuffer};
4056 use arrow_data::ArrayData;
4057 use arrow_schema::{DataType, Field, IntervalUnit, Schema, TimeUnit};
4058
4059 use crate::reader::FileReader;
4060 use crate::writer::FileWriter;
4061
4062 fn make_array_for_type(data_type: DataType) -> ArrayRef {
4064 let len = 3;
4065
4066 if matches!(data_type, DataType::Boolean) {
4067 return Arc::new(BooleanArray::from(vec![true, false, true]));
4068 }
4069
4070 let width = data_type.primitive_width().unwrap();
4071 let data = ArrayData::builder(data_type)
4072 .len(len)
4073 .add_buffer(Buffer::from(MutableBuffer::from_len_zeroed(len * width)))
4074 .build()
4075 .unwrap();
4076
4077 make_array(data)
4078 }
4079
4080 let data_types = vec![
4082 DataType::Boolean,
4083 DataType::Int8,
4084 DataType::Int16,
4085 DataType::Int32,
4086 DataType::Int64,
4087 DataType::UInt8,
4088 DataType::UInt16,
4089 DataType::UInt32,
4090 DataType::UInt64,
4091 DataType::Float16,
4092 DataType::Float32,
4093 DataType::Float64,
4094 DataType::Timestamp(TimeUnit::Second, None),
4095 DataType::Date32,
4096 DataType::Date64,
4097 DataType::Time32(TimeUnit::Second),
4098 DataType::Time64(TimeUnit::Microsecond),
4099 DataType::Duration(TimeUnit::Second),
4100 DataType::Interval(IntervalUnit::YearMonth),
4101 DataType::Interval(IntervalUnit::DayTime),
4102 DataType::Interval(IntervalUnit::MonthDayNano),
4103 DataType::Decimal32(9, 2),
4104 DataType::Decimal64(18, 2),
4105 DataType::Decimal128(38, 2),
4106 DataType::Decimal256(76, 2),
4107 ];
4108
4109 for data_type in data_types {
4114 let skipped = make_array_for_type(data_type.clone());
4115 let values = Int32Array::from(vec![10, 20, 30]);
4116
4117 let schema = Arc::new(Schema::new(vec![
4118 Field::new("skipped", data_type, false),
4119 Field::new("values", DataType::Int32, false),
4120 ]));
4121
4122 let batch =
4123 RecordBatch::try_new(schema, vec![skipped, Arc::new(values.clone())]).unwrap();
4124
4125 let mut buf = Vec::new();
4127 {
4128 let mut writer = FileWriter::try_new(&mut buf, &batch.schema()).unwrap();
4129 writer.write(&batch).unwrap();
4130 writer.finish().unwrap();
4131 }
4132
4133 let mut reader = FileReader::try_new(std::io::Cursor::new(buf), Some(vec![1])).unwrap();
4135 let read_batch = reader.next().unwrap().unwrap();
4136
4137 assert_eq!(read_batch.num_columns(), 1);
4139 assert_eq!(read_batch.column(0).as_ref(), &values);
4140 }
4141 }
4142
4143 fn stream_with_declared_body(body_length: i64, body_bytes: usize) -> Vec<u8> {
4147 let mut fbb = flatbuffers::FlatBufferBuilder::new();
4148 let mut message = crate::MessageBuilder::new(&mut fbb);
4149 message.add_version(crate::MetadataVersion::V5);
4150 message.add_header_type(crate::MessageHeader::NONE);
4151 message.add_bodyLength(body_length);
4152 let root = message.finish();
4153 fbb.finish(root, None);
4154 let metadata = fbb.finished_data();
4155
4156 let mut stream = Vec::new();
4157 stream.extend_from_slice(&CONTINUATION_MARKER);
4158 stream.extend_from_slice(&(metadata.len() as i32).to_le_bytes());
4159 stream.extend_from_slice(metadata);
4160 stream.resize(stream.len() + body_bytes, 0xAB);
4161 stream
4162 }
4163
4164 #[test]
4169 fn test_stream_reader_rejects_implausible_body_length() {
4170 for body_length in [i64::MAX, 1 << 50, -1] {
4171 let stream = stream_with_declared_body(body_length, 0);
4172 let err = StreamReader::try_new(std::io::Cursor::new(stream), None).expect_err(
4173 &format!("a message declaring {body_length} body bytes must not be accepted"),
4174 );
4175 let err = err.to_string();
4176 assert!(
4177 err.contains("Invalid IPC message body length")
4178 || err.contains("Unexpected end of stream")
4179 || err.contains("failed to fill whole buffer"),
4180 "unexpected error for body_length {body_length}: {err}"
4181 );
4182 }
4183 }
4184
4185 #[test]
4189 #[cfg_attr(miri, ignore)] fn test_stream_reader_rejects_truncated_body() {
4191 let over_prealloc = MAX_PREALLOC_BYTES as i64 + 1;
4192 for (body_length, body_bytes) in [(1024, 10), (over_prealloc, MAX_PREALLOC_BYTES)] {
4193 let stream = stream_with_declared_body(body_length, body_bytes);
4194 let err = StreamReader::try_new(std::io::Cursor::new(stream), None)
4195 .expect_err("a body backed by fewer bytes than declared must not be accepted");
4196 assert!(
4197 err.to_string().contains("failed to fill whole buffer"),
4198 "unexpected error for body_length {body_length}: {err}"
4199 );
4200 }
4201 }
4202
4203 #[test]
4205 fn test_stream_reader_rejects_unbacked_metadata_length() {
4206 let mut stream = Vec::new();
4207 stream.extend_from_slice(&CONTINUATION_MARKER);
4208 stream.extend_from_slice(&i32::MAX.to_le_bytes());
4209 stream.extend_from_slice(b"not this many bytes");
4210
4211 let err = StreamReader::try_new(std::io::Cursor::new(stream), None)
4212 .expect_err("metadata length must be backed by the stream");
4213 assert!(
4214 err.to_string().contains("Unexpected end of stream"),
4215 "unexpected error: {err}"
4216 );
4217 }
4218}