vortex_array/arrays/chunked/compute/
slice.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::ops::Range;
5
6use itertools::Itertools;
7use vortex_error::VortexResult;
8
9use crate::ArrayRef;
10use crate::ExecutionCtx;
11use crate::IntoArray;
12use crate::arrays::ChunkedArray;
13use crate::arrays::ChunkedVTable;
14use crate::arrays::SliceKernel;
15
16impl SliceKernel for ChunkedVTable {
17    fn slice(
18        array: &Self::Array,
19        range: Range<usize>,
20        _ctx: &mut ExecutionCtx,
21    ) -> VortexResult<Option<ArrayRef>> {
22        assert!(
23            !array.is_empty() || (range.start > 0 && range.end > 0),
24            "Empty chunked array can't be sliced from {} to {}",
25            range.start,
26            range.end
27        );
28
29        if array.is_empty() {
30            // SAFETY: empty chunked array trivially satisfies all validations
31            unsafe {
32                return Ok(Some(
33                    ChunkedArray::new_unchecked(vec![], array.dtype().clone()).into_array(),
34                ));
35            }
36        }
37
38        let (offset_chunk, offset_in_first_chunk) = array.find_chunk_idx(range.start)?;
39        let (length_chunk, length_in_last_chunk) = array.find_chunk_idx(range.end)?;
40
41        if length_chunk == offset_chunk {
42            let chunk = array.chunk(offset_chunk);
43            return Ok(Some(
44                chunk.slice(offset_in_first_chunk..length_in_last_chunk)?,
45            ));
46        }
47
48        let mut chunks = (offset_chunk..length_chunk + 1)
49            .map(|i| array.chunk(i).clone())
50            .collect_vec();
51        if let Some(c) = chunks.first_mut() {
52            *c = c.slice(offset_in_first_chunk..c.len())?;
53        }
54
55        if length_in_last_chunk == 0 {
56            chunks.pop();
57        } else if let Some(c) = chunks.last_mut() {
58            *c = c.slice(0..length_in_last_chunk)?;
59        }
60
61        // SAFETY: chunks are slices of the original valid chunks, preserving their dtype.
62        // All chunks maintain the same dtype as the original array.
63        Ok(Some(unsafe {
64            ChunkedArray::new_unchecked(chunks, array.dtype().clone()).into_array()
65        }))
66    }
67}
vortex_array/arrays/chunked/compute/slice.rs

vortex_array/arrays/chunked/compute/
slice.rs
