vortex_array/builders/
primitive.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::any::Any;
5use std::mem::MaybeUninit;
6
7use vortex_buffer::BufferMut;
8use vortex_dtype::DType;
9use vortex_dtype::NativePType;
10use vortex_dtype::Nullability;
11use vortex_error::VortexResult;
12use vortex_error::vortex_ensure;
13use vortex_mask::Mask;
14use vortex_scalar::PrimitiveScalar;
15use vortex_scalar::Scalar;
16
17use crate::Array;
18use crate::ArrayRef;
19use crate::IntoArray;
20use crate::arrays::PrimitiveArray;
21use crate::builders::ArrayBuilder;
22use crate::builders::DEFAULT_BUILDER_CAPACITY;
23use crate::builders::LazyBitBufferBuilder;
24use crate::canonical::Canonical;
25use crate::canonical::ToCanonical;
26
27/// The builder for building a [`PrimitiveArray`], parametrized by the `PType`.
28pub struct PrimitiveBuilder<T> {
29    dtype: DType,
30    values: BufferMut<T>,
31    nulls: LazyBitBufferBuilder,
32}
33
34impl<T: NativePType> PrimitiveBuilder<T> {
35    /// Creates a new `PrimitiveBuilder` with a capacity of [`DEFAULT_BUILDER_CAPACITY`].
36    pub fn new(nullability: Nullability) -> Self {
37        Self::with_capacity(nullability, DEFAULT_BUILDER_CAPACITY)
38    }
39
40    /// Creates a new `PrimitiveBuilder` with the given `capacity`.
41    pub fn with_capacity(nullability: Nullability, capacity: usize) -> Self {
42        Self {
43            values: BufferMut::with_capacity(capacity),
44            nulls: LazyBitBufferBuilder::new(capacity),
45            dtype: DType::Primitive(T::PTYPE, nullability),
46        }
47    }
48
49    /// Appends a primitive `value` to the builder.
50    pub fn append_value(&mut self, value: T) {
51        self.values.push(value);
52        self.nulls.append_non_null();
53    }
54
55    /// Returns the raw primitive values in this builder as a slice.
56    pub fn values(&self) -> &[T] {
57        self.values.as_ref()
58    }
59
60    /// Create a new handle to the next `len` uninitialized values in the builder.
61    ///
62    /// All reads/writes through the handle to the values buffer or the validity buffer will operate
63    /// on indices relative to the start of the range.
64    ///
65    /// # Panics
66    ///
67    /// Panics if `len` is 0 or if the current length of the builder plus `len` would exceed the
68    /// capacity of the builder's memory.
69    ///
70    /// ## Example
71    ///
72    /// ```
73    /// use std::mem::MaybeUninit;
74    /// use vortex_array::builders::{ArrayBuilder, PrimitiveBuilder};
75    /// use vortex_dtype::Nullability;
76    ///
77    /// // Create a new builder.
78    /// let mut builder: PrimitiveBuilder<i32> =
79    ///     PrimitiveBuilder::with_capacity(Nullability::NonNullable, 5);
80    ///
81    /// // Populate the values.
82    /// let mut uninit_range = builder.uninit_range(5);
83    /// uninit_range.copy_from_slice(0, &[0, 1, 2, 3, 4]);
84    ///
85    /// // SAFETY: We have initialized all 5 values in the range, and since the array builder is
86    /// // non-nullable, we don't need to set any null bits.
87    /// unsafe { uninit_range.finish(); }
88    ///
89    /// let built = builder.finish_into_primitive();
90    ///
91    /// assert_eq!(built.as_slice::<i32>(), &[0i32, 1, 2, 3, 4]);
92    /// ```
93    pub fn uninit_range(&mut self, len: usize) -> UninitRange<'_, T> {
94        assert_ne!(0, len, "cannot create an uninit range of length 0");
95
96        let current_len = self.values.len();
97        assert!(
98            current_len + len <= self.values.capacity(),
99            "uninit_range of len {len} exceeds builder with length {} and capacity {}",
100            current_len,
101            self.values.capacity()
102        );
103
104        UninitRange { len, builder: self }
105    }
106
107    /// Finishes the builder directly into a [`PrimitiveArray`].
108    pub fn finish_into_primitive(&mut self) -> PrimitiveArray {
109        let validity = self
110            .nulls
111            .finish_with_nullability(self.dtype().nullability());
112
113        PrimitiveArray::new(std::mem::take(&mut self.values).freeze(), validity)
114    }
115
116    /// Extends the primitive array with an iterator.
117    pub fn extend_with_iterator(&mut self, iter: impl IntoIterator<Item = T>, mask: Mask) {
118        self.values.extend(iter);
119        self.nulls.append_validity_mask(mask);
120    }
121}
122
123impl<T: NativePType> ArrayBuilder for PrimitiveBuilder<T> {
124    fn as_any(&self) -> &dyn Any {
125        self
126    }
127
128    fn as_any_mut(&mut self) -> &mut dyn Any {
129        self
130    }
131
132    fn dtype(&self) -> &DType {
133        &self.dtype
134    }
135
136    fn len(&self) -> usize {
137        self.values.len()
138    }
139
140    fn append_zeros(&mut self, n: usize) {
141        self.values.push_n(T::default(), n);
142        self.nulls.append_n_non_nulls(n);
143    }
144
145    unsafe fn append_nulls_unchecked(&mut self, n: usize) {
146        self.values.push_n(T::default(), n);
147        self.nulls.append_n_nulls(n);
148    }
149
150    fn append_scalar(&mut self, scalar: &Scalar) -> VortexResult<()> {
151        vortex_ensure!(
152            scalar.dtype() == self.dtype(),
153            "PrimitiveBuilder expected scalar with dtype {:?}, got {:?}",
154            self.dtype(),
155            scalar.dtype()
156        );
157
158        let primitive_scalar = PrimitiveScalar::try_from(scalar)?;
159        match primitive_scalar.pvalue() {
160            Some(pv) => self.append_value(pv.cast::<T>()),
161            None => self.append_null(),
162        }
163
164        Ok(())
165    }
166
167    unsafe fn extend_from_array_unchecked(&mut self, array: &dyn Array) {
168        let array = array.to_primitive();
169
170        // This should be checked in `extend_from_array` but we can check it again.
171        debug_assert_eq!(
172            array.ptype(),
173            T::PTYPE,
174            "Cannot extend from array with different ptype"
175        );
176
177        self.values.extend_from_slice(array.as_slice::<T>());
178        self.nulls.append_validity_mask(array.validity_mask());
179    }
180
181    fn reserve_exact(&mut self, additional: usize) {
182        self.values.reserve(additional);
183        self.nulls.reserve_exact(additional);
184    }
185
186    unsafe fn set_validity_unchecked(&mut self, validity: Mask) {
187        self.nulls = LazyBitBufferBuilder::new(validity.len());
188        self.nulls.append_validity_mask(validity);
189    }
190
191    fn finish(&mut self) -> ArrayRef {
192        self.finish_into_primitive().into_array()
193    }
194
195    fn finish_into_canonical(&mut self) -> Canonical {
196        Canonical::Primitive(self.finish_into_primitive())
197    }
198}
199
200/// A range of uninitialized values in the primitive builder that can be filled.
201pub struct UninitRange<'a, T> {
202    /// The length of the uninitialized range.
203    ///
204    /// This is guaranteed to be within the memory capacity of the builder.
205    len: usize,
206
207    /// A mutable reference to the builder.
208    ///
209    /// Since this is a mutable reference, we can guarantee that nothing else can modify the builder
210    /// while this `UninitRange` exists.
211    builder: &'a mut PrimitiveBuilder<T>,
212}
213
214impl<T> UninitRange<'_, T> {
215    /// Returns the length of this uninitialized range.
216    #[inline]
217    pub fn len(&self) -> usize {
218        self.len
219    }
220
221    /// Returns true if this range has zero length.
222    #[inline]
223    pub fn is_empty(&self) -> bool {
224        self.len == 0
225    }
226
227    /// Set a value at the given index within this range.
228    ///
229    /// # Panics
230    ///
231    /// Panics if the index is out of bounds.
232    #[inline]
233    pub fn set_value(&mut self, index: usize, value: T) {
234        assert!(index < self.len, "index out of bounds");
235        let spare = self.builder.values.spare_capacity_mut();
236        spare[index] = MaybeUninit::new(value);
237    }
238
239    /// Append a [`Mask`] to this builder's null buffer.
240    ///
241    /// # Panics
242    ///
243    /// Panics if the mask length is not equal to the the length of the current `UninitRange`.
244    ///
245    /// # Safety
246    ///
247    /// - The caller must ensure that they safely initialize `mask.len()` primitive values via
248    ///   [`UninitRange::copy_from_slice`].
249    /// - The caller must also ensure that they only call this method once.
250    pub unsafe fn append_mask(&mut self, mask: Mask) {
251        assert_eq!(
252            mask.len(),
253            self.len,
254            "Tried to append a mask to an `UninitRange` that was beyond the allowed range"
255        );
256
257        // TODO(connor): Ideally, we would call this function `set_mask` and directly set all of the
258        // bits (so that we can call this multiple times), but the underlying `BooleanBuffer` does
259        // not have an easy way to do this correctly.
260
261        self.builder.nulls.append_validity_mask(mask);
262    }
263
264    /// Set a validity bit at the given index.
265    ///
266    /// The index is relative to the start of this range (not relative to the values already in the
267    /// builder).
268    ///
269    /// Note that this will have no effect if the builder is non-nullable.
270    pub fn set_validity_bit(&mut self, index: usize, v: bool) {
271        assert!(index < self.len, "set_bit index out of bounds");
272        // Note that this won't panic because we can only create an `UninitRange` within the
273        // capacity of the builder (it will not automatically resize).
274        let absolute_index = self.builder.values.len() + index;
275        self.builder.nulls.set_bit(absolute_index, v);
276    }
277
278    /// Set values from an initialized range.
279    ///
280    /// Note that the input `offset` should be an offset relative to the local `UninitRange`, not
281    /// the entire `PrimitiveBuilder`.
282    pub fn copy_from_slice(&mut self, local_offset: usize, src: &[T])
283    where
284        T: Copy,
285    {
286        debug_assert!(
287            local_offset + src.len() <= self.len,
288            "tried to copy a slice into a `UninitRange` past its boundary"
289        );
290
291        // SAFETY: &[T] and &[MaybeUninit<T>] have the same layout.
292        let uninit_src: &[MaybeUninit<T>] = unsafe { std::mem::transmute(src) };
293
294        // Note: spare_capacity_mut() returns the spare capacity starting from the current length,
295        // so we just use local_offset directly.
296        let dst =
297            &mut self.builder.values.spare_capacity_mut()[local_offset..local_offset + src.len()];
298        dst.copy_from_slice(uninit_src);
299    }
300
301    /// Get a mutable slice of uninitialized memory at the specified offset within this range.
302    ///
303    /// Note that the offsets are relative to this local range, not to the values already in the
304    /// builder.
305    ///
306    /// # Safety
307    ///
308    /// The caller must ensure that they properly initialize the returned memory before calling
309    /// `finish()` on this range.
310    ///
311    /// # Panics
312    ///
313    /// Panics if `offset + len` exceeds the range bounds.
314    pub unsafe fn slice_uninit_mut(&mut self, offset: usize, len: usize) -> &mut [MaybeUninit<T>] {
315        assert!(
316            offset + len <= self.len,
317            "slice_uninit_mut: offset {} + len {} exceeds range length {}",
318            offset,
319            len,
320            self.len
321        );
322        &mut self.builder.values.spare_capacity_mut()[offset..offset + len]
323    }
324
325    /// Finish building this range, marking it as initialized and advancing the length of the
326    /// underlying values buffer.
327    ///
328    /// # Safety
329    ///
330    /// The caller must ensure that they have safely initialized all `len` values via
331    /// [`copy_from_slice()`] or [`set_value()`], as well as correctly set all of the null bits via
332    /// [`set_validity_bit()`] or [`append_mask()`] if the builder is nullable.
333    ///
334    /// [`copy_from_slice()`]: UninitRange::copy_from_slice
335    /// [`set_value()`]: UninitRange::set_value
336    /// [`set_validity_bit()`]: UninitRange::set_validity_bit
337    /// [`append_mask()`]: UninitRange::append_mask
338    pub unsafe fn finish(self) {
339        // SAFETY: constructor enforces that current length + len does not exceed the capacity of the array.
340        let new_len = self.builder.values.len() + self.len;
341        unsafe { self.builder.values.set_len(new_len) };
342    }
343}
344
345#[cfg(test)]
346mod tests {
347    use super::*;
348
349    /// REGRESSION TEST: This test verifies that multiple sequential ranges have correct offsets.
350    ///
351    /// This would have caught the `Deref` bug where it always returned from the start of the
352    /// buffer.
353    #[test]
354    fn test_multiple_uninit_ranges_correct_offsets() {
355        let mut builder = PrimitiveBuilder::<i32>::with_capacity(Nullability::NonNullable, 10);
356
357        // First range.
358        let mut range1 = builder.uninit_range(3);
359        range1.copy_from_slice(0, &[1, 2, 3]);
360
361        // SAFETY: We initialized all 3 values.
362        unsafe {
363            range1.finish();
364        }
365
366        // Verify the builder now has these values.
367        assert_eq!(builder.values(), &[1, 2, 3]);
368
369        // Second range - this would fail with the old Deref implementation.
370        let mut range2 = builder.uninit_range(2);
371
372        // Set values using copy_from_slice.
373        range2.copy_from_slice(0, &[4, 5]);
374
375        // SAFETY: We initialized both values.
376        unsafe {
377            range2.finish();
378        }
379
380        // Verify the builder now has all 5 values.
381        assert_eq!(builder.values(), &[1, 2, 3, 4, 5]);
382
383        let array = builder.finish_into_primitive();
384        assert_eq!(array.as_slice::<i32>(), &[1, 2, 3, 4, 5]);
385    }
386
387    /// REGRESSION TEST: This test verifies that `append_mask` was correctly moved from
388    /// `PrimitiveBuilder` to `UninitRange`.
389    ///
390    /// The old API had `append_mask` on the builder, which was confusing when used with ranges.
391    /// This test ensures the new API works correctly.
392    #[test]
393    fn test_append_mask_on_uninit_range() {
394        let mut builder = PrimitiveBuilder::<i32>::with_capacity(Nullability::Nullable, 5);
395        let mut range = builder.uninit_range(3);
396
397        // Create a mask for 3 values.
398        let mask = Mask::from_iter([true, false, true]);
399
400        // SAFETY: We're about to initialize the values.
401        unsafe {
402            range.append_mask(mask);
403        }
404
405        // Initialize the values.
406        range.copy_from_slice(0, &[10, 20, 30]);
407
408        // SAFETY: We've initialized all values and set the mask.
409        unsafe {
410            range.finish();
411        }
412
413        let array = builder.finish_into_primitive();
414        assert_eq!(array.len(), 3);
415        // Check validity using scalar_at - nulls will return is_null() = true.
416        assert!(!array.scalar_at(0).is_null());
417        assert!(array.scalar_at(1).is_null());
418        assert!(!array.scalar_at(2).is_null());
419    }
420
421    /// REGRESSION TEST: This test verifies that `append_mask` validates the mask length.
422    ///
423    /// This ensures that masks can only be appended if they match the range length.
424    #[test]
425    #[should_panic(
426        expected = "Tried to append a mask to an `UninitRange` that was beyond the allowed range"
427    )]
428    fn test_append_mask_wrong_length_panics() {
429        let mut builder = PrimitiveBuilder::<i32>::with_capacity(Nullability::Nullable, 10);
430        let mut range = builder.uninit_range(5);
431
432        // Try to append a mask with wrong length (3 instead of 5).
433        let wrong_mask = Mask::from_iter([true, false, true]);
434
435        // SAFETY: This is expected to panic due to length mismatch.
436        unsafe {
437            range.append_mask(wrong_mask);
438        }
439    }
440
441    /// Test that `copy_from_slice` works correctly with different offsets.
442    ///
443    /// This verifies the new simplified API without the redundant `len` parameter.
444    #[test]
445    fn test_copy_from_slice_with_offsets() {
446        let mut builder = PrimitiveBuilder::<i32>::with_capacity(Nullability::NonNullable, 10);
447        let mut range = builder.uninit_range(6);
448
449        // Copy to different offsets.
450        range.copy_from_slice(0, &[1, 2]);
451        range.copy_from_slice(2, &[3, 4]);
452        range.copy_from_slice(4, &[5, 6]);
453
454        // SAFETY: We've initialized all 6 values.
455        unsafe {
456            range.finish();
457        }
458
459        let array = builder.finish_into_primitive();
460        assert_eq!(array.as_slice::<i32>(), &[1, 2, 3, 4, 5, 6]);
461    }
462
463    /// Test that `set_bit` uses relative indexing within the range.
464    ///
465    /// Note: `set_bit` requires the null buffer to already be initialized, so we first
466    /// use `append_mask` to set up the buffer, then demonstrate that `set_bit` can
467    /// modify individual bits with relative indexing.
468    #[test]
469    fn test_set_bit_relative_indexing() {
470        let mut builder = PrimitiveBuilder::<i32>::with_capacity(Nullability::Nullable, 10);
471
472        // First add some values to the builder.
473        builder.append_value(100);
474        builder.append_value(200);
475
476        // Create a range for new values.
477        let mut range = builder.uninit_range(3);
478
479        // Use append_mask to initialize the validity buffer for this range.
480        let initial_mask = Mask::from_iter([false, false, false]);
481        // SAFETY: We're about to initialize the values.
482        unsafe {
483            range.append_mask(initial_mask);
484        }
485
486        // Now we can use set_bit to modify individual bits with relative indexing.
487        range.set_validity_bit(0, true); // Change first bit to valid
488        range.set_validity_bit(2, true); // Change third bit to valid
489        // Leave middle bit as false (null)
490
491        // Initialize the values.
492        range.copy_from_slice(0, &[10, 20, 30]);
493
494        // SAFETY: We've initialized all 3 values and set their validity.
495        unsafe {
496            range.finish();
497        }
498
499        let array = builder.finish_into_primitive();
500
501        // Verify the total length and values.
502        assert_eq!(array.len(), 5);
503        assert_eq!(array.as_slice::<i32>(), &[100, 200, 10, 20, 30]);
504
505        // Check validity - the first two should be valid (from append_value).
506        assert!(!array.scalar_at(0).is_null()); // initial value 100
507        assert!(!array.scalar_at(1).is_null()); // initial value 200
508
509        // Check the range items with modified validity.
510        assert!(!array.scalar_at(2).is_null()); // range index 0 - set to valid
511        assert!(array.scalar_at(3).is_null()); // range index 1 - left as null
512        assert!(!array.scalar_at(4).is_null()); // range index 2 - set to valid
513    }
514
515    /// Test that creating a zero-length uninit range panics.
516    #[test]
517    #[should_panic(expected = "cannot create an uninit range of length 0")]
518    fn test_zero_length_uninit_range_panics() {
519        let mut builder = PrimitiveBuilder::<i32>::new(Nullability::NonNullable);
520        let _range = builder.uninit_range(0);
521    }
522
523    /// Test that creating an uninit range exceeding capacity panics.
524    #[test]
525    #[should_panic(
526        expected = "uninit_range of len 10 exceeds builder with length 0 and capacity 6"
527    )]
528    fn test_uninit_range_exceeds_capacity_panics() {
529        let mut builder = PrimitiveBuilder::<i32>::with_capacity(Nullability::NonNullable, 5);
530        let _range = builder.uninit_range(10);
531    }
532
533    /// Test that `copy_from_slice` debug asserts on out-of-bounds access.
534    ///
535    /// Note: This only panics in debug mode due to `debug_assert!`.
536    #[test]
537    #[cfg(debug_assertions)]
538    #[should_panic(expected = "tried to copy a slice into a `UninitRange` past its boundary")]
539    fn test_copy_from_slice_out_of_bounds() {
540        let mut builder = PrimitiveBuilder::<i32>::with_capacity(Nullability::NonNullable, 10);
541        let mut range = builder.uninit_range(3);
542
543        // Try to copy 3 elements starting at offset 1 (would need 4 slots total).
544        range.copy_from_slice(1, &[1, 2, 3]);
545    }
546
547    /// Test that the unsafe contract of `finish` is documented and works correctly.
548    ///
549    /// This test demonstrates proper usage of the unsafe `finish` method.
550    #[test]
551    fn test_finish_unsafe_contract() {
552        let mut builder = PrimitiveBuilder::<i32>::with_capacity(Nullability::Nullable, 5);
553        let mut range = builder.uninit_range(3);
554
555        // Set validity mask.
556        let mask = Mask::from_iter([true, true, false]);
557        // SAFETY: We're about to initialize the matching number of values.
558        unsafe {
559            range.append_mask(mask);
560        }
561
562        // Initialize all values.
563        range.copy_from_slice(0, &[10, 20, 30]);
564
565        // SAFETY: We have initialized all 3 values and set their validity.
566        unsafe {
567            range.finish();
568        }
569
570        let array = builder.finish_into_primitive();
571        assert_eq!(array.len(), 3);
572        assert_eq!(array.as_slice::<i32>(), &[10, 20, 30]);
573    }
574
575    #[test]
576    fn test_append_scalar() {
577        use vortex_dtype::DType;
578        use vortex_scalar::Scalar;
579
580        let mut builder = PrimitiveBuilder::<i32>::with_capacity(Nullability::Nullable, 10);
581
582        // Test appending a valid primitive value.
583        let scalar1 = Scalar::primitive(42i32, Nullability::Nullable);
584        builder.append_scalar(&scalar1).unwrap();
585
586        // Test appending another value.
587        let scalar2 = Scalar::primitive(84i32, Nullability::Nullable);
588        builder.append_scalar(&scalar2).unwrap();
589
590        // Test appending null value.
591        let null_scalar = Scalar::null(DType::Primitive(
592            vortex_dtype::PType::I32,
593            Nullability::Nullable,
594        ));
595        builder.append_scalar(&null_scalar).unwrap();
596
597        let array = builder.finish_into_primitive();
598        assert_eq!(array.len(), 3);
599
600        // Check actual values.
601        let values = array.as_slice::<i32>();
602        assert_eq!(values[0], 42);
603        assert_eq!(values[1], 84);
604        // values[2] might be any value since it's null.
605
606        // Check validity - first two should be valid, third should be null.
607        use crate::vtable::ValidityHelper;
608        assert!(array.validity().is_valid(0));
609        assert!(array.validity().is_valid(1));
610        assert!(!array.validity().is_valid(2));
611
612        // Test wrong dtype error.
613        let mut builder = PrimitiveBuilder::<i32>::with_capacity(Nullability::NonNullable, 10);
614        let wrong_scalar = Scalar::from(true);
615        assert!(builder.append_scalar(&wrong_scalar).is_err());
616    }
617}
vortex_array/builders/primitive.rs

vortex_array/builders/
primitive.rs
