vortex_runend/
array.rs

1use std::fmt::Debug;
2
3use vortex_array::arrays::PrimitiveArray;
4use vortex_array::compute::{
5    SearchSortedSide, scalar_at, search_sorted_usize, search_sorted_usize_many,
6};
7use vortex_array::stats::{ArrayStats, StatsSetRef};
8use vortex_array::variants::{BoolArrayTrait, PrimitiveArrayTrait};
9use vortex_array::vtable::VTableRef;
10use vortex_array::{
11    Array, ArrayCanonicalImpl, ArrayImpl, ArrayRef, ArrayStatisticsImpl, ArrayValidityImpl,
12    ArrayVariantsImpl, Canonical, Encoding, IntoArray, SerdeMetadata, ToCanonical,
13    try_from_array_ref,
14};
15use vortex_buffer::Buffer;
16use vortex_dtype::DType;
17use vortex_error::{VortexExpect as _, VortexResult, vortex_bail};
18use vortex_mask::Mask;
19
20use crate::compress::{runend_decode_bools, runend_decode_primitive, runend_encode};
21use crate::serde::RunEndMetadata;
22
23#[derive(Clone, Debug)]
24pub struct RunEndArray {
25    ends: ArrayRef,
26    values: ArrayRef,
27    offset: usize,
28    length: usize,
29    stats_set: ArrayStats,
30}
31
32try_from_array_ref!(RunEndArray);
33
34pub struct RunEndEncoding;
35impl Encoding for RunEndEncoding {
36    type Array = RunEndArray;
37    type Metadata = SerdeMetadata<RunEndMetadata>;
38}
39
40impl RunEndArray {
41    pub fn try_new(ends: ArrayRef, values: ArrayRef) -> VortexResult<Self> {
42        let length = if ends.is_empty() {
43            0
44        } else {
45            scalar_at(&ends, ends.len() - 1)?.as_ref().try_into()?
46        };
47        Self::with_offset_and_length(ends, values, 0, length)
48    }
49
50    pub(crate) fn with_offset_and_length(
51        ends: ArrayRef,
52        values: ArrayRef,
53        offset: usize,
54        length: usize,
55    ) -> VortexResult<Self> {
56        if !matches!(values.dtype(), &DType::Bool(_) | &DType::Primitive(_, _)) {
57            vortex_bail!(
58                "RunEnd array can only have Bool or Primitive values, {} given",
59                values.dtype()
60            );
61        }
62
63        if offset != 0 {
64            let first_run_end: usize = scalar_at(&ends, 0)?.as_ref().try_into()?;
65            if first_run_end <= offset {
66                vortex_bail!("First run end {first_run_end} must be bigger than offset {offset}");
67            }
68        }
69
70        if !ends.dtype().is_unsigned_int() || ends.dtype().is_nullable() {
71            vortex_bail!(MismatchedTypes: "non-nullable unsigned int", ends.dtype());
72        }
73        if !ends.statistics().compute_is_strict_sorted().unwrap_or(true) {
74            vortex_bail!("Ends array must be strictly sorted");
75        }
76
77        Ok(Self {
78            ends,
79            values,
80            offset,
81            length,
82            stats_set: Default::default(),
83        })
84    }
85
86    /// Convert the given logical index to an index into the `values` array
87    pub fn find_physical_index(&self, index: usize) -> VortexResult<usize> {
88        search_sorted_usize(self.ends(), index + self.offset(), SearchSortedSide::Right)
89            .map(|s| s.to_ends_index(self.ends().len()))
90    }
91
92    /// Convert a batch of logical indices into an index for the values. Expects indices to be adjusted by offset unlike
93    /// [Self::find_physical_index]
94    ///
95    /// See: [find_physical_index][Self::find_physical_index].
96    pub fn find_physical_indices(&self, indices: &[usize]) -> VortexResult<Buffer<u64>> {
97        search_sorted_usize_many(self.ends(), indices, SearchSortedSide::Right).map(|results| {
98            results
99                .into_iter()
100                .map(|result| result.to_ends_index(self.ends().len()) as u64)
101                .collect()
102        })
103    }
104
105    /// Run the array through run-end encoding.
106    pub fn encode(array: ArrayRef) -> VortexResult<Self> {
107        if let Ok(parray) = PrimitiveArray::try_from(array) {
108            let (ends, values) = runend_encode(&parray)?;
109            Self::try_new(ends.into_array(), values)
110        } else {
111            vortex_bail!("REE can only encode primitive arrays")
112        }
113    }
114
115    /// The offset that the `ends` is relative to.
116    ///
117    /// This is generally zero for a "new" array, and non-zero after a slicing operation.
118    #[inline]
119    pub fn offset(&self) -> usize {
120        self.offset
121    }
122
123    /// The encoded "ends" of value runs.
124    ///
125    /// The `i`-th element indicates that there is a run of the same value, beginning
126    /// at `ends[i]` (inclusive) and terminating at `ends[i+1]` (exclusive).
127    #[inline]
128    pub fn ends(&self) -> &ArrayRef {
129        &self.ends
130    }
131
132    /// The scalar values.
133    ///
134    /// The `i`-th element is the scalar value for the `i`-th repeated run. The run begins
135    /// at `ends[i]` (inclusive) and terminates at `ends[i+1]` (exclusive).
136    #[inline]
137    pub fn values(&self) -> &ArrayRef {
138        &self.values
139    }
140}
141
142impl ArrayImpl for RunEndArray {
143    type Encoding = RunEndEncoding;
144
145    fn _len(&self) -> usize {
146        self.length
147    }
148
149    fn _dtype(&self) -> &DType {
150        self.values.dtype()
151    }
152
153    fn _vtable(&self) -> VTableRef {
154        VTableRef::new_ref(&RunEndEncoding)
155    }
156
157    fn _with_children(&self, children: &[ArrayRef]) -> VortexResult<Self> {
158        let ends = children[0].clone();
159        let values = children[1].clone();
160
161        Self::try_new(ends, values)
162    }
163}
164
165impl ArrayVariantsImpl for RunEndArray {
166    fn _as_bool_typed(&self) -> Option<&dyn BoolArrayTrait> {
167        Some(self)
168    }
169
170    fn _as_primitive_typed(&self) -> Option<&dyn PrimitiveArrayTrait> {
171        Some(self)
172    }
173}
174
175impl PrimitiveArrayTrait for RunEndArray {}
176
177impl BoolArrayTrait for RunEndArray {}
178
179impl ArrayValidityImpl for RunEndArray {
180    fn _is_valid(&self, index: usize) -> VortexResult<bool> {
181        let physical_idx = self
182            .find_physical_index(index)
183            .vortex_expect("Invalid index");
184        self.values().is_valid(physical_idx)
185    }
186
187    fn _all_valid(&self) -> VortexResult<bool> {
188        self.values().all_valid()
189    }
190
191    fn _all_invalid(&self) -> VortexResult<bool> {
192        self.values().all_invalid()
193    }
194
195    fn _validity_mask(&self) -> VortexResult<Mask> {
196        Ok(match self.values().validity_mask()? {
197            Mask::AllTrue(_) => Mask::AllTrue(self.len()),
198            Mask::AllFalse(_) => Mask::AllFalse(self.len()),
199            Mask::Values(values) => {
200                let ree_validity = RunEndArray::with_offset_and_length(
201                    self.ends().clone(),
202                    values.into_array(),
203                    self.offset(),
204                    self.len(),
205                )
206                .vortex_expect("invalid array")
207                .into_array();
208                Mask::from_buffer(ree_validity.to_bool()?.boolean_buffer().clone())
209            }
210        })
211    }
212}
213
214impl ArrayCanonicalImpl for RunEndArray {
215    fn _to_canonical(&self) -> VortexResult<Canonical> {
216        let pends = self.ends().to_primitive()?;
217        match self.dtype() {
218            DType::Bool(_) => {
219                let bools = self.values().to_bool()?;
220                runend_decode_bools(pends, bools, self.offset(), self.len()).map(Canonical::Bool)
221            }
222            DType::Primitive(..) => {
223                let pvalues = self.values().to_primitive()?;
224                runend_decode_primitive(pends, pvalues, self.offset(), self.len())
225                    .map(Canonical::Primitive)
226            }
227            _ => vortex_bail!("Only Primitive and Bool values are supported"),
228        }
229    }
230}
231
232impl ArrayStatisticsImpl for RunEndArray {
233    fn _stats_ref(&self) -> StatsSetRef<'_> {
234        self.stats_set.to_ref(self)
235    }
236}
237
238#[cfg(test)]
239mod tests {
240    use vortex_array::compute::scalar_at;
241    use vortex_array::{Array, IntoArray};
242    use vortex_buffer::buffer;
243    use vortex_dtype::{DType, Nullability, PType};
244
245    use crate::RunEndArray;
246
247    #[test]
248    fn test_runend_constructor() {
249        let arr = RunEndArray::try_new(
250            buffer![2u32, 5, 10].into_array(),
251            buffer![1i32, 2, 3].into_array(),
252        )
253        .unwrap();
254        assert_eq!(arr.len(), 10);
255        assert_eq!(
256            arr.dtype(),
257            &DType::Primitive(PType::I32, Nullability::NonNullable)
258        );
259
260        // 0, 1 => 1
261        // 2, 3, 4 => 2
262        // 5, 6, 7, 8, 9 => 3
263        assert_eq!(scalar_at(&arr, 0).unwrap(), 1.into());
264        assert_eq!(scalar_at(&arr, 2).unwrap(), 2.into());
265        assert_eq!(scalar_at(&arr, 5).unwrap(), 3.into());
266        assert_eq!(scalar_at(&arr, 9).unwrap(), 3.into());
267    }
268}
vortex_runend/array.rs

vortex_runend/
array.rs
