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