Chore: Add more docs + from_option_iter (#5031)

Tracking Issue: #5028

Adds documentation examples to everything.

Also removes some of the macros that were only specific to a single
method that are probably not super useful for users.

There is still an issue where `try_into_mut` on `Mask` doesn't really
behave as it should (because we need to still change out arrow's
`BooleanBuffer` for our `BitBuffer` inside `MaskValues`, but I have
TODOs on those doctests.

---------

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

Author: connortsui20

vortex-buffer/src/bit/buf_mut.rs

@@ -1,8 +1,6 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
-// TODO(connor): The API of `BitBufferMut` should probably share more methods with `BitBuffer`.
-
 use arrow_buffer::bit_chunk_iterator::BitChunks;
 use bitvec::view::BitView;
 
@@ -150,12 +148,11 @@ impl BitBufferMut {
 
     /// Reserve additional bit capacity for the buffer.
     pub fn reserve(&mut self, additional: usize) {
-        let required_capacity = (self.offset + self.len + additional).div_ceil(8);
-        let buffer_capacity = self.buffer.capacity();
-        if required_capacity > self.buffer.capacity() {
-            let additional = required_capacity - buffer_capacity;
-            self.buffer.reserve(additional);
-        }
+        let required_bits = self.offset + self.len + additional;
+        let required_bytes = required_bits.div_ceil(8); // Rounds up.
+
+        let additional_bytes = required_bytes.saturating_sub(self.buffer.len());
+        self.buffer.reserve(additional_bytes);
     }
 
     /// Set the bit at `index` to the given boolean value.
@@ -543,6 +540,28 @@ mod tests {
         assert!(bools.value(9));
     }
 
+    #[test]
+    fn test_reserve_ensures_len_plus_additional() {
+        // This test documents the fix for the bug where reserve was incorrectly
+        // calculating additional bytes from capacity instead of len.
+
+        let mut bits = BitBufferMut::with_capacity(10);
+        assert_eq!(bits.len(), 0);
+
+        bits.reserve(100);
+
+        // Should have capacity for at least len + 100 = 0 + 100 = 100 bits.
+        assert!(bits.capacity() >= 100);
+
+        bits.append_n(true, 50);
+        assert_eq!(bits.len(), 50);
+
+        bits.reserve(100);
+
+        // Should have capacity for at least len + 100 = 50 + 100 = 150 bits.
+        assert!(bits.capacity() >= 150);
+    }
+
     #[test]
     fn test_with_offset_zero() {
         // Test basic operations when offset is 0

vortex-mask/src/lib.rs

@@ -203,6 +203,11 @@ impl MaskValues {
             MaskIter::Indices(self.indices())
         }
     }
+
+    /// Extracts the internal [`BitBuffer`].
+    pub(crate) fn into_buffer(self) -> BitBuffer {
+        self.buffer
+    }
 }
 
 impl Mask {

vortex-mask/src/mask_mut.rs

@@ -2,6 +2,7 @@
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
 use std::ops::Sub;
+use std::sync::Arc;
 
 use vortex_buffer::BitBufferMut;
 
@@ -225,9 +226,13 @@ impl Mask {
             Mask::AllTrue(len) => Ok(MaskMut::new_true(len)),
             Mask::AllFalse(len) => Ok(MaskMut::new_false(len)),
             Mask::Values(values) => {
-                // FIXME(ngates): we can never convert Arrow BooleanBuffer to ByteBufferMut,
-                //  so we have to wait until we use BitBuffer internally in MaskValues.
-                Err(Mask::Values(values))
+                // We need to check for uniqueness twice, first for the `Arc` with `try_unwrap`,
+                // then for the internal `BitBuffer` with `try_into_mut`.
+                let owned_values = Arc::try_unwrap(values).map_err(Mask::Values)?;
+                let bit_buffer = owned_values.into_buffer();
+                let mut_buffer = bit_buffer.try_into_mut().map_err(Mask::from_buffer)?;
+
+                Ok(MaskMut(Inner::Builder(mut_buffer)))
             }
         }
     }
@@ -461,8 +466,6 @@ mod tests {
     }
 
     #[test]
-    // TODO(ngates): when mask uses BitBuffer internally, into_mut should succeed
-    #[should_panic]
     fn test_round_trip_split_unsplit() {
         let mut original = MaskMut::with_capacity(20);
         // Pattern: 10 true, 10 false
@@ -531,4 +534,65 @@ mod tests {
         let frozen = mask.freeze();
         assert_eq!(frozen.true_count(), 30);
     }
+
+    #[test]
+    fn test_try_into_mut_all_variants() {
+        // Test AllTrue and AllFalse variants.
+        let mask_true = Mask::new_true(100);
+        let mut_mask_true = mask_true.try_into_mut().unwrap();
+        assert_eq!(mut_mask_true.len(), 100);
+        assert_eq!(mut_mask_true.freeze().true_count(), 100);
+
+        let mask_false = Mask::new_false(50);
+        let mut_mask_false = mask_false.try_into_mut().unwrap();
+        assert_eq!(mut_mask_false.len(), 50);
+        assert_eq!(mut_mask_false.freeze().true_count(), 0);
+    }
+
+    #[test]
+    fn test_try_into_mut_with_references() {
+        // Create a MaskValues variant.
+        let mut mask_mut = MaskMut::with_capacity(10);
+        mask_mut.append_n(true, 5);
+        mask_mut.append_n(false, 5);
+        let mask = mask_mut.freeze();
+
+        // Should succeed with unique reference (no clones).
+        let mask2 = {
+            let mut mask_mut2 = MaskMut::with_capacity(10);
+            mask_mut2.append_n(true, 5);
+            mask_mut2.append_n(false, 5);
+            mask_mut2.freeze()
+        };
+        let result = mask2.try_into_mut();
+        assert!(result.is_ok());
+        assert_eq!(result.unwrap().len(), 10);
+
+        // Should fail with shared references.
+        let _cloned = mask.clone();
+        let result = mask.try_into_mut();
+        assert!(result.is_err());
+        if let Err(returned_mask) = result {
+            assert_eq!(returned_mask.len(), 10);
+            assert_eq!(returned_mask.true_count(), 5);
+        }
+    }
+
+    #[test]
+    fn test_try_into_mut_round_trip() {
+        // Test freeze -> try_into_mut -> modify -> freeze cycle.
+        let mut original = MaskMut::with_capacity(20);
+        original.append_n(true, 10);
+        original.append_n(false, 10);
+
+        let frozen = original.freeze();
+        assert_eq!(frozen.true_count(), 10);
+
+        let mut mut_mask = frozen.try_into_mut().unwrap();
+        mut_mask.append_n(true, 5);
+        assert_eq!(mut_mask.len(), 25);
+
+        let frozen_again = mut_mask.freeze();
+        assert_eq!(frozen_again.true_count(), 15);
+    }
 }

vortex-vector/src/bool/vector.rs

@@ -11,10 +11,15 @@ use crate::VectorOps;
 
 /// An immutable vector of boolean values.
 ///
-/// The mutable equivalent of this type is [`BoolVectorMut`].
+/// `BoolVector` can be considered a borrowed / frozen version of [`BoolVectorMut`], which is
+/// created via the [`freeze`](crate::VectorMutOps::freeze) method.
+///
+/// See the documentation for [`BoolVectorMut`] for more information.
 #[derive(Debug, Clone)]
 pub struct BoolVector {
+    /// The bits that we use to represent booleans.
     pub(super) bits: BitBuffer,
+    /// The validity mask (where `true` represents an element is **not** null).
     pub(super) validity: Mask,
 }
 

vortex-vector/src/bool/vector_mut.rs

@@ -11,10 +11,62 @@ use crate::{VectorMutOps, VectorOps};
 
 /// A mutable vector of boolean values.
 ///
-/// The immutable equivalent of this type is [`BoolVector`].
+/// `BoolVectorMut` is the primary way to construct boolean vectors. It provides efficient methods
+/// for building vectors incrementally before converting them to an immutable [`BoolVector`] using
+/// the [`freeze`](crate::VectorMutOps::freeze) method.
+///
+/// # Examples
+///
+/// ## Extending and appending
+///
+/// ```
+/// use vortex_vector::{BoolVectorMut, VectorMutOps};
+///
+/// let mut vec1 = BoolVectorMut::from_iter([true, false].map(Some));
+/// let vec2 = BoolVectorMut::from_iter([true, true].map(Some)).freeze();
+///
+/// // Extend from another vector.
+/// vec1.extend_from_vector(&vec2);
+/// assert_eq!(vec1.len(), 4);
+///
+/// // Append null values.
+/// vec1.append_nulls(2);
+/// assert_eq!(vec1.len(), 6);
+/// ```
+///
+/// ## Splitting and unsplitting
+///
+/// ```
+/// use vortex_vector::{BoolVectorMut, VectorMutOps};
+///
+/// let mut vec = BoolVectorMut::from_iter([true, false, true, false, true].map(Some));
+///
+/// // Split the vector at index 3.
+/// let mut second_half = vec.split_off(3);
+/// assert_eq!(vec.len(), 3);
+/// assert_eq!(second_half.len(), 2);
+///
+/// // Rejoin the vectors.
+/// vec.unsplit(second_half);
+/// assert_eq!(vec.len(), 5);
+/// ```
+///
+/// ## Converting to immutable
+///
+/// ```
+/// use vortex_vector::{BoolVectorMut, VectorMutOps, VectorOps};
+///
+/// let mut vec = BoolVectorMut::from_iter([true, false, true].map(Some));
+///
+/// // Freeze into an immutable vector.
+/// let immutable = vec.freeze();
+/// assert_eq!(immutable.len(), 3);
+/// ```
 #[derive(Debug, Clone)]
 pub struct BoolVectorMut {
+    /// The mutable bits that we use to represent booleans.
     pub(super) bits: BitBufferMut,
+    /// The validity mask (where `true` represents an element is **not** null).
     pub(super) validity: MaskMut,
 }
 
@@ -28,6 +80,76 @@ impl BoolVectorMut {
     }
 }
 
+impl FromIterator<Option<bool>> for BoolVectorMut {
+    /// Creates a new [`BoolVectorMut`] from an iterator of `Option<bool>` values.
+    ///
+    /// `None` values will be marked as invalid in the validity mask.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// use vortex_vector::{BoolVectorMut, VectorMutOps};
+    ///
+    /// let mut vec = BoolVectorMut::from_iter([Some(true), None, Some(false)]);
+    /// assert_eq!(vec.len(), 3);
+    /// ```
+    fn from_iter<I>(iter: I) -> Self
+    where
+        I: IntoIterator<Item = Option<bool>>,
+    {
+        let iter = iter.into_iter();
+        let (lower_bound, _) = iter.size_hint();
+
+        let mut bits = Vec::with_capacity(lower_bound);
+        let mut validity = MaskMut::with_capacity(lower_bound);
+
+        for opt_val in iter {
+            match opt_val {
+                Some(val) => {
+                    bits.push(val);
+                    validity.append_n(true, 1);
+                }
+                None => {
+                    bits.push(false); // Value doesn't matter for invalid entries.
+                    validity.append_n(false, 1);
+                }
+            }
+        }
+
+        BoolVectorMut {
+            bits: BitBufferMut::from_iter(bits),
+            validity,
+        }
+    }
+}
+
+impl FromIterator<bool> for BoolVectorMut {
+    /// Creates a new [`BoolVectorMut`] from an iterator of `bool` values.
+    ///
+    /// All values will be treated as non-null.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// use vortex_vector::{BoolVectorMut, VectorMutOps};
+    ///
+    /// let mut vec = BoolVectorMut::from_iter([true, false, false, true]);
+    /// assert_eq!(vec.len(), 4);
+    /// ```
+    fn from_iter<I>(iter: I) -> Self
+    where
+        I: IntoIterator<Item = bool>,
+    {
+        let buffer = BitBufferMut::from_iter(iter);
+        let validity = MaskMut::new_true(buffer.len());
+
+        BoolVectorMut {
+            bits: buffer,
+            validity,
+        }
+    }
+}
+
 impl VectorMutOps for BoolVectorMut {
     type Immutable = BoolVector;
 
@@ -75,3 +197,54 @@ impl VectorMutOps for BoolVectorMut {
         self.validity.unsplit(other.validity);
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_from_iter_with_options() {
+        // Test FromIterator<Option<bool>> with nulls and empty.
+        let vec_empty = BoolVectorMut::from_iter(std::iter::empty::<Option<bool>>());
+        assert_eq!(vec_empty.len(), 0);
+
+        let vec = BoolVectorMut::from_iter([Some(true), None, Some(false), None, Some(true)]);
+        assert_eq!(vec.len(), 5);
+        let frozen = vec.freeze();
+        assert_eq!(frozen.validity().true_count(), 3);
+    }
+
+    #[test]
+    fn test_from_iter_non_null() {
+        // Test FromIterator<bool> creates all-valid vector.
+        let vec = BoolVectorMut::from_iter([true, false, true, true, false]);
+        assert_eq!(vec.len(), 5);
+        let frozen = vec.freeze();
+        assert_eq!(frozen.validity().true_count(), 5);
+    }
+
+    #[test]
+    fn test_operations_preserve_validity() {
+        // Comprehensive test for split/unsplit/extend preserving validity.
+        let mut vec = BoolVectorMut::from_iter([Some(true), None, Some(false), None, Some(true)]);
+
+        // Test split.
+        let second_half = vec.split_off(2);
+        assert_eq!(vec.len(), 2);
+        assert_eq!(second_half.len(), 3);
+
+        // Test validity after split.
+        let frozen_first = vec.freeze();
+        assert_eq!(frozen_first.validity().true_count(), 1);
+        let frozen_second = second_half.freeze();
+        assert_eq!(frozen_second.validity().true_count(), 2);
+
+        // Test unsplit.
+        let mut vec1 = BoolVectorMut::from_iter([Some(true), None]);
+        let vec2 = BoolVectorMut::from_iter([Some(false), Some(true)]);
+        vec1.unsplit(vec2);
+        assert_eq!(vec1.len(), 4);
+        let frozen = vec1.freeze();
+        assert_eq!(frozen.validity().true_count(), 3);
+    }
+}