vortex_vector/primitive/
vector_mut.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4//! Definition and implementation of [`PrimitiveVectorMut`].
5
6use vortex_dtype::NativePType;
7use vortex_dtype::PType;
8use vortex_dtype::PTypeDowncast;
9use vortex_dtype::PTypeUpcast;
10use vortex_dtype::half::f16;
11use vortex_error::vortex_panic;
12use vortex_mask::MaskMut;
13
14use crate::VectorMutOps;
15use crate::match_each_pvector_mut;
16use crate::primitive::PVectorMut;
17use crate::primitive::PrimitiveScalar;
18use crate::primitive::PrimitiveVector;
19
20/// A mutable vector of primitive values.
21///
22/// The immutable equivalent of this type is [`PrimitiveVector`].
23///
24/// `PrimitiveVector` is represented by an enum over all possible [`PVectorMut`] types (which are
25/// templated by the types that implement [`NativePType`]).
26///
27/// See the documentation for [`PVectorMut`] for more information.
28#[derive(Debug, Clone)]
29pub enum PrimitiveVectorMut {
30    /// U8
31    U8(PVectorMut<u8>),
32    /// U16
33    U16(PVectorMut<u16>),
34    /// U32
35    U32(PVectorMut<u32>),
36    /// U64
37    U64(PVectorMut<u64>),
38    /// I8
39    I8(PVectorMut<i8>),
40    /// I16
41    I16(PVectorMut<i16>),
42    /// I32
43    I32(PVectorMut<i32>),
44    /// I64
45    I64(PVectorMut<i64>),
46    /// F16
47    F16(PVectorMut<f16>),
48    /// F32
49    F32(PVectorMut<f32>),
50    /// F64
51    F64(PVectorMut<f64>),
52}
53
54impl PrimitiveVectorMut {
55    /// Returns the [`PType`] of this [`PrimitiveVectorMut`].
56    pub fn ptype(&self) -> PType {
57        match self {
58            Self::U8(_) => PType::U8,
59            Self::U16(_) => PType::U16,
60            Self::U32(_) => PType::U32,
61            Self::U64(_) => PType::U64,
62            Self::I8(_) => PType::I8,
63            Self::I16(_) => PType::I16,
64            Self::I32(_) => PType::I32,
65            Self::I64(_) => PType::I64,
66            Self::F16(_) => PType::F16,
67            Self::F32(_) => PType::F32,
68            Self::F64(_) => PType::F64,
69        }
70    }
71
72    /// Create a new mutable primitive vector with the given primitive type and capacity.
73    pub fn with_capacity(ptype: PType, capacity: usize) -> Self {
74        match ptype {
75            PType::U8 => PVectorMut::<u8>::with_capacity(capacity).into(),
76            PType::U16 => PVectorMut::<u16>::with_capacity(capacity).into(),
77            PType::U32 => PVectorMut::<u32>::with_capacity(capacity).into(),
78            PType::U64 => PVectorMut::<u64>::with_capacity(capacity).into(),
79            PType::I8 => PVectorMut::<i8>::with_capacity(capacity).into(),
80            PType::I16 => PVectorMut::<i16>::with_capacity(capacity).into(),
81            PType::I32 => PVectorMut::<i32>::with_capacity(capacity).into(),
82            PType::I64 => PVectorMut::<i64>::with_capacity(capacity).into(),
83            PType::F16 => PVectorMut::<f16>::with_capacity(capacity).into(),
84            PType::F32 => PVectorMut::<f32>::with_capacity(capacity).into(),
85            PType::F64 => PVectorMut::<f64>::with_capacity(capacity).into(),
86        }
87    }
88}
89
90impl VectorMutOps for PrimitiveVectorMut {
91    type Immutable = PrimitiveVector;
92
93    fn len(&self) -> usize {
94        match_each_pvector_mut!(self, |v| { v.len() })
95    }
96
97    fn validity(&self) -> &MaskMut {
98        match_each_pvector_mut!(self, |v| { v.validity() })
99    }
100
101    fn capacity(&self) -> usize {
102        match_each_pvector_mut!(self, |v| { v.capacity() })
103    }
104
105    fn reserve(&mut self, additional: usize) {
106        match_each_pvector_mut!(self, |v| { v.reserve(additional) })
107    }
108
109    fn clear(&mut self) {
110        match_each_pvector_mut!(self, |v| { v.clear() })
111    }
112
113    fn truncate(&mut self, len: usize) {
114        match_each_pvector_mut!(self, |v| { v.truncate(len) })
115    }
116
117    fn extend_from_vector(&mut self, other: &PrimitiveVector) {
118        match (self, other) {
119            (Self::U8(a), PrimitiveVector::U8(b)) => a.extend_from_vector(b),
120            (Self::U16(a), PrimitiveVector::U16(b)) => a.extend_from_vector(b),
121            (Self::U32(a), PrimitiveVector::U32(b)) => a.extend_from_vector(b),
122            (Self::U64(a), PrimitiveVector::U64(b)) => a.extend_from_vector(b),
123            (Self::I8(a), PrimitiveVector::I8(b)) => a.extend_from_vector(b),
124            (Self::I16(a), PrimitiveVector::I16(b)) => a.extend_from_vector(b),
125            (Self::I32(a), PrimitiveVector::I32(b)) => a.extend_from_vector(b),
126            (Self::I64(a), PrimitiveVector::I64(b)) => a.extend_from_vector(b),
127            (Self::F16(a), PrimitiveVector::F16(b)) => a.extend_from_vector(b),
128            (Self::F32(a), PrimitiveVector::F32(b)) => a.extend_from_vector(b),
129            (Self::F64(a), PrimitiveVector::F64(b)) => a.extend_from_vector(b),
130            _ => ::vortex_error::vortex_panic!("Mismatched primitive vector types"),
131        }
132    }
133
134    fn append_nulls(&mut self, n: usize) {
135        match_each_pvector_mut!(self, |v| { v.append_nulls(n) })
136    }
137
138    fn append_zeros(&mut self, n: usize) {
139        match_each_pvector_mut!(self, |v| { v.append_zeros(n) })
140    }
141
142    #[expect(
143        clippy::many_single_char_names,
144        reason = "single-letter names a/b are clear for matching variants"
145    )]
146    fn append_scalars(&mut self, scalar: &PrimitiveScalar, n: usize) {
147        match (self, scalar) {
148            (Self::U8(a), PrimitiveScalar::U8(b)) => a.append_scalars(b, n),
149            (Self::U16(a), PrimitiveScalar::U16(b)) => a.append_scalars(b, n),
150            (Self::U32(a), PrimitiveScalar::U32(b)) => a.append_scalars(b, n),
151            (Self::U64(a), PrimitiveScalar::U64(b)) => a.append_scalars(b, n),
152            (Self::I8(a), PrimitiveScalar::I8(b)) => a.append_scalars(b, n),
153            (Self::I16(a), PrimitiveScalar::I16(b)) => a.append_scalars(b, n),
154            (Self::I32(a), PrimitiveScalar::I32(b)) => a.append_scalars(b, n),
155            (Self::I64(a), PrimitiveScalar::I64(b)) => a.append_scalars(b, n),
156            (Self::F16(a), PrimitiveScalar::F16(b)) => a.append_scalars(b, n),
157            (Self::F32(a), PrimitiveScalar::F32(b)) => a.append_scalars(b, n),
158            (Self::F64(a), PrimitiveScalar::F64(b)) => a.append_scalars(b, n),
159            _ => vortex_panic!("Mismatched primitive vector and scalar types"),
160        }
161    }
162
163    fn freeze(self) -> PrimitiveVector {
164        match_each_pvector_mut!(self, |v| { v.freeze().into() })
165    }
166
167    fn split_off(&mut self, at: usize) -> Self {
168        match_each_pvector_mut!(self, |v| { v.split_off(at).into() })
169    }
170
171    fn unsplit(&mut self, other: Self) {
172        match (self, other) {
173            (Self::U8(a), Self::U8(b)) => a.unsplit(b),
174            (Self::U16(a), Self::U16(b)) => a.unsplit(b),
175            (Self::U32(a), Self::U32(b)) => a.unsplit(b),
176            (Self::U64(a), Self::U64(b)) => a.unsplit(b),
177            (Self::I8(a), Self::I8(b)) => a.unsplit(b),
178            (Self::I16(a), Self::I16(b)) => a.unsplit(b),
179            (Self::I32(a), Self::I32(b)) => a.unsplit(b),
180            (Self::I64(a), Self::I64(b)) => a.unsplit(b),
181            (Self::F16(a), Self::F16(b)) => a.unsplit(b),
182            (Self::F32(a), Self::F32(b)) => a.unsplit(b),
183            (Self::F64(a), Self::F64(b)) => a.unsplit(b),
184            _ => vortex_panic!("Mismatched primitive vector types"),
185        }
186    }
187}
188
189impl PTypeUpcast for PrimitiveVectorMut {
190    type Input<T: NativePType> = PVectorMut<T>;
191
192    fn from_u8(input: Self::Input<u8>) -> Self {
193        Self::U8(input)
194    }
195
196    fn from_u16(input: Self::Input<u16>) -> Self {
197        Self::U16(input)
198    }
199
200    fn from_u32(input: Self::Input<u32>) -> Self {
201        Self::U32(input)
202    }
203
204    fn from_u64(input: Self::Input<u64>) -> Self {
205        Self::U64(input)
206    }
207
208    fn from_i8(input: Self::Input<i8>) -> Self {
209        Self::I8(input)
210    }
211
212    fn from_i16(input: Self::Input<i16>) -> Self {
213        Self::I16(input)
214    }
215
216    fn from_i32(input: Self::Input<i32>) -> Self {
217        Self::I32(input)
218    }
219
220    fn from_i64(input: Self::Input<i64>) -> Self {
221        Self::I64(input)
222    }
223
224    fn from_f16(input: Self::Input<f16>) -> Self {
225        Self::F16(input)
226    }
227
228    fn from_f32(input: Self::Input<f32>) -> Self {
229        Self::F32(input)
230    }
231
232    fn from_f64(input: Self::Input<f64>) -> Self {
233        Self::F64(input)
234    }
235}
236
237impl PTypeDowncast for PrimitiveVectorMut {
238    type Output<T: NativePType> = PVectorMut<T>;
239
240    fn into_u8(self) -> Self::Output<u8> {
241        if let Self::U8(v) = self {
242            return v;
243        }
244        vortex_panic!("Expected PrimitiveVectorMut::U8, got {self:?}");
245    }
246
247    fn into_u16(self) -> Self::Output<u16> {
248        if let Self::U16(v) = self {
249            return v;
250        }
251        vortex_panic!("Expected PrimitiveVectorMut::U16, got {self:?}");
252    }
253
254    fn into_u32(self) -> Self::Output<u32> {
255        if let Self::U32(v) = self {
256            return v;
257        }
258        vortex_panic!("Expected PrimitiveVectorMut::U32, got {self:?}");
259    }
260
261    fn into_u64(self) -> Self::Output<u64> {
262        if let Self::U64(v) = self {
263            return v;
264        }
265        vortex_panic!("Expected PrimitiveVectorMut::U64, got {self:?}");
266    }
267
268    fn into_i8(self) -> Self::Output<i8> {
269        if let Self::I8(v) = self {
270            return v;
271        }
272        vortex_panic!("Expected PrimitiveVectorMut::I8, got {self:?}");
273    }
274
275    fn into_i16(self) -> Self::Output<i16> {
276        if let Self::I16(v) = self {
277            return v;
278        }
279        vortex_panic!("Expected PrimitiveVectorMut::I16, got {self:?}");
280    }
281
282    fn into_i32(self) -> Self::Output<i32> {
283        if let Self::I32(v) = self {
284            return v;
285        }
286        vortex_panic!("Expected PrimitiveVectorMut::I32, got {self:?}");
287    }
288
289    fn into_i64(self) -> Self::Output<i64> {
290        if let Self::I64(v) = self {
291            return v;
292        }
293        vortex_panic!("Expected PrimitiveVectorMut::I64, got {self:?}");
294    }
295
296    fn into_f16(self) -> Self::Output<f16> {
297        if let Self::F16(v) = self {
298            return v;
299        }
300        vortex_panic!("Expected PrimitiveVectorMut::F16, got {self:?}");
301    }
302
303    fn into_f32(self) -> Self::Output<f32> {
304        if let Self::F32(v) = self {
305            return v;
306        }
307        vortex_panic!("Expected PrimitiveVectorMut::F32, got {self:?}");
308    }
309
310    fn into_f64(self) -> Self::Output<f64> {
311        if let Self::F64(v) = self {
312            return v;
313        }
314        vortex_panic!("Expected PrimitiveVectorMut::F64, got {self:?}");
315    }
316}
317
318impl<'a> PTypeDowncast for &'a mut PrimitiveVectorMut {
319    type Output<T: NativePType> = &'a mut PVectorMut<T>;
320
321    fn into_u8(self) -> Self::Output<u8> {
322        match self {
323            PrimitiveVectorMut::U8(v) => v,
324            _ => vortex_panic!("Expected PrimitiveVectorMut::U8, got {self:?}"),
325        }
326    }
327
328    fn into_u16(self) -> Self::Output<u16> {
329        match self {
330            PrimitiveVectorMut::U16(v) => v,
331            _ => vortex_panic!("Expected PrimitiveVectorMut::U16, got {self:?}"),
332        }
333    }
334
335    fn into_u32(self) -> Self::Output<u32> {
336        match self {
337            PrimitiveVectorMut::U32(v) => v,
338            _ => vortex_panic!("Expected PrimitiveVectorMut::U32, got {self:?}"),
339        }
340    }
341
342    fn into_u64(self) -> Self::Output<u64> {
343        match self {
344            PrimitiveVectorMut::U64(v) => v,
345            _ => vortex_panic!("Expected PrimitiveVectorMut::U64, got {self:?}"),
346        }
347    }
348
349    fn into_i8(self) -> Self::Output<i8> {
350        match self {
351            PrimitiveVectorMut::I8(v) => v,
352            _ => vortex_panic!("Expected PrimitiveVectorMut::I8, got {self:?}"),
353        }
354    }
355
356    fn into_i16(self) -> Self::Output<i16> {
357        match self {
358            PrimitiveVectorMut::I16(v) => v,
359            _ => vortex_panic!("Expected PrimitiveVectorMut::I16, got {self:?}"),
360        }
361    }
362
363    fn into_i32(self) -> Self::Output<i32> {
364        match self {
365            PrimitiveVectorMut::I32(v) => v,
366            _ => vortex_panic!("Expected PrimitiveVectorMut::I32, got {self:?}"),
367        }
368    }
369
370    fn into_i64(self) -> Self::Output<i64> {
371        match self {
372            PrimitiveVectorMut::I64(v) => v,
373            _ => vortex_panic!("Expected PrimitiveVectorMut::I64, got {self:?}"),
374        }
375    }
376
377    fn into_f16(self) -> Self::Output<f16> {
378        match self {
379            PrimitiveVectorMut::F16(v) => v,
380            _ => vortex_panic!("Expected PrimitiveVectorMut::F16, got {self:?}"),
381        }
382    }
383
384    fn into_f32(self) -> Self::Output<f32> {
385        match self {
386            PrimitiveVectorMut::F32(v) => v,
387            _ => vortex_panic!("Expected PrimitiveVectorMut::F32, got {self:?}"),
388        }
389    }
390
391    fn into_f64(self) -> Self::Output<f64> {
392        match self {
393            PrimitiveVectorMut::F64(v) => v,
394            _ => vortex_panic!("Expected PrimitiveVectorMut::F64, got {self:?}"),
395        }
396    }
397}
398
399#[cfg(test)]
400mod tests {
401    use super::*;
402    use crate::VectorOps;
403
404    #[test]
405    fn test_from_iter_with_options() {
406        // Test FromIterator<Option<T>> with different types.
407        let vec_i32: PrimitiveVectorMut =
408            PVectorMut::<i32>::from_iter(vec![Some(1), None, Some(3), None, Some(5)]).into();
409        assert_eq!(vec_i32.len(), 5);
410        let frozen = vec_i32.freeze();
411        assert_eq!(frozen.validity().true_count(), 3);
412
413        // Test empty iterator.
414        let vec_empty: PrimitiveVectorMut =
415            PVectorMut::<f64>::from_iter(std::iter::empty::<Option<f64>>()).into();
416        assert_eq!(vec_empty.len(), 0);
417
418        // Test that None values use T::default().
419        let vec_nulls: PrimitiveVectorMut = PVectorMut::<i32>::from_iter([None, None, None]).into();
420        // Check that validity is all false for nulls.
421        let frozen = vec_nulls.freeze();
422        assert_eq!(frozen.validity().true_count(), 0);
423    }
424
425    #[test]
426    fn test_from_iter_non_null() {
427        // Test FromIterator<T> for different primitive types.
428        let vec_f64: PrimitiveVectorMut =
429            PVectorMut::<f64>::from_iter([1.5, 2.5, 3.5, 4.5, 5.5]).into();
430        assert_eq!(vec_f64.len(), 5);
431        let frozen = vec_f64.freeze();
432        assert_eq!(frozen.validity().true_count(), 5); // All valid.
433
434        let vec_u16: PrimitiveVectorMut = PVectorMut::<u16>::from_iter([1u16, 2, 3, 4, 5]).into();
435        assert_eq!(vec_u16.len(), 5);
436        let frozen = vec_u16.freeze();
437        assert_eq!(frozen.validity().true_count(), 5);
438    }
439
440    #[test]
441    fn test_operations_preserve_validity() {
442        // Test split/unsplit/extend with different primitive types.
443        let mut vec: PrimitiveVectorMut =
444            PVectorMut::<i64>::from_iter([Some(100), None, Some(300), None, Some(500)]).into();
445
446        let second_half = vec.split_off(2);
447        assert_eq!(vec.len(), 2);
448        assert_eq!(second_half.len(), 3);
449
450        let first_frozen = vec.freeze();
451        let second_frozen = second_half.freeze();
452        assert_eq!(first_frozen.validity().true_count(), 1);
453        assert_eq!(second_frozen.validity().true_count(), 2);
454
455        // Test unsplit.
456        let mut vec1: PrimitiveVectorMut = PVectorMut::<u32>::from_iter([Some(1000), None]).into();
457        let vec2: PrimitiveVectorMut = PVectorMut::<u32>::from_iter([None, Some(2000)]).into();
458        vec1.unsplit(vec2);
459        assert_eq!(vec1.len(), 4);
460        let frozen = vec1.freeze();
461        assert_eq!(frozen.validity().true_count(), 2);
462    }
463}
vortex_vector/primitive/vector_mut.rs

vortex_vector/primitive/
vector_mut.rs
