vortex_sparse/
ops.rs

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

vortex_sparse/
ops.rs
