vortex_vector/primitive/generic.rs
1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4//! Definition and implementation of [`PVector<T>`].
5
6use std::fmt::Debug;
7use std::ops::RangeBounds;
8
9use vortex_buffer::Buffer;
10use vortex_dtype::NativePType;
11use vortex_error::{VortexExpect, VortexResult, vortex_ensure};
12use vortex_mask::Mask;
13
14use crate::primitive::{PScalar, PVectorMut};
15use crate::{Scalar, VectorOps};
16
17/// An immutable vector of generic primitive values.
18///
19/// `T` is expected to be bound by [`NativePType`], which templates an internal [`Buffer<T>`] that
20/// stores the elements of the vector.
21#[derive(Debug, Clone)]
22pub struct PVector<T> {
23 /// The buffer representing the vector elements.
24 pub(super) elements: Buffer<T>,
25 /// The validity mask (where `true` represents an element is **not** null).
26 pub(super) validity: Mask,
27}
28
29impl<T> PVector<T> {
30 /// Creates a new [`PVector<T>`] from the given elements buffer and validity mask.
31 ///
32 /// # Panics
33 ///
34 /// Panics if the length of the validity mask does not match the length of the elements buffer.
35 pub fn new(elements: Buffer<T>, validity: Mask) -> Self {
36 Self::try_new(elements, validity).vortex_expect("Failed to create `PVector`")
37 }
38
39 /// Tries to create a new [`PVector<T>`] from the given elements buffer and validity mask.
40 ///
41 /// # Errors
42 ///
43 /// Returns an error if the length of the validity mask does not match the length of the
44 /// elements buffer.
45 pub fn try_new(elements: Buffer<T>, validity: Mask) -> VortexResult<Self> {
46 vortex_ensure!(
47 validity.len() == elements.len(),
48 "`PVector` validity mask must have the same length as elements"
49 );
50
51 Ok(Self { elements, validity })
52 }
53
54 /// Creates a new [`PVector<T>`] from the given elements buffer and validity mask without
55 /// validation.
56 ///
57 /// # Safety
58 ///
59 /// The caller must ensure that the validity mask has the same length as the elements buffer.
60 pub unsafe fn new_unchecked(elements: Buffer<T>, validity: Mask) -> Self {
61 if cfg!(debug_assertions) {
62 Self::new(elements, validity)
63 } else {
64 Self { elements, validity }
65 }
66 }
67
68 /// Decomposes the primitive vector into its constituent parts (buffer and validity).
69 pub fn into_parts(self) -> (Buffer<T>, Mask) {
70 (self.elements, self.validity)
71 }
72
73 /// Gets a nullable element at the given index, panicking on out-of-bounds.
74 ///
75 /// If the element at the given index is null, returns `None`. Otherwise, returns `Some(x)`,
76 /// where `x: T`.
77 ///
78 /// Note that this `get` method is different from the standard library [`slice::get`], which
79 /// returns `None` if the index is out of bounds. This method will panic if the index is out of
80 /// bounds, and return `None` if the elements is null.
81 ///
82 /// # Panics
83 ///
84 /// Panics if the index is out of bounds.
85 pub fn get(&self, index: usize) -> Option<&T> {
86 self.validity.value(index).then(|| &self.elements[index])
87 }
88
89 /// Returns the internal [`Buffer`] of the [`PVector`].
90 ///
91 /// Note that the internal buffer may hold garbage data in place of nulls. That information is
92 /// tracked by the [`validity()`](Self::validity).
93 #[inline]
94 pub fn elements(&self) -> &Buffer<T> {
95 &self.elements
96 }
97}
98
99impl<T: NativePType> AsRef<[T]> for PVector<T> {
100 /// Returns an immutable slice over the internal buffer with elements of type `T`.
101 ///
102 /// Note that this slice may contain garbage data where the [`validity()`] mask states that an
103 /// element is invalid.
104 ///
105 /// The caller should check the [`validity()`] before performing any operations.
106 ///
107 /// [`validity()`]: crate::VectorOps::validity
108 #[inline]
109 fn as_ref(&self) -> &[T] {
110 self.elements.as_slice()
111 }
112}
113
114impl<T: NativePType> VectorOps for PVector<T> {
115 type Mutable = PVectorMut<T>;
116
117 fn len(&self) -> usize {
118 self.elements.len()
119 }
120
121 fn validity(&self) -> &Mask {
122 &self.validity
123 }
124
125 fn scalar_at(&self, index: usize) -> Scalar {
126 assert!(index < self.len(), "Index out of bounds in `PVector`");
127 PScalar::<T>::new(self.validity.value(index).then(|| self.elements[index])).into()
128 }
129
130 fn slice(&self, range: impl RangeBounds<usize> + Clone + Debug) -> Self {
131 let elements = self.elements.slice(range.clone());
132 let validity = self.validity.slice(range);
133 Self::new(elements, validity)
134 }
135
136 /// Try to convert self into a mutable vector.
137 fn try_into_mut(self) -> Result<PVectorMut<T>, Self> {
138 let elements = match self.elements.try_into_mut() {
139 Ok(elements) => elements,
140 Err(elements) => {
141 return Err(Self {
142 elements,
143 validity: self.validity,
144 });
145 }
146 };
147
148 match self.validity.try_into_mut() {
149 Ok(validity_mut) => Ok(PVectorMut {
150 elements,
151 validity: validity_mut,
152 }),
153 Err(validity) => Err(Self {
154 elements: elements.freeze(),
155 validity,
156 }),
157 }
158 }
159
160 fn into_mut(self) -> PVectorMut<T> {
161 let elements = self.elements.into_mut();
162 let validity = self.validity.into_mut();
163
164 PVectorMut { elements, validity }
165 }
166}