vortex_array/arrays/chunked/compute/
take.rs

1use vortex_buffer::BufferMut;
2use vortex_dtype::PType;
3use vortex_error::VortexResult;
4
5use crate::arrays::ChunkedVTable;
6use crate::arrays::chunked::ChunkedArray;
7use crate::compute::{TakeKernel, TakeKernelAdapter, cast, take};
8use crate::{Array, ArrayRef, IntoArray, ToCanonical, register_kernel};
9
10impl TakeKernel for ChunkedVTable {
11    fn take(&self, array: &ChunkedArray, indices: &dyn Array) -> VortexResult<ArrayRef> {
12        let indices = cast(indices, PType::U64.into())?.to_primitive()?;
13
14        // While the chunk idx remains the same, accumulate a list of chunk indices.
15        let mut chunks = Vec::new();
16        let mut indices_in_chunk = BufferMut::<u64>::empty();
17        let mut prev_chunk_idx = array
18            .find_chunk_idx(indices.as_slice::<u64>()[0].try_into()?)
19            .0;
20        for idx in indices.as_slice::<u64>() {
21            let idx = usize::try_from(*idx)?;
22            let (chunk_idx, idx_in_chunk) = array.find_chunk_idx(idx);
23
24            if chunk_idx != prev_chunk_idx {
25                // Start a new chunk
26                let indices_in_chunk_array = indices_in_chunk.clone().into_array();
27                chunks.push(take(array.chunk(prev_chunk_idx)?, &indices_in_chunk_array)?);
28                indices_in_chunk.clear();
29            }
30
31            indices_in_chunk.push(idx_in_chunk as u64);
32            prev_chunk_idx = chunk_idx;
33        }
34
35        if !indices_in_chunk.is_empty() {
36            let indices_in_chunk_array = indices_in_chunk.into_array();
37            chunks.push(take(array.chunk(prev_chunk_idx)?, &indices_in_chunk_array)?);
38        }
39
40        Ok(ChunkedArray::new_unchecked(chunks, array.dtype().clone()).into_array())
41    }
42}
43
44register_kernel!(TakeKernelAdapter(ChunkedVTable).lift());
45
46#[cfg(test)]
47mod test {
48    use vortex_buffer::buffer;
49
50    use crate::IntoArray;
51    use crate::array::Array;
52    use crate::arrays::chunked::ChunkedArray;
53    use crate::canonical::ToCanonical;
54    use crate::compute::take;
55
56    #[test]
57    fn test_take() {
58        let a = buffer![1i32, 2, 3].into_array();
59        let arr = ChunkedArray::try_new(vec![a.clone(), a.clone(), a.clone()], a.dtype().clone())
60            .unwrap();
61        assert_eq!(arr.nchunks(), 3);
62        assert_eq!(arr.len(), 9);
63        let indices = buffer![0u64, 0, 6, 4].into_array();
64
65        let result = take(arr.as_ref(), indices.as_ref())
66            .unwrap()
67            .to_primitive()
68            .unwrap();
69        assert_eq!(result.as_slice::<i32>(), &[1, 1, 1, 2]);
70    }
71}
vortex_array/arrays/chunked/compute/take.rs

vortex_array/arrays/chunked/compute/
take.rs
