vortex_compute/take/vector/
dvector.rs

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

vortex_compute/take/vector/
dvector.rs
