Add arithmetic buffer compute

Signed-off-by: Nicholas Gates <[email protected]>

Author: gatesn

vortex-array/src/arrays/primitive/vtable/mod.rs

@@ -8,7 +8,7 @@ use crate::{EncodingId, EncodingRef, vtable};
 mod array;
 mod canonical;
 mod operations;
-mod pipeline;
+mod operator;
 mod serde;
 mod validity;
 mod visitor;
@@ -26,8 +26,8 @@ impl VTable for PrimitiveVTable {
     type VisitorVTable = Self;
     type ComputeVTable = NotSupported;
     type EncodeVTable = NotSupported;
-    type OperatorVTable = Self;
     type SerdeVTable = Self;
+    type OperatorVTable = Self;
 
     fn id(_encoding: &Self::Encoding) -> EncodingId {
         EncodingId::new_ref("vortex.primitive")

vortex-array/src/arrays/primitive/vtable/operator.rs

@@ -0,0 +1,37 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+use vortex_compute::filter::Filter;
+use vortex_dtype::match_each_native_ptype;
+use vortex_error::VortexResult;
+use vortex_vector::PVector;
+
+use crate::ArrayRef;
+use crate::arrays::{PrimitiveArray, PrimitiveVTable};
+use crate::execution::{BatchKernelRef, BindCtx, kernel};
+use crate::vtable::{OperatorVTable, ValidityHelper};
+
+impl OperatorVTable<PrimitiveVTable> for PrimitiveVTable {
+    fn bind(
+        array: &PrimitiveArray,
+        selection: Option<&ArrayRef>,
+        ctx: &mut dyn BindCtx,
+    ) -> VortexResult<BatchKernelRef> {
+        let mask = ctx.bind_selection(array.len(), selection)?;
+        let validity = ctx.bind_validity(array.validity(), array.len(), selection)?;
+
+        match_each_native_ptype!(array.ptype(), |T| {
+            let elements = array.buffer::<T>();
+            Ok(kernel(move || {
+                let mask = mask.execute()?;
+                let validity = validity.execute()?;
+
+                // Note that validity already has the mask applied so we only need to apply it to
+                // the elements.
+                let elements = elements.filter(&mask);
+
+                Ok(PVector::new(elements, validity).into())
+            }))
+        })
+    }
+}

vortex-array/src/arrays/primitive/vtable/pipeline.rs

@@ -31,6 +31,9 @@ use crate::{BitBuffer, BufferMut, ByteBufferMut, buffer_mut};
 #[derive(Debug, Clone, Eq)]
 pub struct BitBufferMut {
     buffer: ByteBufferMut,
+    /// Represents the offset of the bit buffer into the first byte.
+    ///
+    /// This is always less than 8 (for when the bit buffer is not aligned to a byte).
     offset: usize,
     len: usize,
 }
@@ -162,6 +165,13 @@ impl BitBufferMut {
         self.buffer.reserve(additional_bytes);
     }
 
+    /// Clears the bit buffer (but keeps any allocated memory).
+    pub fn clear(&mut self) {
+        // Since there are no items we need to drop, we simply set the length to 0.
+        self.len = 0;
+        self.offset = 0;
+    }
+
     /// Set the bit at `index` to the given boolean value.
     ///
     /// This operation is checked so if `index` exceeds the buffer length, this will panic.

vortex-buffer/src/bit/buf_mut.rs

@@ -5,6 +5,7 @@ use std::ops::Sub;
 use std::sync::Arc;
 
 use vortex_buffer::BitBufferMut;
+use vortex_error::vortex_panic;
 
 use crate::Mask;
 
@@ -56,6 +57,29 @@ impl MaskMut {
         })
     }
 
