vortex_vector/primitive/
generic_mut_impl.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4//! Helper methods for [`PVectorMut<T>`] that mimic the behavior of [`std::vec::Vec`].
5
6use vortex_buffer::BufferMut;
7use vortex_dtype::NativePType;
8
9use crate::VectorMutOps;
10use crate::primitive::PVectorMut;
11
12/// Point operations for [`PVectorMut`].
13impl<T: NativePType> PVectorMut<T> {
14    /// Gets a nullable element at the given index, panicking on out-of-bounds.
15    ///
16    /// If the element at the given index is null, returns `None`. Otherwise, returns `Some(x)`,
17    /// where `x: T`.
18    ///
19    /// Note that this `get` method is different from the standard library [`slice::get`], which
20    /// returns `None` if the index is out of bounds. This method will panic if the index is out of
21    /// bounds, and return `None` if the elements is null.
22    ///
23    /// # Panics
24    ///
25    /// Panics if the index is out of bounds.
26    pub fn get(&self, index: usize) -> Option<T> {
27        self.validity.value(index).then(|| self.elements[index])
28    }
29
30    /// Pushes an element to the back of the vector.
31    ///
32    /// The element is treated as non-null.
33    pub fn push(&mut self, value: T) {
34        self.elements.push(value);
35        self.validity.append_n(true, 1);
36    }
37
38    /// Pushes an element without bounds checking.
39    ///
40    /// The element is treated as non-null.
41    ///
42    /// # Safety
43    ///
44    /// The caller must ensure that there is sufficient capacity in both elements and validity
45    /// buffers.
46    #[inline]
47    pub unsafe fn push_unchecked(&mut self, value: T) {
48        // SAFETY: The caller guarantees there is sufficient capacity in the elements buffer,
49        // so we can write to the spare capacity and increment the length without bounds checks.
50        unsafe {
51            self.elements.spare_capacity_mut()[0].write(value);
52            self.elements.set_len(self.len() + 1);
53        }
54        self.validity.append_n(true, 1);
55    }
56
57    /// Appends an optional element to the back of the vector, where `None` represents a null
58    /// element.
59    pub fn push_opt(&mut self, value: Option<T>) {
60        if let Some(value) = value {
61            self.push(value);
62        } else {
63            self.elements.push(T::default());
64            self.validity.append_n(false, 1);
65        }
66    }
67}
68
69impl<T: NativePType> AsRef<[T]> for PVectorMut<T> {
70    /// Returns an immutable slice over the internal mutable buffer with elements of type `T`.
71    ///
72    /// Note that this slice may contain garbage data where the [`validity()`] mask from the frozen
73    /// [`PVector`](crate::primitive::PVector) type states that an element is invalid.
74    ///
75    /// The caller should check the frozen [`validity()`] before performing any operations.
76    ///
77    /// [`validity()`]: crate::VectorOps::validity
78    #[inline]
79    fn as_ref(&self) -> &[T] {
80        self.elements.as_slice()
81    }
82}
83
84impl<T: NativePType> AsMut<[T]> for PVectorMut<T> {
85    /// Returns a mutable slice over the internal mutable buffer with elements of type `T`.
86    ///
87    /// Note that this slice may contain garbage data where the [`validity()`] mask from the frozen
88    /// [`PVector`](crate::primitive::PVector) type states that an element is invalid.
89    ///
90    /// The caller should check the frozen [`validity()`] before performing any operations.
91    ///
92    /// [`validity()`]: crate::VectorOps::validity
93    #[inline]
94    fn as_mut(&mut self) -> &mut [T] {
95        self.elements.as_mut_slice()
96    }
97}
98
99/// Batch operations for [`PVectorMut`].
100impl<T: NativePType> PVectorMut<T> {
101    /// Returns the internal [`BufferMut`] of the [`PVectorMut`].
102    ///
103    /// Note that the internal buffer may hold garbage data in place of nulls. That information is
104    /// tracked by the [`validity()`](Self::validity).
105    #[inline]
106    pub fn elements(&self) -> &BufferMut<T> {
107        &self.elements
108    }
109
110    /// Resizes the `Vec` in-place so that `len` is equal to `new_len`.
111    ///
112    /// If `new_len` is greater than `len`, the `Vec` is extended by the difference, with each
113    /// additional slot filled with `value`, where `None` represent a null.
114    ///
115    /// If `new_len` is less than `len`, the `Vec` is simply truncated.
116    pub fn resize(&mut self, new_len: usize, value: Option<T>) {
117        let current_len = self.len();
118
119        if new_len < current_len {
120            self.truncate(new_len);
121        } else {
122            let additional = new_len - current_len;
123
124            match value {
125                Some(value) => {
126                    self.elements.push_n(value, additional);
127                    self.validity.append_n(true, additional);
128                }
129                None => {
130                    self.elements.push_n(T::default(), additional);
131                    self.validity.append_n(false, additional);
132                }
133            }
134        }
135    }
136}
137
138#[cfg(test)]
139mod tests {
140    use super::*;
141    use crate::VectorOps;
142
143    #[test]
144    fn test_get_methods() {
145        let vec = PVectorMut::from_iter([Some(1), None, Some(3), None, Some(5)]);
146
147        // Test get_checked - bounds and nulls.
148        assert_eq!(vec.get(0), Some(1));
149        assert_eq!(vec.get(1), None);
150
151        // Test get - nulls.
152        assert_eq!(vec.get(0), Some(1));
153        assert_eq!(vec.get(1), None);
154        assert_eq!(vec.get(2), Some(3));
155
156        assert_eq!(vec.elements()[0], 1);
157        assert_eq!(vec.elements()[2], 3);
158
159        // Also test PVector methods.
160        let frozen = vec.freeze();
161        assert_eq!(frozen.get(0), Some(&1));
162        assert_eq!(frozen.get(1), None);
163        assert_eq!(frozen.get(2), Some(&3));
164    }
165
166    #[test]
167    #[should_panic(expected = "index out of bounds")]
168    fn test_get_panic() {
169        let vec = PVectorMut::from_iter([Some(1), Some(2)]);
170        let _ = vec.get(10);
171    }
172
173    #[test]
174    fn test_push_variants() {
175        let mut vec = PVectorMut::<i32>::with_capacity(10);
176        vec.push(1);
177        vec.push_opt(None);
178        vec.push_opt(Some(3));
179
180        assert_eq!(vec.len(), 3);
181        assert_eq!(vec.get(0), Some(1));
182        assert_eq!(vec.get(1), None);
183        assert_eq!(vec.get(2), Some(3));
184
185        // Test push_unchecked with pre-reserved capacity.
186        vec.reserve(1);
187        unsafe {
188            vec.push_unchecked(4);
189        }
190        assert_eq!(vec.get(3), Some(4));
191    }
192
193    #[test]
194    fn test_resize_operations() {
195        let mut vec = PVectorMut::from_iter([1i32, 2, 3]);
196
197        // Grow with valid values.
198        vec.resize(5, Some(99));
199        assert_eq!(vec.len(), 5);
200        assert_eq!(vec.get(3), Some(99));
201        assert_eq!(vec.get(4), Some(99));
202
203        // Grow with nulls.
204        vec.resize(7, None);
205        assert_eq!(vec.get(5), None);
206        assert_eq!(vec.get(6), None);
207
208        // Shrink.
209        vec.resize(2, Some(0));
210        assert_eq!(vec.len(), 2);
211        assert_eq!(vec.get(0), Some(1));
212        assert_eq!(vec.get(1), Some(2));
213    }
214
215    #[test]
216    fn test_clear_truncate() {
217        let mut vec = PVectorMut::from_iter([Some(1), None, Some(3), None, Some(5)]);
218        let cap = vec.capacity();
219
220        vec.truncate(3);
221        assert_eq!(vec.len(), 3);
222        assert!(vec.capacity() >= cap); // Capacity preserved.
223
224        vec.truncate(10); // Truncate beyond length - no-op.
225        assert_eq!(vec.len(), 3);
226
227        vec.clear();
228        assert_eq!(vec.len(), 0);
229        assert!(vec.capacity() >= cap); // Capacity still preserved.
230    }
231
232    #[test]
233    fn test_slice_access() {
234        let mut vec = PVectorMut::from_iter([Some(1i32), None, Some(3)]);
235        let slice = vec.as_ref();
236        assert_eq!(slice[0], 1);
237        assert_eq!(slice[2], 3);
238        // slice[1] is undefined for null but safe to access.
239
240        let mut_slice = vec.as_mut();
241        mut_slice[0] = 10;
242        assert_eq!(vec.get(0), Some(10));
243
244        let frozen = vec.freeze();
245        assert_eq!(frozen.as_ref()[0], 10);
246    }
247
248    #[test]
249    fn test_from_iter_variants() {
250        // FromIterator<T> - all non-null.
251        let vec1 = PVectorMut::from_iter([1i32, 2, 3]);
252        assert_eq!(vec1.len(), 3);
253        assert!(vec1.freeze().validity().all_true());
254
255        // FromIterator<Option<T>> - mixed null/non-null.
256        let vec2 = PVectorMut::from_iter([Some(1i32), None, Some(3)]);
257        assert_eq!(vec2.len(), 3);
258        assert_eq!(vec2.freeze().validity().true_count(), 2);
259
260        // Empty iterators.
261        let empty1 = PVectorMut::from_iter::<[i32; 0]>([]);
262        let empty2 = PVectorMut::<i32>::from_iter(std::iter::empty::<Option<i32>>());
263        assert_eq!(empty1.len(), 0);
264        assert_eq!(empty2.len(), 0);
265    }
266
267    #[test]
268    fn test_extend_operations() {
269        let mut vec = PVectorMut::from_iter([1i32, 2]);
270
271        // Extend<T> - all non-null.
272        vec.extend([3, 4]);
273        assert_eq!(vec.len(), 4);
274        assert_eq!(vec.get(3), Some(4));
275
276        // Extend<Option<T>> - mixed null/non-null.
277        vec.extend([Some(5), None, Some(7)]);
278        assert_eq!(vec.len(), 7);
279        assert_eq!(vec.get(5), None);
280        assert_eq!(vec.get(6), Some(7));
281
282        // Extend with iterator that has size hint.
283        let iter = 8..10;
284        vec.extend(iter);
285        assert_eq!(vec.get(8), Some(9));
286    }
287
288    #[test]
289    fn test_empty_vector_edge_cases() {
290        let empty = PVectorMut::<i32>::with_capacity(0);
291        assert_eq!(empty.len(), 0);
292        assert_eq!(empty.as_ref().len(), 0);
293
294        let mut mutable_empty = PVectorMut::<i32>::with_capacity(0);
295        mutable_empty.clear(); // No-op on empty.
296        mutable_empty.truncate(0); // No-op.
297        mutable_empty.resize(0, None); // No-op.
298        assert_eq!(mutable_empty.len(), 0);
299    }
300
301    #[test]
302    fn test_complex_workflow() {
303        // Integration test combining multiple operations.
304        let mut vec = PVectorMut::<i32>::with_capacity(2);
305        vec.extend([1, 2]); // Extend<T>.
306        vec.push_opt(None);
307        vec.resize(5, Some(99));
308        vec.truncate(4);
309        vec.extend([Some(10), None]); // Extend<Option<T>>.
310
311        assert_eq!(vec.len(), 6);
312        let frozen = vec.freeze();
313        assert_eq!(frozen.validity().true_count(), 4);
314        assert_eq!(frozen.get(0), Some(&1));
315        assert_eq!(frozen.get(2), None);
316        assert_eq!(frozen.get(3), Some(&99));
317        assert_eq!(frozen.get(5), None);
318    }
319
320    #[test]
321    fn test_into_iter_roundtrip() {
322        // Test that from_iter followed by into_iter preserves the data.
323        let original_data = vec![
324            Some(1i32),
325            None,
326            Some(3),
327            Some(4),
328            None,
329            Some(6),
330            None,
331            Some(8),
332        ];
333
334        // Create vector from iterator.
335        let vec = PVectorMut::<i32>::from_iter(original_data.clone());
336
337        // Convert back to iterator and collect.
338        let roundtrip: Vec<_> = vec.into_iter().collect();
339
340        // Should be identical.
341        assert_eq!(roundtrip, original_data);
342
343        // Also test with all valid values.
344        let all_valid = vec![1, 2, 3, 4, 5];
345        let vec = PVectorMut::<i32>::from_iter(all_valid.clone());
346        let roundtrip: Vec<_> = vec.into_iter().collect();
347        let expected: Vec<_> = all_valid.into_iter().map(Some).collect();
348        assert_eq!(roundtrip, expected);
349
350        // Test with empty.
351        let empty: Vec<Option<i32>> = vec![];
352        let vec = PVectorMut::<i32>::from_iter(empty.clone());
353        let roundtrip: Vec<_> = vec.into_iter().collect();
354        assert_eq!(roundtrip, empty);
355    }
356}
vortex_vector/primitive/generic_mut_impl.rs

vortex_vector/primitive/
generic_mut_impl.rs
