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,
222
223 buffers: Vec<Buffer>,
236
237 child_data: Vec<ArrayData>,
249
250 nulls: Option<NullBuffer>,
258}
259
260pub type ArrayDataRef = Arc<ArrayData>;
262
263fn checked_len_plus_offset(
264 data_type: &DataType,
265 len: usize,
266 offset: usize,
267) -> Result<usize, ArrowError> {
268 len.checked_add(offset).ok_or_else(|| {
269 ArrowError::InvalidArgumentError(format!(
270 "Length {len} with offset {offset} overflows usize for {data_type}"
271 ))
272 })
273}
274
275impl ArrayData {
276 pub unsafe fn new_unchecked(
293 data_type: DataType,
294 len: usize,
295 null_count: Option<usize>,
296 null_bit_buffer: Option<Buffer>,
297 offset: usize,
298 buffers: Vec<Buffer>,
299 child_data: Vec<ArrayData>,
300 ) -> Self {
301 let builder = Self::inner_new_builder(
302 data_type,
303 len,
304 null_count,
305 null_bit_buffer,
306 offset,
307 buffers,
308 child_data,
309 );
310
311 unsafe { builder.build_unchecked() }
313 }
314
315 pub fn try_new(
329 data_type: DataType,
330 len: usize,
331 null_bit_buffer: Option<Buffer>,
332 offset: usize,
333 buffers: Vec<Buffer>,
334 child_data: Vec<ArrayData>,
335 ) -> Result<Self, ArrowError> {
336 let builder = Self::inner_new_builder(
337 data_type,
338 len,
339 None,
340 null_bit_buffer,
341 offset,
342 buffers,
343 child_data,
344 );
345
346 assert!(!builder.skip_validation.get());
347
348 builder.build()
353 }
354
355 fn inner_new_builder(
356 data_type: DataType,
357 len: usize,
358 null_count: Option<usize>,
359 null_bit_buffer: Option<Buffer>,
360 offset: usize,
361 buffers: Vec<Buffer>,
362 child_data: Vec<ArrayData>,
363 ) -> ArrayDataBuilder {
364 ArrayDataBuilder {
365 data_type,
366 len,
367 null_count,
368 null_bit_buffer,
369 nulls: None,
370 offset,
371 buffers,
372 child_data,
373 align_buffers: false,
374 skip_validation: UnsafeFlag::new(),
375 }
376 }
377
378 pub fn into_parts(
384 self,
385 ) -> (
386 DataType,
387 usize,
388 Option<NullBuffer>,
389 usize,
390 Vec<Buffer>,
391 Vec<ArrayData>,
392 ) {
393 let Self {
394 data_type,
395 len,
396 nulls,
397 offset,
398 buffers,
399 child_data,
400 } = self;
401
402 (data_type, len, nulls, offset, buffers, child_data)
403 }
404
405 #[inline]
407 pub const fn builder(data_type: DataType) -> ArrayDataBuilder {
408 ArrayDataBuilder::new(data_type)
409 }
410
411 #[inline]
413 pub const fn data_type(&self) -> &DataType {
414 &self.data_type
415 }
416
417 pub fn buffers(&self) -> &[Buffer] {
419 &self.buffers
420 }
421
422 pub fn child_data(&self) -> &[ArrayData] {
430 &self.child_data[..]
431 }
432
433 #[inline]
435 pub fn is_null(&self, i: usize) -> bool {
436 match &self.nulls {
437 Some(v) => v.is_null(i),
438 None => false,
439 }
440 }
441
442 #[inline]
446 pub fn nulls(&self) -> Option<&NullBuffer> {
447 self.nulls.as_ref()
448 }
449
450 #[inline]
452 pub fn is_valid(&self, i: usize) -> bool {
453 !self.is_null(i)
454 }
455
456 #[inline]
458 pub const fn len(&self) -> usize {
459 self.len
460 }
461
462 #[inline]
464 pub const fn is_empty(&self) -> bool {
465 self.len == 0
466 }
467
468 #[inline]
520 pub const fn offset(&self) -> usize {
521 self.offset
522 }
523
524 #[inline]
526 pub fn null_count(&self) -> usize {
527 self.nulls
528 .as_ref()
529 .map(|x| x.null_count())
530 .unwrap_or_default()
531 }
532
533 pub fn get_buffer_memory_size(&self) -> usize {
545 let mut size = 0;
546 for buffer in &self.buffers {
547 size += buffer.capacity();
548 }
549 if let Some(bitmap) = &self.nulls {
550 size += bitmap.buffer().capacity()
551 }
552 for child in &self.child_data {
553 size += child.get_buffer_memory_size();
554 }
555 size
556 }
557
558 pub fn get_slice_memory_size(&self) -> Result<usize, ArrowError> {
573 let mut result: usize = 0;
574 let layout = layout(&self.data_type);
575
576 for spec in &layout.buffers {
577 match spec {
578 BufferSpec::FixedWidth { byte_width, .. } => {
579 let len = match self.data_type {
582 DataType::Utf8
583 | DataType::LargeUtf8
584 | DataType::Binary
585 | DataType::LargeBinary
586 | DataType::List(_)
587 | DataType::LargeList(_)
588 | DataType::Map(_, _) => self.len + 1,
589 _ => self.len,
590 };
591 let buffer_size = len.checked_mul(*byte_width).ok_or_else(|| {
592 ArrowError::ComputeError(
593 "Integer overflow computing buffer size".to_string(),
594 )
595 })?;
596 result += buffer_size;
597 }
598 BufferSpec::VariableWidth => {
599 let buffer_len = match self.data_type {
600 DataType::Utf8 | DataType::Binary => {
601 let offsets = self.typed_offsets::<i32>()?;
602 (offsets[self.len] - offsets[0]) as usize
603 }
604 DataType::LargeUtf8 | DataType::LargeBinary => {
605 let offsets = self.typed_offsets::<i64>()?;
606 (offsets[self.len] - offsets[0]) as usize
607 }
608 _ => {
609 return Err(ArrowError::NotYetImplemented(format!(
610 "Invalid data type for VariableWidth buffer. Expected Utf8, LargeUtf8, Binary or LargeBinary. Got {}",
611 self.data_type
612 )));
613 }
614 };
615 result += buffer_len;
616 }
617 BufferSpec::BitMap => {
618 let buffer_size = bit_util::ceil(self.len, 8);
619 result += buffer_size;
620 }
621 BufferSpec::AlwaysNull => {
622 }
624 }
625 }
626
627 if layout.variadic {
628 for buffer in self.buffers.iter().skip(layout.buffers.len()) {
630 result += buffer.capacity();
631 }
632 }
633
634 if self.nulls().is_some() {
635 result += bit_util::ceil(self.len, 8);
636 }
637
638 for child in &self.child_data {
639 result += child.get_slice_memory_size()?;
640 }
641 Ok(result)
642 }
643
644 pub fn get_array_memory_size(&self) -> usize {
653 let mut size = mem::size_of_val(self);
654
655 for buffer in &self.buffers {
657 size += mem::size_of::<Buffer>();
658 size += buffer.capacity();
659 }
660 if let Some(nulls) = &self.nulls {
661 size += nulls.buffer().capacity();
662 }
663 for child in &self.child_data {
664 size += child.get_array_memory_size();
665 }
666
667 size
668 }
669
670 pub fn slice(&self, offset: usize, length: usize) -> ArrayData {
678 let end = offset
679 .checked_add(length)
680 .expect("offset + length overflow");
681 assert!(end <= self.len());
682
683 if let DataType::Struct(_) = self.data_type() {
684 let child_offset = self.offset + offset;
689 ArrayData {
690 data_type: self.data_type().clone(),
691 len: length,
692 offset: 0,
693 buffers: self.buffers.clone(),
694 child_data: self
695 .child_data()
696 .iter()
697 .map(|data| data.slice(child_offset, length))
698 .collect(),
699 nulls: self.nulls.as_ref().map(|x| x.slice(offset, length)),
702 }
703 } else {
704 let mut new_data = self.clone();
705
706 new_data.len = length;
707 new_data.offset = offset + self.offset;
708 new_data.nulls = self.nulls.as_ref().map(|x| x.slice(offset, length));
709
710 new_data
711 }
712 }
713
714 pub fn buffer<T: ArrowNativeType>(&self, buffer: usize) -> &[T] {
721 &self.buffers()[buffer].typed_data()[self.offset..]
722 }
723
724 pub fn new_null(data_type: &DataType, len: usize) -> Self {
730 let bit_len = bit_util::ceil(len, 8);
731 let zeroed = |len: usize| Buffer::from(MutableBuffer::from_len_zeroed(len));
732
733 let (buffers, child_data, has_nulls) = match data_type.primitive_width() {
734 Some(width) => (vec![zeroed(width * len)], vec![], true),
735 None => match data_type {
736 DataType::Null => (vec![], vec![], false),
737 DataType::Boolean => (vec![zeroed(bit_len)], vec![], true),
738 DataType::Binary | DataType::Utf8 => {
739 (vec![zeroed((len + 1) * 4), zeroed(0)], vec![], true)
740 }
741 DataType::BinaryView | DataType::Utf8View => (vec![zeroed(len * 16)], vec![], true),
742 DataType::LargeBinary | DataType::LargeUtf8 => {
743 (vec![zeroed((len + 1) * 8), zeroed(0)], vec![], true)
744 }
745 DataType::FixedSizeBinary(i) => {
746 if *i < 0 {
747 panic!("cannot construct null data from FixedSizeBinary({i})");
748 }
749 (vec![zeroed(*i as usize * len)], vec![], true)
750 }
751 DataType::List(f) | DataType::Map(f, _) => (
752 vec![zeroed((len + 1) * 4)],
753 vec![ArrayData::new_empty(f.data_type())],
754 true,
755 ),
756 DataType::LargeList(f) => (
757 vec![zeroed((len + 1) * 8)],
758 vec![ArrayData::new_empty(f.data_type())],
759 true,
760 ),
761 DataType::ListView(f) => (
762 vec![zeroed(len * 4), zeroed(len * 4)],
763 vec![ArrayData::new_empty(f.data_type())],
764 true,
765 ),
766 DataType::LargeListView(f) => (
767 vec![zeroed(len * 8), zeroed(len * 8)],
768 vec![ArrayData::new_empty(f.data_type())],
769 true,
770 ),
771 DataType::FixedSizeList(f, list_len) => (
772 vec![],
773 vec![ArrayData::new_null(f.data_type(), *list_len as usize * len)],
774 true,
775 ),
776 DataType::Struct(fields) => (
777 vec![],
778 fields
779 .iter()
780 .map(|f| Self::new_null(f.data_type(), len))
781 .collect(),
782 true,
783 ),
784 DataType::Dictionary(k, v) => (
785 vec![zeroed(k.primitive_width().unwrap() * len)],
786 vec![ArrayData::new_empty(v.as_ref())],
787 true,
788 ),
789 DataType::Union(f, mode) => {
790 let (id, _) = f.iter().next().unwrap();
791 let ids = Buffer::from_iter(std::iter::repeat_n(id, len));
792 let buffers = match mode {
793 UnionMode::Sparse => vec![ids],
794 UnionMode::Dense => {
795 let end_offset = i32::from_usize(len).unwrap();
796 vec![ids, Buffer::from_iter(0_i32..end_offset)]
797 }
798 };
799
800 let children = f
801 .iter()
802 .enumerate()
803 .map(|(idx, (_, f))| {
804 if idx == 0 || *mode == UnionMode::Sparse {
805 Self::new_null(f.data_type(), len)
806 } else {
807 Self::new_empty(f.data_type())
808 }
809 })
810 .collect();
811
812 (buffers, children, false)
813 }
814 DataType::RunEndEncoded(r, v) => {
815 if len == 0 {
816 let runs = ArrayData::new_empty(r.data_type());
818 let values = ArrayData::new_empty(v.data_type());
819 (vec![], vec![runs, values], false)
820 } else {
821 let runs = match r.data_type() {
822 DataType::Int16 => {
823 let i = i16::from_usize(len).expect("run overflow");
824 Buffer::from_slice_ref([i])
825 }
826 DataType::Int32 => {
827 let i = i32::from_usize(len).expect("run overflow");
828 Buffer::from_slice_ref([i])
829 }
830 DataType::Int64 => {
831 let i = i64::from_usize(len).expect("run overflow");
832 Buffer::from_slice_ref([i])
833 }
834 dt => unreachable!("Invalid run ends data type {dt}"),
835 };
836
837 let builder = ArrayData::builder(r.data_type().clone())
838 .len(1)
839 .buffers(vec![runs]);
840
841 let runs = unsafe { builder.build_unchecked() };
844 (
845 vec![],
846 vec![runs, ArrayData::new_null(v.data_type(), 1)],
847 false,
848 )
849 }
850 }
851 DataType::Int8
853 | DataType::Int16
854 | DataType::Int32
855 | DataType::Int64
856 | DataType::UInt8
857 | DataType::UInt16
858 | DataType::UInt32
859 | DataType::UInt64
860 | DataType::Float16
861 | DataType::Float32
862 | DataType::Float64
863 | DataType::Timestamp(_, _)
864 | DataType::Date32
865 | DataType::Date64
866 | DataType::Time32(_)
867 | DataType::Time64(_)
868 | DataType::Duration(_)
869 | DataType::Interval(_)
870 | DataType::Decimal32(_, _)
871 | DataType::Decimal64(_, _)
872 | DataType::Decimal128(_, _)
873 | DataType::Decimal256(_, _) => unreachable!("{data_type}"),
874 },
875 };
876
877 let mut builder = ArrayDataBuilder::new(data_type.clone())
878 .len(len)
879 .buffers(buffers)
880 .child_data(child_data);
881
882 if has_nulls {
883 builder = builder.nulls(Some(NullBuffer::new_null(len)))
884 }
885
886 unsafe { builder.build_unchecked() }
889 }
890
891 pub fn new_empty(data_type: &DataType) -> Self {
893 Self::new_null(data_type, 0)
894 }
895
896 pub fn align_buffers(&mut self) {
905 let layout = layout(&self.data_type);
906 for (buffer, spec) in self.buffers.iter_mut().zip(&layout.buffers) {
907 if let BufferSpec::FixedWidth { alignment, .. } = spec
908 && buffer.as_ptr().align_offset(*alignment) != 0
909 {
910 *buffer = Buffer::from_slice_ref(buffer.as_ref());
911 }
912 }
913 for data in &mut self.child_data {
915 data.align_buffers()
916 }
917 }
918
919 pub fn validate(&self) -> Result<(), ArrowError> {
930 let len_plus_offset = checked_len_plus_offset(&self.data_type, self.len, self.offset)?;
932
933 let layout = layout(&self.data_type);
935
936 if !layout.can_contain_null_mask && self.nulls.is_some() {
937 return Err(ArrowError::InvalidArgumentError(format!(
938 "Arrays of type {:?} cannot contain a null bitmask",
939 self.data_type,
940 )));
941 }
942
943 if self.buffers.len() < layout.buffers.len()
945 || (!layout.variadic && self.buffers.len() != layout.buffers.len())
946 {
947 return Err(ArrowError::InvalidArgumentError(format!(
948 "Expected {} buffers in array of type {:?}, got {}",
949 layout.buffers.len(),
950 self.data_type,
951 self.buffers.len(),
952 )));
953 }
954
955 for (i, (buffer, spec)) in self.buffers.iter().zip(layout.buffers.iter()).enumerate() {
956 match spec {
957 BufferSpec::FixedWidth {
958 byte_width,
959 alignment,
960 } => {
961 let min_buffer_size = len_plus_offset.saturating_mul(*byte_width);
962
963 if buffer.len() < min_buffer_size {
964 return Err(ArrowError::InvalidArgumentError(format!(
965 "Need at least {} bytes in buffers[{}] in array of type {:?}, but got {}",
966 min_buffer_size,
967 i,
968 self.data_type,
969 buffer.len()
970 )));
971 }
972
973 let align_offset = buffer.as_ptr().align_offset(*alignment);
974 if align_offset != 0 {
975 return Err(ArrowError::InvalidArgumentError(format!(
976 "Misaligned buffers[{i}] in array of type {:?}, offset from expected alignment of {alignment} by {}",
977 self.data_type,
978 align_offset.min(alignment - align_offset)
979 )));
980 }
981 }
982 BufferSpec::VariableWidth => {
983 }
987 BufferSpec::BitMap => {
988 let min_buffer_size = bit_util::ceil(len_plus_offset, 8);
989 if buffer.len() < min_buffer_size {
990 return Err(ArrowError::InvalidArgumentError(format!(
991 "Need at least {} bytes for bitmap in buffers[{}] in array of type {:?}, but got {}",
992 min_buffer_size,
993 i,
994 self.data_type,
995 buffer.len()
996 )));
997 }
998 }
999 BufferSpec::AlwaysNull => {
1000 }
1002 }
1003 }
1004
1005 if let Some(nulls) = self.nulls() {
1007 if nulls.null_count() > self.len {
1008 return Err(ArrowError::InvalidArgumentError(format!(
1009 "null_count {} for an array exceeds length of {} elements",
1010 nulls.null_count(),
1011 self.len
1012 )));
1013 }
1014
1015 if nulls.len() != self.len {
1016 return Err(ArrowError::InvalidArgumentError(format!(
1017 "null buffer incorrect size. got {} expected {}",
1018 nulls.len(),
1019 self.len
1020 )));
1021 }
1022 }
1023
1024 self.validate_child_data()?;
1025
1026 match &self.data_type {
1028 DataType::Utf8 | DataType::Binary => {
1029 self.validate_offsets::<i32>(self.buffers[1].len())?;
1030 }
1031 DataType::LargeUtf8 | DataType::LargeBinary => {
1032 self.validate_offsets::<i64>(self.buffers[1].len())?;
1033 }
1034 DataType::Dictionary(key_type, _value_type) => {
1035 if !DataType::is_dictionary_key_type(key_type) {
1037 return Err(ArrowError::InvalidArgumentError(format!(
1038 "Dictionary key type must be integer, but was {key_type}"
1039 )));
1040 }
1041 }
1042 DataType::RunEndEncoded(run_ends_type, _) => {
1043 if run_ends_type.is_nullable() {
1044 return Err(ArrowError::InvalidArgumentError(
1045 "The nullable should be set to false for the field defining run_ends array.".to_string()
1046 ));
1047 }
1048 if !DataType::is_run_ends_type(run_ends_type.data_type()) {
1049 return Err(ArrowError::InvalidArgumentError(format!(
1050 "RunArray run_ends types must be Int16, Int32 or Int64, but was {}",
1051 run_ends_type.data_type()
1052 )));
1053 }
1054 }
1055 DataType::Map(f, _) if f.is_nullable() => {
1056 return Err(ArrowError::InvalidArgumentError(
1057 "The nullable should be set to false for the map entries field.".to_string(),
1058 ));
1059 }
1060 _ => {}
1061 }
1062
1063 Ok(())
1064 }
1065
1066 fn typed_offsets<T: ArrowNativeType + num_traits::Num>(&self) -> Result<&[T], ArrowError> {
1073 if self.len == 0 && self.buffer_at(0)?.is_empty() {
1075 return Ok(&[]);
1076 }
1077
1078 let len = checked_len_plus_offset(&self.data_type, self.len, 1)?;
1079
1080 self.typed_buffer(0, len)
1081 }
1082
1083 fn typed_buffer<T: ArrowNativeType + num_traits::Num>(
1085 &self,
1086 idx: usize,
1087 len: usize,
1088 ) -> Result<&[T], ArrowError> {
1089 let buffer = self.buffer_at(idx)?;
1090
1091 let required_elements = checked_len_plus_offset(&self.data_type, len, self.offset)?;
1092 let byte_width = mem::size_of::<T>();
1093 let required_len = required_elements.checked_mul(byte_width).ok_or_else(|| {
1094 ArrowError::InvalidArgumentError(format!(
1095 "Buffer {idx} of {} byte length overflow: {} elements of {} bytes exceeds usize",
1096 self.data_type, required_elements, byte_width
1097 ))
1098 })?;
1099
1100 if buffer.len() < required_len {
1101 return Err(ArrowError::InvalidArgumentError(format!(
1102 "Buffer {} of {} isn't large enough. Expected {} bytes got {}",
1103 idx,
1104 self.data_type,
1105 required_len,
1106 buffer.len()
1107 )));
1108 }
1109
1110 Ok(&buffer.typed_data::<T>()[self.offset..required_elements])
1111 }
1112
1113 fn validate_offsets<T: ArrowNativeType + num_traits::Num + std::fmt::Display>(
1116 &self,
1117 values_length: usize,
1118 ) -> Result<(), ArrowError> {
1119 let offsets = self.typed_offsets::<T>()?;
1121 if offsets.is_empty() {
1122 return Ok(());
1123 }
1124
1125 let first_offset = offsets[0].to_usize().ok_or_else(|| {
1126 ArrowError::InvalidArgumentError(format!(
1127 "Error converting offset[0] ({}) to usize for {}",
1128 offsets[0], self.data_type
1129 ))
1130 })?;
1131
1132 let last_offset = offsets[self.len].to_usize().ok_or_else(|| {
1133 ArrowError::InvalidArgumentError(format!(
1134 "Error converting offset[{}] ({}) to usize for {}",
1135 self.len, offsets[self.len], self.data_type
1136 ))
1137 })?;
1138
1139 if first_offset > values_length {
1140 return Err(ArrowError::InvalidArgumentError(format!(
1141 "First offset {} of {} is larger than values length {}",
1142 first_offset, self.data_type, values_length,
1143 )));
1144 }
1145
1146 if last_offset > values_length {
1147 return Err(ArrowError::InvalidArgumentError(format!(
1148 "Last offset {} of {} is larger than values length {}",
1149 last_offset, self.data_type, values_length,
1150 )));
1151 }
1152
1153 if first_offset > last_offset {
1154 return Err(ArrowError::InvalidArgumentError(format!(
1155 "First offset {} in {} is smaller than last offset {}",
1156 first_offset, self.data_type, last_offset,
1157 )));
1158 }
1159
1160 Ok(())
1161 }
1162
1163 fn validate_offsets_and_sizes<T: ArrowNativeType + num_traits::Num + std::fmt::Display>(
1166 &self,
1167 values_length: usize,
1168 ) -> Result<(), ArrowError> {
1169 let offsets: &[T] = self.typed_buffer(0, self.len)?;
1170 let sizes: &[T] = self.typed_buffer(1, self.len)?;
1171 if offsets.len() != sizes.len() {
1172 return Err(ArrowError::ComputeError(format!(
1173 "ListView offsets len {} does not match sizes len {}",
1174 offsets.len(),
1175 sizes.len()
1176 )));
1177 }
1178
1179 for i in 0..sizes.len() {
1180 let size = sizes[i].to_usize().ok_or_else(|| {
1181 ArrowError::InvalidArgumentError(format!(
1182 "Error converting size[{}] ({}) to usize for {}",
1183 i, sizes[i], self.data_type
1184 ))
1185 })?;
1186 let offset = offsets[i].to_usize().ok_or_else(|| {
1187 ArrowError::InvalidArgumentError(format!(
1188 "Error converting offset[{}] ({}) to usize for {}",
1189 i, offsets[i], self.data_type
1190 ))
1191 })?;
1192 if size
1193 .checked_add(offset)
1194 .expect("Offset and size have exceeded the usize boundary")
1195 > values_length
1196 {
1197 return Err(ArrowError::InvalidArgumentError(format!(
1198 "Size {} at index {} is larger than the remaining values for {}",
1199 size, i, self.data_type
1200 )));
1201 }
1202 }
1203 Ok(())
1204 }
1205
1206 fn validate_child_data(&self) -> Result<(), ArrowError> {
1208 match &self.data_type {
1209 DataType::List(field) => {
1210 let values_data = self.get_single_valid_child_data(field.data_type())?;
1211 self.validate_offsets::<i32>(values_data.len)?;
1212 Ok(())
1213 }
1214 DataType::LargeList(field) => {
1215 let values_data = self.get_single_valid_child_data(field.data_type())?;
1216 self.validate_offsets::<i64>(values_data.len)?;
1217 Ok(())
1218 }
1219 DataType::Map(field, _) => {
1220 let DataType::Struct(entries_fields) = field.data_type() else {
1221 return Err(ArrowError::InvalidArgumentError(format!(
1222 "Map field should be a entries struct data type, got {:?} instead",
1223 field.data_type()
1224 )));
1225 };
1226 if entries_fields.len() != 2 {
1227 return Err(ArrowError::InvalidArgumentError(format!(
1228 "Map entries data type should be a struct containing 2 fields, got {} fields",
1229 entries_fields.len()
1230 )));
1231 }
1232
1233 if entries_fields[0].is_nullable() {
1235 return Err(ArrowError::InvalidArgumentError(
1236 "Map key field must not be nullable".to_string(),
1237 ));
1238 }
1239 let values_data = self.get_single_valid_child_data(field.data_type())?;
1240 self.validate_offsets::<i32>(values_data.len)?;
1241 Ok(())
1242 }
1243 DataType::ListView(field) => {
1244 let values_data = self.get_single_valid_child_data(field.data_type())?;
1245 self.validate_offsets_and_sizes::<i32>(values_data.len)?;
1246 Ok(())
1247 }
1248 DataType::LargeListView(field) => {
1249 let values_data = self.get_single_valid_child_data(field.data_type())?;
1250 self.validate_offsets_and_sizes::<i64>(values_data.len)?;
1251 Ok(())
1252 }
1253 DataType::FixedSizeList(field, list_size) => {
1254 let values_data = self.get_single_valid_child_data(field.data_type())?;
1255
1256 let list_size: usize = (*list_size).try_into().map_err(|_| {
1257 ArrowError::InvalidArgumentError(format!(
1258 "{} has a negative list_size {}",
1259 self.data_type, list_size
1260 ))
1261 })?;
1262
1263 let expected_values_len = self.len
1264 .checked_mul(list_size)
1265 .expect("integer overflow computing expected number of expected values in FixedListSize");
1266
1267 if values_data.len < expected_values_len {
1268 return Err(ArrowError::InvalidArgumentError(format!(
1269 "Values length {} is less than the length ({}) multiplied by the value size ({}) for {}",
1270 values_data.len, self.len, list_size, self.data_type
1271 )));
1272 }
1273
1274 Ok(())
1275 }
1276 DataType::Struct(fields) => {
1277 self.validate_num_child_data(fields.len())?;
1278 let len_plus_offset =
1279 checked_len_plus_offset(&self.data_type, self.len, self.offset)?;
1280 for (i, field) in fields.iter().enumerate() {
1281 let field_data = self.get_valid_child_data(i, field.data_type())?;
1282
1283 if field_data.len < len_plus_offset {
1285 return Err(ArrowError::InvalidArgumentError(format!(
1286 "{} child array #{} for field {} has length smaller than expected for struct array ({} < {})",
1287 self.data_type,
1288 i,
1289 field.name(),
1290 field_data.len,
1291 len_plus_offset
1292 )));
1293 }
1294 }
1295 Ok(())
1296 }
1297 DataType::RunEndEncoded(run_ends_field, values_field) => {
1298 self.validate_num_child_data(2)?;
1299 let run_ends_data = self.get_valid_child_data(0, run_ends_field.data_type())?;
1300 let values_data = self.get_valid_child_data(1, values_field.data_type())?;
1301 if run_ends_data.len != values_data.len {
1302 return Err(ArrowError::InvalidArgumentError(format!(
1303 "The run_ends array length should be the same as values array length. Run_ends array length is {}, values array length is {}",
1304 run_ends_data.len, values_data.len
1305 )));
1306 }
1307 if run_ends_data.nulls.is_some() {
1308 return Err(ArrowError::InvalidArgumentError(
1309 "Found null values in run_ends array. The run_ends array should not have null values.".to_string(),
1310 ));
1311 }
1312 Ok(())
1313 }
1314 DataType::Union(fields, mode) => {
1315 self.validate_num_child_data(fields.len())?;
1316
1317 for (i, (_, field)) in fields.iter().enumerate() {
1318 let field_data = self.get_valid_child_data(i, field.data_type())?;
1319
1320 if mode == &UnionMode::Sparse {
1321 let len_plus_offset =
1322 checked_len_plus_offset(&self.data_type, self.len, self.offset)?;
1323 if field_data.len < len_plus_offset {
1324 return Err(ArrowError::InvalidArgumentError(format!(
1325 "Sparse union child array #{} has length smaller than expected for union array ({} < {})",
1326 i, field_data.len, len_plus_offset
1327 )));
1328 }
1329 }
1330 }
1331 Ok(())
1332 }
1333 DataType::Dictionary(_key_type, value_type) => {
1334 self.get_single_valid_child_data(value_type)?;
1335 Ok(())
1336 }
1337 _ => {
1338 if !self.child_data.is_empty() {
1340 return Err(ArrowError::InvalidArgumentError(format!(
1341 "Expected no child arrays for type {} but got {}",
1342 self.data_type,
1343 self.child_data.len()
1344 )));
1345 }
1346 Ok(())
1347 }
1348 }
1349 }
1350
1351 fn get_single_valid_child_data(
1355 &self,
1356 expected_type: &DataType,
1357 ) -> Result<&ArrayData, ArrowError> {
1358 self.validate_num_child_data(1)?;
1359 self.get_valid_child_data(0, expected_type)
1360 }
1361
1362 fn buffer_at(&self, idx: usize) -> Result<&Buffer, ArrowError> {
1367 self.buffers.get(idx).ok_or_else(|| {
1368 ArrowError::InvalidArgumentError(format!(
1369 "{} should contain at least {} buffer(s), had {}",
1370 self.data_type,
1371 idx + 1,
1372 self.buffers.len()
1373 ))
1374 })
1375 }
1376
1377 fn child_at(&self, idx: usize) -> Result<&ArrayData, ArrowError> {
1382 self.child_data.get(idx).ok_or_else(|| {
1383 ArrowError::InvalidArgumentError(format!(
1384 "{} should contain at least {} child data array(s), had {}",
1385 self.data_type,
1386 idx + 1,
1387 self.child_data.len()
1388 ))
1389 })
1390 }
1391
1392 fn validate_num_child_data(&self, expected_len: usize) -> Result<(), ArrowError> {
1394 if self.child_data.len() != expected_len {
1395 Err(ArrowError::InvalidArgumentError(format!(
1396 "Value data for {} should contain {} child data array(s), had {}",
1397 self.data_type,
1398 expected_len,
1399 self.child_data.len()
1400 )))
1401 } else {
1402 Ok(())
1403 }
1404 }
1405
1406 fn get_valid_child_data(
1409 &self,
1410 i: usize,
1411 expected_type: &DataType,
1412 ) -> Result<&ArrayData, ArrowError> {
1413 let values_data = self.child_data.get(i).ok_or_else(|| {
1414 ArrowError::InvalidArgumentError(format!(
1415 "{} did not have enough child arrays. Expected at least {} but had only {}",
1416 self.data_type,
1417 i + 1,
1418 self.child_data.len()
1419 ))
1420 })?;
1421
1422 if expected_type != &values_data.data_type {
1423 return Err(ArrowError::InvalidArgumentError(format!(
1424 "Child type mismatch for {}. Expected {} but child data had {}",
1425 self.data_type, expected_type, values_data.data_type
1426 )));
1427 }
1428
1429 values_data.validate()?;
1430 Ok(values_data)
1431 }
1432
1433 pub fn validate_data(&self) -> Result<(), ArrowError> {
1449 self.validate()?;
1450
1451 self.validate_nulls()?;
1452 self.validate_values()?;
1453 Ok(())
1454 }
1455
1456 pub fn validate_full(&self) -> Result<(), ArrowError> {
1461 self.validate_data()?;
1462 self.child_data
1464 .iter()
1465 .enumerate()
1466 .try_for_each(|(i, child_data)| {
1467 child_data.validate_full().map_err(|e| {
1468 ArrowError::InvalidArgumentError(format!(
1469 "{} child #{} invalid: {}",
1470 self.data_type, i, e
1471 ))
1472 })
1473 })?;
1474 Ok(())
1475 }
1476
1477 pub fn validate_nulls(&self) -> Result<(), ArrowError> {
1487 if let Some(nulls) = &self.nulls {
1488 let actual = nulls.len() - nulls.inner().count_set_bits();
1489 if actual != nulls.null_count() {
1490 return Err(ArrowError::InvalidArgumentError(format!(
1491 "null_count value ({}) doesn't match actual number of nulls in array ({})",
1492 nulls.null_count(),
1493 actual
1494 )));
1495 }
1496 }
1497
1498 match &self.data_type {
1503 DataType::List(f) | DataType::LargeList(f) | DataType::Map(f, _) => {
1504 if !f.is_nullable() {
1505 let child = &self.child_data[0];
1506 self.validate_non_nullable(None, child, child.nulls())?
1507 }
1508 }
1509 DataType::FixedSizeList(field, len) => {
1510 let child = &self.child_data[0];
1511 if !field.is_nullable() {
1512 match &self.nulls {
1513 Some(nulls) => {
1514 let element_len = *len as usize;
1515 let expanded = nulls.expand(element_len);
1516 self.validate_non_nullable(Some(&expanded), child, child.nulls())?;
1517 }
1518 None => self.validate_non_nullable(None, child, child.nulls())?,
1519 }
1520 }
1521 }
1522 DataType::Struct(fields) => {
1523 for (field, child) in fields.iter().zip(&self.child_data) {
1524 if !field.is_nullable() {
1525 let child_nulls = child
1526 .nulls()
1527 .map(|nulls| nulls.slice(self.offset, self.len));
1528 self.validate_non_nullable(self.nulls(), child, child_nulls.as_ref())?
1529 }
1530 }
1531 }
1532 _ => {}
1533 }
1534
1535 Ok(())
1536 }
1537
1538 fn validate_non_nullable(
1540 &self,
1541 mask: Option<&NullBuffer>,
1542 child: &ArrayData,
1543 child_nulls: Option<&NullBuffer>,
1544 ) -> Result<(), ArrowError> {
1545 let Some(mask) = mask else {
1546 return match child_nulls.map(NullBuffer::null_count).unwrap_or_default() {
1547 0 => Ok(()),
1548 _ => Err(ArrowError::InvalidArgumentError(format!(
1549 "non-nullable child of type {} contains nulls not present in parent {}",
1550 child.data_type, self.data_type
1551 ))),
1552 };
1553 };
1554
1555 match child_nulls {
1556 Some(nulls) if !mask.contains(nulls) => Err(ArrowError::InvalidArgumentError(format!(
1557 "non-nullable child of type {} contains nulls not present in parent",
1558 child.data_type
1559 ))),
1560 _ => Ok(()),
1561 }
1562 }
1563
1564 pub fn validate_values(&self) -> Result<(), ArrowError> {
1570 match &self.data_type {
1571 DataType::Utf8 => self.validate_utf8::<i32>(),
1572 DataType::LargeUtf8 => self.validate_utf8::<i64>(),
1573 DataType::Binary => self.validate_offsets_full::<i32>(self.buffer_at(1)?.len()),
1574 DataType::LargeBinary => self.validate_offsets_full::<i64>(self.buffer_at(1)?.len()),
1575 DataType::BinaryView => {
1576 let views = self.typed_buffer::<u128>(0, self.len)?;
1577 validate_binary_view(views, &self.buffers[1..])
1578 }
1579 DataType::Utf8View => {
1580 let views = self.typed_buffer::<u128>(0, self.len)?;
1581 validate_string_view(views, &self.buffers[1..])
1582 }
1583 DataType::List(_) | DataType::Map(_, _) => {
1584 let child = self.child_at(0)?;
1585 self.validate_offsets_full::<i32>(child.len)
1586 }
1587 DataType::LargeList(_) => {
1588 let child = self.child_at(0)?;
1589 self.validate_offsets_full::<i64>(child.len)
1590 }
1591 DataType::Union(_, _) => {
1592 Ok(())
1598 }
1599 DataType::Dictionary(key_type, _value_type) => {
1600 let dictionary_length = self.child_at(0)?.len;
1601 let dictionary_length = i64::try_from(dictionary_length).map_err(|_| {
1602 ArrowError::InvalidArgumentError(format!(
1603 "Dictionary of {dictionary_length} values is too long for an i64"
1604 ))
1605 })?;
1606 let max_value = dictionary_length - 1;
1607 match key_type.as_ref() {
1608 DataType::UInt8 => self.check_bounds::<u8>(max_value),
1609 DataType::UInt16 => self.check_bounds::<u16>(max_value),
1610 DataType::UInt32 => self.check_bounds::<u32>(max_value),
1611 DataType::UInt64 => self.check_bounds::<u64>(max_value),
1612 DataType::Int8 => self.check_bounds::<i8>(max_value),
1613 DataType::Int16 => self.check_bounds::<i16>(max_value),
1614 DataType::Int32 => self.check_bounds::<i32>(max_value),
1615 DataType::Int64 => self.check_bounds::<i64>(max_value),
1616 _ => Err(ArrowError::InvalidArgumentError(format!(
1617 "Dictionary key type must be an integer, got {key_type}"
1618 ))),
1619 }
1620 }
1621 DataType::RunEndEncoded(run_ends, _values) => {
1622 let run_ends_data = self.child_at(0)?;
1623 match run_ends.data_type() {
1624 DataType::Int16 => run_ends_data.check_run_ends::<i16>(),
1625 DataType::Int32 => run_ends_data.check_run_ends::<i32>(),
1626 DataType::Int64 => run_ends_data.check_run_ends::<i64>(),
1627 data_type => Err(ArrowError::InvalidArgumentError(format!(
1628 "Run end type must be Int16, Int32 or Int64, got {data_type}"
1629 ))),
1630 }
1631 }
1632 _ => {
1633 Ok(())
1635 }
1636 }
1637 }
1638
1639 fn validate_each_offset<T, V>(&self, offset_limit: usize, validate: V) -> Result<(), ArrowError>
1650 where
1651 T: ArrowNativeType + TryInto<usize> + num_traits::Num + std::fmt::Display,
1652 V: Fn(usize, Range<usize>) -> Result<(), ArrowError>,
1653 {
1654 self.typed_offsets::<T>()?
1655 .iter()
1656 .enumerate()
1657 .map(|(i, x)| {
1658 let r = x.to_usize().ok_or_else(|| {
1660 ArrowError::InvalidArgumentError(format!(
1661 "Offset invariant failure: Could not convert offset {x} to usize at position {i}"))}
1662 );
1663 match r {
1665 Ok(n) if n <= offset_limit => Ok((i, n)),
1666 Ok(_) => Err(ArrowError::InvalidArgumentError(format!(
1667 "Offset invariant failure: offset at position {i} out of bounds: {x} > {offset_limit}"))
1668 ),
1669 Err(e) => Err(e),
1670 }
1671 })
1672 .scan(0_usize, |start, end| {
1673 match end {
1675 Ok((i, end)) if *start <= end => {
1676 let range = Some(Ok((i, *start..end)));
1677 *start = end;
1678 range
1679 }
1680 Ok((i, end)) => Some(Err(ArrowError::InvalidArgumentError(format!(
1681 "Offset invariant failure: non-monotonic offset at slot {}: {} > {}",
1682 i - 1, start, end))
1683 )),
1684 Err(err) => Some(Err(err)),
1685 }
1686 })
1687 .skip(1) .try_for_each(|res: Result<(usize, Range<usize>), ArrowError>| {
1689 let (item_index, range) = res?;
1690 validate(item_index-1, range)
1691 })
1692 }
1693
1694 fn validate_utf8<T>(&self) -> Result<(), ArrowError>
1697 where
1698 T: ArrowNativeType + TryInto<usize> + num_traits::Num + std::fmt::Display,
1699 {
1700 let values_buffer = &self.buffer_at(1)?.as_slice();
1701 if let Ok(values_str) = std::str::from_utf8(values_buffer) {
1702 self.validate_each_offset::<T, _>(values_buffer.len(), |string_index, range| {
1704 if !values_str.is_char_boundary(range.start)
1705 || !values_str.is_char_boundary(range.end)
1706 {
1707 return Err(ArrowError::InvalidArgumentError(format!(
1708 "incomplete utf-8 byte sequence from index {string_index}"
1709 )));
1710 }
1711 Ok(())
1712 })
1713 } else {
1714 self.validate_each_offset::<T, _>(values_buffer.len(), |string_index, range| {
1716 std::str::from_utf8(&values_buffer[range.clone()]).map_err(|e| {
1717 ArrowError::InvalidArgumentError(format!(
1718 "Invalid UTF8 sequence at string index {string_index} ({range:?}): {e}"
1719 ))
1720 })?;
1721 Ok(())
1722 })
1723 }
1724 }
1725
1726 fn validate_offsets_full<T>(&self, offset_limit: usize) -> Result<(), ArrowError>
1729 where
1730 T: ArrowNativeType + TryInto<usize> + num_traits::Num + std::fmt::Display,
1731 {
1732 self.validate_each_offset::<T, _>(offset_limit, |_string_index, _range| {
1733 Ok(())
1736 })
1737 }
1738
1739 fn check_bounds<T>(&self, max_value: i64) -> Result<(), ArrowError>
1742 where
1743 T: ArrowNativeType + TryInto<i64> + num_traits::Num + std::fmt::Display,
1744 {
1745 let indexes: &[T] = self.typed_buffer::<T>(0, self.len)?;
1748
1749 indexes.iter().enumerate().try_for_each(|(i, &dict_index)| {
1750 if self.is_null(i) {
1752 return Ok(());
1753 }
1754 let dict_index: i64 = dict_index.try_into().map_err(|_| {
1755 ArrowError::InvalidArgumentError(format!(
1756 "Value at position {i} out of bounds: {dict_index} (can not convert to i64)"
1757 ))
1758 })?;
1759
1760 if dict_index < 0 || dict_index > max_value {
1761 return Err(ArrowError::InvalidArgumentError(format!(
1762 "Value at position {i} out of bounds: {dict_index} (should be in [0, {max_value}])"
1763 )));
1764 }
1765 Ok(())
1766 })
1767 }
1768
1769 fn check_run_ends<T>(&self) -> Result<(), ArrowError>
1771 where
1772 T: ArrowNativeType + TryInto<i64> + num_traits::Num + std::fmt::Display,
1773 {
1774 let values = self.typed_buffer::<T>(0, self.len)?;
1775 let mut prev_value = 0_i64;
1776 values.iter().enumerate().try_for_each(|(ix, &inp_value)| {
1777 let value: i64 = inp_value.try_into().map_err(|_| {
1778 ArrowError::InvalidArgumentError(format!(
1779 "Value at position {ix} out of bounds: {inp_value} (can not convert to i64)"
1780 ))
1781 })?;
1782 if value <= 0_i64 {
1783 return Err(ArrowError::InvalidArgumentError(format!(
1784 "The values in run_ends array should be strictly positive. Found value {value} at index {ix} that does not match the criteria."
1785 )));
1786 }
1787 if ix > 0 && value <= prev_value {
1788 return Err(ArrowError::InvalidArgumentError(format!(
1789 "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."
1790 )));
1791 }
1792
1793 prev_value = value;
1794 Ok(())
1795 })?;
1796
1797 let len_plus_offset = checked_len_plus_offset(&self.data_type, self.len, self.offset)?;
1798 if prev_value.as_usize() < len_plus_offset {
1799 return Err(ArrowError::InvalidArgumentError(format!(
1800 "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 {len_plus_offset}."
1801 )));
1802 }
1803 Ok(())
1804 }
1805
1806 pub fn ptr_eq(&self, other: &Self) -> bool {
1810 if self.offset != other.offset
1811 || self.len != other.len
1812 || self.data_type != other.data_type
1813 || self.buffers.len() != other.buffers.len()
1814 || self.child_data.len() != other.child_data.len()
1815 {
1816 return false;
1817 }
1818
1819 match (&self.nulls, &other.nulls) {
1820 (Some(a), Some(b)) if !a.inner().ptr_eq(b.inner()) => return false,
1821 (Some(_), None) | (None, Some(_)) => return false,
1822 _ => {}
1823 }
1824
1825 if !self
1826 .buffers
1827 .iter()
1828 .zip(other.buffers.iter())
1829 .all(|(a, b)| a.as_ptr() == b.as_ptr())
1830 {
1831 return false;
1832 }
1833
1834 self.child_data
1835 .iter()
1836 .zip(other.child_data.iter())
1837 .all(|(a, b)| a.ptr_eq(b))
1838 }
1839
1840 pub fn into_builder(self) -> ArrayDataBuilder {
1842 self.into()
1843 }
1844
1845 #[cfg(feature = "pool")]
1852 pub fn claim(&self, pool: &dyn arrow_buffer::MemoryPool) {
1853 for buffer in &self.buffers {
1855 buffer.claim(pool);
1856 }
1857
1858 if let Some(nulls) = &self.nulls {
1860 nulls.claim(pool);
1861 }
1862
1863 for child in &self.child_data {
1865 child.claim(pool);
1866 }
1867 }
1868}
1869
1870pub fn layout(data_type: &DataType) -> DataTypeLayout {
1873 use arrow_schema::IntervalUnit::*;
1876
1877 match data_type {
1878 DataType::Null => DataTypeLayout {
1879 buffers: vec![],
1880 can_contain_null_mask: false,
1881 variadic: false,
1882 },
1883 DataType::Boolean => DataTypeLayout {
1884 buffers: vec![BufferSpec::BitMap],
1885 can_contain_null_mask: true,
1886 variadic: false,
1887 },
1888 DataType::Int8 => DataTypeLayout::new_fixed_width::<i8>(),
1889 DataType::Int16 => DataTypeLayout::new_fixed_width::<i16>(),
1890 DataType::Int32 => DataTypeLayout::new_fixed_width::<i32>(),
1891 DataType::Int64 => DataTypeLayout::new_fixed_width::<i64>(),
1892 DataType::UInt8 => DataTypeLayout::new_fixed_width::<u8>(),
1893 DataType::UInt16 => DataTypeLayout::new_fixed_width::<u16>(),
1894 DataType::UInt32 => DataTypeLayout::new_fixed_width::<u32>(),
1895 DataType::UInt64 => DataTypeLayout::new_fixed_width::<u64>(),
1896 DataType::Float16 => DataTypeLayout::new_fixed_width::<half::f16>(),
1897 DataType::Float32 => DataTypeLayout::new_fixed_width::<f32>(),
1898 DataType::Float64 => DataTypeLayout::new_fixed_width::<f64>(),
1899 DataType::Timestamp(_, _) => DataTypeLayout::new_fixed_width::<i64>(),
1900 DataType::Date32 => DataTypeLayout::new_fixed_width::<i32>(),
1901 DataType::Date64 => DataTypeLayout::new_fixed_width::<i64>(),
1902 DataType::Time32(_) => DataTypeLayout::new_fixed_width::<i32>(),
1903 DataType::Time64(_) => DataTypeLayout::new_fixed_width::<i64>(),
1904 DataType::Interval(YearMonth) => DataTypeLayout::new_fixed_width::<i32>(),
1905 DataType::Interval(DayTime) => DataTypeLayout::new_fixed_width::<IntervalDayTime>(),
1906 DataType::Interval(MonthDayNano) => {
1907 DataTypeLayout::new_fixed_width::<IntervalMonthDayNano>()
1908 }
1909 DataType::Duration(_) => DataTypeLayout::new_fixed_width::<i64>(),
1910 DataType::Decimal32(_, _) => DataTypeLayout::new_fixed_width::<i32>(),
1911 DataType::Decimal64(_, _) => DataTypeLayout::new_fixed_width::<i64>(),
1912 DataType::Decimal128(_, _) => DataTypeLayout::new_fixed_width::<i128>(),
1913 DataType::Decimal256(_, _) => DataTypeLayout::new_fixed_width::<i256>(),
1914 DataType::FixedSizeBinary(size) => {
1915 let spec = BufferSpec::FixedWidth {
1916 byte_width: (*size).try_into().unwrap(),
1917 alignment: mem::align_of::<u8>(),
1918 };
1919 DataTypeLayout {
1920 buffers: vec![spec],
1921 can_contain_null_mask: true,
1922 variadic: false,
1923 }
1924 }
1925 DataType::Binary => DataTypeLayout::new_binary::<i32>(),
1926 DataType::LargeBinary => DataTypeLayout::new_binary::<i64>(),
1927 DataType::Utf8 => DataTypeLayout::new_binary::<i32>(),
1928 DataType::LargeUtf8 => DataTypeLayout::new_binary::<i64>(),
1929 DataType::BinaryView | DataType::Utf8View => DataTypeLayout::new_view(),
1930 DataType::FixedSizeList(_, _) => DataTypeLayout::new_nullable_empty(), DataType::List(_) => DataTypeLayout::new_fixed_width::<i32>(),
1932 DataType::ListView(_) => DataTypeLayout::new_list_view::<i32>(),
1933 DataType::LargeListView(_) => DataTypeLayout::new_list_view::<i64>(),
1934 DataType::LargeList(_) => DataTypeLayout::new_fixed_width::<i64>(),
1935 DataType::Map(_, _) => DataTypeLayout::new_fixed_width::<i32>(),
1936 DataType::Struct(_) => DataTypeLayout::new_nullable_empty(), DataType::RunEndEncoded(_, _) => DataTypeLayout::new_empty(), DataType::Union(_, mode) => {
1939 let type_ids = BufferSpec::FixedWidth {
1940 byte_width: mem::size_of::<i8>(),
1941 alignment: mem::align_of::<i8>(),
1942 };
1943
1944 DataTypeLayout {
1945 buffers: match mode {
1946 UnionMode::Sparse => {
1947 vec![type_ids]
1948 }
1949 UnionMode::Dense => {
1950 vec![
1951 type_ids,
1952 BufferSpec::FixedWidth {
1953 byte_width: mem::size_of::<i32>(),
1954 alignment: mem::align_of::<i32>(),
1955 },
1956 ]
1957 }
1958 },
1959 can_contain_null_mask: false,
1960 variadic: false,
1961 }
1962 }
1963 DataType::Dictionary(key_type, _value_type) => layout(key_type),
1964 }
1965}
1966
1967#[derive(Debug, PartialEq, Eq)]
1969pub struct DataTypeLayout {
1971 pub buffers: Vec<BufferSpec>,
1973
1974 pub can_contain_null_mask: bool,
1976
1977 pub variadic: bool,
1981}
1982
1983impl DataTypeLayout {
1984 pub fn new_fixed_width<T>() -> Self {
1986 Self {
1987 buffers: vec![BufferSpec::FixedWidth {
1988 byte_width: mem::size_of::<T>(),
1989 alignment: mem::align_of::<T>(),
1990 }],
1991 can_contain_null_mask: true,
1992 variadic: false,
1993 }
1994 }
1995
1996 pub fn new_nullable_empty() -> Self {
1999 Self {
2000 buffers: vec![],
2001 can_contain_null_mask: true,
2002 variadic: false,
2003 }
2004 }
2005
2006 pub fn new_empty() -> Self {
2009 Self {
2010 buffers: vec![],
2011 can_contain_null_mask: false,
2012 variadic: false,
2013 }
2014 }
2015
2016 pub fn new_binary<T>() -> Self {
2020 Self {
2021 buffers: vec![
2022 BufferSpec::FixedWidth {
2024 byte_width: mem::size_of::<T>(),
2025 alignment: mem::align_of::<T>(),
2026 },
2027 BufferSpec::VariableWidth,
2029 ],
2030 can_contain_null_mask: true,
2031 variadic: false,
2032 }
2033 }
2034
2035 pub fn new_view() -> Self {
2037 Self {
2038 buffers: vec![BufferSpec::FixedWidth {
2039 byte_width: mem::size_of::<u128>(),
2040 alignment: mem::align_of::<u128>(),
2041 }],
2042 can_contain_null_mask: true,
2043 variadic: true,
2044 }
2045 }
2046
2047 pub fn new_list_view<T>() -> Self {
2049 Self {
2050 buffers: vec![
2051 BufferSpec::FixedWidth {
2052 byte_width: mem::size_of::<T>(),
2053 alignment: mem::align_of::<T>(),
2054 },
2055 BufferSpec::FixedWidth {
2056 byte_width: mem::size_of::<T>(),
2057 alignment: mem::align_of::<T>(),
2058 },
2059 ],
2060 can_contain_null_mask: true,
2061 variadic: false,
2062 }
2063 }
2064}
2065
2066#[derive(Debug, PartialEq, Eq)]
2068pub enum BufferSpec {
2069 FixedWidth {
2080 byte_width: usize,
2082 alignment: usize,
2084 },
2085 VariableWidth,
2087 BitMap,
2093 AlwaysNull,
2096}
2097
2098impl PartialEq for ArrayData {
2099 fn eq(&self, other: &Self) -> bool {
2100 equal::equal(self, other)
2101 }
2102}
2103
2104#[derive(Debug, Clone)]
2123#[doc(hidden)]
2124pub struct UnsafeFlag(bool);
2125
2126impl UnsafeFlag {
2127 #[inline]
2131 pub const fn new() -> Self {
2132 Self(false)
2133 }
2134
2135 #[inline]
2145 pub unsafe fn set(&mut self, val: bool) {
2146 self.0 = val;
2147 }
2148
2149 #[inline]
2151 pub fn get(&self) -> bool {
2152 self.0
2153 }
2154}
2155
2156impl Default for UnsafeFlag {
2158 fn default() -> Self {
2159 Self::new()
2160 }
2161}
2162
2163#[derive(Debug)]
2165pub struct ArrayDataBuilder {
2166 data_type: DataType,
2167 len: usize,
2168 null_count: Option<usize>,
2169 null_bit_buffer: Option<Buffer>,
2170 nulls: Option<NullBuffer>,
2171 offset: usize,
2172 buffers: Vec<Buffer>,
2173 child_data: Vec<ArrayData>,
2174 align_buffers: bool,
2178 skip_validation: UnsafeFlag,
2188}
2189
2190impl ArrayDataBuilder {
2191 #[inline]
2192 pub const fn new(data_type: DataType) -> Self {
2194 Self {
2195 data_type,
2196 len: 0,
2197 null_count: None,
2198 null_bit_buffer: None,
2199 nulls: None,
2200 offset: 0,
2201 buffers: vec![],
2202 child_data: vec![],
2203 align_buffers: false,
2204 skip_validation: UnsafeFlag::new(),
2205 }
2206 }
2207
2208 pub fn data_type(self, data_type: DataType) -> Self {
2210 Self { data_type, ..self }
2211 }
2212
2213 #[inline]
2214 pub const fn len(mut self, n: usize) -> Self {
2216 self.len = n;
2217 self
2218 }
2219
2220 pub fn nulls(mut self, nulls: Option<NullBuffer>) -> Self {
2222 self.nulls = nulls;
2223 self.null_count = None;
2224 self.null_bit_buffer = None;
2225 self
2226 }
2227
2228 pub fn null_count(mut self, null_count: usize) -> Self {
2230 self.null_count = Some(null_count);
2231 self
2232 }
2233
2234 pub fn null_bit_buffer(mut self, buf: Option<Buffer>) -> Self {
2236 self.nulls = None;
2237 self.null_bit_buffer = buf;
2238 self
2239 }
2240
2241 #[inline]
2243 pub const fn offset(mut self, n: usize) -> Self {
2244 self.offset = n;
2245 self
2246 }
2247
2248 pub fn buffers(mut self, v: Vec<Buffer>) -> Self {
2250 self.buffers = v;
2251 self
2252 }
2253
2254 pub fn add_buffer(mut self, b: Buffer) -> Self {
2256 self.buffers.push(b);
2257 self
2258 }
2259
2260 pub fn add_buffers<I: IntoIterator<Item = Buffer>>(mut self, bs: I) -> Self {
2262 self.buffers.extend(bs);
2263 self
2264 }
2265
2266 pub fn child_data(mut self, v: Vec<ArrayData>) -> Self {
2268 self.child_data = v;
2269 self
2270 }
2271
2272 pub fn add_child_data(mut self, r: ArrayData) -> Self {
2274 self.child_data.push(r);
2275 self
2276 }
2277
2278 pub unsafe fn build_unchecked(self) -> ArrayData {
2294 unsafe { self.skip_validation(true) }.build().unwrap()
2295 }
2296
2297 pub fn build(self) -> Result<ArrayData, ArrowError> {
2306 let Self {
2307 data_type,
2308 len,
2309 null_count,
2310 null_bit_buffer,
2311 nulls,
2312 offset,
2313 buffers,
2314 child_data,
2315 align_buffers,
2316 skip_validation,
2317 } = self;
2318
2319 let validate = !skip_validation.get() || cfg!(feature = "force_validate");
2321 if validate && let Some(buffer) = null_bit_buffer.as_ref() {
2322 let len_plus_offset = checked_len_plus_offset(&data_type, len, offset)?;
2324 let needed_len = bit_util::ceil(len_plus_offset, 8);
2325 if buffer.len() < needed_len {
2326 return Err(ArrowError::InvalidArgumentError(format!(
2327 "null_bit_buffer size too small. got {} needed {}",
2328 buffer.len(),
2329 needed_len
2330 )));
2331 }
2332 }
2333
2334 let nulls = nulls
2335 .or_else(|| {
2336 let buffer = null_bit_buffer?;
2337 let buffer = BooleanBuffer::new(buffer, offset, len);
2338 Some(match null_count {
2339 Some(n) => {
2340 unsafe { NullBuffer::new_unchecked(buffer, n) }
2342 }
2343 None => NullBuffer::new(buffer),
2344 })
2345 })
2346 .filter(|b| b.null_count() != 0);
2347
2348 let mut data = ArrayData {
2349 data_type,
2350 len,
2351 offset,
2352 buffers,
2353 child_data,
2354 nulls,
2355 };
2356
2357 if align_buffers {
2358 data.align_buffers();
2359 }
2360
2361 if validate {
2362 data.validate_data()?;
2363 }
2364 Ok(data)
2365 }
2366
2367 pub fn align_buffers(mut self, align_buffers: bool) -> Self {
2383 self.align_buffers = align_buffers;
2384 self
2385 }
2386
2387 pub unsafe fn skip_validation(mut self, skip_validation: bool) -> Self {
2401 unsafe {
2402 self.skip_validation.set(skip_validation);
2403 }
2404 self
2405 }
2406}
2407
2408impl From<ArrayData> for ArrayDataBuilder {
2409 fn from(d: ArrayData) -> Self {
2410 Self {
2411 data_type: d.data_type,
2412 len: d.len,
2413 offset: d.offset,
2414 buffers: d.buffers,
2415 child_data: d.child_data,
2416 nulls: d.nulls,
2417 null_bit_buffer: None,
2418 null_count: None,
2419 align_buffers: false,
2420 skip_validation: UnsafeFlag::new(),
2421 }
2422 }
2423}
2424
2425pub(crate) fn get_fixed_size_binary_width(data_type: &DataType) -> usize {
2430 match data_type {
2431 DataType::FixedSizeBinary(i) => {
2432 if *i < 0 {
2433 panic!("cannot compare FixedSizeBinary({})", *i);
2434 }
2435 *i as usize
2436 }
2437 _ => unreachable!(),
2438 }
2439}
2440
2441#[cfg(test)]
2442mod tests {
2443 use super::*;
2444 use crate::ByteView;
2445 use crate::transform::MutableArrayData;
2446 use arrow_buffer::{OffsetBuffer, ScalarBuffer};
2447 use arrow_schema::{Field, Fields};
2448
2449 fn make_i32_buffer(n: usize) -> Buffer {
2453 Buffer::from_slice_ref(vec![42i32; n])
2454 }
2455
2456 fn make_f32_buffer(n: usize) -> Buffer {
2458 Buffer::from_slice_ref(vec![42f32; n])
2459 }
2460
2461 #[test]
2462 fn test_builder() {
2463 let v = (0..25).collect::<Vec<i32>>();
2465 let b1 = Buffer::from_slice_ref(&v);
2466 let arr_data = ArrayData::builder(DataType::Int32)
2467 .len(20)
2468 .offset(5)
2469 .add_buffer(b1)
2470 .null_bit_buffer(Some(Buffer::from([
2471 0b01011111, 0b10110101, 0b01100011, 0b00011110,
2472 ])))
2473 .build()
2474 .unwrap();
2475
2476 assert_eq!(20, arr_data.len());
2477 assert_eq!(10, arr_data.null_count());
2478 assert_eq!(5, arr_data.offset());
2479 assert_eq!(1, arr_data.buffers().len());
2480 assert_eq!(
2481 Buffer::from_slice_ref(&v).as_slice(),
2482 arr_data.buffers()[0].as_slice()
2483 );
2484 }
2485
2486 #[test]
2487 fn test_builder_with_child_data() {
2488 let child_arr_data = ArrayData::try_new(
2489 DataType::Int32,
2490 5,
2491 None,
2492 0,
2493 vec![Buffer::from_slice_ref([1i32, 2, 3, 4, 5])],
2494 vec![],
2495 )
2496 .unwrap();
2497
2498 let field = Arc::new(Field::new("x", DataType::Int32, true));
2499 let data_type = DataType::Struct(vec![field].into());
2500
2501 let arr_data = ArrayData::builder(data_type)
2502 .len(5)
2503 .offset(0)
2504 .add_child_data(child_arr_data.clone())
2505 .build()
2506 .unwrap();
2507
2508 assert_eq!(5, arr_data.len());
2509 assert_eq!(1, arr_data.child_data().len());
2510 assert_eq!(child_arr_data, arr_data.child_data()[0]);
2511 }
2512
2513 #[test]
2514 fn test_struct_validation_accounts_for_parent_offset() {
2515 let data_type =
2516 DataType::Struct(Fields::from(vec![Field::new("x", DataType::Int32, false)]));
2517 let child = ArrayData::builder(DataType::Int32)
2518 .len(5)
2519 .add_buffer(Buffer::from_slice_ref([0, 1, 2, 3, 4]))
2520 .build()
2521 .unwrap();
2522
2523 let err = ArrayData::builder(data_type)
2525 .len(5)
2526 .offset(1)
2527 .add_child_data(child)
2528 .build()
2529 .unwrap_err()
2530 .to_string();
2531
2532 assert!(err.contains(
2533 "child array #0 for field x has length smaller than expected for struct array (5 < 6)"
2534 ));
2535 }
2536
2537 #[test]
2538 fn test_struct_non_nullable_child_nulls_account_for_parent_offset() {
2539 let build = |parent_nulls| {
2540 let child = ArrayData::builder(DataType::Int32)
2541 .len(5)
2542 .add_buffer(Buffer::from_slice_ref([0, 1, 2, 3, 4]))
2543 .nulls(Some(NullBuffer::new(BooleanBuffer::from(vec![
2544 true, true, false, true, true,
2545 ]))))
2546 .build()
2547 .unwrap();
2548
2549 ArrayData::builder(DataType::Struct(Fields::from(vec![Field::new(
2550 "x",
2551 DataType::Int32,
2552 false,
2553 )])))
2554 .len(4)
2555 .offset(1)
2556 .nulls(Some(NullBuffer::new(BooleanBuffer::from(parent_nulls))))
2557 .add_child_data(child)
2558 .build()
2559 };
2560
2561 assert!(build(vec![true, false, true, true]).is_ok());
2562 assert!(build(vec![true, true, false, true]).is_err());
2563 }
2564
2565 #[test]
2566 fn test_struct_equal_accounts_for_parent_offset() {
2567 let data_type =
2568 DataType::Struct(Fields::from(vec![Field::new("x", DataType::Int32, false)]));
2569
2570 let child1 = ArrayData::builder(DataType::Int32)
2571 .len(5)
2572 .add_buffer(Buffer::from_slice_ref([0, 1, 2, 3, 4]))
2573 .build()
2574 .unwrap();
2575 let child2 = child1.slice(1, 4);
2576
2577 let data1 = ArrayData::builder(data_type.clone())
2579 .len(4)
2580 .offset(1)
2581 .add_child_data(child1)
2582 .build()
2583 .unwrap();
2584 let data2 = ArrayData::builder(data_type)
2585 .len(4)
2586 .add_child_data(child2)
2587 .build()
2588 .unwrap();
2589
2590 assert_eq!(data1, data2);
2591 }
2592
2593 #[test]
2594 fn test_extend_struct_accounts_for_parent_offset() {
2595 let data_type =
2596 DataType::Struct(Fields::from(vec![Field::new("x", DataType::Int32, false)]));
2597 let child = ArrayData::builder(DataType::Int32)
2598 .len(5)
2599 .add_buffer(Buffer::from_slice_ref([0, 1, 2, 3, 4]))
2600 .build()
2601 .unwrap();
2602
2603 let data = ArrayData::builder(data_type)
2604 .len(4)
2605 .offset(1)
2606 .add_child_data(child)
2607 .build()
2608 .unwrap();
2609
2610 let mut mutable = MutableArrayData::new(vec![&data], false, data.len());
2611 mutable.try_extend(0, 0, data.len()).unwrap();
2612 let output = mutable.freeze();
2613
2614 assert_eq!(output.child_data()[0].buffer::<i32>(0), &[1, 2, 3, 4]);
2615 }
2616
2617 #[test]
2618 fn test_null_count() {
2619 let mut bit_v: [u8; 2] = [0; 2];
2620 bit_util::set_bit(&mut bit_v, 0);
2621 bit_util::set_bit(&mut bit_v, 3);
2622 bit_util::set_bit(&mut bit_v, 10);
2623 let arr_data = ArrayData::builder(DataType::Int32)
2624 .len(16)
2625 .add_buffer(make_i32_buffer(16))
2626 .null_bit_buffer(Some(Buffer::from(bit_v)))
2627 .build()
2628 .unwrap();
2629 assert_eq!(13, arr_data.null_count());
2630
2631 let mut bit_v: [u8; 2] = [0; 2];
2633 bit_util::set_bit(&mut bit_v, 0);
2634 bit_util::set_bit(&mut bit_v, 3);
2635 bit_util::set_bit(&mut bit_v, 10);
2636 let arr_data = ArrayData::builder(DataType::Int32)
2637 .len(12)
2638 .offset(2)
2639 .add_buffer(make_i32_buffer(14)) .null_bit_buffer(Some(Buffer::from(bit_v)))
2641 .build()
2642 .unwrap();
2643 assert_eq!(10, arr_data.null_count());
2644 }
2645
2646 #[test]
2647 fn test_null_buffer_ref() {
2648 let mut bit_v: [u8; 2] = [0; 2];
2649 bit_util::set_bit(&mut bit_v, 0);
2650 bit_util::set_bit(&mut bit_v, 3);
2651 bit_util::set_bit(&mut bit_v, 10);
2652 let arr_data = ArrayData::builder(DataType::Int32)
2653 .len(16)
2654 .add_buffer(make_i32_buffer(16))
2655 .null_bit_buffer(Some(Buffer::from(bit_v)))
2656 .build()
2657 .unwrap();
2658 assert!(arr_data.nulls().is_some());
2659 assert_eq!(&bit_v, arr_data.nulls().unwrap().validity());
2660 }
2661
2662 #[test]
2663 fn test_slice() {
2664 let mut bit_v: [u8; 2] = [0; 2];
2665 bit_util::set_bit(&mut bit_v, 0);
2666 bit_util::set_bit(&mut bit_v, 3);
2667 bit_util::set_bit(&mut bit_v, 10);
2668 let data = ArrayData::builder(DataType::Int32)
2669 .len(16)
2670 .add_buffer(make_i32_buffer(16))
2671 .null_bit_buffer(Some(Buffer::from(bit_v)))
2672 .build()
2673 .unwrap();
2674 let new_data = data.slice(1, 15);
2675 assert_eq!(data.len() - 1, new_data.len());
2676 assert_eq!(1, new_data.offset());
2677 assert_eq!(data.null_count(), new_data.null_count());
2678
2679 let new_data = new_data.slice(1, 14);
2681 assert_eq!(data.len() - 2, new_data.len());
2682 assert_eq!(2, new_data.offset());
2683 assert_eq!(data.null_count() - 1, new_data.null_count());
2684 }
2685
2686 #[test]
2687 #[should_panic(expected = "offset + length overflow")]
2688 fn test_slice_panics_on_offset_length_overflow() {
2689 let data = ArrayData::builder(DataType::Int32)
2690 .len(4)
2691 .add_buffer(make_i32_buffer(4))
2692 .build()
2693 .unwrap();
2694 let sliced = data.slice(1, 3);
2695
2696 sliced.slice(1, usize::MAX);
2697 }
2698
2699 #[test]
2700 fn test_typed_offsets_length_overflow() {
2701 let data = ArrayData {
2702 data_type: DataType::Binary,
2703 len: usize::MAX,
2704 offset: 0,
2705 buffers: vec![Buffer::from_slice_ref([0_i32])],
2706 child_data: vec![],
2707 nulls: None,
2708 };
2709 let err = data.typed_offsets::<i32>().unwrap_err();
2710
2711 assert_eq!(
2712 err.to_string(),
2713 format!(
2714 "Invalid argument error: Length {} with offset 1 overflows usize for Binary",
2715 usize::MAX
2716 )
2717 );
2718 }
2719
2720 #[test]
2721 fn test_validate_typed_buffer_length_overflow() {
2722 let data = ArrayData {
2723 data_type: DataType::Binary,
2724 len: 0,
2725 offset: 2,
2726 buffers: vec![Buffer::from_slice_ref([0_i32])],
2727 child_data: vec![],
2728 nulls: None,
2729 };
2730 let err = data.typed_buffer::<i32>(0, usize::MAX).unwrap_err();
2731
2732 assert_eq!(
2733 err.to_string(),
2734 format!(
2735 "Invalid argument error: Length {} with offset 2 overflows usize for Binary",
2736 usize::MAX
2737 )
2738 );
2739 }
2740
2741 fn try_new_binary_length_offset_overflow() -> Result<ArrayData, ArrowError> {
2743 ArrayData::try_new(
2744 DataType::Binary,
2745 usize::MAX,
2746 None,
2747 1,
2748 vec![
2749 Buffer::from_slice_ref([0_i32]),
2750 Buffer::from_iter(std::iter::empty::<u8>()),
2751 ],
2752 vec![],
2753 )
2754 }
2755
2756 #[cfg(not(feature = "force_validate"))]
2757 #[test]
2758 fn test_try_new_length_offset_overflow() {
2759 let err = try_new_binary_length_offset_overflow().unwrap_err();
2760
2761 assert_eq!(
2762 err.to_string(),
2763 format!(
2764 "Invalid argument error: Length {} with offset 1 overflows usize for Binary",
2765 usize::MAX
2766 )
2767 );
2768 }
2769
2770 #[cfg(feature = "force_validate")]
2771 #[test]
2772 #[should_panic(
2773 expected = "Length 18446744073709551615 with offset 1 overflows usize for Binary"
2774 )]
2775 fn test_try_new_length_offset_overflow_force_validate() {
2776 try_new_binary_length_offset_overflow().unwrap();
2777 }
2778
2779 #[test]
2780 fn test_equality() {
2781 let int_data = ArrayData::builder(DataType::Int32)
2782 .len(1)
2783 .add_buffer(make_i32_buffer(1))
2784 .build()
2785 .unwrap();
2786
2787 let float_data = ArrayData::builder(DataType::Float32)
2788 .len(1)
2789 .add_buffer(make_f32_buffer(1))
2790 .build()
2791 .unwrap();
2792 assert_ne!(int_data, float_data);
2793 assert!(!int_data.ptr_eq(&float_data));
2794 assert!(int_data.ptr_eq(&int_data));
2795
2796 let int_data_clone = int_data.clone();
2797 assert_eq!(int_data, int_data_clone);
2798 assert!(int_data.ptr_eq(&int_data_clone));
2799 assert!(int_data_clone.ptr_eq(&int_data));
2800
2801 let int_data_slice = int_data_clone.slice(1, 0);
2802 assert!(int_data_slice.ptr_eq(&int_data_slice));
2803 assert!(!int_data.ptr_eq(&int_data_slice));
2804 assert!(!int_data_slice.ptr_eq(&int_data));
2805
2806 let data_buffer = Buffer::from_slice_ref(b"abcdef");
2807 let offsets_buffer = Buffer::from_slice_ref([0_i32, 2_i32, 2_i32, 5_i32]);
2808 let string_data = ArrayData::try_new(
2809 DataType::Utf8,
2810 3,
2811 Some(Buffer::from_iter(vec![true, false, true])),
2812 0,
2813 vec![offsets_buffer, data_buffer],
2814 vec![],
2815 )
2816 .unwrap();
2817
2818 assert_ne!(float_data, string_data);
2819 assert!(!float_data.ptr_eq(&string_data));
2820
2821 assert!(string_data.ptr_eq(&string_data));
2822
2823 let string_data_cloned = string_data.clone();
2824 assert!(string_data_cloned.ptr_eq(&string_data));
2825 assert!(string_data.ptr_eq(&string_data_cloned));
2826
2827 let string_data_slice = string_data.slice(1, 2);
2828 assert!(string_data_slice.ptr_eq(&string_data_slice));
2829 assert!(!string_data_slice.ptr_eq(&string_data))
2830 }
2831
2832 #[test]
2833 fn test_slice_memory_size_view_payload_buffers() {
2834 for data_type in [DataType::Utf8View, DataType::BinaryView] {
2835 let inline_only = ArrayData::builder(data_type.clone())
2836 .len(2)
2837 .add_buffer(Buffer::from_vec(vec![0_u128; 2]))
2838 .build()
2839 .unwrap();
2840 assert_eq!(
2841 inline_only.get_slice_memory_size().unwrap(),
2842 2 * mem::size_of::<u128>()
2843 );
2844
2845 let mut first_payload = Vec::with_capacity(32);
2846 first_payload.extend_from_slice(b"first payload");
2847 let first_view =
2848 ByteView::new(first_payload.len().try_into().unwrap(), &first_payload[..4])
2849 .as_u128();
2850 let first_payload = Buffer::from_vec(first_payload);
2851 assert!(first_payload.capacity() > first_payload.len());
2852 let first_payload_capacity = first_payload.capacity();
2853
2854 let mut second_payload = Vec::with_capacity(64);
2855 second_payload.extend_from_slice(b"second payload");
2856 let second_view = ByteView::new(
2857 second_payload.len().try_into().unwrap(),
2858 &second_payload[..4],
2859 )
2860 .with_buffer_index(1)
2861 .as_u128();
2862 let second_payload = Buffer::from_vec(second_payload);
2863 assert!(second_payload.capacity() > second_payload.len());
2864 let second_payload_capacity = second_payload.capacity();
2865
2866 let data = ArrayData::builder(data_type)
2867 .len(3)
2868 .add_buffer(Buffer::from_vec(vec![first_view, 0_u128, second_view]))
2869 .add_buffer(first_payload)
2870 .add_buffer(second_payload)
2871 .build()
2872 .unwrap();
2873 let sliced = data.slice(1, 1);
2874
2875 assert_eq!(
2876 sliced.get_slice_memory_size().unwrap(),
2877 mem::size_of::<u128>() + first_payload_capacity + second_payload_capacity
2878 );
2879 }
2880 }
2881
2882 #[test]
2883 fn test_slice_memory_size_utf8_offset_buffer_len_plus_one() {
2884 let data_buffer = Buffer::from_slice_ref(b"helloworld");
2886 let offsets_buffer = Buffer::from_slice_ref([0_i32, 5_i32, 10_i32]);
2889 let array = ArrayData::try_new(
2890 DataType::Utf8,
2891 2,
2892 None,
2893 0,
2894 vec![offsets_buffer, data_buffer],
2895 vec![],
2896 )
2897 .unwrap();
2898 assert_eq!(array.get_slice_memory_size().unwrap(), 22); }
2900
2901 #[test]
2902 fn test_slice_memory_size_binary_offset_buffer_len_plus_one() {
2903 let data_buffer = Buffer::from_slice_ref([0u8, 1, 2, 3, 4]);
2906 let offsets_buffer = Buffer::from_slice_ref([0_i32, 2_i32, 5_i32]);
2908 let array = ArrayData::try_new(
2909 DataType::Binary,
2910 2,
2911 None,
2912 0,
2913 vec![offsets_buffer, data_buffer],
2914 vec![],
2915 )
2916 .unwrap();
2917 assert_eq!(array.get_slice_memory_size().unwrap(), 17); }
2919
2920 #[test]
2921 fn test_slice_memory_size() {
2922 let mut bit_v: [u8; 2] = [0; 2];
2923 bit_util::set_bit(&mut bit_v, 0);
2924 bit_util::set_bit(&mut bit_v, 3);
2925 bit_util::set_bit(&mut bit_v, 10);
2926 let data = ArrayData::builder(DataType::Int32)
2927 .len(16)
2928 .add_buffer(make_i32_buffer(16))
2929 .null_bit_buffer(Some(Buffer::from(bit_v)))
2930 .build()
2931 .unwrap();
2932 let new_data = data.slice(1, 14);
2933 assert_eq!(
2934 data.get_slice_memory_size().unwrap() - 8,
2935 new_data.get_slice_memory_size().unwrap()
2936 );
2937 let data_buffer = Buffer::from_slice_ref(b"abcdef");
2938 let offsets_buffer = Buffer::from_slice_ref([0_i32, 2_i32, 2_i32, 5_i32]);
2939 let string_data = ArrayData::try_new(
2940 DataType::Utf8,
2941 3,
2942 Some(Buffer::from_iter(vec![true, false, true])),
2943 0,
2944 vec![offsets_buffer, data_buffer],
2945 vec![],
2946 )
2947 .unwrap();
2948 let string_data_slice = string_data.slice(1, 2);
2949 assert_eq!(
2951 string_data.get_slice_memory_size().unwrap() - 6,
2952 string_data_slice.get_slice_memory_size().unwrap()
2953 );
2954 }
2955
2956 #[test]
2957 fn test_builder_rejects_short_null_bit_buffer() {
2958 for (len, offset) in [(8000, 0), (8, 1)] {
2959 let err = ArrayData::builder(DataType::Int32)
2960 .len(len)
2961 .offset(offset)
2962 .add_buffer(make_i32_buffer(len + offset))
2963 .null_bit_buffer(Some(Buffer::from([0_u8])))
2964 .build()
2965 .unwrap_err();
2966 assert_eq!(
2967 err.to_string(),
2968 format!(
2969 "Invalid argument error: null_bit_buffer size too small. got 1 needed {}",
2970 bit_util::ceil(len + offset, 8)
2971 )
2972 );
2973 }
2974 }
2975
2976 #[test]
2977 fn test_builder_null_bit_buffer_length_overflow() {
2978 let err = ArrayData::builder(DataType::Int32)
2979 .len(usize::MAX)
2980 .offset(1)
2981 .null_bit_buffer(Some(Buffer::default()))
2982 .build()
2983 .unwrap_err();
2984 assert_eq!(
2985 err.to_string(),
2986 format!(
2987 "Invalid argument error: Length {} with offset 1 overflows usize for Int32",
2988 usize::MAX
2989 )
2990 );
2991 }
2992
2993 #[test]
2994 fn test_builder_accepts_valid_null_bit_buffer() {
2995 for (len, offset) in [(8, 0), (7, 1)] {
2996 let data = ArrayData::builder(DataType::Int32)
2997 .len(len)
2998 .offset(offset)
2999 .add_buffer(make_i32_buffer(len + offset))
3000 .null_bit_buffer(Some(Buffer::from([0_u8])))
3001 .build()
3002 .unwrap();
3003 assert_eq!(data.len(), len);
3004 assert_eq!(data.offset(), offset);
3005 assert_eq!(data.null_count(), len);
3006 }
3007 }
3008
3009 #[test]
3010 fn test_count_nulls() {
3011 let buffer = Buffer::from([0b00010110, 0b10011111]);
3012 let buffer = NullBuffer::new(BooleanBuffer::new(buffer, 0, 16));
3013 let count = count_nulls(Some(&buffer), 0, 16);
3014 assert_eq!(count, 7);
3015
3016 let count = count_nulls(Some(&buffer), 4, 8);
3017 assert_eq!(count, 3);
3018 }
3019
3020 #[test]
3021 fn test_contains_nulls() {
3022 let buffer: Buffer =
3023 MutableBuffer::from_iter([false, false, false, true, true, false]).into();
3024 let buffer = NullBuffer::new(BooleanBuffer::new(buffer, 0, 6));
3025 assert!(contains_nulls(Some(&buffer), 0, 6));
3026 assert!(contains_nulls(Some(&buffer), 0, 3));
3027 assert!(!contains_nulls(Some(&buffer), 3, 2));
3028 assert!(!contains_nulls(Some(&buffer), 0, 0));
3029 }
3030
3031 #[test]
3032 fn test_alignment() {
3033 let buffer = Buffer::from_vec(vec![1_i32, 2_i32, 3_i32]);
3034 let sliced = buffer.slice(1);
3035
3036 let mut data = ArrayData {
3037 data_type: DataType::Int32,
3038 len: 0,
3039 offset: 0,
3040 buffers: vec![buffer],
3041 child_data: vec![],
3042 nulls: None,
3043 };
3044 data.validate_full().unwrap();
3045
3046 data.buffers[0] = sliced;
3048 let err = data.validate().unwrap_err();
3049
3050 assert_eq!(
3051 err.to_string(),
3052 "Invalid argument error: Misaligned buffers[0] in array of type Int32, offset from expected alignment of 4 by 1"
3053 );
3054
3055 data.align_buffers();
3056 data.validate_full().unwrap();
3057 }
3058
3059 #[test]
3060 fn test_alignment_struct() {
3061 let buffer = Buffer::from_vec(vec![1_i32, 2_i32, 3_i32]);
3062 let sliced = buffer.slice(1);
3063
3064 let child_data = ArrayData {
3065 data_type: DataType::Int32,
3066 len: 0,
3067 offset: 0,
3068 buffers: vec![buffer],
3069 child_data: vec![],
3070 nulls: None,
3071 };
3072
3073 let schema = DataType::Struct(Fields::from(vec![Field::new("a", DataType::Int32, false)]));
3074 let mut data = ArrayData {
3075 data_type: schema,
3076 len: 0,
3077 offset: 0,
3078 buffers: vec![],
3079 child_data: vec![child_data],
3080 nulls: None,
3081 };
3082 data.validate_full().unwrap();
3083
3084 data.child_data[0].buffers[0] = sliced;
3086 let err = data.validate().unwrap_err();
3087
3088 assert_eq!(
3089 err.to_string(),
3090 "Invalid argument error: Misaligned buffers[0] in array of type Int32, offset from expected alignment of 4 by 1"
3091 );
3092
3093 data.align_buffers();
3094 data.validate_full().unwrap();
3095 }
3096
3097 #[test]
3098 fn test_null_view_types() {
3099 let array_len = 32;
3100 let array = ArrayData::new_null(&DataType::BinaryView, array_len);
3101 assert_eq!(array.len(), array_len);
3102 for i in 0..array.len() {
3103 assert!(array.is_null(i));
3104 }
3105
3106 let array = ArrayData::new_null(&DataType::Utf8View, array_len);
3107 assert_eq!(array.len(), array_len);
3108 for i in 0..array.len() {
3109 assert!(array.is_null(i));
3110 }
3111
3112 let array = ArrayData::new_null(
3113 &DataType::ListView(Arc::new(Field::new_list_field(DataType::Int32, true))),
3114 array_len,
3115 );
3116 assert_eq!(array.len(), array_len);
3117 for i in 0..array.len() {
3118 assert!(array.is_null(i));
3119 }
3120
3121 let array = ArrayData::new_null(
3122 &DataType::LargeListView(Arc::new(Field::new_list_field(DataType::Int32, true))),
3123 array_len,
3124 );
3125 assert_eq!(array.len(), array_len);
3126 for i in 0..array.len() {
3127 assert!(array.is_null(i));
3128 }
3129 }
3130
3131 #[test]
3133 fn test_dont_panic_on_bad_input_when_using_try_new() {
3134 let empty_bytes = Buffer::default();
3135
3136 let array_data = ArrayData::try_new(
3137 DataType::Utf8,
3138 1, None,
3140 0,
3141 vec![Buffer::from_vec(vec![0i32, 2i32]), empty_bytes],
3143 vec![],
3144 );
3145
3146 let res = array_data.expect_err("should get error");
3147
3148 assert_eq!(
3149 res.to_string(),
3150 "Invalid argument error: Last offset 2 of Utf8 is larger than values length 0"
3151 );
3152 }
3153
3154 #[test]
3157 #[cfg(not(feature = "force_validate"))]
3158 fn test_validate_values_rejects_a_non_integer_dictionary_key() {
3159 let values = valid_non_nullable_int32_array_data(2);
3160 let data_type = DataType::Dictionary(Box::new(DataType::Utf8), Box::new(DataType::Int32));
3161 let dictionary = unsafe {
3162 ArrayData::builder(data_type)
3163 .len(1)
3164 .add_child_data(values)
3165 .build_unchecked()
3166 };
3167
3168 let err = dictionary.validate_values().expect_err("should get error");
3169 assert_eq!(
3170 err.to_string(),
3171 "Invalid argument error: Dictionary key type must be an integer, got Utf8"
3172 );
3173 }
3174
3175 #[test]
3176 #[cfg(not(feature = "force_validate"))]
3177 fn test_validate_values_rejects_a_non_integer_run_end() {
3178 let data_type = DataType::RunEndEncoded(
3179 Arc::new(Field::new(
3180 Field::REE_RUN_ENDS_FIELD_DEFAULT_NAME,
3181 DataType::Utf8,
3182 false,
3183 )),
3184 Arc::new(Field::new(
3185 Field::REE_VALUES_FIELD_DEFAULT_NAME,
3186 DataType::Int32,
3187 true,
3188 )),
3189 );
3190 let run_end_encoded = unsafe {
3191 ArrayData::builder(data_type)
3192 .len(1)
3193 .add_child_data(valid_non_nullable_int32_array_data(1))
3194 .add_child_data(valid_non_nullable_int32_array_data(1))
3195 .build_unchecked()
3196 };
3197
3198 let err = run_end_encoded
3199 .validate_values()
3200 .expect_err("should get error");
3201 assert_eq!(
3202 err.to_string(),
3203 "Invalid argument error: Run end type must be Int16, Int32 or Int64, got Utf8"
3204 );
3205 }
3206
3207 #[test]
3209 #[cfg(not(feature = "force_validate"))]
3210 fn test_validate_values_rejects_missing_child_data() {
3211 let int32 = Box::new(DataType::Int32);
3212 let field = || Arc::new(Field::new("f", DataType::Int32, true));
3213 let data_types = [
3214 DataType::Dictionary(int32.clone(), int32.clone()),
3215 DataType::List(field()),
3216 DataType::LargeList(field()),
3217 DataType::RunEndEncoded(field(), field()),
3218 ];
3219
3220 for data_type in data_types {
3221 let data = unsafe {
3222 ArrayData::builder(data_type.clone())
3223 .len(1)
3224 .build_unchecked()
3225 };
3226 let err = data.validate_values().expect_err("should get error");
3227 assert_eq!(
3228 err.to_string(),
3229 format!(
3230 "Invalid argument error: {data_type} should contain at least 1 child data array(s), had 0"
3231 )
3232 );
3233 }
3234 }
3235
3236 #[test]
3238 #[cfg(not(feature = "force_validate"))]
3239 fn test_validate_values_rejects_missing_buffers() {
3240 let cases = [
3242 (DataType::Utf8, 1),
3243 (DataType::LargeUtf8, 1),
3244 (DataType::Binary, 1),
3245 (DataType::LargeBinary, 1),
3246 (DataType::BinaryView, 0),
3247 (DataType::Utf8View, 0),
3248 ];
3249
3250 for (data_type, missing) in cases {
3251 let data = unsafe {
3252 ArrayData::builder(data_type.clone())
3253 .len(1)
3254 .build_unchecked()
3255 };
3256 let err = data.validate_values().expect_err("should get error");
3257 assert_eq!(
3258 err.to_string(),
3259 format!(
3260 "Invalid argument error: {data_type} should contain at least {} buffer(s), had 0",
3261 missing + 1
3262 )
3263 );
3264 }
3265 }
3266
3267 #[test]
3269 #[cfg(not(feature = "force_validate"))]
3270 fn test_validate_values_rejects_a_short_dictionary_keys_buffer() {
3271 let data_type = DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Int32));
3272 let dictionary = unsafe {
3273 ArrayData::builder(data_type)
3274 .len(4)
3275 .add_buffer(Buffer::from_slice_ref([1_i32, 0]))
3276 .add_child_data(valid_non_nullable_int32_array_data(2))
3277 .build_unchecked()
3278 };
3279
3280 let err = dictionary.validate_values().expect_err("should get error");
3281 assert_eq!(
3282 err.to_string(),
3283 "Invalid argument error: Buffer 0 of Dictionary(Int32, Int32) isn't large enough. Expected 16 bytes got 8"
3284 );
3285 }
3286
3287 #[test]
3288 fn should_fail_validation_when_having_map_field_type_is_not_struct() {
3289 let map_field = Field::new("key", DataType::Int32, false);
3290
3291 let map_field_data = valid_non_nullable_int32_array_data(2);
3292
3293 let results = test_both_builder_and_array_data(
3294 DataType::Map(map_field.into(), false),
3295 1,
3296 None,
3297 0,
3298 vec![
3299 OffsetBuffer::<i32>::from_lengths(vec![2])
3300 .into_inner()
3301 .into(),
3302 ],
3303 vec![map_field_data],
3304 );
3305
3306 for result in results {
3307 let array_data_err = result.expect_err("should fail for non struct field");
3308
3309 match array_data_err {
3310 ArrowError::InvalidArgumentError(msg) => {
3311 assert_eq!(
3312 msg,
3313 "Map field should be a entries struct data type, got Int32 instead"
3314 )
3315 }
3316 _ => panic!("unexpected error type {array_data_err}"),
3317 }
3318 }
3319 }
3320
3321 #[test]
3322 fn should_fail_validation_when_having_map_entries_only_have_1_field() {
3323 let struct_data_type = DataType::Struct(Fields::from(vec![Field::new(
3324 Field::MAP_KEY_FIELD_DEFAULT_NAME,
3325 DataType::Int32,
3326 false,
3327 )]));
3328
3329 let key_array_data = valid_non_nullable_int32_array_data(2);
3330
3331 let struct_data = {
3332 let builder = ArrayDataBuilder::new(struct_data_type.clone())
3333 .len(2)
3334 .nulls(None)
3335 .child_data(vec![key_array_data]);
3336
3337 builder.build().unwrap()
3338 };
3339
3340 let results = test_both_builder_and_array_data(
3341 DataType::Map(
3342 Field::new(
3343 Field::MAP_ENTRIES_FIELD_DEFAULT_NAME,
3344 struct_data_type,
3345 false,
3346 )
3347 .into(),
3348 false,
3349 ),
3350 1,
3351 None,
3352 0,
3353 vec![
3354 OffsetBuffer::<i32>::from_lengths(vec![2])
3355 .into_inner()
3356 .into(),
3357 ],
3358 vec![struct_data],
3359 );
3360
3361 for result in results {
3362 let array_data_err = result.expect_err("should fail for nullable key");
3363
3364 match array_data_err {
3365 ArrowError::InvalidArgumentError(msg) => {
3366 assert_eq!(
3367 msg,
3368 "Map entries data type should be a struct containing 2 fields, got 1 fields"
3369 )
3370 }
3371 _ => panic!("unexpected error type {array_data_err}"),
3372 }
3373 }
3374 }
3375
3376 #[test]
3377 fn should_fail_validation_when_having_map_entries_have_3_fields() {
3378 let struct_data_type = DataType::Struct(Fields::from(vec![
3379 Field::new(Field::MAP_KEY_FIELD_DEFAULT_NAME, DataType::Int32, false),
3380 Field::new(Field::MAP_VALUE_FIELD_DEFAULT_NAME, DataType::Utf8, true),
3381 Field::new("other", DataType::Int32, true),
3382 ]));
3383
3384 let key_array_data = valid_non_nullable_int32_array_data(2);
3385
3386 let values_array_data = valid_string_array_data(2);
3387
3388 let other_array_data = key_array_data.clone();
3389
3390 let struct_data = {
3391 let builder = ArrayDataBuilder::new(struct_data_type.clone())
3392 .len(2)
3393 .nulls(None)
3394 .child_data(vec![key_array_data, values_array_data, other_array_data]);
3395
3396 builder.build().unwrap()
3397 };
3398
3399 let results = test_both_builder_and_array_data(
3400 DataType::Map(
3401 Field::new(
3402 Field::MAP_ENTRIES_FIELD_DEFAULT_NAME,
3403 struct_data_type,
3404 false,
3405 )
3406 .into(),
3407 false,
3408 ),
3409 1,
3410 None,
3411 0,
3412 vec![
3413 OffsetBuffer::<i32>::from_lengths(vec![2])
3414 .into_inner()
3415 .into(),
3416 ],
3417 vec![struct_data],
3418 );
3419
3420 for result in results {
3421 let array_data_err = result.expect_err("should fail for nullable key");
3422
3423 match array_data_err {
3424 ArrowError::InvalidArgumentError(msg) => {
3425 assert_eq!(
3426 msg,
3427 "Map entries data type should be a struct containing 2 fields, got 3 fields"
3428 )
3429 }
3430 _ => panic!("unexpected error type {array_data_err}"),
3431 }
3432 }
3433 }
3434
3435 #[test]
3436 fn should_fail_validation_when_having_nullable_map_keys() {
3437 let struct_data_type = DataType::Struct(Fields::from(vec![
3438 Field::new(Field::MAP_KEY_FIELD_DEFAULT_NAME, DataType::Int32, true),
3439 Field::new(Field::MAP_VALUE_FIELD_DEFAULT_NAME, DataType::Utf8, true),
3440 ]));
3441
3442 let key_array_data = valid_non_nullable_int32_array_data(2);
3443 let values_array_data = valid_string_array_data(2);
3444
3445 let struct_data = {
3446 let builder = ArrayDataBuilder::new(struct_data_type.clone())
3447 .len(2)
3448 .nulls(None)
3449 .child_data(vec![key_array_data, values_array_data]);
3450
3451 builder.build().unwrap()
3452 };
3453
3454 let results = test_both_builder_and_array_data(
3455 DataType::Map(
3456 Field::new(
3457 Field::MAP_ENTRIES_FIELD_DEFAULT_NAME,
3458 struct_data_type,
3459 false,
3460 )
3461 .into(),
3462 false,
3463 ),
3464 1,
3465 None,
3466 0,
3467 vec![
3468 OffsetBuffer::<i32>::from_lengths(vec![2])
3469 .into_inner()
3470 .into(),
3471 ],
3472 vec![struct_data],
3473 );
3474
3475 for result in results {
3476 let array_data_err = result.expect_err("should fail for nullable key");
3477
3478 match array_data_err {
3479 ArrowError::InvalidArgumentError(msg) => {
3480 assert_eq!(msg, "Map key field must not be nullable")
3481 }
3482 _ => panic!("unexpected error type {array_data_err}"),
3483 }
3484 }
3485 }
3486
3487 #[test]
3488 fn should_fail_validation_when_having_entries_is_nullable_for_map() {
3489 let struct_data_type = DataType::Struct(Fields::from(vec![
3490 Field::new(Field::MAP_KEY_FIELD_DEFAULT_NAME, DataType::Int32, false),
3491 Field::new(Field::MAP_VALUE_FIELD_DEFAULT_NAME, DataType::Utf8, true),
3492 ]));
3493
3494 let key_array_data = valid_non_nullable_int32_array_data(2);
3495
3496 let values_array_data = valid_string_array_data(2);
3497
3498 let struct_data = {
3499 let builder = ArrayDataBuilder::new(struct_data_type.clone())
3500 .len(2)
3501 .nulls(None)
3502 .child_data(vec![key_array_data, values_array_data]);
3503
3504 builder.build().unwrap()
3505 };
3506
3507 let results = test_both_builder_and_array_data(
3508 DataType::Map(
3509 Field::new(
3510 Field::MAP_ENTRIES_FIELD_DEFAULT_NAME,
3511 struct_data_type,
3512 true,
3513 )
3514 .into(),
3515 false,
3516 ),
3517 1,
3518 None,
3519 0,
3520 vec![
3521 OffsetBuffer::<i32>::from_lengths(vec![2])
3522 .into_inner()
3523 .into(),
3524 ],
3525 vec![struct_data],
3526 );
3527
3528 for result in results {
3529 let array_data_err = result.expect_err("should fail for nullable entries");
3530
3531 match array_data_err {
3532 ArrowError::InvalidArgumentError(msg) => assert_eq!(
3533 msg,
3534 "The nullable should be set to false for the map entries field."
3535 ),
3536 _ => panic!("unexpected error type {array_data_err}"),
3537 }
3538 }
3539 }
3540
3541 #[test]
3542 fn should_allow_to_create_map_from_data() {
3543 let struct_data_type = DataType::Struct(Fields::from(vec![
3544 Field::new(Field::MAP_KEY_FIELD_DEFAULT_NAME, DataType::Int32, false),
3545 Field::new(Field::MAP_VALUE_FIELD_DEFAULT_NAME, DataType::Utf8, true),
3546 ]));
3547
3548 let key_array_data = valid_non_nullable_int32_array_data(2);
3549 let values_array_data = valid_string_array_data(2);
3550
3551 let struct_data = {
3552 let builder = ArrayDataBuilder::new(struct_data_type.clone())
3553 .len(2)
3554 .nulls(None)
3555 .child_data(vec![key_array_data, values_array_data]);
3556
3557 builder.build().unwrap()
3558 };
3559
3560 let results = test_both_builder_and_array_data(
3561 DataType::Map(
3562 Field::new(
3563 Field::MAP_ENTRIES_FIELD_DEFAULT_NAME,
3564 struct_data_type,
3565 false,
3566 )
3567 .into(),
3568 false,
3569 ),
3570 1,
3571 None,
3572 0,
3573 vec![
3574 OffsetBuffer::<i32>::from_lengths(vec![2])
3575 .into_inner()
3576 .into(),
3577 ],
3578 vec![struct_data],
3579 );
3580
3581 for result in results {
3582 result.expect("should be able to create map ArrayData");
3583 }
3584 }
3585
3586 fn valid_string_array_data(length: usize) -> ArrayData {
3587 let offsets = OffsetBuffer::<i32>::from_lengths(vec![0; length])
3588 .into_inner()
3589 .into_inner();
3590 let empty_bytes = Buffer::default();
3591
3592 let builder = ArrayDataBuilder::new(DataType::Utf8)
3593 .len(length)
3594 .buffers(vec![offsets, empty_bytes])
3595 .nulls(None);
3596
3597 builder.build().unwrap()
3598 }
3599
3600 fn valid_non_nullable_int32_array_data(length: usize) -> ArrayData {
3601 let builder = ArrayDataBuilder::new(DataType::Int32)
3602 .len(length)
3603 .nulls(None)
3604 .buffers(vec![
3605 ScalarBuffer::<i32>::from(vec![1; length]).into_inner(),
3606 ]);
3607
3608 builder.build().unwrap()
3609 }
3610
3611 #[test]
3612 fn empty_and_null_map_array_should_pass_validation() {
3613 let dt = DataType::Map(
3614 Field::new(
3615 Field::MAP_ENTRIES_FIELD_DEFAULT_NAME,
3616 DataType::Struct(Fields::from(vec![
3617 Field::new(Field::MAP_KEY_FIELD_DEFAULT_NAME, DataType::Int32, false),
3618 Field::new(Field::MAP_VALUE_FIELD_DEFAULT_NAME, DataType::Utf8, true),
3619 ])),
3620 false,
3621 )
3622 .into(),
3623 false,
3624 );
3625
3626 ArrayData::new_empty(&dt).validate_full().unwrap();
3627 ArrayData::new_null(&dt, 1).validate_full().unwrap();
3628 }
3629
3630 #[test]
3631 fn null_buffer_offset_is_independent_of_data_offset() {
3632 let int_data = ArrayData::builder(DataType::UInt32)
3634 .offset(50)
3635 .len(50)
3636 .add_buffer(Buffer::from_vec(vec![0_u32; 100]))
3637 .build()
3638 .unwrap();
3639 int_data.validate().unwrap();
3640
3641 let nulls = NullBuffer::new(BooleanBuffer::from(vec![false; 100]).slice(0, 50));
3643 let with_sliced_nulls = int_data
3644 .clone()
3645 .into_builder()
3646 .nulls(Some(nulls))
3647 .build()
3648 .unwrap();
3649 with_sliced_nulls.validate().unwrap();
3650
3651 let nulls = NullBuffer::new(BooleanBuffer::from(vec![false; 50]));
3654 let with_unsliced_nulls = int_data.into_builder().nulls(Some(nulls)).build().unwrap();
3655 with_unsliced_nulls.validate().unwrap();
3656 assert_eq!(with_unsliced_nulls.null_count(), 50);
3657 }
3658
3659 fn test_both_builder_and_array_data(
3660 data_type: DataType,
3661 len: usize,
3662 null_bit_buffer: Option<Buffer>,
3663 offset: usize,
3664 buffers: Vec<Buffer>,
3665 child_data: Vec<ArrayData>,
3666 ) -> [Result<ArrayData, ArrowError>; 2] {
3667 let from_builder_res = ArrayData::builder(data_type.clone())
3668 .len(len)
3669 .add_buffers(buffers.clone())
3670 .null_bit_buffer(null_bit_buffer.clone())
3671 .offset(offset)
3672 .child_data(child_data.clone())
3673 .build();
3674
3675 let from_try_new_res =
3676 ArrayData::try_new(data_type, len, null_bit_buffer, offset, buffers, child_data);
3677
3678 [from_builder_res, from_try_new_res]
3679 }
3680}