vortex_vector/
vector_mut.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4//! Definition of the [`VectorMut`] type, which represents mutable and fully decompressed
5//! (canonical) array data.
6//!
7//! [`VectorMut`] can be frozen into the [`Vector`] type.
8
9use vortex_dtype::DType;
10use vortex_error::vortex_panic;
11use vortex_mask::MaskMut;
12
13use crate::Vector;
14use crate::VectorMutOps;
15use crate::VectorOps;
16use crate::binaryview::BinaryVectorMut;
17use crate::binaryview::StringVectorMut;
18use crate::bool::BoolVectorMut;
19use crate::decimal::DecimalVectorMut;
20use crate::fixed_size_list::FixedSizeListVectorMut;
21use crate::listview::ListViewVectorMut;
22use crate::match_each_vector_mut;
23use crate::match_vector_pair;
24use crate::null::NullVectorMut;
25use crate::primitive::PrimitiveVectorMut;
26use crate::struct_::StructVectorMut;
27
28/// An enum over all kinds of mutable vectors, which represent fully decompressed (canonical) array
29/// data.
30///
31/// Most of the behavior of `VectorMut` is described by the [`VectorMutOps`] trait. Note that
32/// vectors are **always** considered as nullable, and it is the responsibility of the user to not
33/// add any nullable data to a vector they want to keep as non-nullable.
34///
35/// The immutable equivalent of this type is [`Vector`], which implements the [`VectorOps`] trait.
36#[derive(Debug, Clone)]
37pub enum VectorMut {
38    /// Mutable Null vectors.
39    Null(NullVectorMut),
40    /// Mutable Boolean vectors.
41    Bool(BoolVectorMut),
42    /// Mutable Decimal vectors.
43    ///
44    /// Note that [`DecimalVectorMut`] is an enum over the different possible (generic)
45    /// [`DVectorMut<D>`](crate::decimal::DVectorMut)s.
46    ///
47    /// See the [documentation](crate::decimal) for more information.
48    Decimal(DecimalVectorMut),
49    /// Mutable Primitive vectors.
50    ///
51    /// Note that [`PrimitiveVectorMut`] is an enum over the different possible (generic)
52    /// [`PVectorMut<T>`](crate::primitive::PVectorMut)s.
53    ///
54    /// See the documentation for more information.
55    Primitive(PrimitiveVectorMut),
56    /// Mutable String vectors.
57    String(StringVectorMut),
58    /// Mutable Binary vectors.
59    Binary(BinaryVectorMut),
60    /// Mutable vectors of Lists with variable sizes.
61    List(ListViewVectorMut),
62    /// Mutable vectors of Lists with fixed sizes.
63    FixedSizeList(FixedSizeListVectorMut),
64    /// Mutable vectors of Struct elements.
65    Struct(StructVectorMut),
66}
67
68impl VectorMut {
69    /// Create a new mutable vector with the given capacity and dtype.
70    pub fn with_capacity(dtype: &DType, capacity: usize) -> Self {
71        match dtype {
72            DType::Null => NullVectorMut::new(0).into(),
73            DType::Bool(_) => BoolVectorMut::with_capacity(capacity).into(),
74            DType::Primitive(ptype, _) => {
75                PrimitiveVectorMut::with_capacity(*ptype, capacity).into()
76            }
77            DType::FixedSizeList(elem_dtype, list_size, _) => {
78                FixedSizeListVectorMut::with_capacity(elem_dtype, *list_size, capacity).into()
79            }
80            DType::Struct(struct_fields, _) => {
81                StructVectorMut::with_capacity(struct_fields, capacity).into()
82            }
83            DType::Decimal(decimal_dtype, _) => {
84                DecimalVectorMut::with_capacity(decimal_dtype, capacity).into()
85            }
86            DType::Utf8(..) => StringVectorMut::with_capacity(capacity).into(),
87            DType::Binary(..) => BinaryVectorMut::with_capacity(capacity).into(),
88            DType::Extension(ext) => VectorMut::with_capacity(ext.storage_dtype(), capacity),
89            DType::List(elem, ..) => {
90                // We arbitrarily choose 2 times the number of list scalars for the capacity of the
91                // elements since we cannot know this ahead of time.
92                ListViewVectorMut::with_capacity(elem.as_ref(), capacity, 2 * capacity).into()
93            }
94        }
95    }
96}
97
98impl VectorMutOps for VectorMut {
99    type Immutable = Vector;
100
101    fn len(&self) -> usize {
102        match_each_vector_mut!(self, |v| { v.len() })
103    }
104
105    fn validity(&self) -> &MaskMut {
106        match_each_vector_mut!(self, |v| { v.validity() })
107    }
108
109    fn capacity(&self) -> usize {
110        match_each_vector_mut!(self, |v| { v.capacity() })
111    }
112
113    fn reserve(&mut self, additional: usize) {
114        match_each_vector_mut!(self, |v| { v.reserve(additional) })
115    }
116
117    fn clear(&mut self) {
118        match_each_vector_mut!(self, |v| { v.clear() })
119    }
120
121    fn truncate(&mut self, len: usize) {
122        match_each_vector_mut!(self, |v| { v.truncate(len) })
123    }
124
125    fn extend_from_vector(&mut self, other: &Vector) {
126        match_vector_pair!(self, other, |a: VectorMut, b: Vector| {
127            a.extend_from_vector(b)
128        })
129    }
130
131    fn append_nulls(&mut self, n: usize) {
132        match_each_vector_mut!(self, |v| { v.append_nulls(n) })
133    }
134
135    fn append_zeros(&mut self, n: usize) {
136        match_each_vector_mut!(self, |v| { v.append_zeros(n) })
137    }
138
139    fn append_scalars(&mut self, scalar: &<Self::Immutable as VectorOps>::Scalar, n: usize) {
140        match_vector_pair!(self, scalar, |a: VectorMut, b: Scalar| {
141            a.append_scalars(b, n)
142        })
143    }
144
145    fn freeze(self) -> Vector {
146        match_each_vector_mut!(self, |v| { v.freeze().into() })
147    }
148
149    fn split_off(&mut self, at: usize) -> Self {
150        match_each_vector_mut!(self, |v| { v.split_off(at).into() })
151    }
152
153    fn unsplit(&mut self, other: Self) {
154        match_vector_pair!(self, other, |a: VectorMut, b: VectorMut| a.unsplit(b))
155    }
156}
157
158impl VectorMut {
159    /// Returns a reference to the inner [`NullVectorMut`] if `self` is of that variant.
160    pub fn as_null_mut(&mut self) -> &mut NullVectorMut {
161        if let VectorMut::Null(v) = self {
162            return v;
163        }
164        vortex_panic!("Expected NullVectorMut, got {self:?}");
165    }
166
167    /// Returns a reference to the inner [`BoolVectorMut`] if `self` is of that variant.
168    pub fn as_bool_mut(&mut self) -> &mut BoolVectorMut {
169        if let VectorMut::Bool(v) = self {
170            return v;
171        }
172        vortex_panic!("Expected BoolVectorMut, got {self:?}");
173    }
174
175    /// Returns a reference to the inner [`PrimitiveVectorMut`] if `self` is of that variant.
176    pub fn as_primitive_mut(&mut self) -> &mut PrimitiveVectorMut {
177        if let VectorMut::Primitive(v) = self {
178            return v;
179        }
180        vortex_panic!("Expected PrimitiveVectorMut, got {self:?}");
181    }
182
183    /// Returns a reference to the inner [`StringVectorMut`] if `self` is of that variant.
184    pub fn as_string_mut(&mut self) -> &mut StringVectorMut {
185        if let VectorMut::String(v) = self {
186            return v;
187        }
188        vortex_panic!("Expected StringVectorMut, got {self:?}");
189    }
190
191    /// Returns a reference to the inner [`BinaryVectorMut`] if `self` is of that variant.
192    pub fn as_binary_mut(&mut self) -> &mut BinaryVectorMut {
193        if let VectorMut::Binary(v) = self {
194            return v;
195        }
196        vortex_panic!("Expected BinaryVectorMut, got {self:?}");
197    }
198
199    /// Returns a reference to the inner [`ListViewVectorMut`] if `self` is of that variant.
200    pub fn as_list_mut(&mut self) -> &mut ListViewVectorMut {
201        if let VectorMut::List(v) = self {
202            return v;
203        }
204        vortex_panic!("Expected ListViewVectorMut, got {self:?}");
205    }
206
207    /// Returns a reference to the inner [`FixedSizeListVectorMut`] if `self` is of that variant.
208    pub fn as_fixed_size_list_mut(&mut self) -> &mut FixedSizeListVectorMut {
209        if let VectorMut::FixedSizeList(v) = self {
210            return v;
211        }
212        vortex_panic!("Expected FixedSizeListVectorMut, got {self:?}");
213    }
214
215    /// Returns a reference to the inner [`StructVectorMut`] if `self` is of that variant.
216    pub fn as_struct_mut(&mut self) -> &mut StructVectorMut {
217        if let VectorMut::Struct(v) = self {
218            return v;
219        }
220        vortex_panic!("Expected StructVectorMut, got {self:?}");
221    }
222
223    /// Consumes `self` and returns the inner [`NullVectorMut`] if `self` is of that variant.
224    pub fn into_null(self) -> NullVectorMut {
225        if let VectorMut::Null(v) = self {
226            return v;
227        }
228        vortex_panic!("Expected NullVectorMut, got {self:?}");
229    }
230
231    /// Consumes `self` and returns the inner [`BoolVectorMut`] if `self` is of that variant.
232    pub fn into_bool(self) -> BoolVectorMut {
233        if let VectorMut::Bool(v) = self {
234            return v;
235        }
236        vortex_panic!("Expected BoolVectorMut, got {self:?}");
237    }
238
239    /// Consumes `self` and returns the inner [`PrimitiveVectorMut`] if `self` is of that variant.
240    pub fn into_primitive(self) -> PrimitiveVectorMut {
241        if let VectorMut::Primitive(v) = self {
242            return v;
243        }
244        vortex_panic!("Expected PrimitiveVectorMut, got {self:?}");
245    }
246
247    /// Consumes `self` and returns the inner [`PrimitiveVectorMut`] if `self` is of that variant.
248    pub fn into_decimal(self) -> DecimalVectorMut {
249        if let VectorMut::Decimal(v) = self {
250            return v;
251        }
252        vortex_panic!("Expected DecimalVectorMut, got {self:?}");
253    }
254
255    /// Consumes `self` and returns the inner [`StringVectorMut`] if `self` is of that variant.
256    #[expect(
257        clippy::same_name_method,
258        reason = "intentionally shadows VarBinTypeDowncast method"
259    )]
260    pub fn into_string(self) -> StringVectorMut {
261        if let VectorMut::String(v) = self {
262            return v;
263        }
264        vortex_panic!("Expected StringVectorMut, got {self:?}");
265    }
266
267    /// Consumes `self` and returns the inner [`BinaryVectorMut`] if `self` is of that variant.
268    #[expect(
269        clippy::same_name_method,
270        reason = "intentionally shadows VarBinTypeDowncast method"
271    )]
272    pub fn into_binary(self) -> BinaryVectorMut {
273        if let VectorMut::Binary(v) = self {
274            return v;
275        }
276        vortex_panic!("Expected BinaryVectorMut, got {self:?}");
277    }
278
279    /// Consumes `self` and returns the inner [`ListViewVectorMut`] if `self` is of that variant.
280    pub fn into_list(self) -> ListViewVectorMut {
281        if let VectorMut::List(v) = self {
282            return v;
283        }
284        vortex_panic!("Expected ListViewVectorMut, got {self:?}");
285    }
286
287    /// Consumes `self` and returns the inner [`FixedSizeListVectorMut`] if `self` is of that
288    /// variant.
289    pub fn into_fixed_size_list(self) -> FixedSizeListVectorMut {
290        if let VectorMut::FixedSizeList(v) = self {
291            return v;
292        }
293        vortex_panic!("Expected FixedSizeListVectorMut, got {self:?}");
294    }
295
296    /// Consumes `self` and returns the inner [`StructVectorMut`] if `self` is of that variant.
297    pub fn into_struct(self) -> StructVectorMut {
298        if let VectorMut::Struct(v) = self {
299            return v;
300        }
301        vortex_panic!("Expected StructVectorMut, got {self:?}");
302    }
303}
304
305#[cfg(test)]
306mod tests {
307    use vortex_dtype::DecimalDType;
308    use vortex_dtype::Nullability;
309    use vortex_dtype::PType;
310
311    use super::*;
312    use crate::VectorOps;
313    use crate::decimal::DecimalVectorMut;
314    use crate::primitive::PVectorMut;
315
316    #[test]
317    fn test_with_capacity() {
318        // Test capacity allocation for different types.
319        let null_vec = VectorMut::with_capacity(&DType::Null, 10);
320        assert_eq!(null_vec.capacity(), usize::MAX); // Null vectors have unlimited capacity.
321
322        let bool_vec = VectorMut::with_capacity(&DType::Bool(Nullability::Nullable), 100);
323        assert!(bool_vec.capacity() >= 100);
324
325        let prim_vec =
326            VectorMut::with_capacity(&DType::Primitive(PType::I32, Nullability::Nullable), 50);
327        assert!(prim_vec.capacity() >= 50);
328    }
329
330    #[test]
331    fn test_with_capacity_decimal() {
332        // Test decimal vectors with different precisions that map to different internal types.
333        // Precision 1-2 uses i8, 3-4 uses i16, 5-9 uses i32, 10-18 uses i64,
334        // 19-38 uses i128, 39-76 uses i256.
335
336        // Test precision 4 (uses i16 internally).
337        let decimal_dtype = DType::Decimal(DecimalDType::new(4, 2), Nullability::Nullable);
338        let decimal_vec = VectorMut::with_capacity(&decimal_dtype, 50);
339
340        match decimal_vec {
341            VectorMut::Decimal(dec_vec) => {
342                assert_eq!(dec_vec.len(), 0, "New vector should be empty");
343                assert!(dec_vec.capacity() >= 50, "Capacity should be at least 50");
344
345                // Verify it's using D16 variant internally.
346                assert!(
347                    matches!(dec_vec, DecimalVectorMut::D16(_)),
348                    "Precision 4 should use D16 variant"
349                );
350            }
351            _ => panic!("Expected decimal vector for decimal dtype"),
352        }
353
354        // Test precision 9 (uses i32 internally).
355        let decimal_dtype = DType::Decimal(DecimalDType::new(9, 0), Nullability::NonNullable);
356        let decimal_vec = VectorMut::with_capacity(&decimal_dtype, 100);
357
358        match decimal_vec {
359            VectorMut::Decimal(dec_vec) => {
360                assert_eq!(dec_vec.len(), 0, "New vector should be empty");
361                assert!(dec_vec.capacity() >= 100, "Capacity should be at least 100");
362
363                // Verify it's using D32 variant internally.
364                assert!(
365                    matches!(dec_vec, DecimalVectorMut::D32(_)),
366                    "Precision 9 should use D32 variant"
367                );
368            }
369            _ => panic!("Expected decimal vector for decimal dtype"),
370        }
371
372        // Test precision 18 (uses i64 internally).
373        let decimal_dtype = DType::Decimal(DecimalDType::new(18, -2), Nullability::Nullable);
374        let decimal_vec = VectorMut::with_capacity(&decimal_dtype, 75);
375
376        match decimal_vec {
377            VectorMut::Decimal(dec_vec) => {
378                assert_eq!(dec_vec.len(), 0, "New vector should be empty");
379                assert!(dec_vec.capacity() >= 75, "Capacity should be at least 75");
380
381                // Verify it's using D64 variant internally.
382                assert!(
383                    matches!(dec_vec, DecimalVectorMut::D64(_)),
384                    "Precision 18 should use D64 variant"
385                );
386            }
387            _ => panic!("Expected decimal vector for decimal dtype"),
388        }
389
390        // Test precision 38 (uses i128 internally).
391        let decimal_dtype = DType::Decimal(DecimalDType::new(38, 10), Nullability::NonNullable);
392        let decimal_vec = VectorMut::with_capacity(&decimal_dtype, 25);
393
394        match decimal_vec {
395            VectorMut::Decimal(dec_vec) => {
396                assert_eq!(dec_vec.len(), 0, "New vector should be empty");
397                assert!(dec_vec.capacity() >= 25, "Capacity should be at least 25");
398
399                // Verify it's using D128 variant internally.
400                assert!(
401                    matches!(dec_vec, DecimalVectorMut::D128(_)),
402                    "Precision 38 should use D128 variant"
403                );
404            }
405            _ => panic!("Expected decimal vector for decimal dtype"),
406        }
407    }
408
409    #[test]
410    #[should_panic(expected = "Mismatched vector types")]
411    fn test_type_mismatch_panics() {
412        // Test that operations between mismatched types panic.
413        let mut vec1 = VectorMut::with_capacity(&DType::Bool(Nullability::Nullable), 10);
414        let vec2 =
415            VectorMut::with_capacity(&DType::Primitive(PType::I32, Nullability::Nullable), 10);
416
417        vec1.unsplit(vec2); // Should panic.
418    }
419
420    #[test]
421    fn test_split_and_unsplit() {
422        // Test split at various positions.
423        let mut vec: VectorMut = BoolVectorMut::from_iter([true, false, true].map(Some)).into();
424
425        // Split at beginning.
426        let second = vec.split_off(0);
427        assert_eq!(vec.len(), 0);
428        assert_eq!(second.len(), 3);
429
430        // Unsplit to restore.
431        vec.unsplit(second);
432        assert_eq!(vec.len(), 3);
433
434        // Split at end.
435        let second = vec.split_off(3);
436        assert_eq!(vec.len(), 3);
437        assert_eq!(second.len(), 0);
438    }
439
440    #[test]
441    fn test_reserve_ensures_len_plus_additional() {
442        // Test that reserve ensures capacity >= len + additional.
443        // This specifically tests the fix for the BitBufferMut::reserve bug.
444        let mut bool_vec: VectorMut = BoolVectorMut::with_capacity(10).into();
445        let initial_len = bool_vec.len();
446        assert_eq!(initial_len, 0);
447
448        // Reserve 100 additional capacity.
449        bool_vec.reserve(100);
450
451        // Should have capacity for at least len + 100.
452        assert!(bool_vec.capacity() >= initial_len + 100);
453        assert!(bool_vec.capacity() >= 100); // Since len is 0.
454
455        // Test with primitive vector too.
456        let mut prim_vec: VectorMut = PVectorMut::<i32>::with_capacity(10).into();
457        prim_vec.reserve(100);
458        assert!(prim_vec.capacity() >= prim_vec.len() + 100);
459
460        // Test with non-empty vector.
461        let mut vec: VectorMut = BoolVectorMut::from_iter([true, false, true].map(Some)).into();
462        let len = vec.len();
463        assert_eq!(len, 3);
464        vec.reserve(50);
465        assert!(vec.capacity() >= len + 50);
466        assert!(vec.capacity() >= 53);
467    }
468
469    #[test]
470    fn test_append_nulls_preserves_validity() {
471        // Test that appending nulls preserves existing validity.
472        let mut vec: VectorMut = BoolVectorMut::from_iter([true].map(Some)).into();
473        vec.append_nulls(2);
474        assert_eq!(vec.len(), 3);
475
476        let frozen = vec.freeze();
477        assert_eq!(frozen.validity().true_count(), 1); // Only first element is valid.
478    }
479
480    #[test]
481    fn test_extend_from_vector() {
482        // Test extending a primitive vector with data from another vector.
483        let mut vec: VectorMut = PVectorMut::<i32>::from_iter([1, 2, 3].map(Some)).into();
484        assert_eq!(vec.len(), 3);
485
486        // Create an immutable vector to extend from.
487        let to_append: Vector = PVectorMut::<i32>::from_iter([4, 5, 6].map(Some))
488            .freeze()
489            .into();
490        assert_eq!(to_append.len(), 3);
491
492        // Extend the mutable vector.
493        vec.extend_from_vector(&to_append);
494
495        // Verify the length is the sum of both vectors.
496        assert_eq!(vec.len(), 6);
497
498        // Verify validity is preserved (all elements are valid).
499        let frozen = vec.freeze();
500        assert_eq!(frozen.validity().true_count(), 6);
501    }
502}
vortex_vector/vector_mut.rs

vortex_vector/
vector_mut.rs
