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!(arr.to_primitive().as_slice::<i32>(), vec![2, 2, 3, 3, 3]);
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        assert_eq!(arr.len(), 5);
78
79        let doubly_sliced = arr.slice(0..3);
80
81        assert_eq!(
82            doubly_sliced.to_primitive().as_slice::<i32>(),
83            vec![2, 2, 3]
84        );
85    }
86
87    #[test]
88    fn slice_end_inclusive() {
89        let arr = RunEndArray::try_new(
90            buffer![2u32, 5, 10].into_array(),
91            buffer![1i32, 2, 3].into_array(),
92        )
93        .unwrap()
94        .slice(4..10);
95        assert_eq!(
96            arr.dtype(),
97            &DType::Primitive(PType::I32, Nullability::NonNullable)
98        );
99        assert_eq!(arr.len(), 6);
100
101        assert_eq!(arr.to_primitive().as_slice::<i32>(), vec![2, 3, 3, 3, 3, 3]);
102    }
103
104    #[test]
105    fn slice_at_end() {
106        let re_array = RunEndArray::try_new(
107            buffer![7_u64, 10].into_array(),
108            buffer![2_u64, 3].into_array(),
109        )
110        .unwrap();
111
112        assert_eq!(re_array.len(), 10);
113
114        let sliced_array = re_array.slice(re_array.len()..re_array.len());
115        assert!(sliced_array.is_empty());
116    }
117
118    #[test]
119    fn slice_single_end() {
120        let re_array = RunEndArray::try_new(
121            buffer![7_u64, 10].into_array(),
122            buffer![2_u64, 3].into_array(),
123        )
124        .unwrap();
125
126        assert_eq!(re_array.len(), 10);
127
128        let sliced_array = re_array.slice(2..5);
129
130        assert!(sliced_array.is_constant())
131    }
132
133    #[test]
134    fn ree_scalar_at_end() {
135        let scalar = RunEndArray::encode(
136            PrimitiveArray::from_iter([1, 1, 1, 4, 4, 4, 2, 2, 5, 5, 5, 5]).into_array(),
137        )
138        .unwrap()
139        .scalar_at(11);
140        assert_eq!(scalar, 5.into());
141    }
142}
vortex_runend/ops.rs

vortex_runend/
ops.rs
