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::array_slots;
25use vortex_array::arrays::DecimalArray;
26use vortex_array::arrays::ListViewArray;
27use vortex_array::arrays::Primitive;
28use vortex_array::arrays::PrimitiveArray;
29use vortex_array::arrays::VarBinViewArray;
30use vortex_array::arrays::listview::ListViewArraySlotsExt;
31use vortex_array::buffer::BufferHandle;
32use vortex_array::dtype::DType;
33use vortex_array::dtype::Nullability;
34use vortex_array::dtype::PType;
35use vortex_array::legacy_session;
36use vortex_array::serde::ArrayChildren;
37use vortex_array::validity::Validity;
38use vortex_array::vtable::VTable;
39use vortex_array::vtable::ValidityVTable;
40use vortex_error::VortexExpect as _;
41use vortex_error::VortexResult;
42use vortex_error::vortex_bail;
43use vortex_error::vortex_ensure;
44use vortex_error::vortex_panic;
45use vortex_session::VortexSession;
46use vortex_session::registry::CachedId;
47
48use crate::compress::runend_decode_decimal;
49use crate::compress::runend_decode_primitive;
50use crate::compress::runend_decode_varbinview;
51use crate::compress::runend_encode;
52use crate::decompress_bool::runend_decode_bools;
53use crate::ops::find_physical_index;
54use crate::ops::find_slice_end_index;
55use crate::rules::RULES;
56
57/// A [`RunEnd`]-encoded Vortex array.
58pub type RunEndArray = Array<RunEnd>;
59
60#[derive(Clone, prost::Message)]
61pub struct RunEndMetadata {
62    #[prost(enumeration = "PType", tag = "1")]
63    pub ends_ptype: i32,
64    #[prost(uint64, tag = "2")]
65    pub num_runs: u64,
66    #[prost(uint64, tag = "3")]
67    pub offset: u64,
68}
69
70impl ArrayHash for RunEndData {
71    fn array_hash<H: Hasher>(&self, state: &mut H, _accuracy: EqMode) {
72        self.offset.hash(state);
73    }
74}
75
76impl ArrayEq for RunEndData {
77    fn array_eq(&self, other: &Self, _accuracy: EqMode) -> bool {
78        self.offset == other.offset
79    }
80}
81
82impl VTable for RunEnd {
83    type TypedArrayData = RunEndData;
84
85    type OperationsVTable = Self;
86    type ValidityVTable = Self;
87
88    fn id(&self) -> ArrayId {
89        static ID: CachedId = CachedId::new("vortex.runend");
90        *ID
91    }
92
93    #[allow(clippy::disallowed_methods)]
94    fn validate(
95        &self,
96        data: &Self::TypedArrayData,
97        dtype: &DType,
98        len: usize,
99        slots: &[Option<ArrayRef>],
100    ) -> VortexResult<()> {
101        let run_end_slots = RunEndSlotsView::from_slots(slots);
102        let ends = run_end_slots.ends;
103        let values = run_end_slots.values;
104        // TODO(ctx): trait fixes - VTable::validate has a fixed signature.
105        let mut ctx = legacy_session().create_execution_ctx();
106        RunEndData::validate_parts(ends, values, data.offset, len, &mut ctx)?;
107        vortex_ensure!(
108            values.dtype() == dtype,
109            "expected dtype {}, got {}",
110            dtype,
111            values.dtype()
112        );
113        Ok(())
114    }
115
116    fn nbuffers(_array: ArrayView<'_, Self>) -> usize {
117        0
118    }
119
120    fn buffer(_array: ArrayView<'_, Self>, idx: usize) -> BufferHandle {
121        vortex_panic!("RunEndArray buffer index {idx} out of bounds")
122    }
123
124    fn buffer_name(_array: ArrayView<'_, Self>, idx: usize) -> Option<String> {
125        vortex_panic!("RunEndArray buffer_name index {idx} out of bounds")
126    }
127
128    fn with_buffers(
129        &self,
130        array: ArrayView<'_, Self>,
131        buffers: &[BufferHandle],
132    ) -> VortexResult<ArrayParts<Self>> {
133        vortex_array::vtable::with_empty_buffers(self, array, buffers)
134    }
135
136    fn serialize(
137        array: ArrayView<'_, Self>,
138        _session: &VortexSession,
139    ) -> VortexResult<Option<Vec<u8>>> {
140        Ok(Some(
141            RunEndMetadata {
142                ends_ptype: PType::try_from(array.ends().dtype())
143                    .vortex_expect("Must be a valid PType") as i32,
144                num_runs: array.ends().len() as u64,
145                offset: array.offset() as u64,
146            }
147            .encode_to_vec(),
148        ))
149    }
150
151    fn deserialize(
152        &self,
153        dtype: &DType,
154        len: usize,
155        metadata: &[u8],
156        _buffers: &[BufferHandle],
157        children: &dyn ArrayChildren,
158        _session: &VortexSession,
159    ) -> VortexResult<ArrayParts<Self>> {
160        let metadata = RunEndMetadata::decode(metadata)?;
161        let ends_dtype = DType::Primitive(metadata.ends_ptype(), Nullability::NonNullable);
162        let runs = usize::try_from(metadata.num_runs).vortex_expect("Must be a valid usize");
163        let ends = children.get(0, &ends_dtype, runs)?;
164
165        let values = children.get(1, dtype, runs)?;
166        let offset = usize::try_from(metadata.offset).vortex_expect("Offset must be a valid usize");
167        let slots = RunEndSlots { ends, values }.into_slots();
168        let data = RunEndData::new(offset);
169        Ok(ArrayParts::new(self.clone(), dtype.clone(), len, data).with_slots(slots))
170    }
171
172    fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String {
173        RunEndSlots::NAMES[idx].to_string()
174    }
175
176    fn reduce_parent(
177        array: ArrayView<'_, Self>,
178        parent: &ArrayRef,
179        child_idx: usize,
180    ) -> VortexResult<Option<ArrayRef>> {
181        RULES.evaluate(array, parent, child_idx)
182    }
183
184    fn execute(array: Array<Self>, ctx: &mut ExecutionCtx) -> VortexResult<ExecutionResult> {
185        run_end_canonicalize(&array, ctx).map(ExecutionResult::done)
186    }
187}
188
189#[array_slots(RunEnd)]
190pub struct RunEndSlots {
191    /// The run-end positions marking where each run terminates.
192    #[slot(0)]
193    pub ends: ArrayRef,
194    /// The values for each run.
195    #[slot(1)]
196    pub values: ArrayRef,
197}
198
199#[derive(Clone, Debug)]
200pub struct RunEndData {
201    offset: usize,
202}
203
204impl Display for RunEndData {
205    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
206        write!(f, "offset: {}", self.offset)
207    }
208}
209
210pub struct RunEndDataParts {
211    pub ends: ArrayRef,
212    pub values: ArrayRef,
213    pub offset: usize,
214}
215
216pub trait RunEndArrayExt: RunEndArraySlotsExt {
217    fn offset(&self) -> usize {
218        self.offset
219    }
220
221    fn dtype(&self) -> &DType {
222        self.values().dtype()
223    }
224
225    fn find_physical_index(&self, index: usize, ctx: &mut ExecutionCtx) -> VortexResult<usize> {
226        find_physical_index(self.ends(), index + self.offset(), ctx)
227    }
228
229    fn find_slice_end_index(&self, index: usize, ctx: &mut ExecutionCtx) -> VortexResult<usize> {
230        find_slice_end_index(self.ends(), index + self.offset(), ctx)
231    }
232}
233
234impl<T: TypedArrayRef<RunEnd>> RunEndArrayExt for T {}
235
236#[derive(Clone, Debug)]
237pub struct RunEnd;
238
239impl RunEnd {
240    /// Build a new [`RunEndArray`] without validation.
241    ///
242    /// # Safety
243    /// See [`RunEndData::new_unchecked`] for preconditions.
244    pub unsafe fn new_unchecked(
245        ends: ArrayRef,
246        values: ArrayRef,
247        offset: usize,
248        length: usize,
249    ) -> RunEndArray {
250        let dtype = values.dtype().clone();
251        let slots = RunEndSlots { ends, values }.into_slots();
252        let data = unsafe { RunEndData::new_unchecked(offset) };
253        unsafe {
254            Array::from_parts_unchecked(
255                ArrayParts::new(RunEnd, dtype, length, data).with_slots(slots),
256            )
257        }
258    }
259
260    /// Build a new [`RunEndArray`] from ends and values.
261    pub fn try_new(
262        ends: ArrayRef,
263        values: ArrayRef,
264        ctx: &mut ExecutionCtx,
265    ) -> VortexResult<RunEndArray> {
266        let len = RunEndData::logical_len_from_ends(&ends, ctx)?;
267        RunEndData::validate_parts(&ends, &values, 0, len, ctx)?;
268        let dtype = values.dtype().clone();
269        let slots = RunEndSlots { ends, values }.into_slots();
270        let data = RunEndData::new(0);
271        Array::try_from_parts(ArrayParts::new(RunEnd, dtype, len, data).with_slots(slots))
272    }
273
274    /// Build a new [`RunEndArray`] from ends, values, offset, and length.
275    pub fn try_new_offset_length(
276        ends: ArrayRef,
277        values: ArrayRef,
278        offset: usize,
279        length: usize,
280        ctx: &mut ExecutionCtx,
281    ) -> VortexResult<RunEndArray> {
282        RunEndData::validate_parts(&ends, &values, offset, length, ctx)?;
283        let dtype = values.dtype().clone();
284        let slots = RunEndSlots { ends, values }.into_slots();
285        let data = RunEndData::new(offset);
286        Array::try_from_parts(ArrayParts::new(RunEnd, dtype, length, data).with_slots(slots))
287    }
288
289    /// Build a new [`RunEndArray`] from ends and values (panics on invalid input).
290    pub fn new(ends: ArrayRef, values: ArrayRef, ctx: &mut ExecutionCtx) -> RunEndArray {
291        Self::try_new(ends, values, ctx).vortex_expect("RunEndData is always valid")
292    }
293
294    /// Run the array through run-end encoding.
295    pub fn encode(array: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult<RunEndArray> {
296        if let Some(parray) = array.as_opt::<Primitive>() {
297            let (ends, values) = runend_encode(parray, ctx);
298            let ends = ends.into_array();
299            let len = array.len();
300            let dtype = values.dtype().clone();
301            let slots = RunEndSlots { ends, values }.into_slots();
302            let data = unsafe { RunEndData::new_unchecked(0) };
303            Array::try_from_parts(ArrayParts::new(RunEnd, dtype, len, data).with_slots(slots))
304        } else {
305            vortex_bail!("REE can only encode primitive arrays")
306        }
307    }
308}
309
310impl RunEndData {
311    fn logical_len_from_ends(ends: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult<usize> {
312        if ends.is_empty() {
313            Ok(0)
314        } else {
315            usize::try_from(&ends.execute_scalar(ends.len() - 1, ctx)?)
316        }
317    }
318
319    /// Validate that `ends` and `values` form a well-formed run-end array covering
320    /// `offset..offset + length`.
321    pub fn validate_parts(
322        ends: &ArrayRef,
323        values: &ArrayRef,
324        offset: usize,
325        length: usize,
326        ctx: &mut ExecutionCtx,
327    ) -> VortexResult<()> {
328        // DType validation
329        vortex_ensure!(
330            ends.dtype().is_unsigned_int(),
331            "run ends must be unsigned integers, was {}",
332            ends.dtype(),
333        );
334        vortex_ensure!(
335            ends.len() == values.len(),
336            "run ends len != run values len, {} != {}",
337            ends.len(),
338            values.len()
339        );
340
341        // Handle empty run-ends
342        if ends.is_empty() {
343            vortex_ensure!(
344                offset == 0,
345                "non-zero offset provided for empty RunEndArray"
346            );
347            return Ok(());
348        }
349
350        // Zero-length logical slices may retain run metadata from the source array.
351        if length == 0 {
352            return Ok(());
353        }
354
355        #[cfg(debug_assertions)]
356        {
357            // Run ends must be strictly sorted for binary search to work correctly.
358            let pre_validation = ends.statistics().to_owned();
359
360            let is_sorted = ends
361                .statistics()
362                .compute_is_strict_sorted(ctx)
363                .unwrap_or(false);
364
365            // Preserve the original statistics since compute_is_strict_sorted may have mutated them.
366            // We don't want to run with different stats in debug mode and outside.
367            ends.statistics().inherit(pre_validation.iter());
368            debug_assert!(is_sorted);
369        }
370
371        // Skip host-only validation when ends are not host-resident.
372        if !ends.is_host() {
373            return Ok(());
374        }
375
376        // Validate the offset and length are valid for the given ends and values
377        if offset != 0 && length != 0 {
378            let first_run_end = usize::try_from(&ends.execute_scalar(0, ctx)?)?;
379            if first_run_end < offset {
380                vortex_bail!("First run end {first_run_end} must be >= offset {offset}");
381            }
382        }
383
384        let last_run_end = usize::try_from(&ends.execute_scalar(ends.len() - 1, ctx)?)?;
385        let min_required_end = offset + length;
386        if last_run_end < min_required_end {
387            vortex_bail!("Last run end {last_run_end} must be >= offset+length {min_required_end}");
388        }
389
390        Ok(())
391    }
392}
393
394impl RunEndData {
395    /// Build a new `RunEndArray` from an array of run `ends` and an array of `values`.
396    ///
397    /// Panics if any of the validation conditions described in [`RunEnd::try_new`] is
398    /// not satisfied.
399    ///
400    /// # Examples
401    ///
402    /// ```
403    /// # use vortex_array::arrays::BoolArray;
404    /// # use vortex_array::IntoArray;
405    /// # use vortex_array::VortexSessionExecute;
406    /// # use vortex_buffer::buffer;
407    /// # use vortex_error::VortexResult;
408    /// # use vortex_runend::RunEnd;
409    /// # fn main() -> VortexResult<()> {
410    /// let session = vortex_array::array_session();
411    /// vortex_runend::initialize(&session);
412    /// let mut ctx = session.create_execution_ctx();
413    /// let ends = buffer![2u8, 3u8].into_array();
414    /// let values = BoolArray::from_iter([false, true]).into_array();
415    /// let run_end = RunEnd::new(ends, values, &mut ctx);
416    ///
417    /// // Array encodes
418    /// assert_eq!(run_end.execute_scalar(0, &mut ctx)?, false.into());
419    /// assert_eq!(run_end.execute_scalar(1, &mut ctx)?, false.into());
420    /// assert_eq!(run_end.execute_scalar(2, &mut ctx)?, true.into());
421    /// # Ok(())
422    /// # }
423    /// ```
424    pub fn new(offset: usize) -> Self {
425        Self { offset }
426    }
427
428    /// Build a new `RunEndArray` without validation.
429    ///
430    /// # Safety
431    ///
432    /// The caller must ensure that all the validation performed in
433    /// [`RunEnd::try_new_offset_length`] is
434    /// satisfied before calling this function.
435    ///
436    /// See [`RunEnd::try_new_offset_length`] for the preconditions needed to build a new array.
437    pub unsafe fn new_unchecked(offset: usize) -> Self {
438        Self { offset }
439    }
440
441    /// Run the array through run-end encoding.
442    pub fn encode(array: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult<Self> {
443        if let Some(parray) = array.as_opt::<Primitive>() {
444            let (_ends, _values) = runend_encode(parray, ctx);
445            // SAFETY: runend_encode handles this
446            unsafe { Ok(Self::new_unchecked(0)) }
447        } else {
448            vortex_bail!("REE can only encode primitive arrays")
449        }
450    }
451
452    pub fn into_parts(self, ends: ArrayRef, values: ArrayRef) -> RunEndDataParts {
453        RunEndDataParts {
454            ends,
455            values,
456            offset: self.offset,
457        }
458    }
459}
460
461impl ValidityVTable<RunEnd> for RunEnd {
462    fn validity(array: ArrayView<'_, RunEnd>) -> VortexResult<Validity> {
463        Ok(match array.values().validity()? {
464            Validity::NonNullable | Validity::AllValid => Validity::AllValid,
465            Validity::AllInvalid => Validity::AllInvalid,
466            Validity::Array(values_validity) => Validity::Array(unsafe {
467                RunEnd::new_unchecked(
468                    array.ends().clone(),
469                    values_validity,
470                    array.offset(),
471                    array.len(),
472                )
473                .into_array()
474            }),
475        })
476    }
477}
478
479pub(super) fn run_end_canonicalize(
480    array: &RunEndArray,
481    ctx: &mut ExecutionCtx,
482) -> VortexResult<ArrayRef> {
483    let pends = array.ends().clone().execute_as("ends", ctx)?;
484
485    Ok(match array.dtype() {
486        DType::Bool(_) => {
487            let bools = array.values().clone().execute_as("values", ctx)?;
488            runend_decode_bools(pends, bools, array.offset(), array.len(), ctx)?
489        }
490        DType::Primitive(..) => {
491            let pvalues = array.values().clone().execute_as("values", ctx)?;
492            runend_decode_primitive(pends, pvalues, array.offset(), array.len(), ctx)?.into_array()
493        }
494        DType::Decimal(..) => {
495            let values = array
496                .values()
497                .clone()
498                .execute_as::<DecimalArray>("values", ctx)?;
499            runend_decode_decimal(pends, values, 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        DType::List(..) => {
509            let values = array
510                .values()
511                .clone()
512                .execute_as::<ListViewArray>("values", ctx)?;
513            runend_decode_listview(pends, values, array.offset(), array.len())?.into_array()
514        }
515        _ => vortex_bail!("Unsupported RunEnd value type: {}", array.dtype()),
516    })
517}
518
519fn runend_decode_listview(
520    ends: PrimitiveArray,
521    values: ListViewArray,
522    offset: usize,
523    length: usize,
524) -> VortexResult<ListViewArray> {
525    let validity = match values.validity()? {
526        Validity::NonNullable => Validity::NonNullable,
527        Validity::AllValid => Validity::AllValid,
528        Validity::AllInvalid => Validity::AllInvalid,
529        Validity::Array(validity) => Validity::Array(unsafe {
530            RunEnd::new_unchecked(ends.clone().into_array(), validity, offset, length).into_array()
531        }),
532    };
533
534    // SAFETY: the `RunEndArray`s re-express valid per-run ListView metadata over the logical output
535    // length. The original `elements` child is reused, so every view still points at a valid range.
536    Ok(unsafe {
537        ListViewArray::new_unchecked(
538            values.elements().clone(),
539            RunEnd::new_unchecked(
540                ends.clone().into_array(),
541                values.offsets().clone(),
542                offset,
543                length,
544            )
545            .into_array(),
546            RunEnd::new_unchecked(ends.into_array(), values.sizes().clone(), offset, length)
547                .into_array(),
548            validity,
549        )
550    })
551}
552
553#[cfg(test)]
554mod tests {
555    use std::sync::Arc;
556    use std::sync::LazyLock;
557
558    use vortex_array::IntoArray;
559    use vortex_array::VortexSessionExecute;
560    use vortex_array::arrays::DecimalArray;
561    use vortex_array::arrays::DictArray;
562    use vortex_array::arrays::ListArray;
563    use vortex_array::arrays::ListViewArray;
564    use vortex_array::arrays::VarBinViewArray;
565    use vortex_array::arrays::listview::ListViewArraySlotsExt;
566    use vortex_array::assert_arrays_eq;
567    use vortex_array::builders::VarBinBuilder;
568    use vortex_array::dtype::DType;
569    use vortex_array::dtype::DecimalDType;
570    use vortex_array::dtype::Nullability;
571    use vortex_array::dtype::PType;
572    use vortex_array::dtype::i256;
573    use vortex_array::validity::Validity;
574    use vortex_buffer::buffer;
575    use vortex_error::VortexResult;
576    use vortex_session::VortexSession;
577
578    use crate::RunEnd;
579
580    static SESSION: LazyLock<VortexSession> = LazyLock::new(|| {
581        let session = vortex_array::array_session();
582        crate::initialize(&session);
583        session
584    });
585
586    #[test]
587    fn test_runend_constructor() {
588        let mut ctx = SESSION.create_execution_ctx();
589        let arr = RunEnd::new(
590            buffer![2u32, 5, 10].into_array(),
591            buffer![1i32, 2, 3].into_array(),
592            &mut ctx,
593        );
594        assert_eq!(arr.len(), 10);
595        assert_eq!(
596            arr.dtype(),
597            &DType::Primitive(PType::I32, Nullability::NonNullable)
598        );
599
600        // 0, 1 => 1
601        // 2, 3, 4 => 2
602        // 5, 6, 7, 8, 9 => 3
603        let expected = buffer![1, 1, 2, 2, 2, 3, 3, 3, 3, 3].into_array();
604        assert_arrays_eq!(arr.into_array(), expected, &mut ctx);
605    }
606
607    #[test]
608    fn test_runend_utf8() {
609        let mut ctx = SESSION.create_execution_ctx();
610        let values =
611            VarBinViewArray::from_iter_nullable_str([Some("a"), None, Some("c")]).into_array();
612        let arr = RunEnd::new(buffer![2u32, 5, 10].into_array(), values, &mut ctx);
613        assert_eq!(arr.len(), 10);
614        assert_eq!(arr.dtype(), &DType::Utf8(Nullability::Nullable));
615
616        let expected = VarBinViewArray::from_iter_nullable_str([
617            Some("a"),
618            Some("a"),
619            None,
620            None,
621            None,
622            Some("c"),
623            Some("c"),
624            Some("c"),
625            Some("c"),
626            Some("c"),
627        ])
628        .into_array();
629        let mut builder = VarBinBuilder::<i32>::with_capacity_in(
630            arr.dtype().clone(),
631            arr.len(),
632            vortex_buffer::BufferAllocatorRef::static_ref(),
633        );
634        arr.append_to_builder(&mut builder, &mut ctx).unwrap();
635        assert_arrays_eq!(builder.finish_into_varbin(), expected, &mut ctx);
636        assert_arrays_eq!(arr.into_array(), expected, &mut ctx);
637    }
638
639    #[test]
640    fn test_runend_decimal() {
641        let mut ctx = SESSION.create_execution_ctx();
642        let decimal_dtype = DecimalDType::new(10, 2);
643        let values = DecimalArray::from_iter([12345i64, 67890, -12300], decimal_dtype).into_array();
644        let arr = RunEnd::new(buffer![2u32, 5, 10].into_array(), values, &mut ctx);
645        assert_eq!(arr.len(), 10);
646        assert_eq!(
647            arr.dtype(),
648            &DType::Decimal(decimal_dtype, Nullability::NonNullable)
649        );
650
651        let expected = DecimalArray::from_iter(
652            [
653                12345i64, 12345, 67890, 67890, 67890, -12300, -12300, -12300, -12300, -12300,
654            ],
655            decimal_dtype,
656        )
657        .into_array();
658        assert_arrays_eq!(arr.into_array(), expected, &mut ctx);
659    }
660
661    #[test]
662    fn test_runend_list_i64() {
663        let mut ctx = SESSION.create_execution_ctx();
664        let values = ListArray::from_iter_slow::<u32, _>(
665            vec![vec![1i64, 2], vec![3], vec![4, 5, 6]],
666            Arc::new(DType::Primitive(PType::I64, Nullability::NonNullable)),
667        )
668        .unwrap()
669        .into_array();
670        let arr = RunEnd::new(buffer![2u32, 5, 10].into_array(), values, &mut ctx);
671
672        let expected = ListArray::from_iter_slow::<u32, _>(
673            vec![
674                vec![1i64, 2],
675                vec![1, 2],
676                vec![3],
677                vec![3],
678                vec![3],
679                vec![4, 5, 6],
680                vec![4, 5, 6],
681                vec![4, 5, 6],
682                vec![4, 5, 6],
683                vec![4, 5, 6],
684            ],
685            Arc::new(DType::Primitive(PType::I64, Nullability::NonNullable)),
686        )
687        .unwrap()
688        .into_array();
689        assert_arrays_eq!(arr.into_array(), expected, &mut ctx);
690    }
691
692    #[test]
693    fn test_runend_nullable_decimal() {
694        let mut ctx = SESSION.create_execution_ctx();
695        let decimal_dtype = DecimalDType::new(10, 2);
696        let values =
697            DecimalArray::from_option_iter([Some(12345i64), None, Some(-12300)], decimal_dtype)
698                .into_array();
699        let arr = RunEnd::new(buffer![2u32, 5, 10].into_array(), values, &mut ctx);
700        assert_eq!(arr.len(), 10);
701        assert_eq!(
702            arr.dtype(),
703            &DType::Decimal(decimal_dtype, Nullability::Nullable)
704        );
705
706        let expected = DecimalArray::from_option_iter(
707            [
708                Some(12345i64),
709                Some(12345),
710                None,
711                None,
712                None,
713                Some(-12300),
714                Some(-12300),
715                Some(-12300),
716                Some(-12300),
717                Some(-12300),
718            ],
719            decimal_dtype,
720        )
721        .into_array();
722        assert_arrays_eq!(arr.into_array(), expected, &mut ctx);
723    }
724
725    #[test]
726    fn test_runend_list_bool() {
727        let mut ctx = SESSION.create_execution_ctx();
728        let values = ListArray::from_iter_slow::<u32, _>(
729            vec![vec![true, false], vec![false], vec![true, true, false]],
730            Arc::new(DType::Bool(Nullability::NonNullable)),
731        )
732        .unwrap()
733        .into_array();
734        let arr = RunEnd::new(buffer![2u32, 5, 10].into_array(), values, &mut ctx);
735
736        let expected = ListArray::from_iter_slow::<u32, _>(
737            vec![
738                vec![true, false],
739                vec![true, false],
740                vec![false],
741                vec![false],
742                vec![false],
743                vec![true, true, false],
744                vec![true, true, false],
745                vec![true, true, false],
746                vec![true, true, false],
747                vec![true, true, false],
748            ],
749            Arc::new(DType::Bool(Nullability::NonNullable)),
750        )
751        .unwrap()
752        .into_array();
753        assert_arrays_eq!(arr.into_array(), expected, &mut ctx);
754    }
755
756    #[test]
757    fn test_runend_list_utf8() {
758        let mut ctx = SESSION.create_execution_ctx();
759        let values = ListArray::try_new(
760            VarBinViewArray::from_iter_str(["a", "b", "c", "d", "e", "f"]).into_array(),
761            buffer![0u32, 2, 3, 6].into_array(),
762            Validity::NonNullable,
763        )
764        .unwrap()
765        .into_array();
766        let arr = RunEnd::new(buffer![2u32, 5, 10].into_array(), values, &mut ctx);
767
768        let expected = ListArray::try_new(
769            VarBinViewArray::from_iter_str([
770                "a", "b", "a", "b", "c", "c", "c", "d", "e", "f", "d", "e", "f", "d", "e", "f",
771                "d", "e", "f", "d", "e", "f",
772            ])
773            .into_array(),
774            buffer![0u32, 2, 4, 5, 6, 7, 10, 13, 16, 19, 22].into_array(),
775            Validity::NonNullable,
776        )
777        .unwrap()
778        .into_array();
779        assert_arrays_eq!(arr.into_array(), expected, &mut ctx);
780    }
781
782    #[test]
783    fn test_runend_list_canonicalizes_to_runend_listview_slots() -> VortexResult<()> {
784        let mut ctx = SESSION.create_execution_ctx();
785        let values = ListArray::try_new(
786            buffer![1i64, 2, 3, 4, 5, 6].into_array(),
787            buffer![0u32, 2, 3, 6].into_array(),
788            Validity::from_iter([true, false, true]),
789        )?
790        .into_array();
791        let arr = RunEnd::try_new(buffer![2u32, 5, 6].into_array(), values, &mut ctx)?;
792
793        let listview = arr
794            .clone()
795            .into_array()
796            .execute::<ListViewArray>(&mut ctx)?;
797        assert!(listview.offsets().is::<RunEnd>());
798        assert!(listview.sizes().is::<RunEnd>());
799        match listview.validity()? {
800            Validity::Array(validity) => assert!(validity.is::<RunEnd>()),
801            validity => panic!("expected array-backed validity, got {validity:?}"),
802        }
803
804        let expected = ListArray::try_new(
805            buffer![1i64, 2, 1, 2, 3, 3, 3, 4, 5, 6].into_array(),
806            buffer![0u32, 2, 4, 5, 6, 7, 10].into_array(),
807            Validity::from_iter([true, true, false, false, false, true]),
808        )?
809        .into_array();
810        assert_arrays_eq!(arr.into_array(), expected, &mut ctx);
811        Ok(())
812    }
813
814    #[test]
815    fn test_runend_dict() {
816        let mut ctx = SESSION.create_execution_ctx();
817        let dict_values = VarBinViewArray::from_iter_str(["x", "y", "z"]).into_array();
818        let dict_codes = buffer![0u32, 1, 2].into_array();
819        let dict = DictArray::try_new(dict_codes, dict_values).unwrap();
820
821        let arr = RunEnd::try_new(
822            buffer![2u32, 5, 10].into_array(),
823            dict.into_array(),
824            &mut ctx,
825        )
826        .unwrap();
827        assert_eq!(arr.len(), 10);
828
829        let expected =
830            VarBinViewArray::from_iter_str(["x", "x", "y", "y", "y", "z", "z", "z", "z", "z"])
831                .into_array();
832        assert_arrays_eq!(arr.into_array(), expected, &mut ctx);
833    }
834
835    #[test]
836    fn test_runend_decimal_i128() -> VortexResult<()> {
837        let mut ctx = SESSION.create_execution_ctx();
838        let decimal_dtype = DecimalDType::new(20, 2);
839        let values = DecimalArray::from_iter([12_345i128, -67_890, 100], decimal_dtype);
840        let arr = RunEnd::try_new(
841            buffer![2u32, 5, 6].into_array(),
842            values.into_array(),
843            &mut ctx,
844        )?;
845
846        let decoded = arr.into_array().execute::<DecimalArray>(&mut ctx)?;
847        let expected = DecimalArray::from_iter(
848            [12_345i128, 12_345, -67_890, -67_890, -67_890, 100],
849            decimal_dtype,
850        );
851        assert_arrays_eq!(decoded, expected, &mut ctx);
852        Ok(())
853    }
854
855    #[test]
856    fn test_runend_decimal_nullable() -> VortexResult<()> {
857        let mut ctx = SESSION.create_execution_ctx();
858        let decimal_dtype = DecimalDType::new(20, 2);
859        let values =
860            DecimalArray::from_option_iter([Some(12_345i128), None, Some(-67_890)], decimal_dtype);
861        let arr = RunEnd::try_new(
862            buffer![2u32, 5, 7].into_array(),
863            values.into_array(),
864            &mut ctx,
865        )?;
866
867        let decoded = arr.into_array().execute::<DecimalArray>(&mut ctx)?;
868        let expected = DecimalArray::from_option_iter(
869            [
870                Some(12_345i128),
871                Some(12_345),
872                None,
873                None,
874                None,
875                Some(-67_890),
876                Some(-67_890),
877            ],
878            decimal_dtype,
879        );
880        assert_arrays_eq!(decoded, expected, &mut ctx);
881        Ok(())
882    }
883
884    #[test]
885    fn test_runend_decimal_slice() -> VortexResult<()> {
886        let mut ctx = SESSION.create_execution_ctx();
887        let decimal_dtype = DecimalDType::new(20, 2);
888        let values = DecimalArray::from_iter([100i128, 200, 300], decimal_dtype);
889        let arr = RunEnd::try_new(
890            buffer![3u32, 5, 10].into_array(),
891            values.into_array(),
892            &mut ctx,
893        )?;
894
895        let sliced = arr.slice(2..8)?;
896        let decoded = sliced.execute::<DecimalArray>(&mut ctx)?;
897        let expected = DecimalArray::from_iter([100i128, 200, 200, 300, 300, 300], decimal_dtype);
898        assert_arrays_eq!(decoded, expected, &mut ctx);
899        Ok(())
900    }
901
902    #[test]
903    fn test_runend_decimal_i256() -> VortexResult<()> {
904        let mut ctx = SESSION.create_execution_ctx();
905        let decimal_dtype = DecimalDType::new(40, 4);
906        let first = i256::from_i128(123_456);
907        let second = i256::from_i128(-789_012);
908        let values = DecimalArray::from_iter([first, second], decimal_dtype);
909        let arr = RunEnd::try_new(buffer![2u32, 5].into_array(), values.into_array(), &mut ctx)?;
910
911        let decoded = arr.into_array().execute::<DecimalArray>(&mut ctx)?;
912        let expected =
913            DecimalArray::from_iter([first, first, second, second, second], decimal_dtype);
914        assert_arrays_eq!(decoded, expected, &mut ctx);
915        Ok(())
916    }
917}
vortex_runend/array.rs

vortex_runend/
array.rs
