vortex_compute/take/vector/
listview.rs

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

vortex_compute/take/vector/
listview.rs
