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 [MutableBuffer::new(capacity * *size as usize), empty_buffer]
138 }
139 DataType::Dictionary(k, _) => [
140 MutableBuffer::new(capacity * k.primitive_width().unwrap()),
141 empty_buffer,
142 ],
143 DataType::FixedSizeList(_, _) | DataType::Struct(_) | DataType::RunEndEncoded(_, _) => {
144 [empty_buffer, MutableBuffer::new(0)]
145 }
146 DataType::Union(_, mode) => {
147 let type_ids = MutableBuffer::new(capacity * mem::size_of::<i8>());
148 match mode {
149 UnionMode::Sparse => [type_ids, empty_buffer],
150 UnionMode::Dense => {
151 let offsets = MutableBuffer::new(capacity * mem::size_of::<i32>());
152 [type_ids, offsets]
153 }
154 }
155 }
156 }
157}
158
159#[derive(Debug, Clone)]
205pub struct ArrayData {
206 data_type: DataType,
208
209 len: usize,
211
212 offset: usize,
217
218 buffers: Vec<Buffer>,
231
232 child_data: Vec<ArrayData>,
242
243 nulls: Option<NullBuffer>,
251}
252
253pub type ArrayDataRef = Arc<ArrayData>;
255
256fn checked_len_plus_offset(
257 data_type: &DataType,
258 len: usize,
259 offset: usize,
260) -> Result<usize, ArrowError> {
261 len.checked_add(offset).ok_or_else(|| {
262 ArrowError::InvalidArgumentError(format!(
263 "Length {len} with offset {offset} overflows usize for {data_type}"
264 ))
265 })
266}
267
268impl ArrayData {
269 pub unsafe fn new_unchecked(
286 data_type: DataType,
287 len: usize,
288 null_count: Option<usize>,
289 null_bit_buffer: Option<Buffer>,
290 offset: usize,
291 buffers: Vec<Buffer>,
292 child_data: Vec<ArrayData>,
293 ) -> Self {
294 let mut skip_validation = UnsafeFlag::new();
295 unsafe { skip_validation.set(true) };
297
298 ArrayDataBuilder {
299 data_type,
300 len,
301 null_count,
302 null_bit_buffer,
303 nulls: None,
304 offset,
305 buffers,
306 child_data,
307 align_buffers: false,
308 skip_validation,
309 }
310 .build()
311 .unwrap()
312 }
313
314 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 let new_self = unsafe {
348 Self::new_unchecked(
349 data_type,
350 len,
351 None,
352 null_bit_buffer,
353 offset,
354 buffers,
355 child_data,
356 )
357 };
358
359 new_self.validate_data()?;
364 Ok(new_self)
365 }
366
367 #[inline]
369 pub const fn builder(data_type: DataType) -> ArrayDataBuilder {
370 ArrayDataBuilder::new(data_type)
371 }
372
373 #[inline]
375 pub const fn data_type(&self) -> &DataType {
376 &self.data_type
377 }
378
379 pub fn buffers(&self) -> &[Buffer] {
381 &self.buffers
382 }
383
384 pub fn child_data(&self) -> &[ArrayData] {
387 &self.child_data[..]
388 }
389
390 #[inline]
392 pub fn is_null(&self, i: usize) -> bool {
393 match &self.nulls {
394 Some(v) => v.is_null(i),
395 None => false,
396 }
397 }
398
399 #[inline]
403 pub fn nulls(&self) -> Option<&NullBuffer> {
404 self.nulls.as_ref()
405 }
406
407 #[inline]
409 pub fn is_valid(&self, i: usize) -> bool {
410 !self.is_null(i)
411 }
412
413 #[inline]
415 pub const fn len(&self) -> usize {
416 self.len
417 }
418
419 #[inline]
421 pub const fn is_empty(&self) -> bool {
422 self.len == 0
423 }
424
425 #[inline]
427 pub const fn offset(&self) -> usize {
428 self.offset
429 }
430
431 #[inline]
433 pub fn null_count(&self) -> usize {
434 self.nulls
435 .as_ref()
436 .map(|x| x.null_count())
437 .unwrap_or_default()
438 }
439
440 pub fn get_buffer_memory_size(&self) -> usize {
452 let mut size = 0;
453 for buffer in &self.buffers {
454 size += buffer.capacity();
455 }
456 if let Some(bitmap) = &self.nulls {
457 size += bitmap.buffer().capacity()
458 }
459 for child in &self.child_data {
460 size += child.get_buffer_memory_size();
461 }
462 size
463 }
464
465 pub fn get_slice_memory_size(&self) -> Result<usize, ArrowError> {
478 let mut result: usize = 0;
479 let layout = layout(&self.data_type);
480
481 for spec in layout.buffers.iter() {
482 match spec {
483 BufferSpec::FixedWidth { byte_width, .. } => {
484 let buffer_size = self.len.checked_mul(*byte_width).ok_or_else(|| {
485 ArrowError::ComputeError(
486 "Integer overflow computing buffer size".to_string(),
487 )
488 })?;
489 result += buffer_size;
490 }
491 BufferSpec::VariableWidth => {
492 let buffer_len = match self.data_type {
493 DataType::Utf8 | DataType::Binary => {
494 let offsets = self.typed_offsets::<i32>()?;
495 (offsets[self.len] - offsets[0]) as usize
496 }
497 DataType::LargeUtf8 | DataType::LargeBinary => {
498 let offsets = self.typed_offsets::<i64>()?;
499 (offsets[self.len] - offsets[0]) as usize
500 }
501 _ => {
502 return Err(ArrowError::NotYetImplemented(format!(
503 "Invalid data type for VariableWidth buffer. Expected Utf8, LargeUtf8, Binary or LargeBinary. Got {}",
504 self.data_type
505 )));
506 }
507 };
508 result += buffer_len;
509 }
510 BufferSpec::BitMap => {
511 let buffer_size = bit_util::ceil(self.len, 8);
512 result += buffer_size;
513 }
514 BufferSpec::AlwaysNull => {
515 }
517 }
518 }
519
520 if self.nulls().is_some() {
521 result += bit_util::ceil(self.len, 8);
522 }
523
524 for child in &self.child_data {
525 result += child.get_slice_memory_size()?;
526 }
527 Ok(result)
528 }
529
530 pub fn get_array_memory_size(&self) -> usize {
539 let mut size = mem::size_of_val(self);
540
541 for buffer in &self.buffers {
543 size += mem::size_of::<Buffer>();
544 size += buffer.capacity();
545 }
546 if let Some(nulls) = &self.nulls {
547 size += nulls.buffer().capacity();
548 }
549 for child in &self.child_data {
550 size += child.get_array_memory_size();
551 }
552
553 size
554 }
555
556 pub fn slice(&self, offset: usize, length: usize) -> ArrayData {
564 let end = offset
565 .checked_add(length)
566 .expect("offset + length overflow");
567 assert!(end <= self.len());
568
569 if let DataType::Struct(_) = self.data_type() {
570 let new_offset = self.offset + offset;
572 ArrayData {
573 data_type: self.data_type().clone(),
574 len: length,
575 offset: new_offset,
576 buffers: self.buffers.clone(),
577 child_data: self
579 .child_data()
580 .iter()
581 .map(|data| data.slice(offset, length))
582 .collect(),
583 nulls: self.nulls.as_ref().map(|x| x.slice(offset, length)),
584 }
585 } else {
586 let mut new_data = self.clone();
587
588 new_data.len = length;
589 new_data.offset = offset + self.offset;
590 new_data.nulls = self.nulls.as_ref().map(|x| x.slice(offset, length));
591
592 new_data
593 }
594 }
595
596 pub fn buffer<T: ArrowNativeType>(&self, buffer: usize) -> &[T] {
603 &self.buffers()[buffer].typed_data()[self.offset..]
604 }
605
606 pub fn new_null(data_type: &DataType, len: usize) -> Self {
608 let bit_len = bit_util::ceil(len, 8);
609 let zeroed = |len: usize| Buffer::from(MutableBuffer::from_len_zeroed(len));
610
611 let (buffers, child_data, has_nulls) = match data_type.primitive_width() {
612 Some(width) => (vec![zeroed(width * len)], vec![], true),
613 None => match data_type {
614 DataType::Null => (vec![], vec![], false),
615 DataType::Boolean => (vec![zeroed(bit_len)], vec![], true),
616 DataType::Binary | DataType::Utf8 => {
617 (vec![zeroed((len + 1) * 4), zeroed(0)], vec![], true)
618 }
619 DataType::BinaryView | DataType::Utf8View => (vec![zeroed(len * 16)], vec![], true),
620 DataType::LargeBinary | DataType::LargeUtf8 => {
621 (vec![zeroed((len + 1) * 8), zeroed(0)], vec![], true)
622 }
623 DataType::FixedSizeBinary(i) => (vec![zeroed(*i as usize * len)], vec![], true),
624 DataType::List(f) | DataType::Map(f, _) => (
625 vec![zeroed((len + 1) * 4)],
626 vec![ArrayData::new_empty(f.data_type())],
627 true,
628 ),
629 DataType::LargeList(f) => (
630 vec![zeroed((len + 1) * 8)],
631 vec![ArrayData::new_empty(f.data_type())],
632 true,
633 ),
634 DataType::ListView(f) => (
635 vec![zeroed(len * 4), zeroed(len * 4)],
636 vec![ArrayData::new_empty(f.data_type())],
637 true,
638 ),
639 DataType::LargeListView(f) => (
640 vec![zeroed(len * 8), zeroed(len * 8)],
641 vec![ArrayData::new_empty(f.data_type())],
642 true,
643 ),
644 DataType::FixedSizeList(f, list_len) => (
645 vec![],
646 vec![ArrayData::new_null(f.data_type(), *list_len as usize * len)],
647 true,
648 ),
649 DataType::Struct(fields) => (
650 vec![],
651 fields
652 .iter()
653 .map(|f| Self::new_null(f.data_type(), len))
654 .collect(),
655 true,
656 ),
657 DataType::Dictionary(k, v) => (
658 vec![zeroed(k.primitive_width().unwrap() * len)],
659 vec![ArrayData::new_empty(v.as_ref())],
660 true,
661 ),
662 DataType::Union(f, mode) => {
663 let (id, _) = f.iter().next().unwrap();
664 let ids = Buffer::from_iter(std::iter::repeat_n(id, len));
665 let buffers = match mode {
666 UnionMode::Sparse => vec![ids],
667 UnionMode::Dense => {
668 let end_offset = i32::from_usize(len).unwrap();
669 vec![ids, Buffer::from_iter(0_i32..end_offset)]
670 }
671 };
672
673 let children = f
674 .iter()
675 .enumerate()
676 .map(|(idx, (_, f))| {
677 if idx == 0 || *mode == UnionMode::Sparse {
678 Self::new_null(f.data_type(), len)
679 } else {
680 Self::new_empty(f.data_type())
681 }
682 })
683 .collect();
684
685 (buffers, children, false)
686 }
687 DataType::RunEndEncoded(r, v) => {
688 let runs = match r.data_type() {
689 DataType::Int16 => {
690 let i = i16::from_usize(len).expect("run overflow");
691 Buffer::from_slice_ref([i])
692 }
693 DataType::Int32 => {
694 let i = i32::from_usize(len).expect("run overflow");
695 Buffer::from_slice_ref([i])
696 }
697 DataType::Int64 => {
698 let i = i64::from_usize(len).expect("run overflow");
699 Buffer::from_slice_ref([i])
700 }
701 dt => unreachable!("Invalid run ends data type {dt}"),
702 };
703
704 let builder = ArrayData::builder(r.data_type().clone())
705 .len(1)
706 .buffers(vec![runs]);
707
708 let runs = unsafe { builder.build_unchecked() };
711 (
712 vec![],
713 vec![runs, ArrayData::new_null(v.data_type(), 1)],
714 false,
715 )
716 }
717 d => unreachable!("{d}"),
718 },
719 };
720
721 let mut builder = ArrayDataBuilder::new(data_type.clone())
722 .len(len)
723 .buffers(buffers)
724 .child_data(child_data);
725
726 if has_nulls {
727 builder = builder.nulls(Some(NullBuffer::new_null(len)))
728 }
729
730 unsafe { builder.build_unchecked() }
733 }
734
735 pub fn new_empty(data_type: &DataType) -> Self {
737 Self::new_null(data_type, 0)
738 }
739
740 pub fn align_buffers(&mut self) {
749 let layout = layout(&self.data_type);
750 for (buffer, spec) in self.buffers.iter_mut().zip(&layout.buffers) {
751 if let BufferSpec::FixedWidth { alignment, .. } = spec {
752 if buffer.as_ptr().align_offset(*alignment) != 0 {
753 *buffer = Buffer::from_slice_ref(buffer.as_ref());
754 }
755 }
756 }
757 for data in self.child_data.iter_mut() {
759 data.align_buffers()
760 }
761 }
762
763 pub fn validate(&self) -> Result<(), ArrowError> {
774 let len_plus_offset = checked_len_plus_offset(&self.data_type, self.len, self.offset)?;
776
777 let layout = layout(&self.data_type);
779
780 if !layout.can_contain_null_mask && self.nulls.is_some() {
781 return Err(ArrowError::InvalidArgumentError(format!(
782 "Arrays of type {:?} cannot contain a null bitmask",
783 self.data_type,
784 )));
785 }
786
787 if self.buffers.len() < layout.buffers.len()
789 || (!layout.variadic && self.buffers.len() != layout.buffers.len())
790 {
791 return Err(ArrowError::InvalidArgumentError(format!(
792 "Expected {} buffers in array of type {:?}, got {}",
793 layout.buffers.len(),
794 self.data_type,
795 self.buffers.len(),
796 )));
797 }
798
799 for (i, (buffer, spec)) in self.buffers.iter().zip(layout.buffers.iter()).enumerate() {
800 match spec {
801 BufferSpec::FixedWidth {
802 byte_width,
803 alignment,
804 } => {
805 let min_buffer_size = len_plus_offset.saturating_mul(*byte_width);
806
807 if buffer.len() < min_buffer_size {
808 return Err(ArrowError::InvalidArgumentError(format!(
809 "Need at least {} bytes in buffers[{}] in array of type {:?}, but got {}",
810 min_buffer_size,
811 i,
812 self.data_type,
813 buffer.len()
814 )));
815 }
816
817 let align_offset = buffer.as_ptr().align_offset(*alignment);
818 if align_offset != 0 {
819 return Err(ArrowError::InvalidArgumentError(format!(
820 "Misaligned buffers[{i}] in array of type {:?}, offset from expected alignment of {alignment} by {}",
821 self.data_type,
822 align_offset.min(alignment - align_offset)
823 )));
824 }
825 }
826 BufferSpec::VariableWidth => {
827 }
831 BufferSpec::BitMap => {
832 let min_buffer_size = bit_util::ceil(len_plus_offset, 8);
833 if buffer.len() < min_buffer_size {
834 return Err(ArrowError::InvalidArgumentError(format!(
835 "Need at least {} bytes for bitmap in buffers[{}] in array of type {:?}, but got {}",
836 min_buffer_size,
837 i,
838 self.data_type,
839 buffer.len()
840 )));
841 }
842 }
843 BufferSpec::AlwaysNull => {
844 }
846 }
847 }
848
849 if let Some(nulls) = self.nulls() {
851 if nulls.null_count() > self.len {
852 return Err(ArrowError::InvalidArgumentError(format!(
853 "null_count {} for an array exceeds length of {} elements",
854 nulls.null_count(),
855 self.len
856 )));
857 }
858
859 let actual_len = nulls.validity().len();
860 let needed_len = bit_util::ceil(len_plus_offset, 8);
861 if actual_len < needed_len {
862 return Err(ArrowError::InvalidArgumentError(format!(
863 "null_bit_buffer size too small. got {actual_len} needed {needed_len}",
864 )));
865 }
866
867 if nulls.len() != self.len {
868 return Err(ArrowError::InvalidArgumentError(format!(
869 "null buffer incorrect size. got {} expected {}",
870 nulls.len(),
871 self.len
872 )));
873 }
874 }
875
876 self.validate_child_data()?;
877
878 match &self.data_type {
880 DataType::Utf8 | DataType::Binary => {
881 self.validate_offsets::<i32>(self.buffers[1].len())?;
882 }
883 DataType::LargeUtf8 | DataType::LargeBinary => {
884 self.validate_offsets::<i64>(self.buffers[1].len())?;
885 }
886 DataType::Dictionary(key_type, _value_type) => {
887 if !DataType::is_dictionary_key_type(key_type) {
889 return Err(ArrowError::InvalidArgumentError(format!(
890 "Dictionary key type must be integer, but was {key_type}"
891 )));
892 }
893 }
894 DataType::RunEndEncoded(run_ends_type, _) => {
895 if run_ends_type.is_nullable() {
896 return Err(ArrowError::InvalidArgumentError(
897 "The nullable should be set to false for the field defining run_ends array.".to_string()
898 ));
899 }
900 if !DataType::is_run_ends_type(run_ends_type.data_type()) {
901 return Err(ArrowError::InvalidArgumentError(format!(
902 "RunArray run_ends types must be Int16, Int32 or Int64, but was {}",
903 run_ends_type.data_type()
904 )));
905 }
906 }
907 _ => {}
908 };
909
910 Ok(())
911 }
912
913 fn typed_offsets<T: ArrowNativeType + num_traits::Num>(&self) -> Result<&[T], ArrowError> {
920 if self.len == 0 && self.buffers[0].is_empty() {
922 return Ok(&[]);
923 }
924
925 let len = checked_len_plus_offset(&self.data_type, self.len, 1)?;
926
927 self.typed_buffer(0, len)
928 }
929
930 fn typed_buffer<T: ArrowNativeType + num_traits::Num>(
932 &self,
933 idx: usize,
934 len: usize,
935 ) -> Result<&[T], ArrowError> {
936 let buffer = &self.buffers[idx];
937
938 let required_elements = checked_len_plus_offset(&self.data_type, len, self.offset)?;
939 let byte_width = mem::size_of::<T>();
940 let required_len = required_elements.checked_mul(byte_width).ok_or_else(|| {
941 ArrowError::InvalidArgumentError(format!(
942 "Buffer {idx} of {} byte length overflow: {} elements of {} bytes exceeds usize",
943 self.data_type, required_elements, byte_width
944 ))
945 })?;
946
947 if buffer.len() < required_len {
948 return Err(ArrowError::InvalidArgumentError(format!(
949 "Buffer {} of {} isn't large enough. Expected {} bytes got {}",
950 idx,
951 self.data_type,
952 required_len,
953 buffer.len()
954 )));
955 }
956
957 Ok(&buffer.typed_data::<T>()[self.offset..required_elements])
958 }
959
960 fn validate_offsets<T: ArrowNativeType + num_traits::Num + std::fmt::Display>(
963 &self,
964 values_length: usize,
965 ) -> Result<(), ArrowError> {
966 let offsets = self.typed_offsets::<T>()?;
968 if offsets.is_empty() {
969 return Ok(());
970 }
971
972 let first_offset = offsets[0].to_usize().ok_or_else(|| {
973 ArrowError::InvalidArgumentError(format!(
974 "Error converting offset[0] ({}) to usize for {}",
975 offsets[0], self.data_type
976 ))
977 })?;
978
979 let last_offset = offsets[self.len].to_usize().ok_or_else(|| {
980 ArrowError::InvalidArgumentError(format!(
981 "Error converting offset[{}] ({}) to usize for {}",
982 self.len, offsets[self.len], self.data_type
983 ))
984 })?;
985
986 if first_offset > values_length {
987 return Err(ArrowError::InvalidArgumentError(format!(
988 "First offset {} of {} is larger than values length {}",
989 first_offset, self.data_type, values_length,
990 )));
991 }
992
993 if last_offset > values_length {
994 return Err(ArrowError::InvalidArgumentError(format!(
995 "Last offset {} of {} is larger than values length {}",
996 last_offset, self.data_type, values_length,
997 )));
998 }
999
1000 if first_offset > last_offset {
1001 return Err(ArrowError::InvalidArgumentError(format!(
1002 "First offset {} in {} is smaller than last offset {}",
1003 first_offset, self.data_type, last_offset,
1004 )));
1005 }
1006
1007 Ok(())
1008 }
1009
1010 fn validate_offsets_and_sizes<T: ArrowNativeType + num_traits::Num + std::fmt::Display>(
1013 &self,
1014 values_length: usize,
1015 ) -> Result<(), ArrowError> {
1016 let offsets: &[T] = self.typed_buffer(0, self.len)?;
1017 let sizes: &[T] = self.typed_buffer(1, self.len)?;
1018 if offsets.len() != sizes.len() {
1019 return Err(ArrowError::ComputeError(format!(
1020 "ListView offsets len {} does not match sizes len {}",
1021 offsets.len(),
1022 sizes.len()
1023 )));
1024 }
1025
1026 for i in 0..sizes.len() {
1027 let size = sizes[i].to_usize().ok_or_else(|| {
1028 ArrowError::InvalidArgumentError(format!(
1029 "Error converting size[{}] ({}) to usize for {}",
1030 i, sizes[i], self.data_type
1031 ))
1032 })?;
1033 let offset = offsets[i].to_usize().ok_or_else(|| {
1034 ArrowError::InvalidArgumentError(format!(
1035 "Error converting offset[{}] ({}) to usize for {}",
1036 i, offsets[i], self.data_type
1037 ))
1038 })?;
1039 if size
1040 .checked_add(offset)
1041 .expect("Offset and size have exceeded the usize boundary")
1042 > values_length
1043 {
1044 return Err(ArrowError::InvalidArgumentError(format!(
1045 "Size {} at index {} is larger than the remaining values for {}",
1046 size, i, self.data_type
1047 )));
1048 }
1049 }
1050 Ok(())
1051 }
1052
1053 fn validate_child_data(&self) -> Result<(), ArrowError> {
1055 match &self.data_type {
1056 DataType::List(field) | DataType::Map(field, _) => {
1057 let values_data = self.get_single_valid_child_data(field.data_type())?;
1058 self.validate_offsets::<i32>(values_data.len)?;
1059 Ok(())
1060 }
1061 DataType::LargeList(field) => {
1062 let values_data = self.get_single_valid_child_data(field.data_type())?;
1063 self.validate_offsets::<i64>(values_data.len)?;
1064 Ok(())
1065 }
1066 DataType::ListView(field) => {
1067 let values_data = self.get_single_valid_child_data(field.data_type())?;
1068 self.validate_offsets_and_sizes::<i32>(values_data.len)?;
1069 Ok(())
1070 }
1071 DataType::LargeListView(field) => {
1072 let values_data = self.get_single_valid_child_data(field.data_type())?;
1073 self.validate_offsets_and_sizes::<i64>(values_data.len)?;
1074 Ok(())
1075 }
1076 DataType::FixedSizeList(field, list_size) => {
1077 let values_data = self.get_single_valid_child_data(field.data_type())?;
1078
1079 let list_size: usize = (*list_size).try_into().map_err(|_| {
1080 ArrowError::InvalidArgumentError(format!(
1081 "{} has a negative list_size {}",
1082 self.data_type, list_size
1083 ))
1084 })?;
1085
1086 let expected_values_len = self.len
1087 .checked_mul(list_size)
1088 .expect("integer overflow computing expected number of expected values in FixedListSize");
1089
1090 if values_data.len < expected_values_len {
1091 return Err(ArrowError::InvalidArgumentError(format!(
1092 "Values length {} is less than the length ({}) multiplied by the value size ({}) for {}",
1093 values_data.len, self.len, list_size, self.data_type
1094 )));
1095 }
1096
1097 Ok(())
1098 }
1099 DataType::Struct(fields) => {
1100 self.validate_num_child_data(fields.len())?;
1101 for (i, field) in fields.iter().enumerate() {
1102 let field_data = self.get_valid_child_data(i, field.data_type())?;
1103
1104 if field_data.len < self.len {
1106 return Err(ArrowError::InvalidArgumentError(format!(
1107 "{} child array #{} for field {} has length smaller than expected for struct array ({} < {})",
1108 self.data_type,
1109 i,
1110 field.name(),
1111 field_data.len,
1112 self.len
1113 )));
1114 }
1115 }
1116 Ok(())
1117 }
1118 DataType::RunEndEncoded(run_ends_field, values_field) => {
1119 self.validate_num_child_data(2)?;
1120 let run_ends_data = self.get_valid_child_data(0, run_ends_field.data_type())?;
1121 let values_data = self.get_valid_child_data(1, values_field.data_type())?;
1122 if run_ends_data.len != values_data.len {
1123 return Err(ArrowError::InvalidArgumentError(format!(
1124 "The run_ends array length should be the same as values array length. Run_ends array length is {}, values array length is {}",
1125 run_ends_data.len, values_data.len
1126 )));
1127 }
1128 if run_ends_data.nulls.is_some() {
1129 return Err(ArrowError::InvalidArgumentError(
1130 "Found null values in run_ends array. The run_ends array should not have null values.".to_string(),
1131 ));
1132 }
1133 Ok(())
1134 }
1135 DataType::Union(fields, mode) => {
1136 self.validate_num_child_data(fields.len())?;
1137
1138 for (i, (_, field)) in fields.iter().enumerate() {
1139 let field_data = self.get_valid_child_data(i, field.data_type())?;
1140
1141 if mode == &UnionMode::Sparse {
1142 let len_plus_offset =
1143 checked_len_plus_offset(&self.data_type, self.len, self.offset)?;
1144 if field_data.len < len_plus_offset {
1145 return Err(ArrowError::InvalidArgumentError(format!(
1146 "Sparse union child array #{} has length smaller than expected for union array ({} < {})",
1147 i, field_data.len, len_plus_offset
1148 )));
1149 }
1150 }
1151 }
1152 Ok(())
1153 }
1154 DataType::Dictionary(_key_type, value_type) => {
1155 self.get_single_valid_child_data(value_type)?;
1156 Ok(())
1157 }
1158 _ => {
1159 if !self.child_data.is_empty() {
1161 return Err(ArrowError::InvalidArgumentError(format!(
1162 "Expected no child arrays for type {} but got {}",
1163 self.data_type,
1164 self.child_data.len()
1165 )));
1166 }
1167 Ok(())
1168 }
1169 }
1170 }
1171
1172 fn get_single_valid_child_data(
1176 &self,
1177 expected_type: &DataType,
1178 ) -> Result<&ArrayData, ArrowError> {
1179 self.validate_num_child_data(1)?;
1180 self.get_valid_child_data(0, expected_type)
1181 }
1182
1183 fn validate_num_child_data(&self, expected_len: usize) -> Result<(), ArrowError> {
1185 if self.child_data.len() != expected_len {
1186 Err(ArrowError::InvalidArgumentError(format!(
1187 "Value data for {} should contain {} child data array(s), had {}",
1188 self.data_type,
1189 expected_len,
1190 self.child_data.len()
1191 )))
1192 } else {
1193 Ok(())
1194 }
1195 }
1196
1197 fn get_valid_child_data(
1200 &self,
1201 i: usize,
1202 expected_type: &DataType,
1203 ) -> Result<&ArrayData, ArrowError> {
1204 let values_data = self.child_data.get(i).ok_or_else(|| {
1205 ArrowError::InvalidArgumentError(format!(
1206 "{} did not have enough child arrays. Expected at least {} but had only {}",
1207 self.data_type,
1208 i + 1,
1209 self.child_data.len()
1210 ))
1211 })?;
1212
1213 if expected_type != &values_data.data_type {
1214 return Err(ArrowError::InvalidArgumentError(format!(
1215 "Child type mismatch for {}. Expected {} but child data had {}",
1216 self.data_type, expected_type, values_data.data_type
1217 )));
1218 }
1219
1220 values_data.validate()?;
1221 Ok(values_data)
1222 }
1223
1224 pub fn validate_data(&self) -> Result<(), ArrowError> {
1240 self.validate()?;
1241
1242 self.validate_nulls()?;
1243 self.validate_values()?;
1244 Ok(())
1245 }
1246
1247 pub fn validate_full(&self) -> Result<(), ArrowError> {
1252 self.validate_data()?;
1253 self.child_data
1255 .iter()
1256 .enumerate()
1257 .try_for_each(|(i, child_data)| {
1258 child_data.validate_full().map_err(|e| {
1259 ArrowError::InvalidArgumentError(format!(
1260 "{} child #{} invalid: {}",
1261 self.data_type, i, e
1262 ))
1263 })
1264 })?;
1265 Ok(())
1266 }
1267
1268 pub fn validate_nulls(&self) -> Result<(), ArrowError> {
1278 if let Some(nulls) = &self.nulls {
1279 let actual = nulls.len() - nulls.inner().count_set_bits();
1280 if actual != nulls.null_count() {
1281 return Err(ArrowError::InvalidArgumentError(format!(
1282 "null_count value ({}) doesn't match actual number of nulls in array ({})",
1283 nulls.null_count(),
1284 actual
1285 )));
1286 }
1287 }
1288
1289 match &self.data_type {
1294 DataType::List(f) | DataType::LargeList(f) | DataType::Map(f, _) => {
1295 if !f.is_nullable() {
1296 self.validate_non_nullable(None, &self.child_data[0])?
1297 }
1298 }
1299 DataType::FixedSizeList(field, len) => {
1300 let child = &self.child_data[0];
1301 if !field.is_nullable() {
1302 match &self.nulls {
1303 Some(nulls) => {
1304 let element_len = *len as usize;
1305 let expanded = nulls.expand(element_len);
1306 self.validate_non_nullable(Some(&expanded), child)?;
1307 }
1308 None => self.validate_non_nullable(None, child)?,
1309 }
1310 }
1311 }
1312 DataType::Struct(fields) => {
1313 for (field, child) in fields.iter().zip(&self.child_data) {
1314 if !field.is_nullable() {
1315 self.validate_non_nullable(self.nulls(), child)?
1316 }
1317 }
1318 }
1319 _ => {}
1320 }
1321
1322 Ok(())
1323 }
1324
1325 fn validate_non_nullable(
1327 &self,
1328 mask: Option<&NullBuffer>,
1329 child: &ArrayData,
1330 ) -> Result<(), ArrowError> {
1331 let mask = match mask {
1332 Some(mask) => mask,
1333 None => {
1334 return match child.null_count() {
1335 0 => Ok(()),
1336 _ => Err(ArrowError::InvalidArgumentError(format!(
1337 "non-nullable child of type {} contains nulls not present in parent {}",
1338 child.data_type, self.data_type
1339 ))),
1340 };
1341 }
1342 };
1343
1344 match child.nulls() {
1345 Some(nulls) if !mask.contains(nulls) => Err(ArrowError::InvalidArgumentError(format!(
1346 "non-nullable child of type {} contains nulls not present in parent",
1347 child.data_type
1348 ))),
1349 _ => Ok(()),
1350 }
1351 }
1352
1353 pub fn validate_values(&self) -> Result<(), ArrowError> {
1359 match &self.data_type {
1360 DataType::Utf8 => self.validate_utf8::<i32>(),
1361 DataType::LargeUtf8 => self.validate_utf8::<i64>(),
1362 DataType::Binary => self.validate_offsets_full::<i32>(self.buffers[1].len()),
1363 DataType::LargeBinary => self.validate_offsets_full::<i64>(self.buffers[1].len()),
1364 DataType::BinaryView => {
1365 let views = self.typed_buffer::<u128>(0, self.len)?;
1366 validate_binary_view(views, &self.buffers[1..])
1367 }
1368 DataType::Utf8View => {
1369 let views = self.typed_buffer::<u128>(0, self.len)?;
1370 validate_string_view(views, &self.buffers[1..])
1371 }
1372 DataType::List(_) | DataType::Map(_, _) => {
1373 let child = &self.child_data[0];
1374 self.validate_offsets_full::<i32>(child.len)
1375 }
1376 DataType::LargeList(_) => {
1377 let child = &self.child_data[0];
1378 self.validate_offsets_full::<i64>(child.len)
1379 }
1380 DataType::Union(fields, mode) => {
1381 let mut child_lens: Vec<Option<usize>> = vec![None; 128];
1382 for (index, (type_id, _)) in fields.iter().enumerate() {
1383 let child = &self.child_data[index];
1384 child_lens[type_id as usize] = Some(child.len());
1385 }
1386
1387 let type_ids = self.typed_buffer::<i8>(0, self.len)?;
1389 for (index, &type_id) in type_ids.iter().enumerate() {
1390 if type_id < 0 {
1391 return Err(ArrowError::InvalidArgumentError(format!(
1392 "Type id values cannot be negative, got {type_id} at index {index}"
1393 )));
1394 }
1395 let type_id = type_id as usize;
1396 if child_lens[type_id].is_none() {
1397 return Err(ArrowError::InvalidArgumentError(format!(
1398 "Type id values must match one of the field type ids, got {type_id} at index {index}"
1399 )));
1400 }
1401 }
1402
1403 if mode == &UnionMode::Dense {
1405 let offsets = self.typed_buffer::<i32>(1, self.len)?;
1406
1407 if type_ids
1408 .iter()
1409 .zip(offsets.iter())
1410 .any(|(&type_id, &offset)| {
1411 offset < 0 || offset as usize >= child_lens[type_id as usize].unwrap()
1412 })
1413 {
1414 return Err(ArrowError::InvalidArgumentError(
1415 "Offsets must be non-negative and within the length of the child array"
1416 .to_string(),
1417 ));
1418 }
1419 }
1420
1421 Ok(())
1422 }
1423 DataType::Dictionary(key_type, _value_type) => {
1424 let dictionary_length: i64 = self.child_data[0].len.try_into().unwrap();
1425 let max_value = dictionary_length - 1;
1426 match key_type.as_ref() {
1427 DataType::UInt8 => self.check_bounds::<u8>(max_value),
1428 DataType::UInt16 => self.check_bounds::<u16>(max_value),
1429 DataType::UInt32 => self.check_bounds::<u32>(max_value),
1430 DataType::UInt64 => self.check_bounds::<u64>(max_value),
1431 DataType::Int8 => self.check_bounds::<i8>(max_value),
1432 DataType::Int16 => self.check_bounds::<i16>(max_value),
1433 DataType::Int32 => self.check_bounds::<i32>(max_value),
1434 DataType::Int64 => self.check_bounds::<i64>(max_value),
1435 _ => unreachable!(),
1436 }
1437 }
1438 DataType::RunEndEncoded(run_ends, _values) => {
1439 let run_ends_data = self.child_data()[0].clone();
1440 let len_plus_offset =
1441 checked_len_plus_offset(&self.data_type, self.len, self.offset)?;
1442 match run_ends.data_type() {
1443 DataType::Int16 => {
1444 run_ends_data.check_run_ends::<i16>(self.len, len_plus_offset)
1445 }
1446 DataType::Int32 => {
1447 run_ends_data.check_run_ends::<i32>(self.len, len_plus_offset)
1448 }
1449 DataType::Int64 => {
1450 run_ends_data.check_run_ends::<i64>(self.len, len_plus_offset)
1451 }
1452 _ => unreachable!(),
1453 }
1454 }
1455 _ => {
1456 Ok(())
1458 }
1459 }
1460 }
1461
1462 fn validate_each_offset<T, V>(&self, offset_limit: usize, validate: V) -> Result<(), ArrowError>
1473 where
1474 T: ArrowNativeType + TryInto<usize> + num_traits::Num + std::fmt::Display,
1475 V: Fn(usize, Range<usize>) -> Result<(), ArrowError>,
1476 {
1477 self.typed_offsets::<T>()?
1478 .iter()
1479 .enumerate()
1480 .map(|(i, x)| {
1481 let r = x.to_usize().ok_or_else(|| {
1483 ArrowError::InvalidArgumentError(format!(
1484 "Offset invariant failure: Could not convert offset {x} to usize at position {i}"))}
1485 );
1486 match r {
1488 Ok(n) if n <= offset_limit => Ok((i, n)),
1489 Ok(_) => Err(ArrowError::InvalidArgumentError(format!(
1490 "Offset invariant failure: offset at position {i} out of bounds: {x} > {offset_limit}"))
1491 ),
1492 Err(e) => Err(e),
1493 }
1494 })
1495 .scan(0_usize, |start, end| {
1496 match end {
1498 Ok((i, end)) if *start <= end => {
1499 let range = Some(Ok((i, *start..end)));
1500 *start = end;
1501 range
1502 }
1503 Ok((i, end)) => Some(Err(ArrowError::InvalidArgumentError(format!(
1504 "Offset invariant failure: non-monotonic offset at slot {}: {} > {}",
1505 i - 1, start, end))
1506 )),
1507 Err(err) => Some(Err(err)),
1508 }
1509 })
1510 .skip(1) .try_for_each(|res: Result<(usize, Range<usize>), ArrowError>| {
1512 let (item_index, range) = res?;
1513 validate(item_index-1, range)
1514 })
1515 }
1516
1517 fn validate_utf8<T>(&self) -> Result<(), ArrowError>
1520 where
1521 T: ArrowNativeType + TryInto<usize> + num_traits::Num + std::fmt::Display,
1522 {
1523 let values_buffer = &self.buffers[1].as_slice();
1524 if let Ok(values_str) = std::str::from_utf8(values_buffer) {
1525 self.validate_each_offset::<T, _>(values_buffer.len(), |string_index, range| {
1527 if !values_str.is_char_boundary(range.start)
1528 || !values_str.is_char_boundary(range.end)
1529 {
1530 return Err(ArrowError::InvalidArgumentError(format!(
1531 "incomplete utf-8 byte sequence from index {string_index}"
1532 )));
1533 }
1534 Ok(())
1535 })
1536 } else {
1537 self.validate_each_offset::<T, _>(values_buffer.len(), |string_index, range| {
1539 std::str::from_utf8(&values_buffer[range.clone()]).map_err(|e| {
1540 ArrowError::InvalidArgumentError(format!(
1541 "Invalid UTF8 sequence at string index {string_index} ({range:?}): {e}"
1542 ))
1543 })?;
1544 Ok(())
1545 })
1546 }
1547 }
1548
1549 fn validate_offsets_full<T>(&self, offset_limit: usize) -> Result<(), ArrowError>
1552 where
1553 T: ArrowNativeType + TryInto<usize> + num_traits::Num + std::fmt::Display,
1554 {
1555 self.validate_each_offset::<T, _>(offset_limit, |_string_index, _range| {
1556 Ok(())
1559 })
1560 }
1561
1562 fn check_bounds<T>(&self, max_value: i64) -> Result<(), ArrowError>
1565 where
1566 T: ArrowNativeType + TryInto<i64> + num_traits::Num + std::fmt::Display,
1567 {
1568 let required_len = checked_len_plus_offset(&self.data_type, self.len, self.offset)?;
1569 let buffer = &self.buffers[0];
1570
1571 assert!(buffer.len() / mem::size_of::<T>() >= required_len);
1574
1575 let indexes: &[T] = &buffer.typed_data::<T>()[self.offset..required_len];
1577
1578 indexes.iter().enumerate().try_for_each(|(i, &dict_index)| {
1579 if self.is_null(i) {
1581 return Ok(());
1582 }
1583 let dict_index: i64 = dict_index.try_into().map_err(|_| {
1584 ArrowError::InvalidArgumentError(format!(
1585 "Value at position {i} out of bounds: {dict_index} (can not convert to i64)"
1586 ))
1587 })?;
1588
1589 if dict_index < 0 || dict_index > max_value {
1590 return Err(ArrowError::InvalidArgumentError(format!(
1591 "Value at position {i} out of bounds: {dict_index} (should be in [0, {max_value}])"
1592 )));
1593 }
1594 Ok(())
1595 })
1596 }
1597
1598 fn check_run_ends<T>(
1600 &self,
1601 parent_len: usize,
1602 parent_len_plus_offset: usize,
1603 ) -> Result<(), ArrowError>
1604 where
1605 T: ArrowNativeType + TryInto<i64> + num_traits::Num + std::fmt::Display,
1606 {
1607 let values = self.typed_buffer::<T>(0, self.len)?;
1608 let mut prev_value: i64 = 0_i64;
1609 values.iter().enumerate().try_for_each(|(ix, &inp_value)| {
1610 let value: i64 = inp_value.try_into().map_err(|_| {
1611 ArrowError::InvalidArgumentError(format!(
1612 "Value at position {ix} out of bounds: {inp_value} (can not convert to i64)"
1613 ))
1614 })?;
1615 if value <= 0_i64 {
1616 return Err(ArrowError::InvalidArgumentError(format!(
1617 "The values in run_ends array should be strictly positive. Found value {value} at index {ix} that does not match the criteria."
1618 )));
1619 }
1620 if ix > 0 && value <= prev_value {
1621 return Err(ArrowError::InvalidArgumentError(format!(
1622 "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."
1623 )));
1624 }
1625
1626 prev_value = value;
1627 Ok(())
1628 })?;
1629
1630 if parent_len > 0 && prev_value.as_usize() < parent_len_plus_offset {
1631 return Err(ArrowError::InvalidArgumentError(format!(
1632 "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 {parent_len_plus_offset}."
1633 )));
1634 }
1635 Ok(())
1636 }
1637
1638 pub fn ptr_eq(&self, other: &Self) -> bool {
1642 if self.offset != other.offset
1643 || self.len != other.len
1644 || self.data_type != other.data_type
1645 || self.buffers.len() != other.buffers.len()
1646 || self.child_data.len() != other.child_data.len()
1647 {
1648 return false;
1649 }
1650
1651 match (&self.nulls, &other.nulls) {
1652 (Some(a), Some(b)) if !a.inner().ptr_eq(b.inner()) => return false,
1653 (Some(_), None) | (None, Some(_)) => return false,
1654 _ => {}
1655 };
1656
1657 if !self
1658 .buffers
1659 .iter()
1660 .zip(other.buffers.iter())
1661 .all(|(a, b)| a.as_ptr() == b.as_ptr())
1662 {
1663 return false;
1664 }
1665
1666 self.child_data
1667 .iter()
1668 .zip(other.child_data.iter())
1669 .all(|(a, b)| a.ptr_eq(b))
1670 }
1671
1672 pub fn into_builder(self) -> ArrayDataBuilder {
1674 self.into()
1675 }
1676}
1677
1678pub fn layout(data_type: &DataType) -> DataTypeLayout {
1681 use arrow_schema::IntervalUnit::*;
1684
1685 match data_type {
1686 DataType::Null => DataTypeLayout {
1687 buffers: vec![],
1688 can_contain_null_mask: false,
1689 variadic: false,
1690 },
1691 DataType::Boolean => DataTypeLayout {
1692 buffers: vec![BufferSpec::BitMap],
1693 can_contain_null_mask: true,
1694 variadic: false,
1695 },
1696 DataType::Int8 => DataTypeLayout::new_fixed_width::<i8>(),
1697 DataType::Int16 => DataTypeLayout::new_fixed_width::<i16>(),
1698 DataType::Int32 => DataTypeLayout::new_fixed_width::<i32>(),
1699 DataType::Int64 => DataTypeLayout::new_fixed_width::<i64>(),
1700 DataType::UInt8 => DataTypeLayout::new_fixed_width::<u8>(),
1701 DataType::UInt16 => DataTypeLayout::new_fixed_width::<u16>(),
1702 DataType::UInt32 => DataTypeLayout::new_fixed_width::<u32>(),
1703 DataType::UInt64 => DataTypeLayout::new_fixed_width::<u64>(),
1704 DataType::Float16 => DataTypeLayout::new_fixed_width::<half::f16>(),
1705 DataType::Float32 => DataTypeLayout::new_fixed_width::<f32>(),
1706 DataType::Float64 => DataTypeLayout::new_fixed_width::<f64>(),
1707 DataType::Timestamp(_, _) => DataTypeLayout::new_fixed_width::<i64>(),
1708 DataType::Date32 => DataTypeLayout::new_fixed_width::<i32>(),
1709 DataType::Date64 => DataTypeLayout::new_fixed_width::<i64>(),
1710 DataType::Time32(_) => DataTypeLayout::new_fixed_width::<i32>(),
1711 DataType::Time64(_) => DataTypeLayout::new_fixed_width::<i64>(),
1712 DataType::Interval(YearMonth) => DataTypeLayout::new_fixed_width::<i32>(),
1713 DataType::Interval(DayTime) => DataTypeLayout::new_fixed_width::<IntervalDayTime>(),
1714 DataType::Interval(MonthDayNano) => {
1715 DataTypeLayout::new_fixed_width::<IntervalMonthDayNano>()
1716 }
1717 DataType::Duration(_) => DataTypeLayout::new_fixed_width::<i64>(),
1718 DataType::Decimal32(_, _) => DataTypeLayout::new_fixed_width::<i32>(),
1719 DataType::Decimal64(_, _) => DataTypeLayout::new_fixed_width::<i64>(),
1720 DataType::Decimal128(_, _) => DataTypeLayout::new_fixed_width::<i128>(),
1721 DataType::Decimal256(_, _) => DataTypeLayout::new_fixed_width::<i256>(),
1722 DataType::FixedSizeBinary(size) => {
1723 let spec = BufferSpec::FixedWidth {
1724 byte_width: (*size).try_into().unwrap(),
1725 alignment: mem::align_of::<u8>(),
1726 };
1727 DataTypeLayout {
1728 buffers: vec![spec],
1729 can_contain_null_mask: true,
1730 variadic: false,
1731 }
1732 }
1733 DataType::Binary => DataTypeLayout::new_binary::<i32>(),
1734 DataType::LargeBinary => DataTypeLayout::new_binary::<i64>(),
1735 DataType::Utf8 => DataTypeLayout::new_binary::<i32>(),
1736 DataType::LargeUtf8 => DataTypeLayout::new_binary::<i64>(),
1737 DataType::BinaryView | DataType::Utf8View => DataTypeLayout::new_view(),
1738 DataType::FixedSizeList(_, _) => DataTypeLayout::new_nullable_empty(), DataType::List(_) => DataTypeLayout::new_fixed_width::<i32>(),
1740 DataType::ListView(_) => DataTypeLayout::new_list_view::<i32>(),
1741 DataType::LargeListView(_) => DataTypeLayout::new_list_view::<i64>(),
1742 DataType::LargeList(_) => DataTypeLayout::new_fixed_width::<i64>(),
1743 DataType::Map(_, _) => DataTypeLayout::new_fixed_width::<i32>(),
1744 DataType::Struct(_) => DataTypeLayout::new_nullable_empty(), DataType::RunEndEncoded(_, _) => DataTypeLayout::new_empty(), DataType::Union(_, mode) => {
1747 let type_ids = BufferSpec::FixedWidth {
1748 byte_width: mem::size_of::<i8>(),
1749 alignment: mem::align_of::<i8>(),
1750 };
1751
1752 DataTypeLayout {
1753 buffers: match mode {
1754 UnionMode::Sparse => {
1755 vec![type_ids]
1756 }
1757 UnionMode::Dense => {
1758 vec![
1759 type_ids,
1760 BufferSpec::FixedWidth {
1761 byte_width: mem::size_of::<i32>(),
1762 alignment: mem::align_of::<i32>(),
1763 },
1764 ]
1765 }
1766 },
1767 can_contain_null_mask: false,
1768 variadic: false,
1769 }
1770 }
1771 DataType::Dictionary(key_type, _value_type) => layout(key_type),
1772 }
1773}
1774
1775#[derive(Debug, PartialEq, Eq)]
1777pub struct DataTypeLayout {
1779 pub buffers: Vec<BufferSpec>,
1781
1782 pub can_contain_null_mask: bool,
1784
1785 pub variadic: bool,
1789}
1790
1791impl DataTypeLayout {
1792 pub fn new_fixed_width<T>() -> Self {
1794 Self {
1795 buffers: vec![BufferSpec::FixedWidth {
1796 byte_width: mem::size_of::<T>(),
1797 alignment: mem::align_of::<T>(),
1798 }],
1799 can_contain_null_mask: true,
1800 variadic: false,
1801 }
1802 }
1803
1804 pub fn new_nullable_empty() -> Self {
1807 Self {
1808 buffers: vec![],
1809 can_contain_null_mask: true,
1810 variadic: false,
1811 }
1812 }
1813
1814 pub fn new_empty() -> Self {
1817 Self {
1818 buffers: vec![],
1819 can_contain_null_mask: false,
1820 variadic: false,
1821 }
1822 }
1823
1824 pub fn new_binary<T>() -> Self {
1828 Self {
1829 buffers: vec![
1830 BufferSpec::FixedWidth {
1832 byte_width: mem::size_of::<T>(),
1833 alignment: mem::align_of::<T>(),
1834 },
1835 BufferSpec::VariableWidth,
1837 ],
1838 can_contain_null_mask: true,
1839 variadic: false,
1840 }
1841 }
1842
1843 pub fn new_view() -> Self {
1845 Self {
1846 buffers: vec![BufferSpec::FixedWidth {
1847 byte_width: mem::size_of::<u128>(),
1848 alignment: mem::align_of::<u128>(),
1849 }],
1850 can_contain_null_mask: true,
1851 variadic: true,
1852 }
1853 }
1854
1855 pub fn new_list_view<T>() -> Self {
1857 Self {
1858 buffers: vec![
1859 BufferSpec::FixedWidth {
1860 byte_width: mem::size_of::<T>(),
1861 alignment: mem::align_of::<T>(),
1862 },
1863 BufferSpec::FixedWidth {
1864 byte_width: mem::size_of::<T>(),
1865 alignment: mem::align_of::<T>(),
1866 },
1867 ],
1868 can_contain_null_mask: true,
1869 variadic: false,
1870 }
1871 }
1872}
1873
1874#[derive(Debug, PartialEq, Eq)]
1876pub enum BufferSpec {
1877 FixedWidth {
1888 byte_width: usize,
1890 alignment: usize,
1892 },
1893 VariableWidth,
1895 BitMap,
1901 #[allow(dead_code)]
1904 AlwaysNull,
1905}
1906
1907impl PartialEq for ArrayData {
1908 fn eq(&self, other: &Self) -> bool {
1909 equal::equal(self, other)
1910 }
1911}
1912
1913#[derive(Debug, Clone)]
1932#[doc(hidden)]
1933pub struct UnsafeFlag(bool);
1934
1935impl UnsafeFlag {
1936 #[inline]
1940 pub const fn new() -> Self {
1941 Self(false)
1942 }
1943
1944 #[inline]
1954 pub unsafe fn set(&mut self, val: bool) {
1955 self.0 = val;
1956 }
1957
1958 #[inline]
1960 pub fn get(&self) -> bool {
1961 self.0
1962 }
1963}
1964
1965impl Default for UnsafeFlag {
1967 fn default() -> Self {
1968 Self::new()
1969 }
1970}
1971
1972#[derive(Debug)]
1974pub struct ArrayDataBuilder {
1975 data_type: DataType,
1976 len: usize,
1977 null_count: Option<usize>,
1978 null_bit_buffer: Option<Buffer>,
1979 nulls: Option<NullBuffer>,
1980 offset: usize,
1981 buffers: Vec<Buffer>,
1982 child_data: Vec<ArrayData>,
1983 align_buffers: bool,
1987 skip_validation: UnsafeFlag,
1997}
1998
1999impl ArrayDataBuilder {
2000 #[inline]
2001 pub const fn new(data_type: DataType) -> Self {
2003 Self {
2004 data_type,
2005 len: 0,
2006 null_count: None,
2007 null_bit_buffer: None,
2008 nulls: None,
2009 offset: 0,
2010 buffers: vec![],
2011 child_data: vec![],
2012 align_buffers: false,
2013 skip_validation: UnsafeFlag::new(),
2014 }
2015 }
2016
2017 pub fn data_type(self, data_type: DataType) -> Self {
2019 Self { data_type, ..self }
2020 }
2021
2022 #[inline]
2023 #[allow(clippy::len_without_is_empty)]
2024 pub const fn len(mut self, n: usize) -> Self {
2026 self.len = n;
2027 self
2028 }
2029
2030 pub fn nulls(mut self, nulls: Option<NullBuffer>) -> Self {
2032 self.nulls = nulls;
2033 self.null_count = None;
2034 self.null_bit_buffer = None;
2035 self
2036 }
2037
2038 pub fn null_count(mut self, null_count: usize) -> Self {
2040 self.null_count = Some(null_count);
2041 self
2042 }
2043
2044 pub fn null_bit_buffer(mut self, buf: Option<Buffer>) -> Self {
2046 self.nulls = None;
2047 self.null_bit_buffer = buf;
2048 self
2049 }
2050
2051 #[inline]
2053 pub const fn offset(mut self, n: usize) -> Self {
2054 self.offset = n;
2055 self
2056 }
2057
2058 pub fn buffers(mut self, v: Vec<Buffer>) -> Self {
2060 self.buffers = v;
2061 self
2062 }
2063
2064 pub fn add_buffer(mut self, b: Buffer) -> Self {
2066 self.buffers.push(b);
2067 self
2068 }
2069
2070 pub fn add_buffers<I: IntoIterator<Item = Buffer>>(mut self, bs: I) -> Self {
2072 self.buffers.extend(bs);
2073 self
2074 }
2075
2076 pub fn child_data(mut self, v: Vec<ArrayData>) -> Self {
2078 self.child_data = v;
2079 self
2080 }
2081
2082 pub fn add_child_data(mut self, r: ArrayData) -> Self {
2084 self.child_data.push(r);
2085 self
2086 }
2087
2088 pub unsafe fn build_unchecked(self) -> ArrayData {
2104 unsafe { self.skip_validation(true) }.build().unwrap()
2105 }
2106
2107 pub fn build(self) -> Result<ArrayData, ArrowError> {
2116 let Self {
2117 data_type,
2118 len,
2119 null_count,
2120 null_bit_buffer,
2121 nulls,
2122 offset,
2123 buffers,
2124 child_data,
2125 align_buffers,
2126 skip_validation,
2127 } = self;
2128
2129 let nulls = nulls
2130 .or_else(|| {
2131 let buffer = null_bit_buffer?;
2132 let buffer = BooleanBuffer::new(buffer, offset, len);
2133 Some(match null_count {
2134 Some(n) => {
2135 unsafe { NullBuffer::new_unchecked(buffer, n) }
2137 }
2138 None => NullBuffer::new(buffer),
2139 })
2140 })
2141 .filter(|b| b.null_count() != 0);
2142
2143 let mut data = ArrayData {
2144 data_type,
2145 len,
2146 offset,
2147 buffers,
2148 child_data,
2149 nulls,
2150 };
2151
2152 if align_buffers {
2153 data.align_buffers();
2154 }
2155
2156 if !skip_validation.get() || cfg!(feature = "force_validate") {
2158 data.validate_data()?;
2159 }
2160 Ok(data)
2161 }
2162
2163 #[deprecated(since = "54.1.0", note = "Use ArrayData::align_buffers instead")]
2165 pub fn build_aligned(self) -> Result<ArrayData, ArrowError> {
2166 self.align_buffers(true).build()
2167 }
2168
2169 pub fn align_buffers(mut self, align_buffers: bool) -> Self {
2185 self.align_buffers = align_buffers;
2186 self
2187 }
2188
2189 pub unsafe fn skip_validation(mut self, skip_validation: bool) -> Self {
2203 unsafe {
2204 self.skip_validation.set(skip_validation);
2205 }
2206 self
2207 }
2208}
2209
2210impl From<ArrayData> for ArrayDataBuilder {
2211 fn from(d: ArrayData) -> Self {
2212 Self {
2213 data_type: d.data_type,
2214 len: d.len,
2215 offset: d.offset,
2216 buffers: d.buffers,
2217 child_data: d.child_data,
2218 nulls: d.nulls,
2219 null_bit_buffer: None,
2220 null_count: None,
2221 align_buffers: false,
2222 skip_validation: UnsafeFlag::new(),
2223 }
2224 }
2225}
2226
2227#[cfg(test)]
2228mod tests {
2229 use crate::transform::MutableArrayData;
2230
2231 use super::*;
2232 use arrow_schema::{Field, Fields, UnionFields};
2233
2234 fn make_i32_buffer(n: usize) -> Buffer {
2238 Buffer::from_slice_ref(vec![42i32; n])
2239 }
2240
2241 fn make_f32_buffer(n: usize) -> Buffer {
2243 Buffer::from_slice_ref(vec![42f32; n])
2244 }
2245
2246 #[test]
2247 fn test_builder() {
2248 let v = (0..25).collect::<Vec<i32>>();
2250 let b1 = Buffer::from_slice_ref(&v);
2251 let arr_data = ArrayData::builder(DataType::Int32)
2252 .len(20)
2253 .offset(5)
2254 .add_buffer(b1)
2255 .null_bit_buffer(Some(Buffer::from([
2256 0b01011111, 0b10110101, 0b01100011, 0b00011110,
2257 ])))
2258 .build()
2259 .unwrap();
2260
2261 assert_eq!(20, arr_data.len());
2262 assert_eq!(10, arr_data.null_count());
2263 assert_eq!(5, arr_data.offset());
2264 assert_eq!(1, arr_data.buffers().len());
2265 assert_eq!(
2266 Buffer::from_slice_ref(&v).as_slice(),
2267 arr_data.buffers()[0].as_slice()
2268 );
2269 }
2270
2271 #[test]
2272 fn test_builder_with_child_data() {
2273 let child_arr_data = ArrayData::try_new(
2274 DataType::Int32,
2275 5,
2276 None,
2277 0,
2278 vec![Buffer::from_slice_ref([1i32, 2, 3, 4, 5])],
2279 vec![],
2280 )
2281 .unwrap();
2282
2283 let field = Arc::new(Field::new("x", DataType::Int32, true));
2284 let data_type = DataType::Struct(vec![field].into());
2285
2286 let arr_data = ArrayData::builder(data_type)
2287 .len(5)
2288 .offset(0)
2289 .add_child_data(child_arr_data.clone())
2290 .build()
2291 .unwrap();
2292
2293 assert_eq!(5, arr_data.len());
2294 assert_eq!(1, arr_data.child_data().len());
2295 assert_eq!(child_arr_data, arr_data.child_data()[0]);
2296 }
2297
2298 #[test]
2299 fn test_extend_fixed_size_list_accounts_for_parent_offset() {
2300 let field = Arc::new(Field::new("item", DataType::Int32, false));
2301 let child = ArrayData::builder(DataType::Int32)
2302 .len(6)
2303 .add_buffer(Buffer::from_slice_ref([0, 1, 2, 3, 4, 5]))
2304 .build()
2305 .unwrap();
2306
2307 let data = ArrayData::builder(DataType::FixedSizeList(field, 2))
2308 .len(2)
2309 .offset(1)
2310 .add_child_data(child)
2311 .build()
2312 .unwrap();
2313
2314 let mut mutable = MutableArrayData::new(vec![&data], false, data.len());
2315 mutable.extend(0, 0, data.len());
2316 let output = mutable.freeze();
2317
2318 assert_eq!(output.child_data()[0].buffer::<i32>(0), &[2, 3, 4, 5]);
2319 }
2320
2321 #[test]
2322 fn test_null_count() {
2323 let mut bit_v: [u8; 2] = [0; 2];
2324 bit_util::set_bit(&mut bit_v, 0);
2325 bit_util::set_bit(&mut bit_v, 3);
2326 bit_util::set_bit(&mut bit_v, 10);
2327 let arr_data = ArrayData::builder(DataType::Int32)
2328 .len(16)
2329 .add_buffer(make_i32_buffer(16))
2330 .null_bit_buffer(Some(Buffer::from(bit_v)))
2331 .build()
2332 .unwrap();
2333 assert_eq!(13, arr_data.null_count());
2334
2335 let mut bit_v: [u8; 2] = [0; 2];
2337 bit_util::set_bit(&mut bit_v, 0);
2338 bit_util::set_bit(&mut bit_v, 3);
2339 bit_util::set_bit(&mut bit_v, 10);
2340 let arr_data = ArrayData::builder(DataType::Int32)
2341 .len(12)
2342 .offset(2)
2343 .add_buffer(make_i32_buffer(14)) .null_bit_buffer(Some(Buffer::from(bit_v)))
2345 .build()
2346 .unwrap();
2347 assert_eq!(10, arr_data.null_count());
2348 }
2349
2350 #[test]
2351 fn test_null_buffer_ref() {
2352 let mut bit_v: [u8; 2] = [0; 2];
2353 bit_util::set_bit(&mut bit_v, 0);
2354 bit_util::set_bit(&mut bit_v, 3);
2355 bit_util::set_bit(&mut bit_v, 10);
2356 let arr_data = ArrayData::builder(DataType::Int32)
2357 .len(16)
2358 .add_buffer(make_i32_buffer(16))
2359 .null_bit_buffer(Some(Buffer::from(bit_v)))
2360 .build()
2361 .unwrap();
2362 assert!(arr_data.nulls().is_some());
2363 assert_eq!(&bit_v, arr_data.nulls().unwrap().validity());
2364 }
2365
2366 #[test]
2367 fn test_slice() {
2368 let mut bit_v: [u8; 2] = [0; 2];
2369 bit_util::set_bit(&mut bit_v, 0);
2370 bit_util::set_bit(&mut bit_v, 3);
2371 bit_util::set_bit(&mut bit_v, 10);
2372 let data = ArrayData::builder(DataType::Int32)
2373 .len(16)
2374 .add_buffer(make_i32_buffer(16))
2375 .null_bit_buffer(Some(Buffer::from(bit_v)))
2376 .build()
2377 .unwrap();
2378 let new_data = data.slice(1, 15);
2379 assert_eq!(data.len() - 1, new_data.len());
2380 assert_eq!(1, new_data.offset());
2381 assert_eq!(data.null_count(), new_data.null_count());
2382
2383 let new_data = new_data.slice(1, 14);
2385 assert_eq!(data.len() - 2, new_data.len());
2386 assert_eq!(2, new_data.offset());
2387 assert_eq!(data.null_count() - 1, new_data.null_count());
2388 }
2389
2390 #[test]
2391 #[should_panic(expected = "offset + length overflow")]
2392 fn test_slice_panics_on_offset_length_overflow() {
2393 let data = ArrayData::builder(DataType::Int32)
2394 .len(4)
2395 .add_buffer(make_i32_buffer(4))
2396 .build()
2397 .unwrap();
2398 let sliced = data.slice(1, 3);
2399
2400 sliced.slice(1, usize::MAX);
2401 }
2402
2403 #[test]
2404 fn test_typed_offsets_length_overflow() {
2405 let data = ArrayData {
2406 data_type: DataType::Binary,
2407 len: usize::MAX,
2408 offset: 0,
2409 buffers: vec![Buffer::from_slice_ref([0_i32])],
2410 child_data: vec![],
2411 nulls: None,
2412 };
2413 let err = data.typed_offsets::<i32>().unwrap_err();
2414
2415 assert_eq!(
2416 err.to_string(),
2417 format!(
2418 "Invalid argument error: Length {} with offset 1 overflows usize for Binary",
2419 usize::MAX
2420 )
2421 );
2422 }
2423
2424 #[test]
2425 fn test_validate_typed_buffer_length_overflow() {
2426 let data = ArrayData {
2427 data_type: DataType::Binary,
2428 len: 0,
2429 offset: 2,
2430 buffers: vec![Buffer::from_slice_ref([0_i32])],
2431 child_data: vec![],
2432 nulls: None,
2433 };
2434 let err = data.typed_buffer::<i32>(0, usize::MAX).unwrap_err();
2435
2436 assert_eq!(
2437 err.to_string(),
2438 format!(
2439 "Invalid argument error: Length {} with offset 2 overflows usize for Binary",
2440 usize::MAX
2441 )
2442 );
2443 }
2444
2445 fn try_new_binary_length_offset_overflow() -> Result<ArrayData, ArrowError> {
2447 ArrayData::try_new(
2448 DataType::Binary,
2449 usize::MAX,
2450 None,
2451 1,
2452 vec![
2453 Buffer::from_slice_ref([0_i32]),
2454 Buffer::from_iter(std::iter::empty::<u8>()),
2455 ],
2456 vec![],
2457 )
2458 }
2459
2460 #[cfg(not(feature = "force_validate"))]
2461 #[test]
2462 fn test_try_new_length_offset_overflow() {
2463 let err = try_new_binary_length_offset_overflow().unwrap_err();
2464
2465 assert_eq!(
2466 err.to_string(),
2467 format!(
2468 "Invalid argument error: Length {} with offset 1 overflows usize for Binary",
2469 usize::MAX
2470 )
2471 );
2472 }
2473
2474 #[cfg(feature = "force_validate")]
2475 #[test]
2476 #[should_panic(
2477 expected = "Length 18446744073709551615 with offset 1 overflows usize for Binary"
2478 )]
2479 fn test_try_new_length_offset_overflow_force_validate() {
2480 try_new_binary_length_offset_overflow().unwrap();
2481 }
2482
2483 #[test]
2484 fn test_equality() {
2485 let int_data = ArrayData::builder(DataType::Int32)
2486 .len(1)
2487 .add_buffer(make_i32_buffer(1))
2488 .build()
2489 .unwrap();
2490
2491 let float_data = ArrayData::builder(DataType::Float32)
2492 .len(1)
2493 .add_buffer(make_f32_buffer(1))
2494 .build()
2495 .unwrap();
2496 assert_ne!(int_data, float_data);
2497 assert!(!int_data.ptr_eq(&float_data));
2498 assert!(int_data.ptr_eq(&int_data));
2499
2500 #[allow(clippy::redundant_clone)]
2501 let int_data_clone = int_data.clone();
2502 assert_eq!(int_data, int_data_clone);
2503 assert!(int_data.ptr_eq(&int_data_clone));
2504 assert!(int_data_clone.ptr_eq(&int_data));
2505
2506 let int_data_slice = int_data_clone.slice(1, 0);
2507 assert!(int_data_slice.ptr_eq(&int_data_slice));
2508 assert!(!int_data.ptr_eq(&int_data_slice));
2509 assert!(!int_data_slice.ptr_eq(&int_data));
2510
2511 let data_buffer = Buffer::from_slice_ref("abcdef".as_bytes());
2512 let offsets_buffer = Buffer::from_slice_ref([0_i32, 2_i32, 2_i32, 5_i32]);
2513 let string_data = ArrayData::try_new(
2514 DataType::Utf8,
2515 3,
2516 Some(Buffer::from_iter(vec![true, false, true])),
2517 0,
2518 vec![offsets_buffer, data_buffer],
2519 vec![],
2520 )
2521 .unwrap();
2522
2523 assert_ne!(float_data, string_data);
2524 assert!(!float_data.ptr_eq(&string_data));
2525
2526 assert!(string_data.ptr_eq(&string_data));
2527
2528 #[allow(clippy::redundant_clone)]
2529 let string_data_cloned = string_data.clone();
2530 assert!(string_data_cloned.ptr_eq(&string_data));
2531 assert!(string_data.ptr_eq(&string_data_cloned));
2532
2533 let string_data_slice = string_data.slice(1, 2);
2534 assert!(string_data_slice.ptr_eq(&string_data_slice));
2535 assert!(!string_data_slice.ptr_eq(&string_data))
2536 }
2537
2538 #[test]
2539 fn test_slice_memory_size() {
2540 let mut bit_v: [u8; 2] = [0; 2];
2541 bit_util::set_bit(&mut bit_v, 0);
2542 bit_util::set_bit(&mut bit_v, 3);
2543 bit_util::set_bit(&mut bit_v, 10);
2544 let data = ArrayData::builder(DataType::Int32)
2545 .len(16)
2546 .add_buffer(make_i32_buffer(16))
2547 .null_bit_buffer(Some(Buffer::from(bit_v)))
2548 .build()
2549 .unwrap();
2550 let new_data = data.slice(1, 14);
2551 assert_eq!(
2552 data.get_slice_memory_size().unwrap() - 8,
2553 new_data.get_slice_memory_size().unwrap()
2554 );
2555 let data_buffer = Buffer::from_slice_ref("abcdef".as_bytes());
2556 let offsets_buffer = Buffer::from_slice_ref([0_i32, 2_i32, 2_i32, 5_i32]);
2557 let string_data = ArrayData::try_new(
2558 DataType::Utf8,
2559 3,
2560 Some(Buffer::from_iter(vec![true, false, true])),
2561 0,
2562 vec![offsets_buffer, data_buffer],
2563 vec![],
2564 )
2565 .unwrap();
2566 let string_data_slice = string_data.slice(1, 2);
2567 assert_eq!(
2569 string_data.get_slice_memory_size().unwrap() - 6,
2570 string_data_slice.get_slice_memory_size().unwrap()
2571 );
2572 }
2573
2574 #[test]
2575 fn test_count_nulls() {
2576 let buffer = Buffer::from([0b00010110, 0b10011111]);
2577 let buffer = NullBuffer::new(BooleanBuffer::new(buffer, 0, 16));
2578 let count = count_nulls(Some(&buffer), 0, 16);
2579 assert_eq!(count, 7);
2580
2581 let count = count_nulls(Some(&buffer), 4, 8);
2582 assert_eq!(count, 3);
2583 }
2584
2585 #[test]
2586 fn test_contains_nulls() {
2587 let buffer: Buffer =
2588 MutableBuffer::from_iter([false, false, false, true, true, false]).into();
2589 let buffer = NullBuffer::new(BooleanBuffer::new(buffer, 0, 6));
2590 assert!(contains_nulls(Some(&buffer), 0, 6));
2591 assert!(contains_nulls(Some(&buffer), 0, 3));
2592 assert!(!contains_nulls(Some(&buffer), 3, 2));
2593 assert!(!contains_nulls(Some(&buffer), 0, 0));
2594 }
2595
2596 #[test]
2597 fn test_alignment() {
2598 let buffer = Buffer::from_vec(vec![1_i32, 2_i32, 3_i32]);
2599 let sliced = buffer.slice(1);
2600
2601 let mut data = ArrayData {
2602 data_type: DataType::Int32,
2603 len: 0,
2604 offset: 0,
2605 buffers: vec![buffer],
2606 child_data: vec![],
2607 nulls: None,
2608 };
2609 data.validate_full().unwrap();
2610
2611 data.buffers[0] = sliced;
2613 let err = data.validate().unwrap_err();
2614
2615 assert_eq!(
2616 err.to_string(),
2617 "Invalid argument error: Misaligned buffers[0] in array of type Int32, offset from expected alignment of 4 by 1"
2618 );
2619
2620 data.align_buffers();
2621 data.validate_full().unwrap();
2622 }
2623
2624 #[test]
2625 fn test_alignment_struct() {
2626 let buffer = Buffer::from_vec(vec![1_i32, 2_i32, 3_i32]);
2627 let sliced = buffer.slice(1);
2628
2629 let child_data = ArrayData {
2630 data_type: DataType::Int32,
2631 len: 0,
2632 offset: 0,
2633 buffers: vec![buffer],
2634 child_data: vec![],
2635 nulls: None,
2636 };
2637
2638 let schema = DataType::Struct(Fields::from(vec![Field::new("a", DataType::Int32, false)]));
2639 let mut data = ArrayData {
2640 data_type: schema,
2641 len: 0,
2642 offset: 0,
2643 buffers: vec![],
2644 child_data: vec![child_data],
2645 nulls: None,
2646 };
2647 data.validate_full().unwrap();
2648
2649 data.child_data[0].buffers[0] = sliced;
2651 let err = data.validate().unwrap_err();
2652
2653 assert_eq!(
2654 err.to_string(),
2655 "Invalid argument error: Misaligned buffers[0] in array of type Int32, offset from expected alignment of 4 by 1"
2656 );
2657
2658 data.align_buffers();
2659 data.validate_full().unwrap();
2660 }
2661
2662 #[test]
2663 fn test_null_view_types() {
2664 let array_len = 32;
2665 let array = ArrayData::new_null(&DataType::BinaryView, array_len);
2666 assert_eq!(array.len(), array_len);
2667 for i in 0..array.len() {
2668 assert!(array.is_null(i));
2669 }
2670
2671 let array = ArrayData::new_null(&DataType::Utf8View, array_len);
2672 assert_eq!(array.len(), array_len);
2673 for i in 0..array.len() {
2674 assert!(array.is_null(i));
2675 }
2676
2677 let array = ArrayData::new_null(
2678 &DataType::ListView(Arc::new(Field::new_list_field(DataType::Int32, true))),
2679 array_len,
2680 );
2681 assert_eq!(array.len(), array_len);
2682 for i in 0..array.len() {
2683 assert!(array.is_null(i));
2684 }
2685
2686 let array = ArrayData::new_null(
2687 &DataType::LargeListView(Arc::new(Field::new_list_field(DataType::Int32, true))),
2688 array_len,
2689 );
2690 assert_eq!(array.len(), array_len);
2691 for i in 0..array.len() {
2692 assert!(array.is_null(i));
2693 }
2694 }
2695
2696 #[test]
2697 fn test_union_validate_type_ids() {
2698 let dt = DataType::Union(
2699 UnionFields::try_new(
2700 vec![5],
2701 vec![Arc::new(Field::new("a", DataType::Int32, true))],
2702 )
2703 .unwrap(),
2704 UnionMode::Sparse,
2705 );
2706
2707 let child = ArrayData::builder(DataType::Int32)
2708 .len(2)
2709 .add_buffer(make_i32_buffer(2))
2710 .build()
2711 .unwrap();
2712
2713 let err = ArrayDataBuilder::new(dt)
2714 .len(2)
2715 .offset(0)
2716 .buffers(vec![vec![5_i8, 2].into()])
2717 .child_data(vec![child])
2718 .build()
2719 .unwrap_err();
2720 assert_eq!(
2721 err.to_string(),
2722 "Invalid argument error: Type id values must match one of the field type ids, got 2 at index 1"
2723 );
2724 }
2725
2726 #[test]
2727 fn test_union_validate_dense_offsets() {
2728 let dt = DataType::Union(
2729 UnionFields::try_new(
2730 vec![5],
2731 vec![Arc::new(Field::new("a", DataType::Int32, true))],
2732 )
2733 .unwrap(),
2734 UnionMode::Dense,
2735 );
2736
2737 let child = ArrayData::builder(DataType::Int32)
2738 .len(2)
2739 .add_buffer(make_i32_buffer(2))
2740 .build()
2741 .unwrap();
2742
2743 let err = ArrayDataBuilder::new(dt.clone())
2744 .len(2)
2745 .buffers(vec![vec![5_i8, 5].into(), vec![0_i32, 10].into()])
2747 .child_data(vec![child.clone()])
2748 .build()
2749 .unwrap_err();
2750 assert_eq!(
2751 err.to_string(),
2752 "Invalid argument error: Offsets must be non-negative and within the length of the child array"
2753 );
2754
2755 let err = ArrayDataBuilder::new(dt)
2756 .len(2)
2757 .buffers(vec![vec![5_i8, 5].into(), vec![0_i32, -1].into()])
2759 .child_data(vec![child])
2760 .build()
2761 .unwrap_err();
2762 assert_eq!(
2763 err.to_string(),
2764 "Invalid argument error: Offsets must be non-negative and within the length of the child array"
2765 );
2766 }
2767
2768 fn run_array_test_children() -> (DataType, Vec<ArrayData>) {
2771 let dt = DataType::RunEndEncoded(
2772 Arc::new(Field::new("run_ends", DataType::Int32, false)),
2773 Arc::new(Field::new("values", DataType::Float32, true)),
2774 );
2775
2776 let runs = ArrayData::builder(DataType::Int32)
2777 .len(2)
2778 .add_buffer(vec![1_i32, 3].into())
2779 .build()
2780 .unwrap();
2781 let values = ArrayData::builder(DataType::Float32)
2782 .len(2)
2783 .add_buffer(vec![1.5_f32, 2.5].into())
2784 .build()
2785 .unwrap();
2786
2787 (dt, vec![runs, values])
2788 }
2789
2790 #[test]
2791 fn test_run_array_validates_last_run() {
2792 let (dt, children) = run_array_test_children();
2793
2794 let err = ArrayDataBuilder::new(dt)
2795 .len(4)
2796 .offset(0)
2797 .child_data(children)
2798 .build()
2799 .unwrap_err();
2800 assert_eq!(
2801 err.to_string(),
2802 "Invalid argument error: 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 3 and offset + length of array is 4."
2803 );
2804 }
2805
2806 #[test]
2807 fn test_run_array_empty_validates() {
2808 let (dt, children) = run_array_test_children();
2809
2810 let data = ArrayDataBuilder::new(dt)
2811 .len(0)
2812 .offset(4)
2813 .child_data(children)
2814 .build()
2815 .unwrap();
2816 assert_eq!(data.len(), 0);
2817 }
2818
2819 #[test]
2822 fn test_run_array_validates_last_run_with_offset() {
2823 let (dt, children) = run_array_test_children();
2824
2825 let err = ArrayDataBuilder::new(dt)
2826 .len(2)
2827 .offset(2)
2828 .child_data(children)
2829 .build()
2830 .unwrap_err();
2831 assert_eq!(
2832 err.to_string(),
2833 "Invalid argument error: 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 3 and offset + length of array is 4."
2834 );
2835 }
2836
2837 #[test]
2840 fn test_run_array_sliced_validates() {
2841 let (dt, children) = run_array_test_children();
2842
2843 let data = ArrayDataBuilder::new(dt)
2844 .len(2)
2845 .offset(1)
2846 .child_data(children)
2847 .build()
2848 .unwrap();
2849 assert_eq!(data.len(), 2);
2850 assert_eq!(data.offset(), 1);
2851 }
2852}