vortex_compute/take/vector/
binaryview.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4//! Implementations of `take` on [`BinaryViewVector`].
5//!
6//! take` on the binary view simply performs a `take` on the views of the vector, which means the
7//! resulting vector may have "garbage" data in its child buffers.
8//!
9//! Note that it is on the outer array type to perform compaction / garbage collection.
10
11use num_traits::AsPrimitive;
12use vortex_buffer::Buffer;
13use vortex_dtype::UnsignedPType;
14use vortex_vector::VectorOps;
15use vortex_vector::binaryview::BinaryView;
16use vortex_vector::binaryview::BinaryViewType;
17use vortex_vector::binaryview::BinaryViewVector;
18use vortex_vector::primitive::PVector;
19
20use crate::take::Take;
21
22impl<T: BinaryViewType, I: UnsignedPType> Take<PVector<I>> for &BinaryViewVector<T> {
23    type Output = BinaryViewVector<T>;
24
25    fn take(self, indices: &PVector<I>) -> BinaryViewVector<T> {
26        if indices.validity().all_true() {
27            self.take(indices.elements().as_slice())
28        } else {
29            take_nullable(self, indices)
30        }
31    }
32}
33
34impl<T: BinaryViewType, I: UnsignedPType> Take<[I]> for &BinaryViewVector<T> {
35    type Output = BinaryViewVector<T>;
36
37    fn take(self, indices: &[I]) -> BinaryViewVector<T> {
38        let taken_views = take_views(self.views(), indices);
39        let taken_validity = self.validity().take(indices);
40
41        debug_assert_eq!(taken_views.len(), taken_validity.len());
42
43        // SAFETY: We called take on views and validity with the same indices, so the new components
44        // must have the same length. The views still point into the same buffers which we clone via
45        // Arc, so all view references remain valid.
46        unsafe {
47            BinaryViewVector::new_unchecked(taken_views, self.buffers().clone(), taken_validity)
48        }
49    }
50}
51
52fn take_nullable<T: BinaryViewType, I: UnsignedPType>(
53    binary_view: &BinaryViewVector<T>,
54    indices: &PVector<I>,
55) -> BinaryViewVector<T> {
56    // We ignore nullability when taking the views since we can let the `Mask` implementation
57    // determine which elements are null.
58    let taken_views = take_views(binary_view.views(), indices.elements().as_slice());
59
60    // Note that this is **not** the same as the `indices: &[I]` `take` implementation above.
61    let taken_validity = binary_view.validity().take(indices);
62
63    debug_assert_eq!(taken_views.len(), taken_validity.len());
64
65    // SAFETY: We used the same indices to take from both components, so they should still have the
66    // same length. The views still point into the same buffers which we clone via Arc, so all view
67    // references remain valid.
68    unsafe {
69        BinaryViewVector::new_unchecked(taken_views, binary_view.buffers().clone(), taken_validity)
70    }
71}
72
73/// Takes views at the given indices.
74fn take_views<I: AsPrimitive<usize>>(
75    views: &Buffer<BinaryView>,
76    indices: &[I],
77) -> Buffer<BinaryView> {
78    Buffer::<BinaryView>::from_trusted_len_iter(indices.iter().map(|i| views[i.as_()]))
79}
vortex_compute/take/vector/binaryview.rs

vortex_compute/take/vector/
binaryview.rs
