diff --git a/encodings/dict/src/variants.rs b/encodings/dict/src/variants.rs index 31abbb7bfc..843086b18f 100644 --- 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 index 29b830d132..02dfff90ef 100644 --- 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 index e04f4e6761..28c9a22509 100644 --- a/vortex-array/src/array/constant/variants.rs +++ b/vortex-array/src/array/constant/variants.rs @@ -1,3 +1,4 @@ +use vortex_dtype::DType; use vortex_scalar::StructScalar; use crate::array::constant::ConstantArray; @@ -5,40 +6,72 @@ 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 index fa3fd3d0fa..298c3cefa3 100644 --- a/vortex-array/src/array/sparse/variants.rs +++ b/vortex-array/src/array/sparse/variants.rs @@ -1,3 +1,4 @@ +use vortex_dtype::DType; use vortex_scalar::StructScalar; use crate::array::sparse::SparseArray; @@ -5,40 +6,72 @@ 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 index 85fcad30eb..b387e0978a 100644 --- 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 index 22e74e1d49..0a387e8966 100644 --- 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 + } } }