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::from_validity_mask(validity);
212 }
213
214 fn finish(&mut self) -> ArrayRef {
215 self.finish_into_primitive().into_array()
216 }
217
218 fn finish_into_canonical(&mut self) -> Canonical {
219 Canonical::Primitive(self.finish_into_primitive())
220 }
221}
222
223pub struct UninitRange<'a, T> {
225 len: usize,
229
230 builder: &'a mut PrimitiveBuilder<T>,
235}
236
237impl<T> UninitRange<'_, T> {
238 #[inline]
240 pub fn len(&self) -> usize {
241 self.len
242 }
243
244 #[inline]
246 pub fn is_empty(&self) -> bool {
247 self.len == 0
248 }
249
250 #[inline]
256 pub fn set_value(&mut self, index: usize, value: T) {
257 assert!(index < self.len, "index out of bounds");
258 let spare = self.builder.values.spare_capacity_mut();
259 spare[index] = MaybeUninit::new(value);
260 }
261
262 pub unsafe fn append_mask(&mut self, mask: &Mask) {
274 assert_eq!(
275 mask.len(),
276 self.len,
277 "Tried to append a mask to an `UninitRange` that was beyond the allowed range"
278 );
279
280 self.builder.nulls.append_validity_mask(mask);
285 }
286
287 pub fn set_validity_bit(&mut self, index: usize, v: bool) {
294 assert!(index < self.len, "set_bit index out of bounds");
295 let absolute_index = self.builder.values.len() + index;
298 self.builder.nulls.set_bit(absolute_index, v);
299 }
300
301 pub fn copy_from_slice(&mut self, local_offset: usize, src: &[T])
306 where
307 T: Copy,
308 {
309 debug_assert!(
310 local_offset + src.len() <= self.len,
311 "tried to copy a slice into a `UninitRange` past its boundary"
312 );
313
314 let uninit_src: &[MaybeUninit<T>] = unsafe { std::mem::transmute(src) };
316
317 let dst =
320 &mut self.builder.values.spare_capacity_mut()[local_offset..local_offset + src.len()];
321 dst.copy_from_slice(uninit_src);
322 }
323
324 pub unsafe fn slice_uninit_mut(&mut self, offset: usize, len: usize) -> &mut [MaybeUninit<T>] {
338 assert!(
339 offset + len <= self.len,
340 "slice_uninit_mut: offset {} + len {} exceeds range length {}",
341 offset,
342 len,
343 self.len
344 );
345 &mut self.builder.values.spare_capacity_mut()[offset..offset + len]
346 }
347
348 pub unsafe fn finish(self) {
362 let new_len = self.builder.values.len() + self.len;
364 unsafe { self.builder.values.set_len(new_len) };
365 }
366}
367
368#[cfg(test)]
369mod tests {
370 use vortex_error::VortexExpect;
371
372 use super::*;
373 use crate::assert_arrays_eq;
374
375 #[test]
380 fn test_multiple_uninit_ranges_correct_offsets() {
381 let mut builder = PrimitiveBuilder::<i32>::with_capacity(Nullability::NonNullable, 10);
382
383 let mut range1 = builder.uninit_range(3);
385 range1.copy_from_slice(0, &[1, 2, 3]);
386
387 unsafe {
389 range1.finish();
390 }
391
392 assert_eq!(builder.values(), &[1, 2, 3]);
394
395 let mut range2 = builder.uninit_range(2);
397
398 range2.copy_from_slice(0, &[4, 5]);
400
401 unsafe {
403 range2.finish();
404 }
405
406 assert_eq!(builder.values(), &[1, 2, 3, 4, 5]);
408
409 let array = builder.finish_into_primitive();
410 assert_arrays_eq!(array, PrimitiveArray::from_iter([1i32, 2, 3, 4, 5]));
411 }
412
413 #[test]
419 fn test_append_mask_on_uninit_range() {
420 let mut builder = PrimitiveBuilder::<i32>::with_capacity(Nullability::Nullable, 5);
421 let mut range = builder.uninit_range(3);
422
423 let mask = Mask::from_iter([true, false, true]);
425
426 unsafe {
428 range.append_mask(&mask);
429 }
430
431 range.copy_from_slice(0, &[10, 20, 30]);
433
434 unsafe {
436 range.finish();
437 }
438
439 let array = builder.finish_into_primitive();
440 assert_eq!(array.len(), 3);
441 assert!(
443 !array
444 .execute_scalar(0, &mut LEGACY_SESSION.create_execution_ctx())
445 .unwrap()
446 .is_null()
447 );
448 assert!(
449 array
450 .execute_scalar(1, &mut LEGACY_SESSION.create_execution_ctx())
451 .unwrap()
452 .is_null()
453 );
454 assert!(
455 !array
456 .execute_scalar(2, &mut LEGACY_SESSION.create_execution_ctx())
457 .unwrap()
458 .is_null()
459 );
460 }
461
462 #[test]
466 #[should_panic(
467 expected = "Tried to append a mask to an `UninitRange` that was beyond the allowed range"
468 )]
469 fn test_append_mask_wrong_length_panics() {
470 let mut builder = PrimitiveBuilder::<i32>::with_capacity(Nullability::Nullable, 10);
471 let mut range = builder.uninit_range(5);
472
473 let wrong_mask = Mask::from_iter([true, false, true]);
475
476 unsafe {
478 range.append_mask(&wrong_mask);
479 }
480 }
481
482 #[test]
486 fn test_copy_from_slice_with_offsets() {
487 let mut builder = PrimitiveBuilder::<i32>::with_capacity(Nullability::NonNullable, 10);
488 let mut range = builder.uninit_range(6);
489
490 range.copy_from_slice(0, &[1, 2]);
492 range.copy_from_slice(2, &[3, 4]);
493 range.copy_from_slice(4, &[5, 6]);
494
495 unsafe {
497 range.finish();
498 }
499
500 let array = builder.finish_into_primitive();
501 assert_arrays_eq!(array, PrimitiveArray::from_iter([1i32, 2, 3, 4, 5, 6]));
502 }
503
504 #[test]
510 fn test_set_bit_relative_indexing() {
511 let mut builder = PrimitiveBuilder::<i32>::with_capacity(Nullability::Nullable, 10);
512
513 builder.append_value(100);
515 builder.append_value(200);
516
517 let mut range = builder.uninit_range(3);
519
520 let initial_mask = Mask::from_iter([false, false, false]);
522 unsafe {
524 range.append_mask(&initial_mask);
525 }
526
527 range.set_validity_bit(0, true); range.set_validity_bit(2, true); range.copy_from_slice(0, &[10, 20, 30]);
534
535 unsafe {
537 range.finish();
538 }
539
540 let array = builder.finish_into_primitive();
541
542 assert_eq!(array.len(), 5);
544 assert_eq!(array.as_slice::<i32>(), &[100, 200, 10, 20, 30]);
545
546 assert!(
548 !array
549 .execute_scalar(0, &mut LEGACY_SESSION.create_execution_ctx())
550 .unwrap()
551 .is_null()
552 ); assert!(
554 !array
555 .execute_scalar(1, &mut LEGACY_SESSION.create_execution_ctx())
556 .unwrap()
557 .is_null()
558 ); assert!(
562 !array
563 .execute_scalar(2, &mut LEGACY_SESSION.create_execution_ctx())
564 .unwrap()
565 .is_null()
566 ); assert!(
568 array
569 .execute_scalar(3, &mut LEGACY_SESSION.create_execution_ctx())
570 .unwrap()
571 .is_null()
572 ); assert!(
574 !array
575 .execute_scalar(4, &mut LEGACY_SESSION.create_execution_ctx())
576 .unwrap()
577 .is_null()
578 ); }
580
581 #[test]
583 #[should_panic(expected = "cannot create an uninit range of length 0")]
584 fn test_zero_length_uninit_range_panics() {
585 let mut builder = PrimitiveBuilder::<i32>::new(Nullability::NonNullable);
586 let _range = builder.uninit_range(0);
587 }
588
589 #[test]
591 #[should_panic(
592 expected = "uninit_range of len 10 exceeds builder with length 0 and capacity 6"
593 )]
594 fn test_uninit_range_exceeds_capacity_panics() {
595 let mut builder = PrimitiveBuilder::<i32>::with_capacity(Nullability::NonNullable, 5);
596 let _range = builder.uninit_range(10);
597 }
598
599 #[test]
603 #[cfg(debug_assertions)]
604 #[should_panic(expected = "tried to copy a slice into a `UninitRange` past its boundary")]
605 fn test_copy_from_slice_out_of_bounds() {
606 let mut builder = PrimitiveBuilder::<i32>::with_capacity(Nullability::NonNullable, 10);
607 let mut range = builder.uninit_range(3);
608
609 range.copy_from_slice(1, &[1, 2, 3]);
611 }
612
613 #[test]
617 fn test_finish_unsafe_contract() {
618 let mut builder = PrimitiveBuilder::<i32>::with_capacity(Nullability::Nullable, 5);
619 let mut range = builder.uninit_range(3);
620
621 let mask = Mask::from_iter([true, true, false]);
623 unsafe {
625 range.append_mask(&mask);
626 }
627
628 range.copy_from_slice(0, &[10, 20, 30]);
630
631 unsafe {
633 range.finish();
634 }
635
636 let array = builder.finish_into_primitive();
637 assert_eq!(array.len(), 3);
638 assert_eq!(array.as_slice::<i32>(), &[10, 20, 30]);
639 }
640
641 #[test]
642 fn test_append_scalar() {
643 use crate::dtype::DType;
644 use crate::scalar::Scalar;
645
646 let mut builder = PrimitiveBuilder::<i32>::with_capacity(Nullability::Nullable, 10);
647
648 let scalar1 = Scalar::primitive(42i32, Nullability::Nullable);
650 builder.append_scalar(&scalar1).unwrap();
651
652 let scalar2 = Scalar::primitive(84i32, Nullability::Nullable);
654 builder.append_scalar(&scalar2).unwrap();
655
656 let null_scalar = Scalar::null(DType::Primitive(
658 crate::dtype::PType::I32,
659 Nullability::Nullable,
660 ));
661 builder.append_scalar(&null_scalar).unwrap();
662
663 let array = builder.finish_into_primitive();
664 assert_eq!(array.len(), 3);
665
666 let values = array.as_slice::<i32>();
668 assert_eq!(values[0], 42);
669 assert_eq!(values[1], 84);
670 let mut ctx = LEGACY_SESSION.create_execution_ctx();
674 assert!(
675 array
676 .validity()
677 .vortex_expect("primitive validity should be derivable")
678 .execute_is_valid(0, &mut ctx)
679 .unwrap()
680 );
681 assert!(
682 array
683 .validity()
684 .vortex_expect("primitive validity should be derivable")
685 .execute_is_valid(1, &mut ctx)
686 .unwrap()
687 );
688 assert!(
689 !array
690 .validity()
691 .vortex_expect("primitive validity should be derivable")
692 .execute_is_valid(2, &mut ctx)
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}