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::{IntoArray, ToCanonical};
27    use vortex_buffer::buffer;
28
29    use super::*;
30
31    #[test]
32    fn slice_partially_invalid() {
33        let values = buffer![0u64].into_array();
34        let indices = buffer![0u8].into_array();
35
36        let sparse = SparseArray::try_new(indices, values, 1000, 999u64.into()).unwrap();
37        let sliced = slice(&sparse, 0, 1000).unwrap();
38        let mut expected = vec![999u64; 1000];
39        expected[0] = 0;
40
41        let values = sliced.to_primitive().unwrap();
42        assert_eq!(values.as_slice::<u64>(), expected);
43    }
44}