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(
328 data_type: DataType,
329 len: usize,
330 null_bit_buffer: Option<Buffer>,
331 offset: usize,
332 buffers: Vec<Buffer>,
333 child_data: Vec<ArrayData>,
334 ) -> Result<Self, ArrowError> {
335 if let Some(null_bit_buffer) = null_bit_buffer.as_ref() {
339 let len_plus_offset = checked_len_plus_offset(&data_type, len, offset)?;
340 let needed_len = bit_util::ceil(len_plus_offset, 8);
341 if null_bit_buffer.len() < needed_len {
342 return Err(ArrowError::InvalidArgumentError(format!(
343 "null_bit_buffer size too small. got {} needed {}",
344 null_bit_buffer.len(),
345 needed_len
346 )));
347 }
348 }
349 let new_self = unsafe {
351 Self::new_unchecked(
352 data_type,
353 len,
354 None,
355 null_bit_buffer,
356 offset,
357 buffers,
358 child_data,
359 )
360 };
361
362 new_self.validate_data()?;
367 Ok(new_self)
368 }
369
370 pub fn into_parts(
376 self,
377 ) -> (
378 DataType,
379 usize,
380 Option<NullBuffer>,
381 usize,
382 Vec<Buffer>,
383 Vec<ArrayData>,
384 ) {
385 let Self {
386 data_type,
387 len,
388 nulls,
389 offset,
390 buffers,
391 child_data,
392 } = self;
393
394 (data_type, len, nulls, offset, buffers, child_data)
395 }
396
397 #[inline]
399 pub const fn builder(data_type: DataType) -> ArrayDataBuilder {
400 ArrayDataBuilder::new(data_type)
401 }
402
403 #[inline]
405 pub const fn data_type(&self) -> &DataType {
406 &self.data_type
407 }
408
409 pub fn buffers(&self) -> &[Buffer] {
411 &self.buffers
412 }
413
414 pub fn child_data(&self) -> &[ArrayData] {
417 &self.child_data[..]
418 }
419
420 #[inline]
422 pub fn is_null(&self, i: usize) -> bool {
423 match &self.nulls {
424 Some(v) => v.is_null(i),
425 None => false,
426 }
427 }
428
429 #[inline]
433 pub fn nulls(&self) -> Option<&NullBuffer> {
434 self.nulls.as_ref()
435 }
436
437 #[inline]
439 pub fn is_valid(&self, i: usize) -> bool {
440 !self.is_null(i)
441 }
442
443 #[inline]
445 pub const fn len(&self) -> usize {
446 self.len
447 }
448
449 #[inline]
451 pub const fn is_empty(&self) -> bool {
452 self.len == 0
453 }
454
455 #[inline]
457 pub const fn offset(&self) -> usize {
458 self.offset
459 }
460
461 #[inline]
463 pub fn null_count(&self) -> usize {
464 self.nulls
465 .as_ref()
466 .map(|x| x.null_count())
467 .unwrap_or_default()
468 }
469
470 pub fn get_buffer_memory_size(&self) -> usize {
482 let mut size = 0;
483 for buffer in &self.buffers {
484 size += buffer.capacity();
485 }
486 if let Some(bitmap) = &self.nulls {
487 size += bitmap.buffer().capacity()
488 }
489 for child in &self.child_data {
490 size += child.get_buffer_memory_size();
491 }
492 size
493 }
494
495 pub fn get_slice_memory_size(&self) -> Result<usize, ArrowError> {
508 let mut result: usize = 0;
509 let layout = layout(&self.data_type);
510
511 for spec in layout.buffers.iter() {
512 match spec {
513 BufferSpec::FixedWidth { byte_width, .. } => {
514 let buffer_size = self.len.checked_mul(*byte_width).ok_or_else(|| {
515 ArrowError::ComputeError(
516 "Integer overflow computing buffer size".to_string(),
517 )
518 })?;
519 result += buffer_size;
520 }
521 BufferSpec::VariableWidth => {
522 let buffer_len = match self.data_type {
523 DataType::Utf8 | DataType::Binary => {
524 let offsets = self.typed_offsets::<i32>()?;
525 (offsets[self.len] - offsets[0]) as usize
526 }
527 DataType::LargeUtf8 | DataType::LargeBinary => {
528 let offsets = self.typed_offsets::<i64>()?;
529 (offsets[self.len] - offsets[0]) as usize
530 }
531 _ => {
532 return Err(ArrowError::NotYetImplemented(format!(
533 "Invalid data type for VariableWidth buffer. Expected Utf8, LargeUtf8, Binary or LargeBinary. Got {}",
534 self.data_type
535 )));
536 }
537 };
538 result += buffer_len;
539 }
540 BufferSpec::BitMap => {
541 let buffer_size = bit_util::ceil(self.len, 8);
542 result += buffer_size;
543 }
544 BufferSpec::AlwaysNull => {
545 }
547 }
548 }
549
550 if self.nulls().is_some() {
551 result += bit_util::ceil(self.len, 8);
552 }
553
554 for child in &self.child_data {
555 result += child.get_slice_memory_size()?;
556 }
557 Ok(result)
558 }
559
560 pub fn get_array_memory_size(&self) -> usize {
569 let mut size = mem::size_of_val(self);
570
571 for buffer in &self.buffers {
573 size += mem::size_of::<Buffer>();
574 size += buffer.capacity();
575 }
576 if let Some(nulls) = &self.nulls {
577 size += nulls.buffer().capacity();
578 }
579 for child in &self.child_data {
580 size += child.get_array_memory_size();
581 }
582
583 size
584 }
585
586 pub fn slice(&self, offset: usize, length: usize) -> ArrayData {
594 let end = offset
595 .checked_add(length)
596 .expect("offset + length overflow");
597 assert!(end <= self.len());
598
599 if let DataType::Struct(_) = self.data_type() {
600 let new_offset = self.offset + offset;
602 ArrayData {
603 data_type: self.data_type().clone(),
604 len: length,
605 offset: new_offset,
606 buffers: self.buffers.clone(),
607 child_data: self
609 .child_data()
610 .iter()
611 .map(|data| data.slice(offset, length))
612 .collect(),
613 nulls: self.nulls.as_ref().map(|x| x.slice(offset, length)),
614 }
615 } else {
616 let mut new_data = self.clone();
617
618 new_data.len = length;
619 new_data.offset = offset + self.offset;
620 new_data.nulls = self.nulls.as_ref().map(|x| x.slice(offset, length));
621
622 new_data
623 }
624 }
625
626 pub fn buffer<T: ArrowNativeType>(&self, buffer: usize) -> &[T] {
633 &self.buffers()[buffer].typed_data()[self.offset..]
634 }
635
636 pub fn new_null(data_type: &DataType, len: usize) -> Self {
638 let bit_len = bit_util::ceil(len, 8);
639 let zeroed = |len: usize| Buffer::from(MutableBuffer::from_len_zeroed(len));
640
641 let (buffers, child_data, has_nulls) = match data_type.primitive_width() {
642 Some(width) => (vec![zeroed(width * len)], vec![], true),
643 None => match data_type {
644 DataType::Null => (vec![], vec![], false),
645 DataType::Boolean => (vec![zeroed(bit_len)], vec![], true),
646 DataType::Binary | DataType::Utf8 => {
647 (vec![zeroed((len + 1) * 4), zeroed(0)], vec![], true)
648 }
649 DataType::BinaryView | DataType::Utf8View => (vec![zeroed(len * 16)], vec![], true),
650 DataType::LargeBinary | DataType::LargeUtf8 => {
651 (vec![zeroed((len + 1) * 8), zeroed(0)], vec![], true)
652 }
653 DataType::FixedSizeBinary(i) => (vec![zeroed(*i as usize * len)], vec![], true),
654 DataType::List(f) | DataType::Map(f, _) => (
655 vec![zeroed((len + 1) * 4)],
656 vec![ArrayData::new_empty(f.data_type())],
657 true,
658 ),
659 DataType::LargeList(f) => (
660 vec![zeroed((len + 1) * 8)],
661 vec![ArrayData::new_empty(f.data_type())],
662 true,
663 ),
664 DataType::ListView(f) => (
665 vec![zeroed(len * 4), zeroed(len * 4)],
666 vec![ArrayData::new_empty(f.data_type())],
667 true,
668 ),
669 DataType::LargeListView(f) => (
670 vec![zeroed(len * 8), zeroed(len * 8)],
671 vec![ArrayData::new_empty(f.data_type())],
672 true,
673 ),
674 DataType::FixedSizeList(f, list_len) => (
675 vec![],
676 vec![ArrayData::new_null(f.data_type(), *list_len as usize * len)],
677 true,
678 ),
679 DataType::Struct(fields) => (
680 vec![],
681 fields
682 .iter()
683 .map(|f| Self::new_null(f.data_type(), len))
684 .collect(),
685 true,
686 ),
687 DataType::Dictionary(k, v) => (
688 vec![zeroed(k.primitive_width().unwrap() * len)],
689 vec![ArrayData::new_empty(v.as_ref())],
690 true,
691 ),
692 DataType::Union(f, mode) => {
693 let (id, _) = f.iter().next().unwrap();
694 let ids = Buffer::from_iter(std::iter::repeat_n(id, len));
695 let buffers = match mode {
696 UnionMode::Sparse => vec![ids],
697 UnionMode::Dense => {
698 let end_offset = i32::from_usize(len).unwrap();
699 vec![ids, Buffer::from_iter(0_i32..end_offset)]
700 }
701 };
702
703 let children = f
704 .iter()
705 .enumerate()
706 .map(|(idx, (_, f))| {
707 if idx == 0 || *mode == UnionMode::Sparse {
708 Self::new_null(f.data_type(), len)
709 } else {
710 Self::new_empty(f.data_type())
711 }
712 })
713 .collect();
714
715 (buffers, children, false)
716 }
717 DataType::RunEndEncoded(r, v) => {
718 if len == 0 {
719 let runs = ArrayData::new_empty(r.data_type());
721 let values = ArrayData::new_empty(v.data_type());
722 (vec![], vec![runs, values], false)
723 } else {
724 let runs = match r.data_type() {
725 DataType::Int16 => {
726 let i = i16::from_usize(len).expect("run overflow");
727 Buffer::from_slice_ref([i])
728 }
729 DataType::Int32 => {
730 let i = i32::from_usize(len).expect("run overflow");
731 Buffer::from_slice_ref([i])
732 }
733 DataType::Int64 => {
734 let i = i64::from_usize(len).expect("run overflow");
735 Buffer::from_slice_ref([i])
736 }
737 dt => unreachable!("Invalid run ends data type {dt}"),
738 };
739
740 let builder = ArrayData::builder(r.data_type().clone())
741 .len(1)
742 .buffers(vec![runs]);
743
744 let runs = unsafe { builder.build_unchecked() };
747 (
748 vec![],
749 vec![runs, ArrayData::new_null(v.data_type(), 1)],
750 false,
751 )
752 }
753 }
754 DataType::Int8
756 | DataType::Int16
757 | DataType::Int32
758 | DataType::Int64
759 | DataType::UInt8
760 | DataType::UInt16
761 | DataType::UInt32
762 | DataType::UInt64
763 | DataType::Float16
764 | DataType::Float32
765 | DataType::Float64
766 | DataType::Timestamp(_, _)
767 | DataType::Date32
768 | DataType::Date64
769 | DataType::Time32(_)
770 | DataType::Time64(_)
771 | DataType::Duration(_)
772 | DataType::Interval(_)
773 | DataType::Decimal32(_, _)
774 | DataType::Decimal64(_, _)
775 | DataType::Decimal128(_, _)
776 | DataType::Decimal256(_, _) => unreachable!("{data_type}"),
777 },
778 };
779
780 let mut builder = ArrayDataBuilder::new(data_type.clone())
781 .len(len)
782 .buffers(buffers)
783 .child_data(child_data);
784
785 if has_nulls {
786 builder = builder.nulls(Some(NullBuffer::new_null(len)))
787 }
788
789 unsafe { builder.build_unchecked() }
792 }
793
794 pub fn new_empty(data_type: &DataType) -> Self {
796 Self::new_null(data_type, 0)
797 }
798
799 pub fn align_buffers(&mut self) {
808 let layout = layout(&self.data_type);
809 for (buffer, spec) in self.buffers.iter_mut().zip(&layout.buffers) {
810 if let BufferSpec::FixedWidth { alignment, .. } = spec {
811 if buffer.as_ptr().align_offset(*alignment) != 0 {
812 *buffer = Buffer::from_slice_ref(buffer.as_ref());
813 }
814 }
815 }
816 for data in self.child_data.iter_mut() {
818 data.align_buffers()
819 }
820 }
821
822 pub fn validate(&self) -> Result<(), ArrowError> {
833 let len_plus_offset = checked_len_plus_offset(&self.data_type, self.len, self.offset)?;
835
836 let layout = layout(&self.data_type);
838
839 if !layout.can_contain_null_mask && self.nulls.is_some() {
840 return Err(ArrowError::InvalidArgumentError(format!(
841 "Arrays of type {:?} cannot contain a null bitmask",
842 self.data_type,
843 )));
844 }
845
846 if self.buffers.len() < layout.buffers.len()
848 || (!layout.variadic && self.buffers.len() != layout.buffers.len())
849 {
850 return Err(ArrowError::InvalidArgumentError(format!(
851 "Expected {} buffers in array of type {:?}, got {}",
852 layout.buffers.len(),
853 self.data_type,
854 self.buffers.len(),
855 )));
856 }
857
858 for (i, (buffer, spec)) in self.buffers.iter().zip(layout.buffers.iter()).enumerate() {
859 match spec {
860 BufferSpec::FixedWidth {
861 byte_width,
862 alignment,
863 } => {
864 let min_buffer_size = len_plus_offset.saturating_mul(*byte_width);
865
866 if buffer.len() < min_buffer_size {
867 return Err(ArrowError::InvalidArgumentError(format!(
868 "Need at least {} bytes in buffers[{}] in array of type {:?}, but got {}",
869 min_buffer_size,
870 i,
871 self.data_type,
872 buffer.len()
873 )));
874 }
875
876 let align_offset = buffer.as_ptr().align_offset(*alignment);
877 if align_offset != 0 {
878 return Err(ArrowError::InvalidArgumentError(format!(
879 "Misaligned buffers[{i}] in array of type {:?}, offset from expected alignment of {alignment} by {}",
880 self.data_type,
881 align_offset.min(alignment - align_offset)
882 )));
883 }
884 }
885 BufferSpec::VariableWidth => {
886 }
890 BufferSpec::BitMap => {
891 let min_buffer_size = bit_util::ceil(len_plus_offset, 8);
892 if buffer.len() < min_buffer_size {
893 return Err(ArrowError::InvalidArgumentError(format!(
894 "Need at least {} bytes for bitmap in buffers[{}] in array of type {:?}, but got {}",
895 min_buffer_size,
896 i,
897 self.data_type,
898 buffer.len()
899 )));
900 }
901 }
902 BufferSpec::AlwaysNull => {
903 }
905 }
906 }
907
908 if let Some(nulls) = self.nulls() {
910 if nulls.null_count() > self.len {
911 return Err(ArrowError::InvalidArgumentError(format!(
912 "null_count {} for an array exceeds length of {} elements",
913 nulls.null_count(),
914 self.len
915 )));
916 }
917
918 let actual_len = nulls.validity().len();
919 let needed_len = bit_util::ceil(len_plus_offset, 8);
920 if actual_len < needed_len {
921 return Err(ArrowError::InvalidArgumentError(format!(
922 "null_bit_buffer size too small. got {actual_len} needed {needed_len}",
923 )));
924 }
925
926 if nulls.len() != self.len {
927 return Err(ArrowError::InvalidArgumentError(format!(
928 "null buffer incorrect size. got {} expected {}",
929 nulls.len(),
930 self.len
931 )));
932 }
933 }
934
935 self.validate_child_data()?;
936
937 match &self.data_type {
939 DataType::Utf8 | DataType::Binary => {
940 self.validate_offsets::<i32>(self.buffers[1].len())?;
941 }
942 DataType::LargeUtf8 | DataType::LargeBinary => {
943 self.validate_offsets::<i64>(self.buffers[1].len())?;
944 }
945 DataType::Dictionary(key_type, _value_type) => {
946 if !DataType::is_dictionary_key_type(key_type) {
948 return Err(ArrowError::InvalidArgumentError(format!(
949 "Dictionary key type must be integer, but was {key_type}"
950 )));
951 }
952 }
953 DataType::RunEndEncoded(run_ends_type, _) => {
954 if run_ends_type.is_nullable() {
955 return Err(ArrowError::InvalidArgumentError(
956 "The nullable should be set to false for the field defining run_ends array.".to_string()
957 ));
958 }
959 if !DataType::is_run_ends_type(run_ends_type.data_type()) {
960 return Err(ArrowError::InvalidArgumentError(format!(
961 "RunArray run_ends types must be Int16, Int32 or Int64, but was {}",
962 run_ends_type.data_type()
963 )));
964 }
965 }
966 _ => {}
967 };
968
969 Ok(())
970 }
971
972 fn typed_offsets<T: ArrowNativeType + num_traits::Num>(&self) -> Result<&[T], ArrowError> {
979 if self.len == 0 && self.buffers[0].is_empty() {
981 return Ok(&[]);
982 }
983
984 let len = checked_len_plus_offset(&self.data_type, self.len, 1)?;
985
986 self.typed_buffer(0, len)
987 }
988
989 fn typed_buffer<T: ArrowNativeType + num_traits::Num>(
991 &self,
992 idx: usize,
993 len: usize,
994 ) -> Result<&[T], ArrowError> {
995 let buffer = &self.buffers[idx];
996
997 let required_elements = checked_len_plus_offset(&self.data_type, len, self.offset)?;
998 let byte_width = mem::size_of::<T>();
999 let required_len = required_elements.checked_mul(byte_width).ok_or_else(|| {
1000 ArrowError::InvalidArgumentError(format!(
1001 "Buffer {idx} of {} byte length overflow: {} elements of {} bytes exceeds usize",
1002 self.data_type, required_elements, byte_width
1003 ))
1004 })?;
1005
1006 if buffer.len() < required_len {
1007 return Err(ArrowError::InvalidArgumentError(format!(
1008 "Buffer {} of {} isn't large enough. Expected {} bytes got {}",
1009 idx,
1010 self.data_type,
1011 required_len,
1012 buffer.len()
1013 )));
1014 }
1015
1016 Ok(&buffer.typed_data::<T>()[self.offset..required_elements])
1017 }
1018
1019 fn validate_offsets<T: ArrowNativeType + num_traits::Num + std::fmt::Display>(
1022 &self,
1023 values_length: usize,
1024 ) -> Result<(), ArrowError> {
1025 let offsets = self.typed_offsets::<T>()?;
1027 if offsets.is_empty() {
1028 return Ok(());
1029 }
1030
1031 let first_offset = offsets[0].to_usize().ok_or_else(|| {
1032 ArrowError::InvalidArgumentError(format!(
1033 "Error converting offset[0] ({}) to usize for {}",
1034 offsets[0], self.data_type
1035 ))
1036 })?;
1037
1038 let last_offset = offsets[self.len].to_usize().ok_or_else(|| {
1039 ArrowError::InvalidArgumentError(format!(
1040 "Error converting offset[{}] ({}) to usize for {}",
1041 self.len, offsets[self.len], self.data_type
1042 ))
1043 })?;
1044
1045 if first_offset > values_length {
1046 return Err(ArrowError::InvalidArgumentError(format!(
1047 "First offset {} of {} is larger than values length {}",
1048 first_offset, self.data_type, values_length,
1049 )));
1050 }
1051
1052 if last_offset > values_length {
1053 return Err(ArrowError::InvalidArgumentError(format!(
1054 "Last offset {} of {} is larger than values length {}",
1055 last_offset, self.data_type, values_length,
1056 )));
1057 }
1058
1059 if first_offset > last_offset {
1060 return Err(ArrowError::InvalidArgumentError(format!(
1061 "First offset {} in {} is smaller than last offset {}",
1062 first_offset, self.data_type, last_offset,
1063 )));
1064 }
1065
1066 Ok(())
1067 }
1068
1069 fn validate_offsets_and_sizes<T: ArrowNativeType + num_traits::Num + std::fmt::Display>(
1072 &self,
1073 values_length: usize,
1074 ) -> Result<(), ArrowError> {
1075 let offsets: &[T] = self.typed_buffer(0, self.len)?;
1076 let sizes: &[T] = self.typed_buffer(1, self.len)?;
1077 if offsets.len() != sizes.len() {
1078 return Err(ArrowError::ComputeError(format!(
1079 "ListView offsets len {} does not match sizes len {}",
1080 offsets.len(),
1081 sizes.len()
1082 )));
1083 }
1084
1085 for i in 0..sizes.len() {
1086 let size = sizes[i].to_usize().ok_or_else(|| {
1087 ArrowError::InvalidArgumentError(format!(
1088 "Error converting size[{}] ({}) to usize for {}",
1089 i, sizes[i], self.data_type
1090 ))
1091 })?;
1092 let offset = offsets[i].to_usize().ok_or_else(|| {
1093 ArrowError::InvalidArgumentError(format!(
1094 "Error converting offset[{}] ({}) to usize for {}",
1095 i, offsets[i], self.data_type
1096 ))
1097 })?;
1098 if size
1099 .checked_add(offset)
1100 .expect("Offset and size have exceeded the usize boundary")
1101 > values_length
1102 {
1103 return Err(ArrowError::InvalidArgumentError(format!(
1104 "Size {} at index {} is larger than the remaining values for {}",
1105 size, i, self.data_type
1106 )));
1107 }
1108 }
1109 Ok(())
1110 }
1111
1112 fn validate_child_data(&self) -> Result<(), ArrowError> {
1114 match &self.data_type {
1115 DataType::List(field) | DataType::Map(field, _) => {
1116 let values_data = self.get_single_valid_child_data(field.data_type())?;
1117 self.validate_offsets::<i32>(values_data.len)?;
1118 Ok(())
1119 }
1120 DataType::LargeList(field) => {
1121 let values_data = self.get_single_valid_child_data(field.data_type())?;
1122 self.validate_offsets::<i64>(values_data.len)?;
1123 Ok(())
1124 }
1125 DataType::ListView(field) => {
1126 let values_data = self.get_single_valid_child_data(field.data_type())?;
1127 self.validate_offsets_and_sizes::<i32>(values_data.len)?;
1128 Ok(())
1129 }
1130 DataType::LargeListView(field) => {
1131 let values_data = self.get_single_valid_child_data(field.data_type())?;
1132 self.validate_offsets_and_sizes::<i64>(values_data.len)?;
1133 Ok(())
1134 }
1135 DataType::FixedSizeList(field, list_size) => {
1136 let values_data = self.get_single_valid_child_data(field.data_type())?;
1137
1138 let list_size: usize = (*list_size).try_into().map_err(|_| {
1139 ArrowError::InvalidArgumentError(format!(
1140 "{} has a negative list_size {}",
1141 self.data_type, list_size
1142 ))
1143 })?;
1144
1145 let expected_values_len = self.len
1146 .checked_mul(list_size)
1147 .expect("integer overflow computing expected number of expected values in FixedListSize");
1148
1149 if values_data.len < expected_values_len {
1150 return Err(ArrowError::InvalidArgumentError(format!(
1151 "Values length {} is less than the length ({}) multiplied by the value size ({}) for {}",
1152 values_data.len, self.len, list_size, self.data_type
1153 )));
1154 }
1155
1156 Ok(())
1157 }
1158 DataType::Struct(fields) => {
1159 self.validate_num_child_data(fields.len())?;
1160 for (i, field) in fields.iter().enumerate() {
1161 let field_data = self.get_valid_child_data(i, field.data_type())?;
1162
1163 if field_data.len < self.len {
1165 return Err(ArrowError::InvalidArgumentError(format!(
1166 "{} child array #{} for field {} has length smaller than expected for struct array ({} < {})",
1167 self.data_type,
1168 i,
1169 field.name(),
1170 field_data.len,
1171 self.len
1172 )));
1173 }
1174 }
1175 Ok(())
1176 }
1177 DataType::RunEndEncoded(run_ends_field, values_field) => {
1178 self.validate_num_child_data(2)?;
1179 let run_ends_data = self.get_valid_child_data(0, run_ends_field.data_type())?;
1180 let values_data = self.get_valid_child_data(1, values_field.data_type())?;
1181 if run_ends_data.len != values_data.len {
1182 return Err(ArrowError::InvalidArgumentError(format!(
1183 "The run_ends array length should be the same as values array length. Run_ends array length is {}, values array length is {}",
1184 run_ends_data.len, values_data.len
1185 )));
1186 }
1187 if run_ends_data.nulls.is_some() {
1188 return Err(ArrowError::InvalidArgumentError(
1189 "Found null values in run_ends array. The run_ends array should not have null values.".to_string(),
1190 ));
1191 }
1192 Ok(())
1193 }
1194 DataType::Union(fields, mode) => {
1195 self.validate_num_child_data(fields.len())?;
1196
1197 for (i, (_, field)) in fields.iter().enumerate() {
1198 let field_data = self.get_valid_child_data(i, field.data_type())?;
1199
1200 if mode == &UnionMode::Sparse {
1201 let len_plus_offset =
1202 checked_len_plus_offset(&self.data_type, self.len, self.offset)?;
1203 if field_data.len < len_plus_offset {
1204 return Err(ArrowError::InvalidArgumentError(format!(
1205 "Sparse union child array #{} has length smaller than expected for union array ({} < {})",
1206 i, field_data.len, len_plus_offset
1207 )));
1208 }
1209 }
1210 }
1211 Ok(())
1212 }
1213 DataType::Dictionary(_key_type, value_type) => {
1214 self.get_single_valid_child_data(value_type)?;
1215 Ok(())
1216 }
1217 _ => {
1218 if !self.child_data.is_empty() {
1220 return Err(ArrowError::InvalidArgumentError(format!(
1221 "Expected no child arrays for type {} but got {}",
1222 self.data_type,
1223 self.child_data.len()
1224 )));
1225 }
1226 Ok(())
1227 }
1228 }
1229 }
1230
1231 fn get_single_valid_child_data(
1235 &self,
1236 expected_type: &DataType,
1237 ) -> Result<&ArrayData, ArrowError> {
1238 self.validate_num_child_data(1)?;
1239 self.get_valid_child_data(0, expected_type)
1240 }
1241
1242 fn validate_num_child_data(&self, expected_len: usize) -> Result<(), ArrowError> {
1244 if self.child_data.len() != expected_len {
1245 Err(ArrowError::InvalidArgumentError(format!(
1246 "Value data for {} should contain {} child data array(s), had {}",
1247 self.data_type,
1248 expected_len,
1249 self.child_data.len()
1250 )))
1251 } else {
1252 Ok(())
1253 }
1254 }
1255
1256 fn get_valid_child_data(
1259 &self,
1260 i: usize,
1261 expected_type: &DataType,
1262 ) -> Result<&ArrayData, ArrowError> {
1263 let values_data = self.child_data.get(i).ok_or_else(|| {
1264 ArrowError::InvalidArgumentError(format!(
1265 "{} did not have enough child arrays. Expected at least {} but had only {}",
1266 self.data_type,
1267 i + 1,
1268 self.child_data.len()
1269 ))
1270 })?;
1271
1272 if expected_type != &values_data.data_type {
1273 return Err(ArrowError::InvalidArgumentError(format!(
1274 "Child type mismatch for {}. Expected {} but child data had {}",
1275 self.data_type, expected_type, values_data.data_type
1276 )));
1277 }
1278
1279 values_data.validate()?;
1280 Ok(values_data)
1281 }
1282
1283 pub fn validate_data(&self) -> Result<(), ArrowError> {
1299 self.validate()?;
1300
1301 self.validate_nulls()?;
1302 self.validate_values()?;
1303 Ok(())
1304 }
1305
1306 pub fn validate_full(&self) -> Result<(), ArrowError> {
1311 self.validate_data()?;
1312 self.child_data
1314 .iter()
1315 .enumerate()
1316 .try_for_each(|(i, child_data)| {
1317 child_data.validate_full().map_err(|e| {
1318 ArrowError::InvalidArgumentError(format!(
1319 "{} child #{} invalid: {}",
1320 self.data_type, i, e
1321 ))
1322 })
1323 })?;
1324 Ok(())
1325 }
1326
1327 pub fn validate_nulls(&self) -> Result<(), ArrowError> {
1337 if let Some(nulls) = &self.nulls {
1338 let actual = nulls.len() - nulls.inner().count_set_bits();
1339 if actual != nulls.null_count() {
1340 return Err(ArrowError::InvalidArgumentError(format!(
1341 "null_count value ({}) doesn't match actual number of nulls in array ({})",
1342 nulls.null_count(),
1343 actual
1344 )));
1345 }
1346 }
1347
1348 match &self.data_type {
1353 DataType::List(f) | DataType::LargeList(f) | DataType::Map(f, _) => {
1354 if !f.is_nullable() {
1355 self.validate_non_nullable(None, &self.child_data[0])?
1356 }
1357 }
1358 DataType::FixedSizeList(field, len) => {
1359 let child = &self.child_data[0];
1360 if !field.is_nullable() {
1361 match &self.nulls {
1362 Some(nulls) => {
1363 let element_len = *len as usize;
1364 let expanded = nulls.expand(element_len);
1365 self.validate_non_nullable(Some(&expanded), child)?;
1366 }
1367 None => self.validate_non_nullable(None, child)?,
1368 }
1369 }
1370 }
1371 DataType::Struct(fields) => {
1372 for (field, child) in fields.iter().zip(&self.child_data) {
1373 if !field.is_nullable() {
1374 self.validate_non_nullable(self.nulls(), child)?
1375 }
1376 }
1377 }
1378 _ => {}
1379 }
1380
1381 Ok(())
1382 }
1383
1384 fn validate_non_nullable(
1386 &self,
1387 mask: Option<&NullBuffer>,
1388 child: &ArrayData,
1389 ) -> Result<(), ArrowError> {
1390 let mask = match mask {
1391 Some(mask) => mask,
1392 None => {
1393 return match child.null_count() {
1394 0 => Ok(()),
1395 _ => Err(ArrowError::InvalidArgumentError(format!(
1396 "non-nullable child of type {} contains nulls not present in parent {}",
1397 child.data_type, self.data_type
1398 ))),
1399 };
1400 }
1401 };
1402
1403 match child.nulls() {
1404 Some(nulls) if !mask.contains(nulls) => Err(ArrowError::InvalidArgumentError(format!(
1405 "non-nullable child of type {} contains nulls not present in parent",
1406 child.data_type
1407 ))),
1408 _ => Ok(()),
1409 }
1410 }
1411
1412 pub fn validate_values(&self) -> Result<(), ArrowError> {
1418 match &self.data_type {
1419 DataType::Utf8 => self.validate_utf8::<i32>(),
1420 DataType::LargeUtf8 => self.validate_utf8::<i64>(),
1421 DataType::Binary => self.validate_offsets_full::<i32>(self.buffers[1].len()),
1422 DataType::LargeBinary => self.validate_offsets_full::<i64>(self.buffers[1].len()),
1423 DataType::BinaryView => {
1424 let views = self.typed_buffer::<u128>(0, self.len)?;
1425 validate_binary_view(views, &self.buffers[1..])
1426 }
1427 DataType::Utf8View => {
1428 let views = self.typed_buffer::<u128>(0, self.len)?;
1429 validate_string_view(views, &self.buffers[1..])
1430 }
1431 DataType::List(_) | DataType::Map(_, _) => {
1432 let child = &self.child_data[0];
1433 self.validate_offsets_full::<i32>(child.len)
1434 }
1435 DataType::LargeList(_) => {
1436 let child = &self.child_data[0];
1437 self.validate_offsets_full::<i64>(child.len)
1438 }
1439 DataType::Union(fields, mode) => {
1440 let mut child_lens: Vec<Option<usize>> = vec![None; 128];
1441 for (index, (type_id, _)) in fields.iter().enumerate() {
1442 let child = &self.child_data[index];
1443 child_lens[type_id as usize] = Some(child.len());
1444 }
1445
1446 let type_ids = self.typed_buffer::<i8>(0, self.len)?;
1448 for (index, &type_id) in type_ids.iter().enumerate() {
1449 if type_id < 0 {
1450 return Err(ArrowError::InvalidArgumentError(format!(
1451 "Type id values cannot be negative, got {type_id} at index {index}"
1452 )));
1453 }
1454 let type_id = type_id as usize;
1455 if child_lens[type_id].is_none() {
1456 return Err(ArrowError::InvalidArgumentError(format!(
1457 "Type id values must match one of the field type ids, got {type_id} at index {index}"
1458 )));
1459 }
1460 }
1461
1462 if mode == &UnionMode::Dense {
1464 let offsets = self.typed_buffer::<i32>(1, self.len)?;
1465
1466 if type_ids
1467 .iter()
1468 .zip(offsets.iter())
1469 .any(|(&type_id, &offset)| {
1470 offset < 0 || offset as usize >= child_lens[type_id as usize].unwrap()
1471 })
1472 {
1473 return Err(ArrowError::InvalidArgumentError(
1474 "Offsets must be non-negative and within the length of the child array"
1475 .to_string(),
1476 ));
1477 }
1478 }
1479
1480 Ok(())
1481 }
1482 DataType::Dictionary(key_type, _value_type) => {
1483 let dictionary_length: i64 = self.child_data[0].len.try_into().unwrap();
1484 let max_value = dictionary_length - 1;
1485 match key_type.as_ref() {
1486 DataType::UInt8 => self.check_bounds::<u8>(max_value),
1487 DataType::UInt16 => self.check_bounds::<u16>(max_value),
1488 DataType::UInt32 => self.check_bounds::<u32>(max_value),
1489 DataType::UInt64 => self.check_bounds::<u64>(max_value),
1490 DataType::Int8 => self.check_bounds::<i8>(max_value),
1491 DataType::Int16 => self.check_bounds::<i16>(max_value),
1492 DataType::Int32 => self.check_bounds::<i32>(max_value),
1493 DataType::Int64 => self.check_bounds::<i64>(max_value),
1494 _ => unreachable!(),
1495 }
1496 }
1497 DataType::RunEndEncoded(run_ends, _values) => {
1498 let len_plus_offset =
1499 checked_len_plus_offset(&self.data_type, self.len, self.offset)?;
1500 let run_ends_data = self.child_data()[0].clone();
1501 match run_ends.data_type() {
1502 DataType::Int16 => {
1503 run_ends_data.check_run_ends::<i16>(self.len, len_plus_offset)
1504 }
1505 DataType::Int32 => {
1506 run_ends_data.check_run_ends::<i32>(self.len, len_plus_offset)
1507 }
1508 DataType::Int64 => {
1509 run_ends_data.check_run_ends::<i64>(self.len, len_plus_offset)
1510 }
1511 _ => unreachable!(),
1512 }
1513 }
1514 _ => {
1515 Ok(())
1517 }
1518 }
1519 }
1520
1521 fn validate_each_offset<T, V>(&self, offset_limit: usize, validate: V) -> Result<(), ArrowError>
1532 where
1533 T: ArrowNativeType + TryInto<usize> + num_traits::Num + std::fmt::Display,
1534 V: Fn(usize, Range<usize>) -> Result<(), ArrowError>,
1535 {
1536 self.typed_offsets::<T>()?
1537 .iter()
1538 .enumerate()
1539 .map(|(i, x)| {
1540 let r = x.to_usize().ok_or_else(|| {
1542 ArrowError::InvalidArgumentError(format!(
1543 "Offset invariant failure: Could not convert offset {x} to usize at position {i}"))}
1544 );
1545 match r {
1547 Ok(n) if n <= offset_limit => Ok((i, n)),
1548 Ok(_) => Err(ArrowError::InvalidArgumentError(format!(
1549 "Offset invariant failure: offset at position {i} out of bounds: {x} > {offset_limit}"))
1550 ),
1551 Err(e) => Err(e),
1552 }
1553 })
1554 .scan(0_usize, |start, end| {
1555 match end {
1557 Ok((i, end)) if *start <= end => {
1558 let range = Some(Ok((i, *start..end)));
1559 *start = end;
1560 range
1561 }
1562 Ok((i, end)) => Some(Err(ArrowError::InvalidArgumentError(format!(
1563 "Offset invariant failure: non-monotonic offset at slot {}: {} > {}",
1564 i - 1, start, end))
1565 )),
1566 Err(err) => Some(Err(err)),
1567 }
1568 })
1569 .skip(1) .try_for_each(|res: Result<(usize, Range<usize>), ArrowError>| {
1571 let (item_index, range) = res?;
1572 validate(item_index-1, range)
1573 })
1574 }
1575
1576 fn validate_utf8<T>(&self) -> Result<(), ArrowError>
1579 where
1580 T: ArrowNativeType + TryInto<usize> + num_traits::Num + std::fmt::Display,
1581 {
1582 let values_buffer = &self.buffers[1].as_slice();
1583 if let Ok(values_str) = std::str::from_utf8(values_buffer) {
1584 self.validate_each_offset::<T, _>(values_buffer.len(), |string_index, range| {
1586 if !values_str.is_char_boundary(range.start)
1587 || !values_str.is_char_boundary(range.end)
1588 {
1589 return Err(ArrowError::InvalidArgumentError(format!(
1590 "incomplete utf-8 byte sequence from index {string_index}"
1591 )));
1592 }
1593 Ok(())
1594 })
1595 } else {
1596 self.validate_each_offset::<T, _>(values_buffer.len(), |string_index, range| {
1598 std::str::from_utf8(&values_buffer[range.clone()]).map_err(|e| {
1599 ArrowError::InvalidArgumentError(format!(
1600 "Invalid UTF8 sequence at string index {string_index} ({range:?}): {e}"
1601 ))
1602 })?;
1603 Ok(())
1604 })
1605 }
1606 }
1607
1608 fn validate_offsets_full<T>(&self, offset_limit: usize) -> Result<(), ArrowError>
1611 where
1612 T: ArrowNativeType + TryInto<usize> + num_traits::Num + std::fmt::Display,
1613 {
1614 self.validate_each_offset::<T, _>(offset_limit, |_string_index, _range| {
1615 Ok(())
1618 })
1619 }
1620
1621 fn check_bounds<T>(&self, max_value: i64) -> Result<(), ArrowError>
1624 where
1625 T: ArrowNativeType + TryInto<i64> + num_traits::Num + std::fmt::Display,
1626 {
1627 let required_len = checked_len_plus_offset(&self.data_type, self.len, self.offset)?;
1628 let buffer = &self.buffers[0];
1629
1630 assert!(buffer.len() / mem::size_of::<T>() >= required_len);
1633
1634 let indexes: &[T] = &buffer.typed_data::<T>()[self.offset..required_len];
1636
1637 indexes.iter().enumerate().try_for_each(|(i, &dict_index)| {
1638 if self.is_null(i) {
1640 return Ok(());
1641 }
1642 let dict_index: i64 = dict_index.try_into().map_err(|_| {
1643 ArrowError::InvalidArgumentError(format!(
1644 "Value at position {i} out of bounds: {dict_index} (can not convert to i64)"
1645 ))
1646 })?;
1647
1648 if dict_index < 0 || dict_index > max_value {
1649 return Err(ArrowError::InvalidArgumentError(format!(
1650 "Value at position {i} out of bounds: {dict_index} (should be in [0, {max_value}])"
1651 )));
1652 }
1653 Ok(())
1654 })
1655 }
1656
1657 fn check_run_ends<T>(
1659 &self,
1660 parent_len: usize,
1661 parent_len_plus_offset: usize,
1662 ) -> Result<(), ArrowError>
1663 where
1664 T: ArrowNativeType + TryInto<i64> + num_traits::Num + std::fmt::Display,
1665 {
1666 let values = self.typed_buffer::<T>(0, self.len)?;
1667 let mut prev_value: i64 = 0_i64;
1668 values.iter().enumerate().try_for_each(|(ix, &inp_value)| {
1669 let value: i64 = inp_value.try_into().map_err(|_| {
1670 ArrowError::InvalidArgumentError(format!(
1671 "Value at position {ix} out of bounds: {inp_value} (can not convert to i64)"
1672 ))
1673 })?;
1674 if value <= 0_i64 {
1675 return Err(ArrowError::InvalidArgumentError(format!(
1676 "The values in run_ends array should be strictly positive. Found value {value} at index {ix} that does not match the criteria."
1677 )));
1678 }
1679 if ix > 0 && value <= prev_value {
1680 return Err(ArrowError::InvalidArgumentError(format!(
1681 "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."
1682 )));
1683 }
1684
1685 prev_value = value;
1686 Ok(())
1687 })?;
1688
1689 if parent_len > 0 && prev_value.as_usize() < parent_len_plus_offset {
1690 return Err(ArrowError::InvalidArgumentError(format!(
1691 "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}."
1692 )));
1693 }
1694 Ok(())
1695 }
1696
1697 pub fn ptr_eq(&self, other: &Self) -> bool {
1701 if self.offset != other.offset
1702 || self.len != other.len
1703 || self.data_type != other.data_type
1704 || self.buffers.len() != other.buffers.len()
1705 || self.child_data.len() != other.child_data.len()
1706 {
1707 return false;
1708 }
1709
1710 match (&self.nulls, &other.nulls) {
1711 (Some(a), Some(b)) if !a.inner().ptr_eq(b.inner()) => return false,
1712 (Some(_), None) | (None, Some(_)) => return false,
1713 _ => {}
1714 };
1715
1716 if !self
1717 .buffers
1718 .iter()
1719 .zip(other.buffers.iter())
1720 .all(|(a, b)| a.as_ptr() == b.as_ptr())
1721 {
1722 return false;
1723 }
1724
1725 self.child_data
1726 .iter()
1727 .zip(other.child_data.iter())
1728 .all(|(a, b)| a.ptr_eq(b))
1729 }
1730
1731 pub fn into_builder(self) -> ArrayDataBuilder {
1733 self.into()
1734 }
1735
1736 #[cfg(feature = "pool")]
1743 pub fn claim(&self, pool: &dyn arrow_buffer::MemoryPool) {
1744 for buffer in &self.buffers {
1746 buffer.claim(pool);
1747 }
1748
1749 if let Some(nulls) = &self.nulls {
1751 nulls.claim(pool);
1752 }
1753
1754 for child in &self.child_data {
1756 child.claim(pool);
1757 }
1758 }
1759}
1760
1761pub fn layout(data_type: &DataType) -> DataTypeLayout {
1764 use arrow_schema::IntervalUnit::*;
1767
1768 match data_type {
1769 DataType::Null => DataTypeLayout {
1770 buffers: vec![],
1771 can_contain_null_mask: false,
1772 variadic: false,
1773 },
1774 DataType::Boolean => DataTypeLayout {
1775 buffers: vec![BufferSpec::BitMap],
1776 can_contain_null_mask: true,
1777 variadic: false,
1778 },
1779 DataType::Int8 => DataTypeLayout::new_fixed_width::<i8>(),
1780 DataType::Int16 => DataTypeLayout::new_fixed_width::<i16>(),
1781 DataType::Int32 => DataTypeLayout::new_fixed_width::<i32>(),
1782 DataType::Int64 => DataTypeLayout::new_fixed_width::<i64>(),
1783 DataType::UInt8 => DataTypeLayout::new_fixed_width::<u8>(),
1784 DataType::UInt16 => DataTypeLayout::new_fixed_width::<u16>(),
1785 DataType::UInt32 => DataTypeLayout::new_fixed_width::<u32>(),
1786 DataType::UInt64 => DataTypeLayout::new_fixed_width::<u64>(),
1787 DataType::Float16 => DataTypeLayout::new_fixed_width::<half::f16>(),
1788 DataType::Float32 => DataTypeLayout::new_fixed_width::<f32>(),
1789 DataType::Float64 => DataTypeLayout::new_fixed_width::<f64>(),
1790 DataType::Timestamp(_, _) => DataTypeLayout::new_fixed_width::<i64>(),
1791 DataType::Date32 => DataTypeLayout::new_fixed_width::<i32>(),
1792 DataType::Date64 => DataTypeLayout::new_fixed_width::<i64>(),
1793 DataType::Time32(_) => DataTypeLayout::new_fixed_width::<i32>(),
1794 DataType::Time64(_) => DataTypeLayout::new_fixed_width::<i64>(),
1795 DataType::Interval(YearMonth) => DataTypeLayout::new_fixed_width::<i32>(),
1796 DataType::Interval(DayTime) => DataTypeLayout::new_fixed_width::<IntervalDayTime>(),
1797 DataType::Interval(MonthDayNano) => {
1798 DataTypeLayout::new_fixed_width::<IntervalMonthDayNano>()
1799 }
1800 DataType::Duration(_) => DataTypeLayout::new_fixed_width::<i64>(),
1801 DataType::Decimal32(_, _) => DataTypeLayout::new_fixed_width::<i32>(),
1802 DataType::Decimal64(_, _) => DataTypeLayout::new_fixed_width::<i64>(),
1803 DataType::Decimal128(_, _) => DataTypeLayout::new_fixed_width::<i128>(),
1804 DataType::Decimal256(_, _) => DataTypeLayout::new_fixed_width::<i256>(),
1805 DataType::FixedSizeBinary(size) => {
1806 let spec = BufferSpec::FixedWidth {
1807 byte_width: (*size).try_into().unwrap(),
1808 alignment: mem::align_of::<u8>(),
1809 };
1810 DataTypeLayout {
1811 buffers: vec![spec],
1812 can_contain_null_mask: true,
1813 variadic: false,
1814 }
1815 }
1816 DataType::Binary => DataTypeLayout::new_binary::<i32>(),
1817 DataType::LargeBinary => DataTypeLayout::new_binary::<i64>(),
1818 DataType::Utf8 => DataTypeLayout::new_binary::<i32>(),
1819 DataType::LargeUtf8 => DataTypeLayout::new_binary::<i64>(),
1820 DataType::BinaryView | DataType::Utf8View => DataTypeLayout::new_view(),
1821 DataType::FixedSizeList(_, _) => DataTypeLayout::new_nullable_empty(), DataType::List(_) => DataTypeLayout::new_fixed_width::<i32>(),
1823 DataType::ListView(_) => DataTypeLayout::new_list_view::<i32>(),
1824 DataType::LargeListView(_) => DataTypeLayout::new_list_view::<i64>(),
1825 DataType::LargeList(_) => DataTypeLayout::new_fixed_width::<i64>(),
1826 DataType::Map(_, _) => DataTypeLayout::new_fixed_width::<i32>(),
1827 DataType::Struct(_) => DataTypeLayout::new_nullable_empty(), DataType::RunEndEncoded(_, _) => DataTypeLayout::new_empty(), DataType::Union(_, mode) => {
1830 let type_ids = BufferSpec::FixedWidth {
1831 byte_width: mem::size_of::<i8>(),
1832 alignment: mem::align_of::<i8>(),
1833 };
1834
1835 DataTypeLayout {
1836 buffers: match mode {
1837 UnionMode::Sparse => {
1838 vec![type_ids]
1839 }
1840 UnionMode::Dense => {
1841 vec![
1842 type_ids,
1843 BufferSpec::FixedWidth {
1844 byte_width: mem::size_of::<i32>(),
1845 alignment: mem::align_of::<i32>(),
1846 },
1847 ]
1848 }
1849 },
1850 can_contain_null_mask: false,
1851 variadic: false,
1852 }
1853 }
1854 DataType::Dictionary(key_type, _value_type) => layout(key_type),
1855 }
1856}
1857
1858#[derive(Debug, PartialEq, Eq)]
1860pub struct DataTypeLayout {
1862 pub buffers: Vec<BufferSpec>,
1864
1865 pub can_contain_null_mask: bool,
1867
1868 pub variadic: bool,
1872}
1873
1874impl DataTypeLayout {
1875 pub fn new_fixed_width<T>() -> Self {
1877 Self {
1878 buffers: vec![BufferSpec::FixedWidth {
1879 byte_width: mem::size_of::<T>(),
1880 alignment: mem::align_of::<T>(),
1881 }],
1882 can_contain_null_mask: true,
1883 variadic: false,
1884 }
1885 }
1886
1887 pub fn new_nullable_empty() -> Self {
1890 Self {
1891 buffers: vec![],
1892 can_contain_null_mask: true,
1893 variadic: false,
1894 }
1895 }
1896
1897 pub fn new_empty() -> Self {
1900 Self {
1901 buffers: vec![],
1902 can_contain_null_mask: false,
1903 variadic: false,
1904 }
1905 }
1906
1907 pub fn new_binary<T>() -> Self {
1911 Self {
1912 buffers: vec![
1913 BufferSpec::FixedWidth {
1915 byte_width: mem::size_of::<T>(),
1916 alignment: mem::align_of::<T>(),
1917 },
1918 BufferSpec::VariableWidth,
1920 ],
1921 can_contain_null_mask: true,
1922 variadic: false,
1923 }
1924 }
1925
1926 pub fn new_view() -> Self {
1928 Self {
1929 buffers: vec![BufferSpec::FixedWidth {
1930 byte_width: mem::size_of::<u128>(),
1931 alignment: mem::align_of::<u128>(),
1932 }],
1933 can_contain_null_mask: true,
1934 variadic: true,
1935 }
1936 }
1937
1938 pub fn new_list_view<T>() -> Self {
1940 Self {
1941 buffers: vec![
1942 BufferSpec::FixedWidth {
1943 byte_width: mem::size_of::<T>(),
1944 alignment: mem::align_of::<T>(),
1945 },
1946 BufferSpec::FixedWidth {
1947 byte_width: mem::size_of::<T>(),
1948 alignment: mem::align_of::<T>(),
1949 },
1950 ],
1951 can_contain_null_mask: true,
1952 variadic: false,
1953 }
1954 }
1955}
1956
1957#[derive(Debug, PartialEq, Eq)]
1959pub enum BufferSpec {
1960 FixedWidth {
1971 byte_width: usize,
1973 alignment: usize,
1975 },
1976 VariableWidth,
1978 BitMap,
1984 #[allow(dead_code)]
1987 AlwaysNull,
1988}
1989
1990impl PartialEq for ArrayData {
1991 fn eq(&self, other: &Self) -> bool {
1992 equal::equal(self, other)
1993 }
1994}
1995
1996#[derive(Debug, Clone)]
2015#[doc(hidden)]
2016pub struct UnsafeFlag(bool);
2017
2018impl UnsafeFlag {
2019 #[inline]
2023 pub const fn new() -> Self {
2024 Self(false)
2025 }
2026
2027 #[inline]
2037 pub unsafe fn set(&mut self, val: bool) {
2038 self.0 = val;
2039 }
2040
2041 #[inline]
2043 pub fn get(&self) -> bool {
2044 self.0
2045 }
2046}
2047
2048impl Default for UnsafeFlag {
2050 fn default() -> Self {
2051 Self::new()
2052 }
2053}
2054
2055#[derive(Debug)]
2057pub struct ArrayDataBuilder {
2058 data_type: DataType,
2059 len: usize,
2060 null_count: Option<usize>,
2061 null_bit_buffer: Option<Buffer>,
2062 nulls: Option<NullBuffer>,
2063 offset: usize,
2064 buffers: Vec<Buffer>,
2065 child_data: Vec<ArrayData>,
2066 align_buffers: bool,
2070 skip_validation: UnsafeFlag,
2080}
2081
2082impl ArrayDataBuilder {
2083 #[inline]
2084 pub const fn new(data_type: DataType) -> Self {
2086 Self {
2087 data_type,
2088 len: 0,
2089 null_count: None,
2090 null_bit_buffer: None,
2091 nulls: None,
2092 offset: 0,
2093 buffers: vec![],
2094 child_data: vec![],
2095 align_buffers: false,
2096 skip_validation: UnsafeFlag::new(),
2097 }
2098 }
2099
2100 pub fn data_type(self, data_type: DataType) -> Self {
2102 Self { data_type, ..self }
2103 }
2104
2105 #[inline]
2106 #[allow(clippy::len_without_is_empty)]
2107 pub const fn len(mut self, n: usize) -> Self {
2109 self.len = n;
2110 self
2111 }
2112
2113 pub fn nulls(mut self, nulls: Option<NullBuffer>) -> Self {
2115 self.nulls = nulls;
2116 self.null_count = None;
2117 self.null_bit_buffer = None;
2118 self
2119 }
2120
2121 pub fn null_count(mut self, null_count: usize) -> Self {
2123 self.null_count = Some(null_count);
2124 self
2125 }
2126
2127 pub fn null_bit_buffer(mut self, buf: Option<Buffer>) -> Self {
2129 self.nulls = None;
2130 self.null_bit_buffer = buf;
2131 self
2132 }
2133
2134 #[inline]
2136 pub const fn offset(mut self, n: usize) -> Self {
2137 self.offset = n;
2138 self
2139 }
2140
2141 pub fn buffers(mut self, v: Vec<Buffer>) -> Self {
2143 self.buffers = v;
2144 self
2145 }
2146
2147 pub fn add_buffer(mut self, b: Buffer) -> Self {
2149 self.buffers.push(b);
2150 self
2151 }
2152
2153 pub fn add_buffers<I: IntoIterator<Item = Buffer>>(mut self, bs: I) -> Self {
2155 self.buffers.extend(bs);
2156 self
2157 }
2158
2159 pub fn child_data(mut self, v: Vec<ArrayData>) -> Self {
2161 self.child_data = v;
2162 self
2163 }
2164
2165 pub fn add_child_data(mut self, r: ArrayData) -> Self {
2167 self.child_data.push(r);
2168 self
2169 }
2170
2171 pub unsafe fn build_unchecked(self) -> ArrayData {
2187 unsafe { self.skip_validation(true) }.build().unwrap()
2188 }
2189
2190 pub fn build(self) -> Result<ArrayData, ArrowError> {
2199 let Self {
2200 data_type,
2201 len,
2202 null_count,
2203 null_bit_buffer,
2204 nulls,
2205 offset,
2206 buffers,
2207 child_data,
2208 align_buffers,
2209 skip_validation,
2210 } = self;
2211
2212 let nulls = nulls
2213 .or_else(|| {
2214 let buffer = null_bit_buffer?;
2215 let buffer = BooleanBuffer::new(buffer, offset, len);
2216 Some(match null_count {
2217 Some(n) => {
2218 unsafe { NullBuffer::new_unchecked(buffer, n) }
2220 }
2221 None => NullBuffer::new(buffer),
2222 })
2223 })
2224 .filter(|b| b.null_count() != 0);
2225
2226 let mut data = ArrayData {
2227 data_type,
2228 len,
2229 offset,
2230 buffers,
2231 child_data,
2232 nulls,
2233 };
2234
2235 if align_buffers {
2236 data.align_buffers();
2237 }
2238
2239 if !skip_validation.get() || cfg!(feature = "force_validate") {
2241 data.validate_data()?;
2242 }
2243 Ok(data)
2244 }
2245
2246 #[deprecated(since = "54.1.0", note = "Use ArrayData::align_buffers instead")]
2248 pub fn build_aligned(self) -> Result<ArrayData, ArrowError> {
2249 self.align_buffers(true).build()
2250 }
2251
2252 pub fn align_buffers(mut self, align_buffers: bool) -> Self {
2268 self.align_buffers = align_buffers;
2269 self
2270 }
2271
2272 pub unsafe fn skip_validation(mut self, skip_validation: bool) -> Self {
2286 unsafe {
2287 self.skip_validation.set(skip_validation);
2288 }
2289 self
2290 }
2291}
2292
2293impl From<ArrayData> for ArrayDataBuilder {
2294 fn from(d: ArrayData) -> Self {
2295 Self {
2296 data_type: d.data_type,
2297 len: d.len,
2298 offset: d.offset,
2299 buffers: d.buffers,
2300 child_data: d.child_data,
2301 nulls: d.nulls,
2302 null_bit_buffer: None,
2303 null_count: None,
2304 align_buffers: false,
2305 skip_validation: UnsafeFlag::new(),
2306 }
2307 }
2308}
2309
2310#[cfg(test)]
2311mod tests {
2312 use crate::transform::MutableArrayData;
2313
2314 use super::*;
2315 use arrow_schema::{Field, Fields, UnionFields};
2316
2317 fn make_i32_buffer(n: usize) -> Buffer {
2321 Buffer::from_slice_ref(vec![42i32; n])
2322 }
2323
2324 fn make_f32_buffer(n: usize) -> Buffer {
2326 Buffer::from_slice_ref(vec![42f32; n])
2327 }
2328
2329 #[test]
2330 fn test_builder() {
2331 let v = (0..25).collect::<Vec<i32>>();
2333 let b1 = Buffer::from_slice_ref(&v);
2334 let arr_data = ArrayData::builder(DataType::Int32)
2335 .len(20)
2336 .offset(5)
2337 .add_buffer(b1)
2338 .null_bit_buffer(Some(Buffer::from([
2339 0b01011111, 0b10110101, 0b01100011, 0b00011110,
2340 ])))
2341 .build()
2342 .unwrap();
2343
2344 assert_eq!(20, arr_data.len());
2345 assert_eq!(10, arr_data.null_count());
2346 assert_eq!(5, arr_data.offset());
2347 assert_eq!(1, arr_data.buffers().len());
2348 assert_eq!(
2349 Buffer::from_slice_ref(&v).as_slice(),
2350 arr_data.buffers()[0].as_slice()
2351 );
2352 }
2353
2354 #[test]
2355 fn test_builder_with_child_data() {
2356 let child_arr_data = ArrayData::try_new(
2357 DataType::Int32,
2358 5,
2359 None,
2360 0,
2361 vec![Buffer::from_slice_ref([1i32, 2, 3, 4, 5])],
2362 vec![],
2363 )
2364 .unwrap();
2365
2366 let field = Arc::new(Field::new("x", DataType::Int32, true));
2367 let data_type = DataType::Struct(vec![field].into());
2368
2369 let arr_data = ArrayData::builder(data_type)
2370 .len(5)
2371 .offset(0)
2372 .add_child_data(child_arr_data.clone())
2373 .build()
2374 .unwrap();
2375
2376 assert_eq!(5, arr_data.len());
2377 assert_eq!(1, arr_data.child_data().len());
2378 assert_eq!(child_arr_data, arr_data.child_data()[0]);
2379 }
2380
2381 #[test]
2382 fn test_extend_fixed_size_list_accounts_for_parent_offset() {
2383 let field = Arc::new(Field::new("item", DataType::Int32, false));
2384 let child = ArrayData::builder(DataType::Int32)
2385 .len(6)
2386 .add_buffer(Buffer::from_slice_ref([0, 1, 2, 3, 4, 5]))
2387 .build()
2388 .unwrap();
2389
2390 let data = ArrayData::builder(DataType::FixedSizeList(field, 2))
2391 .len(2)
2392 .offset(1)
2393 .add_child_data(child)
2394 .build()
2395 .unwrap();
2396
2397 let mut mutable = MutableArrayData::new(vec![&data], false, data.len());
2398 mutable.extend(0, 0, data.len());
2399 let output = mutable.freeze();
2400
2401 assert_eq!(output.child_data()[0].buffer::<i32>(0), &[2, 3, 4, 5]);
2402 }
2403
2404 #[test]
2405 fn test_null_count() {
2406 let mut bit_v: [u8; 2] = [0; 2];
2407 bit_util::set_bit(&mut bit_v, 0);
2408 bit_util::set_bit(&mut bit_v, 3);
2409 bit_util::set_bit(&mut bit_v, 10);
2410 let arr_data = ArrayData::builder(DataType::Int32)
2411 .len(16)
2412 .add_buffer(make_i32_buffer(16))
2413 .null_bit_buffer(Some(Buffer::from(bit_v)))
2414 .build()
2415 .unwrap();
2416 assert_eq!(13, arr_data.null_count());
2417
2418 let mut bit_v: [u8; 2] = [0; 2];
2420 bit_util::set_bit(&mut bit_v, 0);
2421 bit_util::set_bit(&mut bit_v, 3);
2422 bit_util::set_bit(&mut bit_v, 10);
2423 let arr_data = ArrayData::builder(DataType::Int32)
2424 .len(12)
2425 .offset(2)
2426 .add_buffer(make_i32_buffer(14)) .null_bit_buffer(Some(Buffer::from(bit_v)))
2428 .build()
2429 .unwrap();
2430 assert_eq!(10, arr_data.null_count());
2431 }
2432
2433 #[test]
2434 fn test_null_buffer_ref() {
2435 let mut bit_v: [u8; 2] = [0; 2];
2436 bit_util::set_bit(&mut bit_v, 0);
2437 bit_util::set_bit(&mut bit_v, 3);
2438 bit_util::set_bit(&mut bit_v, 10);
2439 let arr_data = ArrayData::builder(DataType::Int32)
2440 .len(16)
2441 .add_buffer(make_i32_buffer(16))
2442 .null_bit_buffer(Some(Buffer::from(bit_v)))
2443 .build()
2444 .unwrap();
2445 assert!(arr_data.nulls().is_some());
2446 assert_eq!(&bit_v, arr_data.nulls().unwrap().validity());
2447 }
2448
2449 #[test]
2450 fn test_slice() {
2451 let mut bit_v: [u8; 2] = [0; 2];
2452 bit_util::set_bit(&mut bit_v, 0);
2453 bit_util::set_bit(&mut bit_v, 3);
2454 bit_util::set_bit(&mut bit_v, 10);
2455 let data = ArrayData::builder(DataType::Int32)
2456 .len(16)
2457 .add_buffer(make_i32_buffer(16))
2458 .null_bit_buffer(Some(Buffer::from(bit_v)))
2459 .build()
2460 .unwrap();
2461 let new_data = data.slice(1, 15);
2462 assert_eq!(data.len() - 1, new_data.len());
2463 assert_eq!(1, new_data.offset());
2464 assert_eq!(data.null_count(), new_data.null_count());
2465
2466 let new_data = new_data.slice(1, 14);
2468 assert_eq!(data.len() - 2, new_data.len());
2469 assert_eq!(2, new_data.offset());
2470 assert_eq!(data.null_count() - 1, new_data.null_count());
2471 }
2472
2473 #[test]
2474 #[should_panic(expected = "offset + length overflow")]
2475 fn test_slice_panics_on_offset_length_overflow() {
2476 let data = ArrayData::builder(DataType::Int32)
2477 .len(4)
2478 .add_buffer(make_i32_buffer(4))
2479 .build()
2480 .unwrap();
2481 let sliced = data.slice(1, 3);
2482
2483 sliced.slice(1, usize::MAX);
2484 }
2485
2486 #[test]
2487 fn test_typed_offsets_length_overflow() {
2488 let data = ArrayData {
2489 data_type: DataType::Binary,
2490 len: usize::MAX,
2491 offset: 0,
2492 buffers: vec![Buffer::from_slice_ref([0_i32])],
2493 child_data: vec![],
2494 nulls: None,
2495 };
2496 let err = data.typed_offsets::<i32>().unwrap_err();
2497
2498 assert_eq!(
2499 err.to_string(),
2500 format!(
2501 "Invalid argument error: Length {} with offset 1 overflows usize for Binary",
2502 usize::MAX
2503 )
2504 );
2505 }
2506
2507 #[test]
2508 fn test_validate_typed_buffer_length_overflow() {
2509 let data = ArrayData {
2510 data_type: DataType::Binary,
2511 len: 0,
2512 offset: 2,
2513 buffers: vec![Buffer::from_slice_ref([0_i32])],
2514 child_data: vec![],
2515 nulls: None,
2516 };
2517 let err = data.typed_buffer::<i32>(0, usize::MAX).unwrap_err();
2518
2519 assert_eq!(
2520 err.to_string(),
2521 format!(
2522 "Invalid argument error: Length {} with offset 2 overflows usize for Binary",
2523 usize::MAX
2524 )
2525 );
2526 }
2527
2528 fn try_new_binary_length_offset_overflow() -> Result<ArrayData, ArrowError> {
2530 ArrayData::try_new(
2531 DataType::Binary,
2532 usize::MAX,
2533 None,
2534 1,
2535 vec![
2536 Buffer::from_slice_ref([0_i32]),
2537 Buffer::from_iter(std::iter::empty::<u8>()),
2538 ],
2539 vec![],
2540 )
2541 }
2542
2543 #[cfg(not(feature = "force_validate"))]
2544 #[test]
2545 fn test_try_new_length_offset_overflow() {
2546 let err = try_new_binary_length_offset_overflow().unwrap_err();
2547
2548 assert_eq!(
2549 err.to_string(),
2550 format!(
2551 "Invalid argument error: Length {} with offset 1 overflows usize for Binary",
2552 usize::MAX
2553 )
2554 );
2555 }
2556
2557 #[cfg(feature = "force_validate")]
2558 #[test]
2559 #[should_panic(
2560 expected = "Length 18446744073709551615 with offset 1 overflows usize for Binary"
2561 )]
2562 fn test_try_new_length_offset_overflow_force_validate() {
2563 try_new_binary_length_offset_overflow().unwrap();
2564 }
2565
2566 #[test]
2567 fn test_equality() {
2568 let int_data = ArrayData::builder(DataType::Int32)
2569 .len(1)
2570 .add_buffer(make_i32_buffer(1))
2571 .build()
2572 .unwrap();
2573
2574 let float_data = ArrayData::builder(DataType::Float32)
2575 .len(1)
2576 .add_buffer(make_f32_buffer(1))
2577 .build()
2578 .unwrap();
2579 assert_ne!(int_data, float_data);
2580 assert!(!int_data.ptr_eq(&float_data));
2581 assert!(int_data.ptr_eq(&int_data));
2582
2583 #[allow(clippy::redundant_clone)]
2584 let int_data_clone = int_data.clone();
2585 assert_eq!(int_data, int_data_clone);
2586 assert!(int_data.ptr_eq(&int_data_clone));
2587 assert!(int_data_clone.ptr_eq(&int_data));
2588
2589 let int_data_slice = int_data_clone.slice(1, 0);
2590 assert!(int_data_slice.ptr_eq(&int_data_slice));
2591 assert!(!int_data.ptr_eq(&int_data_slice));
2592 assert!(!int_data_slice.ptr_eq(&int_data));
2593
2594 let data_buffer = Buffer::from_slice_ref("abcdef".as_bytes());
2595 let offsets_buffer = Buffer::from_slice_ref([0_i32, 2_i32, 2_i32, 5_i32]);
2596 let string_data = ArrayData::try_new(
2597 DataType::Utf8,
2598 3,
2599 Some(Buffer::from_iter(vec![true, false, true])),
2600 0,
2601 vec![offsets_buffer, data_buffer],
2602 vec![],
2603 )
2604 .unwrap();
2605
2606 assert_ne!(float_data, string_data);
2607 assert!(!float_data.ptr_eq(&string_data));
2608
2609 assert!(string_data.ptr_eq(&string_data));
2610
2611 #[allow(clippy::redundant_clone)]
2612 let string_data_cloned = string_data.clone();
2613 assert!(string_data_cloned.ptr_eq(&string_data));
2614 assert!(string_data.ptr_eq(&string_data_cloned));
2615
2616 let string_data_slice = string_data.slice(1, 2);
2617 assert!(string_data_slice.ptr_eq(&string_data_slice));
2618 assert!(!string_data_slice.ptr_eq(&string_data))
2619 }
2620
2621 #[test]
2622 fn test_slice_memory_size() {
2623 let mut bit_v: [u8; 2] = [0; 2];
2624 bit_util::set_bit(&mut bit_v, 0);
2625 bit_util::set_bit(&mut bit_v, 3);
2626 bit_util::set_bit(&mut bit_v, 10);
2627 let data = ArrayData::builder(DataType::Int32)
2628 .len(16)
2629 .add_buffer(make_i32_buffer(16))
2630 .null_bit_buffer(Some(Buffer::from(bit_v)))
2631 .build()
2632 .unwrap();
2633 let new_data = data.slice(1, 14);
2634 assert_eq!(
2635 data.get_slice_memory_size().unwrap() - 8,
2636 new_data.get_slice_memory_size().unwrap()
2637 );
2638 let data_buffer = Buffer::from_slice_ref("abcdef".as_bytes());
2639 let offsets_buffer = Buffer::from_slice_ref([0_i32, 2_i32, 2_i32, 5_i32]);
2640 let string_data = ArrayData::try_new(
2641 DataType::Utf8,
2642 3,
2643 Some(Buffer::from_iter(vec![true, false, true])),
2644 0,
2645 vec![offsets_buffer, data_buffer],
2646 vec![],
2647 )
2648 .unwrap();
2649 let string_data_slice = string_data.slice(1, 2);
2650 assert_eq!(
2652 string_data.get_slice_memory_size().unwrap() - 6,
2653 string_data_slice.get_slice_memory_size().unwrap()
2654 );
2655 }
2656
2657 #[test]
2658 fn test_count_nulls() {
2659 let buffer = Buffer::from([0b00010110, 0b10011111]);
2660 let buffer = NullBuffer::new(BooleanBuffer::new(buffer, 0, 16));
2661 let count = count_nulls(Some(&buffer), 0, 16);
2662 assert_eq!(count, 7);
2663
2664 let count = count_nulls(Some(&buffer), 4, 8);
2665 assert_eq!(count, 3);
2666 }
2667
2668 #[test]
2669 fn test_contains_nulls() {
2670 let buffer: Buffer =
2671 MutableBuffer::from_iter([false, false, false, true, true, false]).into();
2672 let buffer = NullBuffer::new(BooleanBuffer::new(buffer, 0, 6));
2673 assert!(contains_nulls(Some(&buffer), 0, 6));
2674 assert!(contains_nulls(Some(&buffer), 0, 3));
2675 assert!(!contains_nulls(Some(&buffer), 3, 2));
2676 assert!(!contains_nulls(Some(&buffer), 0, 0));
2677 }
2678
2679 #[test]
2680 fn test_alignment() {
2681 let buffer = Buffer::from_vec(vec![1_i32, 2_i32, 3_i32]);
2682 let sliced = buffer.slice(1);
2683
2684 let mut data = ArrayData {
2685 data_type: DataType::Int32,
2686 len: 0,
2687 offset: 0,
2688 buffers: vec![buffer],
2689 child_data: vec![],
2690 nulls: None,
2691 };
2692 data.validate_full().unwrap();
2693
2694 data.buffers[0] = sliced;
2696 let err = data.validate().unwrap_err();
2697
2698 assert_eq!(
2699 err.to_string(),
2700 "Invalid argument error: Misaligned buffers[0] in array of type Int32, offset from expected alignment of 4 by 1"
2701 );
2702
2703 data.align_buffers();
2704 data.validate_full().unwrap();
2705 }
2706
2707 #[test]
2708 fn test_alignment_struct() {
2709 let buffer = Buffer::from_vec(vec![1_i32, 2_i32, 3_i32]);
2710 let sliced = buffer.slice(1);
2711
2712 let child_data = ArrayData {
2713 data_type: DataType::Int32,
2714 len: 0,
2715 offset: 0,
2716 buffers: vec![buffer],
2717 child_data: vec![],
2718 nulls: None,
2719 };
2720
2721 let schema = DataType::Struct(Fields::from(vec![Field::new("a", DataType::Int32, false)]));
2722 let mut data = ArrayData {
2723 data_type: schema,
2724 len: 0,
2725 offset: 0,
2726 buffers: vec![],
2727 child_data: vec![child_data],
2728 nulls: None,
2729 };
2730 data.validate_full().unwrap();
2731
2732 data.child_data[0].buffers[0] = sliced;
2734 let err = data.validate().unwrap_err();
2735
2736 assert_eq!(
2737 err.to_string(),
2738 "Invalid argument error: Misaligned buffers[0] in array of type Int32, offset from expected alignment of 4 by 1"
2739 );
2740
2741 data.align_buffers();
2742 data.validate_full().unwrap();
2743 }
2744
2745 #[test]
2746 fn test_null_view_types() {
2747 let array_len = 32;
2748 let array = ArrayData::new_null(&DataType::BinaryView, array_len);
2749 assert_eq!(array.len(), array_len);
2750 for i in 0..array.len() {
2751 assert!(array.is_null(i));
2752 }
2753
2754 let array = ArrayData::new_null(&DataType::Utf8View, array_len);
2755 assert_eq!(array.len(), array_len);
2756 for i in 0..array.len() {
2757 assert!(array.is_null(i));
2758 }
2759
2760 let array = ArrayData::new_null(
2761 &DataType::ListView(Arc::new(Field::new_list_field(DataType::Int32, true))),
2762 array_len,
2763 );
2764 assert_eq!(array.len(), array_len);
2765 for i in 0..array.len() {
2766 assert!(array.is_null(i));
2767 }
2768
2769 let array = ArrayData::new_null(
2770 &DataType::LargeListView(Arc::new(Field::new_list_field(DataType::Int32, true))),
2771 array_len,
2772 );
2773 assert_eq!(array.len(), array_len);
2774 for i in 0..array.len() {
2775 assert!(array.is_null(i));
2776 }
2777 }
2778
2779 #[test]
2780 fn test_union_validate_type_ids() {
2781 let dt = DataType::Union(
2782 UnionFields::try_new(
2783 vec![5],
2784 vec![Arc::new(Field::new("a", DataType::Int32, true))],
2785 )
2786 .unwrap(),
2787 UnionMode::Sparse,
2788 );
2789
2790 let child = ArrayData::builder(DataType::Int32)
2791 .len(2)
2792 .add_buffer(make_i32_buffer(2))
2793 .build()
2794 .unwrap();
2795
2796 let err = ArrayDataBuilder::new(dt)
2797 .len(2)
2798 .offset(0)
2799 .buffers(vec![vec![5_i8, 2].into()])
2800 .child_data(vec![child])
2801 .build()
2802 .unwrap_err();
2803 assert_eq!(
2804 err.to_string(),
2805 "Invalid argument error: Type id values must match one of the field type ids, got 2 at index 1"
2806 );
2807 }
2808
2809 #[test]
2810 fn test_union_validate_dense_offsets() {
2811 let dt = DataType::Union(
2812 UnionFields::try_new(
2813 vec![5],
2814 vec![Arc::new(Field::new("a", DataType::Int32, true))],
2815 )
2816 .unwrap(),
2817 UnionMode::Dense,
2818 );
2819
2820 let child = ArrayData::builder(DataType::Int32)
2821 .len(2)
2822 .add_buffer(make_i32_buffer(2))
2823 .build()
2824 .unwrap();
2825
2826 let err = ArrayDataBuilder::new(dt.clone())
2827 .len(2)
2828 .buffers(vec![vec![5_i8, 5].into(), vec![0_i32, 10].into()])
2830 .child_data(vec![child.clone()])
2831 .build()
2832 .unwrap_err();
2833 assert_eq!(
2834 err.to_string(),
2835 "Invalid argument error: Offsets must be non-negative and within the length of the child array"
2836 );
2837
2838 let err = ArrayDataBuilder::new(dt)
2839 .len(2)
2840 .buffers(vec![vec![5_i8, 5].into(), vec![0_i32, -1].into()])
2842 .child_data(vec![child])
2843 .build()
2844 .unwrap_err();
2845 assert_eq!(
2846 err.to_string(),
2847 "Invalid argument error: Offsets must be non-negative and within the length of the child array"
2848 );
2849 }
2850
2851 fn run_array_test_children() -> (DataType, Vec<ArrayData>) {
2854 let dt = DataType::RunEndEncoded(
2855 Arc::new(Field::new("run_ends", DataType::Int32, false)),
2856 Arc::new(Field::new("values", DataType::Float32, true)),
2857 );
2858
2859 let runs = ArrayData::builder(DataType::Int32)
2860 .len(2)
2861 .add_buffer(vec![1_i32, 3].into())
2862 .build()
2863 .unwrap();
2864 let values = ArrayData::builder(DataType::Float32)
2865 .len(2)
2866 .add_buffer(vec![1.5_f32, 2.5].into())
2867 .build()
2868 .unwrap();
2869
2870 (dt, vec![runs, values])
2871 }
2872
2873 #[test]
2874 fn test_run_array_validates_last_run() {
2875 let (dt, children) = run_array_test_children();
2876
2877 let err = ArrayDataBuilder::new(dt)
2878 .len(4)
2879 .offset(0)
2880 .child_data(children)
2881 .build()
2882 .unwrap_err();
2883 assert_eq!(
2884 err.to_string(),
2885 "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."
2886 );
2887 }
2888
2889 #[test]
2890 fn test_run_array_empty_validates() {
2891 let (dt, children) = run_array_test_children();
2892
2893 let data = ArrayDataBuilder::new(dt)
2894 .len(0)
2895 .offset(4)
2896 .child_data(children)
2897 .build()
2898 .unwrap();
2899 assert_eq!(data.len(), 0);
2900 }
2901
2902 #[test]
2905 fn test_run_array_validates_last_run_with_offset() {
2906 let (dt, children) = run_array_test_children();
2907
2908 let err = ArrayDataBuilder::new(dt)
2909 .len(2)
2910 .offset(2)
2911 .child_data(children)
2912 .build()
2913 .unwrap_err();
2914 assert_eq!(
2915 err.to_string(),
2916 "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."
2917 );
2918 }
2919
2920 #[test]
2923 fn test_run_array_sliced_validates() {
2924 let (dt, children) = run_array_test_children();
2925
2926 let data = ArrayDataBuilder::new(dt)
2927 .len(2)
2928 .offset(1)
2929 .child_data(children)
2930 .build()
2931 .unwrap();
2932 assert_eq!(data.len(), 2);
2933 assert_eq!(data.offset(), 1);
2934 }
2935}