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