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}