add take for bit buffer, mask, and primitive vector

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

Author: connortsui20

vortex-compute/src/take/bit_buffer.rs

@@ -0,0 +1,57 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+use vortex_buffer::BitBuffer;
+use vortex_buffer::get_bit;
+use vortex_dtype::UnsignedPType;
+
+use crate::take::Take;
+
+impl<I: UnsignedPType> Take<[I]> for &BitBuffer {
+    type Output = BitBuffer;
+
+    fn take(self, indices: &[I]) -> BitBuffer {
+        // For boolean arrays that roughly fit into a single page (at least, on Linux), it's worth
+        // the overhead to convert to a `Vec<bool>`.
+        if self.len() <= 4096 {
+            let bools = self.iter().collect();
+            take_byte_bool(bools, indices)
+        } else {
+            take_bool(self, indices)
+        }
+    }
+}
+
+// NB: We do NOT implement `impl<I: UnsignedPType> Take<PVector<I>> for &BitBuffer`, specifically
+// because there is a very similar implementation on `Mask` that has special logic for working with
+// null indices. That logic could also be implemented on `BitBuffer`, but since it is not
+// immediately clear what should happen in the case of a null index when taking a `BitBuffer` (do
+// you set it to true or false?), we do not implement this at all.
+
+/// # Panics
+///
+/// Panics if an index is out of bounds.
+fn take_byte_bool<I: UnsignedPType>(bools: Vec<bool>, indices: &[I]) -> BitBuffer {
+    BitBuffer::collect_bool(indices.len(), |idx| {
+        // SAFETY: We are iterating within the bounds of the `indices` array, so we are always
+        // within bounds of `indices`.
+        let bool_idx = unsafe { indices.get_unchecked(idx).as_() };
+        bools[bool_idx]
+    })
+}
+
+/// # Panics
+///
+/// Panics if an index is out of bounds.
+fn take_bool<I: UnsignedPType>(bools: &BitBuffer, indices: &[I]) -> BitBuffer {
+    // We dereference to the underlying buffer to avoid incurring an access cost on every index.
+    let buffer = bools.inner().as_ref();
+    let offset = bools.offset();
+
+    BitBuffer::collect_bool(indices.len(), |idx| {
+        // SAFETY: We are iterating within the bounds of the `indices` array, so we are always
+        // within bounds.
+        let bool_idx = unsafe { indices.get_unchecked(idx).as_() };
+        get_bit(buffer, offset + bool_idx)
+    })
+}

vortex-compute/src/take/mask.rs

@@ -0,0 +1,88 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+use vortex_buffer::BitBuffer;
+use vortex_buffer::get_bit;
+use vortex_dtype::UnsignedPType;
+use vortex_mask::Mask;
+use vortex_vector::VectorOps;
+use vortex_vector::primitive::PVector;
+
+use crate::take::Take;
+
+impl<I: UnsignedPType> Take<[I]> for &Mask {
+    type Output = Mask;
+
+    fn take(self, indices: &[I]) -> Mask {
+        match self {
+            Mask::AllTrue(_) => Mask::AllTrue(indices.len()),
+            Mask::AllFalse(_) => Mask::AllFalse(indices.len()),
+            Mask::Values(mask_values) => {
+                let taken_bit_buffer = mask_values.bit_buffer().take(indices);
+                Mask::from_buffer(taken_bit_buffer)
+            }
+        }
+    }
+}
+
+impl<I: UnsignedPType> Take<PVector<I>> for &Mask {
+    type Output = Mask;
+
+    /// Implementation of take on [`Mask`] that is null-aware.
+    ///
+    /// If an index is specified as null by the [`PVector`], then the taken mask value is set to
+    /// `false`.
+    ///
+    /// This is useful for many of the `take` implementations for vectors.
+    fn take(self, indices: &PVector<I>) -> Mask {
+        let indices_validity = indices.validity();
+        let indices_len = indices.len();
+
+        match indices_validity {
+            Mask::AllTrue(_) => return self.take(indices.elements().as_slice()),
+            Mask::AllFalse(_) => return Mask::AllFalse(indices_len),
+            Mask::Values(_) => (),
+        };
+
+        let Mask::Values(indices_validity_values) = indices_validity else {
+            unreachable!("we just matched on the other cases above");
+        };
+
+        match self {
+            // Since all the values are true, the only false values will be from the indices.
+            Mask::AllTrue(_) => Mask::Values(indices_validity_values.clone()),
+            // Since all the values are already false, the indices nullability wont change anything.
+            Mask::AllFalse(_) => Mask::AllFalse(indices_len),
+            Mask::Values(mask_values) => {
+                // For boolean arrays that roughly fit into a single page (at least, on Linux), it's
+                // worth the overhead to convert to a `Vec<bool>`.
+                if self.len() <= 4096 {
+                    let bools = mask_values.bit_buffer().iter().collect();
+                    Mask::from_buffer(take_byte_bool_nullable(bools, indices))
+                } else {
+                    Mask::from_buffer(take_bool_nullable(mask_values.bit_buffer(), indices))
+                }
+            }
+        }
+    }
+}
+
+fn take_byte_bool_nullable<I: UnsignedPType>(bools: Vec<bool>, indices: &PVector<I>) -> BitBuffer {
+    BitBuffer::collect_bool(indices.len(), |idx| {
+        indices
+            .get(idx)
+            .is_some_and(|bool_idx| bools[bool_idx.as_()])
+    })
+}
+
+fn take_bool_nullable<I: UnsignedPType>(bools: &BitBuffer, indices: &PVector<I>) -> BitBuffer {
+    // We dereference to the underlying buffer to avoid incurring an access cost on every index.
+    let buffer = bools.inner().as_ref();
+    let offset = bools.offset();
+
+    BitBuffer::collect_bool(indices.len(), |idx| {
+        indices
+            .get(idx)
+            .is_some_and(|bool_idx| get_bit(buffer, offset + bool_idx.as_()))
+    })
+}

