Feature: add take for the view vectors (#5682)

Tracking Issue: #5652

I separated this out from
#5653 just in case we wanted
to discuss how to handle `take` on the view type vectors.

See the module docs for more detail.

After this gets merged I can implement `DictArray` `batch_execute`!

Signed-off-by: Connor Tsui <[email protected]>

Author: connortsui20

vortex-compute/src/take/vector/binaryview.rs

@@ -1,13 +1,18 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
-// use std::ops::Deref;
-
-// use num_traits::AsPrimitive;
-// use vortex_buffer::Buffer;
+//! Implementations of `take` on [`BinaryViewVector`].
+//!
+//! take` on the binary view simply performs a `take` on the views of the vector, which means the
+//! resulting vector may have "garbage" data in its child buffers.
+//!
+//! Note that it is on the outer array type to perform compaction / garbage collection.
+
+use num_traits::AsPrimitive;
+use vortex_buffer::Buffer;
 use vortex_dtype::UnsignedPType;
 use vortex_vector::VectorOps;
-// use vortex_vector::binaryview::BinaryView;
+use vortex_vector::binaryview::BinaryView;
 use vortex_vector::binaryview::BinaryViewType;
 use vortex_vector::binaryview::BinaryViewVector;
 use vortex_vector::primitive::PVector;
@@ -29,11 +34,7 @@ impl<T: BinaryViewType, I: UnsignedPType> Take<PVector<I>> for &BinaryViewVector
 impl<T: BinaryViewType, I: UnsignedPType> Take<[I]> for &BinaryViewVector<T> {
     type Output = BinaryViewVector<T>;
 
-    fn take(self, _indices: &[I]) -> BinaryViewVector<T> {
-        todo!("TODO(connor): Implement `take` for `BinaryViewVector` and figure out rebuilding");
-
-        /*
-
+    fn take(self, indices: &[I]) -> BinaryViewVector<T> {
         let taken_views = take_views(self.views(), indices);
         let taken_validity = self.validity().take(indices);
 
@@ -45,45 +46,34 @@ impl<T: BinaryViewType, I: UnsignedPType> Take<[I]> for &BinaryViewVector<T> {
         unsafe {
             BinaryViewVector::new_unchecked(taken_views, self.buffers().clone(), taken_validity)
         }
-
-        */
     }
 }
 
 fn take_nullable<T: BinaryViewType, I: UnsignedPType>(
-    _bvector: &BinaryViewVector<T>,
-    _indices: &PVector<I>,
+    binary_view: &BinaryViewVector<T>,
+    indices: &PVector<I>,
 ) -> BinaryViewVector<T> {
-    todo!("TODO(connor): Implement `take` for `BinaryViewVector` and figure out rebuilding");
-
-    /*
-
     // We ignore nullability when taking the views since we can let the `Mask` implementation
     // determine which elements are null.
-    let taken_views = take_views(bvector.views(), indices.elements().as_slice());
-    let taken_validity = bvector.validity().take(indices);
+    let taken_views = take_views(binary_view.views(), indices.elements().as_slice());
+
+    // Note that this is **not** the same as the `indices: &[I]` `take` implementation above.
+    let taken_validity = binary_view.validity().take(indices);
 
     debug_assert_eq!(taken_views.len(), taken_validity.len());
 
     // SAFETY: We used the same indices to take from both components, so they should still have the
     // same length. The views still point into the same buffers which we clone via Arc, so all view
     // references remain valid.
     unsafe {
-        BinaryViewVector::new_unchecked(taken_views, bvector.buffers().clone(), taken_validity)
+        BinaryViewVector::new_unchecked(taken_views, binary_view.buffers().clone(), taken_validity)
     }
-
-    */
 }
 
-/*
-
 /// Takes views at the given indices.
 fn take_views<I: AsPrimitive<usize>>(
     views: &Buffer<BinaryView>,
     indices: &[I],
 ) -> Buffer<BinaryView> {
-    let views_ref = views.deref();
-    Buffer::<BinaryView>::from_trusted_len_iter(indices.iter().map(|i| views_ref[(*i).as_()]))
+    Buffer::<BinaryView>::from_trusted_len_iter(indices.iter().map(|i| views[i.as_()]))
 }
-
-*/

vortex-compute/src/take/vector/listview.rs

@@ -1,6 +1,13 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
+//! Implementations of `take` on [`ListViewVector`].
+//!
+//! take` on the list view simply performs a `take` on the views of the vector, which means the
+//! resulting vector may have "garbage" data in its `elements` child vector.
+//!
+//! Note that it is on the outer array type to perform compaction / garbage collection.
+
 use vortex_dtype::UnsignedPType;
 use vortex_vector::VectorOps;
 use vortex_vector::listview::ListViewVector;
@@ -23,11 +30,7 @@ impl<I: UnsignedPType> Take<PVector<I>> for &ListViewVector {
 impl<I: UnsignedPType> Take<[I]> for &ListViewVector {
     type Output = ListViewVector;
 
-    fn take(self, _indices: &[I]) -> ListViewVector {
-        todo!("TODO(connor): Implement `take` for `ListViewVector` and figure out rebuilding");
-
-        /*
-
+    fn take(self, indices: &[I]) -> ListViewVector {
         let taken_offsets = self.offsets().take(indices);
         let taken_sizes = self.sizes().take(indices);
         let taken_validity = self.validity().take(indices);
@@ -46,24 +49,20 @@ impl<I: UnsignedPType> Take<[I]> for &ListViewVector {
                 taken_validity,
             )
         }
-
-        */
     }
 }
 
 fn take_nullable<I: UnsignedPType>(
-    _lvector: &ListViewVector,
-    _indices: &PVector<I>,
+    list_view: &ListViewVector,
+    indices: &PVector<I>,
 ) -> ListViewVector {
-    todo!("TODO(connor): Implement `take` for `ListViewVector` and figure out rebuilding");
-
-    /*
-
     // We ignore nullability when taking the offsets and sizes since we can let the `Mask`
     // implementation determine which elements are null.
-    let taken_offsets = lvector.offsets().take(indices.elements().as_slice());
-    let taken_sizes = lvector.sizes().take(indices.elements().as_slice());
-    let taken_validity = lvector.validity().take(indices);
+    let taken_offsets = list_view.offsets().take(indices.elements().as_slice());
+    let taken_sizes = list_view.sizes().take(indices.elements().as_slice());
+
+    // Note that this is **not** the same as the `indices: &[I]` `take` implementation above.
+    let taken_validity = list_view.validity().take(indices);
 
     debug_assert_eq!(taken_offsets.len(), taken_validity.len());
     debug_assert_eq!(taken_sizes.len(), taken_validity.len());
@@ -73,12 +72,10 @@ fn take_nullable<I: UnsignedPType>(
     // via Arc, so all view references remain valid.
     unsafe {
         ListViewVector::new_unchecked(
-            lvector.elements().clone(),
+            list_view.elements().clone(),
             taken_offsets,
             taken_sizes,
             taken_validity,
         )
     }
-
-    */
 }