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

vortex_vector/
vector_mut.rs
