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) => match f.iter().next() {
790 None => {
794 assert_eq!(
795 len, 0,
796 "cannot construct null data from an empty union of length {len}, a slot has no type id to carry"
797 );
798 let buffers = match mode {
799 UnionMode::Sparse => vec![zeroed(0)],
800 UnionMode::Dense => vec![zeroed(0), zeroed(0)],
801 };
802 (buffers, vec![], false)
803 }
804 Some((id, _)) => {
805 let ids = Buffer::from_iter(std::iter::repeat_n(id, len));
806 let buffers = match mode {
807 UnionMode::Sparse => vec![ids],
808 UnionMode::Dense => {
809 let end_offset = i32::from_usize(len).unwrap();
810 vec![ids, Buffer::from_iter(0_i32..end_offset)]
811 }
812 };
813
814 let children = f
815 .iter()
816 .enumerate()
817 .map(|(idx, (_, f))| {
818 if idx == 0 || *mode == UnionMode::Sparse {
819 Self::new_null(f.data_type(), len)
820 } else {
821 Self::new_empty(f.data_type())
822 }
823 })
824 .collect();
825
826 (buffers, children, false)
827 }
828 },
829 DataType::RunEndEncoded(r, v) => {
830 if len == 0 {
831 let runs = ArrayData::new_empty(r.data_type());
833 let values = ArrayData::new_empty(v.data_type());
834 (vec![], vec![runs, values], false)
835 } else {
836 let runs = match r.data_type() {
837 DataType::Int16 => {
838 let i = i16::from_usize(len).expect("run overflow");
839 Buffer::from_slice_ref([i])
840 }
841 DataType::Int32 => {
842 let i = i32::from_usize(len).expect("run overflow");
843 Buffer::from_slice_ref([i])
844 }
845 DataType::Int64 => {
846 let i = i64::from_usize(len).expect("run overflow");
847 Buffer::from_slice_ref([i])
848 }
849 dt => unreachable!("Invalid run ends data type {dt}"),
850 };
851
852 let builder = ArrayData::builder(r.data_type().clone())
853 .len(1)
854 .buffers(vec![runs]);
855
856 let runs = unsafe { builder.build_unchecked() };
859 (
860 vec![],
861 vec![runs, ArrayData::new_null(v.data_type(), 1)],
862 false,
863 )
864 }
865 }
866 DataType::Int8
868 | DataType::Int16
869 | DataType::Int32
870 | DataType::Int64
871 | DataType::UInt8
872 | DataType::UInt16
873 | DataType::UInt32
874 | DataType::UInt64
875 | DataType::Float16
876 | DataType::Float32
877 | DataType::Float64
878 | DataType::Timestamp(_, _)
879 | DataType::Date32
880 | DataType::Date64
881 | DataType::Time32(_)
882 | DataType::Time64(_)
883 | DataType::Duration(_)
884 | DataType::Interval(_)
885 | DataType::Decimal32(_, _)
886 | DataType::Decimal64(_, _)
887 | DataType::Decimal128(_, _)
888 | DataType::Decimal256(_, _) => unreachable!("{data_type}"),
889 },
890 };
891
892 let mut builder = ArrayDataBuilder::new(data_type.clone())
893 .len(len)
894 .buffers(buffers)
895 .child_data(child_data);
896
897 if has_nulls {
898 builder = builder.nulls(Some(NullBuffer::new_null(len)))
899 }
900
901 unsafe { builder.build_unchecked() }
904 }
905
906 pub fn new_empty(data_type: &DataType) -> Self {
908 Self::new_null(data_type, 0)
909 }
910
911 pub fn align_buffers(&mut self) {
920 let layout = layout(&self.data_type);
921 for (buffer, spec) in self.buffers.iter_mut().zip(&layout.buffers) {
922 if let BufferSpec::FixedWidth { alignment, .. } = spec
923 && buffer.as_ptr().align_offset(*alignment) != 0
924 {
925 *buffer = Buffer::from_slice_ref(buffer.as_ref());
926 }
927 }
928 for data in &mut self.child_data {
930 data.align_buffers()
931 }
932 }
933
934 pub fn validate(&self) -> Result<(), ArrowError> {
945 let len_plus_offset = checked_len_plus_offset(&self.data_type, self.len, self.offset)?;
947
948 let layout = layout(&self.data_type);
950
951 if !layout.can_contain_null_mask && self.nulls.is_some() {
952 return Err(ArrowError::InvalidArgumentError(format!(
953 "Arrays of type {:?} cannot contain a null bitmask",
954 self.data_type,
955 )));
956 }
957
958 if self.buffers.len() < layout.buffers.len()
960 || (!layout.variadic && self.buffers.len() != layout.buffers.len())
961 {
962 return Err(ArrowError::InvalidArgumentError(format!(
963 "Expected {} buffers in array of type {:?}, got {}",
964 layout.buffers.len(),
965 self.data_type,
966 self.buffers.len(),
967 )));
968 }
969
970 for (i, (buffer, spec)) in self.buffers.iter().zip(layout.buffers.iter()).enumerate() {
971 match spec {
972 BufferSpec::FixedWidth {
973 byte_width,
974 alignment,
975 } => {
976 let min_buffer_size = len_plus_offset.saturating_mul(*byte_width);
977
978 if buffer.len() < min_buffer_size {
979 return Err(ArrowError::InvalidArgumentError(format!(
980 "Need at least {} bytes in buffers[{}] in array of type {:?}, but got {}",
981 min_buffer_size,
982 i,
983 self.data_type,
984 buffer.len()
985 )));
986 }
987
988 let align_offset = buffer.as_ptr().align_offset(*alignment);
989 if align_offset != 0 {
990 return Err(ArrowError::InvalidArgumentError(format!(
991 "Misaligned buffers[{i}] in array of type {:?}, offset from expected alignment of {alignment} by {}",
992 self.data_type,
993 align_offset.min(alignment - align_offset)
994 )));
995 }
996 }
997 BufferSpec::VariableWidth => {
998 }
1002 BufferSpec::BitMap => {
1003 let min_buffer_size = bit_util::ceil(len_plus_offset, 8);
1004 if buffer.len() < min_buffer_size {
1005 return Err(ArrowError::InvalidArgumentError(format!(
1006 "Need at least {} bytes for bitmap in buffers[{}] in array of type {:?}, but got {}",
1007 min_buffer_size,
1008 i,
1009 self.data_type,
1010 buffer.len()
1011 )));
1012 }
1013 }
1014 BufferSpec::AlwaysNull => {
1015 }
1017 }
1018 }
1019
1020 if let Some(nulls) = self.nulls() {
1022 if nulls.null_count() > self.len {
1023 return Err(ArrowError::InvalidArgumentError(format!(
1024 "null_count {} for an array exceeds length of {} elements",
1025 nulls.null_count(),
1026 self.len
1027 )));
1028 }
1029
1030 if nulls.len() != self.len {
1031 return Err(ArrowError::InvalidArgumentError(format!(
1032 "null buffer incorrect size. got {} expected {}",
1033 nulls.len(),
1034 self.len
1035 )));
1036 }
1037 }
1038
1039 self.validate_child_data()?;
1040
1041 match &self.data_type {
1043 DataType::Utf8 | DataType::Binary => {
1044 self.validate_offsets::<i32>(self.buffers[1].len())?;
1045 }
1046 DataType::LargeUtf8 | DataType::LargeBinary => {
1047 self.validate_offsets::<i64>(self.buffers[1].len())?;
1048 }
1049 DataType::Dictionary(key_type, _value_type) => {
1050 if !DataType::is_dictionary_key_type(key_type) {
1052 return Err(ArrowError::InvalidArgumentError(format!(
1053 "Dictionary key type must be integer, but was {key_type}"
1054 )));
1055 }
1056 }
1057 DataType::RunEndEncoded(run_ends_type, _) => {
1058 if run_ends_type.is_nullable() {
1059 return Err(ArrowError::InvalidArgumentError(
1060 "The nullable should be set to false for the field defining run_ends array.".to_string()
1061 ));
1062 }
1063 if !DataType::is_run_ends_type(run_ends_type.data_type()) {
1064 return Err(ArrowError::InvalidArgumentError(format!(
1065 "RunArray run_ends types must be Int16, Int32 or Int64, but was {}",
1066 run_ends_type.data_type()
1067 )));
1068 }
1069 }
1070 DataType::Map(f, _) if f.is_nullable() => {
1071 return Err(ArrowError::InvalidArgumentError(
1072 "The nullable should be set to false for the map entries field.".to_string(),
1073 ));
1074 }
1075 _ => {}
1076 }
1077
1078 Ok(())
1079 }
1080
1081 fn typed_offsets<T: ArrowNativeType + num_traits::Num>(&self) -> Result<&[T], ArrowError> {
1088 if self.len == 0 && self.buffer_at(0)?.is_empty() {
1090 return Ok(&[]);
1091 }
1092
1093 let len = checked_len_plus_offset(&self.data_type, self.len, 1)?;
1094
1095 self.typed_buffer(0, len)
1096 }
1097
1098 fn typed_buffer<T: ArrowNativeType + num_traits::Num>(
1100 &self,
1101 idx: usize,
1102 len: usize,
1103 ) -> Result<&[T], ArrowError> {
1104 let buffer = self.buffer_at(idx)?;
1105
1106 let required_elements = checked_len_plus_offset(&self.data_type, len, self.offset)?;
1107 let byte_width = mem::size_of::<T>();
1108 let required_len = required_elements.checked_mul(byte_width).ok_or_else(|| {
1109 ArrowError::InvalidArgumentError(format!(
1110 "Buffer {idx} of {} byte length overflow: {} elements of {} bytes exceeds usize",
1111 self.data_type, required_elements, byte_width
1112 ))
1113 })?;
1114
1115 if buffer.len() < required_len {
1116 return Err(ArrowError::InvalidArgumentError(format!(
1117 "Buffer {} of {} isn't large enough. Expected {} bytes got {}",
1118 idx,
1119 self.data_type,
1120 required_len,
1121 buffer.len()
1122 )));
1123 }
1124
1125 Ok(&buffer.typed_data::<T>()[self.offset..required_elements])
1126 }
1127
1128 fn validate_offsets<T: ArrowNativeType + num_traits::Num + std::fmt::Display>(
1131 &self,
1132 values_length: usize,
1133 ) -> Result<(), ArrowError> {
1134 let offsets = self.typed_offsets::<T>()?;
1136 if offsets.is_empty() {
1137 return Ok(());
1138 }
1139
1140 let first_offset = offsets[0].to_usize().ok_or_else(|| {
1141 ArrowError::InvalidArgumentError(format!(
1142 "Error converting offset[0] ({}) to usize for {}",
1143 offsets[0], self.data_type
1144 ))
1145 })?;
1146
1147 let last_offset = offsets[self.len].to_usize().ok_or_else(|| {
1148 ArrowError::InvalidArgumentError(format!(
1149 "Error converting offset[{}] ({}) to usize for {}",
1150 self.len, offsets[self.len], self.data_type
1151 ))
1152 })?;
1153
1154 if first_offset > values_length {
1155 return Err(ArrowError::InvalidArgumentError(format!(
1156 "First offset {} of {} is larger than values length {}",
1157 first_offset, self.data_type, values_length,
1158 )));
1159 }
1160
1161 if last_offset > values_length {
1162 return Err(ArrowError::InvalidArgumentError(format!(
1163 "Last offset {} of {} is larger than values length {}",
1164 last_offset, self.data_type, values_length,
1165 )));
1166 }
1167
1168 if first_offset > last_offset {
1169 return Err(ArrowError::InvalidArgumentError(format!(
1170 "First offset {} in {} is smaller than last offset {}",
1171 first_offset, self.data_type, last_offset,
1172 )));
1173 }
1174
1175 Ok(())
1176 }
1177
1178 fn validate_offsets_and_sizes<T: ArrowNativeType + num_traits::Num + std::fmt::Display>(
1181 &self,
1182 values_length: usize,
1183 ) -> Result<(), ArrowError> {
1184 let offsets: &[T] = self.typed_buffer(0, self.len)?;
1185 let sizes: &[T] = self.typed_buffer(1, self.len)?;
1186 if offsets.len() != sizes.len() {
1187 return Err(ArrowError::ComputeError(format!(
1188 "ListView offsets len {} does not match sizes len {}",
1189 offsets.len(),
1190 sizes.len()
1191 )));
1192 }
1193
1194 for i in 0..sizes.len() {
1195 let size = sizes[i].to_usize().ok_or_else(|| {
1196 ArrowError::InvalidArgumentError(format!(
1197 "Error converting size[{}] ({}) to usize for {}",
1198 i, sizes[i], self.data_type
1199 ))
1200 })?;
1201 let offset = offsets[i].to_usize().ok_or_else(|| {
1202 ArrowError::InvalidArgumentError(format!(
1203 "Error converting offset[{}] ({}) to usize for {}",
1204 i, offsets[i], self.data_type
1205 ))
1206 })?;
1207 if size
1208 .checked_add(offset)
1209 .expect("Offset and size have exceeded the usize boundary")
1210 > values_length
1211 {
1212 return Err(ArrowError::InvalidArgumentError(format!(
1213 "Size {} at index {} is larger than the remaining values for {}",
1214 size, i, self.data_type
1215 )));
1216 }
1217 }
1218 Ok(())
1219 }
1220
1221 fn validate_child_data(&self) -> Result<(), ArrowError> {
1223 match &self.data_type {
1224 DataType::List(field) => {
1225 let values_data = self.get_single_valid_child_data(field.data_type())?;
1226 self.validate_offsets::<i32>(values_data.len)?;
1227 }
1228 DataType::LargeList(field) => {
1229 let values_data = self.get_single_valid_child_data(field.data_type())?;
1230 self.validate_offsets::<i64>(values_data.len)?;
1231 }
1232 DataType::Map(field, _) => {
1233 let DataType::Struct(entries_fields) = field.data_type() else {
1234 return Err(ArrowError::InvalidArgumentError(format!(
1235 "Map field should be a entries struct data type, got {:?} instead",
1236 field.data_type()
1237 )));
1238 };
1239 if entries_fields.len() != 2 {
1240 return Err(ArrowError::InvalidArgumentError(format!(
1241 "Map entries data type should be a struct containing 2 fields, got {} fields",
1242 entries_fields.len()
1243 )));
1244 }
1245
1246 if entries_fields[0].is_nullable() {
1248 return Err(ArrowError::InvalidArgumentError(
1249 "Map key field must not be nullable".to_string(),
1250 ));
1251 }
1252 let values_data = self.get_single_valid_child_data(field.data_type())?;
1253 self.validate_offsets::<i32>(values_data.len)?;
1254 }
1255 DataType::ListView(field) => {
1256 let values_data = self.get_single_valid_child_data(field.data_type())?;
1257 self.validate_offsets_and_sizes::<i32>(values_data.len)?;
1258 }
1259 DataType::LargeListView(field) => {
1260 let values_data = self.get_single_valid_child_data(field.data_type())?;
1261 self.validate_offsets_and_sizes::<i64>(values_data.len)?;
1262 }
1263 DataType::FixedSizeList(field, list_size) => {
1264 let values_data = self.get_single_valid_child_data(field.data_type())?;
1265
1266 let list_size: usize = (*list_size).try_into().map_err(|_| {
1267 ArrowError::InvalidArgumentError(format!(
1268 "{} has a negative list_size {}",
1269 self.data_type, list_size
1270 ))
1271 })?;
1272
1273 let expected_values_len = self.len
1274 .checked_mul(list_size)
1275 .expect("integer overflow computing expected number of expected values in FixedListSize");
1276
1277 if values_data.len < expected_values_len {
1278 return Err(ArrowError::InvalidArgumentError(format!(
1279 "Values length {} is less than the length ({}) multiplied by the value size ({}) for {}",
1280 values_data.len, self.len, list_size, self.data_type
1281 )));
1282 }
1283 }
1284 DataType::Struct(fields) => {
1285 self.validate_num_child_data(fields.len())?;
1286 let len_plus_offset =
1287 checked_len_plus_offset(&self.data_type, self.len, self.offset)?;
1288 for (i, field) in fields.iter().enumerate() {
1289 let field_data = self.get_valid_child_data(i, field.data_type())?;
1290
1291 if field_data.len < len_plus_offset {
1293 return Err(ArrowError::InvalidArgumentError(format!(
1294 "{} child array #{} for field {} has length smaller than expected for struct array ({} < {})",
1295 self.data_type,
1296 i,
1297 field.name(),
1298 field_data.len,
1299 len_plus_offset
1300 )));
1301 }
1302 }
1303 }
1304 DataType::RunEndEncoded(run_ends_field, values_field) => {
1305 self.validate_num_child_data(2)?;
1306 let run_ends_data = self.get_valid_child_data(0, run_ends_field.data_type())?;
1307 let values_data = self.get_valid_child_data(1, values_field.data_type())?;
1308 if run_ends_data.len != values_data.len {
1309 return Err(ArrowError::InvalidArgumentError(format!(
1310 "The run_ends array length should be the same as values array length. Run_ends array length is {}, values array length is {}",
1311 run_ends_data.len, values_data.len
1312 )));
1313 }
1314 if run_ends_data.nulls.is_some() {
1315 return Err(ArrowError::InvalidArgumentError(
1316 "Found null values in run_ends array. The run_ends array should not have null values.".to_string(),
1317 ));
1318 }
1319 }
1320 DataType::Union(fields, mode) => {
1321 self.validate_num_child_data(fields.len())?;
1322
1323 for (i, (_, field)) in fields.iter().enumerate() {
1324 let field_data = self.get_valid_child_data(i, field.data_type())?;
1325
1326 if mode == &UnionMode::Sparse {
1327 let len_plus_offset =
1328 checked_len_plus_offset(&self.data_type, self.len, self.offset)?;
1329 if field_data.len < len_plus_offset {
1330 return Err(ArrowError::InvalidArgumentError(format!(
1331 "Sparse union child array #{} has length smaller than expected for union array ({} < {})",
1332 i, field_data.len, len_plus_offset
1333 )));
1334 }
1335 }
1336 }
1337 }
1338 DataType::Dictionary(_key_type, value_type) => {
1339 self.get_single_valid_child_data(value_type)?;
1340 }
1341 _ => {
1342 if !self.child_data.is_empty() {
1344 return Err(ArrowError::InvalidArgumentError(format!(
1345 "Expected no child arrays for type {} but got {}",
1346 self.data_type,
1347 self.child_data.len()
1348 )));
1349 }
1350 }
1351 }
1352 Ok(())
1353 }
1354
1355 fn get_single_valid_child_data(
1359 &self,
1360 expected_type: &DataType,
1361 ) -> Result<&ArrayData, ArrowError> {
1362 self.validate_num_child_data(1)?;
1363 self.get_valid_child_data(0, expected_type)
1364 }
1365
1366 fn buffer_at(&self, idx: usize) -> Result<&Buffer, ArrowError> {
1371 self.buffers.get(idx).ok_or_else(|| {
1372 ArrowError::InvalidArgumentError(format!(
1373 "{} should contain at least {} buffer(s), had {}",
1374 self.data_type,
1375 idx + 1,
1376 self.buffers.len()
1377 ))
1378 })
1379 }
1380
1381 fn child_at(&self, idx: usize) -> Result<&ArrayData, ArrowError> {
1386 self.child_data.get(idx).ok_or_else(|| {
1387 ArrowError::InvalidArgumentError(format!(
1388 "{} should contain at least {} child data array(s), had {}",
1389 self.data_type,
1390 idx + 1,
1391 self.child_data.len()
1392 ))
1393 })
1394 }
1395
1396 fn validate_num_child_data(&self, expected_len: usize) -> Result<(), ArrowError> {
1398 if self.child_data.len() != expected_len {
1399 Err(ArrowError::InvalidArgumentError(format!(
1400 "Value data for {} should contain {} child data array(s), had {}",
1401 self.data_type,
1402 expected_len,
1403 self.child_data.len()
1404 )))
1405 } else {
1406 Ok(())
1407 }
1408 }
1409
1410 fn get_valid_child_data(
1413 &self,
1414 i: usize,
1415 expected_type: &DataType,
1416 ) -> Result<&ArrayData, ArrowError> {
1417 let values_data = self.child_data.get(i).ok_or_else(|| {
1418 ArrowError::InvalidArgumentError(format!(
1419 "{} did not have enough child arrays. Expected at least {} but had only {}",
1420 self.data_type,
1421 i + 1,
1422 self.child_data.len()
1423 ))
1424 })?;
1425
1426 if expected_type != &values_data.data_type {
1427 return Err(ArrowError::InvalidArgumentError(format!(
1428 "Child type mismatch for {}. Expected {} but child data had {}",
1429 self.data_type, expected_type, values_data.data_type
1430 )));
1431 }
1432
1433 values_data.validate()?;
1434 Ok(values_data)
1435 }
1436
1437 pub fn validate_data(&self) -> Result<(), ArrowError> {
1453 self.validate()?;
1454
1455 self.validate_nulls()?;
1456 self.validate_values()?;
1457 Ok(())
1458 }
1459
1460 pub fn validate_full(&self) -> Result<(), ArrowError> {
1465 self.validate_data()?;
1466 self.child_data
1468 .iter()
1469 .enumerate()
1470 .try_for_each(|(i, child_data)| {
1471 child_data.validate_full().map_err(|e| {
1472 ArrowError::InvalidArgumentError(format!(
1473 "{} child #{} invalid: {}",
1474 self.data_type, i, e
1475 ))
1476 })
1477 })?;
1478 Ok(())
1479 }
1480
1481 pub fn validate_nulls(&self) -> Result<(), ArrowError> {
1491 if let Some(nulls) = &self.nulls {
1492 let actual = nulls.len() - nulls.inner().count_set_bits();
1493 if actual != nulls.null_count() {
1494 return Err(ArrowError::InvalidArgumentError(format!(
1495 "null_count value ({}) doesn't match actual number of nulls in array ({})",
1496 nulls.null_count(),
1497 actual
1498 )));
1499 }
1500 }
1501
1502 match &self.data_type {
1507 DataType::List(f) | DataType::LargeList(f) | DataType::Map(f, _) => {
1508 if !f.is_nullable() {
1509 let child = &self.child_data[0];
1510 self.validate_non_nullable(None, child, child.nulls())?
1511 }
1512 }
1513 DataType::FixedSizeList(field, len) => {
1514 let child = &self.child_data[0];
1515 if !field.is_nullable() {
1516 match &self.nulls {
1517 Some(nulls) => {
1518 let element_len = *len as usize;
1519 let expanded = nulls.expand(element_len);
1520 self.validate_non_nullable(Some(&expanded), child, child.nulls())?;
1521 }
1522 None => self.validate_non_nullable(None, child, child.nulls())?,
1523 }
1524 }
1525 }
1526 DataType::Struct(fields) => {
1527 for (field, child) in fields.iter().zip(&self.child_data) {
1528 if !field.is_nullable() {
1529 let child_nulls = child
1530 .nulls()
1531 .map(|nulls| nulls.slice(self.offset, self.len));
1532 self.validate_non_nullable(self.nulls(), child, child_nulls.as_ref())?
1533 }
1534 }
1535 }
1536 _ => {}
1537 }
1538
1539 Ok(())
1540 }
1541
1542 fn validate_non_nullable(
1544 &self,
1545 mask: Option<&NullBuffer>,
1546 child: &ArrayData,
1547 child_nulls: Option<&NullBuffer>,
1548 ) -> Result<(), ArrowError> {
1549 let Some(mask) = mask else {
1550 return match child_nulls.map(NullBuffer::null_count).unwrap_or_default() {
1551 0 => Ok(()),
1552 _ => Err(ArrowError::InvalidArgumentError(format!(
1553 "non-nullable child of type {} contains nulls not present in parent {}",
1554 child.data_type, self.data_type
1555 ))),
1556 };
1557 };
1558
1559 match child_nulls {
1560 Some(nulls) if !mask.contains(nulls) => Err(ArrowError::InvalidArgumentError(format!(
1561 "non-nullable child of type {} contains nulls not present in parent",
1562 child.data_type
1563 ))),
1564 _ => Ok(()),
1565 }
1566 }
1567
1568 pub fn validate_values(&self) -> Result<(), ArrowError> {
1574 match &self.data_type {
1575 DataType::Utf8 => self.validate_utf8::<i32>(),
1576 DataType::LargeUtf8 => self.validate_utf8::<i64>(),
1577 DataType::Binary => self.validate_offsets_full::<i32>(self.buffer_at(1)?.len()),
1578 DataType::LargeBinary => self.validate_offsets_full::<i64>(self.buffer_at(1)?.len()),
1579 DataType::BinaryView => {
1580 let views = self.typed_buffer::<u128>(0, self.len)?;
1581 validate_binary_view(views, &self.buffers[1..])
1582 }
1583 DataType::Utf8View => {
1584 let views = self.typed_buffer::<u128>(0, self.len)?;
1585 validate_string_view(views, &self.buffers[1..])
1586 }
1587 DataType::List(_) | DataType::Map(_, _) => {
1588 let child = self.child_at(0)?;
1589 self.validate_offsets_full::<i32>(child.len)
1590 }
1591 DataType::LargeList(_) => {
1592 let child = self.child_at(0)?;
1593 self.validate_offsets_full::<i64>(child.len)
1594 }
1595 DataType::Union(_, _) => {
1596 Ok(())
1602 }
1603 DataType::Dictionary(key_type, _value_type) => {
1604 let dictionary_length = self.child_at(0)?.len;
1605 let dictionary_length = i64::try_from(dictionary_length).map_err(|_| {
1606 ArrowError::InvalidArgumentError(format!(
1607 "Dictionary of {dictionary_length} values is too long for an i64"
1608 ))
1609 })?;
1610 let max_value = dictionary_length - 1;
1611 match key_type.as_ref() {
1612 DataType::UInt8 => self.check_bounds::<u8>(max_value),
1613 DataType::UInt16 => self.check_bounds::<u16>(max_value),
1614 DataType::UInt32 => self.check_bounds::<u32>(max_value),
1615 DataType::UInt64 => self.check_bounds::<u64>(max_value),
1616 DataType::Int8 => self.check_bounds::<i8>(max_value),
1617 DataType::Int16 => self.check_bounds::<i16>(max_value),
1618 DataType::Int32 => self.check_bounds::<i32>(max_value),
1619 DataType::Int64 => self.check_bounds::<i64>(max_value),
1620 _ => Err(ArrowError::InvalidArgumentError(format!(
1621 "Dictionary key type must be an integer, got {key_type}"
1622 ))),
1623 }
1624 }
1625 DataType::RunEndEncoded(run_ends, _values) => {
1626 let run_ends_data = self.child_at(0)?;
1627 match run_ends.data_type() {
1628 DataType::Int16 => run_ends_data.check_run_ends::<i16>(),
1629 DataType::Int32 => run_ends_data.check_run_ends::<i32>(),
1630 DataType::Int64 => run_ends_data.check_run_ends::<i64>(),
1631 data_type => Err(ArrowError::InvalidArgumentError(format!(
1632 "Run end type must be Int16, Int32 or Int64, got {data_type}"
1633 ))),
1634 }
1635 }
1636 _ => {
1637 Ok(())
1639 }
1640 }
1641 }
1642
1643 fn validate_each_offset<T, V>(&self, offset_limit: usize, validate: V) -> Result<(), ArrowError>
1654 where
1655 T: ArrowNativeType + TryInto<usize> + num_traits::Num + std::fmt::Display,
1656 V: Fn(usize, Range<usize>) -> Result<(), ArrowError>,
1657 {
1658 self.typed_offsets::<T>()?
1659 .iter()
1660 .enumerate()
1661 .map(|(i, x)| {
1662 let r = x.to_usize().ok_or_else(|| {
1664 ArrowError::InvalidArgumentError(format!(
1665 "Offset invariant failure: Could not convert offset {x} to usize at position {i}"))}
1666 );
1667 match r {
1669 Ok(n) if n <= offset_limit => Ok((i, n)),
1670 Ok(_) => Err(ArrowError::InvalidArgumentError(format!(
1671 "Offset invariant failure: offset at position {i} out of bounds: {x} > {offset_limit}"))
1672 ),
1673 Err(e) => Err(e),
1674 }
1675 })
1676 .scan(0_usize, |start, end| {
1677 match end {
1679 Ok((i, end)) if *start <= end => {
1680 let range = Some(Ok((i, *start..end)));
1681 *start = end;
1682 range
1683 }
1684 Ok((i, end)) => Some(Err(ArrowError::InvalidArgumentError(format!(
1685 "Offset invariant failure: non-monotonic offset at slot {}: {} > {}",
1686 i - 1, start, end))
1687 )),
1688 Err(err) => Some(Err(err)),
1689 }
1690 })
1691 .skip(1) .try_for_each(|res: Result<(usize, Range<usize>), ArrowError>| {
1693 let (item_index, range) = res?;
1694 validate(item_index-1, range)
1695 })
1696 }
1697
1698 fn validate_utf8<T>(&self) -> Result<(), ArrowError>
1701 where
1702 T: ArrowNativeType + TryInto<usize> + num_traits::Num + std::fmt::Display,
1703 {
1704 let values_buffer = &self.buffer_at(1)?.as_slice();
1705 if let Ok(values_str) = std::str::from_utf8(values_buffer) {
1706 self.validate_each_offset::<T, _>(values_buffer.len(), |string_index, range| {
1708 if !values_str.is_char_boundary(range.start)
1709 || !values_str.is_char_boundary(range.end)
1710 {
1711 return Err(ArrowError::InvalidArgumentError(format!(
1712 "incomplete utf-8 byte sequence from index {string_index}"
1713 )));
1714 }
1715 Ok(())
1716 })
1717 } else {
1718 self.validate_each_offset::<T, _>(values_buffer.len(), |string_index, range| {
1720 std::str::from_utf8(&values_buffer[range.clone()]).map_err(|e| {
1721 ArrowError::InvalidArgumentError(format!(
1722 "Invalid UTF8 sequence at string index {string_index} ({range:?}): {e}"
1723 ))
1724 })?;
1725 Ok(())
1726 })
1727 }
1728 }
1729
1730 fn validate_offsets_full<T>(&self, offset_limit: usize) -> Result<(), ArrowError>
1733 where
1734 T: ArrowNativeType + TryInto<usize> + num_traits::Num + std::fmt::Display,
1735 {
1736 self.validate_each_offset::<T, _>(offset_limit, |_string_index, _range| {
1737 Ok(())
1740 })
1741 }
1742
1743 fn check_bounds<T>(&self, max_value: i64) -> Result<(), ArrowError>
1746 where
1747 T: ArrowNativeType + TryInto<i64> + num_traits::Num + std::fmt::Display,
1748 {
1749 let indexes: &[T] = self.typed_buffer::<T>(0, self.len)?;
1752
1753 indexes.iter().enumerate().try_for_each(|(i, &dict_index)| {
1754 if self.is_null(i) {
1756 return Ok(());
1757 }
1758 let dict_index: i64 = dict_index.try_into().map_err(|_| {
1759 ArrowError::InvalidArgumentError(format!(
1760 "Value at position {i} out of bounds: {dict_index} (can not convert to i64)"
1761 ))
1762 })?;
1763
1764 if dict_index < 0 || dict_index > max_value {
1765 return Err(ArrowError::InvalidArgumentError(format!(
1766 "Value at position {i} out of bounds: {dict_index} (should be in [0, {max_value}])"
1767 )));
1768 }
1769 Ok(())
1770 })
1771 }
1772
1773 fn check_run_ends<T>(&self) -> Result<(), ArrowError>
1775 where
1776 T: ArrowNativeType + TryInto<i64> + num_traits::Num + std::fmt::Display,
1777 {
1778 let values = self.typed_buffer::<T>(0, self.len)?;
1779 let mut prev_value = 0_i64;
1780 values.iter().enumerate().try_for_each(|(ix, &inp_value)| {
1781 let value: i64 = inp_value.try_into().map_err(|_| {
1782 ArrowError::InvalidArgumentError(format!(
1783 "Value at position {ix} out of bounds: {inp_value} (can not convert to i64)"
1784 ))
1785 })?;
1786 if value <= 0_i64 {
1787 return Err(ArrowError::InvalidArgumentError(format!(
1788 "The values in run_ends array should be strictly positive. Found value {value} at index {ix} that does not match the criteria."
1789 )));
1790 }
1791 if ix > 0 && value <= prev_value {
1792 return Err(ArrowError::InvalidArgumentError(format!(
1793 "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."
1794 )));
1795 }
1796
1797 prev_value = value;
1798 Ok(())
1799 })?;
1800
1801 let len_plus_offset = checked_len_plus_offset(&self.data_type, self.len, self.offset)?;
1802 if prev_value.as_usize() < len_plus_offset {
1803 return Err(ArrowError::InvalidArgumentError(format!(
1804 "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}."
1805 )));
1806 }
1807 Ok(())
1808 }
1809
1810 pub fn ptr_eq(&self, other: &Self) -> bool {
1814 if self.offset != other.offset
1815 || self.len != other.len
1816 || self.data_type != other.data_type
1817 || self.buffers.len() != other.buffers.len()
1818 || self.child_data.len() != other.child_data.len()
1819 {
1820 return false;
1821 }
1822
1823 match (&self.nulls, &other.nulls) {
1824 (Some(a), Some(b)) if !a.inner().ptr_eq(b.inner()) => return false,
1825 (Some(_), None) | (None, Some(_)) => return false,
1826 _ => {}
1827 }
1828
1829 if !self
1830 .buffers
1831 .iter()
1832 .zip(other.buffers.iter())
1833 .all(|(a, b)| a.as_ptr() == b.as_ptr())
1834 {
1835 return false;
1836 }
1837
1838 self.child_data
1839 .iter()
1840 .zip(other.child_data.iter())
1841 .all(|(a, b)| a.ptr_eq(b))
1842 }
1843
1844 pub fn into_builder(self) -> ArrayDataBuilder {
1846 self.into()
1847 }
1848
1849 #[cfg(feature = "pool")]
1856 pub fn claim(&self, pool: &dyn arrow_buffer::MemoryPool) {
1857 for buffer in &self.buffers {
1859 buffer.claim(pool);
1860 }
1861
1862 if let Some(nulls) = &self.nulls {
1864 nulls.claim(pool);
1865 }
1866
1867 for child in &self.child_data {
1869 child.claim(pool);
1870 }
1871 }
1872}
1873
1874pub fn layout(data_type: &DataType) -> DataTypeLayout {
1877 use arrow_schema::IntervalUnit::*;
1880
1881 match data_type {
1882 DataType::Null => DataTypeLayout {
1883 buffers: vec![],
1884 can_contain_null_mask: false,
1885 variadic: false,
1886 },
1887 DataType::Boolean => DataTypeLayout {
1888 buffers: vec![BufferSpec::BitMap],
1889 can_contain_null_mask: true,
1890 variadic: false,
1891 },
1892 DataType::Int8 => DataTypeLayout::new_fixed_width::<i8>(),
1893 DataType::Int16 => DataTypeLayout::new_fixed_width::<i16>(),
1894 DataType::Int32 => DataTypeLayout::new_fixed_width::<i32>(),
1895 DataType::Int64 => DataTypeLayout::new_fixed_width::<i64>(),
1896 DataType::UInt8 => DataTypeLayout::new_fixed_width::<u8>(),
1897 DataType::UInt16 => DataTypeLayout::new_fixed_width::<u16>(),
1898 DataType::UInt32 => DataTypeLayout::new_fixed_width::<u32>(),
1899 DataType::UInt64 => DataTypeLayout::new_fixed_width::<u64>(),
1900 DataType::Float16 => DataTypeLayout::new_fixed_width::<half::f16>(),
1901 DataType::Float32 => DataTypeLayout::new_fixed_width::<f32>(),
1902 DataType::Float64 => DataTypeLayout::new_fixed_width::<f64>(),
1903 DataType::Timestamp(_, _) => DataTypeLayout::new_fixed_width::<i64>(),
1904 DataType::Date32 => DataTypeLayout::new_fixed_width::<i32>(),
1905 DataType::Date64 => DataTypeLayout::new_fixed_width::<i64>(),
1906 DataType::Time32(_) => DataTypeLayout::new_fixed_width::<i32>(),
1907 DataType::Time64(_) => DataTypeLayout::new_fixed_width::<i64>(),
1908 DataType::Interval(YearMonth) => DataTypeLayout::new_fixed_width::<i32>(),
1909 DataType::Interval(DayTime) => DataTypeLayout::new_fixed_width::<IntervalDayTime>(),
1910 DataType::Interval(MonthDayNano) => {
1911 DataTypeLayout::new_fixed_width::<IntervalMonthDayNano>()
1912 }
1913 DataType::Duration(_) => DataTypeLayout::new_fixed_width::<i64>(),
1914 DataType::Decimal32(_, _) => DataTypeLayout::new_fixed_width::<i32>(),
1915 DataType::Decimal64(_, _) => DataTypeLayout::new_fixed_width::<i64>(),
1916 DataType::Decimal128(_, _) => DataTypeLayout::new_fixed_width::<i128>(),
1917 DataType::Decimal256(_, _) => DataTypeLayout::new_fixed_width::<i256>(),
1918 DataType::FixedSizeBinary(size) => {
1919 let spec = BufferSpec::FixedWidth {
1920 byte_width: (*size).try_into().unwrap(),
1921 alignment: mem::align_of::<u8>(),
1922 };
1923 DataTypeLayout {
1924 buffers: vec![spec],
1925 can_contain_null_mask: true,
1926 variadic: false,
1927 }
1928 }
1929 DataType::Binary => DataTypeLayout::new_binary::<i32>(),
1930 DataType::LargeBinary => DataTypeLayout::new_binary::<i64>(),
1931 DataType::Utf8 => DataTypeLayout::new_binary::<i32>(),
1932 DataType::LargeUtf8 => DataTypeLayout::new_binary::<i64>(),
1933 DataType::BinaryView | DataType::Utf8View => DataTypeLayout::new_view(),
1934 DataType::FixedSizeList(_, _) => DataTypeLayout::new_nullable_empty(), DataType::List(_) => DataTypeLayout::new_fixed_width::<i32>(),
1936 DataType::ListView(_) => DataTypeLayout::new_list_view::<i32>(),
1937 DataType::LargeListView(_) => DataTypeLayout::new_list_view::<i64>(),
1938 DataType::LargeList(_) => DataTypeLayout::new_fixed_width::<i64>(),
1939 DataType::Map(_, _) => DataTypeLayout::new_fixed_width::<i32>(),
1940 DataType::Struct(_) => DataTypeLayout::new_nullable_empty(), DataType::RunEndEncoded(_, _) => DataTypeLayout::new_empty(), DataType::Union(_, mode) => {
1943 let type_ids = BufferSpec::FixedWidth {
1944 byte_width: mem::size_of::<i8>(),
1945 alignment: mem::align_of::<i8>(),
1946 };
1947
1948 DataTypeLayout {
1949 buffers: match mode {
1950 UnionMode::Sparse => {
1951 vec![type_ids]
1952 }
1953 UnionMode::Dense => {
1954 vec![
1955 type_ids,
1956 BufferSpec::FixedWidth {
1957 byte_width: mem::size_of::<i32>(),
1958 alignment: mem::align_of::<i32>(),
1959 },
1960 ]
1961 }
1962 },
1963 can_contain_null_mask: false,
1964 variadic: false,
1965 }
1966 }
1967 DataType::Dictionary(key_type, _value_type) => layout(key_type),
1968 }
1969}
1970
1971#[derive(Debug, PartialEq, Eq)]
1973pub struct DataTypeLayout {
1975 pub buffers: Vec<BufferSpec>,
1977
1978 pub can_contain_null_mask: bool,
1980
1981 pub variadic: bool,
1985}
1986
1987impl DataTypeLayout {
1988 pub fn new_fixed_width<T>() -> Self {
1990 Self {
1991 buffers: vec![BufferSpec::FixedWidth {
1992 byte_width: mem::size_of::<T>(),
1993 alignment: mem::align_of::<T>(),
1994 }],
1995 can_contain_null_mask: true,
1996 variadic: false,
1997 }
1998 }
1999
2000 pub fn new_nullable_empty() -> Self {
2003 Self {
2004 buffers: vec![],
2005 can_contain_null_mask: true,
2006 variadic: false,
2007 }
2008 }
2009
2010 pub fn new_empty() -> Self {
2013 Self {
2014 buffers: vec![],
2015 can_contain_null_mask: false,
2016 variadic: false,
2017 }
2018 }
2019
2020 pub fn new_binary<T>() -> Self {
2024 Self {
2025 buffers: vec![
2026 BufferSpec::FixedWidth {
2028 byte_width: mem::size_of::<T>(),
2029 alignment: mem::align_of::<T>(),
2030 },
2031 BufferSpec::VariableWidth,
2033 ],
2034 can_contain_null_mask: true,
2035 variadic: false,
2036 }
2037 }
2038
2039 pub fn new_view() -> Self {
2041 Self {
2042 buffers: vec![BufferSpec::FixedWidth {
2043 byte_width: mem::size_of::<u128>(),
2044 alignment: mem::align_of::<u128>(),
2045 }],
2046 can_contain_null_mask: true,
2047 variadic: true,
2048 }
2049 }
2050
2051 pub fn new_list_view<T>() -> Self {
2053 Self {
2054 buffers: vec![
2055 BufferSpec::FixedWidth {
2056 byte_width: mem::size_of::<T>(),
2057 alignment: mem::align_of::<T>(),
2058 },
2059 BufferSpec::FixedWidth {
2060 byte_width: mem::size_of::<T>(),
2061 alignment: mem::align_of::<T>(),
2062 },
2063 ],
2064 can_contain_null_mask: true,
2065 variadic: false,
2066 }
2067 }
2068}
2069
2070#[derive(Debug, PartialEq, Eq)]
2072pub enum BufferSpec {
2073 FixedWidth {
2084 byte_width: usize,
2086 alignment: usize,
2088 },
2089 VariableWidth,
2091 BitMap,
2097 AlwaysNull,
2100}
2101
2102impl PartialEq for ArrayData {
2103 fn eq(&self, other: &Self) -> bool {
2104 equal::equal(self, other)
2105 }
2106}
2107
2108#[derive(Debug, Clone)]
2127#[doc(hidden)]
2128pub struct UnsafeFlag(bool);
2129
2130impl UnsafeFlag {
2131 #[inline]
2135 pub const fn new() -> Self {
2136 Self(false)
2137 }
2138
2139 #[inline]
2149 pub unsafe fn set(&mut self, val: bool) {
2150 self.0 = val;
2151 }
2152
2153 #[inline]
2155 pub fn get(&self) -> bool {
2156 self.0
2157 }
2158}
2159
2160impl Default for UnsafeFlag {
2162 fn default() -> Self {
2163 Self::new()
2164 }
2165}
2166
2167#[derive(Debug)]
2169pub struct ArrayDataBuilder {
2170 data_type: DataType,
2171 len: usize,
2172 null_count: Option<usize>,
2173 null_bit_buffer: Option<Buffer>,
2174 nulls: Option<NullBuffer>,
2175 offset: usize,
2176 buffers: Vec<Buffer>,
2177 child_data: Vec<ArrayData>,
2178 align_buffers: bool,
2182 skip_validation: UnsafeFlag,
2192}
2193
2194impl ArrayDataBuilder {
2195 #[inline]
2196 pub const fn new(data_type: DataType) -> Self {
2198 Self {
2199 data_type,
2200 len: 0,
2201 null_count: None,
2202 null_bit_buffer: None,
2203 nulls: None,
2204 offset: 0,
2205 buffers: vec![],
2206 child_data: vec![],
2207 align_buffers: false,
2208 skip_validation: UnsafeFlag::new(),
2209 }
2210 }
2211
2212 pub fn data_type(self, data_type: DataType) -> Self {
2214 Self { data_type, ..self }
2215 }
2216
2217 #[inline]
2218 pub const fn len(mut self, n: usize) -> Self {
2220 self.len = n;
2221 self
2222 }
2223
2224 pub fn nulls(mut self, nulls: Option<NullBuffer>) -> Self {
2226 self.nulls = nulls;
2227 self.null_count = None;
2228 self.null_bit_buffer = None;
2229 self
2230 }
2231
2232 pub fn null_count(mut self, null_count: usize) -> Self {
2234 self.null_count = Some(null_count);
2235 self
2236 }
2237
2238 pub fn null_bit_buffer(mut self, buf: Option<Buffer>) -> Self {
2240 self.nulls = None;
2241 self.null_bit_buffer = buf;
2242 self
2243 }
2244
2245 #[inline]
2247 pub const fn offset(mut self, n: usize) -> Self {
2248 self.offset = n;
2249 self
2250 }
2251
2252 pub fn buffers(mut self, v: Vec<Buffer>) -> Self {
2254 self.buffers = v;
2255 self
2256 }
2257
2258 pub fn add_buffer(mut self, b: Buffer) -> Self {
2260 self.buffers.push(b);
2261 self
2262 }
2263
2264 pub fn add_buffers<I: IntoIterator<Item = Buffer>>(mut self, bs: I) -> Self {
2266 self.buffers.extend(bs);
2267 self
2268 }
2269
2270 pub fn child_data(mut self, v: Vec<ArrayData>) -> Self {
2272 self.child_data = v;
2273 self
2274 }
2275
2276 pub fn add_child_data(mut self, r: ArrayData) -> Self {
2278 self.child_data.push(r);
2279 self
2280 }
2281
2282 pub unsafe fn build_unchecked(self) -> ArrayData {
2298 unsafe { self.skip_validation(true) }.build().unwrap()
2299 }
2300
2301 pub fn build(self) -> Result<ArrayData, ArrowError> {
2310 let Self {
2311 data_type,
2312 len,
2313 null_count,
2314 null_bit_buffer,
2315 nulls,
2316 offset,
2317 buffers,
2318 child_data,
2319 align_buffers,
2320 skip_validation,
2321 } = self;
2322
2323 let validate = !skip_validation.get() || cfg!(feature = "force_validate");
2325 if validate && let Some(buffer) = null_bit_buffer.as_ref() {
2326 let len_plus_offset = checked_len_plus_offset(&data_type, len, offset)?;
2328 let needed_len = bit_util::ceil(len_plus_offset, 8);
2329 if buffer.len() < needed_len {
2330 return Err(ArrowError::InvalidArgumentError(format!(
2331 "null_bit_buffer size too small. got {} needed {}",
2332 buffer.len(),
2333 needed_len
2334 )));
2335 }
2336 }
2337
2338 let nulls = nulls
2339 .or_else(|| {
2340 let buffer = null_bit_buffer?;
2341 let buffer = BooleanBuffer::new(buffer, offset, len);
2342 Some(match null_count {
2343 Some(n) => {
2344 unsafe { NullBuffer::new_unchecked(buffer, n) }
2346 }
2347 None => NullBuffer::new(buffer),
2348 })
2349 })
2350 .filter(|b| b.null_count() != 0);
2351
2352 let mut data = ArrayData {
2353 data_type,
2354 len,
2355 offset,
2356 buffers,
2357 child_data,
2358 nulls,
2359 };
2360
2361 if align_buffers {
2362 data.align_buffers();
2363 }
2364
2365 if validate {
2366 data.validate_data()?;
2367 }
2368 Ok(data)
2369 }
2370
2371 pub fn align_buffers(mut self, align_buffers: bool) -> Self {
2387 self.align_buffers = align_buffers;
2388 self
2389 }
2390
2391 pub unsafe fn skip_validation(mut self, skip_validation: bool) -> Self {
2405 unsafe {
2406 self.skip_validation.set(skip_validation);
2407 }
2408 self
2409 }
2410}
2411
2412impl From<ArrayData> for ArrayDataBuilder {
2413 fn from(d: ArrayData) -> Self {
2414 Self {
2415 data_type: d.data_type,
2416 len: d.len,
2417 offset: d.offset,
2418 buffers: d.buffers,
2419 child_data: d.child_data,
2420 nulls: d.nulls,
2421 null_bit_buffer: None,
2422 null_count: None,
2423 align_buffers: false,
2424 skip_validation: UnsafeFlag::new(),
2425 }
2426 }
2427}
2428
2429pub(crate) fn get_fixed_size_binary_width(data_type: &DataType) -> usize {
2434 match data_type {
2435 DataType::FixedSizeBinary(i) => {
2436 if *i < 0 {
2437 panic!("cannot compare FixedSizeBinary({})", *i);
2438 }
2439 *i as usize
2440 }
2441 _ => unreachable!(),
2442 }
2443}
2444
2445#[cfg(test)]
2446mod tests {
2447 use arrow_schema::UnionFields;
2448
2449 use super::*;
2450 use crate::ByteView;
2451 use crate::transform::MutableArrayData;
2452 use arrow_buffer::{OffsetBuffer, ScalarBuffer};
2453 use arrow_schema::{Field, Fields};
2454
2455 fn make_i32_buffer(n: usize) -> Buffer {
2459 Buffer::from_slice_ref(vec![42i32; n])
2460 }
2461
2462 fn make_f32_buffer(n: usize) -> Buffer {
2464 Buffer::from_slice_ref(vec![42f32; n])
2465 }
2466
2467 #[test]
2468 fn test_builder() {
2469 let v = (0..25).collect::<Vec<i32>>();
2471 let b1 = Buffer::from_slice_ref(&v);
2472 let arr_data = ArrayData::builder(DataType::Int32)
2473 .len(20)
2474 .offset(5)
2475 .add_buffer(b1)
2476 .null_bit_buffer(Some(Buffer::from([
2477 0b01011111, 0b10110101, 0b01100011, 0b00011110,
2478 ])))
2479 .build()
2480 .unwrap();
2481
2482 assert_eq!(20, arr_data.len());
2483 assert_eq!(10, arr_data.null_count());
2484 assert_eq!(5, arr_data.offset());
2485 assert_eq!(1, arr_data.buffers().len());
2486 assert_eq!(
2487 Buffer::from_slice_ref(&v).as_slice(),
2488 arr_data.buffers()[0].as_slice()
2489 );
2490 }
2491
2492 #[test]
2493 fn test_builder_with_child_data() {
2494 let child_arr_data = ArrayData::try_new(
2495 DataType::Int32,
2496 5,
2497 None,
2498 0,
2499 vec![Buffer::from_slice_ref([1i32, 2, 3, 4, 5])],
2500 vec![],
2501 )
2502 .unwrap();
2503
2504 let field = Arc::new(Field::new("x", DataType::Int32, true));
2505 let data_type = DataType::Struct(vec![field].into());
2506
2507 let arr_data = ArrayData::builder(data_type)
2508 .len(5)
2509 .offset(0)
2510 .add_child_data(child_arr_data.clone())
2511 .build()
2512 .unwrap();
2513
2514 assert_eq!(5, arr_data.len());
2515 assert_eq!(1, arr_data.child_data().len());
2516 assert_eq!(child_arr_data, arr_data.child_data()[0]);
2517 }
2518
2519 #[test]
2520 fn test_struct_validation_accounts_for_parent_offset() {
2521 let data_type =
2522 DataType::Struct(Fields::from(vec![Field::new("x", DataType::Int32, false)]));
2523 let child = ArrayData::builder(DataType::Int32)
2524 .len(5)
2525 .add_buffer(Buffer::from_slice_ref([0, 1, 2, 3, 4]))
2526 .build()
2527 .unwrap();
2528
2529 let err = ArrayData::builder(data_type)
2531 .len(5)
2532 .offset(1)
2533 .add_child_data(child)
2534 .build()
2535 .unwrap_err()
2536 .to_string();
2537
2538 assert!(err.contains(
2539 "child array #0 for field x has length smaller than expected for struct array (5 < 6)"
2540 ));
2541 }
2542
2543 #[test]
2544 fn test_struct_non_nullable_child_nulls_account_for_parent_offset() {
2545 let build = |parent_nulls| {
2546 let child = ArrayData::builder(DataType::Int32)
2547 .len(5)
2548 .add_buffer(Buffer::from_slice_ref([0, 1, 2, 3, 4]))
2549 .nulls(Some(NullBuffer::new(BooleanBuffer::from(vec![
2550 true, true, false, true, true,
2551 ]))))
2552 .build()
2553 .unwrap();
2554
2555 ArrayData::builder(DataType::Struct(Fields::from(vec![Field::new(
2556 "x",
2557 DataType::Int32,
2558 false,
2559 )])))
2560 .len(4)
2561 .offset(1)
2562 .nulls(Some(NullBuffer::new(BooleanBuffer::from(parent_nulls))))
2563 .add_child_data(child)
2564 .build()
2565 };
2566
2567 assert!(build(vec![true, false, true, true]).is_ok());
2568 assert!(build(vec![true, true, false, true]).is_err());
2569 }
2570
2571 #[test]
2572 fn test_struct_equal_accounts_for_parent_offset() {
2573 let data_type =
2574 DataType::Struct(Fields::from(vec![Field::new("x", DataType::Int32, false)]));
2575
2576 let child1 = ArrayData::builder(DataType::Int32)
2577 .len(5)
2578 .add_buffer(Buffer::from_slice_ref([0, 1, 2, 3, 4]))
2579 .build()
2580 .unwrap();
2581 let child2 = child1.slice(1, 4);
2582
2583 let data1 = ArrayData::builder(data_type.clone())
2585 .len(4)
2586 .offset(1)
2587 .add_child_data(child1)
2588 .build()
2589 .unwrap();
2590 let data2 = ArrayData::builder(data_type)
2591 .len(4)
2592 .add_child_data(child2)
2593 .build()
2594 .unwrap();
2595
2596 assert_eq!(data1, data2);
2597 }
2598
2599 #[test]
2600 fn test_extend_struct_accounts_for_parent_offset() {
2601 let data_type =
2602 DataType::Struct(Fields::from(vec![Field::new("x", DataType::Int32, false)]));
2603 let child = ArrayData::builder(DataType::Int32)
2604 .len(5)
2605 .add_buffer(Buffer::from_slice_ref([0, 1, 2, 3, 4]))
2606 .build()
2607 .unwrap();
2608
2609 let data = ArrayData::builder(data_type)
2610 .len(4)
2611 .offset(1)
2612 .add_child_data(child)
2613 .build()
2614 .unwrap();
2615
2616 let mut mutable = MutableArrayData::new(vec![&data], false, data.len());
2617 mutable.try_extend(0, 0, data.len()).unwrap();
2618 let output = mutable.freeze();
2619
2620 assert_eq!(output.child_data()[0].buffer::<i32>(0), &[1, 2, 3, 4]);
2621 }
2622
2623 #[test]
2624 fn test_null_count() {
2625 let mut bit_v: [u8; 2] = [0; 2];
2626 bit_util::set_bit(&mut bit_v, 0);
2627 bit_util::set_bit(&mut bit_v, 3);
2628 bit_util::set_bit(&mut bit_v, 10);
2629 let arr_data = ArrayData::builder(DataType::Int32)
2630 .len(16)
2631 .add_buffer(make_i32_buffer(16))
2632 .null_bit_buffer(Some(Buffer::from(bit_v)))
2633 .build()
2634 .unwrap();
2635 assert_eq!(13, arr_data.null_count());
2636
2637 let mut bit_v: [u8; 2] = [0; 2];
2639 bit_util::set_bit(&mut bit_v, 0);
2640 bit_util::set_bit(&mut bit_v, 3);
2641 bit_util::set_bit(&mut bit_v, 10);
2642 let arr_data = ArrayData::builder(DataType::Int32)
2643 .len(12)
2644 .offset(2)
2645 .add_buffer(make_i32_buffer(14)) .null_bit_buffer(Some(Buffer::from(bit_v)))
2647 .build()
2648 .unwrap();
2649 assert_eq!(10, arr_data.null_count());
2650 }
2651
2652 #[test]
2653 fn test_null_buffer_ref() {
2654 let mut bit_v: [u8; 2] = [0; 2];
2655 bit_util::set_bit(&mut bit_v, 0);
2656 bit_util::set_bit(&mut bit_v, 3);
2657 bit_util::set_bit(&mut bit_v, 10);
2658 let arr_data = ArrayData::builder(DataType::Int32)
2659 .len(16)
2660 .add_buffer(make_i32_buffer(16))
2661 .null_bit_buffer(Some(Buffer::from(bit_v)))
2662 .build()
2663 .unwrap();
2664 assert!(arr_data.nulls().is_some());
2665 assert_eq!(&bit_v, arr_data.nulls().unwrap().validity());
2666 }
2667
2668 #[test]
2669 fn test_slice() {
2670 let mut bit_v: [u8; 2] = [0; 2];
2671 bit_util::set_bit(&mut bit_v, 0);
2672 bit_util::set_bit(&mut bit_v, 3);
2673 bit_util::set_bit(&mut bit_v, 10);
2674 let data = ArrayData::builder(DataType::Int32)
2675 .len(16)
2676 .add_buffer(make_i32_buffer(16))
2677 .null_bit_buffer(Some(Buffer::from(bit_v)))
2678 .build()
2679 .unwrap();
2680 let new_data = data.slice(1, 15);
2681 assert_eq!(data.len() - 1, new_data.len());
2682 assert_eq!(1, new_data.offset());
2683 assert_eq!(data.null_count(), new_data.null_count());
2684
2685 let new_data = new_data.slice(1, 14);
2687 assert_eq!(data.len() - 2, new_data.len());
2688 assert_eq!(2, new_data.offset());
2689 assert_eq!(data.null_count() - 1, new_data.null_count());
2690 }
2691
2692 #[test]
2693 #[should_panic(expected = "offset + length overflow")]
2694 fn test_slice_panics_on_offset_length_overflow() {
2695 let data = ArrayData::builder(DataType::Int32)
2696 .len(4)
2697 .add_buffer(make_i32_buffer(4))
2698 .build()
2699 .unwrap();
2700 let sliced = data.slice(1, 3);
2701
2702 sliced.slice(1, usize::MAX);
2703 }
2704
2705 #[test]
2706 fn test_typed_offsets_length_overflow() {
2707 let data = ArrayData {
2708 data_type: DataType::Binary,
2709 len: usize::MAX,
2710 offset: 0,
2711 buffers: vec![Buffer::from_slice_ref([0_i32])],
2712 child_data: vec![],
2713 nulls: None,
2714 };
2715 let err = data.typed_offsets::<i32>().unwrap_err();
2716
2717 assert_eq!(
2718 err.to_string(),
2719 format!(
2720 "Invalid argument error: Length {} with offset 1 overflows usize for Binary",
2721 usize::MAX
2722 )
2723 );
2724 }
2725
2726 #[test]
2727 fn test_validate_typed_buffer_length_overflow() {
2728 let data = ArrayData {
2729 data_type: DataType::Binary,
2730 len: 0,
2731 offset: 2,
2732 buffers: vec![Buffer::from_slice_ref([0_i32])],
2733 child_data: vec![],
2734 nulls: None,
2735 };
2736 let err = data.typed_buffer::<i32>(0, usize::MAX).unwrap_err();
2737
2738 assert_eq!(
2739 err.to_string(),
2740 format!(
2741 "Invalid argument error: Length {} with offset 2 overflows usize for Binary",
2742 usize::MAX
2743 )
2744 );
2745 }
2746
2747 fn try_new_binary_length_offset_overflow() -> Result<ArrayData, ArrowError> {
2749 ArrayData::try_new(
2750 DataType::Binary,
2751 usize::MAX,
2752 None,
2753 1,
2754 vec![
2755 Buffer::from_slice_ref([0_i32]),
2756 Buffer::from_iter(std::iter::empty::<u8>()),
2757 ],
2758 vec![],
2759 )
2760 }
2761
2762 #[cfg(not(feature = "force_validate"))]
2763 #[test]
2764 fn test_try_new_length_offset_overflow() {
2765 let err = try_new_binary_length_offset_overflow().unwrap_err();
2766
2767 assert_eq!(
2768 err.to_string(),
2769 format!(
2770 "Invalid argument error: Length {} with offset 1 overflows usize for Binary",
2771 usize::MAX
2772 )
2773 );
2774 }
2775
2776 #[cfg(feature = "force_validate")]
2777 #[test]
2778 #[should_panic(
2779 expected = "Length 18446744073709551615 with offset 1 overflows usize for Binary"
2780 )]
2781 fn test_try_new_length_offset_overflow_force_validate() {
2782 try_new_binary_length_offset_overflow().unwrap();
2783 }
2784
2785 #[test]
2786 fn test_equality() {
2787 let int_data = ArrayData::builder(DataType::Int32)
2788 .len(1)
2789 .add_buffer(make_i32_buffer(1))
2790 .build()
2791 .unwrap();
2792
2793 let float_data = ArrayData::builder(DataType::Float32)
2794 .len(1)
2795 .add_buffer(make_f32_buffer(1))
2796 .build()
2797 .unwrap();
2798 assert_ne!(int_data, float_data);
2799 assert!(!int_data.ptr_eq(&float_data));
2800 assert!(int_data.ptr_eq(&int_data));
2801
2802 let int_data_clone = int_data.clone();
2803 assert_eq!(int_data, int_data_clone);
2804 assert!(int_data.ptr_eq(&int_data_clone));
2805 assert!(int_data_clone.ptr_eq(&int_data));
2806
2807 let int_data_slice = int_data_clone.slice(1, 0);
2808 assert!(int_data_slice.ptr_eq(&int_data_slice));
2809 assert!(!int_data.ptr_eq(&int_data_slice));
2810 assert!(!int_data_slice.ptr_eq(&int_data));
2811
2812 let data_buffer = Buffer::from_slice_ref(b"abcdef");
2813 let offsets_buffer = Buffer::from_slice_ref([0_i32, 2_i32, 2_i32, 5_i32]);
2814 let string_data = ArrayData::try_new(
2815 DataType::Utf8,
2816 3,
2817 Some(Buffer::from_iter(vec![true, false, true])),
2818 0,
2819 vec![offsets_buffer, data_buffer],
2820 vec![],
2821 )
2822 .unwrap();
2823
2824 assert_ne!(float_data, string_data);
2825 assert!(!float_data.ptr_eq(&string_data));
2826
2827 assert!(string_data.ptr_eq(&string_data));
2828
2829 let string_data_cloned = string_data.clone();
2830 assert!(string_data_cloned.ptr_eq(&string_data));
2831 assert!(string_data.ptr_eq(&string_data_cloned));
2832
2833 let string_data_slice = string_data.slice(1, 2);
2834 assert!(string_data_slice.ptr_eq(&string_data_slice));
2835 assert!(!string_data_slice.ptr_eq(&string_data))
2836 }
2837
2838 #[test]
2839 fn test_slice_memory_size_view_payload_buffers() {
2840 for data_type in [DataType::Utf8View, DataType::BinaryView] {
2841 let inline_only = ArrayData::builder(data_type.clone())
2842 .len(2)
2843 .add_buffer(Buffer::from_vec(vec![0_u128; 2]))
2844 .build()
2845 .unwrap();
2846 assert_eq!(
2847 inline_only.get_slice_memory_size().unwrap(),
2848 2 * mem::size_of::<u128>()
2849 );
2850
2851 let mut first_payload = Vec::with_capacity(32);
2852 first_payload.extend_from_slice(b"first payload");
2853 let first_view =
2854 ByteView::new(first_payload.len().try_into().unwrap(), &first_payload[..4])
2855 .as_u128();
2856 let first_payload = Buffer::from_vec(first_payload);
2857 assert!(first_payload.capacity() > first_payload.len());
2858 let first_payload_capacity = first_payload.capacity();
2859
2860 let mut second_payload = Vec::with_capacity(64);
2861 second_payload.extend_from_slice(b"second payload");
2862 let second_view = ByteView::new(
2863 second_payload.len().try_into().unwrap(),
2864 &second_payload[..4],
2865 )
2866 .with_buffer_index(1)
2867 .as_u128();
2868 let second_payload = Buffer::from_vec(second_payload);
2869 assert!(second_payload.capacity() > second_payload.len());
2870 let second_payload_capacity = second_payload.capacity();
2871
2872 let data = ArrayData::builder(data_type)
2873 .len(3)
2874 .add_buffer(Buffer::from_vec(vec![first_view, 0_u128, second_view]))
2875 .add_buffer(first_payload)
2876 .add_buffer(second_payload)
2877 .build()
2878 .unwrap();
2879 let sliced = data.slice(1, 1);
2880
2881 assert_eq!(
2882 sliced.get_slice_memory_size().unwrap(),
2883 mem::size_of::<u128>() + first_payload_capacity + second_payload_capacity
2884 );
2885 }
2886 }
2887
2888 #[test]
2889 fn test_slice_memory_size_utf8_offset_buffer_len_plus_one() {
2890 let data_buffer = Buffer::from_slice_ref(b"helloworld");
2892 let offsets_buffer = Buffer::from_slice_ref([0_i32, 5_i32, 10_i32]);
2895 let array = ArrayData::try_new(
2896 DataType::Utf8,
2897 2,
2898 None,
2899 0,
2900 vec![offsets_buffer, data_buffer],
2901 vec![],
2902 )
2903 .unwrap();
2904 assert_eq!(array.get_slice_memory_size().unwrap(), 22); }
2906
2907 #[test]
2908 fn test_slice_memory_size_binary_offset_buffer_len_plus_one() {
2909 let data_buffer = Buffer::from_slice_ref([0u8, 1, 2, 3, 4]);
2912 let offsets_buffer = Buffer::from_slice_ref([0_i32, 2_i32, 5_i32]);
2914 let array = ArrayData::try_new(
2915 DataType::Binary,
2916 2,
2917 None,
2918 0,
2919 vec![offsets_buffer, data_buffer],
2920 vec![],
2921 )
2922 .unwrap();
2923 assert_eq!(array.get_slice_memory_size().unwrap(), 17); }
2925
2926 #[test]
2927 fn test_slice_memory_size() {
2928 let mut bit_v: [u8; 2] = [0; 2];
2929 bit_util::set_bit(&mut bit_v, 0);
2930 bit_util::set_bit(&mut bit_v, 3);
2931 bit_util::set_bit(&mut bit_v, 10);
2932 let data = ArrayData::builder(DataType::Int32)
2933 .len(16)
2934 .add_buffer(make_i32_buffer(16))
2935 .null_bit_buffer(Some(Buffer::from(bit_v)))
2936 .build()
2937 .unwrap();
2938 let new_data = data.slice(1, 14);
2939 assert_eq!(
2940 data.get_slice_memory_size().unwrap() - 8,
2941 new_data.get_slice_memory_size().unwrap()
2942 );
2943 let data_buffer = Buffer::from_slice_ref(b"abcdef");
2944 let offsets_buffer = Buffer::from_slice_ref([0_i32, 2_i32, 2_i32, 5_i32]);
2945 let string_data = ArrayData::try_new(
2946 DataType::Utf8,
2947 3,
2948 Some(Buffer::from_iter(vec![true, false, true])),
2949 0,
2950 vec![offsets_buffer, data_buffer],
2951 vec![],
2952 )
2953 .unwrap();
2954 let string_data_slice = string_data.slice(1, 2);
2955 assert_eq!(
2957 string_data.get_slice_memory_size().unwrap() - 6,
2958 string_data_slice.get_slice_memory_size().unwrap()
2959 );
2960 }
2961
2962 #[test]
2963 fn test_builder_rejects_short_null_bit_buffer() {
2964 for (len, offset) in [(8000, 0), (8, 1)] {
2965 let err = ArrayData::builder(DataType::Int32)
2966 .len(len)
2967 .offset(offset)
2968 .add_buffer(make_i32_buffer(len + offset))
2969 .null_bit_buffer(Some(Buffer::from([0_u8])))
2970 .build()
2971 .unwrap_err();
2972 assert_eq!(
2973 err.to_string(),
2974 format!(
2975 "Invalid argument error: null_bit_buffer size too small. got 1 needed {}",
2976 bit_util::ceil(len + offset, 8)
2977 )
2978 );
2979 }
2980 }
2981
2982 #[test]
2983 fn test_builder_null_bit_buffer_length_overflow() {
2984 let err = ArrayData::builder(DataType::Int32)
2985 .len(usize::MAX)
2986 .offset(1)
2987 .null_bit_buffer(Some(Buffer::default()))
2988 .build()
2989 .unwrap_err();
2990 assert_eq!(
2991 err.to_string(),
2992 format!(
2993 "Invalid argument error: Length {} with offset 1 overflows usize for Int32",
2994 usize::MAX
2995 )
2996 );
2997 }
2998
2999 #[test]
3000 fn test_builder_accepts_valid_null_bit_buffer() {
3001 for (len, offset) in [(8, 0), (7, 1)] {
3002 let data = ArrayData::builder(DataType::Int32)
3003 .len(len)
3004 .offset(offset)
3005 .add_buffer(make_i32_buffer(len + offset))
3006 .null_bit_buffer(Some(Buffer::from([0_u8])))
3007 .build()
3008 .unwrap();
3009 assert_eq!(data.len(), len);
3010 assert_eq!(data.offset(), offset);
3011 assert_eq!(data.null_count(), len);
3012 }
3013 }
3014
3015 #[test]
3016 fn test_count_nulls() {
3017 let buffer = Buffer::from([0b00010110, 0b10011111]);
3018 let buffer = NullBuffer::new(BooleanBuffer::new(buffer, 0, 16));
3019 let count = count_nulls(Some(&buffer), 0, 16);
3020 assert_eq!(count, 7);
3021
3022 let count = count_nulls(Some(&buffer), 4, 8);
3023 assert_eq!(count, 3);
3024 }
3025
3026 #[test]
3027 fn test_contains_nulls() {
3028 let buffer: Buffer =
3029 MutableBuffer::from_iter([false, false, false, true, true, false]).into();
3030 let buffer = NullBuffer::new(BooleanBuffer::new(buffer, 0, 6));
3031 assert!(contains_nulls(Some(&buffer), 0, 6));
3032 assert!(contains_nulls(Some(&buffer), 0, 3));
3033 assert!(!contains_nulls(Some(&buffer), 3, 2));
3034 assert!(!contains_nulls(Some(&buffer), 0, 0));
3035 }
3036
3037 #[test]
3038 fn test_alignment() {
3039 let buffer = Buffer::from_vec(vec![1_i32, 2_i32, 3_i32]);
3040 let sliced = buffer.slice(1);
3041
3042 let mut data = ArrayData {
3043 data_type: DataType::Int32,
3044 len: 0,
3045 offset: 0,
3046 buffers: vec![buffer],
3047 child_data: vec![],
3048 nulls: None,
3049 };
3050 data.validate_full().unwrap();
3051
3052 data.buffers[0] = sliced;
3054 let err = data.validate().unwrap_err();
3055
3056 assert_eq!(
3057 err.to_string(),
3058 "Invalid argument error: Misaligned buffers[0] in array of type Int32, offset from expected alignment of 4 by 1"
3059 );
3060
3061 data.align_buffers();
3062 data.validate_full().unwrap();
3063 }
3064
3065 #[test]
3066 fn test_alignment_struct() {
3067 let buffer = Buffer::from_vec(vec![1_i32, 2_i32, 3_i32]);
3068 let sliced = buffer.slice(1);
3069
3070 let child_data = ArrayData {
3071 data_type: DataType::Int32,
3072 len: 0,
3073 offset: 0,
3074 buffers: vec![buffer],
3075 child_data: vec![],
3076 nulls: None,
3077 };
3078
3079 let schema = DataType::Struct(Fields::from(vec![Field::new("a", DataType::Int32, false)]));
3080 let mut data = ArrayData {
3081 data_type: schema,
3082 len: 0,
3083 offset: 0,
3084 buffers: vec![],
3085 child_data: vec![child_data],
3086 nulls: None,
3087 };
3088 data.validate_full().unwrap();
3089
3090 data.child_data[0].buffers[0] = sliced;
3092 let err = data.validate().unwrap_err();
3093
3094 assert_eq!(
3095 err.to_string(),
3096 "Invalid argument error: Misaligned buffers[0] in array of type Int32, offset from expected alignment of 4 by 1"
3097 );
3098
3099 data.align_buffers();
3100 data.validate_full().unwrap();
3101 }
3102
3103 #[test]
3104 fn test_null_view_types() {
3105 let array_len = 32;
3106 let array = ArrayData::new_null(&DataType::BinaryView, 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(&DataType::Utf8View, array_len);
3113 assert_eq!(array.len(), array_len);
3114 for i in 0..array.len() {
3115 assert!(array.is_null(i));
3116 }
3117
3118 let array = ArrayData::new_null(
3119 &DataType::ListView(Arc::new(Field::new_list_field(DataType::Int32, true))),
3120 array_len,
3121 );
3122 assert_eq!(array.len(), array_len);
3123 for i in 0..array.len() {
3124 assert!(array.is_null(i));
3125 }
3126
3127 let array = ArrayData::new_null(
3128 &DataType::LargeListView(Arc::new(Field::new_list_field(DataType::Int32, true))),
3129 array_len,
3130 );
3131 assert_eq!(array.len(), array_len);
3132 for i in 0..array.len() {
3133 assert!(array.is_null(i));
3134 }
3135 }
3136
3137 #[test]
3139 fn test_dont_panic_on_bad_input_when_using_try_new() {
3140 let empty_bytes = Buffer::default();
3141
3142 let array_data = ArrayData::try_new(
3143 DataType::Utf8,
3144 1, None,
3146 0,
3147 vec![Buffer::from_vec(vec![0i32, 2i32]), empty_bytes],
3149 vec![],
3150 );
3151
3152 let res = array_data.expect_err("should get error");
3153
3154 assert_eq!(
3155 res.to_string(),
3156 "Invalid argument error: Last offset 2 of Utf8 is larger than values length 0"
3157 );
3158 }
3159
3160 #[test]
3163 #[cfg(not(feature = "force_validate"))]
3164 fn test_validate_values_rejects_a_non_integer_dictionary_key() {
3165 let values = valid_non_nullable_int32_array_data(2);
3166 let data_type = DataType::Dictionary(Box::new(DataType::Utf8), Box::new(DataType::Int32));
3167 let dictionary = unsafe {
3168 ArrayData::builder(data_type)
3169 .len(1)
3170 .add_child_data(values)
3171 .build_unchecked()
3172 };
3173
3174 let err = dictionary.validate_values().expect_err("should get error");
3175 assert_eq!(
3176 err.to_string(),
3177 "Invalid argument error: Dictionary key type must be an integer, got Utf8"
3178 );
3179 }
3180
3181 #[test]
3182 #[cfg(not(feature = "force_validate"))]
3183 fn test_validate_values_rejects_a_non_integer_run_end() {
3184 let data_type = DataType::RunEndEncoded(
3185 Arc::new(Field::new(
3186 Field::REE_RUN_ENDS_FIELD_DEFAULT_NAME,
3187 DataType::Utf8,
3188 false,
3189 )),
3190 Arc::new(Field::new(
3191 Field::REE_VALUES_FIELD_DEFAULT_NAME,
3192 DataType::Int32,
3193 true,
3194 )),
3195 );
3196 let run_end_encoded = unsafe {
3197 ArrayData::builder(data_type)
3198 .len(1)
3199 .add_child_data(valid_non_nullable_int32_array_data(1))
3200 .add_child_data(valid_non_nullable_int32_array_data(1))
3201 .build_unchecked()
3202 };
3203
3204 let err = run_end_encoded
3205 .validate_values()
3206 .expect_err("should get error");
3207 assert_eq!(
3208 err.to_string(),
3209 "Invalid argument error: Run end type must be Int16, Int32 or Int64, got Utf8"
3210 );
3211 }
3212
3213 #[test]
3215 #[cfg(not(feature = "force_validate"))]
3216 fn test_validate_values_rejects_missing_child_data() {
3217 let int32 = Box::new(DataType::Int32);
3218 let field = || Arc::new(Field::new("f", DataType::Int32, true));
3219 let data_types = [
3220 DataType::Dictionary(int32.clone(), int32.clone()),
3221 DataType::List(field()),
3222 DataType::LargeList(field()),
3223 DataType::RunEndEncoded(field(), field()),
3224 ];
3225
3226 for data_type in data_types {
3227 let data = unsafe {
3228 ArrayData::builder(data_type.clone())
3229 .len(1)
3230 .build_unchecked()
3231 };
3232 let err = data.validate_values().expect_err("should get error");
3233 assert_eq!(
3234 err.to_string(),
3235 format!(
3236 "Invalid argument error: {data_type} should contain at least 1 child data array(s), had 0"
3237 )
3238 );
3239 }
3240 }
3241
3242 #[test]
3244 #[cfg(not(feature = "force_validate"))]
3245 fn test_validate_values_rejects_missing_buffers() {
3246 let cases = [
3248 (DataType::Utf8, 1),
3249 (DataType::LargeUtf8, 1),
3250 (DataType::Binary, 1),
3251 (DataType::LargeBinary, 1),
3252 (DataType::BinaryView, 0),
3253 (DataType::Utf8View, 0),
3254 ];
3255
3256 for (data_type, missing) in cases {
3257 let data = unsafe {
3258 ArrayData::builder(data_type.clone())
3259 .len(1)
3260 .build_unchecked()
3261 };
3262 let err = data.validate_values().expect_err("should get error");
3263 assert_eq!(
3264 err.to_string(),
3265 format!(
3266 "Invalid argument error: {data_type} should contain at least {} buffer(s), had 0",
3267 missing + 1
3268 )
3269 );
3270 }
3271 }
3272
3273 #[test]
3275 #[cfg(not(feature = "force_validate"))]
3276 fn test_validate_values_rejects_a_short_dictionary_keys_buffer() {
3277 let data_type = DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Int32));
3278 let dictionary = unsafe {
3279 ArrayData::builder(data_type)
3280 .len(4)
3281 .add_buffer(Buffer::from_slice_ref([1_i32, 0]))
3282 .add_child_data(valid_non_nullable_int32_array_data(2))
3283 .build_unchecked()
3284 };
3285
3286 let err = dictionary.validate_values().expect_err("should get error");
3287 assert_eq!(
3288 err.to_string(),
3289 "Invalid argument error: Buffer 0 of Dictionary(Int32, Int32) isn't large enough. Expected 16 bytes got 8"
3290 );
3291 }
3292
3293 #[test]
3294 fn should_fail_validation_when_having_map_field_type_is_not_struct() {
3295 let map_field = Field::new("key", DataType::Int32, false);
3296
3297 let map_field_data = valid_non_nullable_int32_array_data(2);
3298
3299 let results = test_both_builder_and_array_data(
3300 DataType::Map(map_field.into(), false),
3301 1,
3302 None,
3303 0,
3304 vec![
3305 OffsetBuffer::<i32>::from_lengths(vec![2])
3306 .into_inner()
3307 .into(),
3308 ],
3309 vec![map_field_data],
3310 );
3311
3312 for result in results {
3313 let array_data_err = result.expect_err("should fail for non struct field");
3314
3315 match array_data_err {
3316 ArrowError::InvalidArgumentError(msg) => {
3317 assert_eq!(
3318 msg,
3319 "Map field should be a entries struct data type, got Int32 instead"
3320 )
3321 }
3322 _ => panic!("unexpected error type {array_data_err}"),
3323 }
3324 }
3325 }
3326
3327 #[test]
3328 fn should_fail_validation_when_having_map_entries_only_have_1_field() {
3329 let struct_data_type = DataType::Struct(Fields::from(vec![Field::new(
3330 Field::MAP_KEY_FIELD_DEFAULT_NAME,
3331 DataType::Int32,
3332 false,
3333 )]));
3334
3335 let key_array_data = valid_non_nullable_int32_array_data(2);
3336
3337 let struct_data = {
3338 let builder = ArrayDataBuilder::new(struct_data_type.clone())
3339 .len(2)
3340 .nulls(None)
3341 .child_data(vec![key_array_data]);
3342
3343 builder.build().unwrap()
3344 };
3345
3346 let results = test_both_builder_and_array_data(
3347 DataType::Map(
3348 Field::new(
3349 Field::MAP_ENTRIES_FIELD_DEFAULT_NAME,
3350 struct_data_type,
3351 false,
3352 )
3353 .into(),
3354 false,
3355 ),
3356 1,
3357 None,
3358 0,
3359 vec![
3360 OffsetBuffer::<i32>::from_lengths(vec![2])
3361 .into_inner()
3362 .into(),
3363 ],
3364 vec![struct_data],
3365 );
3366
3367 for result in results {
3368 let array_data_err = result.expect_err("should fail for nullable key");
3369
3370 match array_data_err {
3371 ArrowError::InvalidArgumentError(msg) => {
3372 assert_eq!(
3373 msg,
3374 "Map entries data type should be a struct containing 2 fields, got 1 fields"
3375 )
3376 }
3377 _ => panic!("unexpected error type {array_data_err}"),
3378 }
3379 }
3380 }
3381
3382 #[test]
3383 fn should_fail_validation_when_having_map_entries_have_3_fields() {
3384 let struct_data_type = DataType::Struct(Fields::from(vec![
3385 Field::new(Field::MAP_KEY_FIELD_DEFAULT_NAME, DataType::Int32, false),
3386 Field::new(Field::MAP_VALUE_FIELD_DEFAULT_NAME, DataType::Utf8, true),
3387 Field::new("other", DataType::Int32, true),
3388 ]));
3389
3390 let key_array_data = valid_non_nullable_int32_array_data(2);
3391
3392 let values_array_data = valid_string_array_data(2);
3393
3394 let other_array_data = key_array_data.clone();
3395
3396 let struct_data = {
3397 let builder = ArrayDataBuilder::new(struct_data_type.clone())
3398 .len(2)
3399 .nulls(None)
3400 .child_data(vec![key_array_data, values_array_data, other_array_data]);
3401
3402 builder.build().unwrap()
3403 };
3404
3405 let results = test_both_builder_and_array_data(
3406 DataType::Map(
3407 Field::new(
3408 Field::MAP_ENTRIES_FIELD_DEFAULT_NAME,
3409 struct_data_type,
3410 false,
3411 )
3412 .into(),
3413 false,
3414 ),
3415 1,
3416 None,
3417 0,
3418 vec![
3419 OffsetBuffer::<i32>::from_lengths(vec![2])
3420 .into_inner()
3421 .into(),
3422 ],
3423 vec![struct_data],
3424 );
3425
3426 for result in results {
3427 let array_data_err = result.expect_err("should fail for nullable key");
3428
3429 match array_data_err {
3430 ArrowError::InvalidArgumentError(msg) => {
3431 assert_eq!(
3432 msg,
3433 "Map entries data type should be a struct containing 2 fields, got 3 fields"
3434 )
3435 }
3436 _ => panic!("unexpected error type {array_data_err}"),
3437 }
3438 }
3439 }
3440
3441 #[test]
3442 fn should_fail_validation_when_having_nullable_map_keys() {
3443 let struct_data_type = DataType::Struct(Fields::from(vec![
3444 Field::new(Field::MAP_KEY_FIELD_DEFAULT_NAME, DataType::Int32, true),
3445 Field::new(Field::MAP_VALUE_FIELD_DEFAULT_NAME, DataType::Utf8, true),
3446 ]));
3447
3448 let key_array_data = valid_non_nullable_int32_array_data(2);
3449 let values_array_data = valid_string_array_data(2);
3450
3451 let struct_data = {
3452 let builder = ArrayDataBuilder::new(struct_data_type.clone())
3453 .len(2)
3454 .nulls(None)
3455 .child_data(vec![key_array_data, values_array_data]);
3456
3457 builder.build().unwrap()
3458 };
3459
3460 let results = test_both_builder_and_array_data(
3461 DataType::Map(
3462 Field::new(
3463 Field::MAP_ENTRIES_FIELD_DEFAULT_NAME,
3464 struct_data_type,
3465 false,
3466 )
3467 .into(),
3468 false,
3469 ),
3470 1,
3471 None,
3472 0,
3473 vec![
3474 OffsetBuffer::<i32>::from_lengths(vec![2])
3475 .into_inner()
3476 .into(),
3477 ],
3478 vec![struct_data],
3479 );
3480
3481 for result in results {
3482 let array_data_err = result.expect_err("should fail for nullable key");
3483
3484 match array_data_err {
3485 ArrowError::InvalidArgumentError(msg) => {
3486 assert_eq!(msg, "Map key field must not be nullable")
3487 }
3488 _ => panic!("unexpected error type {array_data_err}"),
3489 }
3490 }
3491 }
3492
3493 #[test]
3494 fn should_fail_validation_when_having_entries_is_nullable_for_map() {
3495 let struct_data_type = DataType::Struct(Fields::from(vec![
3496 Field::new(Field::MAP_KEY_FIELD_DEFAULT_NAME, DataType::Int32, false),
3497 Field::new(Field::MAP_VALUE_FIELD_DEFAULT_NAME, DataType::Utf8, true),
3498 ]));
3499
3500 let key_array_data = valid_non_nullable_int32_array_data(2);
3501
3502 let values_array_data = valid_string_array_data(2);
3503
3504 let struct_data = {
3505 let builder = ArrayDataBuilder::new(struct_data_type.clone())
3506 .len(2)
3507 .nulls(None)
3508 .child_data(vec![key_array_data, values_array_data]);
3509
3510 builder.build().unwrap()
3511 };
3512
3513 let results = test_both_builder_and_array_data(
3514 DataType::Map(
3515 Field::new(
3516 Field::MAP_ENTRIES_FIELD_DEFAULT_NAME,
3517 struct_data_type,
3518 true,
3519 )
3520 .into(),
3521 false,
3522 ),
3523 1,
3524 None,
3525 0,
3526 vec![
3527 OffsetBuffer::<i32>::from_lengths(vec![2])
3528 .into_inner()
3529 .into(),
3530 ],
3531 vec![struct_data],
3532 );
3533
3534 for result in results {
3535 let array_data_err = result.expect_err("should fail for nullable entries");
3536
3537 match array_data_err {
3538 ArrowError::InvalidArgumentError(msg) => assert_eq!(
3539 msg,
3540 "The nullable should be set to false for the map entries field."
3541 ),
3542 _ => panic!("unexpected error type {array_data_err}"),
3543 }
3544 }
3545 }
3546
3547 #[test]
3548 fn should_allow_to_create_map_from_data() {
3549 let struct_data_type = DataType::Struct(Fields::from(vec![
3550 Field::new(Field::MAP_KEY_FIELD_DEFAULT_NAME, DataType::Int32, false),
3551 Field::new(Field::MAP_VALUE_FIELD_DEFAULT_NAME, DataType::Utf8, true),
3552 ]));
3553
3554 let key_array_data = valid_non_nullable_int32_array_data(2);
3555 let values_array_data = valid_string_array_data(2);
3556
3557 let struct_data = {
3558 let builder = ArrayDataBuilder::new(struct_data_type.clone())
3559 .len(2)
3560 .nulls(None)
3561 .child_data(vec![key_array_data, values_array_data]);
3562
3563 builder.build().unwrap()
3564 };
3565
3566 let results = test_both_builder_and_array_data(
3567 DataType::Map(
3568 Field::new(
3569 Field::MAP_ENTRIES_FIELD_DEFAULT_NAME,
3570 struct_data_type,
3571 false,
3572 )
3573 .into(),
3574 false,
3575 ),
3576 1,
3577 None,
3578 0,
3579 vec![
3580 OffsetBuffer::<i32>::from_lengths(vec![2])
3581 .into_inner()
3582 .into(),
3583 ],
3584 vec![struct_data],
3585 );
3586
3587 for result in results {
3588 result.expect("should be able to create map ArrayData");
3589 }
3590 }
3591
3592 fn valid_string_array_data(length: usize) -> ArrayData {
3593 let offsets = OffsetBuffer::<i32>::from_lengths(vec![0; length])
3594 .into_inner()
3595 .into_inner();
3596 let empty_bytes = Buffer::default();
3597
3598 let builder = ArrayDataBuilder::new(DataType::Utf8)
3599 .len(length)
3600 .buffers(vec![offsets, empty_bytes])
3601 .nulls(None);
3602
3603 builder.build().unwrap()
3604 }
3605
3606 fn valid_non_nullable_int32_array_data(length: usize) -> ArrayData {
3607 let builder = ArrayDataBuilder::new(DataType::Int32)
3608 .len(length)
3609 .nulls(None)
3610 .buffers(vec![
3611 ScalarBuffer::<i32>::from(vec![1; length]).into_inner(),
3612 ]);
3613
3614 builder.build().unwrap()
3615 }
3616
3617 #[test]
3618 fn empty_and_null_map_array_should_pass_validation() {
3619 let dt = DataType::Map(
3620 Field::new(
3621 Field::MAP_ENTRIES_FIELD_DEFAULT_NAME,
3622 DataType::Struct(Fields::from(vec![
3623 Field::new(Field::MAP_KEY_FIELD_DEFAULT_NAME, DataType::Int32, false),
3624 Field::new(Field::MAP_VALUE_FIELD_DEFAULT_NAME, DataType::Utf8, true),
3625 ])),
3626 false,
3627 )
3628 .into(),
3629 false,
3630 );
3631
3632 ArrayData::new_empty(&dt).validate_full().unwrap();
3633 ArrayData::new_null(&dt, 1).validate_full().unwrap();
3634 }
3635
3636 #[test]
3637 fn null_buffer_offset_is_independent_of_data_offset() {
3638 let int_data = ArrayData::builder(DataType::UInt32)
3640 .offset(50)
3641 .len(50)
3642 .add_buffer(Buffer::from_vec(vec![0_u32; 100]))
3643 .build()
3644 .unwrap();
3645 int_data.validate().unwrap();
3646
3647 let nulls = NullBuffer::new(BooleanBuffer::from(vec![false; 100]).slice(0, 50));
3649 let with_sliced_nulls = int_data
3650 .clone()
3651 .into_builder()
3652 .nulls(Some(nulls))
3653 .build()
3654 .unwrap();
3655 with_sliced_nulls.validate().unwrap();
3656
3657 let nulls = NullBuffer::new(BooleanBuffer::from(vec![false; 50]));
3660 let with_unsliced_nulls = int_data.into_builder().nulls(Some(nulls)).build().unwrap();
3661 with_unsliced_nulls.validate().unwrap();
3662 assert_eq!(with_unsliced_nulls.null_count(), 50);
3663 }
3664
3665 fn test_both_builder_and_array_data(
3666 data_type: DataType,
3667 len: usize,
3668 null_bit_buffer: Option<Buffer>,
3669 offset: usize,
3670 buffers: Vec<Buffer>,
3671 child_data: Vec<ArrayData>,
3672 ) -> [Result<ArrayData, ArrowError>; 2] {
3673 let from_builder_res = ArrayData::builder(data_type.clone())
3674 .len(len)
3675 .add_buffers(buffers.clone())
3676 .null_bit_buffer(null_bit_buffer.clone())
3677 .offset(offset)
3678 .child_data(child_data.clone())
3679 .build();
3680
3681 let from_try_new_res =
3682 ArrayData::try_new(data_type, len, null_bit_buffer, offset, buffers, child_data);
3683
3684 [from_builder_res, from_try_new_res]
3685 }
3686
3687 #[test]
3688 fn test_new_null_empty_union() {
3689 for mode in [UnionMode::Sparse, UnionMode::Dense] {
3690 let data_type = DataType::Union(UnionFields::empty(), mode);
3691 let data = ArrayData::new_null(&data_type, 0);
3692 data.validate_full()
3693 .expect("an empty union of length zero is valid");
3694 assert_eq!(data.len(), 0);
3695 assert!(data.child_data().is_empty());
3696 }
3697 }
3698
3699 #[test]
3700 #[should_panic(expected = "cannot construct null data from an empty union")]
3701 fn test_new_null_empty_union_with_slots() {
3702 let data_type = DataType::Union(UnionFields::empty(), UnionMode::Dense);
3703 let _ = ArrayData::new_null(&data_type, 1);
3704 }
3705}