add stubs for take on vector + impl FSL, null, and struct

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

Author: connortsui20

Cargo.lock

@@ -30,6 +30,7 @@ arrow-array = { workspace = true, optional = true }
 arrow-buffer = { workspace = true, optional = true }
 arrow-schema = { workspace = true, optional = true }
 half = { workspace = true }
+itertools = { workspace = true }
 log = { workspace = true }
 multiversion = { workspace = true }
 num-traits = { workspace = true }

vortex-compute/Cargo.toml

@@ -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/bit_buffer.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/binaryview.rs

@@ -0,0 +1,31 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+use vortex_dtype::UnsignedPType;
+use vortex_vector::VectorOps;
+use vortex_vector::decimal::DecimalVector;
+use vortex_vector::match_each_dvector;
+use vortex_vector::primitive::PVector;
+
+use crate::take::Take;
+
+impl<I: UnsignedPType> Take<PVector<I>> for &DecimalVector {
+    type Output = DecimalVector;
+
+    fn take(self, indices: &PVector<I>) -> DecimalVector {
+        // If all the indices are valid, we can delegate to the slice indices implementation.
+        if indices.validity().all_true() {
+            return self.take(indices.elements().as_slice());
+        }
+
+        match_each_dvector!(self, |v| { v.take(indices).into() })
+    }
+}
+
+impl<I: UnsignedPType> Take<[I]> for &DecimalVector {
+    type Output = DecimalVector;
+
+    fn take(self, indices: &[I]) -> DecimalVector {
+        match_each_dvector!(self, |v| { v.take(indices).into() })
+    }
+}

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

@@ -0,0 +1,67 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+use vortex_dtype::NativeDecimalType;
+use vortex_dtype::UnsignedPType;
+use vortex_vector::VectorOps;
+use vortex_vector::decimal::DVector;
+use vortex_vector::primitive::PVector;
+
+use crate::take::Take;
+
+impl<D: NativeDecimalType, I: UnsignedPType> Take<PVector<I>> for &DVector<D> {
+    type Output = DVector<D>;
+
+    fn take(self, indices: &PVector<I>) -> DVector<D> {
+        if indices.validity().all_true() {
+            self.take(indices.elements().as_slice())
+        } else {
+            take_nullable(self, indices)
+        }
+    }
+}
+
+impl<D: NativeDecimalType, I: UnsignedPType> Take<[I]> for &DVector<D> {
+    type Output = DVector<D>;
+
+    fn take(self, _indices: &[I]) -> DVector<D> {
+        todo!("TODO(connor): Implement `take` for `DVector` and figure out trait bounds");
+
+        /*
+
+        let taken_elements = self.elements().take(indices);
+        let taken_validity = self.validity().take(indices);
+
+        debug_assert_eq!(taken_elements.len(), taken_validity.len());
+
+        // SAFETY: We called take on both components of the vector with the same indices, so the new
+        // components must have the same length. The elements are unchanged, so they must still be
+        // within the precision/scale bounds.
+        unsafe { DVector::new_unchecked(self.precision_scale(), taken_elements, taken_validity) }
+
+        */
+    }
+}
+
+fn take_nullable<D: NativeDecimalType, I: UnsignedPType>(
+    _dvector: &DVector<D>,
+    _indices: &PVector<I>,
+) -> DVector<D> {
+    todo!("TODO(connor): Implement `take` for `DVector` and figure out trait bounds");
+
+    /*
+
+    // We ignore nullability when taking the elements since we can let the `Mask` implementation
+    // determine which elements are null.
+    let taken_elements = dvector.elements().take(indices.elements().as_slice());
+    let taken_validity = dvector.validity().take(indices);
+
+    debug_assert_eq!(taken_elements.len(), taken_validity.len());
+
+    // SAFETY: We used the same indices to take from both components, so they should still have the
+    // same length. The elements are unchanged, so they must still be within the precision/scale
+    // bounds.
+    unsafe { DVector::new_unchecked(dvector.precision_scale(), taken_elements, taken_validity) }
+
+    */
+}