+    /// Returns the boolean value at a given index.
+    ///
+    /// # Panics
+    ///
+    /// Panics if the index is out of bounds.
+    pub fn value(&self, index: usize) -> bool {
+        match &self.0 {
+            Inner::Empty { .. } => {
+                vortex_panic!("index out of bounds: the length is 0 but the index is {index}")
+            }
+            Inner::Constant { value, len, .. } => {
+                assert!(
+                    index < *len,
+                    "index out of bounds: the length is {} but the index is {index}",
+                    *len
+                );
+
+                *value
+            }
+            Inner::Builder(bit_buffer) => bit_buffer.value(index),
+        }
+    }
+
     /// Reserve capacity for at least `additional` more values to be appended.
     pub fn reserve(&mut self, additional: usize) {
         match &mut self.0 {
@@ -71,6 +95,39 @@ impl MaskMut {
         }
     }
 
+    /// Clears the mask.
+    ///
+    /// Note that this method has no effect on the allocated capacity of the mask.
+    pub fn clear(&mut self) {
+        match &mut self.0 {
+            Inner::Empty { .. } => {}
+            Inner::Constant { capacity, .. } => {
+                self.0 = Inner::Empty {
+                    capacity: *capacity,
+                }
+            }
+            Inner::Builder(bit_buffer) => bit_buffer.clear(),
+        };
+    }
+
+    /// Shortens the mask, keeping the first `len` bits.
+    ///
+    /// If `len` is greater or equal to the vector’s current length, this has no effect.
+    ///
+    /// Note that this method has no effect on the allocated capacity of the mask.
+    pub fn truncate(&mut self, len: usize) {
+        let truncated_len = len;
+        if truncated_len > self.len() {
+            return;
+        }
+
+        match &mut self.0 {
+            Inner::Empty { .. } => {}
+            Inner::Constant { len, .. } => *len = truncated_len.min(*len),
+            Inner::Builder(bit_buffer) => bit_buffer.truncate(truncated_len),
+        };
+    }
+
     /// Append n values to the mask.
     pub fn append_n(&mut self, new_value: bool, n: usize) {
         match &mut self.0 {

vortex-mask/src/mask_mut.rs

@@ -1,12 +1,12 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
-//! [`FromIterator`] and related implementations for [`BoolVectorMut`].
+//! Iterator implementations for [`BoolVectorMut`].
 
 use vortex_buffer::BitBufferMut;
 use vortex_mask::MaskMut;
 
-use crate::BoolVectorMut;
+use crate::{BoolVectorMut, VectorMutOps};
 
 impl FromIterator<Option<bool>> for BoolVectorMut {
     /// Creates a new [`BoolVectorMut`] from an iterator of `Option<bool>` values.
@@ -78,3 +78,63 @@ impl FromIterator<bool> for BoolVectorMut {
         }
     }
 }
+
+/// Iterator over a [`BoolVectorMut`] that yields [`Option<bool>`] values.
+///
+/// This iterator is created by calling [`IntoIterator::into_iter`] on a [`BoolVectorMut`].
+///
+/// It consumes the mutable vector and iterates over the elements, yielding `None` for null values
+/// and `Some(value)` for valid values.
+pub struct BoolVectorMutIterator {
+    /// The vector being iterated over.
+    vector: BoolVectorMut,
+    /// The current index into the vector.
+    index: usize,
+}
+
+impl Iterator for BoolVectorMutIterator {
+    type Item = Option<bool>;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        (self.index < self.vector.len()).then(|| {
+            let value = self
+                .vector
+                .validity
+                .value(self.index)
+                .then(|| self.vector.bits.value(self.index));
+            self.index += 1;
+            value
+        })
+    }
+
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        let remaining = self.vector.len() - self.index;
+        (remaining, Some(remaining))
+    }
+}
+
+impl IntoIterator for BoolVectorMut {
+    type Item = Option<bool>;
+    type IntoIter = BoolVectorMutIterator;
+
+    /// Converts the mutable vector into an iterator over `Option<bool>` values.
+    ///
+    /// This method consumes the `BoolVectorMut` and returns an iterator that yields `None` for
+    /// null values and `Some(value)` for valid values.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// use vortex_vector::BoolVectorMut;
+    ///
+    /// let vec = BoolVectorMut::from_iter([Some(true), None, Some(false), Some(true)]);
+    /// let collected: Vec<_> = vec.into_iter().collect();
+    /// assert_eq!(collected, vec![Some(true), None, Some(false), Some(true)]);
+    /// ```
+    fn into_iter(self) -> Self::IntoIter {
+        BoolVectorMutIterator {
+            vector: self,
+            index: 0,
+        }
+    }
+}

vortex-vector/src/bool/from_iter.rs vortex-vector/src/bool/iter.rsvortex-vector/src/bool/from_iter.rs renamed to vortex-vector/src/bool/iter.rs

@@ -9,7 +9,7 @@ pub use vector::BoolVector;
 mod vector_mut;
 pub use vector_mut::BoolVectorMut;
 
-mod from_iter;
+mod iter;
 
 use crate::{Vector, VectorMut};
 

vortex-vector/src/bool/mod.rs

@@ -224,4 +224,41 @@ mod tests {
         let frozen = vec1.freeze();
         assert_eq!(frozen.validity().true_count(), 3);
     }
+
+    #[test]
+    fn test_into_iter_roundtrip() {
+        // Test that from_iter followed by into_iter preserves the data.
+        let original_data = vec![
+            Some(true),
+            None,
+            Some(false),
+            Some(true),
+            None,
+            Some(false),
+            None,
+            Some(true),
+        ];
+
+        // Create vector from iterator.
+        let vec = BoolVectorMut::from_iter(original_data.clone());
+
+        // Convert back to iterator and collect.
+        let roundtrip: Vec<_> = vec.into_iter().collect();
+
+        // Should be identical.
+        assert_eq!(roundtrip, original_data);
+
+        // Also test with all valid values.
+        let all_valid = vec![true, false, true, false, true];
+        let vec = BoolVectorMut::from_iter(all_valid.clone());
+        let roundtrip: Vec<_> = vec.into_iter().collect();
+        let expected: Vec<_> = all_valid.into_iter().map(Some).collect();
+        assert_eq!(roundtrip, expected);
+
+        // Test with empty.
+        let empty: Vec<Option<bool>> = vec![];
+        let vec = BoolVectorMut::from_iter(empty.clone());
+        let roundtrip: Vec<_> = vec.into_iter().collect();
+        assert_eq!(roundtrip, empty);
+    }
 }