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 {
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 &mut self.child_data {
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 Some(mask) = mask else {
1469 return match child.null_count() {
1470 0 => Ok(()),
1471 _ => Err(ArrowError::InvalidArgumentError(format!(
1472 "non-nullable child of type {} contains nulls not present in parent {}",
1473 child.data_type, self.data_type
1474 ))),
1475 };
1476 };
1477
1478 match child.nulls() {
1479 Some(nulls) if !mask.contains(nulls) => Err(ArrowError::InvalidArgumentError(format!(
1480 "non-nullable child of type {} contains nulls not present in parent",
1481 child.data_type
1482 ))),
1483 _ => Ok(()),
1484 }
1485 }
1486
1487 pub fn validate_values(&self) -> Result<(), ArrowError> {
1493 match &self.data_type {
1494 DataType::Utf8 => self.validate_utf8::<i32>(),
1495 DataType::LargeUtf8 => self.validate_utf8::<i64>(),
1496 DataType::Binary => self.validate_offsets_full::<i32>(self.buffers[1].len()),
1497 DataType::LargeBinary => self.validate_offsets_full::<i64>(self.buffers[1].len()),
1498 DataType::BinaryView => {
1499 let views = self.typed_buffer::<u128>(0, self.len)?;
1500 validate_binary_view(views, &self.buffers[1..])
1501 }
1502 DataType::Utf8View => {
1503 let views = self.typed_buffer::<u128>(0, self.len)?;
1504 validate_string_view(views, &self.buffers[1..])
1505 }
1506 DataType::List(_) | DataType::Map(_, _) => {
1507 let child = &self.child_data[0];
1508 self.validate_offsets_full::<i32>(child.len)
1509 }
1510 DataType::LargeList(_) => {
1511 let child = &self.child_data[0];
1512 self.validate_offsets_full::<i64>(child.len)
1513 }
1514 DataType::Union(_, _) => {
1515 Ok(())
1521 }
1522 DataType::Dictionary(key_type, _value_type) => {
1523 let dictionary_length: i64 = self.child_data[0].len.try_into().unwrap();
1524 let max_value = dictionary_length - 1;
1525 match key_type.as_ref() {
1526 DataType::UInt8 => self.check_bounds::<u8>(max_value),
1527 DataType::UInt16 => self.check_bounds::<u16>(max_value),
1528 DataType::UInt32 => self.check_bounds::<u32>(max_value),
1529 DataType::UInt64 => self.check_bounds::<u64>(max_value),
1530 DataType::Int8 => self.check_bounds::<i8>(max_value),
1531 DataType::Int16 => self.check_bounds::<i16>(max_value),
1532 DataType::Int32 => self.check_bounds::<i32>(max_value),
1533 DataType::Int64 => self.check_bounds::<i64>(max_value),
1534 _ => unreachable!(),
1535 }
1536 }
1537 DataType::RunEndEncoded(run_ends, _values) => {
1538 let run_ends_data = self.child_data()[0].clone();
1539 match run_ends.data_type() {
1540 DataType::Int16 => run_ends_data.check_run_ends::<i16>(),
1541 DataType::Int32 => run_ends_data.check_run_ends::<i32>(),
1542 DataType::Int64 => run_ends_data.check_run_ends::<i64>(),
1543 _ => unreachable!(),
1544 }
1545 }
1546 _ => {
1547 Ok(())
1549 }
1550 }
1551 }
1552
1553 fn validate_each_offset<T, V>(&self, offset_limit: usize, validate: V) -> Result<(), ArrowError>
1564 where
1565 T: ArrowNativeType + TryInto<usize> + num_traits::Num + std::fmt::Display,
1566 V: Fn(usize, Range<usize>) -> Result<(), ArrowError>,
1567 {
1568 self.typed_offsets::<T>()?
1569 .iter()
1570 .enumerate()
1571 .map(|(i, x)| {
1572 let r = x.to_usize().ok_or_else(|| {
1574 ArrowError::InvalidArgumentError(format!(
1575 "Offset invariant failure: Could not convert offset {x} to usize at position {i}"))}
1576 );
1577 match r {
1579 Ok(n) if n <= offset_limit => Ok((i, n)),
1580 Ok(_) => Err(ArrowError::InvalidArgumentError(format!(
1581 "Offset invariant failure: offset at position {i} out of bounds: {x} > {offset_limit}"))
1582 ),
1583 Err(e) => Err(e),
1584 }
1585 })
1586 .scan(0_usize, |start, end| {
1587 match end {
1589 Ok((i, end)) if *start <= end => {
1590 let range = Some(Ok((i, *start..end)));
1591 *start = end;
1592 range
1593 }
1594 Ok((i, end)) => Some(Err(ArrowError::InvalidArgumentError(format!(
1595 "Offset invariant failure: non-monotonic offset at slot {}: {} > {}",
1596 i - 1, start, end))
1597 )),
1598 Err(err) => Some(Err(err)),
1599 }
1600 })
1601 .skip(1) .try_for_each(|res: Result<(usize, Range<usize>), ArrowError>| {
1603 let (item_index, range) = res?;
1604 validate(item_index-1, range)
1605 })
1606 }
1607
1608 fn validate_utf8<T>(&self) -> Result<(), ArrowError>
1611 where
1612 T: ArrowNativeType + TryInto<usize> + num_traits::Num + std::fmt::Display,
1613 {
1614 let values_buffer = &self.buffers[1].as_slice();
1615 if let Ok(values_str) = std::str::from_utf8(values_buffer) {
1616 self.validate_each_offset::<T, _>(values_buffer.len(), |string_index, range| {
1618 if !values_str.is_char_boundary(range.start)
1619 || !values_str.is_char_boundary(range.end)
1620 {
1621 return Err(ArrowError::InvalidArgumentError(format!(
1622 "incomplete utf-8 byte sequence from index {string_index}"
1623 )));
1624 }
1625 Ok(())
1626 })
1627 } else {
1628 self.validate_each_offset::<T, _>(values_buffer.len(), |string_index, range| {
1630 std::str::from_utf8(&values_buffer[range.clone()]).map_err(|e| {
1631 ArrowError::InvalidArgumentError(format!(
1632 "Invalid UTF8 sequence at string index {string_index} ({range:?}): {e}"
1633 ))
1634 })?;
1635 Ok(())
1636 })
1637 }
1638 }
1639
1640 fn validate_offsets_full<T>(&self, offset_limit: usize) -> Result<(), ArrowError>
1643 where
1644 T: ArrowNativeType + TryInto<usize> + num_traits::Num + std::fmt::Display,
1645 {
1646 self.validate_each_offset::<T, _>(offset_limit, |_string_index, _range| {
1647 Ok(())
1650 })
1651 }
1652
1653 fn check_bounds<T>(&self, max_value: i64) -> Result<(), ArrowError>
1656 where
1657 T: ArrowNativeType + TryInto<i64> + num_traits::Num + std::fmt::Display,
1658 {
1659 let required_len = checked_len_plus_offset(&self.data_type, self.len, self.offset)?;
1660 let buffer = &self.buffers[0];
1661
1662 assert!(buffer.len() / mem::size_of::<T>() >= required_len);
1665
1666 let indexes: &[T] = &buffer.typed_data::<T>()[self.offset..required_len];
1668
1669 indexes.iter().enumerate().try_for_each(|(i, &dict_index)| {
1670 if self.is_null(i) {
1672 return Ok(());
1673 }
1674 let dict_index: i64 = dict_index.try_into().map_err(|_| {
1675 ArrowError::InvalidArgumentError(format!(
1676 "Value at position {i} out of bounds: {dict_index} (can not convert to i64)"
1677 ))
1678 })?;
1679
1680 if dict_index < 0 || dict_index > max_value {
1681 return Err(ArrowError::InvalidArgumentError(format!(
1682 "Value at position {i} out of bounds: {dict_index} (should be in [0, {max_value}])"
1683 )));
1684 }
1685 Ok(())
1686 })
1687 }
1688
1689 fn check_run_ends<T>(&self) -> Result<(), ArrowError>
1691 where
1692 T: ArrowNativeType + TryInto<i64> + num_traits::Num + std::fmt::Display,
1693 {
1694 let values = self.typed_buffer::<T>(0, self.len)?;
1695 let mut prev_value = 0_i64;
1696 values.iter().enumerate().try_for_each(|(ix, &inp_value)| {
1697 let value: i64 = inp_value.try_into().map_err(|_| {
1698 ArrowError::InvalidArgumentError(format!(
1699 "Value at position {ix} out of bounds: {inp_value} (can not convert to i64)"
1700 ))
1701 })?;
1702 if value <= 0_i64 {
1703 return Err(ArrowError::InvalidArgumentError(format!(
1704 "The values in run_ends array should be strictly positive. Found value {value} at index {ix} that does not match the criteria."
1705 )));
1706 }
1707 if ix > 0 && value <= prev_value {
1708 return Err(ArrowError::InvalidArgumentError(format!(
1709 "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."
1710 )));
1711 }
1712
1713 prev_value = value;
1714 Ok(())
1715 })?;
1716
1717 let len_plus_offset = checked_len_plus_offset(&self.data_type, self.len, self.offset)?;
1718 if prev_value.as_usize() < len_plus_offset {
1719 return Err(ArrowError::InvalidArgumentError(format!(
1720 "The offset + length of array should be less or equal to last value in the run_ends array. The last value of run_ends array is {prev_value} and offset + length of array is {len_plus_offset}."
1721 )));
1722 }
1723 Ok(())
1724 }
1725
1726 pub fn ptr_eq(&self, other: &Self) -> bool {
1730 if self.offset != other.offset
1731 || self.len != other.len
1732 || self.data_type != other.data_type
1733 || self.buffers.len() != other.buffers.len()
1734 || self.child_data.len() != other.child_data.len()
1735 {
1736 return false;
1737 }
1738
1739 match (&self.nulls, &other.nulls) {
1740 (Some(a), Some(b)) if !a.inner().ptr_eq(b.inner()) => return false,
1741 (Some(_), None) | (None, Some(_)) => return false,
1742 _ => {}
1743 }
1744
1745 if !self
1746 .buffers
1747 .iter()
1748 .zip(other.buffers.iter())
1749 .all(|(a, b)| a.as_ptr() == b.as_ptr())
1750 {
1751 return false;
1752 }
1753
1754 self.child_data
1755 .iter()
1756 .zip(other.child_data.iter())
1757 .all(|(a, b)| a.ptr_eq(b))
1758 }
1759
1760 pub fn into_builder(self) -> ArrayDataBuilder {
1762 self.into()
1763 }
1764
1765 #[cfg(feature = "pool")]
1772 pub fn claim(&self, pool: &dyn arrow_buffer::MemoryPool) {
1773 for buffer in &self.buffers {
1775 buffer.claim(pool);
1776 }
1777
1778 if let Some(nulls) = &self.nulls {
1780 nulls.claim(pool);
1781 }
1782
1783 for child in &self.child_data {
1785 child.claim(pool);
1786 }
1787 }
1788}
1789
1790pub fn layout(data_type: &DataType) -> DataTypeLayout {
1793 use arrow_schema::IntervalUnit::*;
1796
1797 match data_type {
1798 DataType::Null => DataTypeLayout {
1799 buffers: vec![],
1800 can_contain_null_mask: false,
1801 variadic: false,
1802 },
1803 DataType::Boolean => DataTypeLayout {
1804 buffers: vec![BufferSpec::BitMap],
1805 can_contain_null_mask: true,
1806 variadic: false,
1807 },
1808 DataType::Int8 => DataTypeLayout::new_fixed_width::<i8>(),
1809 DataType::Int16 => DataTypeLayout::new_fixed_width::<i16>(),
1810 DataType::Int32 => DataTypeLayout::new_fixed_width::<i32>(),
1811 DataType::Int64 => DataTypeLayout::new_fixed_width::<i64>(),
1812 DataType::UInt8 => DataTypeLayout::new_fixed_width::<u8>(),
1813 DataType::UInt16 => DataTypeLayout::new_fixed_width::<u16>(),
1814 DataType::UInt32 => DataTypeLayout::new_fixed_width::<u32>(),
1815 DataType::UInt64 => DataTypeLayout::new_fixed_width::<u64>(),
1816 DataType::Float16 => DataTypeLayout::new_fixed_width::<half::f16>(),
1817 DataType::Float32 => DataTypeLayout::new_fixed_width::<f32>(),
1818 DataType::Float64 => DataTypeLayout::new_fixed_width::<f64>(),
1819 DataType::Timestamp(_, _) => DataTypeLayout::new_fixed_width::<i64>(),
1820 DataType::Date32 => DataTypeLayout::new_fixed_width::<i32>(),
1821 DataType::Date64 => DataTypeLayout::new_fixed_width::<i64>(),
1822 DataType::Time32(_) => DataTypeLayout::new_fixed_width::<i32>(),
1823 DataType::Time64(_) => DataTypeLayout::new_fixed_width::<i64>(),
1824 DataType::Interval(YearMonth) => DataTypeLayout::new_fixed_width::<i32>(),
1825 DataType::Interval(DayTime) => DataTypeLayout::new_fixed_width::<IntervalDayTime>(),
1826 DataType::Interval(MonthDayNano) => {
1827 DataTypeLayout::new_fixed_width::<IntervalMonthDayNano>()
1828 }
1829 DataType::Duration(_) => DataTypeLayout::new_fixed_width::<i64>(),
1830 DataType::Decimal32(_, _) => DataTypeLayout::new_fixed_width::<i32>(),
1831 DataType::Decimal64(_, _) => DataTypeLayout::new_fixed_width::<i64>(),
1832 DataType::Decimal128(_, _) => DataTypeLayout::new_fixed_width::<i128>(),
1833 DataType::Decimal256(_, _) => DataTypeLayout::new_fixed_width::<i256>(),
1834 DataType::FixedSizeBinary(size) => {
1835 let spec = BufferSpec::FixedWidth {
1836 byte_width: (*size).try_into().unwrap(),
1837 alignment: mem::align_of::<u8>(),
1838 };
1839 DataTypeLayout {
1840 buffers: vec![spec],
1841 can_contain_null_mask: true,
1842 variadic: false,
1843 }
1844 }
1845 DataType::Binary => DataTypeLayout::new_binary::<i32>(),
1846 DataType::LargeBinary => DataTypeLayout::new_binary::<i64>(),
1847 DataType::Utf8 => DataTypeLayout::new_binary::<i32>(),
1848 DataType::LargeUtf8 => DataTypeLayout::new_binary::<i64>(),
1849 DataType::BinaryView | DataType::Utf8View => DataTypeLayout::new_view(),
1850 DataType::FixedSizeList(_, _) => DataTypeLayout::new_nullable_empty(), DataType::List(_) => DataTypeLayout::new_fixed_width::<i32>(),
1852 DataType::ListView(_) => DataTypeLayout::new_list_view::<i32>(),
1853 DataType::LargeListView(_) => DataTypeLayout::new_list_view::<i64>(),
1854 DataType::LargeList(_) => DataTypeLayout::new_fixed_width::<i64>(),
1855 DataType::Map(_, _) => DataTypeLayout::new_fixed_width::<i32>(),
1856 DataType::Struct(_) => DataTypeLayout::new_nullable_empty(), DataType::RunEndEncoded(_, _) => DataTypeLayout::new_empty(), DataType::Union(_, mode) => {
1859 let type_ids = BufferSpec::FixedWidth {
1860 byte_width: mem::size_of::<i8>(),
1861 alignment: mem::align_of::<i8>(),
1862 };
1863
1864 DataTypeLayout {
1865 buffers: match mode {
1866 UnionMode::Sparse => {
1867 vec![type_ids]
1868 }
1869 UnionMode::Dense => {
1870 vec![
1871 type_ids,
1872 BufferSpec::FixedWidth {
1873 byte_width: mem::size_of::<i32>(),
1874 alignment: mem::align_of::<i32>(),
1875 },
1876 ]
1877 }
1878 },
1879 can_contain_null_mask: false,
1880 variadic: false,
1881 }
1882 }
1883 DataType::Dictionary(key_type, _value_type) => layout(key_type),
1884 }
1885}
1886
1887#[derive(Debug, PartialEq, Eq)]
1889pub struct DataTypeLayout {
1891 pub buffers: Vec<BufferSpec>,
1893
1894 pub can_contain_null_mask: bool,
1896
1897 pub variadic: bool,
1901}
1902
1903impl DataTypeLayout {
1904 pub fn new_fixed_width<T>() -> Self {
1906 Self {
1907 buffers: vec![BufferSpec::FixedWidth {
1908 byte_width: mem::size_of::<T>(),
1909 alignment: mem::align_of::<T>(),
1910 }],
1911 can_contain_null_mask: true,
1912 variadic: false,
1913 }
1914 }
1915
1916 pub fn new_nullable_empty() -> Self {
1919 Self {
1920 buffers: vec![],
1921 can_contain_null_mask: true,
1922 variadic: false,
1923 }
1924 }
1925
1926 pub fn new_empty() -> Self {
1929 Self {
1930 buffers: vec![],
1931 can_contain_null_mask: false,
1932 variadic: false,
1933 }
1934 }
1935
1936 pub fn new_binary<T>() -> Self {
1940 Self {
1941 buffers: vec![
1942 BufferSpec::FixedWidth {
1944 byte_width: mem::size_of::<T>(),
1945 alignment: mem::align_of::<T>(),
1946 },
1947 BufferSpec::VariableWidth,
1949 ],
1950 can_contain_null_mask: true,
1951 variadic: false,
1952 }
1953 }
1954
1955 pub fn new_view() -> Self {
1957 Self {
1958 buffers: vec![BufferSpec::FixedWidth {
1959 byte_width: mem::size_of::<u128>(),
1960 alignment: mem::align_of::<u128>(),
1961 }],
1962 can_contain_null_mask: true,
1963 variadic: true,
1964 }
1965 }
1966
1967 pub fn new_list_view<T>() -> Self {
1969 Self {
1970 buffers: vec![
1971 BufferSpec::FixedWidth {
1972 byte_width: mem::size_of::<T>(),
1973 alignment: mem::align_of::<T>(),
1974 },
1975 BufferSpec::FixedWidth {
1976 byte_width: mem::size_of::<T>(),
1977 alignment: mem::align_of::<T>(),
1978 },
1979 ],
1980 can_contain_null_mask: true,
1981 variadic: false,
1982 }
1983 }
1984}
1985
1986#[derive(Debug, PartialEq, Eq)]
1988pub enum BufferSpec {
1989 FixedWidth {
2000 byte_width: usize,
2002 alignment: usize,
2004 },
2005 VariableWidth,
2007 BitMap,
2013 AlwaysNull,
2016}
2017
2018impl PartialEq for ArrayData {
2019 fn eq(&self, other: &Self) -> bool {
2020 equal::equal(self, other)
2021 }
2022}
2023
2024#[derive(Debug, Clone)]
2043#[doc(hidden)]
2044pub struct UnsafeFlag(bool);
2045
2046impl UnsafeFlag {
2047 #[inline]
2051 pub const fn new() -> Self {
2052 Self(false)
2053 }
2054
2055 #[inline]
2065 pub unsafe fn set(&mut self, val: bool) {
2066 self.0 = val;
2067 }
2068
2069 #[inline]
2071 pub fn get(&self) -> bool {
2072 self.0
2073 }
2074}
2075
2076impl Default for UnsafeFlag {
2078 fn default() -> Self {
2079 Self::new()
2080 }
2081}
2082
2083#[derive(Debug)]
2085pub struct ArrayDataBuilder {
2086 data_type: DataType,
2087 len: usize,
2088 null_count: Option<usize>,
2089 null_bit_buffer: Option<Buffer>,
2090 nulls: Option<NullBuffer>,
2091 offset: usize,
2092 buffers: Vec<Buffer>,
2093 child_data: Vec<ArrayData>,
2094 align_buffers: bool,
2098 skip_validation: UnsafeFlag,
2108}
2109
2110impl ArrayDataBuilder {
2111 #[inline]
2112 pub const fn new(data_type: DataType) -> Self {
2114 Self {
2115 data_type,
2116 len: 0,
2117 null_count: None,
2118 null_bit_buffer: None,
2119 nulls: None,
2120 offset: 0,
2121 buffers: vec![],
2122 child_data: vec![],
2123 align_buffers: false,
2124 skip_validation: UnsafeFlag::new(),
2125 }
2126 }
2127
2128 pub fn data_type(self, data_type: DataType) -> Self {
2130 Self { data_type, ..self }
2131 }
2132
2133 #[inline]
2134 pub const fn len(mut self, n: usize) -> Self {
2136 self.len = n;
2137 self
2138 }
2139
2140 pub fn nulls(mut self, nulls: Option<NullBuffer>) -> Self {
2142 self.nulls = nulls;
2143 self.null_count = None;
2144 self.null_bit_buffer = None;
2145 self
2146 }
2147
2148 pub fn null_count(mut self, null_count: usize) -> Self {
2150 self.null_count = Some(null_count);
2151 self
2152 }
2153
2154 pub fn null_bit_buffer(mut self, buf: Option<Buffer>) -> Self {
2156 self.nulls = None;
2157 self.null_bit_buffer = buf;
2158 self
2159 }
2160
2161 #[inline]
2163 pub const fn offset(mut self, n: usize) -> Self {
2164 self.offset = n;
2165 self
2166 }
2167
2168 pub fn buffers(mut self, v: Vec<Buffer>) -> Self {
2170 self.buffers = v;
2171 self
2172 }
2173
2174 pub fn add_buffer(mut self, b: Buffer) -> Self {
2176 self.buffers.push(b);
2177 self
2178 }
2179
2180 pub fn add_buffers<I: IntoIterator<Item = Buffer>>(mut self, bs: I) -> Self {
2182 self.buffers.extend(bs);
2183 self
2184 }
2185
2186 pub fn child_data(mut self, v: Vec<ArrayData>) -> Self {
2188 self.child_data = v;
2189 self
2190 }
2191
2192 pub fn add_child_data(mut self, r: ArrayData) -> Self {
2194 self.child_data.push(r);
2195 self
2196 }
2197
2198 pub unsafe fn build_unchecked(self) -> ArrayData {
2214 unsafe { self.skip_validation(true) }.build().unwrap()
2215 }
2216
2217 pub fn build(self) -> Result<ArrayData, ArrowError> {
2226 let Self {
2227 data_type,
2228 len,
2229 null_count,
2230 null_bit_buffer,
2231 nulls,
2232 offset,
2233 buffers,
2234 child_data,
2235 align_buffers,
2236 skip_validation,
2237 } = self;
2238
2239 let nulls = nulls
2240 .or_else(|| {
2241 let buffer = null_bit_buffer?;
2242 let buffer = BooleanBuffer::new(buffer, offset, len);
2243 Some(match null_count {
2244 Some(n) => {
2245 unsafe { NullBuffer::new_unchecked(buffer, n) }
2247 }
2248 None => NullBuffer::new(buffer),
2249 })
2250 })
2251 .filter(|b| b.null_count() != 0);
2252
2253 let mut data = ArrayData {
2254 data_type,
2255 len,
2256 offset,
2257 buffers,
2258 child_data,
2259 nulls,
2260 };
2261
2262 if align_buffers {
2263 data.align_buffers();
2264 }
2265
2266 if !skip_validation.get() || cfg!(feature = "force_validate") {
2268 data.validate_data()?;
2269 }
2270 Ok(data)
2271 }
2272
2273 pub fn align_buffers(mut self, align_buffers: bool) -> Self {
2289 self.align_buffers = align_buffers;
2290 self
2291 }
2292
2293 pub unsafe fn skip_validation(mut self, skip_validation: bool) -> Self {
2307 unsafe {
2308 self.skip_validation.set(skip_validation);
2309 }
2310 self
2311 }
2312}
2313
2314impl From<ArrayData> for ArrayDataBuilder {
2315 fn from(d: ArrayData) -> Self {
2316 Self {
2317 data_type: d.data_type,
2318 len: d.len,
2319 offset: d.offset,
2320 buffers: d.buffers,
2321 child_data: d.child_data,
2322 nulls: d.nulls,
2323 null_bit_buffer: None,
2324 null_count: None,
2325 align_buffers: false,
2326 skip_validation: UnsafeFlag::new(),
2327 }
2328 }
2329}
2330
2331pub(crate) fn get_fixed_size_binary_width(data_type: &DataType) -> usize {
2336 match data_type {
2337 DataType::FixedSizeBinary(i) => {
2338 if *i < 0 {
2339 panic!("cannot compare FixedSizeBinary({})", *i);
2340 }
2341 *i as usize
2342 }
2343 _ => unreachable!(),
2344 }
2345}
2346
2347#[cfg(test)]
2348mod tests {
2349 use super::*;
2350 use crate::ByteView;
2351 use arrow_buffer::{OffsetBuffer, ScalarBuffer};
2352 use arrow_schema::{Field, Fields};
2353
2354 fn make_i32_buffer(n: usize) -> Buffer {
2358 Buffer::from_slice_ref(vec![42i32; n])
2359 }
2360
2361 fn make_f32_buffer(n: usize) -> Buffer {
2363 Buffer::from_slice_ref(vec![42f32; n])
2364 }
2365
2366 #[test]
2367 fn test_builder() {
2368 let v = (0..25).collect::<Vec<i32>>();
2370 let b1 = Buffer::from_slice_ref(&v);
2371 let arr_data = ArrayData::builder(DataType::Int32)
2372 .len(20)
2373 .offset(5)
2374 .add_buffer(b1)
2375 .null_bit_buffer(Some(Buffer::from([
2376 0b01011111, 0b10110101, 0b01100011, 0b00011110,
2377 ])))
2378 .build()
2379 .unwrap();
2380
2381 assert_eq!(20, arr_data.len());
2382 assert_eq!(10, arr_data.null_count());
2383 assert_eq!(5, arr_data.offset());
2384 assert_eq!(1, arr_data.buffers().len());
2385 assert_eq!(
2386 Buffer::from_slice_ref(&v).as_slice(),
2387 arr_data.buffers()[0].as_slice()
2388 );
2389 }
2390
2391 #[test]
2392 fn test_builder_with_child_data() {
2393 let child_arr_data = ArrayData::try_new(
2394 DataType::Int32,
2395 5,
2396 None,
2397 0,
2398 vec![Buffer::from_slice_ref([1i32, 2, 3, 4, 5])],
2399 vec![],
2400 )
2401 .unwrap();
2402
2403 let field = Arc::new(Field::new("x", DataType::Int32, true));
2404 let data_type = DataType::Struct(vec![field].into());
2405
2406 let arr_data = ArrayData::builder(data_type)
2407 .len(5)
2408 .offset(0)
2409 .add_child_data(child_arr_data.clone())
2410 .build()
2411 .unwrap();
2412
2413 assert_eq!(5, arr_data.len());
2414 assert_eq!(1, arr_data.child_data().len());
2415 assert_eq!(child_arr_data, arr_data.child_data()[0]);
2416 }
2417
2418 #[test]
2419 fn test_null_count() {
2420 let mut bit_v: [u8; 2] = [0; 2];
2421 bit_util::set_bit(&mut bit_v, 0);
2422 bit_util::set_bit(&mut bit_v, 3);
2423 bit_util::set_bit(&mut bit_v, 10);
2424 let arr_data = ArrayData::builder(DataType::Int32)
2425 .len(16)
2426 .add_buffer(make_i32_buffer(16))
2427 .null_bit_buffer(Some(Buffer::from(bit_v)))
2428 .build()
2429 .unwrap();
2430 assert_eq!(13, arr_data.null_count());
2431
2432 let mut bit_v: [u8; 2] = [0; 2];
2434 bit_util::set_bit(&mut bit_v, 0);
2435 bit_util::set_bit(&mut bit_v, 3);
2436 bit_util::set_bit(&mut bit_v, 10);
2437 let arr_data = ArrayData::builder(DataType::Int32)
2438 .len(12)
2439 .offset(2)
2440 .add_buffer(make_i32_buffer(14)) .null_bit_buffer(Some(Buffer::from(bit_v)))
2442 .build()
2443 .unwrap();
2444 assert_eq!(10, arr_data.null_count());
2445 }
2446
2447 #[test]
2448 fn test_null_buffer_ref() {
2449 let mut bit_v: [u8; 2] = [0; 2];
2450 bit_util::set_bit(&mut bit_v, 0);
2451 bit_util::set_bit(&mut bit_v, 3);
2452 bit_util::set_bit(&mut bit_v, 10);
2453 let arr_data = ArrayData::builder(DataType::Int32)
2454 .len(16)
2455 .add_buffer(make_i32_buffer(16))
2456 .null_bit_buffer(Some(Buffer::from(bit_v)))
2457 .build()
2458 .unwrap();
2459 assert!(arr_data.nulls().is_some());
2460 assert_eq!(&bit_v, arr_data.nulls().unwrap().validity());
2461 }
2462
2463 #[test]
2464 fn test_slice() {
2465 let mut bit_v: [u8; 2] = [0; 2];
2466 bit_util::set_bit(&mut bit_v, 0);
2467 bit_util::set_bit(&mut bit_v, 3);
2468 bit_util::set_bit(&mut bit_v, 10);
2469 let data = ArrayData::builder(DataType::Int32)
2470 .len(16)
2471 .add_buffer(make_i32_buffer(16))
2472 .null_bit_buffer(Some(Buffer::from(bit_v)))
2473 .build()
2474 .unwrap();
2475 let new_data = data.slice(1, 15);
2476 assert_eq!(data.len() - 1, new_data.len());
2477 assert_eq!(1, new_data.offset());
2478 assert_eq!(data.null_count(), new_data.null_count());
2479
2480 let new_data = new_data.slice(1, 14);
2482 assert_eq!(data.len() - 2, new_data.len());
2483 assert_eq!(2, new_data.offset());
2484 assert_eq!(data.null_count() - 1, new_data.null_count());
2485 }
2486
2487 #[test]
2488 #[should_panic(expected = "offset + length overflow")]
2489 fn test_slice_panics_on_offset_length_overflow() {
2490 let data = ArrayData::builder(DataType::Int32)
2491 .len(4)
2492 .add_buffer(make_i32_buffer(4))
2493 .build()
2494 .unwrap();
2495 let sliced = data.slice(1, 3);
2496
2497 sliced.slice(1, usize::MAX);
2498 }
2499
2500 #[test]
2501 fn test_typed_offsets_length_overflow() {
2502 let data = ArrayData {
2503 data_type: DataType::Binary,
2504 len: usize::MAX,
2505 offset: 0,
2506 buffers: vec![Buffer::from_slice_ref([0_i32])],
2507 child_data: vec![],
2508 nulls: None,
2509 };
2510 let err = data.typed_offsets::<i32>().unwrap_err();
2511
2512 assert_eq!(
2513 err.to_string(),
2514 format!(
2515 "Invalid argument error: Length {} with offset 1 overflows usize for Binary",
2516 usize::MAX
2517 )
2518 );
2519 }
2520
2521 #[test]
2522 fn test_validate_typed_buffer_length_overflow() {
2523 let data = ArrayData {
2524 data_type: DataType::Binary,
2525 len: 0,
2526 offset: 2,
2527 buffers: vec![Buffer::from_slice_ref([0_i32])],
2528 child_data: vec![],
2529 nulls: None,
2530 };
2531 let err = data.typed_buffer::<i32>(0, usize::MAX).unwrap_err();
2532
2533 assert_eq!(
2534 err.to_string(),
2535 format!(
2536 "Invalid argument error: Length {} with offset 2 overflows usize for Binary",
2537 usize::MAX
2538 )
2539 );
2540 }
2541
2542 fn try_new_binary_length_offset_overflow() -> Result<ArrayData, ArrowError> {
2544 ArrayData::try_new(
2545 DataType::Binary,
2546 usize::MAX,
2547 None,
2548 1,
2549 vec![
2550 Buffer::from_slice_ref([0_i32]),
2551 Buffer::from_iter(std::iter::empty::<u8>()),
2552 ],
2553 vec![],
2554 )
2555 }
2556
2557 #[cfg(not(feature = "force_validate"))]
2558 #[test]
2559 fn test_try_new_length_offset_overflow() {
2560 let err = try_new_binary_length_offset_overflow().unwrap_err();
2561
2562 assert_eq!(
2563 err.to_string(),
2564 format!(
2565 "Invalid argument error: Length {} with offset 1 overflows usize for Binary",
2566 usize::MAX
2567 )
2568 );
2569 }
2570
2571 #[cfg(feature = "force_validate")]
2572 #[test]
2573 #[should_panic(
2574 expected = "Length 18446744073709551615 with offset 1 overflows usize for Binary"
2575 )]
2576 fn test_try_new_length_offset_overflow_force_validate() {
2577 try_new_binary_length_offset_overflow().unwrap();
2578 }
2579
2580 #[test]
2581 fn test_equality() {
2582 let int_data = ArrayData::builder(DataType::Int32)
2583 .len(1)
2584 .add_buffer(make_i32_buffer(1))
2585 .build()
2586 .unwrap();
2587
2588 let float_data = ArrayData::builder(DataType::Float32)
2589 .len(1)
2590 .add_buffer(make_f32_buffer(1))
2591 .build()
2592 .unwrap();
2593 assert_ne!(int_data, float_data);
2594 assert!(!int_data.ptr_eq(&float_data));
2595 assert!(int_data.ptr_eq(&int_data));
2596
2597 let int_data_clone = int_data.clone();
2598 assert_eq!(int_data, int_data_clone);
2599 assert!(int_data.ptr_eq(&int_data_clone));
2600 assert!(int_data_clone.ptr_eq(&int_data));
2601
2602 let int_data_slice = int_data_clone.slice(1, 0);
2603 assert!(int_data_slice.ptr_eq(&int_data_slice));
2604 assert!(!int_data.ptr_eq(&int_data_slice));
2605 assert!(!int_data_slice.ptr_eq(&int_data));
2606
2607 let data_buffer = Buffer::from_slice_ref(b"abcdef");
2608 let offsets_buffer = Buffer::from_slice_ref([0_i32, 2_i32, 2_i32, 5_i32]);
2609 let string_data = ArrayData::try_new(
2610 DataType::Utf8,
2611 3,
2612 Some(Buffer::from_iter(vec![true, false, true])),
2613 0,
2614 vec![offsets_buffer, data_buffer],
2615 vec![],
2616 )
2617 .unwrap();
2618
2619 assert_ne!(float_data, string_data);
2620 assert!(!float_data.ptr_eq(&string_data));
2621
2622 assert!(string_data.ptr_eq(&string_data));
2623
2624 let string_data_cloned = string_data.clone();
2625 assert!(string_data_cloned.ptr_eq(&string_data));
2626 assert!(string_data.ptr_eq(&string_data_cloned));
2627
2628 let string_data_slice = string_data.slice(1, 2);
2629 assert!(string_data_slice.ptr_eq(&string_data_slice));
2630 assert!(!string_data_slice.ptr_eq(&string_data))
2631 }
2632
2633 #[test]
2634 fn test_slice_memory_size_view_payload_buffers() {
2635 for data_type in [DataType::Utf8View, DataType::BinaryView] {
2636 let inline_only = ArrayData::builder(data_type.clone())
2637 .len(2)
2638 .add_buffer(Buffer::from_vec(vec![0_u128; 2]))
2639 .build()
2640 .unwrap();
2641 assert_eq!(
2642 inline_only.get_slice_memory_size().unwrap(),
2643 2 * mem::size_of::<u128>()
2644 );
2645
2646 let mut first_payload = Vec::with_capacity(32);
2647 first_payload.extend_from_slice(b"first payload");
2648 let first_view =
2649 ByteView::new(first_payload.len().try_into().unwrap(), &first_payload[..4])
2650 .as_u128();
2651 let first_payload = Buffer::from_vec(first_payload);
2652 assert!(first_payload.capacity() > first_payload.len());
2653 let first_payload_capacity = first_payload.capacity();
2654
2655 let mut second_payload = Vec::with_capacity(64);
2656 second_payload.extend_from_slice(b"second payload");
2657 let second_view = ByteView::new(
2658 second_payload.len().try_into().unwrap(),
2659 &second_payload[..4],
2660 )
2661 .with_buffer_index(1)
2662 .as_u128();
2663 let second_payload = Buffer::from_vec(second_payload);
2664 assert!(second_payload.capacity() > second_payload.len());
2665 let second_payload_capacity = second_payload.capacity();
2666
2667 let data = ArrayData::builder(data_type)
2668 .len(3)
2669 .add_buffer(Buffer::from_vec(vec![first_view, 0_u128, second_view]))
2670 .add_buffer(first_payload)
2671 .add_buffer(second_payload)
2672 .build()
2673 .unwrap();
2674 let sliced = data.slice(1, 1);
2675
2676 assert_eq!(
2677 sliced.get_slice_memory_size().unwrap(),
2678 mem::size_of::<u128>() + first_payload_capacity + second_payload_capacity
2679 );
2680 }
2681 }
2682
2683 #[test]
2684 fn test_slice_memory_size_utf8_offset_buffer_len_plus_one() {
2685 let data_buffer = Buffer::from_slice_ref(b"helloworld");
2687 let offsets_buffer = Buffer::from_slice_ref([0_i32, 5_i32, 10_i32]);
2690 let array = ArrayData::try_new(
2691 DataType::Utf8,
2692 2,
2693 None,
2694 0,
2695 vec![offsets_buffer, data_buffer],
2696 vec![],
2697 )
2698 .unwrap();
2699 assert_eq!(array.get_slice_memory_size().unwrap(), 22); }
2701
2702 #[test]
2703 fn test_slice_memory_size_binary_offset_buffer_len_plus_one() {
2704 let data_buffer = Buffer::from_slice_ref([0u8, 1, 2, 3, 4]);
2707 let offsets_buffer = Buffer::from_slice_ref([0_i32, 2_i32, 5_i32]);
2709 let array = ArrayData::try_new(
2710 DataType::Binary,
2711 2,
2712 None,
2713 0,
2714 vec![offsets_buffer, data_buffer],
2715 vec![],
2716 )
2717 .unwrap();
2718 assert_eq!(array.get_slice_memory_size().unwrap(), 17); }
2720
2721 #[test]
2722 fn test_slice_memory_size() {
2723 let mut bit_v: [u8; 2] = [0; 2];
2724 bit_util::set_bit(&mut bit_v, 0);
2725 bit_util::set_bit(&mut bit_v, 3);
2726 bit_util::set_bit(&mut bit_v, 10);
2727 let data = ArrayData::builder(DataType::Int32)
2728 .len(16)
2729 .add_buffer(make_i32_buffer(16))
2730 .null_bit_buffer(Some(Buffer::from(bit_v)))
2731 .build()
2732 .unwrap();
2733 let new_data = data.slice(1, 14);
2734 assert_eq!(
2735 data.get_slice_memory_size().unwrap() - 8,
2736 new_data.get_slice_memory_size().unwrap()
2737 );
2738 let data_buffer = Buffer::from_slice_ref(b"abcdef");
2739 let offsets_buffer = Buffer::from_slice_ref([0_i32, 2_i32, 2_i32, 5_i32]);
2740 let string_data = ArrayData::try_new(
2741 DataType::Utf8,
2742 3,
2743 Some(Buffer::from_iter(vec![true, false, true])),
2744 0,
2745 vec![offsets_buffer, data_buffer],
2746 vec![],
2747 )
2748 .unwrap();
2749 let string_data_slice = string_data.slice(1, 2);
2750 assert_eq!(
2752 string_data.get_slice_memory_size().unwrap() - 6,
2753 string_data_slice.get_slice_memory_size().unwrap()
2754 );
2755 }
2756
2757 #[test]
2758 fn test_count_nulls() {
2759 let buffer = Buffer::from([0b00010110, 0b10011111]);
2760 let buffer = NullBuffer::new(BooleanBuffer::new(buffer, 0, 16));
2761 let count = count_nulls(Some(&buffer), 0, 16);
2762 assert_eq!(count, 7);
2763
2764 let count = count_nulls(Some(&buffer), 4, 8);
2765 assert_eq!(count, 3);
2766 }
2767
2768 #[test]
2769 fn test_contains_nulls() {
2770 let buffer: Buffer =
2771 MutableBuffer::from_iter([false, false, false, true, true, false]).into();
2772 let buffer = NullBuffer::new(BooleanBuffer::new(buffer, 0, 6));
2773 assert!(contains_nulls(Some(&buffer), 0, 6));
2774 assert!(contains_nulls(Some(&buffer), 0, 3));
2775 assert!(!contains_nulls(Some(&buffer), 3, 2));
2776 assert!(!contains_nulls(Some(&buffer), 0, 0));
2777 }
2778
2779 #[test]
2780 fn test_alignment() {
2781 let buffer = Buffer::from_vec(vec![1_i32, 2_i32, 3_i32]);
2782 let sliced = buffer.slice(1);
2783
2784 let mut data = ArrayData {
2785 data_type: DataType::Int32,
2786 len: 0,
2787 offset: 0,
2788 buffers: vec![buffer],
2789 child_data: vec![],
2790 nulls: None,
2791 };
2792 data.validate_full().unwrap();
2793
2794 data.buffers[0] = sliced;
2796 let err = data.validate().unwrap_err();
2797
2798 assert_eq!(
2799 err.to_string(),
2800 "Invalid argument error: Misaligned buffers[0] in array of type Int32, offset from expected alignment of 4 by 1"
2801 );
2802
2803 data.align_buffers();
2804 data.validate_full().unwrap();
2805 }
2806
2807 #[test]
2808 fn test_alignment_struct() {
2809 let buffer = Buffer::from_vec(vec![1_i32, 2_i32, 3_i32]);
2810 let sliced = buffer.slice(1);
2811
2812 let child_data = ArrayData {
2813 data_type: DataType::Int32,
2814 len: 0,
2815 offset: 0,
2816 buffers: vec![buffer],
2817 child_data: vec![],
2818 nulls: None,
2819 };
2820
2821 let schema = DataType::Struct(Fields::from(vec![Field::new("a", DataType::Int32, false)]));
2822 let mut data = ArrayData {
2823 data_type: schema,
2824 len: 0,
2825 offset: 0,
2826 buffers: vec![],
2827 child_data: vec![child_data],
2828 nulls: None,
2829 };
2830 data.validate_full().unwrap();
2831
2832 data.child_data[0].buffers[0] = sliced;
2834 let err = data.validate().unwrap_err();
2835
2836 assert_eq!(
2837 err.to_string(),
2838 "Invalid argument error: Misaligned buffers[0] in array of type Int32, offset from expected alignment of 4 by 1"
2839 );
2840
2841 data.align_buffers();
2842 data.validate_full().unwrap();
2843 }
2844
2845 #[test]
2846 fn test_null_view_types() {
2847 let array_len = 32;
2848 let array = ArrayData::new_null(&DataType::BinaryView, array_len);
2849 assert_eq!(array.len(), array_len);
2850 for i in 0..array.len() {
2851 assert!(array.is_null(i));
2852 }
2853
2854 let array = ArrayData::new_null(&DataType::Utf8View, array_len);
2855 assert_eq!(array.len(), array_len);
2856 for i in 0..array.len() {
2857 assert!(array.is_null(i));
2858 }
2859
2860 let array = ArrayData::new_null(
2861 &DataType::ListView(Arc::new(Field::new_list_field(DataType::Int32, true))),
2862 array_len,
2863 );
2864 assert_eq!(array.len(), array_len);
2865 for i in 0..array.len() {
2866 assert!(array.is_null(i));
2867 }
2868
2869 let array = ArrayData::new_null(
2870 &DataType::LargeListView(Arc::new(Field::new_list_field(DataType::Int32, true))),
2871 array_len,
2872 );
2873 assert_eq!(array.len(), array_len);
2874 for i in 0..array.len() {
2875 assert!(array.is_null(i));
2876 }
2877 }
2878
2879 #[test]
2881 fn test_dont_panic_on_bad_input_when_using_try_new() {
2882 let empty_bytes = Buffer::default();
2883
2884 let array_data = ArrayData::try_new(
2885 DataType::Utf8,
2886 1, None,
2888 0,
2889 vec![Buffer::from_vec(vec![0i32, 2i32]), empty_bytes],
2891 vec![],
2892 );
2893
2894 let res = array_data.expect_err("should get error");
2895
2896 assert_eq!(
2897 res.to_string(),
2898 format!("Invalid argument error: Last offset 2 of Utf8 is larger than values length 0",)
2899 );
2900 }
2901
2902 #[test]
2903 fn should_fail_validation_when_having_map_field_type_is_not_struct() {
2904 let map_field = Field::new("key", DataType::Int32, false);
2905
2906 let map_field_data = valid_non_nullable_int32_array_data(2);
2907
2908 let results = test_both_builder_and_array_data(
2909 DataType::Map(map_field.into(), false),
2910 1,
2911 None,
2912 0,
2913 vec![
2914 OffsetBuffer::<i32>::from_lengths(vec![2])
2915 .into_inner()
2916 .into(),
2917 ],
2918 vec![map_field_data],
2919 );
2920
2921 for result in results {
2922 let array_data_err = result.expect_err("should fail for non struct field");
2923
2924 match array_data_err {
2925 ArrowError::InvalidArgumentError(msg) => {
2926 assert_eq!(
2927 msg,
2928 "Map field should be a entries struct data type, got Int32 instead"
2929 )
2930 }
2931 _ => panic!("unexpected error type {array_data_err}"),
2932 }
2933 }
2934 }
2935
2936 #[test]
2937 fn should_fail_validation_when_having_map_entries_only_have_1_field() {
2938 let struct_data_type = DataType::Struct(Fields::from(vec![Field::new(
2939 Field::MAP_KEY_FIELD_DEFAULT_NAME,
2940 DataType::Int32,
2941 false,
2942 )]));
2943
2944 let key_array_data = valid_non_nullable_int32_array_data(2);
2945
2946 let struct_data = {
2947 let builder = ArrayDataBuilder::new(struct_data_type.clone())
2948 .len(2)
2949 .nulls(None)
2950 .child_data(vec![key_array_data]);
2951
2952 builder.build().unwrap()
2953 };
2954
2955 let results = test_both_builder_and_array_data(
2956 DataType::Map(
2957 Field::new(
2958 Field::MAP_ENTRIES_FIELD_DEFAULT_NAME,
2959 struct_data_type,
2960 false,
2961 )
2962 .into(),
2963 false,
2964 ),
2965 1,
2966 None,
2967 0,
2968 vec![
2969 OffsetBuffer::<i32>::from_lengths(vec![2])
2970 .into_inner()
2971 .into(),
2972 ],
2973 vec![struct_data],
2974 );
2975
2976 for result in results {
2977 let array_data_err = result.expect_err("should fail for nullable key");
2978
2979 match array_data_err {
2980 ArrowError::InvalidArgumentError(msg) => {
2981 assert_eq!(
2982 msg,
2983 "Map entries data type should be a struct containing 2 fields, got 1 fields"
2984 )
2985 }
2986 _ => panic!("unexpected error type {array_data_err}"),
2987 }
2988 }
2989 }
2990
2991 #[test]
2992 fn should_fail_validation_when_having_map_entries_have_3_fields() {
2993 let struct_data_type = DataType::Struct(Fields::from(vec![
2994 Field::new(Field::MAP_KEY_FIELD_DEFAULT_NAME, DataType::Int32, false),
2995 Field::new(Field::MAP_VALUE_FIELD_DEFAULT_NAME, DataType::Utf8, true),
2996 Field::new("other", DataType::Int32, true),
2997 ]));
2998
2999 let key_array_data = valid_non_nullable_int32_array_data(2);
3000
3001 let values_array_data = valid_string_array_data(2);
3002
3003 let other_array_data = key_array_data.clone();
3004
3005 let struct_data = {
3006 let builder = ArrayDataBuilder::new(struct_data_type.clone())
3007 .len(2)
3008 .nulls(None)
3009 .child_data(vec![key_array_data, values_array_data, other_array_data]);
3010
3011 builder.build().unwrap()
3012 };
3013
3014 let results = test_both_builder_and_array_data(
3015 DataType::Map(
3016 Field::new(
3017 Field::MAP_ENTRIES_FIELD_DEFAULT_NAME,
3018 struct_data_type,
3019 false,
3020 )
3021 .into(),
3022 false,
3023 ),
3024 1,
3025 None,
3026 0,
3027 vec![
3028 OffsetBuffer::<i32>::from_lengths(vec![2])
3029 .into_inner()
3030 .into(),
3031 ],
3032 vec![struct_data],
3033 );
3034
3035 for result in results {
3036 let array_data_err = result.expect_err("should fail for nullable key");
3037
3038 match array_data_err {
3039 ArrowError::InvalidArgumentError(msg) => {
3040 assert_eq!(
3041 msg,
3042 "Map entries data type should be a struct containing 2 fields, got 3 fields"
3043 )
3044 }
3045 _ => panic!("unexpected error type {array_data_err}"),
3046 }
3047 }
3048 }
3049
3050 #[test]
3051 fn should_fail_validation_when_having_nullable_map_keys() {
3052 let struct_data_type = DataType::Struct(Fields::from(vec![
3053 Field::new(Field::MAP_KEY_FIELD_DEFAULT_NAME, DataType::Int32, true),
3054 Field::new(Field::MAP_VALUE_FIELD_DEFAULT_NAME, DataType::Utf8, true),
3055 ]));
3056
3057 let key_array_data = valid_non_nullable_int32_array_data(2);
3058 let values_array_data = valid_string_array_data(2);
3059
3060 let struct_data = {
3061 let builder = ArrayDataBuilder::new(struct_data_type.clone())
3062 .len(2)
3063 .nulls(None)
3064 .child_data(vec![key_array_data, values_array_data]);
3065
3066 builder.build().unwrap()
3067 };
3068
3069 let results = test_both_builder_and_array_data(
3070 DataType::Map(
3071 Field::new(
3072 Field::MAP_ENTRIES_FIELD_DEFAULT_NAME,
3073 struct_data_type,
3074 false,
3075 )
3076 .into(),
3077 false,
3078 ),
3079 1,
3080 None,
3081 0,
3082 vec![
3083 OffsetBuffer::<i32>::from_lengths(vec![2])
3084 .into_inner()
3085 .into(),
3086 ],
3087 vec![struct_data],
3088 );
3089
3090 for result in results {
3091 let array_data_err = result.expect_err("should fail for nullable key");
3092
3093 match array_data_err {
3094 ArrowError::InvalidArgumentError(msg) => {
3095 assert_eq!(msg, "Map key field must not be nullable")
3096 }
3097 _ => panic!("unexpected error type {array_data_err}"),
3098 }
3099 }
3100 }
3101
3102 #[test]
3103 fn should_fail_validation_when_having_entries_is_nullable_for_map() {
3104 let struct_data_type = DataType::Struct(Fields::from(vec![
3105 Field::new(Field::MAP_KEY_FIELD_DEFAULT_NAME, DataType::Int32, false),
3106 Field::new(Field::MAP_VALUE_FIELD_DEFAULT_NAME, DataType::Utf8, true),
3107 ]));
3108
3109 let key_array_data = valid_non_nullable_int32_array_data(2);
3110
3111 let values_array_data = valid_string_array_data(2);
3112
3113 let struct_data = {
3114 let builder = ArrayDataBuilder::new(struct_data_type.clone())
3115 .len(2)
3116 .nulls(None)
3117 .child_data(vec![key_array_data, values_array_data]);
3118
3119 builder.build().unwrap()
3120 };
3121
3122 let results = test_both_builder_and_array_data(
3123 DataType::Map(
3124 Field::new(
3125 Field::MAP_ENTRIES_FIELD_DEFAULT_NAME,
3126 struct_data_type,
3127 true,
3128 )
3129 .into(),
3130 false,
3131 ),
3132 1,
3133 None,
3134 0,
3135 vec![
3136 OffsetBuffer::<i32>::from_lengths(vec![2])
3137 .into_inner()
3138 .into(),
3139 ],
3140 vec![struct_data],
3141 );
3142
3143 for result in results {
3144 let array_data_err = result.expect_err("should fail for nullable entries");
3145
3146 match array_data_err {
3147 ArrowError::InvalidArgumentError(msg) => assert_eq!(
3148 msg,
3149 "The nullable should be set to false for the map entries field."
3150 ),
3151 _ => panic!("unexpected error type {array_data_err}"),
3152 }
3153 }
3154 }
3155
3156 #[test]
3157 fn should_allow_to_create_map_from_data() {
3158 let struct_data_type = DataType::Struct(Fields::from(vec![
3159 Field::new(Field::MAP_KEY_FIELD_DEFAULT_NAME, DataType::Int32, false),
3160 Field::new(Field::MAP_VALUE_FIELD_DEFAULT_NAME, DataType::Utf8, true),
3161 ]));
3162
3163 let key_array_data = valid_non_nullable_int32_array_data(2);
3164 let values_array_data = valid_string_array_data(2);
3165
3166 let struct_data = {
3167 let builder = ArrayDataBuilder::new(struct_data_type.clone())
3168 .len(2)
3169 .nulls(None)
3170 .child_data(vec![key_array_data, values_array_data]);
3171
3172 builder.build().unwrap()
3173 };
3174
3175 let results = test_both_builder_and_array_data(
3176 DataType::Map(
3177 Field::new(
3178 Field::MAP_ENTRIES_FIELD_DEFAULT_NAME,
3179 struct_data_type,
3180 false,
3181 )
3182 .into(),
3183 false,
3184 ),
3185 1,
3186 None,
3187 0,
3188 vec![
3189 OffsetBuffer::<i32>::from_lengths(vec![2])
3190 .into_inner()
3191 .into(),
3192 ],
3193 vec![struct_data],
3194 );
3195
3196 for result in results {
3197 result.expect("should be able to create map ArrayData");
3198 }
3199 }
3200
3201 fn valid_string_array_data(length: usize) -> ArrayData {
3202 let offsets = OffsetBuffer::<i32>::from_lengths(vec![0; length])
3203 .into_inner()
3204 .into_inner();
3205 let empty_bytes = Buffer::default();
3206
3207 let builder = ArrayDataBuilder::new(DataType::Utf8)
3208 .len(length)
3209 .buffers(vec![offsets, empty_bytes])
3210 .nulls(None);
3211
3212 builder.build().unwrap()
3213 }
3214
3215 fn valid_non_nullable_int32_array_data(length: usize) -> ArrayData {
3216 let builder = ArrayDataBuilder::new(DataType::Int32)
3217 .len(length)
3218 .nulls(None)
3219 .buffers(vec![
3220 ScalarBuffer::<i32>::from(vec![1; length]).into_inner(),
3221 ]);
3222
3223 builder.build().unwrap()
3224 }
3225
3226 #[test]
3227 fn empty_and_null_map_array_should_pass_validation() {
3228 let dt = DataType::Map(
3229 Field::new(
3230 Field::MAP_ENTRIES_FIELD_DEFAULT_NAME,
3231 DataType::Struct(Fields::from(vec![
3232 Field::new(Field::MAP_KEY_FIELD_DEFAULT_NAME, DataType::Int32, false),
3233 Field::new(Field::MAP_VALUE_FIELD_DEFAULT_NAME, DataType::Utf8, true),
3234 ])),
3235 false,
3236 )
3237 .into(),
3238 false,
3239 );
3240
3241 ArrayData::new_empty(&dt).validate_full().unwrap();
3242 ArrayData::new_null(&dt, 1).validate_full().unwrap();
3243 }
3244
3245 fn test_both_builder_and_array_data(
3246 data_type: DataType,
3247 len: usize,
3248 null_bit_buffer: Option<Buffer>,
3249 offset: usize,
3250 buffers: Vec<Buffer>,
3251 child_data: Vec<ArrayData>,
3252 ) -> [Result<ArrayData, ArrowError>; 2] {
3253 let from_builder_res = ArrayData::builder(data_type.clone())
3254 .len(len)
3255 .add_buffers(buffers.clone())
3256 .null_bit_buffer(null_bit_buffer.clone())
3257 .offset(offset)
3258 .child_data(child_data.clone())
3259 .build();
3260
3261 let from_try_new_res =
3262 ArrayData::try_new(data_type, len, null_bit_buffer, offset, buffers, child_data);
3263
3264 [from_builder_res, from_try_new_res]
3265 }
3266}