Variants VTable (#1501)

encodings/alp/src/alp/array.rs

@@ -5,7 +5,7 @@ use vortex_array::array::PrimitiveArray;
 use vortex_array::encoding::ids;
 use vortex_array::stats::StatisticsVTable;
 use vortex_array::validity::{ArrayValidity, LogicalValidity, ValidityVTable};
-use vortex_array::variants::{ArrayVariants, PrimitiveArrayTrait};
+use vortex_array::variants::{PrimitiveArrayTrait, VariantsVTable};
 use vortex_array::visitor::{ArrayVisitor, VisitorVTable};
 use vortex_array::{
     impl_encoding, ArrayDType, ArrayData, ArrayLen, ArrayTrait, Canonical, IntoArrayData,
@@ -124,9 +124,9 @@ impl ALPArray {
 
 impl ArrayTrait for ALPArray {}
 
-impl ArrayVariants for ALPArray {
-    fn as_primitive_array(&self) -> Option<&dyn PrimitiveArrayTrait> {
-        Some(self)
+impl VariantsVTable<ALPArray> for ALPEncoding {
+    fn as_primitive_array<'a>(&self, array: &'a ALPArray) -> Option<&'a dyn PrimitiveArrayTrait> {
+        Some(array)
     }
 }
 

encodings/alp/src/alp_rd/variants.rs

@@ -1,14 +1,10 @@
-use vortex_array::variants::{ArrayVariants, PrimitiveArrayTrait};
+use vortex_array::variants::{PrimitiveArrayTrait, VariantsVTable};
 
-use crate::ALPRDArray;
+use crate::{ALPRDArray, ALPRDEncoding};
 
-impl ArrayVariants for ALPRDArray {
-    fn as_primitive_array(&self) -> Option<&dyn PrimitiveArrayTrait> {
-        Some(self)
-    }
-
-    fn as_primitive_array_unchecked(&self) -> &dyn PrimitiveArrayTrait {
-        self
+impl VariantsVTable<ALPRDArray> for ALPRDEncoding {
+    fn as_primitive_array<'a>(&self, array: &'a ALPRDArray) -> Option<&'a dyn PrimitiveArrayTrait> {
+        Some(array)
     }
 }
 

encodings/bytebool/src/array.rs

@@ -8,7 +8,7 @@ use vortex_array::array::BoolArray;
 use vortex_array::encoding::ids;
 use vortex_array::stats::StatsSet;
 use vortex_array::validity::{LogicalValidity, Validity, ValidityMetadata, ValidityVTable};
-use vortex_array::variants::{ArrayVariants, BoolArrayTrait};
+use vortex_array::variants::{BoolArrayTrait, VariantsVTable};
 use vortex_array::visitor::{ArrayVisitor, VisitorVTable};
 use vortex_array::{impl_encoding, ArrayData, ArrayLen, ArrayTrait, Canonical, IntoCanonical};
 use vortex_buffer::Buffer;
@@ -84,9 +84,9 @@ impl ByteBoolArray {
 
 impl ArrayTrait for ByteBoolArray {}
 
-impl ArrayVariants for ByteBoolArray {
-    fn as_bool_array(&self) -> Option<&dyn BoolArrayTrait> {
-        Some(self)
+impl VariantsVTable<ByteBoolArray> for ByteBoolEncoding {
+    fn as_bool_array<'a>(&self, array: &'a ByteBoolArray) -> Option<&'a dyn BoolArrayTrait> {
+        Some(array)
     }
 }
 

encodings/datetime-parts/src/array.rs

@@ -6,7 +6,7 @@ use vortex_array::compute::try_cast;
 use vortex_array::encoding::ids;
 use vortex_array::stats::{Stat, StatisticsVTable, StatsSet};
 use vortex_array::validity::{ArrayValidity, LogicalValidity, Validity, ValidityVTable};
-use vortex_array::variants::{ArrayVariants, ExtensionArrayTrait};
+use vortex_array::variants::{ExtensionArrayTrait, VariantsVTable};
 use vortex_array::visitor::{ArrayVisitor, VisitorVTable};
 use vortex_array::{
     impl_encoding, ArrayDType, ArrayData, ArrayLen, ArrayTrait, Canonical, IntoArrayData,
@@ -114,9 +114,12 @@ impl DateTimePartsArray {
 
 impl ArrayTrait for DateTimePartsArray {}
 
-impl ArrayVariants for DateTimePartsArray {
-    fn as_extension_array(&self) -> Option<&dyn ExtensionArrayTrait> {
-        Some(self)
+impl VariantsVTable<DateTimePartsArray> for DateTimePartsEncoding {
+    fn as_extension_array<'a>(
+        &self,
+        array: &'a DateTimePartsArray,
+    ) -> Option<&'a dyn ExtensionArrayTrait> {
+        Some(array)
     }
 }
 

encodings/dict/src/variants.rs

@@ -1,26 +1,24 @@
 use vortex_array::variants::{
-    ArrayVariants, BinaryArrayTrait, BoolArrayTrait, PrimitiveArrayTrait, Utf8ArrayTrait,
+    BinaryArrayTrait, BoolArrayTrait, PrimitiveArrayTrait, Utf8ArrayTrait, VariantsVTable,
 };
-use vortex_array::ArrayDType;
-use vortex_dtype::DType;
 
-use crate::DictArray;
+use crate::{DictArray, DictEncoding};
 
-impl ArrayVariants for DictArray {
-    fn as_bool_array(&self) -> Option<&dyn BoolArrayTrait> {
-        matches!(self.dtype(), DType::Bool(..)).then_some(self)
+impl VariantsVTable<DictArray> for DictEncoding {
+    fn as_bool_array<'a>(&self, array: &'a DictArray) -> Option<&'a dyn BoolArrayTrait> {
+        Some(array)
     }
 
-    fn as_primitive_array(&self) -> Option<&dyn PrimitiveArrayTrait> {
-        matches!(self.dtype(), DType::Primitive(..)).then_some(self)
+    fn as_primitive_array<'a>(&self, array: &'a DictArray) -> Option<&'a dyn PrimitiveArrayTrait> {
+        Some(array)
     }
 
-    fn as_utf8_array(&self) -> Option<&dyn Utf8ArrayTrait> {
-        matches!(self.dtype(), DType::Utf8(..)).then_some(self)
+    fn as_utf8_array<'a>(&self, array: &'a DictArray) -> Option<&'a dyn Utf8ArrayTrait> {
+        Some(array)
     }
 
-    fn as_binary_array(&self) -> Option<&dyn BinaryArrayTrait> {
-        matches!(self.dtype(), DType::Binary(..)).then_some(self)
+    fn as_binary_array<'a>(&self, array: &'a DictArray) -> Option<&'a dyn BinaryArrayTrait> {
+        Some(array)
     }
 }
 

encodings/fastlanes/src/bitpacking/mod.rs

@@ -8,7 +8,7 @@ use vortex_array::array::{PrimitiveArray, SparseArray};
 use vortex_array::encoding::ids;
 use vortex_array::stats::{StatisticsVTable, StatsSet};
 use vortex_array::validity::{LogicalValidity, Validity, ValidityMetadata, ValidityVTable};
-use vortex_array::variants::{ArrayVariants, PrimitiveArrayTrait};
+use vortex_array::variants::{PrimitiveArrayTrait, VariantsVTable};
 use vortex_array::visitor::{ArrayVisitor, VisitorVTable};
 use vortex_array::{
     impl_encoding, ArrayDType, ArrayData, ArrayLen, ArrayTrait, Canonical, IntoCanonical,
@@ -232,9 +232,12 @@ impl StatisticsVTable<BitPackedArray> for BitPackedEncoding {}
 
 impl ArrayTrait for BitPackedArray {}
 
-impl ArrayVariants for BitPackedArray {
-    fn as_primitive_array(&self) -> Option<&dyn PrimitiveArrayTrait> {
-        Some(self)
+impl VariantsVTable<BitPackedArray> for BitPackedEncoding {
+    fn as_primitive_array<'a>(
+        &self,
+        array: &'a BitPackedArray,
+    ) -> Option<&'a dyn PrimitiveArrayTrait> {
+        Some(array)
     }
 }
 

encodings/fastlanes/src/delta/mod.rs

@@ -6,7 +6,7 @@ use vortex_array::array::PrimitiveArray;
 use vortex_array::encoding::ids;
 use vortex_array::stats::{StatisticsVTable, StatsSet};
 use vortex_array::validity::{LogicalValidity, Validity, ValidityMetadata, ValidityVTable};
-use vortex_array::variants::{ArrayVariants, PrimitiveArrayTrait};
+use vortex_array::variants::{PrimitiveArrayTrait, VariantsVTable};
 use vortex_array::visitor::{ArrayVisitor, VisitorVTable};
 use vortex_array::{
     impl_encoding, ArrayDType, ArrayData, ArrayLen, ArrayTrait, Canonical, IntoArrayData,
@@ -218,9 +218,9 @@ impl DeltaArray {
 
 impl ArrayTrait for DeltaArray {}
 
-impl ArrayVariants for DeltaArray {
-    fn as_primitive_array(&self) -> Option<&dyn PrimitiveArrayTrait> {
-        Some(self)
+impl VariantsVTable<DeltaArray> for DeltaEncoding {
+    fn as_primitive_array<'a>(&self, array: &'a DeltaArray) -> Option<&'a dyn PrimitiveArrayTrait> {
+        Some(array)
     }
 }
 

encodings/fastlanes/src/for/mod.rs

@@ -5,7 +5,7 @@ use serde::{Deserialize, Serialize};
 use vortex_array::encoding::ids;
 use vortex_array::stats::{StatisticsVTable, StatsSet};
 use vortex_array::validity::{ArrayValidity, LogicalValidity, ValidityVTable};
-use vortex_array::variants::{ArrayVariants, PrimitiveArrayTrait};
+use vortex_array::variants::{PrimitiveArrayTrait, VariantsVTable};
 use vortex_array::visitor::{ArrayVisitor, VisitorVTable};
 use vortex_array::{
     impl_encoding, ArrayDType, ArrayData, ArrayLen, ArrayTrait, Canonical, IntoCanonical,
@@ -104,9 +104,9 @@ impl StatisticsVTable<FoRArray> for FoREncoding {}
 
 impl ArrayTrait for FoRArray {}
 
-impl ArrayVariants for FoRArray {
-    fn as_primitive_array(&self) -> Option<&dyn PrimitiveArrayTrait> {
-        Some(self)
+impl VariantsVTable<FoRArray> for FoREncoding {
+    fn as_primitive_array<'a>(&self, array: &'a FoRArray) -> Option<&'a dyn PrimitiveArrayTrait> {
+        Some(array)
     }
 }
 

encodings/fsst/src/array.rs

@@ -7,7 +7,7 @@ use vortex_array::array::{VarBin, VarBinArray};
 use vortex_array::encoding::ids;
 use vortex_array::stats::{StatisticsVTable, StatsSet};
 use vortex_array::validity::{ArrayValidity, LogicalValidity, Validity, ValidityVTable};
-use vortex_array::variants::{ArrayVariants, BinaryArrayTrait, Utf8ArrayTrait};
+use vortex_array::variants::{BinaryArrayTrait, Utf8ArrayTrait, VariantsVTable};
 use vortex_array::visitor::{ArrayVisitor, VisitorVTable};
 use vortex_array::{
     impl_encoding, ArrayDType, ArrayData, ArrayDef, ArrayLen, ArrayTrait, IntoCanonical,
@@ -212,13 +212,13 @@ impl ValidityVTable<FSSTArray> for FSSTEncoding {
     }
 }
 
-impl ArrayVariants for FSSTArray {
-    fn as_utf8_array(&self) -> Option<&dyn Utf8ArrayTrait> {
-        Some(self)
+impl VariantsVTable<FSSTArray> for FSSTEncoding {
+    fn as_utf8_array<'a>(&self, array: &'a FSSTArray) -> Option<&'a dyn Utf8ArrayTrait> {
+        Some(array)
     }
 
-    fn as_binary_array(&self) -> Option<&dyn BinaryArrayTrait> {
-        Some(self)
+    fn as_binary_array<'a>(&self, array: &'a FSSTArray) -> Option<&'a dyn BinaryArrayTrait> {
+        Some(array)
     }
 }
 

encodings/roaring/src/boolean/mod.rs

@@ -10,7 +10,7 @@ use vortex_array::array::BoolArray;
 use vortex_array::encoding::ids;
 use vortex_array::stats::StatsSet;
 use vortex_array::validity::{LogicalValidity, ValidityVTable};
-use vortex_array::variants::{ArrayVariants, BoolArrayTrait};
+use vortex_array::variants::{BoolArrayTrait, VariantsVTable};
 use vortex_array::visitor::{ArrayVisitor, VisitorVTable};
 use vortex_array::{
     impl_encoding, ArrayData, ArrayLen, ArrayTrait, Canonical, IntoArrayData, IntoCanonical,
@@ -84,9 +84,9 @@ impl RoaringBoolArray {
 
 impl ArrayTrait for RoaringBoolArray {}
 
-impl ArrayVariants for RoaringBoolArray {
-    fn as_bool_array(&self) -> Option<&dyn BoolArrayTrait> {
-        Some(self)
+impl VariantsVTable<RoaringBoolArray> for RoaringBoolEncoding {
+    fn as_bool_array<'a>(&self, array: &'a RoaringBoolArray) -> Option<&'a dyn BoolArrayTrait> {
+        Some(array)
     }
 }
 

encodings/roaring/src/integer/mod.rs

@@ -9,7 +9,7 @@ use vortex_array::compute::try_cast;
 use vortex_array::encoding::ids;
 use vortex_array::stats::{ArrayStatistics, Stat, StatisticsVTable, StatsSet};
 use vortex_array::validity::{LogicalValidity, Validity, ValidityVTable};
-use vortex_array::variants::{ArrayVariants, PrimitiveArrayTrait};
+use vortex_array::variants::{PrimitiveArrayTrait, VariantsVTable};
 use vortex_array::visitor::{ArrayVisitor, VisitorVTable};
 use vortex_array::{
     impl_encoding, ArrayDType as _, ArrayData, ArrayLen, ArrayTrait, Canonical, IntoArrayData,
@@ -99,9 +99,12 @@ impl RoaringIntArray {
 
 impl ArrayTrait for RoaringIntArray {}
 
-impl ArrayVariants for RoaringIntArray {
-    fn as_primitive_array(&self) -> Option<&dyn PrimitiveArrayTrait> {
-        Some(self)
+impl VariantsVTable<RoaringIntArray> for RoaringIntEncoding {
+    fn as_primitive_array<'a>(
+        &self,
+        array: &'a RoaringIntArray,
+    ) -> Option<&'a dyn PrimitiveArrayTrait> {
+        Some(array)
     }
 }
 

encodings/runend-bool/src/array.rs

@@ -6,7 +6,7 @@ use vortex_array::compute::{scalar_at, search_sorted, SearchSortedSide};
 use vortex_array::encoding::ids;
 use vortex_array::stats::{ArrayStatistics, Stat, StatisticsVTable, StatsSet};
 use vortex_array::validity::{LogicalValidity, Validity, ValidityMetadata, ValidityVTable};
-use vortex_array::variants::{ArrayVariants, BoolArrayTrait, PrimitiveArrayTrait};
+use vortex_array::variants::{BoolArrayTrait, PrimitiveArrayTrait, VariantsVTable};
 use vortex_array::visitor::{ArrayVisitor, VisitorVTable};
 use vortex_array::{
     impl_encoding, ArrayDType, ArrayData, ArrayLen, ArrayTrait, Canonical, IntoArrayData,
@@ -178,9 +178,9 @@ pub(crate) fn value_at_index(idx: usize, start: bool) -> bool {
 
 impl BoolArrayTrait for RunEndBoolArray {}
 
-impl ArrayVariants for RunEndBoolArray {
-    fn as_bool_array(&self) -> Option<&dyn BoolArrayTrait> {
-        Some(self)
+impl VariantsVTable<RunEndBoolArray> for RunEndBoolEncoding {
+    fn as_bool_array<'a>(&self, array: &'a RunEndBoolArray) -> Option<&'a dyn BoolArrayTrait> {
+        Some(array)
     }
 }
 

encodings/runend/src/array.rs

@@ -8,7 +8,7 @@ use vortex_array::stats::{ArrayStatistics, Stat, StatisticsVTable, StatsSet};
 use vortex_array::validity::{
     ArrayValidity, LogicalValidity, Validity, ValidityMetadata, ValidityVTable,
 };
-use vortex_array::variants::{ArrayVariants, BoolArrayTrait, PrimitiveArrayTrait};
+use vortex_array::variants::{BoolArrayTrait, PrimitiveArrayTrait, VariantsVTable};
 use vortex_array::visitor::{ArrayVisitor, VisitorVTable};
 use vortex_array::{
     impl_encoding, ArrayDType, ArrayData, ArrayLen, ArrayTrait, Canonical, IntoArrayData,
@@ -187,9 +187,16 @@ impl RunEndArray {
 
 impl ArrayTrait for RunEndArray {}
 
-impl ArrayVariants for RunEndArray {
-    fn as_primitive_array(&self) -> Option<&dyn PrimitiveArrayTrait> {
-        Some(self)
+impl VariantsVTable<RunEndArray> for RunEndEncoding {
+    fn as_bool_array<'a>(&self, array: &'a RunEndArray) -> Option<&'a dyn BoolArrayTrait> {
+        Some(array)
+    }
+
+    fn as_primitive_array<'a>(
+        &self,
+        array: &'a RunEndArray,
+    ) -> Option<&'a dyn PrimitiveArrayTrait> {
+        Some(array)
     }
 }
 

encodings/zigzag/src/array.rs

@@ -5,7 +5,7 @@ use vortex_array::array::PrimitiveArray;
 use vortex_array::encoding::ids;
 use vortex_array::stats::{ArrayStatistics, Stat, StatisticsVTable, StatsSet};
 use vortex_array::validity::{ArrayValidity, LogicalValidity, ValidityVTable};
-use vortex_array::variants::{ArrayVariants, PrimitiveArrayTrait};
+use vortex_array::variants::{PrimitiveArrayTrait, VariantsVTable};
 use vortex_array::visitor::{ArrayVisitor, VisitorVTable};
 use vortex_array::{
     impl_encoding, ArrayDType, ArrayData, ArrayLen, ArrayTrait, Canonical, IntoArrayVariant,
@@ -71,9 +71,12 @@ impl ZigZagArray {
 
 impl ArrayTrait for ZigZagArray {}
 
-impl ArrayVariants for ZigZagArray {
-    fn as_primitive_array(&self) -> Option<&dyn PrimitiveArrayTrait> {
-        Some(self)
+impl VariantsVTable<ZigZagArray> for ZigZagEncoding {
+    fn as_primitive_array<'a>(
+        &self,
+        array: &'a ZigZagArray,
+    ) -> Option<&'a dyn PrimitiveArrayTrait> {
+        Some(array)
     }
 }
 

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

@@ -13,7 +13,7 @@ use vortex_error::{vortex_bail, VortexExpect as _, VortexResult};
 use crate::encoding::ids;
 use crate::stats::StatsSet;
 use crate::validity::{LogicalValidity, Validity, ValidityMetadata, ValidityVTable};
-use crate::variants::{ArrayVariants, BoolArrayTrait};
+use crate::variants::{BoolArrayTrait, VariantsVTable};
 use crate::visitor::{ArrayVisitor, VisitorVTable};
 use crate::{
     impl_encoding, ArrayData, ArrayLen, ArrayTrait, Canonical, IntoArrayData, IntoCanonical,
@@ -178,9 +178,9 @@ impl BoolArray {
 
 impl ArrayTrait for BoolArray {}
 
-impl ArrayVariants for BoolArray {
-    fn as_bool_array(&self) -> Option<&dyn BoolArrayTrait> {
-        Some(self)
+impl VariantsVTable<BoolArray> for BoolEncoding {
+    fn as_bool_array<'a>(&self, array: &'a BoolArray) -> Option<&'a dyn BoolArrayTrait> {
+        Some(array)
     }
 }
 

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

@@ -1,6 +1,6 @@
 use arrow_buffer::{BooleanBufferBuilder, Buffer, MutableBuffer, ScalarBuffer};
 use vortex_dtype::{DType, PType, StructDType};
-use vortex_error::{vortex_bail, vortex_err, ErrString, VortexResult};
+use vortex_error::{vortex_bail, vortex_err, ErrString, VortexExpect, VortexResult};
 
 use crate::array::chunked::ChunkedArray;
 use crate::array::extension::ExtensionArray;
@@ -131,11 +131,11 @@ fn swizzle_struct_chunks(
     let mut field_arrays = Vec::new();
 
     for (field_idx, field_dtype) in struct_dtype.dtypes().iter().enumerate() {
-        let field_chunks = chunks.iter().map(|c| c.with_dyn(|d|
-            d.as_struct_array_unchecked()
+        let field_chunks = chunks.iter().map(|c| c.as_struct_array()
+                .vortex_expect("Chunk was not a StructArray")
                 .field(field_idx)
-                .ok_or_else(|| vortex_err!("All chunks must have same dtype; missing field at index {}, current chunk dtype: {}", field_idx, c.dtype())),
-        )).collect::<VortexResult<Vec<_>>>()?;
+                .ok_or_else(|| vortex_err!("All chunks must have same dtype; missing field at index {}, current chunk dtype: {}", field_idx, c.dtype()))
+        ).collect::<VortexResult<Vec<_>>>()?;
         let field_array = ChunkedArray::try_new(field_chunks, field_dtype.clone())?;
         field_arrays.push(field_array.into_array());
     }