vortex_array/arrays/null/compute/
take.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use vortex_dtype::match_each_integer_ptype;
5use vortex_error::VortexResult;
6use vortex_error::vortex_bail;
7
8use crate::Array;
9use crate::ArrayRef;
10use crate::IntoArray;
11use crate::ToCanonical;
12use crate::arrays::NullArray;
13use crate::arrays::NullVTable;
14use crate::compute::TakeKernel;
15use crate::compute::TakeKernelAdapter;
16use crate::register_kernel;
17
18impl TakeKernel for NullVTable {
19    #[allow(clippy::cast_possible_truncation)]
20    fn take(&self, array: &NullArray, indices: &dyn Array) -> VortexResult<ArrayRef> {
21        let indices = indices.to_primitive();
22
23        // Enforce all indices are valid
24        match_each_integer_ptype!(indices.ptype(), |T| {
25            for index in indices.as_slice::<T>() {
26                if (*index as usize) >= array.len() {
27                    vortex_bail!(OutOfBounds: *index as usize, 0, array.len());
28                }
29            }
30        });
31
32        Ok(NullArray::new(indices.len()).into_array())
33    }
34}
35
36register_kernel!(TakeKernelAdapter(NullVTable).lift());