vortex_runend/
array.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::fmt::Debug;
5use std::fmt::Display;
6use std::fmt::Formatter;
7use std::hash::Hash;
8use std::hash::Hasher;
9
10use prost::Message;
11use vortex_array::Array;
12use vortex_array::ArrayEq;
13use vortex_array::ArrayHash;
14use vortex_array::ArrayId;
15use vortex_array::ArrayParts;
16use vortex_array::ArrayRef;
17use vortex_array::ArrayView;
18use vortex_array::EqMode;
19use vortex_array::ExecutionCtx;
20use vortex_array::ExecutionResult;
21use vortex_array::IntoArray;
22use vortex_array::TypedArrayRef;
23use vortex_array::VortexSessionExecute;
24use vortex_array::arrays::Primitive;
25use vortex_array::arrays::VarBinViewArray;
26use vortex_array::buffer::BufferHandle;
27use vortex_array::dtype::DType;
28use vortex_array::dtype::Nullability;
29use vortex_array::dtype::PType;
30use vortex_array::legacy_session;
31use vortex_array::serde::ArrayChildren;
32use vortex_array::smallvec::smallvec;
33use vortex_array::validity::Validity;
34use vortex_array::vtable::VTable;
35use vortex_array::vtable::ValidityVTable;
36use vortex_error::VortexExpect as _;
37use vortex_error::VortexResult;
38use vortex_error::vortex_bail;
39use vortex_error::vortex_ensure;
40use vortex_error::vortex_panic;
41use vortex_session::VortexSession;
42use vortex_session::registry::CachedId;
43
44use crate::compress::runend_decode_primitive;
45use crate::compress::runend_decode_varbinview;
46use crate::compress::runend_encode;
47use crate::decompress_bool::runend_decode_bools;
48use crate::ops::find_physical_index;
49use crate::ops::find_slice_end_index;
50use crate::rules::RULES;
51
52/// A [`RunEnd`]-encoded Vortex array.
53pub type RunEndArray = Array<RunEnd>;
54
55#[derive(Clone, prost::Message)]
56pub struct RunEndMetadata {
57    #[prost(enumeration = "PType", tag = "1")]
58    pub ends_ptype: i32,
59    #[prost(uint64, tag = "2")]
60    pub num_runs: u64,
61    #[prost(uint64, tag = "3")]
62    pub offset: u64,
63}
64
65impl ArrayHash for RunEndData {
66    fn array_hash<H: Hasher>(&self, state: &mut H, _accuracy: EqMode) {
67        self.offset.hash(state);
68    }
69}
70
71impl ArrayEq for RunEndData {
72    fn array_eq(&self, other: &Self, _accuracy: EqMode) -> bool {
73        self.offset == other.offset
74    }
75}
76
77impl VTable for RunEnd {
78    type TypedArrayData = RunEndData;
79
80    type OperationsVTable = Self;
81    type ValidityVTable = Self;
82
83    fn id(&self) -> ArrayId {
84        static ID: CachedId = CachedId::new("vortex.runend");
85        *ID
86    }
87
88    #[allow(clippy::disallowed_methods)]
89    fn validate(
90        &self,
91        data: &Self::TypedArrayData,
92        dtype: &DType,
93        len: usize,
94        slots: &[Option<ArrayRef>],
95    ) -> VortexResult<()> {
96        let ends = slots[ENDS_SLOT]
97            .as_ref()
98            .vortex_expect("RunEndArray ends slot");
99        let values = slots[VALUES_SLOT]
100            .as_ref()
101            .vortex_expect("RunEndArray values slot");
102        // TODO(ctx): trait fixes - VTable::validate has a fixed signature.
103        let mut ctx = legacy_session().create_execution_ctx();
104        RunEndData::validate_parts(ends, values, data.offset, len, &mut ctx)?;
105        vortex_ensure!(
106            values.dtype() == dtype,
107            "expected dtype {}, got {}",
108            dtype,
109            values.dtype()
110        );
111        Ok(())
112    }
113
114    fn nbuffers(_array: ArrayView<'_, Self>) -> usize {
115        0
116    }
117
118    fn buffer(_array: ArrayView<'_, Self>, idx: usize) -> BufferHandle {
119        vortex_panic!("RunEndArray buffer index {idx} out of bounds")
120    }
121
122    fn buffer_name(_array: ArrayView<'_, Self>, idx: usize) -> Option<String> {
123        vortex_panic!("RunEndArray buffer_name index {idx} out of bounds")
124    }
125
126    fn with_buffers(
127        &self,
128        array: ArrayView<'_, Self>,
129        buffers: &[BufferHandle],
130    ) -> VortexResult<ArrayParts<Self>> {
131        vortex_array::vtable::with_empty_buffers(self, array, buffers)
132    }
133
134    fn serialize(
135        array: ArrayView<'_, Self>,
136        _session: &VortexSession,
137    ) -> VortexResult<Option<Vec<u8>>> {
138        Ok(Some(
139            RunEndMetadata {
140                ends_ptype: PType::try_from(array.ends().dtype())
141                    .vortex_expect("Must be a valid PType") as i32,
142                num_runs: array.ends().len() as u64,
143                offset: array.offset() as u64,
144            }
145            .encode_to_vec(),
146        ))
147    }
148
149    fn deserialize(
150        &self,
151        dtype: &DType,
152        len: usize,
153        metadata: &[u8],
154        _buffers: &[BufferHandle],
155        children: &dyn ArrayChildren,
156        _session: &VortexSession,
157    ) -> VortexResult<ArrayParts<Self>> {
158        let metadata = RunEndMetadata::decode(metadata)?;
159        let ends_dtype = DType::Primitive(metadata.ends_ptype(), Nullability::NonNullable);
160        let runs = usize::try_from(metadata.num_runs).vortex_expect("Must be a valid usize");
161        let ends = children.get(0, &ends_dtype, runs)?;
162
163        let values = children.get(1, dtype, runs)?;
164        let offset = usize::try_from(metadata.offset).vortex_expect("Offset must be a valid usize");
165        let slots = smallvec![Some(ends), Some(values)];
166        let data = RunEndData::new(offset);
167        Ok(ArrayParts::new(self.clone(), dtype.clone(), len, data).with_slots(slots))
168    }
169
170    fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String {
171        SLOT_NAMES[idx].to_string()
172    }
173
174    fn reduce_parent(
175        array: ArrayView<'_, Self>,
176        parent: &ArrayRef,
177        child_idx: usize,
178    ) -> VortexResult<Option<ArrayRef>> {
179        RULES.evaluate(array, parent, child_idx)
180    }
181
182    fn execute(array: Array<Self>, ctx: &mut ExecutionCtx) -> VortexResult<ExecutionResult> {
183        run_end_canonicalize(&array, ctx).map(ExecutionResult::done)
184    }
185}
186
187/// The run-end positions marking where each run terminates.
188pub(super) const ENDS_SLOT: usize = 0;
189/// The values for each run.
190pub(super) const VALUES_SLOT: usize = 1;
191pub(super) const NUM_SLOTS: usize = 2;
192pub(super) const SLOT_NAMES: [&str; NUM_SLOTS] = ["ends", "values"];
193
194#[derive(Clone, Debug)]
195pub struct RunEndData {
196    offset: usize,
197}
198
199impl Display for RunEndData {
200    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
201        write!(f, "offset: {}", self.offset)
202    }
203}
204
205pub struct RunEndDataParts {
206    pub ends: ArrayRef,
207    pub values: ArrayRef,
208    pub offset: usize,
209}
210
211pub trait RunEndArrayExt: TypedArrayRef<RunEnd> {
212    fn offset(&self) -> usize {
213        self.offset
214    }
215
216    fn ends(&self) -> &ArrayRef {
217        self.as_ref().slots()[ENDS_SLOT]
218            .as_ref()
219            .vortex_expect("RunEndArray ends slot")
220    }
221
222    fn values(&self) -> &ArrayRef {
223        self.as_ref().slots()[VALUES_SLOT]
224            .as_ref()
225            .vortex_expect("RunEndArray values slot")
226    }
227
228    fn dtype(&self) -> &DType {
229        self.values().dtype()
230    }
231
232    fn find_physical_index(&self, index: usize, ctx: &mut ExecutionCtx) -> VortexResult<usize> {
233        find_physical_index(self.ends(), index + self.offset(), ctx)
234    }
235
236    fn find_slice_end_index(&self, index: usize, ctx: &mut ExecutionCtx) -> VortexResult<usize> {
237        find_slice_end_index(self.ends(), index + self.offset(), ctx)
238    }
239}
240
241impl<T: TypedArrayRef<RunEnd>> RunEndArrayExt for T {}
242
243#[derive(Clone, Debug)]
244pub struct RunEnd;
245
246impl RunEnd {
247    /// Build a new [`RunEndArray`] without validation.
248    ///
249    /// # Safety
250    /// See [`RunEndData::new_unchecked`] for preconditions.
251    pub unsafe fn new_unchecked(
252        ends: ArrayRef,
253        values: ArrayRef,
254        offset: usize,
255        length: usize,
256    ) -> RunEndArray {
257        let dtype = values.dtype().clone();
258        let slots = smallvec![Some(ends), Some(values)];
259        let data = unsafe { RunEndData::new_unchecked(offset) };
260        unsafe {
261            Array::from_parts_unchecked(
262                ArrayParts::new(RunEnd, dtype, length, data).with_slots(slots),
263            )
264        }
265    }
266
267    /// Build a new [`RunEndArray`] from ends and values.
268    pub fn try_new(
269        ends: ArrayRef,
270        values: ArrayRef,
271        ctx: &mut ExecutionCtx,
272    ) -> VortexResult<RunEndArray> {
273        let len = RunEndData::logical_len_from_ends(&ends, ctx)?;
274        RunEndData::validate_parts(&ends, &values, 0, len, ctx)?;
275        let dtype = values.dtype().clone();
276        let slots = smallvec![Some(ends), Some(values)];
277        let data = RunEndData::new(0);
278        Array::try_from_parts(ArrayParts::new(RunEnd, dtype, len, data).with_slots(slots))
279    }
280
281    /// Build a new [`RunEndArray`] from ends, values, offset, and length.
282    pub fn try_new_offset_length(
283        ends: ArrayRef,
284        values: ArrayRef,
285        offset: usize,
286        length: usize,
287        ctx: &mut ExecutionCtx,
288    ) -> VortexResult<RunEndArray> {
289        RunEndData::validate_parts(&ends, &values, offset, length, ctx)?;
290        let dtype = values.dtype().clone();
291        let slots = smallvec![Some(ends), Some(values)];
292        let data = RunEndData::new(offset);
293        Array::try_from_parts(ArrayParts::new(RunEnd, dtype, length, data).with_slots(slots))
294    }
295
296    /// Build a new [`RunEndArray`] from ends and values (panics on invalid input).
297    pub fn new(ends: ArrayRef, values: ArrayRef, ctx: &mut ExecutionCtx) -> RunEndArray {
298        Self::try_new(ends, values, ctx).vortex_expect("RunEndData is always valid")
299    }
300
301    /// Run the array through run-end encoding.
302    pub fn encode(array: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult<RunEndArray> {
303        if let Some(parray) = array.as_opt::<Primitive>() {
304            let (ends, values) = runend_encode(parray, ctx);
305            let ends = ends.into_array();
306            let len = array.len();
307            let dtype = values.dtype().clone();
308            let slots = smallvec![Some(ends), Some(values)];
309            let data = unsafe { RunEndData::new_unchecked(0) };
310            Array::try_from_parts(ArrayParts::new(RunEnd, dtype, len, data).with_slots(slots))
311        } else {
312            vortex_bail!("REE can only encode primitive arrays")
313        }
314    }
315}
316
317impl RunEndData {
318    fn logical_len_from_ends(ends: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult<usize> {
319        if ends.is_empty() {
320            Ok(0)
321        } else {
322            usize::try_from(&ends.execute_scalar(ends.len() - 1, ctx)?)
323        }
324    }
325
326    /// Validate that `ends` and `values` form a well-formed run-end array covering
327    /// `offset..offset + length`.
328    pub fn validate_parts(
329        ends: &ArrayRef,
330        values: &ArrayRef,
331        offset: usize,
332        length: usize,
333        ctx: &mut ExecutionCtx,
334    ) -> VortexResult<()> {
335        // DType validation
336        vortex_ensure!(
337            ends.dtype().is_unsigned_int(),
338            "run ends must be unsigned integers, was {}",
339            ends.dtype(),
340        );
341        vortex_ensure!(
342            ends.len() == values.len(),
343            "run ends len != run values len, {} != {}",
344            ends.len(),
345            values.len()
346        );
347
348        // Handle empty run-ends
349        if ends.is_empty() {
350            vortex_ensure!(
351                offset == 0,
352                "non-zero offset provided for empty RunEndArray"
353            );
354            return Ok(());
355        }
356
357        // Zero-length logical slices may retain run metadata from the source array.
358        if length == 0 {
359            return Ok(());
360        }
361
362        #[cfg(debug_assertions)]
363        {
364            // Run ends must be strictly sorted for binary search to work correctly.
365            let pre_validation = ends.statistics().to_owned();
366
367            let is_sorted = ends
368                .statistics()
369                .compute_is_strict_sorted(ctx)
370                .unwrap_or(false);
371
372            // Preserve the original statistics since compute_is_strict_sorted may have mutated them.
373            // We don't want to run with different stats in debug mode and outside.
374            ends.statistics().inherit(pre_validation.iter());
375            debug_assert!(is_sorted);
376        }
377
378        // Skip host-only validation when ends are not host-resident.
379        if !ends.is_host() {
380            return Ok(());
381        }
382
383        // Validate the offset and length are valid for the given ends and values
384        if offset != 0 && length != 0 {
385            let first_run_end = usize::try_from(&ends.execute_scalar(0, ctx)?)?;
386            if first_run_end < offset {
387                vortex_bail!("First run end {first_run_end} must be >= offset {offset}");
388            }
389        }
390
391        let last_run_end = usize::try_from(&ends.execute_scalar(ends.len() - 1, ctx)?)?;
392        let min_required_end = offset + length;
393        if last_run_end < min_required_end {
394            vortex_bail!("Last run end {last_run_end} must be >= offset+length {min_required_end}");
395        }
396
397        Ok(())
398    }
399}
400
401impl RunEndData {
402    /// Build a new `RunEndArray` from an array of run `ends` and an array of `values`.
403    ///
404    /// Panics if any of the validation conditions described in [`RunEnd::try_new`] is
405    /// not satisfied.
406    ///
407    /// # Examples
408    ///
409    /// ```
410    /// # use vortex_array::arrays::BoolArray;
411    /// # use vortex_array::IntoArray;
412    /// # use vortex_array::VortexSessionExecute;
413    /// # use vortex_buffer::buffer;
414    /// # use vortex_error::VortexResult;
415    /// # use vortex_runend::RunEnd;
416    /// # fn main() -> VortexResult<()> {
417    /// let session = vortex_array::array_session();
418    /// vortex_runend::initialize(&session);
419    /// let mut ctx = session.create_execution_ctx();
420    /// let ends = buffer![2u8, 3u8].into_array();
421    /// let values = BoolArray::from_iter([false, true]).into_array();
422    /// let run_end = RunEnd::new(ends, values, &mut ctx);
423    ///
424    /// // Array encodes
425    /// assert_eq!(run_end.execute_scalar(0, &mut ctx)?, false.into());
426    /// assert_eq!(run_end.execute_scalar(1, &mut ctx)?, false.into());
427    /// assert_eq!(run_end.execute_scalar(2, &mut ctx)?, true.into());
428    /// # Ok(())
429    /// # }
430    /// ```
431    pub fn new(offset: usize) -> Self {
432        Self { offset }
433    }
434
435    /// Build a new `RunEndArray` without validation.
436    ///
437    /// # Safety
438    ///
439    /// The caller must ensure that all the validation performed in
440    /// [`RunEnd::try_new_offset_length`] is
441    /// satisfied before calling this function.
442    ///
443    /// See [`RunEnd::try_new_offset_length`] for the preconditions needed to build a new array.
444    pub unsafe fn new_unchecked(offset: usize) -> Self {
445        Self { offset }
446    }
447
448    /// Run the array through run-end encoding.
449    pub fn encode(array: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult<Self> {
450        if let Some(parray) = array.as_opt::<Primitive>() {
451            let (_ends, _values) = runend_encode(parray, ctx);
452            // SAFETY: runend_encode handles this
453            unsafe { Ok(Self::new_unchecked(0)) }
454        } else {
455            vortex_bail!("REE can only encode primitive arrays")
456        }
457    }
458
459    pub fn into_parts(self, ends: ArrayRef, values: ArrayRef) -> RunEndDataParts {
460        RunEndDataParts {
461            ends,
462            values,
463            offset: self.offset,
464        }
465    }
466}
467
468impl ValidityVTable<RunEnd> for RunEnd {
469    fn validity(array: ArrayView<'_, RunEnd>) -> VortexResult<Validity> {
470        Ok(match array.values().validity()? {
471            Validity::NonNullable | Validity::AllValid => Validity::AllValid,
472            Validity::AllInvalid => Validity::AllInvalid,
473            Validity::Array(values_validity) => Validity::Array(unsafe {
474                RunEnd::new_unchecked(
475                    array.ends().clone(),
476                    values_validity,
477                    array.offset(),
478                    array.len(),
479                )
480                .into_array()
481            }),
482        })
483    }
484}
485
486pub(super) fn run_end_canonicalize(
487    array: &RunEndArray,
488    ctx: &mut ExecutionCtx,
489) -> VortexResult<ArrayRef> {
490    let pends = array.ends().clone().execute_as("ends", ctx)?;
491
492    Ok(match array.dtype() {
493        DType::Bool(_) => {
494            let bools = array.values().clone().execute_as("values", ctx)?;
495            runend_decode_bools(pends, bools, array.offset(), array.len(), ctx)?
496        }
497        DType::Primitive(..) => {
498            let pvalues = array.values().clone().execute_as("values", ctx)?;
499            runend_decode_primitive(pends, pvalues, array.offset(), array.len(), ctx)?.into_array()
500        }
501        DType::Utf8(_) | DType::Binary(_) => {
502            let values = array
503                .values()
504                .clone()
505                .execute_as::<VarBinViewArray>("values", ctx)?;
506            runend_decode_varbinview(pends, values, array.offset(), array.len(), ctx)?.into_array()
507        }
508        _ => vortex_bail!("Unsupported RunEnd value type: {}", array.dtype()),
509    })
510}
511
512#[cfg(test)]
513mod tests {
514    use std::sync::LazyLock;
515
516    use vortex_array::IntoArray;
517    use vortex_array::VortexSessionExecute;
518    use vortex_array::arrays::DictArray;
519    use vortex_array::arrays::VarBinViewArray;
520    use vortex_array::assert_arrays_eq;
521    use vortex_array::dtype::DType;
522    use vortex_array::dtype::Nullability;
523    use vortex_array::dtype::PType;
524    use vortex_buffer::buffer;
525    use vortex_session::VortexSession;
526
527    use crate::RunEnd;
528
529    static SESSION: LazyLock<VortexSession> = LazyLock::new(|| {
530        let session = vortex_array::array_session();
531        crate::initialize(&session);
532        session
533    });
534
535    #[test]
536    fn test_runend_constructor() {
537        let mut ctx = SESSION.create_execution_ctx();
538        let arr = RunEnd::new(
539            buffer![2u32, 5, 10].into_array(),
540            buffer![1i32, 2, 3].into_array(),
541            &mut ctx,
542        );
543        assert_eq!(arr.len(), 10);
544        assert_eq!(
545            arr.dtype(),
546            &DType::Primitive(PType::I32, Nullability::NonNullable)
547        );
548
549        // 0, 1 => 1
550        // 2, 3, 4 => 2
551        // 5, 6, 7, 8, 9 => 3
552        let expected = buffer![1, 1, 2, 2, 2, 3, 3, 3, 3, 3].into_array();
553        assert_arrays_eq!(arr.into_array(), expected, &mut ctx);
554    }
555
556    #[test]
557    fn test_runend_utf8() {
558        let mut ctx = SESSION.create_execution_ctx();
559        let values = VarBinViewArray::from_iter_str(["a", "b", "c"]).into_array();
560        let arr = RunEnd::new(buffer![2u32, 5, 10].into_array(), values, &mut ctx);
561        assert_eq!(arr.len(), 10);
562        assert_eq!(arr.dtype(), &DType::Utf8(Nullability::NonNullable));
563
564        let expected =
565            VarBinViewArray::from_iter_str(["a", "a", "b", "b", "b", "c", "c", "c", "c", "c"])
566                .into_array();
567        assert_arrays_eq!(arr.into_array(), expected, &mut ctx);
568    }
569
570    #[test]
571    fn test_runend_dict() {
572        let mut ctx = SESSION.create_execution_ctx();
573        let dict_values = VarBinViewArray::from_iter_str(["x", "y", "z"]).into_array();
574        let dict_codes = buffer![0u32, 1, 2].into_array();
575        let dict = DictArray::try_new(dict_codes, dict_values).unwrap();
576
577        let arr = RunEnd::try_new(
578            buffer![2u32, 5, 10].into_array(),
579            dict.into_array(),
580            &mut ctx,
581        )
582        .unwrap();
583        assert_eq!(arr.len(), 10);
584
585        let expected =
586            VarBinViewArray::from_iter_str(["x", "x", "y", "y", "y", "z", "z", "z", "z", "z"])
587                .into_array();
588        assert_arrays_eq!(arr.into_array(), expected, &mut ctx);
589    }
590}
vortex_runend/array.rs

vortex_runend/
array.rs
