vortex_runend/
ops.rs

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

vortex_runend/
ops.rs
