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 #[allow(dead_code)]
2020 AlwaysNull,
2021}
2022
2023impl PartialEq for ArrayData {
2024 fn eq(&self, other: &Self) -> bool {
2025 equal::equal(self, other)
2026 }
2027}
2028
2029#[derive(Debug, Clone)]
2048#[doc(hidden)]
2049pub struct UnsafeFlag(bool);
2050
2051impl UnsafeFlag {
2052 #[inline]
2056 pub const fn new() -> Self {
2057 Self(false)
2058 }
2059
2060 #[inline]
2070 pub unsafe fn set(&mut self, val: bool) {
2071 self.0 = val;
2072 }
2073
2074 #[inline]
2076 pub fn get(&self) -> bool {
2077 self.0
2078 }
2079}
2080
2081impl Default for UnsafeFlag {
2083 fn default() -> Self {
2084 Self::new()
2085 }
2086}
2087
2088#[derive(Debug)]
2090pub struct ArrayDataBuilder {
2091 data_type: DataType,
2092 len: usize,
2093 null_count: Option<usize>,
2094 null_bit_buffer: Option<Buffer>,
2095 nulls: Option<NullBuffer>,
2096 offset: usize,
2097 buffers: Vec<Buffer>,
2098 child_data: Vec<ArrayData>,
2099 align_buffers: bool,
2103 skip_validation: UnsafeFlag,
2113}
2114
2115impl ArrayDataBuilder {
2116 #[inline]
2117 pub const fn new(data_type: DataType) -> Self {
2119 Self {
2120 data_type,
2121 len: 0,
2122 null_count: None,
2123 null_bit_buffer: None,
2124 nulls: None,
2125 offset: 0,
2126 buffers: vec![],
2127 child_data: vec![],
2128 align_buffers: false,
2129 skip_validation: UnsafeFlag::new(),
2130 }
2131 }
2132
2133 pub fn data_type(self, data_type: DataType) -> Self {
2135 Self { data_type, ..self }
2136 }
2137
2138 #[inline]
2139 #[allow(clippy::len_without_is_empty)]
2140 pub const fn len(mut self, n: usize) -> Self {
2142 self.len = n;
2143 self
2144 }
2145
2146 pub fn nulls(mut self, nulls: Option<NullBuffer>) -> Self {
2148 self.nulls = nulls;
2149 self.null_count = None;
2150 self.null_bit_buffer = None;
2151 self
2152 }
2153
2154 pub fn null_count(mut self, null_count: usize) -> Self {
2156 self.null_count = Some(null_count);
2157 self
2158 }
2159
2160 pub fn null_bit_buffer(mut self, buf: Option<Buffer>) -> Self {
2162 self.nulls = None;
2163 self.null_bit_buffer = buf;
2164 self
2165 }
2166
2167 #[inline]
2169 pub const fn offset(mut self, n: usize) -> Self {
2170 self.offset = n;
2171 self
2172 }
2173
2174 pub fn buffers(mut self, v: Vec<Buffer>) -> Self {
2176 self.buffers = v;
2177 self
2178 }
2179
2180 pub fn add_buffer(mut self, b: Buffer) -> Self {
2182 self.buffers.push(b);
2183 self
2184 }
2185
2186 pub fn add_buffers<I: IntoIterator<Item = Buffer>>(mut self, bs: I) -> Self {
2188 self.buffers.extend(bs);
2189 self
2190 }
2191
2192 pub fn child_data(mut self, v: Vec<ArrayData>) -> Self {
2194 self.child_data = v;
2195 self
2196 }
2197
2198 pub fn add_child_data(mut self, r: ArrayData) -> Self {
2200 self.child_data.push(r);
2201 self
2202 }
2203
2204 pub unsafe fn build_unchecked(self) -> ArrayData {
2220 unsafe { self.skip_validation(true) }.build().unwrap()
2221 }
2222
2223 pub fn build(self) -> Result<ArrayData, ArrowError> {
2232 let Self {
2233 data_type,
2234 len,
2235 null_count,
2236 null_bit_buffer,
2237 nulls,
2238 offset,
2239 buffers,
2240 child_data,
2241 align_buffers,
2242 skip_validation,
2243 } = self;
2244
2245 let nulls = nulls
2246 .or_else(|| {
2247 let buffer = null_bit_buffer?;
2248 let buffer = BooleanBuffer::new(buffer, offset, len);
2249 Some(match null_count {
2250 Some(n) => {
2251 unsafe { NullBuffer::new_unchecked(buffer, n) }
2253 }
2254 None => NullBuffer::new(buffer),
2255 })
2256 })
2257 .filter(|b| b.null_count() != 0);
2258
2259 let mut data = ArrayData {
2260 data_type,
2261 len,
2262 offset,
2263 buffers,
2264 child_data,
2265 nulls,
2266 };
2267
2268 if align_buffers {
2269 data.align_buffers();
2270 }
2271
2272 if !skip_validation.get() || cfg!(feature = "force_validate") {
2274 data.validate_data()?;
2275 }
2276 Ok(data)
2277 }
2278
2279 pub fn align_buffers(mut self, align_buffers: bool) -> Self {
2295 self.align_buffers = align_buffers;
2296 self
2297 }
2298
2299 pub unsafe fn skip_validation(mut self, skip_validation: bool) -> Self {
2313 unsafe {
2314 self.skip_validation.set(skip_validation);
2315 }
2316 self
2317 }
2318}
2319
2320impl From<ArrayData> for ArrayDataBuilder {
2321 fn from(d: ArrayData) -> Self {
2322 Self {
2323 data_type: d.data_type,
2324 len: d.len,
2325 offset: d.offset,
2326 buffers: d.buffers,
2327 child_data: d.child_data,
2328 nulls: d.nulls,
2329 null_bit_buffer: None,
2330 null_count: None,
2331 align_buffers: false,
2332 skip_validation: UnsafeFlag::new(),
2333 }
2334 }
2335}
2336
2337pub(crate) fn get_fixed_size_binary_width(data_type: &DataType) -> usize {
2342 match data_type {
2343 DataType::FixedSizeBinary(i) => {
2344 if *i < 0 {
2345 panic!("cannot compare FixedSizeBinary({})", *i);
2346 }
2347 *i as usize
2348 }
2349 _ => unreachable!(),
2350 }
2351}
2352
2353#[cfg(test)]
2354mod tests {
2355 use super::*;
2356 use crate::ByteView;
2357 use arrow_buffer::{OffsetBuffer, ScalarBuffer};
2358 use arrow_schema::{Field, Fields};
2359
2360 fn make_i32_buffer(n: usize) -> Buffer {
2364 Buffer::from_slice_ref(vec![42i32; n])
2365 }
2366
2367 fn make_f32_buffer(n: usize) -> Buffer {
2369 Buffer::from_slice_ref(vec![42f32; n])
2370 }
2371
2372 #[test]
2373 fn test_builder() {
2374 let v = (0..25).collect::<Vec<i32>>();
2376 let b1 = Buffer::from_slice_ref(&v);
2377 let arr_data = ArrayData::builder(DataType::Int32)
2378 .len(20)
2379 .offset(5)
2380 .add_buffer(b1)
2381 .null_bit_buffer(Some(Buffer::from([
2382 0b01011111, 0b10110101, 0b01100011, 0b00011110,
2383 ])))
2384 .build()
2385 .unwrap();
2386
2387 assert_eq!(20, arr_data.len());
2388 assert_eq!(10, arr_data.null_count());
2389 assert_eq!(5, arr_data.offset());
2390 assert_eq!(1, arr_data.buffers().len());
2391 assert_eq!(
2392 Buffer::from_slice_ref(&v).as_slice(),
2393 arr_data.buffers()[0].as_slice()
2394 );
2395 }
2396
2397 #[test]
2398 fn test_builder_with_child_data() {
2399 let child_arr_data = ArrayData::try_new(
2400 DataType::Int32,
2401 5,
2402 None,
2403 0,
2404 vec![Buffer::from_slice_ref([1i32, 2, 3, 4, 5])],
2405 vec![],
2406 )
2407 .unwrap();
2408
2409 let field = Arc::new(Field::new("x", DataType::Int32, true));
2410 let data_type = DataType::Struct(vec![field].into());
2411
2412 let arr_data = ArrayData::builder(data_type)
2413 .len(5)
2414 .offset(0)
2415 .add_child_data(child_arr_data.clone())
2416 .build()
2417 .unwrap();
2418
2419 assert_eq!(5, arr_data.len());
2420 assert_eq!(1, arr_data.child_data().len());
2421 assert_eq!(child_arr_data, arr_data.child_data()[0]);
2422 }
2423
2424 #[test]
2425 fn test_null_count() {
2426 let mut bit_v: [u8; 2] = [0; 2];
2427 bit_util::set_bit(&mut bit_v, 0);
2428 bit_util::set_bit(&mut bit_v, 3);
2429 bit_util::set_bit(&mut bit_v, 10);
2430 let arr_data = ArrayData::builder(DataType::Int32)
2431 .len(16)
2432 .add_buffer(make_i32_buffer(16))
2433 .null_bit_buffer(Some(Buffer::from(bit_v)))
2434 .build()
2435 .unwrap();
2436 assert_eq!(13, arr_data.null_count());
2437
2438 let mut bit_v: [u8; 2] = [0; 2];
2440 bit_util::set_bit(&mut bit_v, 0);
2441 bit_util::set_bit(&mut bit_v, 3);
2442 bit_util::set_bit(&mut bit_v, 10);
2443 let arr_data = ArrayData::builder(DataType::Int32)
2444 .len(12)
2445 .offset(2)
2446 .add_buffer(make_i32_buffer(14)) .null_bit_buffer(Some(Buffer::from(bit_v)))
2448 .build()
2449 .unwrap();
2450 assert_eq!(10, arr_data.null_count());
2451 }
2452
2453 #[test]
2454 fn test_null_buffer_ref() {
2455 let mut bit_v: [u8; 2] = [0; 2];
2456 bit_util::set_bit(&mut bit_v, 0);
2457 bit_util::set_bit(&mut bit_v, 3);
2458 bit_util::set_bit(&mut bit_v, 10);
2459 let arr_data = ArrayData::builder(DataType::Int32)
2460 .len(16)
2461 .add_buffer(make_i32_buffer(16))
2462 .null_bit_buffer(Some(Buffer::from(bit_v)))
2463 .build()
2464 .unwrap();
2465 assert!(arr_data.nulls().is_some());
2466 assert_eq!(&bit_v, arr_data.nulls().unwrap().validity());
2467 }
2468
2469 #[test]
2470 fn test_slice() {
2471 let mut bit_v: [u8; 2] = [0; 2];
2472 bit_util::set_bit(&mut bit_v, 0);
2473 bit_util::set_bit(&mut bit_v, 3);
2474 bit_util::set_bit(&mut bit_v, 10);
2475 let data = ArrayData::builder(DataType::Int32)
2476 .len(16)
2477 .add_buffer(make_i32_buffer(16))
2478 .null_bit_buffer(Some(Buffer::from(bit_v)))
2479 .build()
2480 .unwrap();
2481 let new_data = data.slice(1, 15);
2482 assert_eq!(data.len() - 1, new_data.len());
2483 assert_eq!(1, new_data.offset());
2484 assert_eq!(data.null_count(), new_data.null_count());
2485
2486 let new_data = new_data.slice(1, 14);
2488 assert_eq!(data.len() - 2, new_data.len());
2489 assert_eq!(2, new_data.offset());
2490 assert_eq!(data.null_count() - 1, new_data.null_count());
2491 }
2492
2493 #[test]
2494 #[should_panic(expected = "offset + length overflow")]
2495 fn test_slice_panics_on_offset_length_overflow() {
2496 let data = ArrayData::builder(DataType::Int32)
2497 .len(4)
2498 .add_buffer(make_i32_buffer(4))
2499 .build()
2500 .unwrap();
2501 let sliced = data.slice(1, 3);
2502
2503 sliced.slice(1, usize::MAX);
2504 }
2505
2506 #[test]
2507 fn test_typed_offsets_length_overflow() {
2508 let data = ArrayData {
2509 data_type: DataType::Binary,
2510 len: usize::MAX,
2511 offset: 0,
2512 buffers: vec![Buffer::from_slice_ref([0_i32])],
2513 child_data: vec![],
2514 nulls: None,
2515 };
2516 let err = data.typed_offsets::<i32>().unwrap_err();
2517
2518 assert_eq!(
2519 err.to_string(),
2520 format!(
2521 "Invalid argument error: Length {} with offset 1 overflows usize for Binary",
2522 usize::MAX
2523 )
2524 );
2525 }
2526
2527 #[test]
2528 fn test_validate_typed_buffer_length_overflow() {
2529 let data = ArrayData {
2530 data_type: DataType::Binary,
2531 len: 0,
2532 offset: 2,
2533 buffers: vec![Buffer::from_slice_ref([0_i32])],
2534 child_data: vec![],
2535 nulls: None,
2536 };
2537 let err = data.typed_buffer::<i32>(0, usize::MAX).unwrap_err();
2538
2539 assert_eq!(
2540 err.to_string(),
2541 format!(
2542 "Invalid argument error: Length {} with offset 2 overflows usize for Binary",
2543 usize::MAX
2544 )
2545 );
2546 }
2547
2548 fn try_new_binary_length_offset_overflow() -> Result<ArrayData, ArrowError> {
2550 ArrayData::try_new(
2551 DataType::Binary,
2552 usize::MAX,
2553 None,
2554 1,
2555 vec![
2556 Buffer::from_slice_ref([0_i32]),
2557 Buffer::from_iter(std::iter::empty::<u8>()),
2558 ],
2559 vec![],
2560 )
2561 }
2562
2563 #[cfg(not(feature = "force_validate"))]
2564 #[test]
2565 fn test_try_new_length_offset_overflow() {
2566 let err = try_new_binary_length_offset_overflow().unwrap_err();
2567
2568 assert_eq!(
2569 err.to_string(),
2570 format!(
2571 "Invalid argument error: Length {} with offset 1 overflows usize for Binary",
2572 usize::MAX
2573 )
2574 );
2575 }
2576
2577 #[cfg(feature = "force_validate")]
2578 #[test]
2579 #[should_panic(
2580 expected = "Length 18446744073709551615 with offset 1 overflows usize for Binary"
2581 )]
2582 fn test_try_new_length_offset_overflow_force_validate() {
2583 try_new_binary_length_offset_overflow().unwrap();
2584 }
2585
2586 #[test]
2587 fn test_equality() {
2588 let int_data = ArrayData::builder(DataType::Int32)
2589 .len(1)
2590 .add_buffer(make_i32_buffer(1))
2591 .build()
2592 .unwrap();
2593
2594 let float_data = ArrayData::builder(DataType::Float32)
2595 .len(1)
2596 .add_buffer(make_f32_buffer(1))
2597 .build()
2598 .unwrap();
2599 assert_ne!(int_data, float_data);
2600 assert!(!int_data.ptr_eq(&float_data));
2601 assert!(int_data.ptr_eq(&int_data));
2602
2603 #[allow(clippy::redundant_clone)]
2604 let int_data_clone = int_data.clone();
2605 assert_eq!(int_data, int_data_clone);
2606 assert!(int_data.ptr_eq(&int_data_clone));
2607 assert!(int_data_clone.ptr_eq(&int_data));
2608
2609 let int_data_slice = int_data_clone.slice(1, 0);
2610 assert!(int_data_slice.ptr_eq(&int_data_slice));
2611 assert!(!int_data.ptr_eq(&int_data_slice));
2612 assert!(!int_data_slice.ptr_eq(&int_data));
2613
2614 let data_buffer = Buffer::from_slice_ref("abcdef".as_bytes());
2615 let offsets_buffer = Buffer::from_slice_ref([0_i32, 2_i32, 2_i32, 5_i32]);
2616 let string_data = ArrayData::try_new(
2617 DataType::Utf8,
2618 3,
2619 Some(Buffer::from_iter(vec![true, false, true])),
2620 0,
2621 vec![offsets_buffer, data_buffer],
2622 vec![],
2623 )
2624 .unwrap();
2625
2626 assert_ne!(float_data, string_data);
2627 assert!(!float_data.ptr_eq(&string_data));
2628
2629 assert!(string_data.ptr_eq(&string_data));
2630
2631 #[allow(clippy::redundant_clone)]
2632 let string_data_cloned = string_data.clone();
2633 assert!(string_data_cloned.ptr_eq(&string_data));
2634 assert!(string_data.ptr_eq(&string_data_cloned));
2635
2636 let string_data_slice = string_data.slice(1, 2);
2637 assert!(string_data_slice.ptr_eq(&string_data_slice));
2638 assert!(!string_data_slice.ptr_eq(&string_data))
2639 }
2640
2641 #[test]
2642 fn test_slice_memory_size_view_payload_buffers() {
2643 for data_type in [DataType::Utf8View, DataType::BinaryView] {
2644 let inline_only = ArrayData::builder(data_type.clone())
2645 .len(2)
2646 .add_buffer(Buffer::from_vec(vec![0_u128; 2]))
2647 .build()
2648 .unwrap();
2649 assert_eq!(
2650 inline_only.get_slice_memory_size().unwrap(),
2651 2 * mem::size_of::<u128>()
2652 );
2653
2654 let mut first_payload = Vec::with_capacity(32);
2655 first_payload.extend_from_slice(b"first payload");
2656 let first_view =
2657 ByteView::new(first_payload.len().try_into().unwrap(), &first_payload[..4])
2658 .as_u128();
2659 let first_payload = Buffer::from_vec(first_payload);
2660 assert!(first_payload.capacity() > first_payload.len());
2661 let first_payload_capacity = first_payload.capacity();
2662
2663 let mut second_payload = Vec::with_capacity(64);
2664 second_payload.extend_from_slice(b"second payload");
2665 let second_view = ByteView::new(
2666 second_payload.len().try_into().unwrap(),
2667 &second_payload[..4],
2668 )
2669 .with_buffer_index(1)
2670 .as_u128();
2671 let second_payload = Buffer::from_vec(second_payload);
2672 assert!(second_payload.capacity() > second_payload.len());
2673 let second_payload_capacity = second_payload.capacity();
2674
2675 let data = ArrayData::builder(data_type)
2676 .len(3)
2677 .add_buffer(Buffer::from_vec(vec![first_view, 0_u128, second_view]))
2678 .add_buffer(first_payload)
2679 .add_buffer(second_payload)
2680 .build()
2681 .unwrap();
2682 let sliced = data.slice(1, 1);
2683
2684 assert_eq!(
2685 sliced.get_slice_memory_size().unwrap(),
2686 mem::size_of::<u128>() + first_payload_capacity + second_payload_capacity
2687 );
2688 }
2689 }
2690
2691 #[test]
2692 fn test_slice_memory_size_utf8_offset_buffer_len_plus_one() {
2693 let data_buffer = Buffer::from_slice_ref("helloworld".as_bytes());
2695 let offsets_buffer = Buffer::from_slice_ref([0_i32, 5_i32, 10_i32]);
2698 let array = ArrayData::try_new(
2699 DataType::Utf8,
2700 2,
2701 None,
2702 0,
2703 vec![offsets_buffer, data_buffer],
2704 vec![],
2705 )
2706 .unwrap();
2707 assert_eq!(array.get_slice_memory_size().unwrap(), 22); }
2709
2710 #[test]
2711 fn test_slice_memory_size_binary_offset_buffer_len_plus_one() {
2712 let data_buffer = Buffer::from_slice_ref([0u8, 1, 2, 3, 4]);
2715 let offsets_buffer = Buffer::from_slice_ref([0_i32, 2_i32, 5_i32]);
2717 let array = ArrayData::try_new(
2718 DataType::Binary,
2719 2,
2720 None,
2721 0,
2722 vec![offsets_buffer, data_buffer],
2723 vec![],
2724 )
2725 .unwrap();
2726 assert_eq!(array.get_slice_memory_size().unwrap(), 17); }
2728
2729 #[test]
2730 fn test_slice_memory_size() {
2731 let mut bit_v: [u8; 2] = [0; 2];
2732 bit_util::set_bit(&mut bit_v, 0);
2733 bit_util::set_bit(&mut bit_v, 3);
2734 bit_util::set_bit(&mut bit_v, 10);
2735 let data = ArrayData::builder(DataType::Int32)
2736 .len(16)
2737 .add_buffer(make_i32_buffer(16))
2738 .null_bit_buffer(Some(Buffer::from(bit_v)))
2739 .build()
2740 .unwrap();
2741 let new_data = data.slice(1, 14);
2742 assert_eq!(
2743 data.get_slice_memory_size().unwrap() - 8,
2744 new_data.get_slice_memory_size().unwrap()
2745 );
2746 let data_buffer = Buffer::from_slice_ref("abcdef".as_bytes());
2747 let offsets_buffer = Buffer::from_slice_ref([0_i32, 2_i32, 2_i32, 5_i32]);
2748 let string_data = ArrayData::try_new(
2749 DataType::Utf8,
2750 3,
2751 Some(Buffer::from_iter(vec![true, false, true])),
2752 0,
2753 vec![offsets_buffer, data_buffer],
2754 vec![],
2755 )
2756 .unwrap();
2757 let string_data_slice = string_data.slice(1, 2);
2758 assert_eq!(
2760 string_data.get_slice_memory_size().unwrap() - 6,
2761 string_data_slice.get_slice_memory_size().unwrap()
2762 );
2763 }
2764
2765 #[test]
2766 fn test_count_nulls() {
2767 let buffer = Buffer::from([0b00010110, 0b10011111]);
2768 let buffer = NullBuffer::new(BooleanBuffer::new(buffer, 0, 16));
2769 let count = count_nulls(Some(&buffer), 0, 16);
2770 assert_eq!(count, 7);
2771
2772 let count = count_nulls(Some(&buffer), 4, 8);
2773 assert_eq!(count, 3);
2774 }
2775
2776 #[test]
2777 fn test_contains_nulls() {
2778 let buffer: Buffer =
2779 MutableBuffer::from_iter([false, false, false, true, true, false]).into();
2780 let buffer = NullBuffer::new(BooleanBuffer::new(buffer, 0, 6));
2781 assert!(contains_nulls(Some(&buffer), 0, 6));
2782 assert!(contains_nulls(Some(&buffer), 0, 3));
2783 assert!(!contains_nulls(Some(&buffer), 3, 2));
2784 assert!(!contains_nulls(Some(&buffer), 0, 0));
2785 }
2786
2787 #[test]
2788 fn test_alignment() {
2789 let buffer = Buffer::from_vec(vec![1_i32, 2_i32, 3_i32]);
2790 let sliced = buffer.slice(1);
2791
2792 let mut data = ArrayData {
2793 data_type: DataType::Int32,
2794 len: 0,
2795 offset: 0,
2796 buffers: vec![buffer],
2797 child_data: vec![],
2798 nulls: None,
2799 };
2800 data.validate_full().unwrap();
2801
2802 data.buffers[0] = sliced;
2804 let err = data.validate().unwrap_err();
2805
2806 assert_eq!(
2807 err.to_string(),
2808 "Invalid argument error: Misaligned buffers[0] in array of type Int32, offset from expected alignment of 4 by 1"
2809 );
2810
2811 data.align_buffers();
2812 data.validate_full().unwrap();
2813 }
2814
2815 #[test]
2816 fn test_alignment_struct() {
2817 let buffer = Buffer::from_vec(vec![1_i32, 2_i32, 3_i32]);
2818 let sliced = buffer.slice(1);
2819
2820 let child_data = ArrayData {
2821 data_type: DataType::Int32,
2822 len: 0,
2823 offset: 0,
2824 buffers: vec![buffer],
2825 child_data: vec![],
2826 nulls: None,
2827 };
2828
2829 let schema = DataType::Struct(Fields::from(vec![Field::new("a", DataType::Int32, false)]));
2830 let mut data = ArrayData {
2831 data_type: schema,
2832 len: 0,
2833 offset: 0,
2834 buffers: vec![],
2835 child_data: vec![child_data],
2836 nulls: None,
2837 };
2838 data.validate_full().unwrap();
2839
2840 data.child_data[0].buffers[0] = sliced;
2842 let err = data.validate().unwrap_err();
2843
2844 assert_eq!(
2845 err.to_string(),
2846 "Invalid argument error: Misaligned buffers[0] in array of type Int32, offset from expected alignment of 4 by 1"
2847 );
2848
2849 data.align_buffers();
2850 data.validate_full().unwrap();
2851 }
2852
2853 #[test]
2854 fn test_null_view_types() {
2855 let array_len = 32;
2856 let array = ArrayData::new_null(&DataType::BinaryView, array_len);
2857 assert_eq!(array.len(), array_len);
2858 for i in 0..array.len() {
2859 assert!(array.is_null(i));
2860 }
2861
2862 let array = ArrayData::new_null(&DataType::Utf8View, array_len);
2863 assert_eq!(array.len(), array_len);
2864 for i in 0..array.len() {
2865 assert!(array.is_null(i));
2866 }
2867
2868 let array = ArrayData::new_null(
2869 &DataType::ListView(Arc::new(Field::new_list_field(DataType::Int32, true))),
2870 array_len,
2871 );
2872 assert_eq!(array.len(), array_len);
2873 for i in 0..array.len() {
2874 assert!(array.is_null(i));
2875 }
2876
2877 let array = ArrayData::new_null(
2878 &DataType::LargeListView(Arc::new(Field::new_list_field(DataType::Int32, true))),
2879 array_len,
2880 );
2881 assert_eq!(array.len(), array_len);
2882 for i in 0..array.len() {
2883 assert!(array.is_null(i));
2884 }
2885 }
2886
2887 #[test]
2889 fn test_dont_panic_on_bad_input_when_using_try_new() {
2890 let empty_bytes = Buffer::default();
2891
2892 let array_data = ArrayData::try_new(
2893 DataType::Utf8,
2894 1, None,
2896 0,
2897 vec![Buffer::from_vec(vec![0i32, 2i32]), empty_bytes],
2899 vec![],
2900 );
2901
2902 let res = array_data.expect_err("should get error");
2903
2904 assert_eq!(
2905 res.to_string(),
2906 format!("Invalid argument error: Last offset 2 of Utf8 is larger than values length 0",)
2907 );
2908 }
2909
2910 #[test]
2911 fn should_fail_validation_when_having_map_field_type_is_not_struct() {
2912 let map_field = Field::new("key", DataType::Int32, false);
2913
2914 let map_field_data = valid_non_nullable_int32_array_data(2);
2915
2916 let results = test_both_builder_and_array_data(
2917 DataType::Map(map_field.into(), false),
2918 1,
2919 None,
2920 0,
2921 vec![
2922 OffsetBuffer::<i32>::from_lengths(vec![2])
2923 .into_inner()
2924 .into(),
2925 ],
2926 vec![map_field_data],
2927 );
2928
2929 for result in results {
2930 let array_data_err = result.expect_err("should fail for non struct field");
2931
2932 match array_data_err {
2933 ArrowError::InvalidArgumentError(msg) => {
2934 assert_eq!(
2935 msg,
2936 "Map field should be a entries struct data type, got Int32 instead"
2937 )
2938 }
2939 _ => panic!("unexpected error type {array_data_err}"),
2940 };
2941 }
2942 }
2943
2944 #[test]
2945 fn should_fail_validation_when_having_map_entries_only_have_1_field() {
2946 let struct_data_type = DataType::Struct(Fields::from(vec![Field::new(
2947 Field::MAP_KEY_FIELD_DEFAULT_NAME,
2948 DataType::Int32,
2949 false,
2950 )]));
2951
2952 let key_array_data = valid_non_nullable_int32_array_data(2);
2953
2954 let struct_data = {
2955 let builder = ArrayDataBuilder::new(struct_data_type.clone())
2956 .len(2)
2957 .nulls(None)
2958 .child_data(vec![key_array_data]);
2959
2960 builder.build().unwrap()
2961 };
2962
2963 let results = test_both_builder_and_array_data(
2964 DataType::Map(
2965 Field::new(
2966 Field::MAP_ENTRIES_FIELD_DEFAULT_NAME,
2967 struct_data_type,
2968 false,
2969 )
2970 .into(),
2971 false,
2972 ),
2973 1,
2974 None,
2975 0,
2976 vec![
2977 OffsetBuffer::<i32>::from_lengths(vec![2])
2978 .into_inner()
2979 .into(),
2980 ],
2981 vec![struct_data],
2982 );
2983
2984 for result in results {
2985 let array_data_err = result.expect_err("should fail for nullable key");
2986
2987 match array_data_err {
2988 ArrowError::InvalidArgumentError(msg) => {
2989 assert_eq!(
2990 msg,
2991 "Map entries data type should be a struct containing 2 fields, got 1 fields"
2992 )
2993 }
2994 _ => panic!("unexpected error type {array_data_err}"),
2995 };
2996 }
2997 }
2998
2999 #[test]
3000 fn should_fail_validation_when_having_map_entries_have_3_fields() {
3001 let struct_data_type = DataType::Struct(Fields::from(vec![
3002 Field::new(Field::MAP_KEY_FIELD_DEFAULT_NAME, DataType::Int32, false),
3003 Field::new(Field::MAP_VALUE_FIELD_DEFAULT_NAME, DataType::Utf8, true),
3004 Field::new("other", DataType::Int32, true),
3005 ]));
3006
3007 let key_array_data = valid_non_nullable_int32_array_data(2);
3008
3009 let values_array_data = valid_string_array_data(2);
3010
3011 let other_array_data = key_array_data.clone();
3012
3013 let struct_data = {
3014 let builder = ArrayDataBuilder::new(struct_data_type.clone())
3015 .len(2)
3016 .nulls(None)
3017 .child_data(vec![key_array_data, values_array_data, other_array_data]);
3018
3019 builder.build().unwrap()
3020 };
3021
3022 let results = test_both_builder_and_array_data(
3023 DataType::Map(
3024 Field::new(
3025 Field::MAP_ENTRIES_FIELD_DEFAULT_NAME,
3026 struct_data_type,
3027 false,
3028 )
3029 .into(),
3030 false,
3031 ),
3032 1,
3033 None,
3034 0,
3035 vec![
3036 OffsetBuffer::<i32>::from_lengths(vec![2])
3037 .into_inner()
3038 .into(),
3039 ],
3040 vec![struct_data],
3041 );
3042
3043 for result in results {
3044 let array_data_err = result.expect_err("should fail for nullable key");
3045
3046 match array_data_err {
3047 ArrowError::InvalidArgumentError(msg) => {
3048 assert_eq!(
3049 msg,
3050 "Map entries data type should be a struct containing 2 fields, got 3 fields"
3051 )
3052 }
3053 _ => panic!("unexpected error type {array_data_err}"),
3054 };
3055 }
3056 }
3057
3058 #[test]
3059 fn should_fail_validation_when_having_nullable_map_keys() {
3060 let struct_data_type = DataType::Struct(Fields::from(vec![
3061 Field::new(Field::MAP_KEY_FIELD_DEFAULT_NAME, DataType::Int32, true),
3062 Field::new(Field::MAP_VALUE_FIELD_DEFAULT_NAME, DataType::Utf8, true),
3063 ]));
3064
3065 let key_array_data = valid_non_nullable_int32_array_data(2);
3066 let values_array_data = valid_string_array_data(2);
3067
3068 let struct_data = {
3069 let builder = ArrayDataBuilder::new(struct_data_type.clone())
3070 .len(2)
3071 .nulls(None)
3072 .child_data(vec![key_array_data, values_array_data]);
3073
3074 builder.build().unwrap()
3075 };
3076
3077 let results = test_both_builder_and_array_data(
3078 DataType::Map(
3079 Field::new(
3080 Field::MAP_ENTRIES_FIELD_DEFAULT_NAME,
3081 struct_data_type,
3082 false,
3083 )
3084 .into(),
3085 false,
3086 ),
3087 1,
3088 None,
3089 0,
3090 vec![
3091 OffsetBuffer::<i32>::from_lengths(vec![2])
3092 .into_inner()
3093 .into(),
3094 ],
3095 vec![struct_data],
3096 );
3097
3098 for result in results {
3099 let array_data_err = result.expect_err("should fail for nullable key");
3100
3101 match array_data_err {
3102 ArrowError::InvalidArgumentError(msg) => {
3103 assert_eq!(msg, "Map key field must not be nullable")
3104 }
3105 _ => panic!("unexpected error type {array_data_err}"),
3106 };
3107 }
3108 }
3109
3110 #[test]
3111 fn should_fail_validation_when_having_entries_is_nullable_for_map() {
3112 let struct_data_type = DataType::Struct(Fields::from(vec![
3113 Field::new(Field::MAP_KEY_FIELD_DEFAULT_NAME, DataType::Int32, false),
3114 Field::new(Field::MAP_VALUE_FIELD_DEFAULT_NAME, DataType::Utf8, true),
3115 ]));
3116
3117 let key_array_data = valid_non_nullable_int32_array_data(2);
3118
3119 let values_array_data = valid_string_array_data(2);
3120
3121 let struct_data = {
3122 let builder = ArrayDataBuilder::new(struct_data_type.clone())
3123 .len(2)
3124 .nulls(None)
3125 .child_data(vec![key_array_data, values_array_data]);
3126
3127 builder.build().unwrap()
3128 };
3129
3130 let results = test_both_builder_and_array_data(
3131 DataType::Map(
3132 Field::new(
3133 Field::MAP_ENTRIES_FIELD_DEFAULT_NAME,
3134 struct_data_type,
3135 true,
3136 )
3137 .into(),
3138 false,
3139 ),
3140 1,
3141 None,
3142 0,
3143 vec![
3144 OffsetBuffer::<i32>::from_lengths(vec![2])
3145 .into_inner()
3146 .into(),
3147 ],
3148 vec![struct_data],
3149 );
3150
3151 for result in results {
3152 let array_data_err = result.expect_err("should fail for nullable entries");
3153
3154 match array_data_err {
3155 ArrowError::InvalidArgumentError(msg) => assert_eq!(
3156 msg,
3157 "The nullable should be set to false for the map entries field."
3158 ),
3159 _ => panic!("unexpected error type {array_data_err}"),
3160 };
3161 }
3162 }
3163
3164 #[test]
3165 fn should_allow_to_create_map_from_data() {
3166 let struct_data_type = DataType::Struct(Fields::from(vec![
3167 Field::new(Field::MAP_KEY_FIELD_DEFAULT_NAME, DataType::Int32, false),
3168 Field::new(Field::MAP_VALUE_FIELD_DEFAULT_NAME, DataType::Utf8, true),
3169 ]));
3170
3171 let key_array_data = valid_non_nullable_int32_array_data(2);
3172 let values_array_data = valid_string_array_data(2);
3173
3174 let struct_data = {
3175 let builder = ArrayDataBuilder::new(struct_data_type.clone())
3176 .len(2)
3177 .nulls(None)
3178 .child_data(vec![key_array_data, values_array_data]);
3179
3180 builder.build().unwrap()
3181 };
3182
3183 let results = test_both_builder_and_array_data(
3184 DataType::Map(
3185 Field::new(
3186 Field::MAP_ENTRIES_FIELD_DEFAULT_NAME,
3187 struct_data_type,
3188 false,
3189 )
3190 .into(),
3191 false,
3192 ),
3193 1,
3194 None,
3195 0,
3196 vec![
3197 OffsetBuffer::<i32>::from_lengths(vec![2])
3198 .into_inner()
3199 .into(),
3200 ],
3201 vec![struct_data],
3202 );
3203
3204 for result in results {
3205 result.expect("should be able to create map ArrayData");
3206 }
3207 }
3208
3209 fn valid_string_array_data(length: usize) -> ArrayData {
3210 let offsets = OffsetBuffer::<i32>::from_lengths(vec![0; length])
3211 .into_inner()
3212 .into_inner();
3213 let empty_bytes = Buffer::default();
3214
3215 let builder = ArrayDataBuilder::new(DataType::Utf8)
3216 .len(length)
3217 .buffers(vec![offsets, empty_bytes])
3218 .nulls(None);
3219
3220 builder.build().unwrap()
3221 }
3222
3223 fn valid_non_nullable_int32_array_data(length: usize) -> ArrayData {
3224 let builder = ArrayDataBuilder::new(DataType::Int32)
3225 .len(length)
3226 .nulls(None)
3227 .buffers(vec![
3228 ScalarBuffer::<i32>::from(vec![1; length]).into_inner(),
3229 ]);
3230
3231 builder.build().unwrap()
3232 }
3233
3234 #[test]
3235 fn empty_and_null_map_array_should_pass_validation() {
3236 let dt = DataType::Map(
3237 Field::new(
3238 Field::MAP_ENTRIES_FIELD_DEFAULT_NAME,
3239 DataType::Struct(Fields::from(vec![
3240 Field::new(Field::MAP_KEY_FIELD_DEFAULT_NAME, DataType::Int32, false),
3241 Field::new(Field::MAP_VALUE_FIELD_DEFAULT_NAME, DataType::Utf8, true),
3242 ])),
3243 false,
3244 )
3245 .into(),
3246 false,
3247 );
3248
3249 ArrayData::new_empty(&dt).validate_full().unwrap();
3250 ArrayData::new_null(&dt, 1).validate_full().unwrap();
3251 }
3252
3253 fn test_both_builder_and_array_data(
3254 data_type: DataType,
3255 len: usize,
3256 null_bit_buffer: Option<Buffer>,
3257 offset: usize,
3258 buffers: Vec<Buffer>,
3259 child_data: Vec<ArrayData>,
3260 ) -> [Result<ArrayData, ArrowError>; 2] {
3261 let from_builder_res = ArrayData::builder(data_type.clone())
3262 .len(len)
3263 .add_buffers(buffers.clone())
3264 .null_bit_buffer(null_bit_buffer.clone())
3265 .offset(offset)
3266 .child_data(child_data.clone())
3267 .build();
3268
3269 let from_try_new_res =
3270 ArrayData::try_new(data_type, len, null_bit_buffer, offset, buffers, child_data);
3271
3272 [from_builder_res, from_try_new_res]
3273 }
3274}