vortex_runend/compute/
take.rs

1use num_traits::AsPrimitive;
2use vortex_array::compute::{TakeFn, take};
3use vortex_array::variants::PrimitiveArrayTrait;
4use vortex_array::{Array, ArrayRef, IntoArray, ToCanonical};
5use vortex_dtype::match_each_integer_ptype;
6use vortex_error::{VortexResult, vortex_bail};
7
8use crate::{RunEndArray, RunEndEncoding};
9
10impl TakeFn<&RunEndArray> for RunEndEncoding {
11    fn take(&self, array: &RunEndArray, indices: &dyn Array) -> VortexResult<ArrayRef> {
12        let primitive_indices = indices.to_primitive()?;
13
14        let checked_indices = match_each_integer_ptype!(primitive_indices.ptype(), |$P| {
15            primitive_indices
16                .as_slice::<$P>()
17                .iter()
18                .copied()
19                .map(|idx| {
20                    let usize_idx = idx as usize;
21                    if usize_idx >= array.len() {
22                        vortex_bail!(OutOfBounds: usize_idx, 0, array.len());
23                    }
24                    Ok(usize_idx)
25                })
26                .collect::<VortexResult<Vec<_>>>()?
27        });
28
29        take_indices_unchecked(array, &checked_indices)
30    }
31}
32
33/// Perform a take operation on a RunEndArray by binary searching for each of the indices.
34pub fn take_indices_unchecked<T: AsPrimitive<usize>>(
35    array: &RunEndArray,
36    indices: &[T],
37) -> VortexResult<ArrayRef> {
38    let adjusted_indices = indices
39        .iter()
40        .map(|idx| idx.as_() + array.offset())
41        .collect::<Vec<_>>();
42    let physical_indices = array.find_physical_indices(&adjusted_indices)?.into_array();
43    take(array.values(), &physical_indices)
44}
45
46#[cfg(test)]
47mod test {
48    use vortex_array::arrays::PrimitiveArray;
49    use vortex_array::compute::{scalar_at, slice, take};
50    use vortex_array::{Array, ToCanonical};
51
52    use crate::RunEndArray;
53
54    fn ree_array() -> RunEndArray {
55        RunEndArray::encode(
56            PrimitiveArray::from_iter([1, 1, 1, 4, 4, 4, 2, 2, 5, 5, 5, 5]).into_array(),
57        )
58        .unwrap()
59    }
60
61    #[test]
62    fn ree_take() {
63        let taken = take(&ree_array(), &PrimitiveArray::from_iter([9, 8, 1, 3])).unwrap();
64        assert_eq!(
65            taken.to_primitive().unwrap().as_slice::<i32>(),
66            &[5, 5, 1, 4]
67        );
68    }
69
70    #[test]
71    fn ree_take_end() {
72        let taken = take(&ree_array(), &PrimitiveArray::from_iter([11])).unwrap();
73        assert_eq!(taken.to_primitive().unwrap().as_slice::<i32>(), &[5]);
74    }
75
76    #[test]
77    #[should_panic]
78    fn ree_take_out_of_bounds() {
79        take(&ree_array(), &PrimitiveArray::from_iter([12])).unwrap();
80    }
81
82    #[test]
83    fn sliced_take() {
84        let sliced = slice(&ree_array(), 4, 9).unwrap();
85        let taken = take(sliced.as_ref(), &PrimitiveArray::from_iter([1, 3, 4])).unwrap();
86
87        assert_eq!(taken.len(), 3);
88        assert_eq!(scalar_at(taken.as_ref(), 0).unwrap(), 4.into());
89        assert_eq!(scalar_at(taken.as_ref(), 1).unwrap(), 2.into());
90        assert_eq!(scalar_at(taken.as_ref(), 2).unwrap(), 5.into());
91    }
92}
vortex_runend/compute/take.rs

vortex_runend/compute/
take.rs
