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;
6
7use vortex_error::VortexExpect;
8use vortex_error::VortexResult;
9use vortex_error::vortex_bail;
10use vortex_error::vortex_panic;
11use vortex_mask::Mask;
12
13use crate::ArrayRef;
14use crate::ExecutionCtx;
15use crate::LEGACY_SESSION;
16use crate::VortexSessionExecute;
17use crate::aggregate_fn::fns::sum::sum;
18use crate::arrays::BoolArray;
19use crate::arrays::bool::BoolArrayExt;
20use crate::builtins::ArrayBuiltins;
21use crate::dtype::DType;
22use crate::dtype::FieldNames;
23use crate::dtype::PType;
24use crate::dtype::extension::ExtDTypeRef;
25use crate::scalar::PValue;
26use crate::scalar::Scalar;
27use crate::search_sorted::IndexOrd;
28
29impl ArrayRef {
30    /// Downcasts the array for null-specific behavior.
31    pub fn as_null_typed(&self) -> NullTyped<'_> {
32        matches!(self.dtype(), DType::Null)
33            .then(|| NullTyped(self))
34            .vortex_expect("Array does not have DType::Null")
35    }
36
37    /// Downcasts the array for bool-specific behavior.
38    pub fn as_bool_typed(&self) -> BoolTyped<'_> {
39        matches!(self.dtype(), DType::Bool(..))
40            .then(|| BoolTyped(self))
41            .vortex_expect("Array does not have DType::Bool")
42    }
43
44    /// Downcasts the array for primitive-specific behavior.
45    pub fn as_primitive_typed(&self) -> PrimitiveTyped<'_> {
46        matches!(self.dtype(), DType::Primitive(..))
47            .then(|| PrimitiveTyped(self))
48            .vortex_expect("Array does not have DType::Primitive")
49    }
50
51    /// Downcasts the array for decimal-specific behavior.
52    pub fn as_decimal_typed(&self) -> DecimalTyped<'_> {
53        matches!(self.dtype(), DType::Decimal(..))
54            .then(|| DecimalTyped(self))
55            .vortex_expect("Array does not have DType::Decimal")
56    }
57
58    /// Downcasts the array for utf8-specific behavior.
59    pub fn as_utf8_typed(&self) -> Utf8Typed<'_> {
60        matches!(self.dtype(), DType::Utf8(..))
61            .then(|| Utf8Typed(self))
62            .vortex_expect("Array does not have DType::Utf8")
63    }
64
65    /// Downcasts the array for binary-specific behavior.
66    pub fn as_binary_typed(&self) -> BinaryTyped<'_> {
67        matches!(self.dtype(), DType::Binary(..))
68            .then(|| BinaryTyped(self))
69            .vortex_expect("Array does not have DType::Binary")
70    }
71
72    /// Downcasts the array for struct-specific behavior.
73    pub fn as_struct_typed(&self) -> StructTyped<'_> {
74        matches!(self.dtype(), DType::Struct(..))
75            .then(|| StructTyped(self))
76            .vortex_expect("Array does not have DType::Struct")
77    }
78
79    /// Downcasts the array for list-specific behavior.
80    pub fn as_list_typed(&self) -> ListTyped<'_> {
81        matches!(self.dtype(), DType::List(..))
82            .then(|| ListTyped(self))
83            .vortex_expect("Array does not have DType::List")
84    }
85
86    /// Downcasts the array for extension-specific behavior.
87    pub fn as_extension_typed(&self) -> ExtensionTyped<'_> {
88        matches!(self.dtype(), DType::Extension(..))
89            .then(|| ExtensionTyped(self))
90            .vortex_expect("Array does not have DType::Extension")
91    }
92
93    pub fn try_to_mask_fill_null_false(&self, ctx: &mut ExecutionCtx) -> VortexResult<Mask> {
94        if !matches!(self.dtype(), DType::Bool(_)) {
95            vortex_bail!("mask must be bool array, has dtype {}", self.dtype());
96        }
97
98        // Convert nulls to false first in case this can be done cheaply by the encoding.
99        let array = self
100            .clone()
101            .fill_null(Scalar::bool(false, self.dtype().nullability()))?;
102
103        Ok(array.execute::<BoolArray>(ctx)?.to_mask_fill_null_false())
104    }
105}
106
107#[expect(dead_code)]
108pub struct NullTyped<'a>(&'a ArrayRef);
109
110pub struct BoolTyped<'a>(&'a ArrayRef);
111
112impl BoolTyped<'_> {
113    pub fn true_count(&self) -> VortexResult<usize> {
114        let mut ctx = LEGACY_SESSION.create_execution_ctx();
115        let true_count = sum(self.0, &mut ctx)?;
116        Ok(true_count
117            .as_primitive()
118            .as_::<usize>()
119            .vortex_expect("true count should never be null"))
120    }
121}
122
123pub struct PrimitiveTyped<'a>(&'a ArrayRef);
124
125impl PrimitiveTyped<'_> {
126    pub fn ptype(&self) -> PType {
127        let DType::Primitive(ptype, _) = self.0.dtype() else {
128            vortex_panic!("Expected Primitive DType")
129        };
130        *ptype
131    }
132
133    /// Return the primitive value at the given index.
134    pub fn value(&self, idx: usize) -> VortexResult<Option<PValue>> {
135        self.0
136            .is_valid(idx)?
137            .then(|| self.value_unchecked(idx))
138            .transpose()
139    }
140
141    /// Return the primitive value at the given index, ignoring nullability.
142    pub fn value_unchecked(&self, idx: usize) -> VortexResult<PValue> {
143        Ok(self
144            .0
145            .scalar_at(idx)?
146            .as_primitive()
147            .pvalue()
148            .unwrap_or_else(|| PValue::zero(&self.ptype())))
149    }
150}
151
152impl IndexOrd<Option<PValue>> for PrimitiveTyped<'_> {
153    fn index_cmp(&self, idx: usize, elem: &Option<PValue>) -> VortexResult<Option<Ordering>> {
154        let value = self.value(idx)?;
155        Ok(value.partial_cmp(elem))
156    }
157
158    fn index_len(&self) -> usize {
159        self.0.len()
160    }
161}
162
163// TODO(ngates): add generics to the `value` function and implement this over T.
164impl IndexOrd<PValue> for PrimitiveTyped<'_> {
165    fn index_cmp(&self, idx: usize, elem: &PValue) -> VortexResult<Option<Ordering>> {
166        assert!(self.0.all_valid()?);
167        let value = self.value_unchecked(idx)?;
168        Ok(value.partial_cmp(elem))
169    }
170
171    fn index_len(&self) -> usize {
172        self.0.len()
173    }
174}
175
176#[expect(dead_code)]
177pub struct Utf8Typed<'a>(&'a ArrayRef);
178
179#[expect(dead_code)]
180pub struct BinaryTyped<'a>(&'a ArrayRef);
181
182#[expect(dead_code)]
183pub struct DecimalTyped<'a>(&'a ArrayRef);
184
185pub struct StructTyped<'a>(&'a ArrayRef);
186
187impl StructTyped<'_> {
188    pub fn names(&self) -> &FieldNames {
189        let DType::Struct(st, _) = self.0.dtype() else {
190            unreachable!()
191        };
192        st.names()
193    }
194
195    pub fn dtypes(&self) -> Vec<DType> {
196        let DType::Struct(st, _) = self.0.dtype() else {
197            unreachable!()
198        };
199        st.fields().collect()
200    }
201
202    pub fn nfields(&self) -> usize {
203        self.names().len()
204    }
205}
206
207#[expect(dead_code)]
208pub struct ListTyped<'a>(&'a ArrayRef);
209
210pub struct ExtensionTyped<'a>(&'a ArrayRef);
211
212impl ExtensionTyped<'_> {
213    /// Returns the extension logical [`DType`].
214    pub fn ext_dtype(&self) -> &ExtDTypeRef {
215        let DType::Extension(ext_dtype) = self.0.dtype() else {
216            vortex_panic!("Expected ExtDType")
217        };
218        ext_dtype
219    }
220}
vortex_array/variants.rs

vortex_array/
variants.rs
