vortex_array/array/
variants.rs

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