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

vortex_vector/primitive/
generic_mut_impl.rs
