Skip to main content

vortex_sparse/
slice.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::ops::Range;
5
6use vortex_array::ArrayRef;
7use vortex_array::ArrayView;
8use vortex_array::ExecutionCtx;
9use vortex_array::IntoArray;
10use vortex_array::arrays::slice::SliceKernel;
11use vortex_error::VortexResult;
12
13use crate::ConstantArray;
14use crate::Sparse;
15use crate::SparseExt as _;
16
17impl SliceKernel for Sparse {
18    fn slice(
19        array: ArrayView<'_, Self>,
20        range: Range<usize>,
21        _ctx: &mut ExecutionCtx,
22    ) -> VortexResult<Option<ArrayRef>> {
23        let Some(new_patches) = array.patches().slice(range.clone())? else {
24            return Ok(Some(
25                ConstantArray::new(array.fill_scalar().clone(), range.len()).into_array(),
26            ));
27        };
28
29        // If the number of values in the sparse array matches the array length, then all
30        // values are in fact patches, since patches are sorted this is the correct values.
31        if new_patches.array_len() == new_patches.values().len() {
32            return Ok(Some(new_patches.into_values()));
33        }
34
35        // SAFETY:
36        // patches slice will ensure that dtype of patches is unchanged and the indices and
37        // values match
38        Ok(Some(
39            unsafe { Sparse::new_unchecked(new_patches, array.fill_scalar().clone()) }.into_array(),
40        ))
41    }
42}