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

vortex_runend/
ops.rs
