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::ExecutionCtx;
8use vortex_array::IntoArray;
9use vortex_array::arrays::slice::SliceKernel;
10use vortex_error::VortexResult;
11
12use crate::ConstantArray;
13use crate::SparseArray;
14use crate::SparseVTable;
15
16impl SliceKernel for SparseVTable {
17    fn slice(
18        array: &SparseArray,
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 { SparseArray::new_unchecked(new_patches, array.fill_scalar().clone()) }
39                .into_array(),
40        ))
41    }
42}