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;
15
16impl SliceKernel for Sparse {
17    fn slice(
18        array: ArrayView<'_, Self>,
19        range: Range<usize>,
20        _ctx: &mut ExecutionCtx,
21    ) -> VortexResult<Option<ArrayRef>> {
22        let Some(new_patches) = array.patches().slice(range.clone())? else {
23            return Ok(Some(
24                ConstantArray::new(array.fill_scalar().clone(), range.len()).into_array(),
25            ));
26        };
27
28        // If the number of values in the sparse array matches the array length, then all
29        // values are in fact patches, since patches are sorted this is the correct values.
30        if new_patches.array_len() == new_patches.values().len() {
31            return Ok(Some(new_patches.into_values()));
32        }
33
34        // SAFETY:
35        // patches slice will ensure that dtype of patches is unchanged and the indices and
36        // values match
37        Ok(Some(
38            unsafe { Sparse::new_unchecked(new_patches, array.fill_scalar().clone()) }.into_array(),
39        ))
40    }
41}