vortex_compute/take/vector/
pvector.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use vortex_dtype::NativePType;
5use vortex_dtype::UnsignedPType;
6use vortex_vector::VectorOps;
7use vortex_vector::primitive::PVector;
8
9use crate::take::Take;
10
11impl<T: NativePType, I: UnsignedPType> Take<PVector<I>> for &PVector<T> {
12    type Output = PVector<T>;
13
14    fn take(self, indices: &PVector<I>) -> PVector<T> {
15        if indices.validity().all_true() {
16            self.take(indices.elements().as_slice())
17        } else {
18            take_nullable(self, indices)
19        }
20    }
21}
22
23impl<T: NativePType, I: UnsignedPType> Take<[I]> for &PVector<T> {
24    type Output = PVector<T>;
25
26    fn take(self, indices: &[I]) -> PVector<T> {
27        let taken_elements = self.elements().take(indices);
28        let taken_validity = self.validity().take(indices);
29
30        debug_assert_eq!(taken_elements.len(), taken_validity.len());
31
32        // SAFETY: we called take on both components of the vector with the same indices, so the new
33        // components must have the same length.
34        unsafe { PVector::new_unchecked(taken_elements, taken_validity) }
35    }
36}
37
38fn take_nullable<T: NativePType, I: UnsignedPType>(
39    pvector: &PVector<T>,
40    indices: &PVector<I>,
41) -> PVector<T> {
42    // We ignore nullability when taking the elements since we can let the `Mask` implementation
43    // determine which elements are null.
44    let taken_elements = pvector.elements().take(indices.elements().as_slice());
45    let taken_validity = pvector.validity().take(indices);
46
47    debug_assert_eq!(taken_elements.len(), taken_validity.len());
48
49    // SAFETY: We used the same indices to take from both components, so they should still have the
50    // same length.
51    unsafe { PVector::new_unchecked(taken_elements, taken_validity) }
52}
vortex_compute/take/vector/pvector.rs

vortex_compute/take/vector/
pvector.rs
