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 if let Some(null_bit_buffer) = null_bit_buffer.as_ref() {
340 let len_plus_offset = checked_len_plus_offset(&data_type, len, offset)?;
341 let needed_len = bit_util::ceil(len_plus_offset, 8);
342 if null_bit_buffer.len() < needed_len {
343 return Err(ArrowError::InvalidArgumentError(format!(
344 "null_bit_buffer size too small. got {} needed {}",
345 null_bit_buffer.len(),
346 needed_len
347 )));
348 }
349 }
350
351 let builder = Self::inner_new_builder(
352 data_type,
353 len,
354 None,
355 null_bit_buffer,
356 offset,
357 buffers,
358 child_data,
359 );
360
361 assert!(!builder.skip_validation.get());
362
363 builder.build()
368 }
369
370 fn inner_new_builder(
371 data_type: DataType,
372 len: usize,
373 null_count: Option<usize>,
374 null_bit_buffer: Option<Buffer>,
375 offset: usize,
376 buffers: Vec<Buffer>,
377 child_data: Vec<ArrayData>,
378 ) -> ArrayDataBuilder {
379 ArrayDataBuilder {
380 data_type,
381 len,
382 null_count,
383 null_bit_buffer,
384 nulls: None,
385 offset,
386 buffers,
387 child_data,
388 align_buffers: false,
389 skip_validation: UnsafeFlag::new(),
390 }
391 }
392
393 pub fn into_parts(
399 self,
400 ) -> (
401 DataType,
402 usize,
403 Option<NullBuffer>,
404 usize,
405 Vec<Buffer>,
406 Vec<ArrayData>,
407 ) {
408 let Self {
409 data_type,
410 len,
411 nulls,
412 offset,
413 buffers,
414 child_data,
415 } = self;
416
417 (data_type, len, nulls, offset, buffers, child_data)
418 }
419
420 #[inline]
422 pub const fn builder(data_type: DataType) -> ArrayDataBuilder {
423 ArrayDataBuilder::new(data_type)
424 }
425
426 #[inline]
428 pub const fn data_type(&self) -> &DataType {
429 &self.data_type
430 }
431
432 pub fn buffers(&self) -> &[Buffer] {
434 &self.buffers
435 }
436
437 pub fn child_data(&self) -> &[ArrayData] {
445 &self.child_data[..]
446 }
447
448 #[inline]
450 pub fn is_null(&self, i: usize) -> bool {
451 match &self.nulls {
452 Some(v) => v.is_null(i),
453 None => false,
454 }
455 }
456
457 #[inline]
461 pub fn nulls(&self) -> Option<&NullBuffer> {
462 self.nulls.as_ref()
463 }
464
465 #[inline]
467 pub fn is_valid(&self, i: usize) -> bool {
468 !self.is_null(i)
469 }
470
471 #[inline]
473 pub const fn len(&self) -> usize {
474 self.len
475 }
476
477 #[inline]
479 pub const fn is_empty(&self) -> bool {
480 self.len == 0
481 }
482
483 #[inline]
535 pub const fn offset(&self) -> usize {
536 self.offset
537 }
538
539 #[inline]
541 pub fn null_count(&self) -> usize {
542 self.nulls
543 .as_ref()
544 .map(|x| x.null_count())
545 .unwrap_or_default()
546 }
547
548 pub fn get_buffer_memory_size(&self) -> usize {
560 let mut size = 0;
561 for buffer in &self.buffers {
562 size += buffer.capacity();
563 }
564 if let Some(bitmap) = &self.nulls {
565 size += bitmap.buffer().capacity()
566 }
567 for child in &self.child_data {
568 size += child.get_buffer_memory_size();
569 }
570 size
571 }
572
573 pub fn get_slice_memory_size(&self) -> Result<usize, ArrowError> {
588 let mut result: usize = 0;
589 let layout = layout(&self.data_type);
590
591 for spec in &layout.buffers {
592 match spec {
593 BufferSpec::FixedWidth { byte_width, .. } => {
594 let len = match self.data_type {
597 DataType::Utf8
598 | DataType::LargeUtf8
599 | DataType::Binary
600 | DataType::LargeBinary
601 | DataType::List(_)
602 | DataType::LargeList(_)
603 | DataType::Map(_, _) => self.len + 1,
604 _ => self.len,
605 };
606 let buffer_size = len.checked_mul(*byte_width).ok_or_else(|| {
607 ArrowError::ComputeError(
608 "Integer overflow computing buffer size".to_string(),
609 )
610 })?;
611 result += buffer_size;
612 }
613 BufferSpec::VariableWidth => {
614 let buffer_len = match self.data_type {
615 DataType::Utf8 | DataType::Binary => {
616 let offsets = self.typed_offsets::<i32>()?;
617 (offsets[self.len] - offsets[0]) as usize
618 }
619 DataType::LargeUtf8 | DataType::LargeBinary => {
620 let offsets = self.typed_offsets::<i64>()?;
621 (offsets[self.len] - offsets[0]) as usize
622 }
623 _ => {
624 return Err(ArrowError::NotYetImplemented(format!(
625 "Invalid data type for VariableWidth buffer. Expected Utf8, LargeUtf8, Binary or LargeBinary. Got {}",
626 self.data_type
627 )));
628 }
629 };
630 result += buffer_len;
631 }
632 BufferSpec::BitMap => {
633 let buffer_size = bit_util::ceil(self.len, 8);
634 result += buffer_size;
635 }
636 BufferSpec::AlwaysNull => {
637 }
639 }
640 }
641
642 if layout.variadic {
643 for buffer in self.buffers.iter().skip(layout.buffers.len()) {
645 result += buffer.capacity();
646 }
647 }
648
649 if self.nulls().is_some() {
650 result += bit_util::ceil(self.len, 8);
651 }
652
653 for child in &self.child_data {
654 result += child.get_slice_memory_size()?;
655 }
656 Ok(result)
657 }
658
659 pub fn get_array_memory_size(&self) -> usize {
668 let mut size = mem::size_of_val(self);
669
670 for buffer in &self.buffers {
672 size += mem::size_of::<Buffer>();
673 size += buffer.capacity();
674 }
675 if let Some(nulls) = &self.nulls {
676 size += nulls.buffer().capacity();
677 }
678 for child in &self.child_data {
679 size += child.get_array_memory_size();
680 }
681
682 size
683 }
684
685 pub fn slice(&self, offset: usize, length: usize) -> ArrayData {
693 let end = offset
694 .checked_add(length)
695 .expect("offset + length overflow");
696 assert!(end <= self.len());
697
698 if let DataType::Struct(_) = self.data_type() {
699 let child_offset = self.offset + offset;
704 ArrayData {
705 data_type: self.data_type().clone(),
706 len: length,
707 offset: 0,
708 buffers: self.buffers.clone(),
709 child_data: self
710 .child_data()
711 .iter()
712 .map(|data| data.slice(child_offset, length))
713 .collect(),
714 nulls: self.nulls.as_ref().map(|x| x.slice(offset, length)),
717 }
718 } else {
719 let mut new_data = self.clone();
720
721 new_data.len = length;
722 new_data.offset = offset + self.offset;
723 new_data.nulls = self.nulls.as_ref().map(|x| x.slice(offset, length));
724
725 new_data
726 }
727 }
728
729 pub fn buffer<T: ArrowNativeType>(&self, buffer: usize) -> &[T] {
736 &self.buffers()[buffer].typed_data()[self.offset..]
737 }
738
739 pub fn new_null(data_type: &DataType, len: usize) -> Self {
745 let bit_len = bit_util::ceil(len, 8);
746 let zeroed = |len: usize| Buffer::from(MutableBuffer::from_len_zeroed(len));
747
748 let (buffers, child_data, has_nulls) = match data_type.primitive_width() {
749 Some(width) => (vec![zeroed(width * len)], vec![], true),
750 None => match data_type {
751 DataType::Null => (vec![], vec![], false),
752 DataType::Boolean => (vec![zeroed(bit_len)], vec![], true),
753 DataType::Binary | DataType::Utf8 => {
754 (vec![zeroed((len + 1) * 4), zeroed(0)], vec![], true)
755 }
756 DataType::BinaryView | DataType::Utf8View => (vec![zeroed(len * 16)], vec![], true),
757 DataType::LargeBinary | DataType::LargeUtf8 => {
758 (vec![zeroed((len + 1) * 8), zeroed(0)], vec![], true)
759 }
760 DataType::FixedSizeBinary(i) => {
761 if *i < 0 {
762 panic!("cannot construct null data from FixedSizeBinary({i})");
763 }
764 (vec![zeroed(*i as usize * len)], vec![], true)
765 }
766 DataType::List(f) | DataType::Map(f, _) => (
767 vec![zeroed((len + 1) * 4)],
768 vec![ArrayData::new_empty(f.data_type())],
769 true,
770 ),
771 DataType::LargeList(f) => (
772 vec![zeroed((len + 1) * 8)],
773 vec![ArrayData::new_empty(f.data_type())],
774 true,
775 ),
776 DataType::ListView(f) => (
777 vec![zeroed(len * 4), zeroed(len * 4)],
778 vec![ArrayData::new_empty(f.data_type())],
779 true,
780 ),
781 DataType::LargeListView(f) => (
782 vec![zeroed(len * 8), zeroed(len * 8)],
783 vec![ArrayData::new_empty(f.data_type())],
784 true,
785 ),
786 DataType::FixedSizeList(f, list_len) => (
787 vec![],
788 vec![ArrayData::new_null(f.data_type(), *list_len as usize * len)],
789 true,
790 ),
791 DataType::Struct(fields) => (
792 vec![],
793 fields
794 .iter()
795 .map(|f| Self::new_null(f.data_type(), len))
796 .collect(),
797 true,
798 ),
799 DataType::Dictionary(k, v) => (
800 vec![zeroed(k.primitive_width().unwrap() * len)],
801 vec![ArrayData::new_empty(v.as_ref())],
802 true,
803 ),
804 DataType::Union(f, mode) => {
805 let (id, _) = f.iter().next().unwrap();
806 let ids = Buffer::from_iter(std::iter::repeat_n(id, len));
807 let buffers = match mode {
808 UnionMode::Sparse => vec![ids],
809 UnionMode::Dense => {
810 let end_offset = i32::from_usize(len).unwrap();
811 vec![ids, Buffer::from_iter(0_i32..end_offset)]
812 }
813 };
814
815 let children = f
816 .iter()
817 .enumerate()
818 .map(|(idx, (_, f))| {
819 if idx == 0 || *mode == UnionMode::Sparse {
820 Self::new_null(f.data_type(), len)
821 } else {
822 Self::new_empty(f.data_type())
823 }
824 })
825 .collect();
826
827 (buffers, children, false)
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 Ok(())
1228 }
1229 DataType::LargeList(field) => {
1230 let values_data = self.get_single_valid_child_data(field.data_type())?;
1231 self.validate_offsets::<i64>(values_data.len)?;
1232 Ok(())
1233 }
1234 DataType::Map(field, _) => {
1235 let DataType::Struct(entries_fields) = field.data_type() else {
1236 return Err(ArrowError::InvalidArgumentError(format!(
1237 "Map field should be a entries struct data type, got {:?} instead",
1238 field.data_type()
1239 )));
1240 };
1241 if entries_fields.len() != 2 {
1242 return Err(ArrowError::InvalidArgumentError(format!(
1243 "Map entries data type should be a struct containing 2 fields, got {} fields",
1244 entries_fields.len()
1245 )));
1246 }
1247
1248 if entries_fields[0].is_nullable() {
1250 return Err(ArrowError::InvalidArgumentError(
1251 "Map key field must not be nullable".to_string(),
1252 ));
1253 }
1254 let values_data = self.get_single_valid_child_data(field.data_type())?;
1255 self.validate_offsets::<i32>(values_data.len)?;
1256 Ok(())
1257 }
1258 DataType::ListView(field) => {
1259 let values_data = self.get_single_valid_child_data(field.data_type())?;
1260 self.validate_offsets_and_sizes::<i32>(values_data.len)?;
1261 Ok(())
1262 }
1263 DataType::LargeListView(field) => {
1264 let values_data = self.get_single_valid_child_data(field.data_type())?;
1265 self.validate_offsets_and_sizes::<i64>(values_data.len)?;
1266 Ok(())
1267 }
1268 DataType::FixedSizeList(field, list_size) => {
1269 let values_data = self.get_single_valid_child_data(field.data_type())?;
1270
1271 let list_size: usize = (*list_size).try_into().map_err(|_| {
1272 ArrowError::InvalidArgumentError(format!(
1273 "{} has a negative list_size {}",
1274 self.data_type, list_size
1275 ))
1276 })?;
1277
1278 let expected_values_len = self.len
1279 .checked_mul(list_size)
1280 .expect("integer overflow computing expected number of expected values in FixedListSize");
1281
1282 if values_data.len < expected_values_len {
1283 return Err(ArrowError::InvalidArgumentError(format!(
1284 "Values length {} is less than the length ({}) multiplied by the value size ({}) for {}",
1285 values_data.len, self.len, list_size, self.data_type
1286 )));
1287 }
1288
1289 Ok(())
1290 }
1291 DataType::Struct(fields) => {
1292 self.validate_num_child_data(fields.len())?;
1293 let len_plus_offset =
1294 checked_len_plus_offset(&self.data_type, self.len, self.offset)?;
1295 for (i, field) in fields.iter().enumerate() {
1296 let field_data = self.get_valid_child_data(i, field.data_type())?;
1297
1298 if field_data.len < len_plus_offset {
1300 return Err(ArrowError::InvalidArgumentError(format!(
1301 "{} child array #{} for field {} has length smaller than expected for struct array ({} < {})",
1302 self.data_type,
1303 i,
1304 field.name(),
1305 field_data.len,
1306 len_plus_offset
1307 )));
1308 }
1309 }
1310 Ok(())
1311 }
1312 DataType::RunEndEncoded(run_ends_field, values_field) => {
1313 self.validate_num_child_data(2)?;
1314 let run_ends_data = self.get_valid_child_data(0, run_ends_field.data_type())?;
1315 let values_data = self.get_valid_child_data(1, values_field.data_type())?;
1316 if run_ends_data.len != values_data.len {
1317 return Err(ArrowError::InvalidArgumentError(format!(
1318 "The run_ends array length should be the same as values array length. Run_ends array length is {}, values array length is {}",
1319 run_ends_data.len, values_data.len
1320 )));
1321 }
1322 if run_ends_data.nulls.is_some() {
1323 return Err(ArrowError::InvalidArgumentError(
1324 "Found null values in run_ends array. The run_ends array should not have null values.".to_string(),
1325 ));
1326 }
1327 Ok(())
1328 }
1329 DataType::Union(fields, mode) => {
1330 self.validate_num_child_data(fields.len())?;
1331
1332 for (i, (_, field)) in fields.iter().enumerate() {
1333 let field_data = self.get_valid_child_data(i, field.data_type())?;
1334
1335 if mode == &UnionMode::Sparse {
1336 let len_plus_offset =
1337 checked_len_plus_offset(&self.data_type, self.len, self.offset)?;
1338 if field_data.len < len_plus_offset {
1339 return Err(ArrowError::InvalidArgumentError(format!(
1340 "Sparse union child array #{} has length smaller than expected for union array ({} < {})",
1341 i, field_data.len, len_plus_offset
1342 )));
1343 }
1344 }
1345 }
1346 Ok(())
1347 }
1348 DataType::Dictionary(_key_type, value_type) => {
1349 self.get_single_valid_child_data(value_type)?;
1350 Ok(())
1351 }
1352 _ => {
1353 if !self.child_data.is_empty() {
1355 return Err(ArrowError::InvalidArgumentError(format!(
1356 "Expected no child arrays for type {} but got {}",
1357 self.data_type,
1358 self.child_data.len()
1359 )));
1360 }
1361 Ok(())
1362 }
1363 }
1364 }
1365
1366 fn get_single_valid_child_data(
1370 &self,
1371 expected_type: &DataType,
1372 ) -> Result<&ArrayData, ArrowError> {
1373 self.validate_num_child_data(1)?;
1374 self.get_valid_child_data(0, expected_type)
1375 }
1376
1377 fn buffer_at(&self, idx: usize) -> Result<&Buffer, ArrowError> {
1382 self.buffers.get(idx).ok_or_else(|| {
1383 ArrowError::InvalidArgumentError(format!(
1384 "{} should contain at least {} buffer(s), had {}",
1385 self.data_type,
1386 idx + 1,
1387 self.buffers.len()
1388 ))
1389 })
1390 }
1391
1392 fn child_at(&self, idx: usize) -> Result<&ArrayData, ArrowError> {
1397 self.child_data.get(idx).ok_or_else(|| {
1398 ArrowError::InvalidArgumentError(format!(
1399 "{} should contain at least {} child data array(s), had {}",
1400 self.data_type,
1401 idx + 1,
1402 self.child_data.len()
1403 ))
1404 })
1405 }
1406
1407 fn validate_num_child_data(&self, expected_len: usize) -> Result<(), ArrowError> {
1409 if self.child_data.len() != expected_len {
1410 Err(ArrowError::InvalidArgumentError(format!(
1411 "Value data for {} should contain {} child data array(s), had {}",
1412 self.data_type,
1413 expected_len,
1414 self.child_data.len()
1415 )))
1416 } else {
1417 Ok(())
1418 }
1419 }
1420
1421 fn get_valid_child_data(
1424 &self,
1425 i: usize,
1426 expected_type: &DataType,
1427 ) -> Result<&ArrayData, ArrowError> {
1428 let values_data = self.child_data.get(i).ok_or_else(|| {
1429 ArrowError::InvalidArgumentError(format!(
1430 "{} did not have enough child arrays. Expected at least {} but had only {}",
1431 self.data_type,
1432 i + 1,
1433 self.child_data.len()
1434 ))
1435 })?;
1436
1437 if expected_type != &values_data.data_type {
1438 return Err(ArrowError::InvalidArgumentError(format!(
1439 "Child type mismatch for {}. Expected {} but child data had {}",
1440 self.data_type, expected_type, values_data.data_type
1441 )));
1442 }
1443
1444 values_data.validate()?;
1445 Ok(values_data)
1446 }
1447
1448 pub fn validate_data(&self) -> Result<(), ArrowError> {
1464 self.validate()?;
1465
1466 self.validate_nulls()?;
1467 self.validate_values()?;
1468 Ok(())
1469 }
1470
1471 pub fn validate_full(&self) -> Result<(), ArrowError> {
1476 self.validate_data()?;
1477 self.child_data
1479 .iter()
1480 .enumerate()
1481 .try_for_each(|(i, child_data)| {
1482 child_data.validate_full().map_err(|e| {
1483 ArrowError::InvalidArgumentError(format!(
1484 "{} child #{} invalid: {}",
1485 self.data_type, i, e
1486 ))
1487 })
1488 })?;
1489 Ok(())
1490 }
1491
1492 pub fn validate_nulls(&self) -> Result<(), ArrowError> {
1502 if let Some(nulls) = &self.nulls {
1503 let actual = nulls.len() - nulls.inner().count_set_bits();
1504 if actual != nulls.null_count() {
1505 return Err(ArrowError::InvalidArgumentError(format!(
1506 "null_count value ({}) doesn't match actual number of nulls in array ({})",
1507 nulls.null_count(),
1508 actual
1509 )));
1510 }
1511 }
1512
1513 match &self.data_type {
1518 DataType::List(f) | DataType::LargeList(f) | DataType::Map(f, _) => {
1519 if !f.is_nullable() {
1520 let child = &self.child_data[0];
1521 self.validate_non_nullable(None, child, child.nulls())?
1522 }
1523 }
1524 DataType::FixedSizeList(field, len) => {
1525 let child = &self.child_data[0];
1526 if !field.is_nullable() {
1527 match &self.nulls {
1528 Some(nulls) => {
1529 let element_len = *len as usize;
1530 let expanded = nulls.expand(element_len);
1531 self.validate_non_nullable(Some(&expanded), child, child.nulls())?;
1532 }
1533 None => self.validate_non_nullable(None, child, child.nulls())?,
1534 }
1535 }
1536 }
1537 DataType::Struct(fields) => {
1538 for (field, child) in fields.iter().zip(&self.child_data) {
1539 if !field.is_nullable() {
1540 let child_nulls = child
1541 .nulls()
1542 .map(|nulls| nulls.slice(self.offset, self.len));
1543 self.validate_non_nullable(self.nulls(), child, child_nulls.as_ref())?
1544 }
1545 }
1546 }
1547 _ => {}
1548 }
1549
1550 Ok(())
1551 }
1552
1553 fn validate_non_nullable(
1555 &self,
1556 mask: Option<&NullBuffer>,
1557 child: &ArrayData,
1558 child_nulls: Option<&NullBuffer>,
1559 ) -> Result<(), ArrowError> {
1560 let Some(mask) = mask else {
1561 return match child_nulls.map(NullBuffer::null_count).unwrap_or_default() {
1562 0 => Ok(()),
1563 _ => Err(ArrowError::InvalidArgumentError(format!(
1564 "non-nullable child of type {} contains nulls not present in parent {}",
1565 child.data_type, self.data_type
1566 ))),
1567 };
1568 };
1569
1570 match child_nulls {
1571 Some(nulls) if !mask.contains(nulls) => Err(ArrowError::InvalidArgumentError(format!(
1572 "non-nullable child of type {} contains nulls not present in parent",
1573 child.data_type
1574 ))),
1575 _ => Ok(()),
1576 }
1577 }
1578
1579 pub fn validate_values(&self) -> Result<(), ArrowError> {
1585 match &self.data_type {
1586 DataType::Utf8 => self.validate_utf8::<i32>(),
1587 DataType::LargeUtf8 => self.validate_utf8::<i64>(),
1588 DataType::Binary => self.validate_offsets_full::<i32>(self.buffer_at(1)?.len()),
1589 DataType::LargeBinary => self.validate_offsets_full::<i64>(self.buffer_at(1)?.len()),
1590 DataType::BinaryView => {
1591 let views = self.typed_buffer::<u128>(0, self.len)?;
1592 validate_binary_view(views, &self.buffers[1..])
1593 }
1594 DataType::Utf8View => {
1595 let views = self.typed_buffer::<u128>(0, self.len)?;
1596 validate_string_view(views, &self.buffers[1..])
1597 }
1598 DataType::List(_) | DataType::Map(_, _) => {
1599 let child = self.child_at(0)?;
1600 self.validate_offsets_full::<i32>(child.len)
1601 }
1602 DataType::LargeList(_) => {
1603 let child = self.child_at(0)?;
1604 self.validate_offsets_full::<i64>(child.len)
1605 }
1606 DataType::Union(_, _) => {
1607 Ok(())
1613 }
1614 DataType::Dictionary(key_type, _value_type) => {
1615 let dictionary_length = self.child_at(0)?.len;
1616 let dictionary_length = i64::try_from(dictionary_length).map_err(|_| {
1617 ArrowError::InvalidArgumentError(format!(
1618 "Dictionary of {dictionary_length} values is too long for an i64"
1619 ))
1620 })?;
1621 let max_value = dictionary_length - 1;
1622 match key_type.as_ref() {
1623 DataType::UInt8 => self.check_bounds::<u8>(max_value),
1624 DataType::UInt16 => self.check_bounds::<u16>(max_value),
1625 DataType::UInt32 => self.check_bounds::<u32>(max_value),
1626 DataType::UInt64 => self.check_bounds::<u64>(max_value),
1627 DataType::Int8 => self.check_bounds::<i8>(max_value),
1628 DataType::Int16 => self.check_bounds::<i16>(max_value),
1629 DataType::Int32 => self.check_bounds::<i32>(max_value),
1630 DataType::Int64 => self.check_bounds::<i64>(max_value),
1631 _ => Err(ArrowError::InvalidArgumentError(format!(
1632 "Dictionary key type must be an integer, got {key_type}"
1633 ))),
1634 }
1635 }
1636 DataType::RunEndEncoded(run_ends, _values) => {
1637 let run_ends_data = self.child_at(0)?;
1638 match run_ends.data_type() {
1639 DataType::Int16 => run_ends_data.check_run_ends::<i16>(),
1640 DataType::Int32 => run_ends_data.check_run_ends::<i32>(),
1641 DataType::Int64 => run_ends_data.check_run_ends::<i64>(),
1642 data_type => Err(ArrowError::InvalidArgumentError(format!(
1643 "Run end type must be Int16, Int32 or Int64, got {data_type}"
1644 ))),
1645 }
1646 }
1647 _ => {
1648 Ok(())
1650 }
1651 }
1652 }
1653
1654 fn validate_each_offset<T, V>(&self, offset_limit: usize, validate: V) -> Result<(), ArrowError>
1665 where
1666 T: ArrowNativeType + TryInto<usize> + num_traits::Num + std::fmt::Display,
1667 V: Fn(usize, Range<usize>) -> Result<(), ArrowError>,
1668 {
1669 self.typed_offsets::<T>()?
1670 .iter()
1671 .enumerate()
1672 .map(|(i, x)| {
1673 let r = x.to_usize().ok_or_else(|| {
1675 ArrowError::InvalidArgumentError(format!(
1676 "Offset invariant failure: Could not convert offset {x} to usize at position {i}"))}
1677 );
1678 match r {
1680 Ok(n) if n <= offset_limit => Ok((i, n)),
1681 Ok(_) => Err(ArrowError::InvalidArgumentError(format!(
1682 "Offset invariant failure: offset at position {i} out of bounds: {x} > {offset_limit}"))
1683 ),
1684 Err(e) => Err(e),
1685 }
1686 })
1687 .scan(0_usize, |start, end| {
1688 match end {
1690 Ok((i, end)) if *start <= end => {
1691 let range = Some(Ok((i, *start..end)));
1692 *start = end;
1693 range
1694 }
1695 Ok((i, end)) => Some(Err(ArrowError::InvalidArgumentError(format!(
1696 "Offset invariant failure: non-monotonic offset at slot {}: {} > {}",
1697 i - 1, start, end))
1698 )),
1699 Err(err) => Some(Err(err)),
1700 }
1701 })
1702 .skip(1) .try_for_each(|res: Result<(usize, Range<usize>), ArrowError>| {
1704 let (item_index, range) = res?;
1705 validate(item_index-1, range)
1706 })
1707 }
1708
1709 fn validate_utf8<T>(&self) -> Result<(), ArrowError>
1712 where
1713 T: ArrowNativeType + TryInto<usize> + num_traits::Num + std::fmt::Display,
1714 {
1715 let values_buffer = &self.buffer_at(1)?.as_slice();
1716 if let Ok(values_str) = std::str::from_utf8(values_buffer) {
1717 self.validate_each_offset::<T, _>(values_buffer.len(), |string_index, range| {
1719 if !values_str.is_char_boundary(range.start)
1720 || !values_str.is_char_boundary(range.end)
1721 {
1722 return Err(ArrowError::InvalidArgumentError(format!(
1723 "incomplete utf-8 byte sequence from index {string_index}"
1724 )));
1725 }
1726 Ok(())
1727 })
1728 } else {
1729 self.validate_each_offset::<T, _>(values_buffer.len(), |string_index, range| {
1731 std::str::from_utf8(&values_buffer[range.clone()]).map_err(|e| {
1732 ArrowError::InvalidArgumentError(format!(
1733 "Invalid UTF8 sequence at string index {string_index} ({range:?}): {e}"
1734 ))
1735 })?;
1736 Ok(())
1737 })
1738 }
1739 }
1740
1741 fn validate_offsets_full<T>(&self, offset_limit: usize) -> Result<(), ArrowError>
1744 where
1745 T: ArrowNativeType + TryInto<usize> + num_traits::Num + std::fmt::Display,
1746 {
1747 self.validate_each_offset::<T, _>(offset_limit, |_string_index, _range| {
1748 Ok(())
1751 })
1752 }
1753
1754 fn check_bounds<T>(&self, max_value: i64) -> Result<(), ArrowError>
1757 where
1758 T: ArrowNativeType + TryInto<i64> + num_traits::Num + std::fmt::Display,
1759 {
1760 let indexes: &[T] = self.typed_buffer::<T>(0, self.len)?;
1763
1764 indexes.iter().enumerate().try_for_each(|(i, &dict_index)| {
1765 if self.is_null(i) {
1767 return Ok(());
1768 }
1769 let dict_index: i64 = dict_index.try_into().map_err(|_| {
1770 ArrowError::InvalidArgumentError(format!(
1771 "Value at position {i} out of bounds: {dict_index} (can not convert to i64)"
1772 ))
1773 })?;
1774
1775 if dict_index < 0 || dict_index > max_value {
1776 return Err(ArrowError::InvalidArgumentError(format!(
1777 "Value at position {i} out of bounds: {dict_index} (should be in [0, {max_value}])"
1778 )));
1779 }
1780 Ok(())
1781 })
1782 }
1783
1784 fn check_run_ends<T>(&self) -> Result<(), ArrowError>
1786 where
1787 T: ArrowNativeType + TryInto<i64> + num_traits::Num + std::fmt::Display,
1788 {
1789 let values = self.typed_buffer::<T>(0, self.len)?;
1790 let mut prev_value = 0_i64;
1791 values.iter().enumerate().try_for_each(|(ix, &inp_value)| {
1792 let value: i64 = inp_value.try_into().map_err(|_| {
1793 ArrowError::InvalidArgumentError(format!(
1794 "Value at position {ix} out of bounds: {inp_value} (can not convert to i64)"
1795 ))
1796 })?;
1797 if value <= 0_i64 {
1798 return Err(ArrowError::InvalidArgumentError(format!(
1799 "The values in run_ends array should be strictly positive. Found value {value} at index {ix} that does not match the criteria."
1800 )));
1801 }
1802 if ix > 0 && value <= prev_value {
1803 return Err(ArrowError::InvalidArgumentError(format!(
1804 "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."
1805 )));
1806 }
1807
1808 prev_value = value;
1809 Ok(())
1810 })?;
1811
1812 let len_plus_offset = checked_len_plus_offset(&self.data_type, self.len, self.offset)?;
1813 if prev_value.as_usize() < len_plus_offset {
1814 return Err(ArrowError::InvalidArgumentError(format!(
1815 "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}."
1816 )));
1817 }
1818 Ok(())
1819 }
1820
1821 pub fn ptr_eq(&self, other: &Self) -> bool {
1825 if self.offset != other.offset
1826 || self.len != other.len
1827 || self.data_type != other.data_type
1828 || self.buffers.len() != other.buffers.len()
1829 || self.child_data.len() != other.child_data.len()
1830 {
1831 return false;
1832 }
1833
1834 match (&self.nulls, &other.nulls) {
1835 (Some(a), Some(b)) if !a.inner().ptr_eq(b.inner()) => return false,
1836 (Some(_), None) | (None, Some(_)) => return false,
1837 _ => {}
1838 }
1839
1840 if !self
1841 .buffers
1842 .iter()
1843 .zip(other.buffers.iter())
1844 .all(|(a, b)| a.as_ptr() == b.as_ptr())
1845 {
1846 return false;
1847 }
1848
1849 self.child_data
1850 .iter()
1851 .zip(other.child_data.iter())
1852 .all(|(a, b)| a.ptr_eq(b))
1853 }
1854
1855 pub fn into_builder(self) -> ArrayDataBuilder {
1857 self.into()
1858 }
1859
1860 #[cfg(feature = "pool")]
1867 pub fn claim(&self, pool: &dyn arrow_buffer::MemoryPool) {
1868 for buffer in &self.buffers {
1870 buffer.claim(pool);
1871 }
1872
1873 if let Some(nulls) = &self.nulls {
1875 nulls.claim(pool);
1876 }
1877
1878 for child in &self.child_data {
1880 child.claim(pool);
1881 }
1882 }
1883}
1884
1885pub fn layout(data_type: &DataType) -> DataTypeLayout {
1888 use arrow_schema::IntervalUnit::*;
1891
1892 match data_type {
1893 DataType::Null => DataTypeLayout {
1894 buffers: vec![],
1895 can_contain_null_mask: false,
1896 variadic: false,
1897 },
1898 DataType::Boolean => DataTypeLayout {
1899 buffers: vec![BufferSpec::BitMap],
1900 can_contain_null_mask: true,
1901 variadic: false,
1902 },
1903 DataType::Int8 => DataTypeLayout::new_fixed_width::<i8>(),
1904 DataType::Int16 => DataTypeLayout::new_fixed_width::<i16>(),
1905 DataType::Int32 => DataTypeLayout::new_fixed_width::<i32>(),
1906 DataType::Int64 => DataTypeLayout::new_fixed_width::<i64>(),
1907 DataType::UInt8 => DataTypeLayout::new_fixed_width::<u8>(),
1908 DataType::UInt16 => DataTypeLayout::new_fixed_width::<u16>(),
1909 DataType::UInt32 => DataTypeLayout::new_fixed_width::<u32>(),
1910 DataType::UInt64 => DataTypeLayout::new_fixed_width::<u64>(),
1911 DataType::Float16 => DataTypeLayout::new_fixed_width::<half::f16>(),
1912 DataType::Float32 => DataTypeLayout::new_fixed_width::<f32>(),
1913 DataType::Float64 => DataTypeLayout::new_fixed_width::<f64>(),
1914 DataType::Timestamp(_, _) => DataTypeLayout::new_fixed_width::<i64>(),
1915 DataType::Date32 => DataTypeLayout::new_fixed_width::<i32>(),
1916 DataType::Date64 => DataTypeLayout::new_fixed_width::<i64>(),
1917 DataType::Time32(_) => DataTypeLayout::new_fixed_width::<i32>(),
1918 DataType::Time64(_) => DataTypeLayout::new_fixed_width::<i64>(),
1919 DataType::Interval(YearMonth) => DataTypeLayout::new_fixed_width::<i32>(),
1920 DataType::Interval(DayTime) => DataTypeLayout::new_fixed_width::<IntervalDayTime>(),
1921 DataType::Interval(MonthDayNano) => {
1922 DataTypeLayout::new_fixed_width::<IntervalMonthDayNano>()
1923 }
1924 DataType::Duration(_) => DataTypeLayout::new_fixed_width::<i64>(),
1925 DataType::Decimal32(_, _) => DataTypeLayout::new_fixed_width::<i32>(),
1926 DataType::Decimal64(_, _) => DataTypeLayout::new_fixed_width::<i64>(),
1927 DataType::Decimal128(_, _) => DataTypeLayout::new_fixed_width::<i128>(),
1928 DataType::Decimal256(_, _) => DataTypeLayout::new_fixed_width::<i256>(),
1929 DataType::FixedSizeBinary(size) => {
1930 let spec = BufferSpec::FixedWidth {
1931 byte_width: (*size).try_into().unwrap(),
1932 alignment: mem::align_of::<u8>(),
1933 };
1934 DataTypeLayout {
1935 buffers: vec![spec],
1936 can_contain_null_mask: true,
1937 variadic: false,
1938 }
1939 }
1940 DataType::Binary => DataTypeLayout::new_binary::<i32>(),
1941 DataType::LargeBinary => DataTypeLayout::new_binary::<i64>(),
1942 DataType::Utf8 => DataTypeLayout::new_binary::<i32>(),
1943 DataType::LargeUtf8 => DataTypeLayout::new_binary::<i64>(),
1944 DataType::BinaryView | DataType::Utf8View => DataTypeLayout::new_view(),
1945 DataType::FixedSizeList(_, _) => DataTypeLayout::new_nullable_empty(), DataType::List(_) => DataTypeLayout::new_fixed_width::<i32>(),
1947 DataType::ListView(_) => DataTypeLayout::new_list_view::<i32>(),
1948 DataType::LargeListView(_) => DataTypeLayout::new_list_view::<i64>(),
1949 DataType::LargeList(_) => DataTypeLayout::new_fixed_width::<i64>(),
1950 DataType::Map(_, _) => DataTypeLayout::new_fixed_width::<i32>(),
1951 DataType::Struct(_) => DataTypeLayout::new_nullable_empty(), DataType::RunEndEncoded(_, _) => DataTypeLayout::new_empty(), DataType::Union(_, mode) => {
1954 let type_ids = BufferSpec::FixedWidth {
1955 byte_width: mem::size_of::<i8>(),
1956 alignment: mem::align_of::<i8>(),
1957 };
1958
1959 DataTypeLayout {
1960 buffers: match mode {
1961 UnionMode::Sparse => {
1962 vec![type_ids]
1963 }
1964 UnionMode::Dense => {
1965 vec![
1966 type_ids,
1967 BufferSpec::FixedWidth {
1968 byte_width: mem::size_of::<i32>(),
1969 alignment: mem::align_of::<i32>(),
1970 },
1971 ]
1972 }
1973 },
1974 can_contain_null_mask: false,
1975 variadic: false,
1976 }
1977 }
1978 DataType::Dictionary(key_type, _value_type) => layout(key_type),
1979 }
1980}
1981
1982#[derive(Debug, PartialEq, Eq)]
1984pub struct DataTypeLayout {
1986 pub buffers: Vec<BufferSpec>,
1988
1989 pub can_contain_null_mask: bool,
1991
1992 pub variadic: bool,
1996}
1997
1998impl DataTypeLayout {
1999 pub fn new_fixed_width<T>() -> Self {
2001 Self {
2002 buffers: vec![BufferSpec::FixedWidth {
2003 byte_width: mem::size_of::<T>(),
2004 alignment: mem::align_of::<T>(),
2005 }],
2006 can_contain_null_mask: true,
2007 variadic: false,
2008 }
2009 }
2010
2011 pub fn new_nullable_empty() -> Self {
2014 Self {
2015 buffers: vec![],
2016 can_contain_null_mask: true,
2017 variadic: false,
2018 }
2019 }
2020
2021 pub fn new_empty() -> Self {
2024 Self {
2025 buffers: vec![],
2026 can_contain_null_mask: false,
2027 variadic: false,
2028 }
2029 }
2030
2031 pub fn new_binary<T>() -> Self {
2035 Self {
2036 buffers: vec![
2037 BufferSpec::FixedWidth {
2039 byte_width: mem::size_of::<T>(),
2040 alignment: mem::align_of::<T>(),
2041 },
2042 BufferSpec::VariableWidth,
2044 ],
2045 can_contain_null_mask: true,
2046 variadic: false,
2047 }
2048 }
2049
2050 pub fn new_view() -> Self {
2052 Self {
2053 buffers: vec![BufferSpec::FixedWidth {
2054 byte_width: mem::size_of::<u128>(),
2055 alignment: mem::align_of::<u128>(),
2056 }],
2057 can_contain_null_mask: true,
2058 variadic: true,
2059 }
2060 }
2061
2062 pub fn new_list_view<T>() -> Self {
2064 Self {
2065 buffers: vec![
2066 BufferSpec::FixedWidth {
2067 byte_width: mem::size_of::<T>(),
2068 alignment: mem::align_of::<T>(),
2069 },
2070 BufferSpec::FixedWidth {
2071 byte_width: mem::size_of::<T>(),
2072 alignment: mem::align_of::<T>(),
2073 },
2074 ],
2075 can_contain_null_mask: true,
2076 variadic: false,
2077 }
2078 }
2079}
2080
2081#[derive(Debug, PartialEq, Eq)]
2083pub enum BufferSpec {
2084 FixedWidth {
2095 byte_width: usize,
2097 alignment: usize,
2099 },
2100 VariableWidth,
2102 BitMap,
2108 AlwaysNull,
2111}
2112
2113impl PartialEq for ArrayData {
2114 fn eq(&self, other: &Self) -> bool {
2115 equal::equal(self, other)
2116 }
2117}
2118
2119#[derive(Debug, Clone)]
2138#[doc(hidden)]
2139pub struct UnsafeFlag(bool);
2140
2141impl UnsafeFlag {
2142 #[inline]
2146 pub const fn new() -> Self {
2147 Self(false)
2148 }
2149
2150 #[inline]
2160 pub unsafe fn set(&mut self, val: bool) {
2161 self.0 = val;
2162 }
2163
2164 #[inline]
2166 pub fn get(&self) -> bool {
2167 self.0
2168 }
2169}
2170
2171impl Default for UnsafeFlag {
2173 fn default() -> Self {
2174 Self::new()
2175 }
2176}
2177
2178#[derive(Debug)]
2180pub struct ArrayDataBuilder {
2181 data_type: DataType,
2182 len: usize,
2183 null_count: Option<usize>,
2184 null_bit_buffer: Option<Buffer>,
2185 nulls: Option<NullBuffer>,
2186 offset: usize,
2187 buffers: Vec<Buffer>,
2188 child_data: Vec<ArrayData>,
2189 align_buffers: bool,
2193 skip_validation: UnsafeFlag,
2203}
2204
2205impl ArrayDataBuilder {
2206 #[inline]
2207 pub const fn new(data_type: DataType) -> Self {
2209 Self {
2210 data_type,
2211 len: 0,
2212 null_count: None,
2213 null_bit_buffer: None,
2214 nulls: None,
2215 offset: 0,
2216 buffers: vec![],
2217 child_data: vec![],
2218 align_buffers: false,
2219 skip_validation: UnsafeFlag::new(),
2220 }
2221 }
2222
2223 pub fn data_type(self, data_type: DataType) -> Self {
2225 Self { data_type, ..self }
2226 }
2227
2228 #[inline]
2229 pub const fn len(mut self, n: usize) -> Self {
2231 self.len = n;
2232 self
2233 }
2234
2235 pub fn nulls(mut self, nulls: Option<NullBuffer>) -> Self {
2237 self.nulls = nulls;
2238 self.null_count = None;
2239 self.null_bit_buffer = None;
2240 self
2241 }
2242
2243 pub fn null_count(mut self, null_count: usize) -> Self {
2245 self.null_count = Some(null_count);
2246 self
2247 }
2248
2249 pub fn null_bit_buffer(mut self, buf: Option<Buffer>) -> Self {
2251 self.nulls = None;
2252 self.null_bit_buffer = buf;
2253 self
2254 }
2255
2256 #[inline]
2258 pub const fn offset(mut self, n: usize) -> Self {
2259 self.offset = n;
2260 self
2261 }
2262
2263 pub fn buffers(mut self, v: Vec<Buffer>) -> Self {
2265 self.buffers = v;
2266 self
2267 }
2268
2269 pub fn add_buffer(mut self, b: Buffer) -> Self {
2271 self.buffers.push(b);
2272 self
2273 }
2274
2275 pub fn add_buffers<I: IntoIterator<Item = Buffer>>(mut self, bs: I) -> Self {
2277 self.buffers.extend(bs);
2278 self
2279 }
2280
2281 pub fn child_data(mut self, v: Vec<ArrayData>) -> Self {
2283 self.child_data = v;
2284 self
2285 }
2286
2287 pub fn add_child_data(mut self, r: ArrayData) -> Self {
2289 self.child_data.push(r);
2290 self
2291 }
2292
2293 pub unsafe fn build_unchecked(self) -> ArrayData {
2309 unsafe { self.skip_validation(true) }.build().unwrap()
2310 }
2311
2312 pub fn build(self) -> Result<ArrayData, ArrowError> {
2321 let Self {
2322 data_type,
2323 len,
2324 null_count,
2325 null_bit_buffer,
2326 nulls,
2327 offset,
2328 buffers,
2329 child_data,
2330 align_buffers,
2331 skip_validation,
2332 } = self;
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 !skip_validation.get() || cfg!(feature = "force_validate") {
2363 data.validate_data()?;
2364 }
2365 Ok(data)
2366 }
2367
2368 pub fn align_buffers(mut self, align_buffers: bool) -> Self {
2384 self.align_buffers = align_buffers;
2385 self
2386 }
2387
2388 pub unsafe fn skip_validation(mut self, skip_validation: bool) -> Self {
2402 unsafe {
2403 self.skip_validation.set(skip_validation);
2404 }
2405 self
2406 }
2407}
2408
2409impl From<ArrayData> for ArrayDataBuilder {
2410 fn from(d: ArrayData) -> Self {
2411 Self {
2412 data_type: d.data_type,
2413 len: d.len,
2414 offset: d.offset,
2415 buffers: d.buffers,
2416 child_data: d.child_data,
2417 nulls: d.nulls,
2418 null_bit_buffer: None,
2419 null_count: None,
2420 align_buffers: false,
2421 skip_validation: UnsafeFlag::new(),
2422 }
2423 }
2424}
2425
2426pub(crate) fn get_fixed_size_binary_width(data_type: &DataType) -> usize {
2431 match data_type {
2432 DataType::FixedSizeBinary(i) => {
2433 if *i < 0 {
2434 panic!("cannot compare FixedSizeBinary({})", *i);
2435 }
2436 *i as usize
2437 }
2438 _ => unreachable!(),
2439 }
2440}
2441
2442#[cfg(test)]
2443mod tests {
2444 use super::*;
2445 use crate::ByteView;
2446 use crate::transform::MutableArrayData;
2447 use arrow_buffer::{OffsetBuffer, ScalarBuffer};
2448 use arrow_schema::{Field, Fields};
2449
2450 fn make_i32_buffer(n: usize) -> Buffer {
2454 Buffer::from_slice_ref(vec![42i32; n])
2455 }
2456
2457 fn make_f32_buffer(n: usize) -> Buffer {
2459 Buffer::from_slice_ref(vec![42f32; n])
2460 }
2461
2462 #[test]
2463 fn test_builder() {
2464 let v = (0..25).collect::<Vec<i32>>();
2466 let b1 = Buffer::from_slice_ref(&v);
2467 let arr_data = ArrayData::builder(DataType::Int32)
2468 .len(20)
2469 .offset(5)
2470 .add_buffer(b1)
2471 .null_bit_buffer(Some(Buffer::from([
2472 0b01011111, 0b10110101, 0b01100011, 0b00011110,
2473 ])))
2474 .build()
2475 .unwrap();
2476
2477 assert_eq!(20, arr_data.len());
2478 assert_eq!(10, arr_data.null_count());
2479 assert_eq!(5, arr_data.offset());
2480 assert_eq!(1, arr_data.buffers().len());
2481 assert_eq!(
2482 Buffer::from_slice_ref(&v).as_slice(),
2483 arr_data.buffers()[0].as_slice()
2484 );
2485 }
2486
2487 #[test]
2488 fn test_builder_with_child_data() {
2489 let child_arr_data = ArrayData::try_new(
2490 DataType::Int32,
2491 5,
2492 None,
2493 0,
2494 vec![Buffer::from_slice_ref([1i32, 2, 3, 4, 5])],
2495 vec![],
2496 )
2497 .unwrap();
2498
2499 let field = Arc::new(Field::new("x", DataType::Int32, true));
2500 let data_type = DataType::Struct(vec![field].into());
2501
2502 let arr_data = ArrayData::builder(data_type)
2503 .len(5)
2504 .offset(0)
2505 .add_child_data(child_arr_data.clone())
2506 .build()
2507 .unwrap();
2508
2509 assert_eq!(5, arr_data.len());
2510 assert_eq!(1, arr_data.child_data().len());
2511 assert_eq!(child_arr_data, arr_data.child_data()[0]);
2512 }
2513
2514 #[test]
2515 fn test_struct_validation_accounts_for_parent_offset() {
2516 let data_type =
2517 DataType::Struct(Fields::from(vec![Field::new("x", DataType::Int32, false)]));
2518 let child = ArrayData::builder(DataType::Int32)
2519 .len(5)
2520 .add_buffer(Buffer::from_slice_ref([0, 1, 2, 3, 4]))
2521 .build()
2522 .unwrap();
2523
2524 let err = ArrayData::builder(data_type)
2526 .len(5)
2527 .offset(1)
2528 .add_child_data(child)
2529 .build()
2530 .unwrap_err()
2531 .to_string();
2532
2533 assert!(err.contains(
2534 "child array #0 for field x has length smaller than expected for struct array (5 < 6)"
2535 ));
2536 }
2537
2538 #[test]
2539 fn test_struct_non_nullable_child_nulls_account_for_parent_offset() {
2540 let build = |parent_nulls| {
2541 let child = ArrayData::builder(DataType::Int32)
2542 .len(5)
2543 .add_buffer(Buffer::from_slice_ref([0, 1, 2, 3, 4]))
2544 .nulls(Some(NullBuffer::new(BooleanBuffer::from(vec![
2545 true, true, false, true, true,
2546 ]))))
2547 .build()
2548 .unwrap();
2549
2550 ArrayData::builder(DataType::Struct(Fields::from(vec![Field::new(
2551 "x",
2552 DataType::Int32,
2553 false,
2554 )])))
2555 .len(4)
2556 .offset(1)
2557 .nulls(Some(NullBuffer::new(BooleanBuffer::from(parent_nulls))))
2558 .add_child_data(child)
2559 .build()
2560 };
2561
2562 assert!(build(vec![true, false, true, true]).is_ok());
2563 assert!(build(vec![true, true, false, true]).is_err());
2564 }
2565
2566 #[test]
2567 fn test_struct_equal_accounts_for_parent_offset() {
2568 let data_type =
2569 DataType::Struct(Fields::from(vec![Field::new("x", DataType::Int32, false)]));
2570
2571 let child1 = ArrayData::builder(DataType::Int32)
2572 .len(5)
2573 .add_buffer(Buffer::from_slice_ref([0, 1, 2, 3, 4]))
2574 .build()
2575 .unwrap();
2576 let child2 = child1.slice(1, 4);
2577
2578 let data1 = ArrayData::builder(data_type.clone())
2580 .len(4)
2581 .offset(1)
2582 .add_child_data(child1)
2583 .build()
2584 .unwrap();
2585 let data2 = ArrayData::builder(data_type)
2586 .len(4)
2587 .add_child_data(child2)
2588 .build()
2589 .unwrap();
2590
2591 assert_eq!(data1, data2);
2592 }
2593
2594 #[test]
2595 fn test_extend_struct_accounts_for_parent_offset() {
2596 let data_type =
2597 DataType::Struct(Fields::from(vec![Field::new("x", DataType::Int32, false)]));
2598 let child = ArrayData::builder(DataType::Int32)
2599 .len(5)
2600 .add_buffer(Buffer::from_slice_ref([0, 1, 2, 3, 4]))
2601 .build()
2602 .unwrap();
2603
2604 let data = ArrayData::builder(data_type)
2605 .len(4)
2606 .offset(1)
2607 .add_child_data(child)
2608 .build()
2609 .unwrap();
2610
2611 let mut mutable = MutableArrayData::new(vec![&data], false, data.len());
2612 mutable.try_extend(0, 0, data.len()).unwrap();
2613 let output = mutable.freeze();
2614
2615 assert_eq!(output.child_data()[0].buffer::<i32>(0), &[1, 2, 3, 4]);
2616 }
2617
2618 #[test]
2619 fn test_null_count() {
2620 let mut bit_v: [u8; 2] = [0; 2];
2621 bit_util::set_bit(&mut bit_v, 0);
2622 bit_util::set_bit(&mut bit_v, 3);
2623 bit_util::set_bit(&mut bit_v, 10);
2624 let arr_data = ArrayData::builder(DataType::Int32)
2625 .len(16)
2626 .add_buffer(make_i32_buffer(16))
2627 .null_bit_buffer(Some(Buffer::from(bit_v)))
2628 .build()
2629 .unwrap();
2630 assert_eq!(13, arr_data.null_count());
2631
2632 let mut bit_v: [u8; 2] = [0; 2];
2634 bit_util::set_bit(&mut bit_v, 0);
2635 bit_util::set_bit(&mut bit_v, 3);
2636 bit_util::set_bit(&mut bit_v, 10);
2637 let arr_data = ArrayData::builder(DataType::Int32)
2638 .len(12)
2639 .offset(2)
2640 .add_buffer(make_i32_buffer(14)) .null_bit_buffer(Some(Buffer::from(bit_v)))
2642 .build()
2643 .unwrap();
2644 assert_eq!(10, arr_data.null_count());
2645 }
2646
2647 #[test]
2648 fn test_null_buffer_ref() {
2649 let mut bit_v: [u8; 2] = [0; 2];
2650 bit_util::set_bit(&mut bit_v, 0);
2651 bit_util::set_bit(&mut bit_v, 3);
2652 bit_util::set_bit(&mut bit_v, 10);
2653 let arr_data = ArrayData::builder(DataType::Int32)
2654 .len(16)
2655 .add_buffer(make_i32_buffer(16))
2656 .null_bit_buffer(Some(Buffer::from(bit_v)))
2657 .build()
2658 .unwrap();
2659 assert!(arr_data.nulls().is_some());
2660 assert_eq!(&bit_v, arr_data.nulls().unwrap().validity());
2661 }
2662
2663 #[test]
2664 fn test_slice() {
2665 let mut bit_v: [u8; 2] = [0; 2];
2666 bit_util::set_bit(&mut bit_v, 0);
2667 bit_util::set_bit(&mut bit_v, 3);
2668 bit_util::set_bit(&mut bit_v, 10);
2669 let data = ArrayData::builder(DataType::Int32)
2670 .len(16)
2671 .add_buffer(make_i32_buffer(16))
2672 .null_bit_buffer(Some(Buffer::from(bit_v)))
2673 .build()
2674 .unwrap();
2675 let new_data = data.slice(1, 15);
2676 assert_eq!(data.len() - 1, new_data.len());
2677 assert_eq!(1, new_data.offset());
2678 assert_eq!(data.null_count(), new_data.null_count());
2679
2680 let new_data = new_data.slice(1, 14);
2682 assert_eq!(data.len() - 2, new_data.len());
2683 assert_eq!(2, new_data.offset());
2684 assert_eq!(data.null_count() - 1, new_data.null_count());
2685 }
2686
2687 #[test]
2688 #[should_panic(expected = "offset + length overflow")]
2689 fn test_slice_panics_on_offset_length_overflow() {
2690 let data = ArrayData::builder(DataType::Int32)
2691 .len(4)
2692 .add_buffer(make_i32_buffer(4))
2693 .build()
2694 .unwrap();
2695 let sliced = data.slice(1, 3);
2696
2697 sliced.slice(1, usize::MAX);
2698 }
2699
2700 #[test]
2701 fn test_typed_offsets_length_overflow() {
2702 let data = ArrayData {
2703 data_type: DataType::Binary,
2704 len: usize::MAX,
2705 offset: 0,
2706 buffers: vec![Buffer::from_slice_ref([0_i32])],
2707 child_data: vec![],
2708 nulls: None,
2709 };
2710 let err = data.typed_offsets::<i32>().unwrap_err();
2711
2712 assert_eq!(
2713 err.to_string(),
2714 format!(
2715 "Invalid argument error: Length {} with offset 1 overflows usize for Binary",
2716 usize::MAX
2717 )
2718 );
2719 }
2720
2721 #[test]
2722 fn test_validate_typed_buffer_length_overflow() {
2723 let data = ArrayData {
2724 data_type: DataType::Binary,
2725 len: 0,
2726 offset: 2,
2727 buffers: vec![Buffer::from_slice_ref([0_i32])],
2728 child_data: vec![],
2729 nulls: None,
2730 };
2731 let err = data.typed_buffer::<i32>(0, usize::MAX).unwrap_err();
2732
2733 assert_eq!(
2734 err.to_string(),
2735 format!(
2736 "Invalid argument error: Length {} with offset 2 overflows usize for Binary",
2737 usize::MAX
2738 )
2739 );
2740 }
2741
2742 fn try_new_binary_length_offset_overflow() -> Result<ArrayData, ArrowError> {
2744 ArrayData::try_new(
2745 DataType::Binary,
2746 usize::MAX,
2747 None,
2748 1,
2749 vec![
2750 Buffer::from_slice_ref([0_i32]),
2751 Buffer::from_iter(std::iter::empty::<u8>()),
2752 ],
2753 vec![],
2754 )
2755 }
2756
2757 #[cfg(not(feature = "force_validate"))]
2758 #[test]
2759 fn test_try_new_length_offset_overflow() {
2760 let err = try_new_binary_length_offset_overflow().unwrap_err();
2761
2762 assert_eq!(
2763 err.to_string(),
2764 format!(
2765 "Invalid argument error: Length {} with offset 1 overflows usize for Binary",
2766 usize::MAX
2767 )
2768 );
2769 }
2770
2771 #[cfg(feature = "force_validate")]
2772 #[test]
2773 #[should_panic(
2774 expected = "Length 18446744073709551615 with offset 1 overflows usize for Binary"
2775 )]
2776 fn test_try_new_length_offset_overflow_force_validate() {
2777 try_new_binary_length_offset_overflow().unwrap();
2778 }
2779
2780 #[test]
2781 fn test_equality() {
2782 let int_data = ArrayData::builder(DataType::Int32)
2783 .len(1)
2784 .add_buffer(make_i32_buffer(1))
2785 .build()
2786 .unwrap();
2787
2788 let float_data = ArrayData::builder(DataType::Float32)
2789 .len(1)
2790 .add_buffer(make_f32_buffer(1))
2791 .build()
2792 .unwrap();
2793 assert_ne!(int_data, float_data);
2794 assert!(!int_data.ptr_eq(&float_data));
2795 assert!(int_data.ptr_eq(&int_data));
2796
2797 let int_data_clone = int_data.clone();
2798 assert_eq!(int_data, int_data_clone);
2799 assert!(int_data.ptr_eq(&int_data_clone));
2800 assert!(int_data_clone.ptr_eq(&int_data));
2801
2802 let int_data_slice = int_data_clone.slice(1, 0);
2803 assert!(int_data_slice.ptr_eq(&int_data_slice));
2804 assert!(!int_data.ptr_eq(&int_data_slice));
2805 assert!(!int_data_slice.ptr_eq(&int_data));
2806
2807 let data_buffer = Buffer::from_slice_ref(b"abcdef");
2808 let offsets_buffer = Buffer::from_slice_ref([0_i32, 2_i32, 2_i32, 5_i32]);
2809 let string_data = ArrayData::try_new(
2810 DataType::Utf8,
2811 3,
2812 Some(Buffer::from_iter(vec![true, false, true])),
2813 0,
2814 vec![offsets_buffer, data_buffer],
2815 vec![],
2816 )
2817 .unwrap();
2818
2819 assert_ne!(float_data, string_data);
2820 assert!(!float_data.ptr_eq(&string_data));
2821
2822 assert!(string_data.ptr_eq(&string_data));
2823
2824 let string_data_cloned = string_data.clone();
2825 assert!(string_data_cloned.ptr_eq(&string_data));
2826 assert!(string_data.ptr_eq(&string_data_cloned));
2827
2828 let string_data_slice = string_data.slice(1, 2);
2829 assert!(string_data_slice.ptr_eq(&string_data_slice));
2830 assert!(!string_data_slice.ptr_eq(&string_data))
2831 }
2832
2833 #[test]
2834 fn test_slice_memory_size_view_payload_buffers() {
2835 for data_type in [DataType::Utf8View, DataType::BinaryView] {
2836 let inline_only = ArrayData::builder(data_type.clone())
2837 .len(2)
2838 .add_buffer(Buffer::from_vec(vec![0_u128; 2]))
2839 .build()
2840 .unwrap();
2841 assert_eq!(
2842 inline_only.get_slice_memory_size().unwrap(),
2843 2 * mem::size_of::<u128>()
2844 );
2845
2846 let mut first_payload = Vec::with_capacity(32);
2847 first_payload.extend_from_slice(b"first payload");
2848 let first_view =
2849 ByteView::new(first_payload.len().try_into().unwrap(), &first_payload[..4])
2850 .as_u128();
2851 let first_payload = Buffer::from_vec(first_payload);
2852 assert!(first_payload.capacity() > first_payload.len());
2853 let first_payload_capacity = first_payload.capacity();
2854
2855 let mut second_payload = Vec::with_capacity(64);
2856 second_payload.extend_from_slice(b"second payload");
2857 let second_view = ByteView::new(
2858 second_payload.len().try_into().unwrap(),
2859 &second_payload[..4],
2860 )
2861 .with_buffer_index(1)
2862 .as_u128();
2863 let second_payload = Buffer::from_vec(second_payload);
2864 assert!(second_payload.capacity() > second_payload.len());
2865 let second_payload_capacity = second_payload.capacity();
2866
2867 let data = ArrayData::builder(data_type)
2868 .len(3)
2869 .add_buffer(Buffer::from_vec(vec![first_view, 0_u128, second_view]))
2870 .add_buffer(first_payload)
2871 .add_buffer(second_payload)
2872 .build()
2873 .unwrap();
2874 let sliced = data.slice(1, 1);
2875
2876 assert_eq!(
2877 sliced.get_slice_memory_size().unwrap(),
2878 mem::size_of::<u128>() + first_payload_capacity + second_payload_capacity
2879 );
2880 }
2881 }
2882
2883 #[test]
2884 fn test_slice_memory_size_utf8_offset_buffer_len_plus_one() {
2885 let data_buffer = Buffer::from_slice_ref(b"helloworld");
2887 let offsets_buffer = Buffer::from_slice_ref([0_i32, 5_i32, 10_i32]);
2890 let array = ArrayData::try_new(
2891 DataType::Utf8,
2892 2,
2893 None,
2894 0,
2895 vec![offsets_buffer, data_buffer],
2896 vec![],
2897 )
2898 .unwrap();
2899 assert_eq!(array.get_slice_memory_size().unwrap(), 22); }
2901
2902 #[test]
2903 fn test_slice_memory_size_binary_offset_buffer_len_plus_one() {
2904 let data_buffer = Buffer::from_slice_ref([0u8, 1, 2, 3, 4]);
2907 let offsets_buffer = Buffer::from_slice_ref([0_i32, 2_i32, 5_i32]);
2909 let array = ArrayData::try_new(
2910 DataType::Binary,
2911 2,
2912 None,
2913 0,
2914 vec![offsets_buffer, data_buffer],
2915 vec![],
2916 )
2917 .unwrap();
2918 assert_eq!(array.get_slice_memory_size().unwrap(), 17); }
2920
2921 #[test]
2922 fn test_slice_memory_size() {
2923 let mut bit_v: [u8; 2] = [0; 2];
2924 bit_util::set_bit(&mut bit_v, 0);
2925 bit_util::set_bit(&mut bit_v, 3);
2926 bit_util::set_bit(&mut bit_v, 10);
2927 let data = ArrayData::builder(DataType::Int32)
2928 .len(16)
2929 .add_buffer(make_i32_buffer(16))
2930 .null_bit_buffer(Some(Buffer::from(bit_v)))
2931 .build()
2932 .unwrap();
2933 let new_data = data.slice(1, 14);
2934 assert_eq!(
2935 data.get_slice_memory_size().unwrap() - 8,
2936 new_data.get_slice_memory_size().unwrap()
2937 );
2938 let data_buffer = Buffer::from_slice_ref(b"abcdef");
2939 let offsets_buffer = Buffer::from_slice_ref([0_i32, 2_i32, 2_i32, 5_i32]);
2940 let string_data = ArrayData::try_new(
2941 DataType::Utf8,
2942 3,
2943 Some(Buffer::from_iter(vec![true, false, true])),
2944 0,
2945 vec![offsets_buffer, data_buffer],
2946 vec![],
2947 )
2948 .unwrap();
2949 let string_data_slice = string_data.slice(1, 2);
2950 assert_eq!(
2952 string_data.get_slice_memory_size().unwrap() - 6,
2953 string_data_slice.get_slice_memory_size().unwrap()
2954 );
2955 }
2956
2957 #[test]
2958 fn test_count_nulls() {
2959 let buffer = Buffer::from([0b00010110, 0b10011111]);
2960 let buffer = NullBuffer::new(BooleanBuffer::new(buffer, 0, 16));
2961 let count = count_nulls(Some(&buffer), 0, 16);
2962 assert_eq!(count, 7);
2963
2964 let count = count_nulls(Some(&buffer), 4, 8);
2965 assert_eq!(count, 3);
2966 }
2967
2968 #[test]
2969 fn test_contains_nulls() {
2970 let buffer: Buffer =
2971 MutableBuffer::from_iter([false, false, false, true, true, false]).into();
2972 let buffer = NullBuffer::new(BooleanBuffer::new(buffer, 0, 6));
2973 assert!(contains_nulls(Some(&buffer), 0, 6));
2974 assert!(contains_nulls(Some(&buffer), 0, 3));
2975 assert!(!contains_nulls(Some(&buffer), 3, 2));
2976 assert!(!contains_nulls(Some(&buffer), 0, 0));
2977 }
2978
2979 #[test]
2980 fn test_alignment() {
2981 let buffer = Buffer::from_vec(vec![1_i32, 2_i32, 3_i32]);
2982 let sliced = buffer.slice(1);
2983
2984 let mut data = ArrayData {
2985 data_type: DataType::Int32,
2986 len: 0,
2987 offset: 0,
2988 buffers: vec![buffer],
2989 child_data: vec![],
2990 nulls: None,
2991 };
2992 data.validate_full().unwrap();
2993
2994 data.buffers[0] = sliced;
2996 let err = data.validate().unwrap_err();
2997
2998 assert_eq!(
2999 err.to_string(),
3000 "Invalid argument error: Misaligned buffers[0] in array of type Int32, offset from expected alignment of 4 by 1"
3001 );
3002
3003 data.align_buffers();
3004 data.validate_full().unwrap();
3005 }
3006
3007 #[test]
3008 fn test_alignment_struct() {
3009 let buffer = Buffer::from_vec(vec![1_i32, 2_i32, 3_i32]);
3010 let sliced = buffer.slice(1);
3011
3012 let child_data = ArrayData {
3013 data_type: DataType::Int32,
3014 len: 0,
3015 offset: 0,
3016 buffers: vec![buffer],
3017 child_data: vec![],
3018 nulls: None,
3019 };
3020
3021 let schema = DataType::Struct(Fields::from(vec![Field::new("a", DataType::Int32, false)]));
3022 let mut data = ArrayData {
3023 data_type: schema,
3024 len: 0,
3025 offset: 0,
3026 buffers: vec![],
3027 child_data: vec![child_data],
3028 nulls: None,
3029 };
3030 data.validate_full().unwrap();
3031
3032 data.child_data[0].buffers[0] = sliced;
3034 let err = data.validate().unwrap_err();
3035
3036 assert_eq!(
3037 err.to_string(),
3038 "Invalid argument error: Misaligned buffers[0] in array of type Int32, offset from expected alignment of 4 by 1"
3039 );
3040
3041 data.align_buffers();
3042 data.validate_full().unwrap();
3043 }
3044
3045 #[test]
3046 fn test_null_view_types() {
3047 let array_len = 32;
3048 let array = ArrayData::new_null(&DataType::BinaryView, array_len);
3049 assert_eq!(array.len(), array_len);
3050 for i in 0..array.len() {
3051 assert!(array.is_null(i));
3052 }
3053
3054 let array = ArrayData::new_null(&DataType::Utf8View, array_len);
3055 assert_eq!(array.len(), array_len);
3056 for i in 0..array.len() {
3057 assert!(array.is_null(i));
3058 }
3059
3060 let array = ArrayData::new_null(
3061 &DataType::ListView(Arc::new(Field::new_list_field(DataType::Int32, true))),
3062 array_len,
3063 );
3064 assert_eq!(array.len(), array_len);
3065 for i in 0..array.len() {
3066 assert!(array.is_null(i));
3067 }
3068
3069 let array = ArrayData::new_null(
3070 &DataType::LargeListView(Arc::new(Field::new_list_field(DataType::Int32, true))),
3071 array_len,
3072 );
3073 assert_eq!(array.len(), array_len);
3074 for i in 0..array.len() {
3075 assert!(array.is_null(i));
3076 }
3077 }
3078
3079 #[test]
3081 fn test_dont_panic_on_bad_input_when_using_try_new() {
3082 let empty_bytes = Buffer::default();
3083
3084 let array_data = ArrayData::try_new(
3085 DataType::Utf8,
3086 1, None,
3088 0,
3089 vec![Buffer::from_vec(vec![0i32, 2i32]), empty_bytes],
3091 vec![],
3092 );
3093
3094 let res = array_data.expect_err("should get error");
3095
3096 assert_eq!(
3097 res.to_string(),
3098 "Invalid argument error: Last offset 2 of Utf8 is larger than values length 0"
3099 );
3100 }
3101
3102 #[test]
3105 #[cfg(not(feature = "force_validate"))]
3106 fn test_validate_values_rejects_a_non_integer_dictionary_key() {
3107 let values = valid_non_nullable_int32_array_data(2);
3108 let data_type = DataType::Dictionary(Box::new(DataType::Utf8), Box::new(DataType::Int32));
3109 let dictionary = unsafe {
3110 ArrayData::builder(data_type)
3111 .len(1)
3112 .add_child_data(values)
3113 .build_unchecked()
3114 };
3115
3116 let err = dictionary.validate_values().expect_err("should get error");
3117 assert_eq!(
3118 err.to_string(),
3119 "Invalid argument error: Dictionary key type must be an integer, got Utf8"
3120 );
3121 }
3122
3123 #[test]
3124 #[cfg(not(feature = "force_validate"))]
3125 fn test_validate_values_rejects_a_non_integer_run_end() {
3126 let data_type = DataType::RunEndEncoded(
3127 Arc::new(Field::new("run_ends", DataType::Utf8, false)),
3128 Arc::new(Field::new("values", DataType::Int32, true)),
3129 );
3130 let run_end_encoded = unsafe {
3131 ArrayData::builder(data_type)
3132 .len(1)
3133 .add_child_data(valid_non_nullable_int32_array_data(1))
3134 .add_child_data(valid_non_nullable_int32_array_data(1))
3135 .build_unchecked()
3136 };
3137
3138 let err = run_end_encoded
3139 .validate_values()
3140 .expect_err("should get error");
3141 assert_eq!(
3142 err.to_string(),
3143 "Invalid argument error: Run end type must be Int16, Int32 or Int64, got Utf8"
3144 );
3145 }
3146
3147 #[test]
3149 #[cfg(not(feature = "force_validate"))]
3150 fn test_validate_values_rejects_missing_child_data() {
3151 let int32 = Box::new(DataType::Int32);
3152 let field = || Arc::new(Field::new("f", DataType::Int32, true));
3153 let data_types = [
3154 DataType::Dictionary(int32.clone(), int32.clone()),
3155 DataType::List(field()),
3156 DataType::LargeList(field()),
3157 DataType::RunEndEncoded(field(), field()),
3158 ];
3159
3160 for data_type in data_types {
3161 let data = unsafe {
3162 ArrayData::builder(data_type.clone())
3163 .len(1)
3164 .build_unchecked()
3165 };
3166 let err = data.validate_values().expect_err("should get error");
3167 assert_eq!(
3168 err.to_string(),
3169 format!(
3170 "Invalid argument error: {data_type} should contain at least 1 child data array(s), had 0"
3171 )
3172 );
3173 }
3174 }
3175
3176 #[test]
3178 #[cfg(not(feature = "force_validate"))]
3179 fn test_validate_values_rejects_missing_buffers() {
3180 let cases = [
3182 (DataType::Utf8, 1),
3183 (DataType::LargeUtf8, 1),
3184 (DataType::Binary, 1),
3185 (DataType::LargeBinary, 1),
3186 (DataType::BinaryView, 0),
3187 (DataType::Utf8View, 0),
3188 ];
3189
3190 for (data_type, missing) in cases {
3191 let data = unsafe {
3192 ArrayData::builder(data_type.clone())
3193 .len(1)
3194 .build_unchecked()
3195 };
3196 let err = data.validate_values().expect_err("should get error");
3197 assert_eq!(
3198 err.to_string(),
3199 format!(
3200 "Invalid argument error: {data_type} should contain at least {} buffer(s), had 0",
3201 missing + 1
3202 )
3203 );
3204 }
3205 }
3206
3207 #[test]
3209 #[cfg(not(feature = "force_validate"))]
3210 fn test_validate_values_rejects_a_short_dictionary_keys_buffer() {
3211 let data_type = DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Int32));
3212 let dictionary = unsafe {
3213 ArrayData::builder(data_type)
3214 .len(4)
3215 .add_buffer(Buffer::from_slice_ref([1_i32, 0]))
3216 .add_child_data(valid_non_nullable_int32_array_data(2))
3217 .build_unchecked()
3218 };
3219
3220 let err = dictionary.validate_values().expect_err("should get error");
3221 assert_eq!(
3222 err.to_string(),
3223 "Invalid argument error: Buffer 0 of Dictionary(Int32, Int32) isn't large enough. Expected 16 bytes got 8"
3224 );
3225 }
3226
3227 #[test]
3228 fn should_fail_validation_when_having_map_field_type_is_not_struct() {
3229 let map_field = Field::new("key", DataType::Int32, false);
3230
3231 let map_field_data = valid_non_nullable_int32_array_data(2);
3232
3233 let results = test_both_builder_and_array_data(
3234 DataType::Map(map_field.into(), false),
3235 1,
3236 None,
3237 0,
3238 vec![
3239 OffsetBuffer::<i32>::from_lengths(vec![2])
3240 .into_inner()
3241 .into(),
3242 ],
3243 vec![map_field_data],
3244 );
3245
3246 for result in results {
3247 let array_data_err = result.expect_err("should fail for non struct field");
3248
3249 match array_data_err {
3250 ArrowError::InvalidArgumentError(msg) => {
3251 assert_eq!(
3252 msg,
3253 "Map field should be a entries struct data type, got Int32 instead"
3254 )
3255 }
3256 _ => panic!("unexpected error type {array_data_err}"),
3257 }
3258 }
3259 }
3260
3261 #[test]
3262 fn should_fail_validation_when_having_map_entries_only_have_1_field() {
3263 let struct_data_type = DataType::Struct(Fields::from(vec![Field::new(
3264 Field::MAP_KEY_FIELD_DEFAULT_NAME,
3265 DataType::Int32,
3266 false,
3267 )]));
3268
3269 let key_array_data = valid_non_nullable_int32_array_data(2);
3270
3271 let struct_data = {
3272 let builder = ArrayDataBuilder::new(struct_data_type.clone())
3273 .len(2)
3274 .nulls(None)
3275 .child_data(vec![key_array_data]);
3276
3277 builder.build().unwrap()
3278 };
3279
3280 let results = test_both_builder_and_array_data(
3281 DataType::Map(
3282 Field::new(
3283 Field::MAP_ENTRIES_FIELD_DEFAULT_NAME,
3284 struct_data_type,
3285 false,
3286 )
3287 .into(),
3288 false,
3289 ),
3290 1,
3291 None,
3292 0,
3293 vec![
3294 OffsetBuffer::<i32>::from_lengths(vec![2])
3295 .into_inner()
3296 .into(),
3297 ],
3298 vec![struct_data],
3299 );
3300
3301 for result in results {
3302 let array_data_err = result.expect_err("should fail for nullable key");
3303
3304 match array_data_err {
3305 ArrowError::InvalidArgumentError(msg) => {
3306 assert_eq!(
3307 msg,
3308 "Map entries data type should be a struct containing 2 fields, got 1 fields"
3309 )
3310 }
3311 _ => panic!("unexpected error type {array_data_err}"),
3312 }
3313 }
3314 }
3315
3316 #[test]
3317 fn should_fail_validation_when_having_map_entries_have_3_fields() {
3318 let struct_data_type = DataType::Struct(Fields::from(vec![
3319 Field::new(Field::MAP_KEY_FIELD_DEFAULT_NAME, DataType::Int32, false),
3320 Field::new(Field::MAP_VALUE_FIELD_DEFAULT_NAME, DataType::Utf8, true),
3321 Field::new("other", DataType::Int32, true),
3322 ]));
3323
3324 let key_array_data = valid_non_nullable_int32_array_data(2);
3325
3326 let values_array_data = valid_string_array_data(2);
3327
3328 let other_array_data = key_array_data.clone();
3329
3330 let struct_data = {
3331 let builder = ArrayDataBuilder::new(struct_data_type.clone())
3332 .len(2)
3333 .nulls(None)
3334 .child_data(vec![key_array_data, values_array_data, other_array_data]);
3335
3336 builder.build().unwrap()
3337 };
3338
3339 let results = test_both_builder_and_array_data(
3340 DataType::Map(
3341 Field::new(
3342 Field::MAP_ENTRIES_FIELD_DEFAULT_NAME,
3343 struct_data_type,
3344 false,
3345 )
3346 .into(),
3347 false,
3348 ),
3349 1,
3350 None,
3351 0,
3352 vec![
3353 OffsetBuffer::<i32>::from_lengths(vec![2])
3354 .into_inner()
3355 .into(),
3356 ],
3357 vec![struct_data],
3358 );
3359
3360 for result in results {
3361 let array_data_err = result.expect_err("should fail for nullable key");
3362
3363 match array_data_err {
3364 ArrowError::InvalidArgumentError(msg) => {
3365 assert_eq!(
3366 msg,
3367 "Map entries data type should be a struct containing 2 fields, got 3 fields"
3368 )
3369 }
3370 _ => panic!("unexpected error type {array_data_err}"),
3371 }
3372 }
3373 }
3374
3375 #[test]
3376 fn should_fail_validation_when_having_nullable_map_keys() {
3377 let struct_data_type = DataType::Struct(Fields::from(vec![
3378 Field::new(Field::MAP_KEY_FIELD_DEFAULT_NAME, DataType::Int32, true),
3379 Field::new(Field::MAP_VALUE_FIELD_DEFAULT_NAME, DataType::Utf8, true),
3380 ]));
3381
3382 let key_array_data = valid_non_nullable_int32_array_data(2);
3383 let values_array_data = valid_string_array_data(2);
3384
3385 let struct_data = {
3386 let builder = ArrayDataBuilder::new(struct_data_type.clone())
3387 .len(2)
3388 .nulls(None)
3389 .child_data(vec![key_array_data, values_array_data]);
3390
3391 builder.build().unwrap()
3392 };
3393
3394 let results = test_both_builder_and_array_data(
3395 DataType::Map(
3396 Field::new(
3397 Field::MAP_ENTRIES_FIELD_DEFAULT_NAME,
3398 struct_data_type,
3399 false,
3400 )
3401 .into(),
3402 false,
3403 ),
3404 1,
3405 None,
3406 0,
3407 vec![
3408 OffsetBuffer::<i32>::from_lengths(vec![2])
3409 .into_inner()
3410 .into(),
3411 ],
3412 vec![struct_data],
3413 );
3414
3415 for result in results {
3416 let array_data_err = result.expect_err("should fail for nullable key");
3417
3418 match array_data_err {
3419 ArrowError::InvalidArgumentError(msg) => {
3420 assert_eq!(msg, "Map key field must not be nullable")
3421 }
3422 _ => panic!("unexpected error type {array_data_err}"),
3423 }
3424 }
3425 }
3426
3427 #[test]
3428 fn should_fail_validation_when_having_entries_is_nullable_for_map() {
3429 let struct_data_type = DataType::Struct(Fields::from(vec![
3430 Field::new(Field::MAP_KEY_FIELD_DEFAULT_NAME, DataType::Int32, false),
3431 Field::new(Field::MAP_VALUE_FIELD_DEFAULT_NAME, DataType::Utf8, true),
3432 ]));
3433
3434 let key_array_data = valid_non_nullable_int32_array_data(2);
3435
3436 let values_array_data = valid_string_array_data(2);
3437
3438 let struct_data = {
3439 let builder = ArrayDataBuilder::new(struct_data_type.clone())
3440 .len(2)
3441 .nulls(None)
3442 .child_data(vec![key_array_data, values_array_data]);
3443
3444 builder.build().unwrap()
3445 };
3446
3447 let results = test_both_builder_and_array_data(
3448 DataType::Map(
3449 Field::new(
3450 Field::MAP_ENTRIES_FIELD_DEFAULT_NAME,
3451 struct_data_type,
3452 true,
3453 )
3454 .into(),
3455 false,
3456 ),
3457 1,
3458 None,
3459 0,
3460 vec![
3461 OffsetBuffer::<i32>::from_lengths(vec![2])
3462 .into_inner()
3463 .into(),
3464 ],
3465 vec![struct_data],
3466 );
3467
3468 for result in results {
3469 let array_data_err = result.expect_err("should fail for nullable entries");
3470
3471 match array_data_err {
3472 ArrowError::InvalidArgumentError(msg) => assert_eq!(
3473 msg,
3474 "The nullable should be set to false for the map entries field."
3475 ),
3476 _ => panic!("unexpected error type {array_data_err}"),
3477 }
3478 }
3479 }
3480
3481 #[test]
3482 fn should_allow_to_create_map_from_data() {
3483 let struct_data_type = DataType::Struct(Fields::from(vec![
3484 Field::new(Field::MAP_KEY_FIELD_DEFAULT_NAME, DataType::Int32, false),
3485 Field::new(Field::MAP_VALUE_FIELD_DEFAULT_NAME, DataType::Utf8, true),
3486 ]));
3487
3488 let key_array_data = valid_non_nullable_int32_array_data(2);
3489 let values_array_data = valid_string_array_data(2);
3490
3491 let struct_data = {
3492 let builder = ArrayDataBuilder::new(struct_data_type.clone())
3493 .len(2)
3494 .nulls(None)
3495 .child_data(vec![key_array_data, values_array_data]);
3496
3497 builder.build().unwrap()
3498 };
3499
3500 let results = test_both_builder_and_array_data(
3501 DataType::Map(
3502 Field::new(
3503 Field::MAP_ENTRIES_FIELD_DEFAULT_NAME,
3504 struct_data_type,
3505 false,
3506 )
3507 .into(),
3508 false,
3509 ),
3510 1,
3511 None,
3512 0,
3513 vec![
3514 OffsetBuffer::<i32>::from_lengths(vec![2])
3515 .into_inner()
3516 .into(),
3517 ],
3518 vec![struct_data],
3519 );
3520
3521 for result in results {
3522 result.expect("should be able to create map ArrayData");
3523 }
3524 }
3525
3526 fn valid_string_array_data(length: usize) -> ArrayData {
3527 let offsets = OffsetBuffer::<i32>::from_lengths(vec![0; length])
3528 .into_inner()
3529 .into_inner();
3530 let empty_bytes = Buffer::default();
3531
3532 let builder = ArrayDataBuilder::new(DataType::Utf8)
3533 .len(length)
3534 .buffers(vec![offsets, empty_bytes])
3535 .nulls(None);
3536
3537 builder.build().unwrap()
3538 }
3539
3540 fn valid_non_nullable_int32_array_data(length: usize) -> ArrayData {
3541 let builder = ArrayDataBuilder::new(DataType::Int32)
3542 .len(length)
3543 .nulls(None)
3544 .buffers(vec![
3545 ScalarBuffer::<i32>::from(vec![1; length]).into_inner(),
3546 ]);
3547
3548 builder.build().unwrap()
3549 }
3550
3551 #[test]
3552 fn empty_and_null_map_array_should_pass_validation() {
3553 let dt = DataType::Map(
3554 Field::new(
3555 Field::MAP_ENTRIES_FIELD_DEFAULT_NAME,
3556 DataType::Struct(Fields::from(vec![
3557 Field::new(Field::MAP_KEY_FIELD_DEFAULT_NAME, DataType::Int32, false),
3558 Field::new(Field::MAP_VALUE_FIELD_DEFAULT_NAME, DataType::Utf8, true),
3559 ])),
3560 false,
3561 )
3562 .into(),
3563 false,
3564 );
3565
3566 ArrayData::new_empty(&dt).validate_full().unwrap();
3567 ArrayData::new_null(&dt, 1).validate_full().unwrap();
3568 }
3569
3570 #[test]
3571 fn null_buffer_offset_is_independent_of_data_offset() {
3572 let int_data = ArrayData::builder(DataType::UInt32)
3574 .offset(50)
3575 .len(50)
3576 .add_buffer(Buffer::from_vec(vec![0_u32; 100]))
3577 .build()
3578 .unwrap();
3579 int_data.validate().unwrap();
3580
3581 let nulls = NullBuffer::new(BooleanBuffer::from(vec![false; 100]).slice(0, 50));
3583 let with_sliced_nulls = int_data
3584 .clone()
3585 .into_builder()
3586 .nulls(Some(nulls))
3587 .build()
3588 .unwrap();
3589 with_sliced_nulls.validate().unwrap();
3590
3591 let nulls = NullBuffer::new(BooleanBuffer::from(vec![false; 50]));
3594 let with_unsliced_nulls = int_data.into_builder().nulls(Some(nulls)).build().unwrap();
3595 with_unsliced_nulls.validate().unwrap();
3596 assert_eq!(with_unsliced_nulls.null_count(), 50);
3597 }
3598
3599 fn test_both_builder_and_array_data(
3600 data_type: DataType,
3601 len: usize,
3602 null_bit_buffer: Option<Buffer>,
3603 offset: usize,
3604 buffers: Vec<Buffer>,
3605 child_data: Vec<ArrayData>,
3606 ) -> [Result<ArrayData, ArrowError>; 2] {
3607 let from_builder_res = ArrayData::builder(data_type.clone())
3608 .len(len)
3609 .add_buffers(buffers.clone())
3610 .null_bit_buffer(null_bit_buffer.clone())
3611 .offset(offset)
3612 .child_data(child_data.clone())
3613 .build();
3614
3615 let from_try_new_res =
3616 ArrayData::try_new(data_type, len, null_bit_buffer, offset, buffers, child_data);
3617
3618 [from_builder_res, from_try_new_res]
3619 }
3620}