vortex_sparse/
ops.rs

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

vortex_sparse/
ops.rs
