DType variant traits (#473)

spiraldb · Jul 17, 2024 · 5f72446 · 5f72446
1 parent 93da8f1
commit 5f72446
Show file tree

Hide file tree

Showing 35 changed files with 705 additions and 43 deletions.
diff --git a/bench-vortex/src/vortex_utils.rs b/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;

diff --git a/encodings/alp/src/array.rs b/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))

diff --git a/encodings/byte_bool/src/lib.rs b/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()

diff --git a/encodings/datetime-parts/src/array.rs b/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(

diff --git a/encodings/dict/src/lib.rs b/encodings/dict/src/lib.rs
@@ -9,3 +9,4 @@ mod compress;
 mod compute;
 mod dict;
 mod stats;
+mod variants;
diff --git a/encodings/dict/src/variants.rs b/encodings/dict/src/variants.rs
@@ -0,0 +1,37 @@
+use vortex::variants::{ArrayVariants, BinaryArrayTrait, PrimitiveArrayTrait, Utf8ArrayTrait};
+use vortex::ArrayDType;
+use vortex_dtype::DType;
+
+use crate::DictArray;
+
+impl ArrayVariants for DictArray {
+    fn as_primitive_array(&self) -> Option<&dyn PrimitiveArrayTrait> {
+        if matches!(self.dtype(), DType::Primitive(..)) {
+            Some(self)
+        } else {
+            None
+        }
+    }
+
+    fn as_utf8_array(&self) -> Option<&dyn Utf8ArrayTrait> {
+        if matches!(self.dtype(), DType::Utf8(..)) {
+            Some(self)
+        } else {
+            None
+        }
+    }
+
+    fn as_binary_array(&self) -> Option<&dyn BinaryArrayTrait> {
+        if matches!(self.dtype(), DType::Binary(..)) {
+            Some(self)
+        } else {
+            None
+        }
+    }
+}
+
+impl PrimitiveArrayTrait for DictArray {}
+
+impl Utf8ArrayTrait for DictArray {}
+
+impl BinaryArrayTrait for DictArray {}
diff --git a/encodings/fastlanes/src/bitpacking/mod.rs b/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;

diff --git a/encodings/fastlanes/src/delta/mod.rs b/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)

diff --git a/encodings/fastlanes/src/for/mod.rs b/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 {}
diff --git a/encodings/roaring/src/boolean/mod.rs b/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?

diff --git a/encodings/roaring/src/integer/mod.rs b/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;
@@ -29,7 +30,7 @@ impl RoaringIntArray {
         let length = bitmap.statistics().cardinality as usize;
         Ok(Self {
             typed: TypedArray::try_from_parts(
-                DType::Bool(NonNullable),
+                DType::Primitive(ptype, NonNullable),
                 length,
                 RoaringIntMetadata { ptype },
                 Some(Buffer::from(bitmap.serialize::<Portable>())),
@@ -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

diff --git a/encodings/runend/src/runend.rs b/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)

diff --git a/encodings/zigzag/src/zigzag.rs b/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))

diff --git a/vortex-array/src/array/bool/mod.rs b/vortex-array/src/array/bool/mod.rs
@@ -5,6 +5,7 @@ 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, Canonical, IntoCanonical};
 
@@ -74,6 +75,14 @@ impl BoolArray {
 
 impl ArrayTrait for BoolArray {}
 
+impl ArrayVariants for BoolArray {
+    fn as_bool_array(&self) -> Option<&dyn BoolArrayTrait> {
+        Some(self)
+    }
+}
+
+impl BoolArrayTrait for BoolArray {}
+
 impl From<BooleanBuffer> for BoolArray {
     fn from(value: BooleanBuffer) -> Self {
         Self::try_new(value, Validity::NonNullable).unwrap()

diff --git a/vortex-array/src/array/chunked/canonical.rs b/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,
 };

diff --git a/vortex-array/src/array/chunked/mod.rs b/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);