vortex_array/array/
implementation.rs

1use std::any::Any;
2use std::fmt::Debug;
3use std::sync::Arc;
4
5use vortex_dtype::DType;
6use vortex_error::{VortexResult, vortex_bail};
7use vortex_mask::Mask;
8
9use crate::array::canonical::ArrayCanonicalImpl;
10use crate::array::validity::ArrayValidityImpl;
11use crate::array::visitor::ArrayVisitorImpl;
12use crate::builders::ArrayBuilder;
13use crate::stats::{Precision, Stat, StatsSetRef};
14use crate::vtable::VTableRef;
15use crate::{
16    Array, ArrayRef, ArrayStatisticsImpl, ArrayVariantsImpl, Canonical, Encoding, EncodingId,
17};
18
19/// A trait used to encapsulate common implementation behaviour for a Vortex [`Array`].
20pub trait ArrayImpl:
21    'static
22    + Send
23    + Sync
24    + Debug
25    + Clone
26    + ArrayCanonicalImpl
27    + ArrayStatisticsImpl
28    + ArrayValidityImpl
29    + ArrayVariantsImpl
30    + ArrayVisitorImpl<<Self::Encoding as Encoding>::Metadata>
31{
32    type Encoding: Encoding;
33
34    fn _len(&self) -> usize;
35    fn _dtype(&self) -> &DType;
36    fn _vtable(&self) -> VTableRef;
37}
38
39impl<A: ArrayImpl + 'static> Array for A {
40    fn as_any(&self) -> &dyn Any {
41        self
42    }
43
44    fn as_any_arc(self: Arc<Self>) -> Arc<dyn Any + Send + Sync> {
45        self
46    }
47
48    fn to_array(&self) -> ArrayRef {
49        Arc::new(self.clone())
50    }
51
52    fn into_array(self) -> ArrayRef
53    where
54        Self: Sized,
55    {
56        Arc::new(self)
57    }
58
59    fn len(&self) -> usize {
60        ArrayImpl::_len(self)
61    }
62
63    fn dtype(&self) -> &DType {
64        ArrayImpl::_dtype(self)
65    }
66
67    fn encoding(&self) -> EncodingId {
68        self.vtable().id()
69    }
70
71    fn vtable(&self) -> VTableRef {
72        ArrayImpl::_vtable(self)
73    }
74
75    /// Returns whether the item at `index` is valid.
76    fn is_valid(&self, index: usize) -> VortexResult<bool> {
77        if index >= self.len() {
78            vortex_bail!("Index out of bounds: {} >= {}", index, self.len());
79        }
80        ArrayValidityImpl::_is_valid(self, index)
81    }
82
83    /// Returns whether the item at `index` is invalid.
84    fn is_invalid(&self, index: usize) -> VortexResult<bool> {
85        self.is_valid(index).map(|valid| !valid)
86    }
87
88    /// Returns whether all items in the array are valid.
89    ///
90    /// This is usually cheaper than computing a precise `valid_count`.
91    fn all_valid(&self) -> VortexResult<bool> {
92        ArrayValidityImpl::_all_valid(self)
93    }
94
95    /// Returns whether the array is all invalid.
96    ///
97    /// This is usually cheaper than computing a precise `invalid_count`.
98    fn all_invalid(&self) -> VortexResult<bool> {
99        ArrayValidityImpl::_all_invalid(self)
100    }
101
102    /// Returns the number of valid elements in the array.
103    fn valid_count(&self) -> VortexResult<usize> {
104        if let Some(Precision::Exact(invalid_count)) =
105            self.statistics().get_as::<usize>(Stat::NullCount)
106        {
107            return Ok(self.len() - invalid_count);
108        }
109
110        let count = ArrayValidityImpl::_valid_count(self)?;
111        assert!(count <= self.len(), "Valid count exceeds array length");
112
113        self.statistics()
114            .set(Stat::NullCount, Precision::exact(self.len() - count));
115
116        Ok(count)
117    }
118
119    /// Returns the number of invalid elements in the array.
120    fn invalid_count(&self) -> VortexResult<usize> {
121        if let Some(Precision::Exact(invalid_count)) =
122            self.statistics().get_as::<usize>(Stat::NullCount)
123        {
124            return Ok(invalid_count);
125        }
126
127        let count = ArrayValidityImpl::_invalid_count(self)?;
128        assert!(count <= self.len(), "Invalid count exceeds array length");
129
130        self.statistics()
131            .set(Stat::NullCount, Precision::exact(count));
132
133        Ok(count)
134    }
135
136    /// Returns the canonical validity mask for the array.
137    fn validity_mask(&self) -> VortexResult<Mask> {
138        let mask = ArrayValidityImpl::_validity_mask(self)?;
139        assert_eq!(mask.len(), self.len(), "Validity mask length mismatch");
140        Ok(mask)
141    }
142
143    /// Returns the canonical representation of the array.
144    fn to_canonical(&self) -> VortexResult<Canonical> {
145        let canonical = ArrayCanonicalImpl::_to_canonical(self)?;
146        assert_eq!(
147            canonical.as_ref().len(),
148            self.len(),
149            "Canonical length mismatch"
150        );
151        assert_eq!(
152            canonical.as_ref().dtype(),
153            self.dtype(),
154            "Canonical dtype mismatch"
155        );
156        canonical.as_ref().statistics().inherit(self.statistics());
157        Ok(canonical)
158    }
159
160    /// Writes the array into the canonical builder.
161    ///
162    /// The [`DType`] of the builder must match that of the array.
163    fn append_to_builder(&self, builder: &mut dyn ArrayBuilder) -> VortexResult<()> {
164        // TODO(ngates): add dtype function to ArrayBuilder
165        // if builder.dtype() != self.dtype() {
166        //     vortex_bail!(
167        //         "Builder dtype mismatch: expected {:?}, got {:?}",
168        //         self.dtype(),
169        //         builder.dtype()
170        //     );
171        // }
172        let len = builder.len();
173        ArrayCanonicalImpl::_append_to_builder(self, builder)?;
174        assert_eq!(
175            len + self.len(),
176            builder.len(),
177            "Builder length mismatch after writing array"
178        );
179        Ok(())
180    }
181
182    fn statistics(&self) -> StatsSetRef<'_> {
183        self._stats_ref()
184    }
185}