add view vector take stubs

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

Author: connortsui20

vortex-compute/src/take/bit_buffer.rs

@@ -58,3 +58,29 @@ fn take_bool<I: UnsignedPType>(bools: &BitBuffer, indices: &[I]) -> BitBuffer {
         get_bit(buffer, offset + bool_idx)
     })
 }
+
+#[cfg(test)]
+mod tests {
+    use crate::take::Take;
+
+    #[test]
+    fn test_bit_buffer_take_small_and_large() {
+        use vortex_buffer::BitBuffer;
+
+        // Small buffer (uses take_byte_bool path).
+        let small: BitBuffer = [true, false, true, true, false, true, false, false]
+            .into_iter()
+            .collect();
+        let result = (&small).take(&[7u32, 0, 2, 5, 1][..]);
+
+        let values: Vec<bool> = (0..result.len()).map(|i| result.value(i)).collect();
+        assert_eq!(values, vec![false, true, true, true, false]);
+
+        // Large buffer (uses take_bool path, len > 4096).
+        let large: BitBuffer = (0..5000).map(|i| i % 3 == 0).collect();
+        let result = (&large).take(&[4999u32, 0, 1, 2, 3, 4998][..]);
+
+        let values: Vec<bool> = (0..result.len()).map(|i| result.value(i)).collect();
+        assert_eq!(values, vec![false, true, false, false, true, true]);
+    }
+}

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

@@ -0,0 +1,89 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+// use std::ops::Deref;
+
+// 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::BinaryViewType;
+use vortex_vector::binaryview::BinaryViewVector;
+use vortex_vector::primitive::PVector;
+
+use crate::take::Take;
+
+impl<T: BinaryViewType, I: UnsignedPType> Take<PVector<I>> for &BinaryViewVector<T> {
+    type Output = BinaryViewVector<T>;
+
+    fn take(self, indices: &PVector<I>) -> BinaryViewVector<T> {
+        if indices.validity().all_true() {
+            self.take(indices.elements().as_slice())
+        } else {
+            take_nullable(self, indices)
+        }
+    }
+}
+
+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");
+
+        /*
+
+        let taken_views = take_views(self.views(), indices);
+        let taken_validity = self.validity().take(indices);
+
+        debug_assert_eq!(taken_views.len(), taken_validity.len());
+
+        // SAFETY: We called take on views and validity with the same indices, so the new components
+        // must 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, self.buffers().clone(), taken_validity)
+        }
+
+        */
+    }
+}
+
+fn take_nullable<T: BinaryViewType, I: UnsignedPType>(
+    _bvector: &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);
+
+    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)
+    }
+
+    */
+}
+
+/*
+
+/// 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_()]))
+}
+
+*/

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

@@ -0,0 +1,84 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+use vortex_dtype::UnsignedPType;
+use vortex_vector::VectorOps;
+use vortex_vector::listview::ListViewVector;
+use vortex_vector::primitive::PVector;
+
+use crate::take::Take;
+
+impl<I: UnsignedPType> Take<PVector<I>> for &ListViewVector {
+    type Output = ListViewVector;
+
+    fn take(self, indices: &PVector<I>) -> ListViewVector {
+        if indices.validity().all_true() {
+            self.take(indices.elements().as_slice())
+        } else {
+            take_nullable(self, indices)
+        }
+    }
+}
+
+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");
+
+        /*
+
+        let taken_offsets = self.offsets().take(indices);
+        let taken_sizes = self.sizes().take(indices);
+        let taken_validity = self.validity().take(indices);
+
+        debug_assert_eq!(taken_offsets.len(), taken_validity.len());
+        debug_assert_eq!(taken_sizes.len(), taken_validity.len());
+
+        // SAFETY: We called take on offsets, sizes, and validity with the same indices, so the new
+        // components must have the same length. The offsets and sizes still point into the same
+        // elements array which we clone via Arc, so all view references remain valid.
+        unsafe {
+            ListViewVector::new_unchecked(
+                self.elements().clone(),
+                taken_offsets,
+                taken_sizes,
+                taken_validity,
+            )
+        }
+
+        */
+    }
+}
+
+fn take_nullable<I: UnsignedPType>(
+    _lvector: &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);
+
+    debug_assert_eq!(taken_offsets.len(), taken_validity.len());
+    debug_assert_eq!(taken_sizes.len(), taken_validity.len());
+
+    // SAFETY: We used the same indices to take from all components, so they should still have the
+    // same length. The offsets and sizes still point into the same elements array which we clone
+    // via Arc, so all view references remain valid.
+    unsafe {
+        ListViewVector::new_unchecked(
+            lvector.elements().clone(),
+            taken_offsets,
+            taken_sizes,
+            taken_validity,
+        )
+    }
+
+    */
+}

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

@@ -1,19 +1,45 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
+use vortex_dtype::UnsignedPType;
+use vortex_vector::Vector;
 use vortex_vector::match_each_unsigned_pvector;
+use vortex_vector::match_each_vector;
 use vortex_vector::primitive::PVector;
 use vortex_vector::primitive::PrimitiveVector;
 
 use crate::take::Take;
 
+mod binaryview;
 mod bool;
+mod decimal;
+mod dvector;
+mod fixed_size_list;
+mod listview;
+mod null;
 mod primitive;
 mod pvector;
+mod struct_;
 
 #[cfg(test)]
 mod tests;
 
+impl<I: UnsignedPType> Take<PVector<I>> for &Vector {
+    type Output = Vector;
+
+    fn take(self, indices: &PVector<I>) -> Vector {
+        match_each_vector!(self, |v| { v.take(indices).into() })
+    }
+}
+
+impl<I: UnsignedPType> Take<[I]> for &Vector {
+    type Output = Vector;
+
+    fn take(self, indices: &[I]) -> Vector {
+        match_each_vector!(self, |v| { v.take(indices).into() })
+    }
+}
+
 impl<T> Take<PrimitiveVector> for &T
 where
     for<'a> &'a T: Take<PVector<u8>, Output = T>,