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