Struct variant

bench-vortex/src/vortex_utils.rs

@@ -4,6 +4,7 @@ use std::path::PathBuf;
 
 use vortex::array::chunked::ChunkedArray;
 use vortex::array::struct_::StructArray;
+use vortex::variants::StructArrayTrait;
 use vortex::ArrayDType;
 use vortex_dtype::DType;
 use vortex_error::VortexResult;

encodings/alp/src/array.rs

@@ -2,6 +2,7 @@ use serde::{Deserialize, Serialize};
 use vortex::array::primitive::PrimitiveArray;
 use vortex::stats::ArrayStatisticsCompute;
 use vortex::validity::{ArrayValidity, LogicalValidity};
+use vortex::variants::{ArrayVariants, PrimitiveArrayTrait};
 use vortex::visitor::{AcceptArrayVisitor, ArrayVisitor};
 use vortex::{impl_encoding, ArrayDType, Canonical, IntoCanonical};
 use vortex_dtype::PType;
@@ -92,6 +93,14 @@ impl ALPArray {
 
 impl ArrayTrait for ALPArray {}
 
+impl ArrayVariants for ALPArray {
+    fn as_primitive_array(&self) -> Option<&dyn PrimitiveArrayTrait> {
+        Some(self)
+    }
+}
+
+impl PrimitiveArrayTrait for ALPArray {}
+
 impl ArrayValidity for ALPArray {
     fn is_valid(&self, index: usize) -> bool {
         self.encoded().with_dyn(|a| a.is_valid(index))

encodings/byte_bool/src/lib.rs

@@ -3,6 +3,7 @@ use std::mem::ManuallyDrop;
 use arrow_buffer::BooleanBuffer;
 use serde::{Deserialize, Serialize};
 use vortex::array::bool::BoolArray;
+use vortex::variants::{ArrayVariants, BoolArrayTrait};
 use vortex::{
     impl_encoding,
     validity::{ArrayValidity, LogicalValidity, Validity, ValidityMetadata},
@@ -73,6 +74,14 @@ impl ByteBoolArray {
 
 impl ArrayTrait for ByteBoolArray {}
 
+impl ArrayVariants for ByteBoolArray {
+    fn as_bool_array(&self) -> Option<&dyn BoolArrayTrait> {
+        Some(self)
+    }
+}
+
+impl BoolArrayTrait for ByteBoolArray {}
+
 impl From<Vec<bool>> for ByteBoolArray {
     fn from(value: Vec<bool>) -> Self {
         Self::try_from_vec(value, Validity::AllValid).unwrap()

encodings/datetime-parts/src/array.rs

@@ -1,6 +1,7 @@
 use serde::{Deserialize, Serialize};
 use vortex::stats::ArrayStatisticsCompute;
 use vortex::validity::{ArrayValidity, LogicalValidity};
+use vortex::variants::{ArrayVariants, ExtensionArrayTrait};
 use vortex::visitor::{AcceptArrayVisitor, ArrayVisitor};
 use vortex::{impl_encoding, ArrayDType, Canonical, IntoCanonical};
 use vortex_error::vortex_bail;
@@ -79,6 +80,14 @@ impl DateTimePartsArray {
 
 impl ArrayTrait for DateTimePartsArray {}
 
+impl ArrayVariants for DateTimePartsArray {
+    fn as_extension_array(&self) -> Option<&dyn ExtensionArrayTrait> {
+        Some(self)
+    }
+}
+
+impl ExtensionArrayTrait for DateTimePartsArray {}
+
 impl IntoCanonical for DateTimePartsArray {
     fn into_canonical(self) -> VortexResult<Canonical> {
         Ok(Canonical::Extension(

encodings/dict/src/lib.rs

@@ -9,3 +9,4 @@ mod compress;
 mod compute;
 mod dict;
 mod stats;
+mod variants;

encodings/dict/src/variants.rs

@@ -0,0 +1,23 @@
+use vortex::variants::{ArrayVariants, BinaryArrayTrait, PrimitiveArrayTrait, Utf8ArrayTrait};
+
+use crate::DictArray;
+
+impl ArrayVariants for DictArray {
+    fn as_primitive_array(&self) -> Option<&dyn PrimitiveArrayTrait> {
+        Some(self)
+    }
+
+    fn as_utf8_array(&self) -> Option<&dyn Utf8ArrayTrait> {
+        Some(self)
+    }
+
+    fn as_binary_array(&self) -> Option<&dyn BinaryArrayTrait> {
+        Some(self)
+    }
+}
+
+impl PrimitiveArrayTrait for DictArray {}
+
+impl Utf8ArrayTrait for DictArray {}
+
+impl BinaryArrayTrait for DictArray {}

encodings/fastlanes/src/bitpacking/mod.rs

@@ -3,6 +3,7 @@ pub use compress::*;
 use vortex::array::primitive::{Primitive, PrimitiveArray};
 use vortex::stats::ArrayStatisticsCompute;
 use vortex::validity::{ArrayValidity, LogicalValidity, Validity, ValidityMetadata};
+use vortex::variants::{ArrayVariants, PrimitiveArrayTrait};
 use vortex::visitor::{AcceptArrayVisitor, ArrayVisitor};
 use vortex::{impl_encoding, ArrayDType, Canonical, IntoCanonical};
 use vortex_dtype::{Nullability, PType};
@@ -201,6 +202,14 @@ impl ArrayTrait for BitPackedArray {
     }
 }
 
+impl ArrayVariants for BitPackedArray {
+    fn as_primitive_array(&self) -> Option<&dyn PrimitiveArrayTrait> {
+        Some(self)
+    }
+}
+
+impl PrimitiveArrayTrait for BitPackedArray {}
+
 #[cfg(test)]
 mod test {
     use vortex::array::primitive::PrimitiveArray;

encodings/fastlanes/src/delta/mod.rs

@@ -3,6 +3,7 @@ use serde::{Deserialize, Serialize};
 use vortex::stats::ArrayStatisticsCompute;
 use vortex::validity::ValidityMetadata;
 use vortex::validity::{ArrayValidity, LogicalValidity, Validity};
+use vortex::variants::{ArrayVariants, PrimitiveArrayTrait};
 use vortex::visitor::{AcceptArrayVisitor, ArrayVisitor};
 use vortex::{impl_encoding, ArrayDType, Canonical, IntoCanonical};
 use vortex_dtype::match_each_unsigned_integer_ptype;
@@ -92,6 +93,14 @@ impl DeltaArray {
 
 impl ArrayTrait for DeltaArray {}
 
+impl ArrayVariants for DeltaArray {
+    fn as_primitive_array(&self) -> Option<&dyn PrimitiveArrayTrait> {
+        Some(self)
+    }
+}
+
+impl PrimitiveArrayTrait for DeltaArray {}
+
 impl IntoCanonical for DeltaArray {
     fn into_canonical(self) -> VortexResult<Canonical> {
         delta_decompress(self).map(Canonical::Primitive)

encodings/fastlanes/src/for/mod.rs

@@ -2,6 +2,7 @@ pub use compress::*;
 use serde::{Deserialize, Serialize};
 use vortex::stats::ArrayStatisticsCompute;
 use vortex::validity::{ArrayValidity, LogicalValidity};
+use vortex::variants::{ArrayVariants, PrimitiveArrayTrait};
 use vortex::visitor::{AcceptArrayVisitor, ArrayVisitor};
 use vortex::{impl_encoding, ArrayDType, Canonical, IntoCanonical};
 use vortex_dtype::PType;
@@ -95,3 +96,11 @@ impl ArrayTrait for FoRArray {
         self.encoded().nbytes()
     }
 }
+
+impl ArrayVariants for FoRArray {
+    fn as_primitive_array(&self) -> Option<&dyn PrimitiveArrayTrait> {
+        Some(self)
+    }
+}
+
+impl PrimitiveArrayTrait for FoRArray {}

encodings/roaring/src/boolean/mod.rs

@@ -6,6 +6,7 @@ use serde::{Deserialize, Serialize};
 use vortex::array::bool::{Bool, BoolArray};
 use vortex::stats::ArrayStatisticsCompute;
 use vortex::validity::{ArrayValidity, LogicalValidity, Validity};
+use vortex::variants::{ArrayVariants, BoolArrayTrait};
 use vortex::visitor::{AcceptArrayVisitor, ArrayVisitor};
 use vortex::{impl_encoding, ArrayDType, Canonical, IntoCanonical};
 use vortex_buffer::Buffer;
@@ -62,6 +63,14 @@ impl RoaringBoolArray {
 
 impl ArrayTrait for RoaringBoolArray {}
 
+impl ArrayVariants for RoaringBoolArray {
+    fn as_bool_array(&self) -> Option<&dyn BoolArrayTrait> {
+        Some(self)
+    }
+}
+
+impl BoolArrayTrait for RoaringBoolArray {}
+
 impl AcceptArrayVisitor for RoaringBoolArray {
     fn accept(&self, _visitor: &mut dyn ArrayVisitor) -> VortexResult<()> {
         // TODO(ngates): should we store a buffer in memory? Or delay serialization?

encodings/roaring/src/integer/mod.rs

@@ -4,6 +4,7 @@ use serde::{Deserialize, Serialize};
 use vortex::array::primitive::{Primitive, PrimitiveArray};
 use vortex::stats::ArrayStatisticsCompute;
 use vortex::validity::{ArrayValidity, LogicalValidity};
+use vortex::variants::{ArrayVariants, PrimitiveArrayTrait};
 use vortex::visitor::{AcceptArrayVisitor, ArrayVisitor};
 use vortex::{impl_encoding, Canonical, IntoCanonical};
 use vortex_buffer::Buffer;
@@ -64,6 +65,14 @@ impl RoaringIntArray {
 
 impl ArrayTrait for RoaringIntArray {}
 
+impl ArrayVariants for RoaringIntArray {
+    fn as_primitive_array(&self) -> Option<&dyn PrimitiveArrayTrait> {
+        Some(self)
+    }
+}
+
+impl PrimitiveArrayTrait for RoaringIntArray {}
+
 impl ArrayValidity for RoaringIntArray {
     fn is_valid(&self, _index: usize) -> bool {
         true

encodings/runend/src/runend.rs

@@ -4,6 +4,7 @@ use vortex::compute::unary::scalar_at::scalar_at;
 use vortex::compute::{search_sorted, SearchSortedSide};
 use vortex::stats::{ArrayStatistics, ArrayStatisticsCompute};
 use vortex::validity::{ArrayValidity, LogicalValidity, Validity, ValidityMetadata};
+use vortex::variants::{ArrayVariants, PrimitiveArrayTrait};
 use vortex::visitor::{AcceptArrayVisitor, ArrayVisitor};
 use vortex::{impl_encoding, ArrayDType, Canonical, IntoArrayVariant, IntoCanonical};
 use vortex_error::vortex_bail;
@@ -107,6 +108,14 @@ impl RunEndArray {
 
 impl ArrayTrait for RunEndArray {}
 
+impl ArrayVariants for RunEndArray {
+    fn as_primitive_array(&self) -> Option<&dyn PrimitiveArrayTrait> {
+        Some(self)
+    }
+}
+
+impl PrimitiveArrayTrait for RunEndArray {}
+
 impl ArrayValidity for RunEndArray {
     fn is_valid(&self, index: usize) -> bool {
         self.validity().is_valid(index)

encodings/zigzag/src/zigzag.rs

@@ -2,6 +2,7 @@ use serde::{Deserialize, Serialize};
 use vortex::array::primitive::PrimitiveArray;
 use vortex::stats::ArrayStatisticsCompute;
 use vortex::validity::{ArrayValidity, LogicalValidity};
+use vortex::variants::{ArrayVariants, PrimitiveArrayTrait};
 use vortex::visitor::{AcceptArrayVisitor, ArrayVisitor};
 use vortex::{impl_encoding, ArrayDType, Canonical, IntoCanonical};
 use vortex_dtype::PType;
@@ -52,6 +53,14 @@ impl ZigZagArray {
 
 impl ArrayTrait for ZigZagArray {}
 
+impl ArrayVariants for ZigZagArray {
+    fn as_primitive_array(&self) -> Option<&dyn PrimitiveArrayTrait> {
+        Some(self)
+    }
+}
+
+impl PrimitiveArrayTrait for ZigZagArray {}
+
 impl ArrayValidity for ZigZagArray {
     fn is_valid(&self, index: usize) -> bool {
         self.encoded().with_dyn(|a| a.is_valid(index))

vortex-array/src/array/bool/mod.rs

@@ -5,8 +5,9 @@ use vortex_buffer::Buffer;
 
 use crate::validity::{ArrayValidity, ValidityMetadata};
 use crate::validity::{LogicalValidity, Validity};
+use crate::variants::{ArrayVariants, BoolArrayTrait};
 use crate::visitor::{AcceptArrayVisitor, ArrayVisitor};
-use crate::{impl_encoding, BoolArrayTrait, Canonical, IntoCanonical};
+use crate::{impl_encoding, Canonical, IntoCanonical};
 
 mod accessors;
 mod compute;
@@ -72,7 +73,9 @@ impl BoolArray {
     }
 }
 
-impl ArrayTrait for BoolArray {
+impl ArrayTrait for BoolArray {}
+
+impl ArrayVariants for BoolArray {
     fn as_bool_array(&self) -> Option<&dyn BoolArrayTrait> {
         Some(self)
     }

vortex-array/src/array/chunked/canonical.rs

@@ -13,6 +13,7 @@ use crate::array::struct_::StructArray;
 use crate::array::varbin::builder::VarBinBuilder;
 use crate::array::varbin::VarBinArray;
 use crate::validity::Validity;
+use crate::variants::StructArrayTrait;
 use crate::{
     Array, ArrayDType, ArrayValidity, Canonical, IntoArray, IntoArrayVariant, IntoCanonical,
 };

vortex-array/src/array/chunked/mod.rs

@@ -22,6 +22,7 @@ use crate::{impl_encoding, ArrayDType};
 mod canonical;
 mod compute;
 mod stats;
+mod variants;
 
 impl_encoding!("vortex.chunked", 11u16, Chunked);
 

vortex-array/src/array/chunked/variants.rs

@@ -0,0 +1,69 @@
+use crate::array::chunked::ChunkedArray;
+use crate::variants::{
+    ArrayVariants, BinaryArrayTrait, BoolArrayTrait, ExtensionArrayTrait, ListArrayTrait,
+    NullArrayTrait, PrimitiveArrayTrait, StructArrayTrait, Utf8ArrayTrait,
+};
+use crate::{Array, ArrayDType, IntoArray};
+
+/// Chunked arrays support all DTypes
+impl ArrayVariants for ChunkedArray {
+    fn as_null_array(&self) -> Option<&dyn NullArrayTrait> {
+        Some(self)
+    }
+
+    fn as_bool_array(&self) -> Option<&dyn BoolArrayTrait> {
+        Some(self)
+    }
+
+    fn as_primitive_array(&self) -> Option<&dyn PrimitiveArrayTrait> {
+        Some(self)
+    }
+
+    fn as_utf8_array(&self) -> Option<&dyn Utf8ArrayTrait> {
+        Some(self)
+    }
+
+    fn as_binary_array(&self) -> Option<&dyn BinaryArrayTrait> {
+        Some(self)
+    }
+
+    fn as_struct_array(&self) -> Option<&dyn StructArrayTrait> {
+        Some(self)
+    }
+
+    fn as_list_array(&self) -> Option<&dyn ListArrayTrait> {
+        Some(self)
+    }
+
+    fn as_extension_array(&self) -> Option<&dyn ExtensionArrayTrait> {
+        Some(self)
+    }
+}
+
+impl NullArrayTrait for ChunkedArray {}
+
+impl BoolArrayTrait for ChunkedArray {}
+
+impl PrimitiveArrayTrait for ChunkedArray {}
+
+impl Utf8ArrayTrait for ChunkedArray {}
+
+impl BinaryArrayTrait for ChunkedArray {}
+
+impl StructArrayTrait for ChunkedArray {
+    fn field(&self, idx: usize) -> Option<Array> {
+        let mut chunks = Vec::with_capacity(self.nchunks());
+        for chunk in self.chunks() {
+            let array = chunk.with_dyn(|a| a.as_struct_array().and_then(|s| s.field(idx)))?;
+            chunks.push(array);
+        }
+        let chunked = ChunkedArray::try_new(chunks, self.dtype().clone())
+            .expect("should be correct dtype")
+            .into_array();
+        Some(chunked)
+    }
+}
+
+impl ListArrayTrait for ChunkedArray {}
+
+impl ExtensionArrayTrait for ChunkedArray {}

vortex-array/src/array/constant/mod.rs

@@ -11,6 +11,7 @@ use crate::visitor::{AcceptArrayVisitor, ArrayVisitor};
 mod canonical;
 mod compute;
 mod stats;
+mod variants;
 
 impl_encoding!("vortex.constant", 10u16, Constant);