1mod binary_array;
21
22use crate::types::*;
23use arrow_buffer::{ArrowNativeType, NullBuffer, OffsetBuffer, ScalarBuffer};
24use arrow_data::ArrayData;
25use arrow_schema::{DataType, IntervalUnit, TimeUnit};
26use std::any::Any;
27use std::sync::Arc;
28
29pub use binary_array::*;
30
31mod boolean_array;
32pub use boolean_array::*;
33
34mod byte_array;
35pub use byte_array::*;
36
37mod dictionary_array;
38pub use dictionary_array::*;
39
40mod fixed_size_binary_array;
41pub use fixed_size_binary_array::*;
42
43mod fixed_size_list_array;
44pub use fixed_size_list_array::*;
45
46mod list_array;
47pub use list_array::*;
48
49mod map_array;
50pub use map_array::*;
51
52mod null_array;
53pub use null_array::*;
54
55mod primitive_array;
56pub use primitive_array::*;
57
58mod string_array;
59pub use string_array::*;
60
61mod struct_array;
62pub use struct_array::*;
63
64mod union_array;
65pub use union_array::*;
66
67mod run_array;
68
69pub use run_array::*;
70
71mod byte_view_array;
72
73pub use byte_view_array::*;
74
75mod list_view_array;
76
77pub use list_view_array::*;
78
79use crate::iterator::ArrayIter;
80
81pub trait Array: std::fmt::Debug + Send + Sync {
83 fn as_any(&self) -> &dyn Any;
106
107 fn to_data(&self) -> ArrayData;
109
110 fn into_data(self) -> ArrayData;
114
115 fn data_type(&self) -> &DataType;
128
129 fn slice(&self, offset: usize, length: usize) -> ArrayRef;
143
144 fn len(&self) -> usize;
156
157 fn is_empty(&self) -> bool;
169
170 fn shrink_to_fit(&mut self) {}
175
176 fn offset(&self) -> usize;
192
193 fn nulls(&self) -> Option<&NullBuffer>;
206
207 fn logical_nulls(&self) -> Option<NullBuffer> {
226 self.nulls().cloned()
227 }
228
229 fn is_null(&self, index: usize) -> bool {
251 self.nulls().map(|n| n.is_null(index)).unwrap_or_default()
252 }
253
254 fn is_valid(&self, index: usize) -> bool {
268 !self.is_null(index)
269 }
270
271 fn null_count(&self) -> usize {
287 self.nulls().map(|n| n.null_count()).unwrap_or_default()
288 }
289
290 fn logical_null_count(&self) -> usize {
307 self.logical_nulls()
308 .map(|n| n.null_count())
309 .unwrap_or_default()
310 }
311
312 fn is_nullable(&self) -> bool {
326 self.logical_null_count() != 0
327 }
328
329 fn get_buffer_memory_size(&self) -> usize;
334
335 fn get_array_memory_size(&self) -> usize;
339}
340
341pub type ArrayRef = Arc<dyn Array>;
343
344impl Array for ArrayRef {
346 fn as_any(&self) -> &dyn Any {
347 self.as_ref().as_any()
348 }
349
350 fn to_data(&self) -> ArrayData {
351 self.as_ref().to_data()
352 }
353
354 fn into_data(self) -> ArrayData {
355 self.to_data()
356 }
357
358 fn data_type(&self) -> &DataType {
359 self.as_ref().data_type()
360 }
361
362 fn slice(&self, offset: usize, length: usize) -> ArrayRef {
363 self.as_ref().slice(offset, length)
364 }
365
366 fn len(&self) -> usize {
367 self.as_ref().len()
368 }
369
370 fn is_empty(&self) -> bool {
371 self.as_ref().is_empty()
372 }
373
374 fn shrink_to_fit(&mut self) {
376 if let Some(slf) = Arc::get_mut(self) {
377 slf.shrink_to_fit();
378 } else {
379 }
381 }
382
383 fn offset(&self) -> usize {
384 self.as_ref().offset()
385 }
386
387 fn nulls(&self) -> Option<&NullBuffer> {
388 self.as_ref().nulls()
389 }
390
391 fn logical_nulls(&self) -> Option<NullBuffer> {
392 self.as_ref().logical_nulls()
393 }
394
395 fn is_null(&self, index: usize) -> bool {
396 self.as_ref().is_null(index)
397 }
398
399 fn is_valid(&self, index: usize) -> bool {
400 self.as_ref().is_valid(index)
401 }
402
403 fn null_count(&self) -> usize {
404 self.as_ref().null_count()
405 }
406
407 fn logical_null_count(&self) -> usize {
408 self.as_ref().logical_null_count()
409 }
410
411 fn is_nullable(&self) -> bool {
412 self.as_ref().is_nullable()
413 }
414
415 fn get_buffer_memory_size(&self) -> usize {
416 self.as_ref().get_buffer_memory_size()
417 }
418
419 fn get_array_memory_size(&self) -> usize {
420 self.as_ref().get_array_memory_size()
421 }
422}
423
424impl<T: Array> Array for &T {
425 fn as_any(&self) -> &dyn Any {
426 T::as_any(self)
427 }
428
429 fn to_data(&self) -> ArrayData {
430 T::to_data(self)
431 }
432
433 fn into_data(self) -> ArrayData {
434 self.to_data()
435 }
436
437 fn data_type(&self) -> &DataType {
438 T::data_type(self)
439 }
440
441 fn slice(&self, offset: usize, length: usize) -> ArrayRef {
442 T::slice(self, offset, length)
443 }
444
445 fn len(&self) -> usize {
446 T::len(self)
447 }
448
449 fn is_empty(&self) -> bool {
450 T::is_empty(self)
451 }
452
453 fn offset(&self) -> usize {
454 T::offset(self)
455 }
456
457 fn nulls(&self) -> Option<&NullBuffer> {
458 T::nulls(self)
459 }
460
461 fn logical_nulls(&self) -> Option<NullBuffer> {
462 T::logical_nulls(self)
463 }
464
465 fn is_null(&self, index: usize) -> bool {
466 T::is_null(self, index)
467 }
468
469 fn is_valid(&self, index: usize) -> bool {
470 T::is_valid(self, index)
471 }
472
473 fn null_count(&self) -> usize {
474 T::null_count(self)
475 }
476
477 fn logical_null_count(&self) -> usize {
478 T::logical_null_count(self)
479 }
480
481 fn is_nullable(&self) -> bool {
482 T::is_nullable(self)
483 }
484
485 fn get_buffer_memory_size(&self) -> usize {
486 T::get_buffer_memory_size(self)
487 }
488
489 fn get_array_memory_size(&self) -> usize {
490 T::get_array_memory_size(self)
491 }
492}
493
494pub trait ArrayAccessor: Array {
575 type Item: Send + Sync;
577
578 fn value(&self, index: usize) -> Self::Item;
582
583 unsafe fn value_unchecked(&self, index: usize) -> Self::Item;
587}
588
589pub trait StringArrayType<'a>: ArrayAccessor<Item = &'a str> + Sized {
597 fn is_ascii(&self) -> bool;
599
600 fn iter(&self) -> ArrayIter<Self>;
602}
603
604impl<'a, O: OffsetSizeTrait> StringArrayType<'a> for &'a GenericStringArray<O> {
605 fn is_ascii(&self) -> bool {
606 GenericStringArray::<O>::is_ascii(self)
607 }
608
609 fn iter(&self) -> ArrayIter<Self> {
610 GenericStringArray::<O>::iter(self)
611 }
612}
613impl<'a> StringArrayType<'a> for &'a StringViewArray {
614 fn is_ascii(&self) -> bool {
615 StringViewArray::is_ascii(self)
616 }
617
618 fn iter(&self) -> ArrayIter<Self> {
619 StringViewArray::iter(self)
620 }
621}
622
623pub trait BinaryArrayType<'a>: ArrayAccessor<Item = &'a [u8]> + Sized {
631 fn iter(&self) -> ArrayIter<Self>;
633}
634
635impl<'a, O: OffsetSizeTrait> BinaryArrayType<'a> for &'a GenericBinaryArray<O> {
636 fn iter(&self) -> ArrayIter<Self> {
637 GenericBinaryArray::<O>::iter(self)
638 }
639}
640impl<'a> BinaryArrayType<'a> for &'a BinaryViewArray {
641 fn iter(&self) -> ArrayIter<Self> {
642 BinaryViewArray::iter(self)
643 }
644}
645
646impl PartialEq for dyn Array + '_ {
647 fn eq(&self, other: &Self) -> bool {
648 self.to_data().eq(&other.to_data())
649 }
650}
651
652impl<T: Array> PartialEq<T> for dyn Array + '_ {
653 fn eq(&self, other: &T) -> bool {
654 self.to_data().eq(&other.to_data())
655 }
656}
657
658impl PartialEq for NullArray {
659 fn eq(&self, other: &NullArray) -> bool {
660 self.to_data().eq(&other.to_data())
661 }
662}
663
664impl<T: ArrowPrimitiveType> PartialEq for PrimitiveArray<T> {
665 fn eq(&self, other: &PrimitiveArray<T>) -> bool {
666 self.to_data().eq(&other.to_data())
667 }
668}
669
670impl<K: ArrowDictionaryKeyType> PartialEq for DictionaryArray<K> {
671 fn eq(&self, other: &Self) -> bool {
672 self.to_data().eq(&other.to_data())
673 }
674}
675
676impl PartialEq for BooleanArray {
677 fn eq(&self, other: &BooleanArray) -> bool {
678 self.to_data().eq(&other.to_data())
679 }
680}
681
682impl<OffsetSize: OffsetSizeTrait> PartialEq for GenericStringArray<OffsetSize> {
683 fn eq(&self, other: &Self) -> bool {
684 self.to_data().eq(&other.to_data())
685 }
686}
687
688impl<OffsetSize: OffsetSizeTrait> PartialEq for GenericBinaryArray<OffsetSize> {
689 fn eq(&self, other: &Self) -> bool {
690 self.to_data().eq(&other.to_data())
691 }
692}
693
694impl PartialEq for FixedSizeBinaryArray {
695 fn eq(&self, other: &Self) -> bool {
696 self.to_data().eq(&other.to_data())
697 }
698}
699
700impl<OffsetSize: OffsetSizeTrait> PartialEq for GenericListArray<OffsetSize> {
701 fn eq(&self, other: &Self) -> bool {
702 self.to_data().eq(&other.to_data())
703 }
704}
705
706impl<OffsetSize: OffsetSizeTrait> PartialEq for GenericListViewArray<OffsetSize> {
707 fn eq(&self, other: &Self) -> bool {
708 self.to_data().eq(&other.to_data())
709 }
710}
711
712impl PartialEq for MapArray {
713 fn eq(&self, other: &Self) -> bool {
714 self.to_data().eq(&other.to_data())
715 }
716}
717
718impl PartialEq for FixedSizeListArray {
719 fn eq(&self, other: &Self) -> bool {
720 self.to_data().eq(&other.to_data())
721 }
722}
723
724impl PartialEq for StructArray {
725 fn eq(&self, other: &Self) -> bool {
726 self.to_data().eq(&other.to_data())
727 }
728}
729
730impl<T: ByteViewType + ?Sized> PartialEq for GenericByteViewArray<T> {
731 fn eq(&self, other: &Self) -> bool {
732 self.to_data().eq(&other.to_data())
733 }
734}
735
736pub fn make_array(data: ArrayData) -> ArrayRef {
739 match data.data_type() {
740 DataType::Boolean => Arc::new(BooleanArray::from(data)) as ArrayRef,
741 DataType::Int8 => Arc::new(Int8Array::from(data)) as ArrayRef,
742 DataType::Int16 => Arc::new(Int16Array::from(data)) as ArrayRef,
743 DataType::Int32 => Arc::new(Int32Array::from(data)) as ArrayRef,
744 DataType::Int64 => Arc::new(Int64Array::from(data)) as ArrayRef,
745 DataType::UInt8 => Arc::new(UInt8Array::from(data)) as ArrayRef,
746 DataType::UInt16 => Arc::new(UInt16Array::from(data)) as ArrayRef,
747 DataType::UInt32 => Arc::new(UInt32Array::from(data)) as ArrayRef,
748 DataType::UInt64 => Arc::new(UInt64Array::from(data)) as ArrayRef,
749 DataType::Float16 => Arc::new(Float16Array::from(data)) as ArrayRef,
750 DataType::Float32 => Arc::new(Float32Array::from(data)) as ArrayRef,
751 DataType::Float64 => Arc::new(Float64Array::from(data)) as ArrayRef,
752 DataType::Date32 => Arc::new(Date32Array::from(data)) as ArrayRef,
753 DataType::Date64 => Arc::new(Date64Array::from(data)) as ArrayRef,
754 DataType::Time32(TimeUnit::Second) => Arc::new(Time32SecondArray::from(data)) as ArrayRef,
755 DataType::Time32(TimeUnit::Millisecond) => {
756 Arc::new(Time32MillisecondArray::from(data)) as ArrayRef
757 }
758 DataType::Time64(TimeUnit::Microsecond) => {
759 Arc::new(Time64MicrosecondArray::from(data)) as ArrayRef
760 }
761 DataType::Time64(TimeUnit::Nanosecond) => {
762 Arc::new(Time64NanosecondArray::from(data)) as ArrayRef
763 }
764 DataType::Timestamp(TimeUnit::Second, _) => {
765 Arc::new(TimestampSecondArray::from(data)) as ArrayRef
766 }
767 DataType::Timestamp(TimeUnit::Millisecond, _) => {
768 Arc::new(TimestampMillisecondArray::from(data)) as ArrayRef
769 }
770 DataType::Timestamp(TimeUnit::Microsecond, _) => {
771 Arc::new(TimestampMicrosecondArray::from(data)) as ArrayRef
772 }
773 DataType::Timestamp(TimeUnit::Nanosecond, _) => {
774 Arc::new(TimestampNanosecondArray::from(data)) as ArrayRef
775 }
776 DataType::Interval(IntervalUnit::YearMonth) => {
777 Arc::new(IntervalYearMonthArray::from(data)) as ArrayRef
778 }
779 DataType::Interval(IntervalUnit::DayTime) => {
780 Arc::new(IntervalDayTimeArray::from(data)) as ArrayRef
781 }
782 DataType::Interval(IntervalUnit::MonthDayNano) => {
783 Arc::new(IntervalMonthDayNanoArray::from(data)) as ArrayRef
784 }
785 DataType::Duration(TimeUnit::Second) => {
786 Arc::new(DurationSecondArray::from(data)) as ArrayRef
787 }
788 DataType::Duration(TimeUnit::Millisecond) => {
789 Arc::new(DurationMillisecondArray::from(data)) as ArrayRef
790 }
791 DataType::Duration(TimeUnit::Microsecond) => {
792 Arc::new(DurationMicrosecondArray::from(data)) as ArrayRef
793 }
794 DataType::Duration(TimeUnit::Nanosecond) => {
795 Arc::new(DurationNanosecondArray::from(data)) as ArrayRef
796 }
797 DataType::Binary => Arc::new(BinaryArray::from(data)) as ArrayRef,
798 DataType::LargeBinary => Arc::new(LargeBinaryArray::from(data)) as ArrayRef,
799 DataType::FixedSizeBinary(_) => Arc::new(FixedSizeBinaryArray::from(data)) as ArrayRef,
800 DataType::BinaryView => Arc::new(BinaryViewArray::from(data)) as ArrayRef,
801 DataType::Utf8 => Arc::new(StringArray::from(data)) as ArrayRef,
802 DataType::LargeUtf8 => Arc::new(LargeStringArray::from(data)) as ArrayRef,
803 DataType::Utf8View => Arc::new(StringViewArray::from(data)) as ArrayRef,
804 DataType::List(_) => Arc::new(ListArray::from(data)) as ArrayRef,
805 DataType::LargeList(_) => Arc::new(LargeListArray::from(data)) as ArrayRef,
806 DataType::ListView(_) => Arc::new(ListViewArray::from(data)) as ArrayRef,
807 DataType::LargeListView(_) => Arc::new(LargeListViewArray::from(data)) as ArrayRef,
808 DataType::Struct(_) => Arc::new(StructArray::from(data)) as ArrayRef,
809 DataType::Map(_, _) => Arc::new(MapArray::from(data)) as ArrayRef,
810 DataType::Union(_, _) => Arc::new(UnionArray::from(data)) as ArrayRef,
811 DataType::FixedSizeList(_, _) => Arc::new(FixedSizeListArray::from(data)) as ArrayRef,
812 DataType::Dictionary(ref key_type, _) => match key_type.as_ref() {
813 DataType::Int8 => Arc::new(DictionaryArray::<Int8Type>::from(data)) as ArrayRef,
814 DataType::Int16 => Arc::new(DictionaryArray::<Int16Type>::from(data)) as ArrayRef,
815 DataType::Int32 => Arc::new(DictionaryArray::<Int32Type>::from(data)) as ArrayRef,
816 DataType::Int64 => Arc::new(DictionaryArray::<Int64Type>::from(data)) as ArrayRef,
817 DataType::UInt8 => Arc::new(DictionaryArray::<UInt8Type>::from(data)) as ArrayRef,
818 DataType::UInt16 => Arc::new(DictionaryArray::<UInt16Type>::from(data)) as ArrayRef,
819 DataType::UInt32 => Arc::new(DictionaryArray::<UInt32Type>::from(data)) as ArrayRef,
820 DataType::UInt64 => Arc::new(DictionaryArray::<UInt64Type>::from(data)) as ArrayRef,
821 dt => panic!("Unexpected dictionary key type {dt:?}"),
822 },
823 DataType::RunEndEncoded(ref run_ends_type, _) => match run_ends_type.data_type() {
824 DataType::Int16 => Arc::new(RunArray::<Int16Type>::from(data)) as ArrayRef,
825 DataType::Int32 => Arc::new(RunArray::<Int32Type>::from(data)) as ArrayRef,
826 DataType::Int64 => Arc::new(RunArray::<Int64Type>::from(data)) as ArrayRef,
827 dt => panic!("Unexpected data type for run_ends array {dt:?}"),
828 },
829 DataType::Null => Arc::new(NullArray::from(data)) as ArrayRef,
830 DataType::Decimal128(_, _) => Arc::new(Decimal128Array::from(data)) as ArrayRef,
831 DataType::Decimal256(_, _) => Arc::new(Decimal256Array::from(data)) as ArrayRef,
832 dt => panic!("Unexpected data type {dt:?}"),
833 }
834}
835
836pub fn new_empty_array(data_type: &DataType) -> ArrayRef {
849 let data = ArrayData::new_empty(data_type);
850 make_array(data)
851}
852
853pub fn new_null_array(data_type: &DataType, length: usize) -> ArrayRef {
867 make_array(ArrayData::new_null(data_type, length))
868}
869
870unsafe fn get_offsets<O: ArrowNativeType>(data: &ArrayData) -> OffsetBuffer<O> {
876 match data.is_empty() && data.buffers()[0].is_empty() {
877 true => OffsetBuffer::new_empty(),
878 false => {
879 let buffer =
880 ScalarBuffer::new(data.buffers()[0].clone(), data.offset(), data.len() + 1);
881 unsafe { OffsetBuffer::new_unchecked(buffer) }
884 }
885 }
886}
887
888fn print_long_array<A, F>(array: &A, f: &mut std::fmt::Formatter, print_item: F) -> std::fmt::Result
890where
891 A: Array,
892 F: Fn(&A, usize, &mut std::fmt::Formatter) -> std::fmt::Result,
893{
894 let head = std::cmp::min(10, array.len());
895
896 for i in 0..head {
897 if array.is_null(i) {
898 writeln!(f, " null,")?;
899 } else {
900 write!(f, " ")?;
901 print_item(array, i, f)?;
902 writeln!(f, ",")?;
903 }
904 }
905 if array.len() > 10 {
906 if array.len() > 20 {
907 writeln!(f, " ...{} elements...,", array.len() - 20)?;
908 }
909
910 let tail = std::cmp::max(head, array.len() - 10);
911
912 for i in tail..array.len() {
913 if array.is_null(i) {
914 writeln!(f, " null,")?;
915 } else {
916 write!(f, " ")?;
917 print_item(array, i, f)?;
918 writeln!(f, ",")?;
919 }
920 }
921 }
922 Ok(())
923}
924
925#[cfg(test)]
926mod tests {
927 use super::*;
928 use crate::cast::{as_union_array, downcast_array};
929 use crate::downcast_run_array;
930 use arrow_buffer::MutableBuffer;
931 use arrow_schema::{Field, Fields, UnionFields, UnionMode};
932
933 #[test]
934 fn test_empty_primitive() {
935 let array = new_empty_array(&DataType::Int32);
936 let a = array.as_any().downcast_ref::<Int32Array>().unwrap();
937 assert_eq!(a.len(), 0);
938 let expected: &[i32] = &[];
939 assert_eq!(a.values(), expected);
940 }
941
942 #[test]
943 fn test_empty_variable_sized() {
944 let array = new_empty_array(&DataType::Utf8);
945 let a = array.as_any().downcast_ref::<StringArray>().unwrap();
946 assert_eq!(a.len(), 0);
947 assert_eq!(a.value_offsets()[0], 0i32);
948 }
949
950 #[test]
951 fn test_empty_list_primitive() {
952 let data_type = DataType::List(Arc::new(Field::new_list_field(DataType::Int32, false)));
953 let array = new_empty_array(&data_type);
954 let a = array.as_any().downcast_ref::<ListArray>().unwrap();
955 assert_eq!(a.len(), 0);
956 assert_eq!(a.value_offsets()[0], 0i32);
957 }
958
959 #[test]
960 fn test_null_boolean() {
961 let array = new_null_array(&DataType::Boolean, 9);
962 let a = array.as_any().downcast_ref::<BooleanArray>().unwrap();
963 assert_eq!(a.len(), 9);
964 for i in 0..9 {
965 assert!(a.is_null(i));
966 }
967 }
968
969 #[test]
970 fn test_null_primitive() {
971 let array = new_null_array(&DataType::Int32, 9);
972 let a = array.as_any().downcast_ref::<Int32Array>().unwrap();
973 assert_eq!(a.len(), 9);
974 for i in 0..9 {
975 assert!(a.is_null(i));
976 }
977 }
978
979 #[test]
980 fn test_null_struct() {
981 let struct_type = DataType::Struct(vec![Field::new("data", DataType::Int64, false)].into());
984 let array = new_null_array(&struct_type, 9);
985
986 let a = array.as_any().downcast_ref::<StructArray>().unwrap();
987 assert_eq!(a.len(), 9);
988 assert_eq!(a.column(0).len(), 9);
989 for i in 0..9 {
990 assert!(a.is_null(i));
991 }
992
993 a.slice(0, 5);
995 }
996
997 #[test]
998 fn test_null_variable_sized() {
999 let array = new_null_array(&DataType::Utf8, 9);
1000 let a = array.as_any().downcast_ref::<StringArray>().unwrap();
1001 assert_eq!(a.len(), 9);
1002 assert_eq!(a.value_offsets()[9], 0i32);
1003 for i in 0..9 {
1004 assert!(a.is_null(i));
1005 }
1006 }
1007
1008 #[test]
1009 fn test_null_list_primitive() {
1010 let data_type = DataType::List(Arc::new(Field::new_list_field(DataType::Int32, true)));
1011 let array = new_null_array(&data_type, 9);
1012 let a = array.as_any().downcast_ref::<ListArray>().unwrap();
1013 assert_eq!(a.len(), 9);
1014 assert_eq!(a.value_offsets()[9], 0i32);
1015 for i in 0..9 {
1016 assert!(a.is_null(i));
1017 }
1018 }
1019
1020 #[test]
1021 fn test_null_map() {
1022 let data_type = DataType::Map(
1023 Arc::new(Field::new(
1024 "entry",
1025 DataType::Struct(Fields::from(vec![
1026 Field::new("key", DataType::Utf8, false),
1027 Field::new("value", DataType::Int32, true),
1028 ])),
1029 false,
1030 )),
1031 false,
1032 );
1033 let array = new_null_array(&data_type, 9);
1034 let a = array.as_any().downcast_ref::<MapArray>().unwrap();
1035 assert_eq!(a.len(), 9);
1036 assert_eq!(a.value_offsets()[9], 0i32);
1037 for i in 0..9 {
1038 assert!(a.is_null(i));
1039 }
1040 }
1041
1042 #[test]
1043 fn test_null_dictionary() {
1044 let values =
1045 vec![None, None, None, None, None, None, None, None, None] as Vec<Option<&str>>;
1046
1047 let array: DictionaryArray<Int8Type> = values.into_iter().collect();
1048 let array = Arc::new(array) as ArrayRef;
1049
1050 let null_array = new_null_array(array.data_type(), 9);
1051 assert_eq!(&array, &null_array);
1052 assert_eq!(
1053 array.to_data().buffers()[0].len(),
1054 null_array.to_data().buffers()[0].len()
1055 );
1056 }
1057
1058 #[test]
1059 fn test_null_union() {
1060 for mode in [UnionMode::Sparse, UnionMode::Dense] {
1061 let data_type = DataType::Union(
1062 UnionFields::new(
1063 vec![2, 1],
1064 vec![
1065 Field::new("foo", DataType::Int32, true),
1066 Field::new("bar", DataType::Int64, true),
1067 ],
1068 ),
1069 mode,
1070 );
1071 let array = new_null_array(&data_type, 4);
1072
1073 let array = as_union_array(array.as_ref());
1074 assert_eq!(array.len(), 4);
1075 assert_eq!(array.null_count(), 0);
1076 assert_eq!(array.logical_null_count(), 4);
1077
1078 for i in 0..4 {
1079 let a = array.value(i);
1080 assert_eq!(a.len(), 1);
1081 assert_eq!(a.null_count(), 1);
1082 assert_eq!(a.logical_null_count(), 1);
1083 assert!(a.is_null(0))
1084 }
1085
1086 array.to_data().validate_full().unwrap();
1087 }
1088 }
1089
1090 #[test]
1091 #[allow(unused_parens)]
1092 fn test_null_runs() {
1093 for r in [DataType::Int16, DataType::Int32, DataType::Int64] {
1094 let data_type = DataType::RunEndEncoded(
1095 Arc::new(Field::new("run_ends", r, false)),
1096 Arc::new(Field::new("values", DataType::Utf8, true)),
1097 );
1098
1099 let array = new_null_array(&data_type, 4);
1100 let array = array.as_ref();
1101
1102 downcast_run_array! {
1103 array => {
1104 assert_eq!(array.len(), 4);
1105 assert_eq!(array.null_count(), 0);
1106 assert_eq!(array.logical_null_count(), 4);
1107 assert_eq!(array.values().len(), 1);
1108 assert_eq!(array.values().null_count(), 1);
1109 assert_eq!(array.run_ends().len(), 4);
1110 assert_eq!(array.run_ends().values(), &[4]);
1111
1112 let idx = array.get_physical_indices(&[0, 1, 2, 3]).unwrap();
1113 assert_eq!(idx, &[0,0,0,0]);
1114 }
1115 d => unreachable!("{d}")
1116 }
1117 }
1118 }
1119
1120 #[test]
1121 fn test_null_fixed_size_binary() {
1122 for size in [1, 2, 7] {
1123 let array = new_null_array(&DataType::FixedSizeBinary(size), 6);
1124 let array = array
1125 .as_ref()
1126 .as_any()
1127 .downcast_ref::<FixedSizeBinaryArray>()
1128 .unwrap();
1129
1130 assert_eq!(array.len(), 6);
1131 assert_eq!(array.null_count(), 6);
1132 assert_eq!(array.logical_null_count(), 6);
1133 array.iter().for_each(|x| assert!(x.is_none()));
1134 }
1135 }
1136
1137 #[test]
1138 fn test_memory_size_null() {
1139 let null_arr = NullArray::new(32);
1140
1141 assert_eq!(0, null_arr.get_buffer_memory_size());
1142 assert_eq!(
1143 std::mem::size_of::<usize>(),
1144 null_arr.get_array_memory_size()
1145 );
1146 }
1147
1148 #[test]
1149 fn test_memory_size_primitive() {
1150 let arr = PrimitiveArray::<Int64Type>::from_iter_values(0..128);
1151 let empty = PrimitiveArray::<Int64Type>::from(ArrayData::new_empty(arr.data_type()));
1152
1153 assert_eq!(
1155 arr.get_array_memory_size() - empty.get_array_memory_size(),
1156 128 * std::mem::size_of::<i64>()
1157 );
1158 }
1159
1160 #[test]
1161 fn test_memory_size_primitive_sliced() {
1162 let arr = PrimitiveArray::<Int64Type>::from_iter_values(0..128);
1163 let slice1 = arr.slice(0, 64);
1164 let slice2 = arr.slice(64, 64);
1165
1166 assert_eq!(slice1.get_array_memory_size(), arr.get_array_memory_size());
1168 assert_eq!(slice2.get_array_memory_size(), arr.get_array_memory_size());
1169 }
1170
1171 #[test]
1172 fn test_memory_size_primitive_nullable() {
1173 let arr: PrimitiveArray<Int64Type> = (0..128)
1174 .map(|i| if i % 20 == 0 { Some(i) } else { None })
1175 .collect();
1176 let empty_with_bitmap = PrimitiveArray::<Int64Type>::from(
1177 ArrayData::builder(arr.data_type().clone())
1178 .add_buffer(MutableBuffer::new(0).into())
1179 .null_bit_buffer(Some(MutableBuffer::new_null(0).into()))
1180 .build()
1181 .unwrap(),
1182 );
1183
1184 assert_eq!(
1188 std::mem::size_of::<PrimitiveArray<Int64Type>>(),
1189 empty_with_bitmap.get_array_memory_size()
1190 );
1191
1192 assert_eq!(
1195 arr.get_array_memory_size() - empty_with_bitmap.get_array_memory_size(),
1196 128 * std::mem::size_of::<i64>() + 64
1197 );
1198 }
1199
1200 #[test]
1201 fn test_memory_size_dictionary() {
1202 let values = PrimitiveArray::<Int64Type>::from_iter_values(0..16);
1203 let keys = PrimitiveArray::<Int16Type>::from_iter_values(
1204 (0..256).map(|i| (i % values.len()) as i16),
1205 );
1206
1207 let dict_data_type = DataType::Dictionary(
1208 Box::new(keys.data_type().clone()),
1209 Box::new(values.data_type().clone()),
1210 );
1211 let dict_data = keys
1212 .into_data()
1213 .into_builder()
1214 .data_type(dict_data_type)
1215 .child_data(vec![values.into_data()])
1216 .build()
1217 .unwrap();
1218
1219 let empty_data = ArrayData::new_empty(&DataType::Dictionary(
1220 Box::new(DataType::Int16),
1221 Box::new(DataType::Int64),
1222 ));
1223
1224 let arr = DictionaryArray::<Int16Type>::from(dict_data);
1225 let empty = DictionaryArray::<Int16Type>::from(empty_data);
1226
1227 let expected_keys_size = 256 * std::mem::size_of::<i16>();
1228 assert_eq!(
1229 arr.keys().get_array_memory_size() - empty.keys().get_array_memory_size(),
1230 expected_keys_size
1231 );
1232
1233 let expected_values_size = 16 * std::mem::size_of::<i64>();
1234 assert_eq!(
1235 arr.values().get_array_memory_size() - empty.values().get_array_memory_size(),
1236 expected_values_size
1237 );
1238
1239 let expected_size = expected_keys_size + expected_values_size;
1240 assert_eq!(
1241 arr.get_array_memory_size() - empty.get_array_memory_size(),
1242 expected_size
1243 );
1244 }
1245
1246 fn compute_my_thing(arr: &dyn Array) -> bool {
1248 !arr.is_empty()
1249 }
1250
1251 #[test]
1252 fn test_array_ref_as_array() {
1253 let arr: Int32Array = vec![1, 2, 3].into_iter().map(Some).collect();
1254
1255 assert!(compute_my_thing(&arr));
1257
1258 let arr: ArrayRef = Arc::new(arr);
1260 assert!(compute_my_thing(&arr));
1261 assert!(compute_my_thing(arr.as_ref()));
1262 }
1263
1264 #[test]
1265 fn test_downcast_array() {
1266 let array: Int32Array = vec![1, 2, 3].into_iter().map(Some).collect();
1267
1268 let boxed: ArrayRef = Arc::new(array);
1269 let array: Int32Array = downcast_array(&boxed);
1270
1271 let expected: Int32Array = vec![1, 2, 3].into_iter().map(Some).collect();
1272 assert_eq!(array, expected);
1273 }
1274}