vortex_array/
variants.rs

1//! This module defines array traits for each Vortex DType.
2//!
3//! When callers only want to make assumptions about the DType, and not about any specific
4//! encoding, they can use these traits to write encoding-agnostic code.
5
6use std::sync::Arc;
7
8use vortex_dtype::{DType, ExtDType, FieldName, FieldNames, PType};
9use vortex_error::{VortexExpect, VortexResult, vortex_err, vortex_panic};
10
11use crate::compute::sum;
12use crate::{Array, ArrayRef};
13
14pub trait NullArrayTrait {}
15
16pub trait BoolArrayTrait: Array {}
17
18impl dyn BoolArrayTrait + '_ {
19    pub fn true_count(&self) -> VortexResult<usize> {
20        let true_count = sum(self)?;
21        Ok(true_count
22            .as_primitive()
23            .as_::<usize>()
24            .vortex_expect("true count should never overflow usize")
25            .vortex_expect("true count should never be null"))
26    }
27}
28
29pub trait PrimitiveArrayTrait: Array {
30    /// The logical primitive type of the array.
31    ///
32    /// This is a type that can safely be converted into a `NativePType` for use in
33    /// `maybe_null_slice` or `into_maybe_null_slice`.
34    fn ptype(&self) -> PType {
35        if let DType::Primitive(ptype, ..) = self.dtype() {
36            *ptype
37        } else {
38            vortex_panic!("array must have primitive data type");
39        }
40    }
41}
42
43pub trait Utf8ArrayTrait {}
44
45pub trait BinaryArrayTrait {}
46
47pub trait StructArrayTrait: Array {
48    fn names(&self) -> &FieldNames {
49        let DType::Struct(st, _) = self.dtype() else {
50            unreachable!()
51        };
52        st.names()
53    }
54
55    fn dtypes(&self) -> Vec<DType> {
56        let DType::Struct(st, _) = self.dtype() else {
57            unreachable!()
58        };
59        st.fields().collect()
60    }
61
62    fn nfields(&self) -> usize {
63        self.names().len()
64    }
65
66    /// Return a field's array by index, ignoring struct nullability
67    fn maybe_null_field_by_idx(&self, idx: usize) -> VortexResult<ArrayRef>;
68
69    /// Return a field's array by name, ignoring struct nullability
70    fn maybe_null_field_by_name(&self, name: &str) -> VortexResult<ArrayRef> {
71        let field_idx = self
72            .names()
73            .iter()
74            .position(|field_name| field_name.as_ref() == name)
75            .ok_or_else(|| vortex_err!("Field not found: {}", name))?;
76        self.maybe_null_field_by_idx(field_idx)
77    }
78
79    fn project(&self, projection: &[FieldName]) -> VortexResult<ArrayRef>;
80}
81
82impl dyn StructArrayTrait + '_ {
83    pub fn fields(&self) -> impl Iterator<Item = ArrayRef> + '_ {
84        (0..self.nfields()).map(|i| {
85            self.maybe_null_field_by_idx(i)
86                .vortex_expect("never out of bounds")
87        })
88    }
89}
90
91pub trait ListArrayTrait {}
92
93pub trait ExtensionArrayTrait: Array {
94    /// Returns the extension logical [`DType`].
95    fn ext_dtype(&self) -> &Arc<ExtDType> {
96        let DType::Extension(ext_dtype) = self.dtype() else {
97            vortex_panic!("Expected ExtDType")
98        };
99        ext_dtype
100    }
101
102    /// Returns the underlying [`ArrayRef`], without the [`ExtDType`].
103    fn storage_data(&self) -> ArrayRef;
104}
vortex_array/variants.rs

vortex_array/
variants.rs
