Feature: Add IntoIter implementation for vectors (#5093)

Adds `IntoIter` impls for primitive and bool vector.

Consolidates all the extend, from_iter, and (new) into_iter stuff into
`iter.rs` files.

The non-cosmetic changes are at the bottom of the `iter.rs` files.

There is a world where we could use the internal `BufferIterator` types
but this was way simpler.

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

Author: connortsui20

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

@@ -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/mod.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/vector_mut.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);
+    }
 }

vortex-vector/src/primitive/generic.rs

@@ -69,6 +69,11 @@ impl<T: NativePType> PVector<T> {
         Self { elements, validity }
     }
 
+    /// Decomposes the primitive vector into its constituent parts.
+    pub fn into_parts(self) -> (Buffer<T>, Mask) {
+        (self.elements, self.validity)
+    }
+
     /// Gets a nullable element at the given index.
     ///
     /// If the element at the given index is null, returns `None`. Otherwise, returns `Some(x)`,

vortex-vector/src/primitive/generic_mut.rs

@@ -159,6 +159,11 @@ impl<T> PVectorMut<T> {
             validity: MaskMut::with_capacity(capacity),
         }
     }
+
+    /// Decomposes the primitive vector into its constituent parts.
+    pub fn into_parts(self) -> (BufferMut<T>, MaskMut) {
+        (self.elements, self.validity)
+    }
 }
 
 impl<T: NativePType> VectorMutOps for PVectorMut<T> {

vortex-vector/src/primitive/iter.rs

@@ -0,0 +1,171 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+//! Iterator implementations for [`PVectorMut`].
+
+use vortex_dtype::NativePType;
+
+use crate::{PVectorMut, VectorMutOps};
+
+impl<T: NativePType> Extend<Option<T>> for PVectorMut<T> {
+    /// Extends the vector from an iterator of optional values.
+    ///
+    /// `None` values will be marked as null in the validity mask.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// use vortex_vector::{PVectorMut, VectorMutOps, VectorOps};
+    ///
+    /// let mut vec = PVectorMut::from_iter([Some(1i32), None]);
+    /// vec.extend([Some(3), None, Some(5)]);
+    /// assert_eq!(vec.len(), 5);
+    ///
+    /// let frozen = vec.freeze();
+    /// assert_eq!(frozen.validity().true_count(), 3); // Only 3 non-null values.
+    /// ```
+    fn extend<I: IntoIterator<Item = Option<T>>>(&mut self, iter: I) {
+        let iter = iter.into_iter();
+        // Since we do not know the length of the iterator, we can only guess how much memory we
+        // need to reserve. Note that these hints may be inaccurate.
+        let (lower_bound, _) = iter.size_hint();
+
+        // We choose not to use the optional upper bound size hint to match the standard library.
+
+        self.reserve(lower_bound);
+
+        // We have to update validity per-element since it depends on Option variant.
+        for opt_val in iter {
+            match opt_val {
+                Some(val) => {
+                    self.elements.push(val);
+                    self.validity.append_n(true, 1);
+                }
+                None => {
+                    self.elements.push(T::default());
+                    self.validity.append_n(false, 1);
+                }
+            }
+        }
+    }
+}
+
+impl<T: NativePType> FromIterator<Option<T>> for PVectorMut<T> {
+    /// Creates a new [`PVectorMut<T>`] from an iterator of `Option<T>` values.
+    ///
+    /// `None` values will be marked as invalid in the validity mask.
+    ///
+    /// Internally, this uses the [`Extend<Option<T>>`] trait implementation.
+    fn from_iter<I>(iter: I) -> Self
+    where
+        I: IntoIterator<Item = Option<T>>,
+    {
+        let iter = iter.into_iter();
+
+        let mut vec = Self::with_capacity(iter.size_hint().0);
+        vec.extend(iter);
+
+        vec
+    }
+}
+
+impl<T: NativePType> Extend<T> for PVectorMut<T> {
+    /// Extends the vector from an iterator of values.
+    ///
+    /// All values from the iterator will be marked as non-null in the validity mask.
+    ///
+    /// Internally, this uses the [`Extend<T>`] trait implementation.
+    fn extend<I: IntoIterator<Item = T>>(&mut self, iter: I) {
+        let start_len = self.len();
+
+        // Allow the `BufferMut` implementation to handle extending efficiently.
+        self.elements.extend(iter);
+        self.validity.append_n(true, self.len() - start_len);
+    }
+}
+
+impl<T: NativePType> FromIterator<T> for PVectorMut<T> {
+    /// Creates a new [`PVectorMut<T>`] from an iterator of `T` values.
+    ///
+    /// All values will be treated as non-null.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// use vortex_vector::{PVectorMut, VectorMutOps};
+    ///
+    /// let mut vec = PVectorMut::from_iter([1i32, 2, 3, 4]);
+    /// assert_eq!(vec.len(), 4);
+    /// ```
+    fn from_iter<I>(iter: I) -> Self
+    where
+        I: IntoIterator<Item = T>,
+    {
+        let iter = iter.into_iter();
+
+        let mut vec = Self::with_capacity(iter.size_hint().0);
+        vec.extend(iter);
+
+        vec
+    }
+}
+
+/// Iterator over a [`PVectorMut<T>`] that yields [`Option<T>`] values.
+///
+/// This iterator is created by calling [`IntoIterator::into_iter`] on a [`PVectorMut<T>`].
+///
+/// It consumes the mutable vector and iterates over the elements, yielding `None` for null values
+/// and `Some(value)` for valid values.
+pub struct PVectorMutIterator<T: NativePType> {
+    /// The vector being iterated over.
+    vector: PVectorMut<T>,
+    /// The current index into the vector.
+    index: usize,
+}
+
+impl<T: NativePType> Iterator for PVectorMutIterator<T> {
+    type Item = Option<T>;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        (self.index < self.vector.len()).then(|| {
+            let value = self
+                .vector
+                .validity
+                .value(self.index)
+                .then(|| self.vector.elements[self.index]);
+            self.index += 1;
+            value
+        })
+    }
+
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        let remaining = self.vector.len() - self.index;
+        (remaining, Some(remaining))
+    }
+}
+
+impl<T: NativePType> IntoIterator for PVectorMut<T> {
+    type Item = Option<T>;
+    type IntoIter = PVectorMutIterator<T>;
+
+    /// Converts the mutable vector into an iterator over [`Option<T>`] values.
+    ///
+    /// This method consumes the [`PVectorMut<T>`] and returns an iterator that yields `None` for
+    /// null values and `Some(value)` for valid values.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// use vortex_vector::PVectorMut;
+    ///
+    /// let vec = PVectorMut::<i32>::from_iter([Some(1), None, Some(3), Some(4)]);
+    /// let collected: Vec<_> = vec.into_iter().collect();
+    /// assert_eq!(collected, vec![Some(1), None, Some(3), Some(4)]);
+    /// ```
+    fn into_iter(self) -> Self::IntoIter {
+        PVectorMutIterator {
+            vector: self,
+            index: 0,
+        }
+    }
+}

vortex-vector/src/primitive/mod.rs

@@ -23,6 +23,7 @@ pub use generic_mut::PVectorMut;
 mod vector;
 pub use vector::PrimitiveVector;
 
+mod iter;
 mod pvector_impl;
 mod vector_mut;
 pub use vector_mut::PrimitiveVectorMut;