vortex_array/arrays/chunked/compute/
slice.rs

1use vortex_error::{VortexResult, vortex_bail};
2
3use crate::arrays::ChunkedEncoding;
4use crate::arrays::chunked::ChunkedArray;
5use crate::compute::{SliceFn, slice};
6use crate::{Array, ArrayRef};
7
8impl SliceFn<&ChunkedArray> for ChunkedEncoding {
9    fn slice(&self, array: &ChunkedArray, start: usize, stop: usize) -> VortexResult<ArrayRef> {
10        let (offset_chunk, offset_in_first_chunk) = array.find_chunk_idx(start);
11        let (length_chunk, length_in_last_chunk) = array.find_chunk_idx(stop);
12
13        if array.is_empty() && (start != 0 || stop != 0) {
14            vortex_bail!(ComputeError: "Empty chunked array can't be sliced from {start} to {stop}");
15        } else if array.is_empty() {
16            return Ok(ChunkedArray::new_unchecked(vec![], array.dtype().clone()).into_array());
17        }
18
19        if length_chunk == offset_chunk {
20            let chunk = array.chunk(offset_chunk)?;
21            return Ok(ChunkedArray::new_unchecked(
22                vec![slice(chunk, offset_in_first_chunk, length_in_last_chunk)?],
23                array.dtype().clone(),
24            )
25            .into_array());
26        }
27
28        let mut chunks = (offset_chunk..length_chunk + 1)
29            .map(|i| array.chunk(i).cloned())
30            .collect::<VortexResult<Vec<_>>>()?;
31        if let Some(c) = chunks.first_mut() {
32            *c = slice(c, offset_in_first_chunk, c.len())?;
33        }
34
35        if length_in_last_chunk == 0 {
36            chunks.pop();
37        } else if let Some(c) = chunks.last_mut() {
38            *c = slice(c, 0, length_in_last_chunk)?;
39        }
40
41        Ok(ChunkedArray::new_unchecked(chunks, array.dtype().clone()).into_array())
42    }
43}
44
45#[cfg(test)]
46mod tests {
47    use vortex_dtype::{DType, NativePType, Nullability, PType};
48
49    use crate::arrays::{ChunkedArray, PrimitiveArray};
50    use crate::canonical::ToCanonical;
51    use crate::compute::slice;
52    use crate::{Array, ArrayExt};
53
54    fn chunked_array() -> ChunkedArray {
55        ChunkedArray::try_new(
56            vec![
57                PrimitiveArray::from_iter([1u64, 2, 3]).into_array(),
58                PrimitiveArray::from_iter([4u64, 5, 6]).into_array(),
59                PrimitiveArray::from_iter([7u64, 8, 9]).into_array(),
60            ],
61            DType::Primitive(PType::U64, Nullability::NonNullable),
62        )
63        .unwrap()
64    }
65
66    fn assert_equal_slices<T: NativePType>(arr: &dyn Array, slice: &[T]) {
67        let mut values = Vec::with_capacity(arr.len());
68        arr.as_::<ChunkedArray>()
69            .chunks()
70            .iter()
71            .map(|a| a.to_primitive().unwrap())
72            .for_each(|a| values.extend_from_slice(a.as_slice::<T>()));
73        assert_eq!(values, slice);
74    }
75
76    #[test]
77    fn slice_middle() {
78        assert_equal_slices(&slice(&chunked_array(), 2, 5).unwrap(), &[3u64, 4, 5])
79    }
80
81    #[test]
82    fn slice_begin() {
83        assert_equal_slices(&slice(&chunked_array(), 1, 3).unwrap(), &[2u64, 3]);
84    }
85
86    #[test]
87    fn slice_aligned() {
88        assert_equal_slices(&slice(&chunked_array(), 3, 6).unwrap(), &[4u64, 5, 6]);
89    }
90
91    #[test]
92    fn slice_many_aligned() {
93        assert_equal_slices(
94            &slice(&chunked_array(), 0, 6).unwrap(),
95            &[1u64, 2, 3, 4, 5, 6],
96        );
97    }
98
99    #[test]
100    fn slice_end() {
101        assert_equal_slices(&slice(&chunked_array(), 7, 8).unwrap(), &[8u64]);
102    }
103
104    #[test]
105    fn slice_exactly_end() {
106        assert_equal_slices(&slice(&chunked_array(), 6, 9).unwrap(), &[7u64, 8, 9]);
107    }
108
109    #[test]
110    fn slice_empty() {
111        let chunked = ChunkedArray::try_new(vec![], PType::U32.into()).unwrap();
112        let sliced = slice(&chunked, 0, 0).unwrap();
113
114        assert!(sliced.is_empty());
115    }
116}
vortex_array/arrays/chunked/compute/slice.rs

vortex_array/arrays/chunked/compute/
slice.rs
