Skip to main content

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::ArrayVTable;
8use vortex_array::ArrayView;
9use vortex_array::ExecutionCtx;
10use vortex_array::IntoArray;
11use vortex_array::arrays::Slice;
12use vortex_array::arrays::slice::SliceExecuteAdaptor;
13use vortex_array::arrays::slice::SliceKernel;
14use vortex_array::optimizer::kernels::ArrayKernelsExt;
15use vortex_error::VortexResult;
16use vortex_session::VortexSession;
17
18use crate::FL_CHUNK_SIZE;
19use crate::RLE;
20use crate::rle::RLEArrayExt;
21use crate::rle::RLEArraySlotsExt;
22
23pub(crate) fn initialize(session: &VortexSession) {
24    let kernels = session.kernels();
25    kernels.register_execute_parent_kernel(Slice.id(), RLE, SliceExecuteAdaptor(RLE));
26}
27
28impl SliceKernel for RLE {
29    fn slice(
30        array: ArrayView<'_, Self>,
31        range: Range<usize>,
32        ctx: &mut ExecutionCtx,
33    ) -> VortexResult<Option<ArrayRef>> {
34        let offset_in_chunk = array.offset();
35        let chunk_start_idx = (offset_in_chunk + range.start) / FL_CHUNK_SIZE;
36        let chunk_end_idx = (offset_in_chunk + range.end).div_ceil(FL_CHUNK_SIZE);
37
38        let values_start_idx = array.values_idx_offset(chunk_start_idx, ctx);
39        let values_end_idx = if chunk_end_idx < array.values_idx_offsets().len() {
40            array.values_idx_offset(chunk_end_idx, ctx)
41        } else {
42            array.values().len()
43        };
44
45        let sliced_values = array.values().slice(values_start_idx..values_end_idx)?;
46
47        let sliced_values_idx_offsets = array
48            .values_idx_offsets()
49            .slice(chunk_start_idx..chunk_end_idx)?;
50
51        let sliced_indices = array
52            .indices()
53            .slice(chunk_start_idx * FL_CHUNK_SIZE..chunk_end_idx * FL_CHUNK_SIZE)?;
54
55        Ok(Some(
56            RLE::try_new(
57                sliced_values,
58                sliced_indices,
59                sliced_values_idx_offsets,
60                // Keep the offset relative to the first chunk.
61                (array.offset() + range.start) % FL_CHUNK_SIZE,
62                range.len(),
63            )?
64            .into_array(),
65        ))
66    }
67}