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

vortex_array/arrays/chunked/compute/
slice.rs
