Skip to content

Commit 3813d48

Browse files
authored
Chore: add more vector constructors (#5070)
Signed-off-by: Connor Tsui <[email protected]>
1 parent b78f2fe commit 3813d48

File tree

4 files changed

+154
-8
lines changed

4 files changed

+154
-8
lines changed

vortex-vector/src/bool/vector.rs

Lines changed: 28 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@
44
//! Definition and implementation of [`BoolVector`].
55
66
use vortex_buffer::BitBuffer;
7+
use vortex_error::{VortexExpect, VortexResult, vortex_ensure};
78
use vortex_mask::Mask;
89

910
use super::BoolVectorMut;
@@ -30,11 +31,36 @@ impl BoolVector {
3031
///
3132
/// Panics if the length of the validity mask does not match the length of the bits.
3233
pub fn new(bits: BitBuffer, validity: Mask) -> Self {
33-
assert_eq!(
34+
Self::try_new(bits, validity)
35+
.vortex_expect("`BoolVector` validity mask must have the same length as bits")
36+
}
37+
38+
/// Tries to create a new [`BoolVector`] from the given bits and validity mask.
39+
///
40+
/// # Errors
41+
///
42+
/// Returns an error if the length of the validity mask does not match the length of the bits.
43+
pub fn try_new(bits: BitBuffer, validity: Mask) -> VortexResult<Self> {
44+
vortex_ensure!(
45+
validity.len() == bits.len(),
46+
"`BoolVector` validity mask must have the same length as bits"
47+
);
48+
49+
Ok(Self { bits, validity })
50+
}
51+
52+
/// Creates a new [`BoolVector`] from the given bits and validity mask without validation.
53+
///
54+
/// # Safety
55+
///
56+
/// The caller must ensure that the validity mask has the same length as the bits.
57+
pub fn new_unchecked(bits: BitBuffer, validity: Mask) -> Self {
58+
debug_assert_eq!(
3459
validity.len(),
3560
bits.len(),
36-
"BoolVector validity mask must have the same length as bits"
61+
"`BoolVector` validity mask must have the same length as bits"
3762
);
63+
3864
Self { bits, validity }
3965
}
4066

vortex-vector/src/bool/vector_mut.rs

Lines changed: 37 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@
44
//! Definition and implementation of [`BoolVectorMut`].
55
66
use vortex_buffer::BitBufferMut;
7+
use vortex_error::{VortexExpect, VortexResult, vortex_ensure};
78
use vortex_mask::MaskMut;
89

910
use super::BoolVector;
@@ -71,15 +72,45 @@ pub struct BoolVectorMut {
7172
}
7273

7374
impl BoolVectorMut {
74-
/// Create a mutable vector from the given parts, without checking lengths or capacities.
75+
/// Creates a new [`BoolVectorMut`] from the given bits and validity mask.
7576
///
76-
/// # SAFETY
77+
/// # Panics
7778
///
78-
/// The caller must ensure both parts have the same length and capacity. Ideally they are
79-
/// taken from `into_parts`, mutated in a way that doesn't re-allocate, and then passed back
80-
/// to this function.
79+
/// Panics if the length of the validity mask does not match the length of the bits.
80+
pub fn new(bits: BitBufferMut, validity: MaskMut) -> Self {
81+
Self::try_new(bits, validity)
82+
.vortex_expect("`BoolVector` validity mask must have the same length as bits")
83+
}
84+
85+
/// Tries to create a new [`BoolVectorMut`] from the given bits and validity mask.
86+
///
87+
/// # Errors
88+
///
89+
/// Returns an error if the length of the validity mask does not match the length of the bits.
90+
pub fn try_new(bits: BitBufferMut, validity: MaskMut) -> VortexResult<Self> {
91+
vortex_ensure!(
92+
validity.len() == bits.len(),
93+
"`BoolVector` validity mask must have the same length as bits"
94+
);
95+
96+
Ok(Self { bits, validity })
97+
}
98+
99+
/// Creates a new [`BoolVectorMut`] from the given bits and validity mask without validation.
100+
///
101+
/// # Safety
102+
///
103+
/// The caller must ensure that the validity mask has the same length as the bits.
104+
///
105+
/// Ideally, they are taken from `into_parts`, mutated in a way that doesn't re-allocate, and
106+
/// then passed back to this function.
81107
pub unsafe fn new_unchecked(bits: BitBufferMut, validity: MaskMut) -> Self {
82-
debug_assert_eq!(bits.len(), validity.len());
108+
debug_assert_eq!(
109+
bits.len(),
110+
validity.len(),
111+
"`BoolVector` validity mask must have the same length as bits"
112+
);
113+
83114
Self { bits, validity }
84115
}
85116

vortex-vector/src/primitive/generic.rs

Lines changed: 44 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@
55
66
use vortex_buffer::Buffer;
77
use vortex_dtype::NativePType;
8+
use vortex_error::{VortexExpect, VortexResult, vortex_ensure};
89
use vortex_mask::Mask;
910

1011
use crate::{PVectorMut, VectorOps};
@@ -26,6 +27,49 @@ pub struct PVector<T> {
2627
pub(super) validity: Mask,
2728
}
2829

30+
impl<T: NativePType> PVector<T> {
31+
/// Creates a new [`PVector<T>`] from the given elements buffer and validity mask.
32+
///
33+
/// # Panics
34+
///
35+
/// Panics if the length of the validity mask does not match the length of the elements buffer.
36+
pub fn new(elements: Buffer<T>, validity: Mask) -> Self {
37+
Self::try_new(elements, validity)
38+
.vortex_expect("`PVector` validity mask must have the same length as elements")
39+
}
40+
41+
/// Tries to create a new [`PVector<T>`] from the given elements buffer and validity mask.
42+
///
43+
/// # Errors
44+
///
45+
/// Returns an error if the length of the validity mask does not match the length of the
46+
/// elements buffer.
47+
pub fn try_new(elements: Buffer<T>, validity: Mask) -> VortexResult<Self> {
48+
vortex_ensure!(
49+
validity.len() == elements.len(),
50+
"`PVector` validity mask must have the same length as elements"
51+
);
52+
53+
Ok(Self { elements, validity })
54+
}
55+
56+
/// Creates a new [`PVector<T>`] from the given elements buffer and validity mask without
57+
/// validation.
58+
///
59+
/// # Safety
60+
///
61+
/// The caller must ensure that the validity mask has the same length as the elements buffer.
62+
pub fn new_unchecked(elements: Buffer<T>, validity: Mask) -> Self {
63+
debug_assert_eq!(
64+
validity.len(),
65+
elements.len(),
66+
"`PVector` validity mask must have the same length as elements"
67+
);
68+
69+
Self { elements, validity }
70+
}
71+
}
72+
2973
impl<T: NativePType> VectorOps for PVector<T> {
3074
type Mutable = PVectorMut<T>;
3175

vortex-vector/src/primitive/generic_mut.rs

Lines changed: 45 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@
55
66
use vortex_buffer::BufferMut;
77
use vortex_dtype::NativePType;
8+
use vortex_error::{VortexExpect, VortexResult, vortex_ensure};
89
use vortex_mask::MaskMut;
910

1011
use crate::{PVector, VectorMutOps, VectorOps};
@@ -107,6 +108,50 @@ pub struct PVectorMut<T> {
107108
}
108109

109110
impl<T> PVectorMut<T> {
111+
/// Creates a new [`PVectorMut<T>`] from the given elements buffer and validity mask.
112+
///
113+
/// # Panics
114+
///
115+
/// Panics if the length of the validity mask does not match the length of the elements buffer.
116+
pub fn new(elements: BufferMut<T>, validity: MaskMut) -> Self {
117+
Self::try_new(elements, validity)
118+
.vortex_expect("`PVectorMut` validity mask must have the same length as elements")
119+
}
120+
121+
/// Tries to create a new [`PVectorMut<T>`] from the given elements buffer and validity mask.
122+
///
123+
/// # Errors
124+
///
125+
/// Returns an error if the length of the validity mask does not match the length of the
126+
/// elements buffer.
127+
pub fn try_new(elements: BufferMut<T>, validity: MaskMut) -> VortexResult<Self> {
128+
vortex_ensure!(
129+
validity.len() == elements.len(),
130+
"`PVectorMut` validity mask must have the same length as elements"
131+
);
132+
133+
Ok(Self { elements, validity })
134+
}
135+
136+
/// Creates a new [`PVectorMut<T>`] from the given elements buffer and validity mask without
137+
/// validation.
138+
///
139+
/// # Safety
140+
///
141+
/// The caller must ensure that the validity mask has the same length as the elements buffer.
142+
///
143+
/// Ideally, they are taken from `into_parts`, mutated in a way that doesn't re-allocate, and
144+
/// then passed back to this function.
145+
pub unsafe fn new_unchecked(elements: BufferMut<T>, validity: MaskMut) -> Self {
146+
debug_assert_eq!(
147+
elements.len(),
148+
validity.len(),
149+
"`PVectorMut` validity mask must have the same length as elements"
150+
);
151+
152+
Self { elements, validity }
153+
}
154+
110155
/// Create a new mutable primitive vector with the given capacity.
111156
pub fn with_capacity(capacity: usize) -> Self {
112157
Self {

0 commit comments

Comments
 (0)