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

vortex_vector/primitive/
vector_mut.rs
