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        // If the number of values in the sparse array matches the array length, then all
17        // values are in fact patches, since patches are sorted this is the correct values.
18        if new_patches.array_len() == new_patches.values().len() {
19            return Ok(new_patches.into_values());
20        }
21
22        Ok(
23            SparseArray::try_new_from_patches(new_patches, array.fill_scalar().clone())?
24                .into_array(),
25        )
26    }
27}
28
29#[cfg(test)]
30mod tests {
31    use vortex_array::compute::slice;
32    use vortex_array::{IntoArray, ToCanonical};
33    use vortex_buffer::buffer;
34
35    use super::*;
36
37    #[test]
38    fn slice_partially_invalid() {
39        let values = buffer![0u64].into_array();
40        let indices = buffer![0u8].into_array();
41
42        let sparse = SparseArray::try_new(indices, values, 1000, 999u64.into()).unwrap();
43        let sliced = slice(&sparse, 0, 1000).unwrap();
44        let mut expected = vec![999u64; 1000];
45        expected[0] = 0;
46
47        let values = sliced.to_primitive().unwrap();
48        assert_eq!(values.as_slice::<u64>(), expected);
49    }
50}
vortex_sparse/compute/slice.rs

vortex_sparse/compute/
slice.rs
