vortex_array/arrays/chunked/compute/
slice.rs

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

vortex_array/arrays/chunked/compute/
slice.rs
