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

vortex_vector/
vector_mut.rs