vortex-compute/src/take/mod.rs

@@ -3,8 +3,11 @@
 
 //! Take function.
 
+mod bit_buffer;
 mod buffer;
+mod mask;
 pub mod slice;
+mod vector;
 
 /// Function for taking based on indices (which can have different representations).
 pub trait Take<Indices: ?Sized> {

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

@@ -0,0 +1,5 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+mod primitive;
+mod pvector;

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

@@ -0,0 +1,39 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+use vortex_dtype::UnsignedPType;
+use vortex_vector::VectorOps;
+use vortex_vector::match_each_pvector;
+use vortex_vector::primitive::PVector;
+use vortex_vector::primitive::PrimitiveVector;
+
+use crate::take::Take;
+
+impl Take<PrimitiveVector> for &PrimitiveVector {
+    type Output = PrimitiveVector;
+
+    fn take(self, indices: &PrimitiveVector) -> PrimitiveVector {
+        match_each_pvector!(self, |v| { v.take(indices).into() })
+    }
+}
+
+impl<I: UnsignedPType> Take<PVector<I>> for &PrimitiveVector {
+    type Output = PrimitiveVector;
+
+    fn take(self, indices: &PVector<I>) -> PrimitiveVector {
+        // 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_pvector!(self, |v| { v.take(indices).into() })
+    }
+}
+
+impl<I: UnsignedPType> Take<[I]> for &PrimitiveVector {
+    type Output = PrimitiveVector;
+
+    fn take(self, indices: &[I]) -> PrimitiveVector {
+        match_each_pvector!(self, |v| { v.take(indices).into() })
+    }
+}

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

@@ -0,0 +1,62 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+use vortex_dtype::NativePType;
+use vortex_dtype::UnsignedPType;
+use vortex_vector::VectorOps;
+use vortex_vector::match_each_unsigned_pvector;
+use vortex_vector::primitive::PVector;
+use vortex_vector::primitive::PrimitiveVector;
+
+use crate::take::Take;
+
+impl<T: NativePType> Take<PrimitiveVector> for &PVector<T> {
+    type Output = PVector<T>;
+
+    fn take(self, indices: &PrimitiveVector) -> PVector<T> {
+        match_each_unsigned_pvector!(indices, |iv| { self.take(iv) })
+    }
+}
+
+impl<T: NativePType, I: UnsignedPType> Take<PVector<I>> for &PVector<T> {
+    type Output = PVector<T>;
+
+    fn take(self, indices: &PVector<I>) -> PVector<T> {
+        if indices.validity().all_true() {
+            self.take(indices.elements().as_slice())
+        } else {
+            take_nullable(self, indices)
+        }
+    }
+}
+
+impl<T: NativePType, I: UnsignedPType> Take<[I]> for &PVector<T> {
+    type Output = PVector<T>;
+
+    fn take(self, indices: &[I]) -> PVector<T> {
+        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.
+        unsafe { PVector::new_unchecked(taken_elements, taken_validity) }
+    }
+}
+
+fn take_nullable<T: NativePType, I: UnsignedPType>(
+    pvector: &PVector<T>,
+    indices: &PVector<I>,
+) -> PVector<T> {
+    // We ignore nullability when taking the elements since we can let the `Mask` implementation
+    // determine which elements are null.
+    let taken_elements = pvector.elements().take(indices.elements().as_slice());
+    let taken_validity = pvector.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.
+    unsafe { PVector::new_unchecked(taken_elements, taken_validity) }
+}