1use crate::bit_iterator::BitSliceIterator;
22use arrow_buffer::buffer::{BooleanBuffer, NullBuffer};
23use arrow_buffer::{
24 ArrowNativeType, Buffer, IntervalDayTime, IntervalMonthDayNano, MutableBuffer, bit_util, i256,
25};
26use arrow_schema::{ArrowError, DataType, UnionMode};
27use std::mem;
28use std::ops::Range;
29use std::sync::Arc;
30
31use crate::{equal, validate_binary_view, validate_string_view};
32
33#[inline]
34pub(crate) fn contains_nulls(
35 null_bit_buffer: Option<&NullBuffer>,
36 offset: usize,
37 len: usize,
38) -> bool {
39 match null_bit_buffer {
40 Some(buffer) => {
41 match BitSliceIterator::new(buffer.validity(), buffer.offset() + offset, len).next() {
42 Some((start, end)) => start != 0 || end != len,
43 None => len != 0, }
45 }
46 None => false, }
48}
49
50#[inline]
51pub(crate) fn count_nulls(
52 null_bit_buffer: Option<&NullBuffer>,
53 offset: usize,
54 len: usize,
55) -> usize {
56 if let Some(buf) = null_bit_buffer {
57 let buffer = buf.buffer();
58 len - buffer.count_set_bits_offset(offset + buf.offset(), len)
59 } else {
60 0
61 }
62}
63
64#[inline]
66pub(crate) fn new_buffers(data_type: &DataType, capacity: usize) -> [MutableBuffer; 2] {
67 let empty_buffer = MutableBuffer::new(0);
68 match data_type {
69 DataType::Null => [empty_buffer, MutableBuffer::new(0)],
70 DataType::Boolean => {
71 let bytes = bit_util::ceil(capacity, 8);
72 let buffer = MutableBuffer::new(bytes);
73 [buffer, empty_buffer]
74 }
75 DataType::UInt8
76 | DataType::UInt16
77 | DataType::UInt32
78 | DataType::UInt64
79 | DataType::Int8
80 | DataType::Int16
81 | DataType::Int32
82 | DataType::Int64
83 | DataType::Float16
84 | DataType::Float32
85 | DataType::Float64
86 | DataType::Decimal32(_, _)
87 | DataType::Decimal64(_, _)
88 | DataType::Decimal128(_, _)
89 | DataType::Decimal256(_, _)
90 | DataType::Date32
91 | DataType::Time32(_)
92 | DataType::Date64
93 | DataType::Time64(_)
94 | DataType::Duration(_)
95 | DataType::Timestamp(_, _)
96 | DataType::Interval(_) => [
97 MutableBuffer::new(capacity * data_type.primitive_width().unwrap()),
98 empty_buffer,
99 ],
100 DataType::Utf8 | DataType::Binary => {
101 let mut buffer = MutableBuffer::new((1 + capacity) * mem::size_of::<i32>());
102 buffer.push(0i32);
104 [buffer, MutableBuffer::new(capacity * mem::size_of::<u8>())]
105 }
106 DataType::LargeUtf8 | DataType::LargeBinary => {
107 let mut buffer = MutableBuffer::new((1 + capacity) * mem::size_of::<i64>());
108 buffer.push(0i64);
110 [buffer, MutableBuffer::new(capacity * mem::size_of::<u8>())]
111 }
112 DataType::BinaryView | DataType::Utf8View => [
113 MutableBuffer::new(capacity * mem::size_of::<u128>()),
114 empty_buffer,
115 ],
116 DataType::List(_) | DataType::Map(_, _) => {
117 let mut buffer = MutableBuffer::new((1 + capacity) * mem::size_of::<i32>());
119 buffer.push(0i32);
120 [buffer, empty_buffer]
121 }
122 DataType::ListView(_) => [
123 MutableBuffer::new(capacity * mem::size_of::<i32>()),
124 MutableBuffer::new(capacity * mem::size_of::<i32>()),
125 ],
126 DataType::LargeList(_) => {
127 let mut buffer = MutableBuffer::new((1 + capacity) * mem::size_of::<i64>());
129 buffer.push(0i64);
130 [buffer, empty_buffer]
131 }
132 DataType::LargeListView(_) => [
133 MutableBuffer::new(capacity * mem::size_of::<i64>()),
134 MutableBuffer::new(capacity * mem::size_of::<i64>()),
135 ],
136 DataType::FixedSizeBinary(size) => {
137 if *size < 0 {
138 panic!("cannot construct buffers from FixedSizeBinary({size})");
139 }
140 [MutableBuffer::new(capacity * *size as usize), empty_buffer]
141 }
142 DataType::Dictionary(k, _) => [
143 MutableBuffer::new(capacity * k.primitive_width().unwrap()),
144 empty_buffer,
145 ],
146 DataType::FixedSizeList(_, _) | DataType::Struct(_) | DataType::RunEndEncoded(_, _) => {
147 [empty_buffer, MutableBuffer::new(0)]
148 }
149 DataType::Union(_, mode) => {
150 let type_ids = MutableBuffer::new(capacity * mem::size_of::<i8>());
151 match mode {
152 UnionMode::Sparse => [type_ids, empty_buffer],
153 UnionMode::Dense => {
154 let offsets = MutableBuffer::new(capacity * mem::size_of::<i32>());
155 [type_ids, offsets]
156 }
157 }
158 }
159 }
160}
161
162#[derive(Debug, Clone)]
208pub struct ArrayData {
209 data_type: DataType,
211
212 len: usize,
214
215 offset: usize,
220
221 buffers: Vec<Buffer>,
234
235 child_data: Vec<ArrayData>,
245
246 nulls: Option<NullBuffer>,
254}
255
256pub type ArrayDataRef = Arc<ArrayData>;
258
259fn checked_len_plus_offset(
260 data_type: &DataType,
261 len: usize,
262 offset: usize,
263) -> Result<usize, ArrowError> {
264 len.checked_add(offset).ok_or_else(|| {
265 ArrowError::InvalidArgumentError(format!(
266 "Length {len} with offset {offset} overflows usize for {data_type}"
267 ))
268 })
269}
270
271impl ArrayData {
272 pub unsafe fn new_unchecked(
289 data_type: DataType,
290 len: usize,
291 null_count: Option<usize>,
292 null_bit_buffer: Option<Buffer>,
293 offset: usize,
294 buffers: Vec<Buffer>,
295 child_data: Vec<ArrayData>,
296 ) -> Self {
297 let builder = Self::inner_new_builder(
298 data_type,
299 len,
300 null_count,
301 null_bit_buffer,
302 offset,
303 buffers,
304 child_data,
305 );
306
307 unsafe { builder.build_unchecked() }
309 }
310
311 pub fn try_new(
325 data_type: DataType,
326 len: usize,
327 null_bit_buffer: Option<Buffer>,
328 offset: usize,
329 buffers: Vec<Buffer>,
330 child_data: Vec<ArrayData>,
331 ) -> Result<Self, ArrowError> {
332 if let Some(null_bit_buffer) = null_bit_buffer.as_ref() {
336 let len_plus_offset = checked_len_plus_offset(&data_type, len, offset)?;
337 let needed_len = bit_util::ceil(len_plus_offset, 8);
338 if null_bit_buffer.len() < needed_len {
339 return Err(ArrowError::InvalidArgumentError(format!(
340 "null_bit_buffer size too small. got {} needed {}",
341 null_bit_buffer.len(),
342 needed_len
343 )));
344 }
345 }
346
347 let builder = Self::inner_new_builder(
348 data_type,
349 len,
350 None,
351 null_bit_buffer,
352 offset,
353 buffers,
354 child_data,
355 );
356
357 assert!(!builder.skip_validation.get());
358
359 builder.build()
364 }
365
366 fn inner_new_builder(
367 data_type: DataType,
368 len: usize,
369 null_count: Option<usize>,
370 null_bit_buffer: Option<Buffer>,
371 offset: usize,
372 buffers: Vec<Buffer>,
373 child_data: Vec<ArrayData>,
374 ) -> ArrayDataBuilder {
375 ArrayDataBuilder {
376 data_type,
377 len,
378 null_count,
379 null_bit_buffer,
380 nulls: None,
381 offset,
382 buffers,
383 child_data,
384 align_buffers: false,
385 skip_validation: UnsafeFlag::new(),
386 }
387 }
388
389 pub fn into_parts(
395 self,
396 ) -> (
397 DataType,
398 usize,
399 Option<NullBuffer>,
400 usize,
401 Vec<Buffer>,
402 Vec<ArrayData>,
403 ) {
404 let Self {
405 data_type,
406 len,
407 nulls,
408 offset,
409 buffers,
410 child_data,
411 } = self;
412
413 (data_type, len, nulls, offset, buffers, child_data)
414 }
415
416 #[inline]
418 pub const fn builder(data_type: DataType) -> ArrayDataBuilder {
419 ArrayDataBuilder::new(data_type)
420 }
421
422 #[inline]
424 pub const fn data_type(&self) -> &DataType {
425 &self.data_type
426 }
427
428 pub fn buffers(&self) -> &[Buffer] {
430 &self.buffers
431 }
432
433 pub fn child_data(&self) -> &[ArrayData] {
436 &self.child_data[..]
437 }
438
439 #[inline]
441 pub fn is_null(&self, i: usize) -> bool {
442 match &self.nulls {
443 Some(v) => v.is_null(i),
444 None => false,
445 }
446 }
447
448 #[inline]
452 pub fn nulls(&self) -> Option<&NullBuffer> {
453 self.nulls.as_ref()
454 }
455
456 #[inline]
458 pub fn is_valid(&self, i: usize) -> bool {
459 !self.is_null(i)
460 }
461
462 #[inline]
464 pub const fn len(&self) -> usize {
465 self.len
466 }
467
468 #[inline]
470 pub const fn is_empty(&self) -> bool {
471 self.len == 0
472 }
473
474 #[inline]
476 pub const fn offset(&self) -> usize {
477 self.offset
478 }
479
480 #[inline]
482 pub fn null_count(&self) -> usize {
483 self.nulls
484 .as_ref()
485 .map(|x| x.null_count())
486 .unwrap_or_default()
487 }
488
489 pub fn get_buffer_memory_size(&self) -> usize {
501 let mut size = 0;
502 for buffer in &self.buffers {
503 size += buffer.capacity();
504 }
505 if let Some(bitmap) = &self.nulls {
506 size += bitmap.buffer().capacity()
507 }
508 for child in &self.child_data {
509 size += child.get_buffer_memory_size();
510 }
511 size
512 }
513
514 pub fn get_slice_memory_size(&self) -> Result<usize, ArrowError> {
527 let mut result: usize = 0;
528 let layout = layout(&self.data_type);
529
530 for spec in layout.buffers.iter() {
531 match spec {
532 BufferSpec::FixedWidth { byte_width, .. } => {
533 let len = match self.data_type {
536 DataType::Utf8
537 | DataType::LargeUtf8
538 | DataType::Binary
539 | DataType::LargeBinary
540 | DataType::List(_)
541 | DataType::LargeList(_)
542 | DataType::Map(_, _) => self.len + 1,
543 _ => self.len,
544 };
545 let buffer_size = len.checked_mul(*byte_width).ok_or_else(|| {
546 ArrowError::ComputeError(
547 "Integer overflow computing buffer size".to_string(),
548 )
549 })?;
550 result += buffer_size;
551 }
552 BufferSpec::VariableWidth => {
553 let buffer_len = match self.data_type {
554 DataType::Utf8 | DataType::Binary => {
555 let offsets = self.typed_offsets::<i32>()?;
556 (offsets[self.len] - offsets[0]) as usize
557 }
558 DataType::LargeUtf8 | DataType::LargeBinary => {
559 let offsets = self.typed_offsets::<i64>()?;
560 (offsets[self.len] - offsets[0]) as usize
561 }
562 _ => {
563 return Err(ArrowError::NotYetImplemented(format!(
564 "Invalid data type for VariableWidth buffer. Expected Utf8, LargeUtf8, Binary or LargeBinary. Got {}",
565 self.data_type
566 )));
567 }
568 };
569 result += buffer_len;
570 }
571 BufferSpec::BitMap => {
572 let buffer_size = bit_util::ceil(self.len, 8);
573 result += buffer_size;
574 }
575 BufferSpec::AlwaysNull => {
576 }
578 }
579 }
580
581 if self.nulls().is_some() {
582 result += bit_util::ceil(self.len, 8);
583 }
584
585 for child in &self.child_data {
586 result += child.get_slice_memory_size()?;
587 }
588 Ok(result)
589 }
590
591 pub fn get_array_memory_size(&self) -> usize {
600 let mut size = mem::size_of_val(self);
601
602 for buffer in &self.buffers {
604 size += mem::size_of::<Buffer>();
605 size += buffer.capacity();
606 }
607 if let Some(nulls) = &self.nulls {
608 size += nulls.buffer().capacity();
609 }
610 for child in &self.child_data {
611 size += child.get_array_memory_size();
612 }
613
614 size
615 }
616
617 pub fn slice(&self, offset: usize, length: usize) -> ArrayData {
625 let end = offset
626 .checked_add(length)
627 .expect("offset + length overflow");
628 assert!(end <= self.len());
629
630 if let DataType::Struct(_) = self.data_type() {
631 let new_offset = self.offset + offset;
633 ArrayData {
634 data_type: self.data_type().clone(),
635 len: length,
636 offset: new_offset,
637 buffers: self.buffers.clone(),
638 child_data: self
640 .child_data()
641 .iter()
642 .map(|data| data.slice(offset, length))
643 .collect(),
644 nulls: self.nulls.as_ref().map(|x| x.slice(offset, length)),
645 }
646 } else {
647 let mut new_data = self.clone();
648
649 new_data.len = length;
650 new_data.offset = offset + self.offset;
651 new_data.nulls = self.nulls.as_ref().map(|x| x.slice(offset, length));
652
653 new_data
654 }
655 }
656
657 pub fn buffer<T: ArrowNativeType>(&self, buffer: usize) -> &[T] {
664 &self.buffers()[buffer].typed_data()[self.offset..]
665 }
666
667 pub fn new_null(data_type: &DataType, len: usize) -> Self {
673 let bit_len = bit_util::ceil(len, 8);
674 let zeroed = |len: usize| Buffer::from(MutableBuffer::from_len_zeroed(len));
675
676 let (buffers, child_data, has_nulls) = match data_type.primitive_width() {
677 Some(width) => (vec![zeroed(width * len)], vec![], true),
678 None => match data_type {
679 DataType::Null => (vec![], vec![], false),
680 DataType::Boolean => (vec![zeroed(bit_len)], vec![], true),
681 DataType::Binary | DataType::Utf8 => {
682 (vec![zeroed((len + 1) * 4), zeroed(0)], vec![], true)
683 }
684 DataType::BinaryView | DataType::Utf8View => (vec![zeroed(len * 16)], vec![], true),
685 DataType::LargeBinary | DataType::LargeUtf8 => {
686 (vec![zeroed((len + 1) * 8), zeroed(0)], vec![], true)
687 }
688 DataType::FixedSizeBinary(i) => {
689 if *i < 0 {
690 panic!("cannot construct null data from FixedSizeBinary({i})");
691 }
692 (vec![zeroed(*i as usize * len)], vec![], true)
693 }
694 DataType::List(f) | DataType::Map(f, _) => (
695 vec![zeroed((len + 1) * 4)],
696 vec![ArrayData::new_empty(f.data_type())],
697 true,
698 ),
699 DataType::LargeList(f) => (
700 vec![zeroed((len + 1) * 8)],
701 vec![ArrayData::new_empty(f.data_type())],
702 true,
703 ),
704 DataType::ListView(f) => (
705 vec![zeroed(len * 4), zeroed(len * 4)],
706 vec![ArrayData::new_empty(f.data_type())],
707 true,
708 ),
709 DataType::LargeListView(f) => (
710 vec![zeroed(len * 8), zeroed(len * 8)],
711 vec![ArrayData::new_empty(f.data_type())],
712 true,
713 ),
714 DataType::FixedSizeList(f, list_len) => (
715 vec![],
716 vec![ArrayData::new_null(f.data_type(), *list_len as usize * len)],
717 true,
718 ),
719 DataType::Struct(fields) => (
720 vec![],
721 fields
722 .iter()
723 .map(|f| Self::new_null(f.data_type(), len))
724 .collect(),
725 true,
726 ),
727 DataType::Dictionary(k, v) => (
728 vec![zeroed(k.primitive_width().unwrap() * len)],
729 vec![ArrayData::new_empty(v.as_ref())],
730 true,
731 ),
732 DataType::Union(f, mode) => {
733 let (id, _) = f.iter().next().unwrap();
734 let ids = Buffer::from_iter(std::iter::repeat_n(id, len));
735 let buffers = match mode {
736 UnionMode::Sparse => vec![ids],
737 UnionMode::Dense => {
738 let end_offset = i32::from_usize(len).unwrap();
739 vec![ids, Buffer::from_iter(0_i32..end_offset)]
740 }
741 };
742
743 let children = f
744 .iter()
745 .enumerate()
746 .map(|(idx, (_, f))| {
747 if idx == 0 || *mode == UnionMode::Sparse {
748 Self::new_null(f.data_type(), len)
749 } else {
750 Self::new_empty(f.data_type())
751 }
752 })
753 .collect();
754
755 (buffers, children, false)
756 }
757 DataType::RunEndEncoded(r, v) => {
758 if len == 0 {
759 let runs = ArrayData::new_empty(r.data_type());
761 let values = ArrayData::new_empty(v.data_type());
762 (vec![], vec![runs, values], false)
763 } else {
764 let runs = match r.data_type() {
765 DataType::Int16 => {
766 let i = i16::from_usize(len).expect("run overflow");
767 Buffer::from_slice_ref([i])
768 }
769 DataType::Int32 => {
770 let i = i32::from_usize(len).expect("run overflow");
771 Buffer::from_slice_ref([i])
772 }
773 DataType::Int64 => {
774 let i = i64::from_usize(len).expect("run overflow");
775 Buffer::from_slice_ref([i])
776 }
777 dt => unreachable!("Invalid run ends data type {dt}"),
778 };
779
780 let builder = ArrayData::builder(r.data_type().clone())
781 .len(1)
782 .buffers(vec![runs]);
783
784 let runs = unsafe { builder.build_unchecked() };
787 (
788 vec![],
789 vec![runs, ArrayData::new_null(v.data_type(), 1)],
790 false,
791 )
792 }
793 }
794 DataType::Int8
796 | DataType::Int16
797 | DataType::Int32
798 | DataType::Int64
799 | DataType::UInt8
800 | DataType::UInt16
801 | DataType::UInt32
802 | DataType::UInt64
803 | DataType::Float16
804 | DataType::Float32
805 | DataType::Float64
806 | DataType::Timestamp(_, _)
807 | DataType::Date32
808 | DataType::Date64
809 | DataType::Time32(_)
810 | DataType::Time64(_)
811 | DataType::Duration(_)
812 | DataType::Interval(_)
813 | DataType::Decimal32(_, _)
814 | DataType::Decimal64(_, _)
815 | DataType::Decimal128(_, _)
816 | DataType::Decimal256(_, _) => unreachable!("{data_type}"),
817 },
818 };
819
820 let mut builder = ArrayDataBuilder::new(data_type.clone())
821 .len(len)
822 .buffers(buffers)
823 .child_data(child_data);
824
825 if has_nulls {
826 builder = builder.nulls(Some(NullBuffer::new_null(len)))
827 }
828
829 unsafe { builder.build_unchecked() }
832 }
833
834 pub fn new_empty(data_type: &DataType) -> Self {
836 Self::new_null(data_type, 0)
837 }
838
839 pub fn align_buffers(&mut self) {
848 let layout = layout(&self.data_type);
849 for (buffer, spec) in self.buffers.iter_mut().zip(&layout.buffers) {
850 if let BufferSpec::FixedWidth { alignment, .. } = spec {
851 if buffer.as_ptr().align_offset(*alignment) != 0 {
852 *buffer = Buffer::from_slice_ref(buffer.as_ref());
853 }
854 }
855 }
856 for data in self.child_data.iter_mut() {
858 data.align_buffers()
859 }
860 }
861
862 pub fn validate(&self) -> Result<(), ArrowError> {
873 let len_plus_offset = checked_len_plus_offset(&self.data_type, self.len, self.offset)?;
875
876 let layout = layout(&self.data_type);
878
879 if !layout.can_contain_null_mask && self.nulls.is_some() {
880 return Err(ArrowError::InvalidArgumentError(format!(
881 "Arrays of type {:?} cannot contain a null bitmask",
882 self.data_type,
883 )));
884 }
885
886 if self.buffers.len() < layout.buffers.len()
888 || (!layout.variadic && self.buffers.len() != layout.buffers.len())
889 {
890 return Err(ArrowError::InvalidArgumentError(format!(
891 "Expected {} buffers in array of type {:?}, got {}",
892 layout.buffers.len(),
893 self.data_type,
894 self.buffers.len(),
895 )));
896 }
897
898 for (i, (buffer, spec)) in self.buffers.iter().zip(layout.buffers.iter()).enumerate() {
899 match spec {
900 BufferSpec::FixedWidth {
901 byte_width,
902 alignment,
903 } => {
904 let min_buffer_size = len_plus_offset.saturating_mul(*byte_width);
905
906 if buffer.len() < min_buffer_size {
907 return Err(ArrowError::InvalidArgumentError(format!(
908 "Need at least {} bytes in buffers[{}] in array of type {:?}, but got {}",
909 min_buffer_size,
910 i,
911 self.data_type,
912 buffer.len()
913 )));
914 }
915
916 let align_offset = buffer.as_ptr().align_offset(*alignment);
917 if align_offset != 0 {
918 return Err(ArrowError::InvalidArgumentError(format!(
919 "Misaligned buffers[{i}] in array of type {:?}, offset from expected alignment of {alignment} by {}",
920 self.data_type,
921 align_offset.min(alignment - align_offset)
922 )));
923 }
924 }
925 BufferSpec::VariableWidth => {
926 }
930 BufferSpec::BitMap => {
931 let min_buffer_size = bit_util::ceil(len_plus_offset, 8);
932 if buffer.len() < min_buffer_size {
933 return Err(ArrowError::InvalidArgumentError(format!(
934 "Need at least {} bytes for bitmap in buffers[{}] in array of type {:?}, but got {}",
935 min_buffer_size,
936 i,
937 self.data_type,
938 buffer.len()
939 )));
940 }
941 }
942 BufferSpec::AlwaysNull => {
943 }
945 }
946 }
947
948 if let Some(nulls) = self.nulls() {
950 if nulls.null_count() > self.len {
951 return Err(ArrowError::InvalidArgumentError(format!(
952 "null_count {} for an array exceeds length of {} elements",
953 nulls.null_count(),
954 self.len
955 )));
956 }
957
958 let actual_len = nulls.validity().len();
959 let needed_len = bit_util::ceil(len_plus_offset, 8);
960 if actual_len < needed_len {
961 return Err(ArrowError::InvalidArgumentError(format!(
962 "null_bit_buffer size too small. got {actual_len} needed {needed_len}",
963 )));
964 }
965
966 if nulls.len() != self.len {
967 return Err(ArrowError::InvalidArgumentError(format!(
968 "null buffer incorrect size. got {} expected {}",
969 nulls.len(),
970 self.len
971 )));
972 }
973 }
974
975 self.validate_child_data()?;
976
977 match &self.data_type {
979 DataType::Utf8 | DataType::Binary => {
980 self.validate_offsets::<i32>(self.buffers[1].len())?;
981 }
982 DataType::LargeUtf8 | DataType::LargeBinary => {
983 self.validate_offsets::<i64>(self.buffers[1].len())?;
984 }
985 DataType::Dictionary(key_type, _value_type) => {
986 if !DataType::is_dictionary_key_type(key_type) {
988 return Err(ArrowError::InvalidArgumentError(format!(
989 "Dictionary key type must be integer, but was {key_type}"
990 )));
991 }
992 }
993 DataType::RunEndEncoded(run_ends_type, _) => {
994 if run_ends_type.is_nullable() {
995 return Err(ArrowError::InvalidArgumentError(
996 "The nullable should be set to false for the field defining run_ends array.".to_string()
997 ));
998 }
999 if !DataType::is_run_ends_type(run_ends_type.data_type()) {
1000 return Err(ArrowError::InvalidArgumentError(format!(
1001 "RunArray run_ends types must be Int16, Int32 or Int64, but was {}",
1002 run_ends_type.data_type()
1003 )));
1004 }
1005 }
1006 DataType::Map(f, _) if f.is_nullable() => {
1007 return Err(ArrowError::InvalidArgumentError(
1008 "The nullable should be set to false for the map entries field.".to_string(),
1009 ));
1010 }
1011 _ => {}
1012 };
1013
1014 Ok(())
1015 }
1016
1017 fn typed_offsets<T: ArrowNativeType + num_traits::Num>(&self) -> Result<&[T], ArrowError> {
1024 if self.len == 0 && self.buffers[0].is_empty() {
1026 return Ok(&[]);
1027 }
1028
1029 let len = checked_len_plus_offset(&self.data_type, self.len, 1)?;
1030
1031 self.typed_buffer(0, len)
1032 }
1033
1034 fn typed_buffer<T: ArrowNativeType + num_traits::Num>(
1036 &self,
1037 idx: usize,
1038 len: usize,
1039 ) -> Result<&[T], ArrowError> {
1040 let buffer = &self.buffers[idx];
1041
1042 let required_elements = checked_len_plus_offset(&self.data_type, len, self.offset)?;
1043 let byte_width = mem::size_of::<T>();
1044 let required_len = required_elements.checked_mul(byte_width).ok_or_else(|| {
1045 ArrowError::InvalidArgumentError(format!(
1046 "Buffer {idx} of {} byte length overflow: {} elements of {} bytes exceeds usize",
1047 self.data_type, required_elements, byte_width
1048 ))
1049 })?;
1050
1051 if buffer.len() < required_len {
1052 return Err(ArrowError::InvalidArgumentError(format!(
1053 "Buffer {} of {} isn't large enough. Expected {} bytes got {}",
1054 idx,
1055 self.data_type,
1056 required_len,
1057 buffer.len()
1058 )));
1059 }
1060
1061 Ok(&buffer.typed_data::<T>()[self.offset..required_elements])
1062 }
1063
1064 fn validate_offsets<T: ArrowNativeType + num_traits::Num + std::fmt::Display>(
1067 &self,
1068 values_length: usize,
1069 ) -> Result<(), ArrowError> {
1070 let offsets = self.typed_offsets::<T>()?;
1072 if offsets.is_empty() {
1073 return Ok(());
1074 }
1075
1076 let first_offset = offsets[0].to_usize().ok_or_else(|| {
1077 ArrowError::InvalidArgumentError(format!(
1078 "Error converting offset[0] ({}) to usize for {}",
1079 offsets[0], self.data_type
1080 ))
1081 })?;
1082
1083 let last_offset = offsets[self.len].to_usize().ok_or_else(|| {
1084 ArrowError::InvalidArgumentError(format!(
1085 "Error converting offset[{}] ({}) to usize for {}",
1086 self.len, offsets[self.len], self.data_type
1087 ))
1088 })?;
1089
1090 if first_offset > values_length {
1091 return Err(ArrowError::InvalidArgumentError(format!(
1092 "First offset {} of {} is larger than values length {}",
1093 first_offset, self.data_type, values_length,
1094 )));
1095 }
1096
1097 if last_offset > values_length {
1098 return Err(ArrowError::InvalidArgumentError(format!(
1099 "Last offset {} of {} is larger than values length {}",
1100 last_offset, self.data_type, values_length,
1101 )));
1102 }
1103
1104 if first_offset > last_offset {
1105 return Err(ArrowError::InvalidArgumentError(format!(
1106 "First offset {} in {} is smaller than last offset {}",
1107 first_offset, self.data_type, last_offset,
1108 )));
1109 }
1110
1111 Ok(())
1112 }
1113
1114 fn validate_offsets_and_sizes<T: ArrowNativeType + num_traits::Num + std::fmt::Display>(
1117 &self,
1118 values_length: usize,
1119 ) -> Result<(), ArrowError> {
1120 let offsets: &[T] = self.typed_buffer(0, self.len)?;
1121 let sizes: &[T] = self.typed_buffer(1, self.len)?;
1122 if offsets.len() != sizes.len() {
1123 return Err(ArrowError::ComputeError(format!(
1124 "ListView offsets len {} does not match sizes len {}",
1125 offsets.len(),
1126 sizes.len()
1127 )));
1128 }
1129
1130 for i in 0..sizes.len() {
1131 let size = sizes[i].to_usize().ok_or_else(|| {
1132 ArrowError::InvalidArgumentError(format!(
1133 "Error converting size[{}] ({}) to usize for {}",
1134 i, sizes[i], self.data_type
1135 ))
1136 })?;
1137 let offset = offsets[i].to_usize().ok_or_else(|| {
1138 ArrowError::InvalidArgumentError(format!(
1139 "Error converting offset[{}] ({}) to usize for {}",
1140 i, offsets[i], self.data_type
1141 ))
1142 })?;
1143 if size
1144 .checked_add(offset)
1145 .expect("Offset and size have exceeded the usize boundary")
1146 > values_length
1147 {
1148 return Err(ArrowError::InvalidArgumentError(format!(
1149 "Size {} at index {} is larger than the remaining values for {}",
1150 size, i, self.data_type
1151 )));
1152 }
1153 }
1154 Ok(())
1155 }
1156
1157 fn validate_child_data(&self) -> Result<(), ArrowError> {
1159 match &self.data_type {
1160 DataType::List(field) => {
1161 let values_data = self.get_single_valid_child_data(field.data_type())?;
1162 self.validate_offsets::<i32>(values_data.len)?;
1163 Ok(())
1164 }
1165 DataType::LargeList(field) => {
1166 let values_data = self.get_single_valid_child_data(field.data_type())?;
1167 self.validate_offsets::<i64>(values_data.len)?;
1168 Ok(())
1169 }
1170 DataType::Map(field, _) => {
1171 let DataType::Struct(entries_fields) = field.data_type() else {
1172 return Err(ArrowError::InvalidArgumentError(format!(
1173 "Map field should be a entries struct data type, got {:?} instead",
1174 field.data_type()
1175 )));
1176 };
1177 if entries_fields.len() != 2 {
1178 return Err(ArrowError::InvalidArgumentError(format!(
1179 "Map entries data type should be a struct containing 2 fields, got {} fields",
1180 entries_fields.len()
1181 )));
1182 }
1183
1184 if entries_fields[0].is_nullable() {
1186 return Err(ArrowError::InvalidArgumentError(
1187 "Map key field must not be nullable".to_string(),
1188 ));
1189 }
1190 let values_data = self.get_single_valid_child_data(field.data_type())?;
1191 self.validate_offsets::<i32>(values_data.len)?;
1192 Ok(())
1193 }
1194 DataType::ListView(field) => {
1195 let values_data = self.get_single_valid_child_data(field.data_type())?;
1196 self.validate_offsets_and_sizes::<i32>(values_data.len)?;
1197 Ok(())
1198 }
1199 DataType::LargeListView(field) => {
1200 let values_data = self.get_single_valid_child_data(field.data_type())?;
1201 self.validate_offsets_and_sizes::<i64>(values_data.len)?;
1202 Ok(())
1203 }
1204 DataType::FixedSizeList(field, list_size) => {
1205 let values_data = self.get_single_valid_child_data(field.data_type())?;
1206
1207 let list_size: usize = (*list_size).try_into().map_err(|_| {
1208 ArrowError::InvalidArgumentError(format!(
1209 "{} has a negative list_size {}",
1210 self.data_type, list_size
1211 ))
1212 })?;
1213
1214 let expected_values_len = self.len
1215 .checked_mul(list_size)
1216 .expect("integer overflow computing expected number of expected values in FixedListSize");
1217
1218 if values_data.len < expected_values_len {
1219 return Err(ArrowError::InvalidArgumentError(format!(
1220 "Values length {} is less than the length ({}) multiplied by the value size ({}) for {}",
1221 values_data.len, self.len, list_size, self.data_type
1222 )));
1223 }
1224
1225 Ok(())
1226 }
1227 DataType::Struct(fields) => {
1228 self.validate_num_child_data(fields.len())?;
1229 for (i, field) in fields.iter().enumerate() {
1230 let field_data = self.get_valid_child_data(i, field.data_type())?;
1231
1232 if field_data.len < self.len {
1234 return Err(ArrowError::InvalidArgumentError(format!(
1235 "{} child array #{} for field {} has length smaller than expected for struct array ({} < {})",
1236 self.data_type,
1237 i,
1238 field.name(),
1239 field_data.len,
1240 self.len
1241 )));
1242 }
1243 }
1244 Ok(())
1245 }
1246 DataType::RunEndEncoded(run_ends_field, values_field) => {
1247 self.validate_num_child_data(2)?;
1248 let run_ends_data = self.get_valid_child_data(0, run_ends_field.data_type())?;
1249 let values_data = self.get_valid_child_data(1, values_field.data_type())?;
1250 if run_ends_data.len != values_data.len {
1251 return Err(ArrowError::InvalidArgumentError(format!(
1252 "The run_ends array length should be the same as values array length. Run_ends array length is {}, values array length is {}",
1253 run_ends_data.len, values_data.len
1254 )));
1255 }
1256 if run_ends_data.nulls.is_some() {
1257 return Err(ArrowError::InvalidArgumentError(
1258 "Found null values in run_ends array. The run_ends array should not have null values.".to_string(),
1259 ));
1260 }
1261 Ok(())
1262 }
1263 DataType::Union(fields, mode) => {
1264 self.validate_num_child_data(fields.len())?;
1265
1266 for (i, (_, field)) in fields.iter().enumerate() {
1267 let field_data = self.get_valid_child_data(i, field.data_type())?;
1268
1269 if mode == &UnionMode::Sparse {
1270 let len_plus_offset =
1271 checked_len_plus_offset(&self.data_type, self.len, self.offset)?;
1272 if field_data.len < len_plus_offset {
1273 return Err(ArrowError::InvalidArgumentError(format!(
1274 "Sparse union child array #{} has length smaller than expected for union array ({} < {})",
1275 i, field_data.len, len_plus_offset
1276 )));
1277 }
1278 }
1279 }
1280 Ok(())
1281 }
1282 DataType::Dictionary(_key_type, value_type) => {
1283 self.get_single_valid_child_data(value_type)?;
1284 Ok(())
1285 }
1286 _ => {
1287 if !self.child_data.is_empty() {
1289 return Err(ArrowError::InvalidArgumentError(format!(
1290 "Expected no child arrays for type {} but got {}",
1291 self.data_type,
1292 self.child_data.len()
1293 )));
1294 }
1295 Ok(())
1296 }
1297 }
1298 }
1299
1300 fn get_single_valid_child_data(
1304 &self,
1305 expected_type: &DataType,
1306 ) -> Result<&ArrayData, ArrowError> {
1307 self.validate_num_child_data(1)?;
1308 self.get_valid_child_data(0, expected_type)
1309 }
1310
1311 fn validate_num_child_data(&self, expected_len: usize) -> Result<(), ArrowError> {
1313 if self.child_data.len() != expected_len {
1314 Err(ArrowError::InvalidArgumentError(format!(
1315 "Value data for {} should contain {} child data array(s), had {}",
1316 self.data_type,
1317 expected_len,
1318 self.child_data.len()
1319 )))
1320 } else {
1321 Ok(())
1322 }
1323 }
1324
1325 fn get_valid_child_data(
1328 &self,
1329 i: usize,
1330 expected_type: &DataType,
1331 ) -> Result<&ArrayData, ArrowError> {
1332 let values_data = self.child_data.get(i).ok_or_else(|| {
1333 ArrowError::InvalidArgumentError(format!(
1334 "{} did not have enough child arrays. Expected at least {} but had only {}",
1335 self.data_type,
1336 i + 1,
1337 self.child_data.len()
1338 ))
1339 })?;
1340
1341 if expected_type != &values_data.data_type {
1342 return Err(ArrowError::InvalidArgumentError(format!(
1343 "Child type mismatch for {}. Expected {} but child data had {}",
1344 self.data_type, expected_type, values_data.data_type
1345 )));
1346 }
1347
1348 values_data.validate()?;
1349 Ok(values_data)
1350 }
1351
1352 pub fn validate_data(&self) -> Result<(), ArrowError> {
1368 self.validate()?;
1369
1370 self.validate_nulls()?;
1371 self.validate_values()?;
1372 Ok(())
1373 }
1374
1375 pub fn validate_full(&self) -> Result<(), ArrowError> {
1380 self.validate_data()?;
1381 self.child_data
1383 .iter()
1384 .enumerate()
1385 .try_for_each(|(i, child_data)| {
1386 child_data.validate_full().map_err(|e| {
1387 ArrowError::InvalidArgumentError(format!(
1388 "{} child #{} invalid: {}",
1389 self.data_type, i, e
1390 ))
1391 })
1392 })?;
1393 Ok(())
1394 }
1395
1396 pub fn validate_nulls(&self) -> Result<(), ArrowError> {
1406 if let Some(nulls) = &self.nulls {
1407 let actual = nulls.len() - nulls.inner().count_set_bits();
1408 if actual != nulls.null_count() {
1409 return Err(ArrowError::InvalidArgumentError(format!(
1410 "null_count value ({}) doesn't match actual number of nulls in array ({})",
1411 nulls.null_count(),
1412 actual
1413 )));
1414 }
1415 }
1416
1417 match &self.data_type {
1422 DataType::List(f) | DataType::LargeList(f) | DataType::Map(f, _) => {
1423 if !f.is_nullable() {
1424 self.validate_non_nullable(None, &self.child_data[0])?
1425 }
1426 }
1427 DataType::FixedSizeList(field, len) => {
1428 let child = &self.child_data[0];
1429 if !field.is_nullable() {
1430 match &self.nulls {
1431 Some(nulls) => {
1432 let element_len = *len as usize;
1433 let expanded = nulls.expand(element_len);
1434 self.validate_non_nullable(Some(&expanded), child)?;
1435 }
1436 None => self.validate_non_nullable(None, child)?,
1437 }
1438 }
1439 }
1440 DataType::Struct(fields) => {
1441 for (field, child) in fields.iter().zip(&self.child_data) {
1442 if !field.is_nullable() {
1443 self.validate_non_nullable(self.nulls(), child)?
1444 }
1445 }
1446 }
1447 _ => {}
1448 }
1449
1450 Ok(())
1451 }
1452
1453 fn validate_non_nullable(
1455 &self,
1456 mask: Option<&NullBuffer>,
1457 child: &ArrayData,
1458 ) -> Result<(), ArrowError> {
1459 let mask = match mask {
1460 Some(mask) => mask,
1461 None => {
1462 return match child.null_count() {
1463 0 => Ok(()),
1464 _ => Err(ArrowError::InvalidArgumentError(format!(
1465 "non-nullable child of type {} contains nulls not present in parent {}",
1466 child.data_type, self.data_type
1467 ))),
1468 };
1469 }
1470 };
1471
1472 match child.nulls() {
1473 Some(nulls) if !mask.contains(nulls) => Err(ArrowError::InvalidArgumentError(format!(
1474 "non-nullable child of type {} contains nulls not present in parent",
1475 child.data_type
1476 ))),
1477 _ => Ok(()),
1478 }
1479 }
1480
1481 pub fn validate_values(&self) -> Result<(), ArrowError> {
1487 match &self.data_type {
1488 DataType::Utf8 => self.validate_utf8::<i32>(),
1489 DataType::LargeUtf8 => self.validate_utf8::<i64>(),
1490 DataType::Binary => self.validate_offsets_full::<i32>(self.buffers[1].len()),
1491 DataType::LargeBinary => self.validate_offsets_full::<i64>(self.buffers[1].len()),
1492 DataType::BinaryView => {
1493 let views = self.typed_buffer::<u128>(0, self.len)?;
1494 validate_binary_view(views, &self.buffers[1..])
1495 }
1496 DataType::Utf8View => {
1497 let views = self.typed_buffer::<u128>(0, self.len)?;
1498 validate_string_view(views, &self.buffers[1..])
1499 }
1500 DataType::List(_) | DataType::Map(_, _) => {
1501 let child = &self.child_data[0];
1502 self.validate_offsets_full::<i32>(child.len)
1503 }
1504 DataType::LargeList(_) => {
1505 let child = &self.child_data[0];
1506 self.validate_offsets_full::<i64>(child.len)
1507 }
1508 DataType::Union(fields, mode) => {
1509 let mut child_lens: Vec<Option<usize>> = vec![None; 128];
1510 for (index, (type_id, _)) in fields.iter().enumerate() {
1511 let child = &self.child_data[index];
1512 child_lens[type_id as usize] = Some(child.len());
1513 }
1514
1515 let type_ids = self.typed_buffer::<i8>(0, self.len)?;
1517 for (index, &type_id) in type_ids.iter().enumerate() {
1518 if type_id < 0 {
1519 return Err(ArrowError::InvalidArgumentError(format!(
1520 "Type id values cannot be negative, got {type_id} at index {index}"
1521 )));
1522 }
1523 let type_id = type_id as usize;
1524 if child_lens[type_id].is_none() {
1525 return Err(ArrowError::InvalidArgumentError(format!(
1526 "Type id values must match one of the field type ids, got {type_id} at index {index}"
1527 )));
1528 }
1529 }
1530
1531 if mode == &UnionMode::Dense {
1533 let offsets = self.typed_buffer::<i32>(1, self.len)?;
1534
1535 if type_ids
1536 .iter()
1537 .zip(offsets.iter())
1538 .any(|(&type_id, &offset)| {
1539 offset < 0 || offset as usize >= child_lens[type_id as usize].unwrap()
1540 })
1541 {
1542 return Err(ArrowError::InvalidArgumentError(
1543 "Offsets must be non-negative and within the length of the child array"
1544 .to_string(),
1545 ));
1546 }
1547 }
1548
1549 Ok(())
1550 }
1551 DataType::Dictionary(key_type, _value_type) => {
1552 let dictionary_length: i64 = self.child_data[0].len.try_into().unwrap();
1553 let max_value = dictionary_length - 1;
1554 match key_type.as_ref() {
1555 DataType::UInt8 => self.check_bounds::<u8>(max_value),
1556 DataType::UInt16 => self.check_bounds::<u16>(max_value),
1557 DataType::UInt32 => self.check_bounds::<u32>(max_value),
1558 DataType::UInt64 => self.check_bounds::<u64>(max_value),
1559 DataType::Int8 => self.check_bounds::<i8>(max_value),
1560 DataType::Int16 => self.check_bounds::<i16>(max_value),
1561 DataType::Int32 => self.check_bounds::<i32>(max_value),
1562 DataType::Int64 => self.check_bounds::<i64>(max_value),
1563 _ => unreachable!(),
1564 }
1565 }
1566 DataType::RunEndEncoded(run_ends, _values) => {
1567 let len_plus_offset =
1568 checked_len_plus_offset(&self.data_type, self.len, self.offset)?;
1569 let run_ends_data = self.child_data()[0].clone();
1570 match run_ends.data_type() {
1571 DataType::Int16 => {
1572 run_ends_data.check_run_ends::<i16>(self.len, len_plus_offset)
1573 }
1574 DataType::Int32 => {
1575 run_ends_data.check_run_ends::<i32>(self.len, len_plus_offset)
1576 }
1577 DataType::Int64 => {
1578 run_ends_data.check_run_ends::<i64>(self.len, len_plus_offset)
1579 }
1580 _ => unreachable!(),
1581 }
1582 }
1583 _ => {
1584 Ok(())
1586 }
1587 }
1588 }
1589
1590 fn validate_each_offset<T, V>(&self, offset_limit: usize, validate: V) -> Result<(), ArrowError>
1601 where
1602 T: ArrowNativeType + TryInto<usize> + num_traits::Num + std::fmt::Display,
1603 V: Fn(usize, Range<usize>) -> Result<(), ArrowError>,
1604 {
1605 self.typed_offsets::<T>()?
1606 .iter()
1607 .enumerate()
1608 .map(|(i, x)| {
1609 let r = x.to_usize().ok_or_else(|| {
1611 ArrowError::InvalidArgumentError(format!(
1612 "Offset invariant failure: Could not convert offset {x} to usize at position {i}"))}
1613 );
1614 match r {
1616 Ok(n) if n <= offset_limit => Ok((i, n)),
1617 Ok(_) => Err(ArrowError::InvalidArgumentError(format!(
1618 "Offset invariant failure: offset at position {i} out of bounds: {x} > {offset_limit}"))
1619 ),
1620 Err(e) => Err(e),
1621 }
1622 })
1623 .scan(0_usize, |start, end| {
1624 match end {
1626 Ok((i, end)) if *start <= end => {
1627 let range = Some(Ok((i, *start..end)));
1628 *start = end;
1629 range
1630 }
1631 Ok((i, end)) => Some(Err(ArrowError::InvalidArgumentError(format!(
1632 "Offset invariant failure: non-monotonic offset at slot {}: {} > {}",
1633 i - 1, start, end))
1634 )),
1635 Err(err) => Some(Err(err)),
1636 }
1637 })
1638 .skip(1) .try_for_each(|res: Result<(usize, Range<usize>), ArrowError>| {
1640 let (item_index, range) = res?;
1641 validate(item_index-1, range)
1642 })
1643 }
1644
1645 fn validate_utf8<T>(&self) -> Result<(), ArrowError>
1648 where
1649 T: ArrowNativeType + TryInto<usize> + num_traits::Num + std::fmt::Display,
1650 {
1651 let values_buffer = &self.buffers[1].as_slice();
1652 if let Ok(values_str) = std::str::from_utf8(values_buffer) {
1653 self.validate_each_offset::<T, _>(values_buffer.len(), |string_index, range| {
1655 if !values_str.is_char_boundary(range.start)
1656 || !values_str.is_char_boundary(range.end)
1657 {
1658 return Err(ArrowError::InvalidArgumentError(format!(
1659 "incomplete utf-8 byte sequence from index {string_index}"
1660 )));
1661 }
1662 Ok(())
1663 })
1664 } else {
1665 self.validate_each_offset::<T, _>(values_buffer.len(), |string_index, range| {
1667 std::str::from_utf8(&values_buffer[range.clone()]).map_err(|e| {
1668 ArrowError::InvalidArgumentError(format!(
1669 "Invalid UTF8 sequence at string index {string_index} ({range:?}): {e}"
1670 ))
1671 })?;
1672 Ok(())
1673 })
1674 }
1675 }
1676
1677 fn validate_offsets_full<T>(&self, offset_limit: usize) -> Result<(), ArrowError>
1680 where
1681 T: ArrowNativeType + TryInto<usize> + num_traits::Num + std::fmt::Display,
1682 {
1683 self.validate_each_offset::<T, _>(offset_limit, |_string_index, _range| {
1684 Ok(())
1687 })
1688 }
1689
1690 fn check_bounds<T>(&self, max_value: i64) -> Result<(), ArrowError>
1693 where
1694 T: ArrowNativeType + TryInto<i64> + num_traits::Num + std::fmt::Display,
1695 {
1696 let required_len = checked_len_plus_offset(&self.data_type, self.len, self.offset)?;
1697 let buffer = &self.buffers[0];
1698
1699 assert!(buffer.len() / mem::size_of::<T>() >= required_len);
1702
1703 let indexes: &[T] = &buffer.typed_data::<T>()[self.offset..required_len];
1705
1706 indexes.iter().enumerate().try_for_each(|(i, &dict_index)| {
1707 if self.is_null(i) {
1709 return Ok(());
1710 }
1711 let dict_index: i64 = dict_index.try_into().map_err(|_| {
1712 ArrowError::InvalidArgumentError(format!(
1713 "Value at position {i} out of bounds: {dict_index} (can not convert to i64)"
1714 ))
1715 })?;
1716
1717 if dict_index < 0 || dict_index > max_value {
1718 return Err(ArrowError::InvalidArgumentError(format!(
1719 "Value at position {i} out of bounds: {dict_index} (should be in [0, {max_value}])"
1720 )));
1721 }
1722 Ok(())
1723 })
1724 }
1725
1726 fn check_run_ends<T>(
1728 &self,
1729 parent_len: usize,
1730 parent_len_plus_offset: usize,
1731 ) -> Result<(), ArrowError>
1732 where
1733 T: ArrowNativeType + TryInto<i64> + num_traits::Num + std::fmt::Display,
1734 {
1735 let values = self.typed_buffer::<T>(0, self.len)?;
1736 let mut prev_value: i64 = 0_i64;
1737 values.iter().enumerate().try_for_each(|(ix, &inp_value)| {
1738 let value: i64 = inp_value.try_into().map_err(|_| {
1739 ArrowError::InvalidArgumentError(format!(
1740 "Value at position {ix} out of bounds: {inp_value} (can not convert to i64)"
1741 ))
1742 })?;
1743 if value <= 0_i64 {
1744 return Err(ArrowError::InvalidArgumentError(format!(
1745 "The values in run_ends array should be strictly positive. Found value {value} at index {ix} that does not match the criteria."
1746 )));
1747 }
1748 if ix > 0 && value <= prev_value {
1749 return Err(ArrowError::InvalidArgumentError(format!(
1750 "The values in run_ends array should be strictly increasing. Found value {value} at index {ix} with previous value {prev_value} that does not match the criteria."
1751 )));
1752 }
1753
1754 prev_value = value;
1755 Ok(())
1756 })?;
1757
1758 if parent_len > 0 && prev_value.as_usize() < parent_len_plus_offset {
1759 return Err(ArrowError::InvalidArgumentError(format!(
1760 "The offset + length of array should be less or equal to last value in the run_ends array. The last value of run_ends array is {prev_value} and offset + length of array is {parent_len_plus_offset}."
1761 )));
1762 }
1763 Ok(())
1764 }
1765
1766 pub fn ptr_eq(&self, other: &Self) -> bool {
1770 if self.offset != other.offset
1771 || self.len != other.len
1772 || self.data_type != other.data_type
1773 || self.buffers.len() != other.buffers.len()
1774 || self.child_data.len() != other.child_data.len()
1775 {
1776 return false;
1777 }
1778
1779 match (&self.nulls, &other.nulls) {
1780 (Some(a), Some(b)) if !a.inner().ptr_eq(b.inner()) => return false,
1781 (Some(_), None) | (None, Some(_)) => return false,
1782 _ => {}
1783 };
1784
1785 if !self
1786 .buffers
1787 .iter()
1788 .zip(other.buffers.iter())
1789 .all(|(a, b)| a.as_ptr() == b.as_ptr())
1790 {
1791 return false;
1792 }
1793
1794 self.child_data
1795 .iter()
1796 .zip(other.child_data.iter())
1797 .all(|(a, b)| a.ptr_eq(b))
1798 }
1799
1800 pub fn into_builder(self) -> ArrayDataBuilder {
1802 self.into()
1803 }
1804
1805 #[cfg(feature = "pool")]
1812 pub fn claim(&self, pool: &dyn arrow_buffer::MemoryPool) {
1813 for buffer in &self.buffers {
1815 buffer.claim(pool);
1816 }
1817
1818 if let Some(nulls) = &self.nulls {
1820 nulls.claim(pool);
1821 }
1822
1823 for child in &self.child_data {
1825 child.claim(pool);
1826 }
1827 }
1828}
1829
1830pub fn layout(data_type: &DataType) -> DataTypeLayout {
1833 use arrow_schema::IntervalUnit::*;
1836
1837 match data_type {
1838 DataType::Null => DataTypeLayout {
1839 buffers: vec![],
1840 can_contain_null_mask: false,
1841 variadic: false,
1842 },
1843 DataType::Boolean => DataTypeLayout {
1844 buffers: vec![BufferSpec::BitMap],
1845 can_contain_null_mask: true,
1846 variadic: false,
1847 },
1848 DataType::Int8 => DataTypeLayout::new_fixed_width::<i8>(),
1849 DataType::Int16 => DataTypeLayout::new_fixed_width::<i16>(),
1850 DataType::Int32 => DataTypeLayout::new_fixed_width::<i32>(),
1851 DataType::Int64 => DataTypeLayout::new_fixed_width::<i64>(),
1852 DataType::UInt8 => DataTypeLayout::new_fixed_width::<u8>(),
1853 DataType::UInt16 => DataTypeLayout::new_fixed_width::<u16>(),
1854 DataType::UInt32 => DataTypeLayout::new_fixed_width::<u32>(),
1855 DataType::UInt64 => DataTypeLayout::new_fixed_width::<u64>(),
1856 DataType::Float16 => DataTypeLayout::new_fixed_width::<half::f16>(),
1857 DataType::Float32 => DataTypeLayout::new_fixed_width::<f32>(),
1858 DataType::Float64 => DataTypeLayout::new_fixed_width::<f64>(),
1859 DataType::Timestamp(_, _) => DataTypeLayout::new_fixed_width::<i64>(),
1860 DataType::Date32 => DataTypeLayout::new_fixed_width::<i32>(),
1861 DataType::Date64 => DataTypeLayout::new_fixed_width::<i64>(),
1862 DataType::Time32(_) => DataTypeLayout::new_fixed_width::<i32>(),
1863 DataType::Time64(_) => DataTypeLayout::new_fixed_width::<i64>(),
1864 DataType::Interval(YearMonth) => DataTypeLayout::new_fixed_width::<i32>(),
1865 DataType::Interval(DayTime) => DataTypeLayout::new_fixed_width::<IntervalDayTime>(),
1866 DataType::Interval(MonthDayNano) => {
1867 DataTypeLayout::new_fixed_width::<IntervalMonthDayNano>()
1868 }
1869 DataType::Duration(_) => DataTypeLayout::new_fixed_width::<i64>(),
1870 DataType::Decimal32(_, _) => DataTypeLayout::new_fixed_width::<i32>(),
1871 DataType::Decimal64(_, _) => DataTypeLayout::new_fixed_width::<i64>(),
1872 DataType::Decimal128(_, _) => DataTypeLayout::new_fixed_width::<i128>(),
1873 DataType::Decimal256(_, _) => DataTypeLayout::new_fixed_width::<i256>(),
1874 DataType::FixedSizeBinary(size) => {
1875 let spec = BufferSpec::FixedWidth {
1876 byte_width: (*size).try_into().unwrap(),
1877 alignment: mem::align_of::<u8>(),
1878 };
1879 DataTypeLayout {
1880 buffers: vec![spec],
1881 can_contain_null_mask: true,
1882 variadic: false,
1883 }
1884 }
1885 DataType::Binary => DataTypeLayout::new_binary::<i32>(),
1886 DataType::LargeBinary => DataTypeLayout::new_binary::<i64>(),
1887 DataType::Utf8 => DataTypeLayout::new_binary::<i32>(),
1888 DataType::LargeUtf8 => DataTypeLayout::new_binary::<i64>(),
1889 DataType::BinaryView | DataType::Utf8View => DataTypeLayout::new_view(),
1890 DataType::FixedSizeList(_, _) => DataTypeLayout::new_nullable_empty(), DataType::List(_) => DataTypeLayout::new_fixed_width::<i32>(),
1892 DataType::ListView(_) => DataTypeLayout::new_list_view::<i32>(),
1893 DataType::LargeListView(_) => DataTypeLayout::new_list_view::<i64>(),
1894 DataType::LargeList(_) => DataTypeLayout::new_fixed_width::<i64>(),
1895 DataType::Map(_, _) => DataTypeLayout::new_fixed_width::<i32>(),
1896 DataType::Struct(_) => DataTypeLayout::new_nullable_empty(), DataType::RunEndEncoded(_, _) => DataTypeLayout::new_empty(), DataType::Union(_, mode) => {
1899 let type_ids = BufferSpec::FixedWidth {
1900 byte_width: mem::size_of::<i8>(),
1901 alignment: mem::align_of::<i8>(),
1902 };
1903
1904 DataTypeLayout {
1905 buffers: match mode {
1906 UnionMode::Sparse => {
1907 vec![type_ids]
1908 }
1909 UnionMode::Dense => {
1910 vec![
1911 type_ids,
1912 BufferSpec::FixedWidth {
1913 byte_width: mem::size_of::<i32>(),
1914 alignment: mem::align_of::<i32>(),
1915 },
1916 ]
1917 }
1918 },
1919 can_contain_null_mask: false,
1920 variadic: false,
1921 }
1922 }
1923 DataType::Dictionary(key_type, _value_type) => layout(key_type),
1924 }
1925}
1926
1927#[derive(Debug, PartialEq, Eq)]
1929pub struct DataTypeLayout {
1931 pub buffers: Vec<BufferSpec>,
1933
1934 pub can_contain_null_mask: bool,
1936
1937 pub variadic: bool,
1941}
1942
1943impl DataTypeLayout {
1944 pub fn new_fixed_width<T>() -> Self {
1946 Self {
1947 buffers: vec![BufferSpec::FixedWidth {
1948 byte_width: mem::size_of::<T>(),
1949 alignment: mem::align_of::<T>(),
1950 }],
1951 can_contain_null_mask: true,
1952 variadic: false,
1953 }
1954 }
1955
1956 pub fn new_nullable_empty() -> Self {
1959 Self {
1960 buffers: vec![],
1961 can_contain_null_mask: true,
1962 variadic: false,
1963 }
1964 }
1965
1966 pub fn new_empty() -> Self {
1969 Self {
1970 buffers: vec![],
1971 can_contain_null_mask: false,
1972 variadic: false,
1973 }
1974 }
1975
1976 pub fn new_binary<T>() -> Self {
1980 Self {
1981 buffers: vec![
1982 BufferSpec::FixedWidth {
1984 byte_width: mem::size_of::<T>(),
1985 alignment: mem::align_of::<T>(),
1986 },
1987 BufferSpec::VariableWidth,
1989 ],
1990 can_contain_null_mask: true,
1991 variadic: false,
1992 }
1993 }
1994
1995 pub fn new_view() -> Self {
1997 Self {
1998 buffers: vec![BufferSpec::FixedWidth {
1999 byte_width: mem::size_of::<u128>(),
2000 alignment: mem::align_of::<u128>(),
2001 }],
2002 can_contain_null_mask: true,
2003 variadic: true,
2004 }
2005 }
2006
2007 pub fn new_list_view<T>() -> Self {
2009 Self {
2010 buffers: vec![
2011 BufferSpec::FixedWidth {
2012 byte_width: mem::size_of::<T>(),
2013 alignment: mem::align_of::<T>(),
2014 },
2015 BufferSpec::FixedWidth {
2016 byte_width: mem::size_of::<T>(),
2017 alignment: mem::align_of::<T>(),
2018 },
2019 ],
2020 can_contain_null_mask: true,
2021 variadic: false,
2022 }
2023 }
2024}
2025
2026#[derive(Debug, PartialEq, Eq)]
2028pub enum BufferSpec {
2029 FixedWidth {
2040 byte_width: usize,
2042 alignment: usize,
2044 },
2045 VariableWidth,
2047 BitMap,
2053 #[allow(dead_code)]
2056 AlwaysNull,
2057}
2058
2059impl PartialEq for ArrayData {
2060 fn eq(&self, other: &Self) -> bool {
2061 equal::equal(self, other)
2062 }
2063}
2064
2065#[derive(Debug, Clone)]
2084#[doc(hidden)]
2085pub struct UnsafeFlag(bool);
2086
2087impl UnsafeFlag {
2088 #[inline]
2092 pub const fn new() -> Self {
2093 Self(false)
2094 }
2095
2096 #[inline]
2106 pub unsafe fn set(&mut self, val: bool) {
2107 self.0 = val;
2108 }
2109
2110 #[inline]
2112 pub fn get(&self) -> bool {
2113 self.0
2114 }
2115}
2116
2117impl Default for UnsafeFlag {
2119 fn default() -> Self {
2120 Self::new()
2121 }
2122}
2123
2124#[derive(Debug)]
2126pub struct ArrayDataBuilder {
2127 data_type: DataType,
2128 len: usize,
2129 null_count: Option<usize>,
2130 null_bit_buffer: Option<Buffer>,
2131 nulls: Option<NullBuffer>,
2132 offset: usize,
2133 buffers: Vec<Buffer>,
2134 child_data: Vec<ArrayData>,
2135 align_buffers: bool,
2139 skip_validation: UnsafeFlag,
2149}
2150
2151impl ArrayDataBuilder {
2152 #[inline]
2153 pub const fn new(data_type: DataType) -> Self {
2155 Self {
2156 data_type,
2157 len: 0,
2158 null_count: None,
2159 null_bit_buffer: None,
2160 nulls: None,
2161 offset: 0,
2162 buffers: vec![],
2163 child_data: vec![],
2164 align_buffers: false,
2165 skip_validation: UnsafeFlag::new(),
2166 }
2167 }
2168
2169 pub fn data_type(self, data_type: DataType) -> Self {
2171 Self { data_type, ..self }
2172 }
2173
2174 #[inline]
2175 #[allow(clippy::len_without_is_empty)]
2176 pub const fn len(mut self, n: usize) -> Self {
2178 self.len = n;
2179 self
2180 }
2181
2182 pub fn nulls(mut self, nulls: Option<NullBuffer>) -> Self {
2184 self.nulls = nulls;
2185 self.null_count = None;
2186 self.null_bit_buffer = None;
2187 self
2188 }
2189
2190 pub fn null_count(mut self, null_count: usize) -> Self {
2192 self.null_count = Some(null_count);
2193 self
2194 }
2195
2196 pub fn null_bit_buffer(mut self, buf: Option<Buffer>) -> Self {
2198 self.nulls = None;
2199 self.null_bit_buffer = buf;
2200 self
2201 }
2202
2203 #[inline]
2205 pub const fn offset(mut self, n: usize) -> Self {
2206 self.offset = n;
2207 self
2208 }
2209
2210 pub fn buffers(mut self, v: Vec<Buffer>) -> Self {
2212 self.buffers = v;
2213 self
2214 }
2215
2216 pub fn add_buffer(mut self, b: Buffer) -> Self {
2218 self.buffers.push(b);
2219 self
2220 }
2221
2222 pub fn add_buffers<I: IntoIterator<Item = Buffer>>(mut self, bs: I) -> Self {
2224 self.buffers.extend(bs);
2225 self
2226 }
2227
2228 pub fn child_data(mut self, v: Vec<ArrayData>) -> Self {
2230 self.child_data = v;
2231 self
2232 }
2233
2234 pub fn add_child_data(mut self, r: ArrayData) -> Self {
2236 self.child_data.push(r);
2237 self
2238 }
2239
2240 pub unsafe fn build_unchecked(self) -> ArrayData {
2256 unsafe { self.skip_validation(true) }.build().unwrap()
2257 }
2258
2259 pub fn build(self) -> Result<ArrayData, ArrowError> {
2268 let Self {
2269 data_type,
2270 len,
2271 null_count,
2272 null_bit_buffer,
2273 nulls,
2274 offset,
2275 buffers,
2276 child_data,
2277 align_buffers,
2278 skip_validation,
2279 } = self;
2280
2281 let nulls = nulls
2282 .or_else(|| {
2283 let buffer = null_bit_buffer?;
2284 let buffer = BooleanBuffer::new(buffer, offset, len);
2285 Some(match null_count {
2286 Some(n) => {
2287 unsafe { NullBuffer::new_unchecked(buffer, n) }
2289 }
2290 None => NullBuffer::new(buffer),
2291 })
2292 })
2293 .filter(|b| b.null_count() != 0);
2294
2295 let mut data = ArrayData {
2296 data_type,
2297 len,
2298 offset,
2299 buffers,
2300 child_data,
2301 nulls,
2302 };
2303
2304 if align_buffers {
2305 data.align_buffers();
2306 }
2307
2308 if !skip_validation.get() || cfg!(feature = "force_validate") {
2310 data.validate_data()?;
2311 }
2312 Ok(data)
2313 }
2314
2315 pub fn align_buffers(mut self, align_buffers: bool) -> Self {
2331 self.align_buffers = align_buffers;
2332 self
2333 }
2334
2335 pub unsafe fn skip_validation(mut self, skip_validation: bool) -> Self {
2349 unsafe {
2350 self.skip_validation.set(skip_validation);
2351 }
2352 self
2353 }
2354}
2355
2356impl From<ArrayData> for ArrayDataBuilder {
2357 fn from(d: ArrayData) -> Self {
2358 Self {
2359 data_type: d.data_type,
2360 len: d.len,
2361 offset: d.offset,
2362 buffers: d.buffers,
2363 child_data: d.child_data,
2364 nulls: d.nulls,
2365 null_bit_buffer: None,
2366 null_count: None,
2367 align_buffers: false,
2368 skip_validation: UnsafeFlag::new(),
2369 }
2370 }
2371}
2372
2373pub(crate) fn get_fixed_size_binary_width(data_type: &DataType) -> usize {
2378 match data_type {
2379 DataType::FixedSizeBinary(i) => {
2380 if *i < 0 {
2381 panic!("cannot compare FixedSizeBinary({})", *i);
2382 }
2383 *i as usize
2384 }
2385 _ => unreachable!(),
2386 }
2387}
2388
2389#[cfg(test)]
2390mod tests {
2391 use crate::transform::MutableArrayData;
2392
2393 use super::*;
2394 use arrow_buffer::{OffsetBuffer, ScalarBuffer};
2395 use arrow_schema::{Field, Fields, UnionFields};
2396
2397 fn make_i32_buffer(n: usize) -> Buffer {
2401 Buffer::from_slice_ref(vec![42i32; n])
2402 }
2403
2404 fn make_f32_buffer(n: usize) -> Buffer {
2406 Buffer::from_slice_ref(vec![42f32; n])
2407 }
2408
2409 #[test]
2410 fn test_builder() {
2411 let v = (0..25).collect::<Vec<i32>>();
2413 let b1 = Buffer::from_slice_ref(&v);
2414 let arr_data = ArrayData::builder(DataType::Int32)
2415 .len(20)
2416 .offset(5)
2417 .add_buffer(b1)
2418 .null_bit_buffer(Some(Buffer::from([
2419 0b01011111, 0b10110101, 0b01100011, 0b00011110,
2420 ])))
2421 .build()
2422 .unwrap();
2423
2424 assert_eq!(20, arr_data.len());
2425 assert_eq!(10, arr_data.null_count());
2426 assert_eq!(5, arr_data.offset());
2427 assert_eq!(1, arr_data.buffers().len());
2428 assert_eq!(
2429 Buffer::from_slice_ref(&v).as_slice(),
2430 arr_data.buffers()[0].as_slice()
2431 );
2432 }
2433
2434 #[test]
2435 fn test_builder_with_child_data() {
2436 let child_arr_data = ArrayData::try_new(
2437 DataType::Int32,
2438 5,
2439 None,
2440 0,
2441 vec![Buffer::from_slice_ref([1i32, 2, 3, 4, 5])],
2442 vec![],
2443 )
2444 .unwrap();
2445
2446 let field = Arc::new(Field::new("x", DataType::Int32, true));
2447 let data_type = DataType::Struct(vec![field].into());
2448
2449 let arr_data = ArrayData::builder(data_type)
2450 .len(5)
2451 .offset(0)
2452 .add_child_data(child_arr_data.clone())
2453 .build()
2454 .unwrap();
2455
2456 assert_eq!(5, arr_data.len());
2457 assert_eq!(1, arr_data.child_data().len());
2458 assert_eq!(child_arr_data, arr_data.child_data()[0]);
2459 }
2460
2461 #[test]
2462 fn test_extend_fixed_size_list_accounts_for_parent_offset() {
2463 let field = Arc::new(Field::new("item", DataType::Int32, false));
2464 let child = ArrayData::builder(DataType::Int32)
2465 .len(6)
2466 .add_buffer(Buffer::from_slice_ref([0, 1, 2, 3, 4, 5]))
2467 .build()
2468 .unwrap();
2469
2470 let data = ArrayData::builder(DataType::FixedSizeList(field, 2))
2471 .len(2)
2472 .offset(1)
2473 .add_child_data(child)
2474 .build()
2475 .unwrap();
2476
2477 let mut mutable = MutableArrayData::new(vec![&data], false, data.len());
2478 mutable.try_extend(0, 0, data.len()).unwrap();
2479 let output = mutable.freeze();
2480
2481 assert_eq!(output.child_data()[0].buffer::<i32>(0), &[2, 3, 4, 5]);
2482 }
2483
2484 #[test]
2485 fn test_null_count() {
2486 let mut bit_v: [u8; 2] = [0; 2];
2487 bit_util::set_bit(&mut bit_v, 0);
2488 bit_util::set_bit(&mut bit_v, 3);
2489 bit_util::set_bit(&mut bit_v, 10);
2490 let arr_data = ArrayData::builder(DataType::Int32)
2491 .len(16)
2492 .add_buffer(make_i32_buffer(16))
2493 .null_bit_buffer(Some(Buffer::from(bit_v)))
2494 .build()
2495 .unwrap();
2496 assert_eq!(13, arr_data.null_count());
2497
2498 let mut bit_v: [u8; 2] = [0; 2];
2500 bit_util::set_bit(&mut bit_v, 0);
2501 bit_util::set_bit(&mut bit_v, 3);
2502 bit_util::set_bit(&mut bit_v, 10);
2503 let arr_data = ArrayData::builder(DataType::Int32)
2504 .len(12)
2505 .offset(2)
2506 .add_buffer(make_i32_buffer(14)) .null_bit_buffer(Some(Buffer::from(bit_v)))
2508 .build()
2509 .unwrap();
2510 assert_eq!(10, arr_data.null_count());
2511 }
2512
2513 #[test]
2514 fn test_null_buffer_ref() {
2515 let mut bit_v: [u8; 2] = [0; 2];
2516 bit_util::set_bit(&mut bit_v, 0);
2517 bit_util::set_bit(&mut bit_v, 3);
2518 bit_util::set_bit(&mut bit_v, 10);
2519 let arr_data = ArrayData::builder(DataType::Int32)
2520 .len(16)
2521 .add_buffer(make_i32_buffer(16))
2522 .null_bit_buffer(Some(Buffer::from(bit_v)))
2523 .build()
2524 .unwrap();
2525 assert!(arr_data.nulls().is_some());
2526 assert_eq!(&bit_v, arr_data.nulls().unwrap().validity());
2527 }
2528
2529 #[test]
2530 fn test_slice() {
2531 let mut bit_v: [u8; 2] = [0; 2];
2532 bit_util::set_bit(&mut bit_v, 0);
2533 bit_util::set_bit(&mut bit_v, 3);
2534 bit_util::set_bit(&mut bit_v, 10);
2535 let data = ArrayData::builder(DataType::Int32)
2536 .len(16)
2537 .add_buffer(make_i32_buffer(16))
2538 .null_bit_buffer(Some(Buffer::from(bit_v)))
2539 .build()
2540 .unwrap();
2541 let new_data = data.slice(1, 15);
2542 assert_eq!(data.len() - 1, new_data.len());
2543 assert_eq!(1, new_data.offset());
2544 assert_eq!(data.null_count(), new_data.null_count());
2545
2546 let new_data = new_data.slice(1, 14);
2548 assert_eq!(data.len() - 2, new_data.len());
2549 assert_eq!(2, new_data.offset());
2550 assert_eq!(data.null_count() - 1, new_data.null_count());
2551 }
2552
2553 #[test]
2554 #[should_panic(expected = "offset + length overflow")]
2555 fn test_slice_panics_on_offset_length_overflow() {
2556 let data = ArrayData::builder(DataType::Int32)
2557 .len(4)
2558 .add_buffer(make_i32_buffer(4))
2559 .build()
2560 .unwrap();
2561 let sliced = data.slice(1, 3);
2562
2563 sliced.slice(1, usize::MAX);
2564 }
2565
2566 #[test]
2567 fn test_typed_offsets_length_overflow() {
2568 let data = ArrayData {
2569 data_type: DataType::Binary,
2570 len: usize::MAX,
2571 offset: 0,
2572 buffers: vec![Buffer::from_slice_ref([0_i32])],
2573 child_data: vec![],
2574 nulls: None,
2575 };
2576 let err = data.typed_offsets::<i32>().unwrap_err();
2577
2578 assert_eq!(
2579 err.to_string(),
2580 format!(
2581 "Invalid argument error: Length {} with offset 1 overflows usize for Binary",
2582 usize::MAX
2583 )
2584 );
2585 }
2586
2587 #[test]
2588 fn test_validate_typed_buffer_length_overflow() {
2589 let data = ArrayData {
2590 data_type: DataType::Binary,
2591 len: 0,
2592 offset: 2,
2593 buffers: vec![Buffer::from_slice_ref([0_i32])],
2594 child_data: vec![],
2595 nulls: None,
2596 };
2597 let err = data.typed_buffer::<i32>(0, usize::MAX).unwrap_err();
2598
2599 assert_eq!(
2600 err.to_string(),
2601 format!(
2602 "Invalid argument error: Length {} with offset 2 overflows usize for Binary",
2603 usize::MAX
2604 )
2605 );
2606 }
2607
2608 fn try_new_binary_length_offset_overflow() -> Result<ArrayData, ArrowError> {
2610 ArrayData::try_new(
2611 DataType::Binary,
2612 usize::MAX,
2613 None,
2614 1,
2615 vec![
2616 Buffer::from_slice_ref([0_i32]),
2617 Buffer::from_iter(std::iter::empty::<u8>()),
2618 ],
2619 vec![],
2620 )
2621 }
2622
2623 #[cfg(not(feature = "force_validate"))]
2624 #[test]
2625 fn test_try_new_length_offset_overflow() {
2626 let err = try_new_binary_length_offset_overflow().unwrap_err();
2627
2628 assert_eq!(
2629 err.to_string(),
2630 format!(
2631 "Invalid argument error: Length {} with offset 1 overflows usize for Binary",
2632 usize::MAX
2633 )
2634 );
2635 }
2636
2637 #[cfg(feature = "force_validate")]
2638 #[test]
2639 #[should_panic(
2640 expected = "Length 18446744073709551615 with offset 1 overflows usize for Binary"
2641 )]
2642 fn test_try_new_length_offset_overflow_force_validate() {
2643 try_new_binary_length_offset_overflow().unwrap();
2644 }
2645
2646 #[test]
2647 fn test_equality() {
2648 let int_data = ArrayData::builder(DataType::Int32)
2649 .len(1)
2650 .add_buffer(make_i32_buffer(1))
2651 .build()
2652 .unwrap();
2653
2654 let float_data = ArrayData::builder(DataType::Float32)
2655 .len(1)
2656 .add_buffer(make_f32_buffer(1))
2657 .build()
2658 .unwrap();
2659 assert_ne!(int_data, float_data);
2660 assert!(!int_data.ptr_eq(&float_data));
2661 assert!(int_data.ptr_eq(&int_data));
2662
2663 #[allow(clippy::redundant_clone)]
2664 let int_data_clone = int_data.clone();
2665 assert_eq!(int_data, int_data_clone);
2666 assert!(int_data.ptr_eq(&int_data_clone));
2667 assert!(int_data_clone.ptr_eq(&int_data));
2668
2669 let int_data_slice = int_data_clone.slice(1, 0);
2670 assert!(int_data_slice.ptr_eq(&int_data_slice));
2671 assert!(!int_data.ptr_eq(&int_data_slice));
2672 assert!(!int_data_slice.ptr_eq(&int_data));
2673
2674 let data_buffer = Buffer::from_slice_ref("abcdef".as_bytes());
2675 let offsets_buffer = Buffer::from_slice_ref([0_i32, 2_i32, 2_i32, 5_i32]);
2676 let string_data = ArrayData::try_new(
2677 DataType::Utf8,
2678 3,
2679 Some(Buffer::from_iter(vec![true, false, true])),
2680 0,
2681 vec![offsets_buffer, data_buffer],
2682 vec![],
2683 )
2684 .unwrap();
2685
2686 assert_ne!(float_data, string_data);
2687 assert!(!float_data.ptr_eq(&string_data));
2688
2689 assert!(string_data.ptr_eq(&string_data));
2690
2691 #[allow(clippy::redundant_clone)]
2692 let string_data_cloned = string_data.clone();
2693 assert!(string_data_cloned.ptr_eq(&string_data));
2694 assert!(string_data.ptr_eq(&string_data_cloned));
2695
2696 let string_data_slice = string_data.slice(1, 2);
2697 assert!(string_data_slice.ptr_eq(&string_data_slice));
2698 assert!(!string_data_slice.ptr_eq(&string_data))
2699 }
2700
2701 #[test]
2702 fn test_slice_memory_size_utf8_offset_buffer_len_plus_one() {
2703 let data_buffer = Buffer::from_slice_ref("helloworld".as_bytes());
2705 let offsets_buffer = Buffer::from_slice_ref([0_i32, 5_i32, 10_i32]);
2708 let array = ArrayData::try_new(
2709 DataType::Utf8,
2710 2,
2711 None,
2712 0,
2713 vec![offsets_buffer, data_buffer],
2714 vec![],
2715 )
2716 .unwrap();
2717 assert_eq!(array.get_slice_memory_size().unwrap(), 22); }
2719
2720 #[test]
2721 fn test_slice_memory_size_binary_offset_buffer_len_plus_one() {
2722 let data_buffer = Buffer::from_slice_ref([0u8, 1, 2, 3, 4]);
2725 let offsets_buffer = Buffer::from_slice_ref([0_i32, 2_i32, 5_i32]);
2727 let array = ArrayData::try_new(
2728 DataType::Binary,
2729 2,
2730 None,
2731 0,
2732 vec![offsets_buffer, data_buffer],
2733 vec![],
2734 )
2735 .unwrap();
2736 assert_eq!(array.get_slice_memory_size().unwrap(), 17); }
2738
2739 #[test]
2740 fn test_slice_memory_size() {
2741 let mut bit_v: [u8; 2] = [0; 2];
2742 bit_util::set_bit(&mut bit_v, 0);
2743 bit_util::set_bit(&mut bit_v, 3);
2744 bit_util::set_bit(&mut bit_v, 10);
2745 let data = ArrayData::builder(DataType::Int32)
2746 .len(16)
2747 .add_buffer(make_i32_buffer(16))
2748 .null_bit_buffer(Some(Buffer::from(bit_v)))
2749 .build()
2750 .unwrap();
2751 let new_data = data.slice(1, 14);
2752 assert_eq!(
2753 data.get_slice_memory_size().unwrap() - 8,
2754 new_data.get_slice_memory_size().unwrap()
2755 );
2756 let data_buffer = Buffer::from_slice_ref("abcdef".as_bytes());
2757 let offsets_buffer = Buffer::from_slice_ref([0_i32, 2_i32, 2_i32, 5_i32]);
2758 let string_data = ArrayData::try_new(
2759 DataType::Utf8,
2760 3,
2761 Some(Buffer::from_iter(vec![true, false, true])),
2762 0,
2763 vec![offsets_buffer, data_buffer],
2764 vec![],
2765 )
2766 .unwrap();
2767 let string_data_slice = string_data.slice(1, 2);
2768 assert_eq!(
2770 string_data.get_slice_memory_size().unwrap() - 6,
2771 string_data_slice.get_slice_memory_size().unwrap()
2772 );
2773 }
2774
2775 #[test]
2776 fn test_count_nulls() {
2777 let buffer = Buffer::from([0b00010110, 0b10011111]);
2778 let buffer = NullBuffer::new(BooleanBuffer::new(buffer, 0, 16));
2779 let count = count_nulls(Some(&buffer), 0, 16);
2780 assert_eq!(count, 7);
2781
2782 let count = count_nulls(Some(&buffer), 4, 8);
2783 assert_eq!(count, 3);
2784 }
2785
2786 #[test]
2787 fn test_contains_nulls() {
2788 let buffer: Buffer =
2789 MutableBuffer::from_iter([false, false, false, true, true, false]).into();
2790 let buffer = NullBuffer::new(BooleanBuffer::new(buffer, 0, 6));
2791 assert!(contains_nulls(Some(&buffer), 0, 6));
2792 assert!(contains_nulls(Some(&buffer), 0, 3));
2793 assert!(!contains_nulls(Some(&buffer), 3, 2));
2794 assert!(!contains_nulls(Some(&buffer), 0, 0));
2795 }
2796
2797 #[test]
2798 fn test_alignment() {
2799 let buffer = Buffer::from_vec(vec![1_i32, 2_i32, 3_i32]);
2800 let sliced = buffer.slice(1);
2801
2802 let mut data = ArrayData {
2803 data_type: DataType::Int32,
2804 len: 0,
2805 offset: 0,
2806 buffers: vec![buffer],
2807 child_data: vec![],
2808 nulls: None,
2809 };
2810 data.validate_full().unwrap();
2811
2812 data.buffers[0] = sliced;
2814 let err = data.validate().unwrap_err();
2815
2816 assert_eq!(
2817 err.to_string(),
2818 "Invalid argument error: Misaligned buffers[0] in array of type Int32, offset from expected alignment of 4 by 1"
2819 );
2820
2821 data.align_buffers();
2822 data.validate_full().unwrap();
2823 }
2824
2825 #[test]
2826 fn test_alignment_struct() {
2827 let buffer = Buffer::from_vec(vec![1_i32, 2_i32, 3_i32]);
2828 let sliced = buffer.slice(1);
2829
2830 let child_data = ArrayData {
2831 data_type: DataType::Int32,
2832 len: 0,
2833 offset: 0,
2834 buffers: vec![buffer],
2835 child_data: vec![],
2836 nulls: None,
2837 };
2838
2839 let schema = DataType::Struct(Fields::from(vec![Field::new("a", DataType::Int32, false)]));
2840 let mut data = ArrayData {
2841 data_type: schema,
2842 len: 0,
2843 offset: 0,
2844 buffers: vec![],
2845 child_data: vec![child_data],
2846 nulls: None,
2847 };
2848 data.validate_full().unwrap();
2849
2850 data.child_data[0].buffers[0] = sliced;
2852 let err = data.validate().unwrap_err();
2853
2854 assert_eq!(
2855 err.to_string(),
2856 "Invalid argument error: Misaligned buffers[0] in array of type Int32, offset from expected alignment of 4 by 1"
2857 );
2858
2859 data.align_buffers();
2860 data.validate_full().unwrap();
2861 }
2862
2863 #[test]
2864 fn test_null_view_types() {
2865 let array_len = 32;
2866 let array = ArrayData::new_null(&DataType::BinaryView, array_len);
2867 assert_eq!(array.len(), array_len);
2868 for i in 0..array.len() {
2869 assert!(array.is_null(i));
2870 }
2871
2872 let array = ArrayData::new_null(&DataType::Utf8View, array_len);
2873 assert_eq!(array.len(), array_len);
2874 for i in 0..array.len() {
2875 assert!(array.is_null(i));
2876 }
2877
2878 let array = ArrayData::new_null(
2879 &DataType::ListView(Arc::new(Field::new_list_field(DataType::Int32, true))),
2880 array_len,
2881 );
2882 assert_eq!(array.len(), array_len);
2883 for i in 0..array.len() {
2884 assert!(array.is_null(i));
2885 }
2886
2887 let array = ArrayData::new_null(
2888 &DataType::LargeListView(Arc::new(Field::new_list_field(DataType::Int32, true))),
2889 array_len,
2890 );
2891 assert_eq!(array.len(), array_len);
2892 for i in 0..array.len() {
2893 assert!(array.is_null(i));
2894 }
2895 }
2896
2897 #[test]
2899 fn test_dont_panic_on_bad_input_when_using_try_new() {
2900 let empty_bytes = Buffer::default();
2901
2902 let array_data = ArrayData::try_new(
2903 DataType::Utf8,
2904 1, None,
2906 0,
2907 vec![Buffer::from_vec(vec![0i32, 2i32]), empty_bytes],
2909 vec![],
2910 );
2911
2912 let res = array_data.expect_err("should get error");
2913
2914 assert_eq!(
2915 res.to_string(),
2916 format!("Invalid argument error: Last offset 2 of Utf8 is larger than values length 0",)
2917 );
2918 }
2919
2920 #[test]
2921 fn should_fail_validation_when_having_map_field_type_is_not_struct() {
2922 let map_field = Field::new("key", DataType::Int32, false);
2923
2924 let map_field_data = valid_non_nullable_int32_array_data(2);
2925
2926 let results = test_both_builder_and_array_data(
2927 DataType::Map(map_field.into(), false),
2928 1,
2929 None,
2930 0,
2931 vec![
2932 OffsetBuffer::<i32>::from_lengths(vec![2])
2933 .into_inner()
2934 .into(),
2935 ],
2936 vec![map_field_data],
2937 );
2938
2939 for result in results {
2940 let array_data_err = result.expect_err("should fail for non struct field");
2941
2942 match array_data_err {
2943 ArrowError::InvalidArgumentError(msg) => {
2944 assert_eq!(
2945 msg,
2946 "Map field should be a entries struct data type, got Int32 instead"
2947 )
2948 }
2949 _ => panic!("unexpected error type {array_data_err}"),
2950 };
2951 }
2952 }
2953
2954 #[test]
2955 fn should_fail_validation_when_having_map_entries_only_have_1_field() {
2956 let struct_data_type = DataType::Struct(Fields::from(vec![Field::new(
2957 "key",
2958 DataType::Int32,
2959 false,
2960 )]));
2961
2962 let key_array_data = valid_non_nullable_int32_array_data(2);
2963
2964 let struct_data = {
2965 let builder = ArrayDataBuilder::new(struct_data_type.clone())
2966 .len(2)
2967 .nulls(None)
2968 .child_data(vec![key_array_data]);
2969
2970 builder.build().unwrap()
2971 };
2972
2973 let results = test_both_builder_and_array_data(
2974 DataType::Map(Field::new("entries", struct_data_type, false).into(), false),
2975 1,
2976 None,
2977 0,
2978 vec![
2979 OffsetBuffer::<i32>::from_lengths(vec![2])
2980 .into_inner()
2981 .into(),
2982 ],
2983 vec![struct_data],
2984 );
2985
2986 for result in results {
2987 let array_data_err = result.expect_err("should fail for nullable key");
2988
2989 match array_data_err {
2990 ArrowError::InvalidArgumentError(msg) => {
2991 assert_eq!(
2992 msg,
2993 "Map entries data type should be a struct containing 2 fields, got 1 fields"
2994 )
2995 }
2996 _ => panic!("unexpected error type {array_data_err}"),
2997 };
2998 }
2999 }
3000
3001 #[test]
3002 fn should_fail_validation_when_having_map_entries_have_3_fields() {
3003 let struct_data_type = DataType::Struct(Fields::from(vec![
3004 Field::new("key", DataType::Int32, false),
3005 Field::new("values", DataType::Utf8, true),
3006 Field::new("other", DataType::Int32, true),
3007 ]));
3008
3009 let key_array_data = valid_non_nullable_int32_array_data(2);
3010
3011 let values_array_data = valid_string_array_data(2);
3012
3013 let other_array_data = key_array_data.clone();
3014
3015 let struct_data = {
3016 let builder = ArrayDataBuilder::new(struct_data_type.clone())
3017 .len(2)
3018 .nulls(None)
3019 .child_data(vec![key_array_data, values_array_data, other_array_data]);
3020
3021 builder.build().unwrap()
3022 };
3023
3024 let results = test_both_builder_and_array_data(
3025 DataType::Map(Field::new("entries", struct_data_type, false).into(), false),
3026 1,
3027 None,
3028 0,
3029 vec![
3030 OffsetBuffer::<i32>::from_lengths(vec![2])
3031 .into_inner()
3032 .into(),
3033 ],
3034 vec![struct_data],
3035 );
3036
3037 for result in results {
3038 let array_data_err = result.expect_err("should fail for nullable key");
3039
3040 match array_data_err {
3041 ArrowError::InvalidArgumentError(msg) => {
3042 assert_eq!(
3043 msg,
3044 "Map entries data type should be a struct containing 2 fields, got 3 fields"
3045 )
3046 }
3047 _ => panic!("unexpected error type {array_data_err}"),
3048 };
3049 }
3050 }
3051
3052 #[test]
3053 fn should_fail_validation_when_having_nullable_map_keys() {
3054 let struct_data_type = DataType::Struct(Fields::from(vec![
3055 Field::new("key", DataType::Int32, true),
3056 Field::new("values", DataType::Utf8, true),
3057 ]));
3058
3059 let key_array_data = valid_non_nullable_int32_array_data(2);
3060 let values_array_data = valid_string_array_data(2);
3061
3062 let struct_data = {
3063 let builder = ArrayDataBuilder::new(struct_data_type.clone())
3064 .len(2)
3065 .nulls(None)
3066 .child_data(vec![key_array_data, values_array_data]);
3067
3068 builder.build().unwrap()
3069 };
3070
3071 let results = test_both_builder_and_array_data(
3072 DataType::Map(Field::new("entries", struct_data_type, false).into(), false),
3073 1,
3074 None,
3075 0,
3076 vec![
3077 OffsetBuffer::<i32>::from_lengths(vec![2])
3078 .into_inner()
3079 .into(),
3080 ],
3081 vec![struct_data],
3082 );
3083
3084 for result in results {
3085 let array_data_err = result.expect_err("should fail for nullable key");
3086
3087 match array_data_err {
3088 ArrowError::InvalidArgumentError(msg) => {
3089 assert_eq!(msg, "Map key field must not be nullable")
3090 }
3091 _ => panic!("unexpected error type {array_data_err}"),
3092 };
3093 }
3094 }
3095
3096 #[test]
3097 fn should_fail_validation_when_having_entries_is_nullable_for_map() {
3098 let struct_data_type = DataType::Struct(Fields::from(vec![
3099 Field::new("key", DataType::Int32, false),
3100 Field::new("values", DataType::Utf8, true),
3101 ]));
3102
3103 let key_array_data = valid_non_nullable_int32_array_data(2);
3104
3105 let values_array_data = valid_string_array_data(2);
3106
3107 let struct_data = {
3108 let builder = ArrayDataBuilder::new(struct_data_type.clone())
3109 .len(2)
3110 .nulls(None)
3111 .child_data(vec![key_array_data, values_array_data]);
3112
3113 builder.build().unwrap()
3114 };
3115
3116 let results = test_both_builder_and_array_data(
3117 DataType::Map(Field::new("entries", struct_data_type, true).into(), false),
3118 1,
3119 None,
3120 0,
3121 vec![
3122 OffsetBuffer::<i32>::from_lengths(vec![2])
3123 .into_inner()
3124 .into(),
3125 ],
3126 vec![struct_data],
3127 );
3128
3129 for result in results {
3130 let array_data_err = result.expect_err("should fail for nullable entries");
3131
3132 match array_data_err {
3133 ArrowError::InvalidArgumentError(msg) => assert_eq!(
3134 msg,
3135 "The nullable should be set to false for the map entries field."
3136 ),
3137 _ => panic!("unexpected error type {array_data_err}"),
3138 };
3139 }
3140 }
3141
3142 #[test]
3143 fn should_allow_to_create_map_from_data() {
3144 let struct_data_type = DataType::Struct(Fields::from(vec![
3145 Field::new("key", DataType::Int32, false),
3146 Field::new("values", DataType::Utf8, true),
3147 ]));
3148
3149 let key_array_data = valid_non_nullable_int32_array_data(2);
3150 let values_array_data = valid_string_array_data(2);
3151
3152 let struct_data = {
3153 let builder = ArrayDataBuilder::new(struct_data_type.clone())
3154 .len(2)
3155 .nulls(None)
3156 .child_data(vec![key_array_data, values_array_data]);
3157
3158 builder.build().unwrap()
3159 };
3160
3161 let results = test_both_builder_and_array_data(
3162 DataType::Map(Field::new("entries", struct_data_type, false).into(), false),
3163 1,
3164 None,
3165 0,
3166 vec![
3167 OffsetBuffer::<i32>::from_lengths(vec![2])
3168 .into_inner()
3169 .into(),
3170 ],
3171 vec![struct_data],
3172 );
3173
3174 for result in results {
3175 result.expect("should be able to create map ArrayData");
3176 }
3177 }
3178
3179 fn valid_string_array_data(length: usize) -> ArrayData {
3180 let offsets = OffsetBuffer::<i32>::from_lengths(vec![0; length])
3181 .into_inner()
3182 .into_inner();
3183 let empty_bytes = Buffer::default();
3184
3185 let builder = ArrayDataBuilder::new(DataType::Utf8)
3186 .len(length)
3187 .buffers(vec![offsets, empty_bytes])
3188 .nulls(None);
3189
3190 builder.build().unwrap()
3191 }
3192
3193 fn valid_non_nullable_int32_array_data(length: usize) -> ArrayData {
3194 let builder = ArrayDataBuilder::new(DataType::Int32)
3195 .len(length)
3196 .nulls(None)
3197 .buffers(vec![
3198 ScalarBuffer::<i32>::from(vec![1; length]).into_inner(),
3199 ]);
3200
3201 builder.build().unwrap()
3202 }
3203
3204 #[test]
3205 fn empty_and_null_map_array_should_pass_validation() {
3206 let dt = DataType::Map(
3207 Field::new(
3208 "entries",
3209 DataType::Struct(Fields::from(vec![
3210 Field::new("key", DataType::Int32, false),
3211 Field::new("values", DataType::Utf8, true),
3212 ])),
3213 false,
3214 )
3215 .into(),
3216 false,
3217 );
3218
3219 ArrayData::new_empty(&dt).validate_full().unwrap();
3220 ArrayData::new_null(&dt, 1).validate_full().unwrap();
3221 }
3222
3223 fn test_both_builder_and_array_data(
3224 data_type: DataType,
3225 len: usize,
3226 null_bit_buffer: Option<Buffer>,
3227 offset: usize,
3228 buffers: Vec<Buffer>,
3229 child_data: Vec<ArrayData>,
3230 ) -> [Result<ArrayData, ArrowError>; 2] {
3231 let from_builder_res = ArrayData::builder(data_type.clone())
3232 .len(len)
3233 .add_buffers(buffers.clone())
3234 .null_bit_buffer(null_bit_buffer.clone())
3235 .offset(offset)
3236 .child_data(child_data.clone())
3237 .build();
3238
3239 let from_try_new_res =
3240 ArrayData::try_new(data_type, len, null_bit_buffer, offset, buffers, child_data);
3241
3242 [from_builder_res, from_try_new_res]
3243 }
3244
3245 #[test]
3246 fn test_union_validate_type_ids() {
3247 let dt = DataType::Union(
3248 UnionFields::try_new(
3249 vec![5],
3250 vec![Arc::new(Field::new("a", DataType::Int32, true))],
3251 )
3252 .unwrap(),
3253 UnionMode::Sparse,
3254 );
3255
3256 let child = ArrayData::builder(DataType::Int32)
3257 .len(2)
3258 .add_buffer(make_i32_buffer(2))
3259 .build()
3260 .unwrap();
3261
3262 let err = ArrayDataBuilder::new(dt)
3263 .len(2)
3264 .offset(0)
3265 .buffers(vec![vec![5_i8, 2].into()])
3266 .child_data(vec![child])
3267 .build()
3268 .unwrap_err();
3269 assert_eq!(
3270 err.to_string(),
3271 "Invalid argument error: Type id values must match one of the field type ids, got 2 at index 1"
3272 );
3273 }
3274
3275 #[test]
3276 fn test_union_validate_dense_offsets() {
3277 let dt = DataType::Union(
3278 UnionFields::try_new(
3279 vec![5],
3280 vec![Arc::new(Field::new("a", DataType::Int32, true))],
3281 )
3282 .unwrap(),
3283 UnionMode::Dense,
3284 );
3285
3286 let child = ArrayData::builder(DataType::Int32)
3287 .len(2)
3288 .add_buffer(make_i32_buffer(2))
3289 .build()
3290 .unwrap();
3291
3292 let err = ArrayDataBuilder::new(dt.clone())
3293 .len(2)
3294 .buffers(vec![vec![5_i8, 5].into(), vec![0_i32, 10].into()])
3296 .child_data(vec![child.clone()])
3297 .build()
3298 .unwrap_err();
3299 assert_eq!(
3300 err.to_string(),
3301 "Invalid argument error: Offsets must be non-negative and within the length of the child array"
3302 );
3303
3304 let err = ArrayDataBuilder::new(dt)
3305 .len(2)
3306 .buffers(vec![vec![5_i8, 5].into(), vec![0_i32, -1].into()])
3308 .child_data(vec![child])
3309 .build()
3310 .unwrap_err();
3311 assert_eq!(
3312 err.to_string(),
3313 "Invalid argument error: Offsets must be non-negative and within the length of the child array"
3314 );
3315 }
3316
3317 fn run_array_test_children() -> (DataType, Vec<ArrayData>) {
3320 let dt = DataType::RunEndEncoded(
3321 Arc::new(Field::new("run_ends", DataType::Int32, false)),
3322 Arc::new(Field::new("values", DataType::Float32, true)),
3323 );
3324
3325 let runs = ArrayData::builder(DataType::Int32)
3326 .len(2)
3327 .add_buffer(vec![1_i32, 3].into())
3328 .build()
3329 .unwrap();
3330 let values = ArrayData::builder(DataType::Float32)
3331 .len(2)
3332 .add_buffer(vec![1.5_f32, 2.5].into())
3333 .build()
3334 .unwrap();
3335
3336 (dt, vec![runs, values])
3337 }
3338
3339 #[test]
3340 fn test_run_array_validates_last_run() {
3341 let (dt, children) = run_array_test_children();
3342
3343 let err = ArrayDataBuilder::new(dt)
3344 .len(4)
3345 .offset(0)
3346 .child_data(children)
3347 .build()
3348 .unwrap_err();
3349 assert_eq!(
3350 err.to_string(),
3351 "Invalid argument error: The offset + length of array should be less or equal to last value in the run_ends array. The last value of run_ends array is 3 and offset + length of array is 4."
3352 );
3353 }
3354
3355 #[test]
3356 fn test_run_array_empty_validates() {
3357 let (dt, children) = run_array_test_children();
3358
3359 let data = ArrayDataBuilder::new(dt)
3360 .len(0)
3361 .offset(4)
3362 .child_data(children)
3363 .build()
3364 .unwrap();
3365 assert_eq!(data.len(), 0);
3366 }
3367
3368 #[test]
3371 fn test_run_array_validates_last_run_with_offset() {
3372 let (dt, children) = run_array_test_children();
3373
3374 let err = ArrayDataBuilder::new(dt)
3375 .len(2)
3376 .offset(2)
3377 .child_data(children)
3378 .build()
3379 .unwrap_err();
3380 assert_eq!(
3381 err.to_string(),
3382 "Invalid argument error: The offset + length of array should be less or equal to last value in the run_ends array. The last value of run_ends array is 3 and offset + length of array is 4."
3383 );
3384 }
3385
3386 #[test]
3389 fn test_run_array_sliced_validates() {
3390 let (dt, children) = run_array_test_children();
3391
3392 let data = ArrayDataBuilder::new(dt)
3393 .len(2)
3394 .offset(1)
3395 .child_data(children)
3396 .build()
3397 .unwrap();
3398 assert_eq!(data.len(), 2);
3399 assert_eq!(data.offset(), 1);
3400 }
3401}