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_error::VortexResult;
8use vortex_scalar::Scalar;
9
10use crate::{SparseArray, SparseVTable};
11
12impl OperationsVTable<SparseVTable> for SparseVTable {
13    fn slice(array: &SparseArray, start: usize, stop: usize) -> VortexResult<ArrayRef> {
14        let new_patches = array.patches().slice(start, stop)?;
15
16        let Some(new_patches) = new_patches else {
17            return Ok(ConstantArray::new(array.fill_scalar().clone(), stop - start).into_array());
18        };
19
20        // If the number of values in the sparse array matches the array length, then all
21        // values are in fact patches, since patches are sorted this is the correct values.
22        if new_patches.array_len() == new_patches.values().len() {
23            return Ok(new_patches.into_values());
24        }
25
26        Ok(
27            SparseArray::try_new_from_patches(new_patches, array.fill_scalar().clone())?
28                .into_array(),
29        )
30    }
31
32    fn scalar_at(array: &SparseArray, index: usize) -> VortexResult<Scalar> {
33        Ok(array
34            .patches()
35            .get_patched(index)?
36            .unwrap_or_else(|| array.fill_scalar().clone()))
37    }
38}
39
40#[cfg(test)]
41mod tests {
42    use vortex_array::{IntoArray, ToCanonical};
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).unwrap();
54        let mut expected = vec![999u64; 1000];
55        expected[0] = 0;
56
57        let values = sliced.to_primitive().unwrap();
58        assert_eq!(values.as_slice::<u64>(), expected);
59    }
60}
vortex_sparse/ops.rs

vortex_sparse/
ops.rs
