vortex_array/
variants.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4//! This module defines extension functionality specific to each Vortex DType.
5use std::cmp::Ordering;
6use std::sync::Arc;
7
8use vortex_dtype::{DType, ExtDType, FieldNames, PType};
9use vortex_error::{VortexExpect, VortexResult, vortex_panic};
10use vortex_scalar::PValue;
11
12use crate::Array;
13use crate::compute::sum;
14use crate::search_sorted::IndexOrd;
15
16impl dyn Array + '_ {
17    /// Downcasts the array for null-specific behavior.
18    pub fn as_null_typed(&self) -> NullTyped<'_> {
19        matches!(self.dtype(), DType::Null)
20            .then(|| NullTyped(self))
21            .vortex_expect("Array does not have DType::Null")
22    }
23
24    /// Downcasts the array for bool-specific behavior.
25    pub fn as_bool_typed(&self) -> BoolTyped<'_> {
26        matches!(self.dtype(), DType::Bool(..))
27            .then(|| BoolTyped(self))
28            .vortex_expect("Array does not have DType::Bool")
29    }
30
31    /// Downcasts the array for primitive-specific behavior.
32    pub fn as_primitive_typed(&self) -> PrimitiveTyped<'_> {
33        matches!(self.dtype(), DType::Primitive(..))
34            .then(|| PrimitiveTyped(self))
35            .vortex_expect("Array does not have DType::Primitive")
36    }
37
38    /// Downcasts the array for decimal-specific behavior.
39    pub fn as_decimal_typed(&self) -> DecimalTyped<'_> {
40        matches!(self.dtype(), DType::Decimal(..))
41            .then(|| DecimalTyped(self))
42            .vortex_expect("Array does not have DType::Decimal")
43    }
44
45    /// Downcasts the array for utf8-specific behavior.
46    pub fn as_utf8_typed(&self) -> Utf8Typed<'_> {
47        matches!(self.dtype(), DType::Utf8(..))
48            .then(|| Utf8Typed(self))
49            .vortex_expect("Array does not have DType::Utf8")
50    }
51
52    /// Downcasts the array for binary-specific behavior.
53    pub fn as_binary_typed(&self) -> BinaryTyped<'_> {
54        matches!(self.dtype(), DType::Binary(..))
55            .then(|| BinaryTyped(self))
56            .vortex_expect("Array does not have DType::Binary")
57    }
58
59    /// Downcasts the array for struct-specific behavior.
60    pub fn as_struct_typed(&self) -> StructTyped<'_> {
61        matches!(self.dtype(), DType::Struct(..))
62            .then(|| StructTyped(self))
63            .vortex_expect("Array does not have DType::Struct")
64    }
65
66    /// Downcasts the array for list-specific behavior.
67    pub fn as_list_typed(&self) -> ListTyped<'_> {
68        matches!(self.dtype(), DType::List(..))
69            .then(|| ListTyped(self))
70            .vortex_expect("Array does not have DType::List")
71    }
72
73    /// Downcasts the array for extension-specific behavior.
74    pub fn as_extension_typed(&self) -> ExtensionTyped<'_> {
75        matches!(self.dtype(), DType::Extension(..))
76            .then(|| ExtensionTyped(self))
77            .vortex_expect("Array does not have DType::Extension")
78    }
79}
80
81#[allow(dead_code)]
82pub struct NullTyped<'a>(&'a dyn Array);
83
84pub struct BoolTyped<'a>(&'a dyn Array);
85
86impl BoolTyped<'_> {
87    pub fn true_count(&self) -> VortexResult<usize> {
88        let true_count = sum(self.0)?;
89        Ok(true_count
90            .as_primitive()
91            .as_::<usize>()
92            .vortex_expect("true count should never be null"))
93    }
94}
95
96pub struct PrimitiveTyped<'a>(&'a dyn Array);
97
98impl PrimitiveTyped<'_> {
99    pub fn ptype(&self) -> PType {
100        let DType::Primitive(ptype, _) = self.0.dtype() else {
101            vortex_panic!("Expected Primitive DType")
102        };
103        *ptype
104    }
105
106    /// Return the primitive value at the given index.
107    pub fn value(&self, idx: usize) -> Option<PValue> {
108        self.0.is_valid(idx).then(|| self.value_unchecked(idx))
109    }
110
111    /// Return the primitive value at the given index, ignoring nullability.
112    pub fn value_unchecked(&self, idx: usize) -> PValue {
113        self.0
114            .scalar_at(idx)
115            .as_primitive()
116            .pvalue()
117            .unwrap_or_else(|| PValue::zero(self.ptype()))
118    }
119}
120
121impl IndexOrd<Option<PValue>> for PrimitiveTyped<'_> {
122    fn index_cmp(&self, idx: usize, elem: &Option<PValue>) -> Option<Ordering> {
123        self.value(idx).partial_cmp(elem)
124    }
125
126    fn index_len(&self) -> usize {
127        self.0.len()
128    }
129}
130
131// TODO(ngates): add generics to the `value` function and implement this over T.
132impl IndexOrd<PValue> for PrimitiveTyped<'_> {
133    fn index_cmp(&self, idx: usize, elem: &PValue) -> Option<Ordering> {
134        assert!(self.0.all_valid());
135        self.value_unchecked(idx).partial_cmp(elem)
136    }
137
138    fn index_len(&self) -> usize {
139        self.0.len()
140    }
141}
142
143#[allow(dead_code)]
144pub struct Utf8Typed<'a>(&'a dyn Array);
145
146#[allow(dead_code)]
147pub struct BinaryTyped<'a>(&'a dyn Array);
148
149#[allow(dead_code)]
150pub struct DecimalTyped<'a>(&'a dyn Array);
151
152pub struct StructTyped<'a>(&'a dyn Array);
153
154impl StructTyped<'_> {
155    pub fn names(&self) -> &FieldNames {
156        let DType::Struct(st, _) = self.0.dtype() else {
157            unreachable!()
158        };
159        st.names()
160    }
161
162    pub fn dtypes(&self) -> Vec<DType> {
163        let DType::Struct(st, _) = self.0.dtype() else {
164            unreachable!()
165        };
166        st.fields().collect()
167    }
168
169    pub fn nfields(&self) -> usize {
170        self.names().len()
171    }
172}
173
174#[allow(dead_code)]
175pub struct ListTyped<'a>(&'a dyn Array);
176
177pub struct ExtensionTyped<'a>(&'a dyn Array);
178
179impl ExtensionTyped<'_> {
180    /// Returns the extension logical [`DType`].
181    pub fn ext_dtype(&self) -> &Arc<ExtDType> {
182        let DType::Extension(ext_dtype) = self.0.dtype() else {
183            vortex_panic!("Expected ExtDType")
184        };
185        ext_dtype
186    }
187}
vortex_array/variants.rs

vortex_array/
variants.rs
