vortex_runend/
ops.rs

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

vortex_runend/
ops.rs
