vortex_array/array/
variants.rs

1use std::sync::Arc;
2
3use vortex_dtype::DType;
4
5use crate::variants::{
6    BinaryArrayTrait, BoolArrayTrait, ExtensionArrayTrait, ListArrayTrait, NullArrayTrait,
7    PrimitiveArrayTrait, StructArrayTrait, Utf8ArrayTrait,
8};
9use crate::{Array, ArrayImpl};
10
11pub trait ArrayVariants {
12    /// Downcasts the array for null-specific behavior.
13    fn as_null_typed(&self) -> Option<&dyn NullArrayTrait>;
14
15    /// Downcasts the array for bool-specific behavior.
16    fn as_bool_typed(&self) -> Option<&dyn BoolArrayTrait>;
17
18    /// Downcasts the array for primitive-specific behavior.
19    fn as_primitive_typed(&self) -> Option<&dyn PrimitiveArrayTrait>;
20
21    /// Downcasts the array for utf8-specific behavior.
22    fn as_utf8_typed(&self) -> Option<&dyn Utf8ArrayTrait>;
23
24    /// Downcasts the array for binary-specific behavior.
25    fn as_binary_typed(&self) -> Option<&dyn BinaryArrayTrait>;
26
27    /// Downcasts the array for struct-specific behavior.
28    fn as_struct_typed(&self) -> Option<&dyn StructArrayTrait>;
29
30    /// Downcasts the array for list-specific behavior.
31    fn as_list_typed(&self) -> Option<&dyn ListArrayTrait>;
32
33    /// Downcasts the array for extension-specific behavior.
34    fn as_extension_typed(&self) -> Option<&dyn ExtensionArrayTrait>;
35}
36
37impl ArrayVariants for Arc<dyn Array> {
38    fn as_null_typed(&self) -> Option<&dyn NullArrayTrait> {
39        self.as_ref().as_null_typed()
40    }
41
42    fn as_bool_typed(&self) -> Option<&dyn BoolArrayTrait> {
43        self.as_ref().as_bool_typed()
44    }
45
46    fn as_primitive_typed(&self) -> Option<&dyn PrimitiveArrayTrait> {
47        self.as_ref().as_primitive_typed()
48    }
49
50    fn as_utf8_typed(&self) -> Option<&dyn Utf8ArrayTrait> {
51        self.as_ref().as_utf8_typed()
52    }
53
54    fn as_binary_typed(&self) -> Option<&dyn BinaryArrayTrait> {
55        self.as_ref().as_binary_typed()
56    }
57
58    fn as_struct_typed(&self) -> Option<&dyn StructArrayTrait> {
59        self.as_ref().as_struct_typed()
60    }
61
62    fn as_list_typed(&self) -> Option<&dyn ListArrayTrait> {
63        self.as_ref().as_list_typed()
64    }
65
66    fn as_extension_typed(&self) -> Option<&dyn ExtensionArrayTrait> {
67        self.as_ref().as_extension_typed()
68    }
69}
70
71/// Implementation trait for downcasting to type-specific traits.
72pub trait ArrayVariantsImpl {
73    /// Downcasts the array for null-specific behavior.
74    fn _as_null_typed(&self) -> Option<&dyn NullArrayTrait> {
75        None
76    }
77
78    /// Downcasts the array for bool-specific behavior.
79    fn _as_bool_typed(&self) -> Option<&dyn BoolArrayTrait> {
80        None
81    }
82
83    /// Downcasts the array for primitive-specific behavior.
84    fn _as_primitive_typed(&self) -> Option<&dyn PrimitiveArrayTrait> {
85        None
86    }
87
88    /// Downcasts the array for utf8-specific behavior.
89    fn _as_utf8_typed(&self) -> Option<&dyn Utf8ArrayTrait> {
90        None
91    }
92
93    /// Downcasts the array for binary-specific behavior.
94    fn _as_binary_typed(&self) -> Option<&dyn BinaryArrayTrait> {
95        None
96    }
97
98    /// Downcasts the array for struct-specific behavior.
99    fn _as_struct_typed(&self) -> Option<&dyn StructArrayTrait> {
100        None
101    }
102
103    /// Downcasts the array for list-specific behavior.
104    fn _as_list_typed(&self) -> Option<&dyn ListArrayTrait> {
105        None
106    }
107
108    /// Downcasts the array for extension-specific behavior.
109    fn _as_extension_typed(&self) -> Option<&dyn ExtensionArrayTrait> {
110        None
111    }
112}
113
114impl<A: ArrayImpl> ArrayVariants for A {
115    /// Downcasts the array for null-specific behavior.
116    fn as_null_typed(&self) -> Option<&dyn NullArrayTrait> {
117        matches!(self.dtype(), DType::Null)
118            .then(|| ArrayVariantsImpl::_as_null_typed(self))
119            .flatten()
120    }
121
122    /// Downcasts the array for bool-specific behavior.
123    fn as_bool_typed(&self) -> Option<&dyn BoolArrayTrait> {
124        matches!(self.dtype(), DType::Bool(..))
125            .then(|| ArrayVariantsImpl::_as_bool_typed(self))
126            .flatten()
127    }
128
129    /// Downcasts the array for primitive-specific behavior.
130    fn as_primitive_typed(&self) -> Option<&dyn PrimitiveArrayTrait> {
131        matches!(self.dtype(), DType::Primitive(..))
132            .then(|| ArrayVariantsImpl::_as_primitive_typed(self))
133            .flatten()
134    }
135
136    /// Downcasts the array for utf8-specific behavior.
137    fn as_utf8_typed(&self) -> Option<&dyn Utf8ArrayTrait> {
138        matches!(self.dtype(), DType::Utf8(..))
139            .then(|| ArrayVariantsImpl::_as_utf8_typed(self))
140            .flatten()
141    }
142
143    /// Downcasts the array for binary-specific behavior.
144    fn as_binary_typed(&self) -> Option<&dyn BinaryArrayTrait> {
145        matches!(self.dtype(), DType::Binary(..))
146            .then(|| ArrayVariantsImpl::_as_binary_typed(self))
147            .flatten()
148    }
149
150    /// Downcasts the array for struct-specific behavior.
151    fn as_struct_typed(&self) -> Option<&dyn StructArrayTrait> {
152        matches!(self.dtype(), DType::Struct(..))
153            .then(|| ArrayVariantsImpl::_as_struct_typed(self))
154            .flatten()
155    }
156
157    /// Downcasts the array for list-specific behavior.
158    fn as_list_typed(&self) -> Option<&dyn ListArrayTrait> {
159        matches!(self.dtype(), DType::List(..))
160            .then(|| ArrayVariantsImpl::_as_list_typed(self))
161            .flatten()
162    }
163
164    /// Downcasts the array for extension-specific behavior.
165    fn as_extension_typed(&self) -> Option<&dyn ExtensionArrayTrait> {
166        matches!(self.dtype(), DType::Extension(..))
167            .then(|| ArrayVariantsImpl::_as_extension_typed(self))
168            .flatten()
169    }
170}