vortex_sparse/compute/
slice.rs

1use vortex_array::Array;
2use vortex_array::arrays::ConstantArray;
3use vortex_error::VortexResult;
4
5use crate::compute::SliceFn;
6use crate::{ArrayRef, SparseArray, SparseEncoding};
7
8impl SliceFn<&SparseArray> for SparseEncoding {
9    fn slice(&self, array: &SparseArray, start: usize, stop: usize) -> VortexResult<ArrayRef> {
10        let new_patches = array.patches().slice(start, stop)?;
11
12        let Some(new_patches) = new_patches else {
13            return Ok(ConstantArray::new(array.fill_scalar().clone(), stop - start).into_array());
14        };
15
16        Ok(
17            SparseArray::try_new_from_patches(new_patches, array.fill_scalar().clone())?
18                .into_array(),
19        )
20    }
21}
22
23#[cfg(test)]
24mod tests {
25    use vortex_array::compute::slice;
26    use vortex_array::{ArrayExt, IntoArray, ToCanonical};
27    use vortex_buffer::buffer;
28
29    use super::*;
30
31    #[test]
32    fn test_slice() {
33        let values = buffer![15_u32, 135, 13531, 42].into_array();
34        let indices = buffer![10_u64, 11, 50, 100].into_array();
35
36        let sparse = SparseArray::try_new(indices, values, 101, 0_u32.into())
37            .unwrap()
38            .into_array();
39
40        let sliced = slice(&sparse, 15, 100).unwrap();
41        assert_eq!(sliced.len(), 100 - 15);
42        let primitive = sliced
43            .as_::<SparseArray>()
44            .patches()
45            .values()
46            .to_primitive()
47            .unwrap();
48
49        assert_eq!(primitive.as_slice::<u32>(), &[13531]);
50    }
51
52    #[test]
53    fn doubly_sliced() {
54        let values = buffer![15_u32, 135, 13531, 42].into_array();
55        let indices = buffer![10_u64, 11, 50, 100].into_array();
56
57        let sparse = SparseArray::try_new(indices, values, 101, 0_u32.into())
58            .unwrap()
59            .into_array();
60
61        let sliced = slice(&sparse, 15, 100).unwrap();
62        assert_eq!(sliced.len(), 100 - 15);
63        let primitive = sliced
64            .as_::<SparseArray>()
65            .patches()
66            .values()
67            .to_primitive()
68            .unwrap();
69
70        assert_eq!(primitive.as_slice::<u32>(), &[13531]);
71
72        let doubly_sliced = slice(&sliced, 35, 36).unwrap();
73        let primitive_doubly_sliced = doubly_sliced
74            .as_::<SparseArray>()
75            .patches()
76            .values()
77            .to_primitive()
78            .unwrap();
79
80        assert_eq!(primitive_doubly_sliced.as_slice::<u32>(), &[13531]);
81    }
82
83    #[test]
84    fn slice_partially_invalid() {
85        let values = buffer![0u64].into_array();
86        let indices = buffer![0u8].into_array();
87
88        let sparse = SparseArray::try_new(indices, values, 1000, 999u64.into()).unwrap();
89        let sliced = slice(&sparse, 0, 1000).unwrap();
90        let mut expected = vec![999u64; 1000];
91        expected[0] = 0;
92
93        let actual = sliced.to_primitive().unwrap().as_slice::<u64>().to_vec();
94        assert_eq!(expected, actual);
95    }
96}
vortex_sparse/compute/slice.rs

vortex_sparse/compute/
slice.rs
