vortex_fastlanes/rle/
kernel.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::SliceExecuteAdaptor;
10use vortex_array::arrays::SliceKernel;
11use vortex_array::kernel::ParentKernelSet;
12use vortex_error::VortexResult;
13
14use crate::FL_CHUNK_SIZE;
15use crate::RLEArray;
16use crate::RLEVTable;
17
18pub(crate) static PARENT_KERNELS: ParentKernelSet<RLEVTable> =
19    ParentKernelSet::new(&[ParentKernelSet::lift(&SliceExecuteAdaptor(RLEVTable))]);
20
21impl SliceKernel for RLEVTable {
22    fn slice(
23        array: &RLEArray,
24        range: Range<usize>,
25        _ctx: &mut ExecutionCtx,
26    ) -> VortexResult<Option<ArrayRef>> {
27        let offset_in_chunk = array.offset();
28        let chunk_start_idx = (offset_in_chunk + range.start) / FL_CHUNK_SIZE;
29        let chunk_end_idx = (offset_in_chunk + range.end).div_ceil(FL_CHUNK_SIZE);
30
31        let values_start_idx = array.values_idx_offset(chunk_start_idx);
32        let values_end_idx = if chunk_end_idx < array.values_idx_offsets().len() {
33            array.values_idx_offset(chunk_end_idx)
34        } else {
35            array.values().len()
36        };
37
38        let sliced_values = array.values().slice(values_start_idx..values_end_idx)?;
39
40        let sliced_values_idx_offsets = array
41            .values_idx_offsets()
42            .slice(chunk_start_idx..chunk_end_idx)?;
43
44        let sliced_indices = array
45            .indices()
46            .slice(chunk_start_idx * FL_CHUNK_SIZE..chunk_end_idx * FL_CHUNK_SIZE)?;
47
48        // SAFETY: Slicing preserves all invariants.
49        Ok(Some(unsafe {
50            RLEArray::new_unchecked(
51                sliced_values,
52                sliced_indices,
53                sliced_values_idx_offsets,
54                array.dtype().clone(),
55                // Keep the offset relative to the first chunk.
56                (array.offset() + range.start) % FL_CHUNK_SIZE,
57                range.len(),
58            )
59            .into_array()
60        }))
61    }
62}
vortex_fastlanes/rle/kernel.rs

vortex_fastlanes/rle/
kernel.rs
