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

vortex_sparse/
ops.rs
