Struct variant

spiraldb · Jul 17, 2024 · eb8a8d2 · eb8a8d2
1 parent 99c3d93
commit eb8a8d2
Show file tree

Hide file tree

Showing 6 changed files with 169 additions and 33 deletions.
diff --git a/encodings/dict/src/variants.rs b/encodings/dict/src/variants.rs
@@ -1,18 +1,32 @@
 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> {
-        Some(self)
+        if matches!(self.dtype(), DType::Primitive(..)) {
+            Some(self)
+        } else {
+            None
+        }
     }
 
     fn as_utf8_array(&self) -> Option<&dyn Utf8ArrayTrait> {
-        Some(self)
+        if matches!(self.dtype(), DType::Utf8(..)) {
+            Some(self)
+        } else {
+            None
+        }
     }
 
     fn as_binary_array(&self) -> Option<&dyn BinaryArrayTrait> {
-        Some(self)
+        if matches!(self.dtype(), DType::Binary(..)) {
+            Some(self)
+        } else {
+            None
+        }
     }
 }
 

diff --git a/vortex-array/src/array/chunked/variants.rs b/vortex-array/src/array/chunked/variants.rs
@@ -1,3 +1,5 @@
+use vortex_dtype::DType;
+
 use crate::array::chunked::ChunkedArray;
 use crate::variants::{
     ArrayVariants, BinaryArrayTrait, BoolArrayTrait, ExtensionArrayTrait, ListArrayTrait,
@@ -8,35 +10,67 @@ use crate::{Array, ArrayDType, IntoArray};
 /// Chunked arrays support all DTypes
 impl ArrayVariants for ChunkedArray {
     fn as_null_array(&self) -> Option<&dyn NullArrayTrait> {
-        Some(self)
+        if matches!(self.dtype(), DType::Null) {
+            Some(self)
+        } else {
+            None
+        }
     }
 
     fn as_bool_array(&self) -> Option<&dyn BoolArrayTrait> {
-        Some(self)
+        if matches!(self.dtype(), DType::Bool(_)) {
+            Some(self)
+        } else {
+            None
+        }
     }
 
     fn as_primitive_array(&self) -> Option<&dyn PrimitiveArrayTrait> {
-        Some(self)
+        if matches!(self.dtype(), DType::Primitive(..)) {
+            Some(self)
+        } else {
+            None
+        }
     }
 
     fn as_utf8_array(&self) -> Option<&dyn Utf8ArrayTrait> {
-        Some(self)
+        if matches!(self.dtype(), DType::Utf8(_)) {
+            Some(self)
+        } else {
+            None
+        }
     }
 
     fn as_binary_array(&self) -> Option<&dyn BinaryArrayTrait> {
-        Some(self)
+        if matches!(self.dtype(), DType::Binary(_)) {
+            Some(self)
+        } else {
+            None
+        }
     }
 
     fn as_struct_array(&self) -> Option<&dyn StructArrayTrait> {
-        Some(self)
+        if matches!(self.dtype(), DType::Struct(..)) {
+            Some(self)
+        } else {
+            None
+        }
     }
 
     fn as_list_array(&self) -> Option<&dyn ListArrayTrait> {
-        Some(self)
+        if matches!(self.dtype(), DType::List(..)) {
+            Some(self)
+        } else {
+            None
+        }
     }
 
     fn as_extension_array(&self) -> Option<&dyn ExtensionArrayTrait> {
-        Some(self)
+        if matches!(self.dtype(), DType::Extension(..)) {
+            Some(self)
+        } else {
+            None
+        }
     }
 }
 

diff --git a/vortex-array/src/array/constant/variants.rs b/vortex-array/src/array/constant/variants.rs
@@ -1,44 +1,77 @@
+use vortex_dtype::DType;
 use vortex_scalar::StructScalar;
 
 use crate::array::constant::ConstantArray;
 use crate::variants::{
     ArrayVariants, BinaryArrayTrait, BoolArrayTrait, ExtensionArrayTrait, ListArrayTrait,
     NullArrayTrait, PrimitiveArrayTrait, StructArrayTrait, Utf8ArrayTrait,
 };
-use crate::{Array, IntoArray};
+use crate::{Array, ArrayDType, IntoArray};
 
 /// Constant arrays support all DTypes
 impl ArrayVariants for ConstantArray {
     fn as_null_array(&self) -> Option<&dyn NullArrayTrait> {
-        Some(self)
+        if matches!(self.dtype(), DType::Null) {
+            Some(self)
+        } else {
+            None
+        }
     }
 
     fn as_bool_array(&self) -> Option<&dyn BoolArrayTrait> {
-        Some(self)
+        if matches!(self.dtype(), DType::Bool(_)) {
+            Some(self)
+        } else {
+            None
+        }
     }
 
     fn as_primitive_array(&self) -> Option<&dyn PrimitiveArrayTrait> {
-        Some(self)
+        if matches!(self.dtype(), DType::Primitive(..)) {
+            Some(self)
+        } else {
+            None
+        }
     }
 
     fn as_utf8_array(&self) -> Option<&dyn Utf8ArrayTrait> {
-        Some(self)
+        if matches!(self.dtype(), DType::Utf8(_)) {
+            Some(self)
+        } else {
+            None
+        }
     }
 
     fn as_binary_array(&self) -> Option<&dyn BinaryArrayTrait> {
-        Some(self)
+        if matches!(self.dtype(), DType::Binary(_)) {
+            Some(self)
+        } else {
+            None
+        }
     }
 
     fn as_struct_array(&self) -> Option<&dyn StructArrayTrait> {
-        Some(self)
+        if matches!(self.dtype(), DType::Struct(..)) {
+            Some(self)
+        } else {
+            None
+        }
     }
 
     fn as_list_array(&self) -> Option<&dyn ListArrayTrait> {
-        Some(self)
+        if matches!(self.dtype(), DType::List(..)) {
+            Some(self)
+        } else {
+            None
+        }
     }
 
     fn as_extension_array(&self) -> Option<&dyn ExtensionArrayTrait> {
-        Some(self)
+        if matches!(self.dtype(), DType::Extension(..)) {
+            Some(self)
+        } else {
+            None
+        }
     }
 }
 

diff --git a/vortex-array/src/array/sparse/variants.rs b/vortex-array/src/array/sparse/variants.rs
@@ -1,44 +1,77 @@
+use vortex_dtype::DType;
 use vortex_scalar::StructScalar;
 
 use crate::array::sparse::SparseArray;
 use crate::variants::{
     ArrayVariants, BinaryArrayTrait, BoolArrayTrait, ExtensionArrayTrait, ListArrayTrait,
     NullArrayTrait, PrimitiveArrayTrait, StructArrayTrait, Utf8ArrayTrait,
 };
-use crate::{Array, IntoArray};
+use crate::{Array, ArrayDType, IntoArray};
 
 /// Sparse arrays support all DTypes
 impl ArrayVariants for SparseArray {
     fn as_null_array(&self) -> Option<&dyn NullArrayTrait> {
-        Some(self)
+        if matches!(self.dtype(), DType::Null) {
+            Some(self)
+        } else {
+            None
+        }
     }
 
     fn as_bool_array(&self) -> Option<&dyn BoolArrayTrait> {
-        Some(self)
+        if matches!(self.dtype(), DType::Bool(_)) {
+            Some(self)
+        } else {
+            None
+        }
     }
 
     fn as_primitive_array(&self) -> Option<&dyn PrimitiveArrayTrait> {
-        Some(self)
+        if matches!(self.dtype(), DType::Primitive(..)) {
+            Some(self)
+        } else {
+            None
+        }
     }
 
     fn as_utf8_array(&self) -> Option<&dyn Utf8ArrayTrait> {
-        Some(self)
+        if matches!(self.dtype(), DType::Utf8(_)) {
+            Some(self)
+        } else {
+            None
+        }
     }
 
     fn as_binary_array(&self) -> Option<&dyn BinaryArrayTrait> {
-        Some(self)
+        if matches!(self.dtype(), DType::Binary(_)) {
+            Some(self)
+        } else {
+            None
+        }
     }
 
     fn as_struct_array(&self) -> Option<&dyn StructArrayTrait> {
-        Some(self)
+        if matches!(self.dtype(), DType::Struct(..)) {
+            Some(self)
+        } else {
+            None
+        }
     }
 
     fn as_list_array(&self) -> Option<&dyn ListArrayTrait> {
-        Some(self)
+        if matches!(self.dtype(), DType::List(..)) {
+            Some(self)
+        } else {
+            None
+        }
     }
 
     fn as_extension_array(&self) -> Option<&dyn ExtensionArrayTrait> {
-        Some(self)
+        if matches!(self.dtype(), DType::Extension(..)) {
+            Some(self)
+        } else {
+            None
+        }
     }
 }
 

diff --git a/vortex-array/src/array/varbin/variants.rs b/vortex-array/src/array/varbin/variants.rs
@@ -1,13 +1,24 @@
+use vortex_dtype::DType;
+
 use crate::array::varbin::VarBinArray;
 use crate::variants::{ArrayVariants, BinaryArrayTrait, Utf8ArrayTrait};
+use crate::ArrayDType;
 
 impl ArrayVariants for VarBinArray {
     fn as_utf8_array(&self) -> Option<&dyn Utf8ArrayTrait> {
-        Some(self)
+        if matches!(self.dtype(), DType::Utf8(..)) {
+            Some(self)
+        } else {
+            None
+        }
     }
 
     fn as_binary_array(&self) -> Option<&dyn BinaryArrayTrait> {
-        Some(self)
+        if matches!(self.dtype(), DType::Binary(..)) {
+            Some(self)
+        } else {
+            None
+        }
     }
 }
 

diff --git a/vortex-array/src/array/varbinview/variants.rs b/vortex-array/src/array/varbinview/variants.rs
@@ -1,13 +1,24 @@
+use vortex_dtype::DType;
+
 use crate::array::varbinview::VarBinViewArray;
 use crate::variants::{ArrayVariants, BinaryArrayTrait, Utf8ArrayTrait};
+use crate::ArrayDType;
 
 impl ArrayVariants for VarBinViewArray {
     fn as_utf8_array(&self) -> Option<&dyn Utf8ArrayTrait> {
-        Some(self)
+        if matches!(self.dtype(), DType::Utf8(..)) {
+            Some(self)
+        } else {
+            None
+        }
     }
 
     fn as_binary_array(&self) -> Option<&dyn BinaryArrayTrait> {
-        Some(self)
+        if matches!(self.dtype(), DType::Binary(..)) {
+            Some(self)
+        } else {
+            None
+        }
     }
 }