vortex_vector/primitive/
iter.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4//! Iterator implementations for [`PVectorMut`].
5
6use vortex_dtype::NativePType;
7
8use crate::VectorMutOps;
9use crate::primitive::PVectorMut;
10
11impl<T: NativePType> Extend<Option<T>> for PVectorMut<T> {
12    /// Extends the vector from an iterator of optional values.
13    ///
14    /// `None` values will be marked as null in the validity mask.
15    ///
16    /// # Examples
17    ///
18    /// ```
19    /// use vortex_vector::primitive::PVectorMut;
20    /// use vortex_vector::{VectorMutOps, VectorOps};
21    ///
22    /// let mut vec = PVectorMut::from_iter([Some(1i32), None]);
23    /// vec.extend([Some(3), None, Some(5)]);
24    /// assert_eq!(vec.len(), 5);
25    ///
26    /// let frozen = vec.freeze();
27    /// assert_eq!(frozen.validity().true_count(), 3); // Only 3 non-null values.
28    /// ```
29    fn extend<I: IntoIterator<Item = Option<T>>>(&mut self, iter: I) {
30        let iter = iter.into_iter();
31        // Since we do not know the length of the iterator, we can only guess how much memory we
32        // need to reserve. Note that these hints may be inaccurate.
33        let (lower_bound, _) = iter.size_hint();
34
35        // We choose not to use the optional upper bound size hint to match the standard library.
36
37        self.reserve(lower_bound);
38
39        // We have to update validity per-element since it depends on Option variant.
40        for opt_val in iter {
41            match opt_val {
42                Some(val) => {
43                    self.elements.push(val);
44                    self.validity.append_n(true, 1);
45                }
46                None => {
47                    self.elements.push(T::default());
48                    self.validity.append_n(false, 1);
49                }
50            }
51        }
52    }
53}
54
55impl<T: NativePType> FromIterator<Option<T>> for PVectorMut<T> {
56    /// Creates a new [`PVectorMut<T>`] from an iterator of `Option<T>` values.
57    ///
58    /// `None` values will be marked as invalid in the validity mask.
59    ///
60    /// Internally, this uses the [`Extend<Option<T>>`] trait implementation.
61    fn from_iter<I>(iter: I) -> Self
62    where
63        I: IntoIterator<Item = Option<T>>,
64    {
65        let iter = iter.into_iter();
66
67        let mut vec = Self::with_capacity(iter.size_hint().0);
68        vec.extend(iter);
69
70        vec
71    }
72}
73
74impl<T: NativePType> Extend<T> for PVectorMut<T> {
75    /// Extends the vector from an iterator of values.
76    ///
77    /// All values from the iterator will be marked as non-null in the validity mask.
78    ///
79    /// Internally, this uses the [`Extend<T>`] trait implementation.
80    fn extend<I: IntoIterator<Item = T>>(&mut self, iter: I) {
81        let start_len = self.len();
82
83        // Allow the `BufferMut` implementation to handle extending efficiently.
84        self.elements.extend(iter);
85        self.validity.append_n(true, self.len() - start_len);
86    }
87}
88
89impl<T: NativePType> FromIterator<T> for PVectorMut<T> {
90    /// Creates a new [`PVectorMut<T>`] from an iterator of `T` values.
91    ///
92    /// All values will be treated as non-null.
93    ///
94    /// # Examples
95    ///
96    /// ```
97    /// use vortex_vector::primitive::PVectorMut;
98    /// use vortex_vector::VectorMutOps;
99    ///
100    /// let mut vec = PVectorMut::from_iter([1i32, 2, 3, 4]);
101    /// assert_eq!(vec.len(), 4);
102    /// ```
103    fn from_iter<I>(iter: I) -> Self
104    where
105        I: IntoIterator<Item = T>,
106    {
107        let iter = iter.into_iter();
108
109        let mut vec = Self::with_capacity(iter.size_hint().0);
110        vec.extend(iter);
111
112        vec
113    }
114}
115
116/// Iterator over a [`PVectorMut<T>`] that yields [`Option<T>`] values.
117///
118/// This iterator is created by calling [`IntoIterator::into_iter`] on a [`PVectorMut<T>`].
119///
120/// It consumes the mutable vector and iterates over the elements, yielding `None` for null values
121/// and `Some(value)` for valid values.
122#[derive(Debug)]
123pub struct PVectorMutIterator<T: NativePType> {
124    /// The vector being iterated over.
125    vector: PVectorMut<T>,
126    /// The current index into the vector.
127    index: usize,
128}
129
130impl<T: NativePType> Iterator for PVectorMutIterator<T> {
131    type Item = Option<T>;
132
133    fn next(&mut self) -> Option<Self::Item> {
134        (self.index < self.vector.len()).then(|| {
135            let value = self
136                .vector
137                .validity
138                .value(self.index)
139                .then(|| self.vector.elements[self.index]);
140            self.index += 1;
141            value
142        })
143    }
144
145    fn size_hint(&self) -> (usize, Option<usize>) {
146        let remaining = self.vector.len() - self.index;
147        (remaining, Some(remaining))
148    }
149}
150
151impl<T: NativePType> IntoIterator for PVectorMut<T> {
152    type Item = Option<T>;
153    type IntoIter = PVectorMutIterator<T>;
154
155    /// Converts the mutable vector into an iterator over [`Option<T>`] values.
156    ///
157    /// This method consumes the [`PVectorMut<T>`] and returns an iterator that yields `None` for
158    /// null values and `Some(value)` for valid values.
159    ///
160    /// # Examples
161    ///
162    /// ```
163    /// use vortex_vector::primitive::PVectorMut;
164    ///
165    /// let vec = PVectorMut::<i32>::from_iter([Some(1), None, Some(3), Some(4)]);
166    /// let collected: Vec<_> = vec.into_iter().collect();
167    /// assert_eq!(collected, vec![Some(1), None, Some(3), Some(4)]);
168    /// ```
169    fn into_iter(self) -> Self::IntoIter {
170        PVectorMutIterator {
171            vector: self,
172            index: 0,
173        }
174    }
175}
vortex_vector/primitive/iter.rs

vortex_vector/primitive/
iter.rs
