vortex_sparse/
ops.rs

1use vortex_array::arrays::ConstantArray;
2use vortex_array::vtable::OperationsVTable;
3use vortex_array::{Array, ArrayRef, IntoArray};
4use vortex_error::VortexResult;
5use vortex_scalar::Scalar;
6
7use crate::{SparseArray, SparseVTable};
8
9impl OperationsVTable<SparseVTable> for SparseVTable {
10    fn slice(array: &SparseArray, start: usize, stop: usize) -> VortexResult<ArrayRef> {
11        let new_patches = array.patches().slice(start, stop)?;
12
13        let Some(new_patches) = new_patches else {
14            return Ok(ConstantArray::new(array.fill_scalar().clone(), stop - start).into_array());
15        };
16
17        // If the number of values in the sparse array matches the array length, then all
18        // values are in fact patches, since patches are sorted this is the correct values.
19        if new_patches.array_len() == new_patches.values().len() {
20            return Ok(new_patches.into_values());
21        }
22
23        Ok(
24            SparseArray::try_new_from_patches(new_patches, array.fill_scalar().clone())?
25                .into_array(),
26        )
27    }
28
29    fn scalar_at(array: &SparseArray, index: usize) -> VortexResult<Scalar> {
30        Ok(array
31            .patches()
32            .get_patched(index)?
33            .unwrap_or_else(|| array.fill_scalar().clone()))
34    }
35}
36
37#[cfg(test)]
38mod tests {
39
40    use vortex_array::{IntoArray, ToCanonical};
41    use vortex_buffer::buffer;
42
43    use super::*;
44
45    #[test]
46    fn slice_partially_invalid() {
47        let values = buffer![0u64].into_array();
48        let indices = buffer![0u8].into_array();
49
50        let sparse = SparseArray::try_new(indices, values, 1000, 999u64.into()).unwrap();
51        let sliced = sparse.slice(0, 1000).unwrap();
52        let mut expected = vec![999u64; 1000];
53        expected[0] = 0;
54
55        let values = sliced.to_primitive().unwrap();
56        assert_eq!(values.as_slice::<u64>(), expected);
57    }
58}