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 DecimalArrayTrait: Array {}
48
49pub trait StructArrayTrait: Array {
50    fn names(&self) -> &FieldNames {
51        let DType::Struct(st, _) = self.dtype() else {
52            unreachable!()
53        };
54        st.names()
55    }
56
57    fn dtypes(&self) -> Vec<DType> {
58        let DType::Struct(st, _) = self.dtype() else {
59            unreachable!()
60        };
61        st.fields().collect()
62    }
63
64    fn nfields(&self) -> usize {
65        self.names().len()
66    }
67
68    /// Return a field's array by index, ignoring struct nullability
69    fn maybe_null_field_by_idx(&self, idx: usize) -> VortexResult<ArrayRef>;
70
71    /// Return a field's array by name, ignoring struct nullability
72    fn maybe_null_field_by_name(&self, name: &str) -> VortexResult<ArrayRef> {
73        let field_idx = self
74            .names()
75            .iter()
76            .position(|field_name| field_name.as_ref() == name)
77            .ok_or_else(|| vortex_err!("Field not found: {}", name))?;
78        self.maybe_null_field_by_idx(field_idx)
79    }
80
81    fn project(&self, projection: &[FieldName]) -> VortexResult<ArrayRef>;
82}
83
84impl dyn StructArrayTrait + '_ {
85    pub fn fields(&self) -> impl Iterator<Item = ArrayRef> + '_ {
86        (0..self.nfields()).map(|i| {
87            self.maybe_null_field_by_idx(i)
88                .vortex_expect("never out of bounds")
89        })
90    }
91}
92
93pub trait ListArrayTrait {}
94
95pub trait ExtensionArrayTrait: Array {
96    /// Returns the extension logical [`DType`].
97    fn ext_dtype(&self) -> &Arc<ExtDType> {
98        let DType::Extension(ext_dtype) = self.dtype() else {
99            vortex_panic!("Expected ExtDType")
100        };
101        ext_dtype
102    }
103
104    /// Returns the underlying [`ArrayRef`], without the [`ExtDType`].
105    fn storage_data(&self) -> ArrayRef;
106}
vortex_array/variants.rs

vortex_array/
variants.rs
