1use std::any::Any;
5use std::mem::MaybeUninit;
6
7use vortex_buffer::BufferMut;
8use vortex_error::VortexExpect;
9use vortex_error::VortexResult;
10use vortex_error::vortex_ensure;
11use vortex_mask::Mask;
12
13use crate::ArrayRef;
14use crate::IntoArray;
15use crate::LEGACY_SESSION;
16use crate::VortexSessionExecute;
17use crate::arrays::PrimitiveArray;
18use crate::builders::ArrayBuilder;
19use crate::builders::DEFAULT_BUILDER_CAPACITY;
20use crate::builders::LazyBitBufferBuilder;
21use crate::canonical::Canonical;
22#[expect(deprecated)]
23use crate::canonical::ToCanonical as _;
24use crate::dtype::DType;
25use crate::dtype::NativePType;
26use crate::dtype::Nullability;
27use crate::scalar::Scalar;
28
29pub struct PrimitiveBuilder<T> {
31 dtype: DType,
32 values: BufferMut<T>,
33 nulls: LazyBitBufferBuilder,
34}
35
36impl<T: NativePType> PrimitiveBuilder<T> {
37 pub fn new(nullability: Nullability) -> Self {
39 Self::with_capacity(nullability, DEFAULT_BUILDER_CAPACITY)
40 }
41
42 pub fn with_capacity(nullability: Nullability, capacity: usize) -> Self {
44 Self {
45 values: BufferMut::with_capacity(capacity),
46 nulls: LazyBitBufferBuilder::new(capacity),
47 dtype: DType::Primitive(T::PTYPE, nullability),
48 }
49 }
50
51 pub fn append_value(&mut self, value: T) {
53 self.values.push(value);
54 self.nulls.append_non_null();
55 }
56
57 pub fn append_n_values(&mut self, value: T, n: usize) {
59 self.values.push_n(value, n);
60 self.nulls.append_n_non_nulls(n);
61 }
62
63 pub fn values(&self) -> &[T] {
65 self.values.as_ref()
66 }
67
68 pub fn values_mut(&mut self) -> &mut [T] {
70 self.values.as_mut()
71 }
72
73 pub fn uninit_range(&mut self, len: usize) -> UninitRange<'_, T> {
107 assert_ne!(0, len, "cannot create an uninit range of length 0");
108
109 let current_len = self.values.len();
110 assert!(
111 current_len + len <= self.values.capacity(),
112 "uninit_range of len {len} exceeds builder with length {} and capacity {}",
113 current_len,
114 self.values.capacity()
115 );
116
117 UninitRange { len, builder: self }
118 }
119
120 pub fn finish_into_primitive(&mut self) -> PrimitiveArray {
122 let validity = self
123 .nulls
124 .finish_with_nullability(self.dtype().nullability());
125
126 PrimitiveArray::new(std::mem::take(&mut self.values).freeze(), validity)
127 }
128
129 pub fn extend_with_iterator(&mut self, iter: impl IntoIterator<Item = T>, mask: Mask) {
131 self.values.extend(iter);
132 self.nulls.append_validity_mask(mask);
133 }
134}
135
136impl<T: NativePType> ArrayBuilder for PrimitiveBuilder<T> {
137 fn as_any(&self) -> &dyn Any {
138 self
139 }
140
141 fn as_any_mut(&mut self) -> &mut dyn Any {
142 self
143 }
144
145 fn dtype(&self) -> &DType {
146 &self.dtype
147 }
148
149 fn len(&self) -> usize {
150 self.values.len()
151 }
152
153 fn append_zeros(&mut self, n: usize) {
154 self.values.push_n(T::default(), n);
155 self.nulls.append_n_non_nulls(n);
156 }
157
158 unsafe fn append_nulls_unchecked(&mut self, n: usize) {
159 self.values.push_n(T::default(), n);
160 self.nulls.append_n_nulls(n);
161 }
162
163 fn append_scalar(&mut self, scalar: &Scalar) -> VortexResult<()> {
164 vortex_ensure!(
165 scalar.dtype() == self.dtype(),
166 "PrimitiveBuilder expected scalar with dtype {}, got {}",
167 self.dtype(),
168 scalar.dtype()
169 );
170
171 if let Some(pv) = scalar.as_primitive().pvalue() {
172 self.append_value(pv.cast::<T>()?)
173 } else {
174 self.append_null()
175 }
176
177 Ok(())
178 }
179
180 unsafe fn extend_from_array_unchecked(&mut self, array: &ArrayRef) {
181 #[expect(deprecated)]
182 let array = array.to_primitive();
183
184 debug_assert_eq!(
186 array.ptype(),
187 T::PTYPE,
188 "Cannot extend from array with different ptype"
189 );
190
191 self.values.extend_from_slice(array.as_slice::<T>());
192 self.nulls.append_validity_mask(
193 array
194 .as_ref()
195 .validity()
196 .vortex_expect("validity_mask")
197 .execute_mask(
198 array.as_ref().len(),
199 &mut LEGACY_SESSION.create_execution_ctx(),
200 )
201 .vortex_expect("Failed to compute validity mask"),
202 );
203 }
204
205 fn reserve_exact(&mut self, additional: usize) {
206 self.values.reserve(additional);
207 self.nulls.reserve_exact(additional);
208 }
209
210 unsafe fn set_validity_unchecked(&mut self, validity: Mask) {
211 self.nulls = LazyBitBufferBuilder::new(validity.len());
212 self.nulls.append_validity_mask(validity);
213 }
214
215 fn finish(&mut self) -> ArrayRef {
216 self.finish_into_primitive().into_array()
217 }
218
219 fn finish_into_canonical(&mut self) -> Canonical {
220 Canonical::Primitive(self.finish_into_primitive())
221 }
222}
223
224pub struct UninitRange<'a, T> {
226 len: usize,
230
231 builder: &'a mut PrimitiveBuilder<T>,
236}
237
238impl<T> UninitRange<'_, T> {
239 #[inline]
241 pub fn len(&self) -> usize {
242 self.len
243 }
244
245 #[inline]
247 pub fn is_empty(&self) -> bool {
248 self.len == 0
249 }
250
251 #[inline]
257 pub fn set_value(&mut self, index: usize, value: T) {
258 assert!(index < self.len, "index out of bounds");
259 let spare = self.builder.values.spare_capacity_mut();
260 spare[index] = MaybeUninit::new(value);
261 }
262
263 pub unsafe fn append_mask(&mut self, mask: Mask) {
275 assert_eq!(
276 mask.len(),
277 self.len,
278 "Tried to append a mask to an `UninitRange` that was beyond the allowed range"
279 );
280
281 self.builder.nulls.append_validity_mask(mask);
286 }
287
288 pub fn set_validity_bit(&mut self, index: usize, v: bool) {
295 assert!(index < self.len, "set_bit index out of bounds");
296 let absolute_index = self.builder.values.len() + index;
299 self.builder.nulls.set_bit(absolute_index, v);
300 }
301
302 pub fn copy_from_slice(&mut self, local_offset: usize, src: &[T])
307 where
308 T: Copy,
309 {
310 debug_assert!(
311 local_offset + src.len() <= self.len,
312 "tried to copy a slice into a `UninitRange` past its boundary"
313 );
314
315 let uninit_src: &[MaybeUninit<T>] = unsafe { std::mem::transmute(src) };
317
318 let dst =
321 &mut self.builder.values.spare_capacity_mut()[local_offset..local_offset + src.len()];
322 dst.copy_from_slice(uninit_src);
323 }
324
325 pub unsafe fn slice_uninit_mut(&mut self, offset: usize, len: usize) -> &mut [MaybeUninit<T>] {
339 assert!(
340 offset + len <= self.len,
341 "slice_uninit_mut: offset {} + len {} exceeds range length {}",
342 offset,
343 len,
344 self.len
345 );
346 &mut self.builder.values.spare_capacity_mut()[offset..offset + len]
347 }
348
349 pub unsafe fn finish(self) {
363 let new_len = self.builder.values.len() + self.len;
365 unsafe { self.builder.values.set_len(new_len) };
366 }
367}
368
369#[cfg(test)]
370mod tests {
371 use vortex_error::VortexExpect;
372
373 use super::*;
374 use crate::assert_arrays_eq;
375
376 #[test]
381 fn test_multiple_uninit_ranges_correct_offsets() {
382 let mut builder = PrimitiveBuilder::<i32>::with_capacity(Nullability::NonNullable, 10);
383
384 let mut range1 = builder.uninit_range(3);
386 range1.copy_from_slice(0, &[1, 2, 3]);
387
388 unsafe {
390 range1.finish();
391 }
392
393 assert_eq!(builder.values(), &[1, 2, 3]);
395
396 let mut range2 = builder.uninit_range(2);
398
399 range2.copy_from_slice(0, &[4, 5]);
401
402 unsafe {
404 range2.finish();
405 }
406
407 assert_eq!(builder.values(), &[1, 2, 3, 4, 5]);
409
410 let array = builder.finish_into_primitive();
411 assert_arrays_eq!(array, PrimitiveArray::from_iter([1i32, 2, 3, 4, 5]));
412 }
413
414 #[test]
420 fn test_append_mask_on_uninit_range() {
421 let mut builder = PrimitiveBuilder::<i32>::with_capacity(Nullability::Nullable, 5);
422 let mut range = builder.uninit_range(3);
423
424 let mask = Mask::from_iter([true, false, true]);
426
427 unsafe {
429 range.append_mask(mask);
430 }
431
432 range.copy_from_slice(0, &[10, 20, 30]);
434
435 unsafe {
437 range.finish();
438 }
439
440 let array = builder.finish_into_primitive();
441 assert_eq!(array.len(), 3);
442 assert!(
444 !array
445 .execute_scalar(0, &mut LEGACY_SESSION.create_execution_ctx())
446 .unwrap()
447 .is_null()
448 );
449 assert!(
450 array
451 .execute_scalar(1, &mut LEGACY_SESSION.create_execution_ctx())
452 .unwrap()
453 .is_null()
454 );
455 assert!(
456 !array
457 .execute_scalar(2, &mut LEGACY_SESSION.create_execution_ctx())
458 .unwrap()
459 .is_null()
460 );
461 }
462
463 #[test]
467 #[should_panic(
468 expected = "Tried to append a mask to an `UninitRange` that was beyond the allowed range"
469 )]
470 fn test_append_mask_wrong_length_panics() {
471 let mut builder = PrimitiveBuilder::<i32>::with_capacity(Nullability::Nullable, 10);
472 let mut range = builder.uninit_range(5);
473
474 let wrong_mask = Mask::from_iter([true, false, true]);
476
477 unsafe {
479 range.append_mask(wrong_mask);
480 }
481 }
482
483 #[test]
487 fn test_copy_from_slice_with_offsets() {
488 let mut builder = PrimitiveBuilder::<i32>::with_capacity(Nullability::NonNullable, 10);
489 let mut range = builder.uninit_range(6);
490
491 range.copy_from_slice(0, &[1, 2]);
493 range.copy_from_slice(2, &[3, 4]);
494 range.copy_from_slice(4, &[5, 6]);
495
496 unsafe {
498 range.finish();
499 }
500
501 let array = builder.finish_into_primitive();
502 assert_arrays_eq!(array, PrimitiveArray::from_iter([1i32, 2, 3, 4, 5, 6]));
503 }
504
505 #[test]
511 fn test_set_bit_relative_indexing() {
512 let mut builder = PrimitiveBuilder::<i32>::with_capacity(Nullability::Nullable, 10);
513
514 builder.append_value(100);
516 builder.append_value(200);
517
518 let mut range = builder.uninit_range(3);
520
521 let initial_mask = Mask::from_iter([false, false, false]);
523 unsafe {
525 range.append_mask(initial_mask);
526 }
527
528 range.set_validity_bit(0, true); range.set_validity_bit(2, true); range.copy_from_slice(0, &[10, 20, 30]);
535
536 unsafe {
538 range.finish();
539 }
540
541 let array = builder.finish_into_primitive();
542
543 assert_eq!(array.len(), 5);
545 assert_eq!(array.as_slice::<i32>(), &[100, 200, 10, 20, 30]);
546
547 assert!(
549 !array
550 .execute_scalar(0, &mut LEGACY_SESSION.create_execution_ctx())
551 .unwrap()
552 .is_null()
553 ); assert!(
555 !array
556 .execute_scalar(1, &mut LEGACY_SESSION.create_execution_ctx())
557 .unwrap()
558 .is_null()
559 ); assert!(
563 !array
564 .execute_scalar(2, &mut LEGACY_SESSION.create_execution_ctx())
565 .unwrap()
566 .is_null()
567 ); assert!(
569 array
570 .execute_scalar(3, &mut LEGACY_SESSION.create_execution_ctx())
571 .unwrap()
572 .is_null()
573 ); assert!(
575 !array
576 .execute_scalar(4, &mut LEGACY_SESSION.create_execution_ctx())
577 .unwrap()
578 .is_null()
579 ); }
581
582 #[test]
584 #[should_panic(expected = "cannot create an uninit range of length 0")]
585 fn test_zero_length_uninit_range_panics() {
586 let mut builder = PrimitiveBuilder::<i32>::new(Nullability::NonNullable);
587 let _range = builder.uninit_range(0);
588 }
589
590 #[test]
592 #[should_panic(
593 expected = "uninit_range of len 10 exceeds builder with length 0 and capacity 6"
594 )]
595 fn test_uninit_range_exceeds_capacity_panics() {
596 let mut builder = PrimitiveBuilder::<i32>::with_capacity(Nullability::NonNullable, 5);
597 let _range = builder.uninit_range(10);
598 }
599
600 #[test]
604 #[cfg(debug_assertions)]
605 #[should_panic(expected = "tried to copy a slice into a `UninitRange` past its boundary")]
606 fn test_copy_from_slice_out_of_bounds() {
607 let mut builder = PrimitiveBuilder::<i32>::with_capacity(Nullability::NonNullable, 10);
608 let mut range = builder.uninit_range(3);
609
610 range.copy_from_slice(1, &[1, 2, 3]);
612 }
613
614 #[test]
618 fn test_finish_unsafe_contract() {
619 let mut builder = PrimitiveBuilder::<i32>::with_capacity(Nullability::Nullable, 5);
620 let mut range = builder.uninit_range(3);
621
622 let mask = Mask::from_iter([true, true, false]);
624 unsafe {
626 range.append_mask(mask);
627 }
628
629 range.copy_from_slice(0, &[10, 20, 30]);
631
632 unsafe {
634 range.finish();
635 }
636
637 let array = builder.finish_into_primitive();
638 assert_eq!(array.len(), 3);
639 assert_eq!(array.as_slice::<i32>(), &[10, 20, 30]);
640 }
641
642 #[test]
643 fn test_append_scalar() {
644 use crate::dtype::DType;
645 use crate::scalar::Scalar;
646
647 let mut builder = PrimitiveBuilder::<i32>::with_capacity(Nullability::Nullable, 10);
648
649 let scalar1 = Scalar::primitive(42i32, Nullability::Nullable);
651 builder.append_scalar(&scalar1).unwrap();
652
653 let scalar2 = Scalar::primitive(84i32, Nullability::Nullable);
655 builder.append_scalar(&scalar2).unwrap();
656
657 let null_scalar = Scalar::null(DType::Primitive(
659 crate::dtype::PType::I32,
660 Nullability::Nullable,
661 ));
662 builder.append_scalar(&null_scalar).unwrap();
663
664 let array = builder.finish_into_primitive();
665 assert_eq!(array.len(), 3);
666
667 let values = array.as_slice::<i32>();
669 assert_eq!(values[0], 42);
670 assert_eq!(values[1], 84);
671 assert!(
675 array
676 .validity()
677 .vortex_expect("primitive validity should be derivable")
678 .is_valid(0)
679 .unwrap()
680 );
681 assert!(
682 array
683 .validity()
684 .vortex_expect("primitive validity should be derivable")
685 .is_valid(1)
686 .unwrap()
687 );
688 assert!(
689 !array
690 .validity()
691 .vortex_expect("primitive validity should be derivable")
692 .is_valid(2)
693 .unwrap()
694 );
695
696 let mut builder = PrimitiveBuilder::<i32>::with_capacity(Nullability::NonNullable, 10);
698 let wrong_scalar = Scalar::from(true);
699 assert!(builder.append_scalar(&wrong_scalar).is_err());
700 }
701}