vortex_runend/
ops.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use vortex_array::arrays::ConstantArray;
5use vortex_array::vtable::OperationsVTable;
6use vortex_array::{Array, ArrayRef, IntoArray};
7use vortex_scalar::Scalar;
8
9use crate::{RunEndArray, RunEndVTable};
10
11impl OperationsVTable<RunEndVTable> for RunEndVTable {
12    fn slice(array: &RunEndArray, start: usize, stop: usize) -> ArrayRef {
13        let new_length = stop - start;
14
15        let slice_begin = array.find_physical_index(start);
16        let slice_end = array.find_physical_index(stop) + 1;
17
18        // If the sliced range contains only a single run, opt to return a ConstantArray.
19        if slice_begin + 1 == slice_end {
20            let value = array.values().scalar_at(slice_begin);
21            return ConstantArray::new(value, new_length).into_array();
22        }
23
24        // SAFETY: we maintain the ends invariant in our slice implementation
25        unsafe {
26            RunEndArray::new_unchecked(
27                array.ends().slice(slice_begin, slice_end),
28                array.values().slice(slice_begin, slice_end),
29                start + array.offset(),
30                new_length,
31            )
32            .into_array()
33        }
34    }
35
36    fn scalar_at(array: &RunEndArray, index: usize) -> Scalar {
37        array.values().scalar_at(array.find_physical_index(index))
38    }
39}
40
41#[cfg(test)]
42mod tests {
43    use vortex_array::arrays::PrimitiveArray;
44    use vortex_array::{Array, IntoArray, ToCanonical};
45    use vortex_buffer::buffer;
46    use vortex_dtype::{DType, Nullability, PType};
47
48    use crate::RunEndArray;
49
50    #[test]
51    fn slice_array() {
52        let arr = RunEndArray::try_new(
53            buffer![2u32, 5, 10].into_array(),
54            buffer![1i32, 2, 3].into_array(),
55        )
56        .unwrap()
57        .slice(3, 8);
58        assert_eq!(
59            arr.dtype(),
60            &DType::Primitive(PType::I32, Nullability::NonNullable)
61        );
62        assert_eq!(arr.len(), 5);
63
64        assert_eq!(
65            arr.to_primitive().unwrap().as_slice::<i32>(),
66            vec![2, 2, 3, 3, 3]
67        );
68    }
69
70    #[test]
71    fn double_slice() {
72        let arr = RunEndArray::try_new(
73            buffer![2u32, 5, 10].into_array(),
74            buffer![1i32, 2, 3].into_array(),
75        )
76        .unwrap()
77        .slice(3, 8);
78        assert_eq!(arr.len(), 5);
79
80        let doubly_sliced = arr.slice(0, 3);
81
82        assert_eq!(
83            doubly_sliced.to_primitive().unwrap().as_slice::<i32>(),
84            vec![2, 2, 3]
85        );
86    }
87
88    #[test]
89    fn slice_end_inclusive() {
90        let arr = RunEndArray::try_new(
91            buffer![2u32, 5, 10].into_array(),
92            buffer![1i32, 2, 3].into_array(),
93        )
94        .unwrap()
95        .slice(4, 10);
96        assert_eq!(
97            arr.dtype(),
98            &DType::Primitive(PType::I32, Nullability::NonNullable)
99        );
100        assert_eq!(arr.len(), 6);
101
102        assert_eq!(
103            arr.to_primitive().unwrap().as_slice::<i32>(),
104            vec![2, 3, 3, 3, 3, 3]
105        );
106    }
107
108    #[test]
109    fn slice_at_end() {
110        let re_array = RunEndArray::try_new(
111            buffer![7_u64, 10].into_array(),
112            buffer![2_u64, 3].into_array(),
113        )
114        .unwrap();
115
116        assert_eq!(re_array.len(), 10);
117
118        let sliced_array = re_array.slice(re_array.len(), re_array.len());
119        assert!(sliced_array.is_empty());
120    }
121
122    #[test]
123    fn slice_single_end() {
124        let re_array = RunEndArray::try_new(
125            buffer![7_u64, 10].into_array(),
126            buffer![2_u64, 3].into_array(),
127        )
128        .unwrap();
129
130        assert_eq!(re_array.len(), 10);
131
132        let sliced_array = re_array.slice(2, 5);
133
134        assert!(sliced_array.is_constant())
135    }
136
137    #[test]
138    fn ree_scalar_at_end() {
139        let scalar = RunEndArray::encode(
140            PrimitiveArray::from_iter([1, 1, 1, 4, 4, 4, 2, 2, 5, 5, 5, 5]).into_array(),
141        )
142        .unwrap()
143        .scalar_at(11);
144        assert_eq!(scalar, 5.into());
145    }
146}
vortex_runend/ops.rs

vortex_runend/
ops.rs
