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

vortex_runend/
ops.rs
