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;
22use crate::canonical::ToCanonical;
23use crate::dtype::DType;
24use crate::dtype::NativePType;
25use crate::dtype::Nullability;
26use crate::scalar::Scalar;
27
28pub struct PrimitiveBuilder<T> {
30 dtype: DType,
31 values: BufferMut<T>,
32 nulls: LazyBitBufferBuilder,
33}
34
35impl<T: NativePType> PrimitiveBuilder<T> {
36 pub fn new(nullability: Nullability) -> Self {
38 Self::with_capacity(nullability, DEFAULT_BUILDER_CAPACITY)
39 }
40
41 pub fn with_capacity(nullability: Nullability, capacity: usize) -> Self {
43 Self {
44 values: BufferMut::with_capacity(capacity),
45 nulls: LazyBitBufferBuilder::new(capacity),
46 dtype: DType::Primitive(T::PTYPE, nullability),
47 }
48 }
49
50 pub fn append_value(&mut self, value: T) {
52 self.values.push(value);
53 self.nulls.append_non_null();
54 }
55
56 pub fn append_n_values(&mut self, value: T, n: usize) {
58 self.values.push_n(value, n);
59 self.nulls.append_n_non_nulls(n);
60 }
61
62 pub fn values(&self) -> &[T] {
64 self.values.as_ref()
65 }
66
67 pub fn values_mut(&mut self) -> &mut [T] {
69 self.values.as_mut()
70 }
71
72 pub fn uninit_range(&mut self, len: usize) -> UninitRange<'_, T> {
106 assert_ne!(0, len, "cannot create an uninit range of length 0");
107
108 let current_len = self.values.len();
109 assert!(
110 current_len + len <= self.values.capacity(),
111 "uninit_range of len {len} exceeds builder with length {} and capacity {}",
112 current_len,
113 self.values.capacity()
114 );
115
116 UninitRange { len, builder: self }
117 }
118
119 pub fn finish_into_primitive(&mut self) -> PrimitiveArray {
121 let validity = self
122 .nulls
123 .finish_with_nullability(self.dtype().nullability());
124
125 PrimitiveArray::new(std::mem::take(&mut self.values).freeze(), validity)
126 }
127
128 pub fn extend_with_iterator(&mut self, iter: impl IntoIterator<Item = T>, mask: Mask) {
130 self.values.extend(iter);
131 self.nulls.append_validity_mask(mask);
132 }
133}
134
135impl<T: NativePType> ArrayBuilder for PrimitiveBuilder<T> {
136 fn as_any(&self) -> &dyn Any {
137 self
138 }
139
140 fn as_any_mut(&mut self) -> &mut dyn Any {
141 self
142 }
143
144 fn dtype(&self) -> &DType {
145 &self.dtype
146 }
147
148 fn len(&self) -> usize {
149 self.values.len()
150 }
151
152 fn append_zeros(&mut self, n: usize) {
153 self.values.push_n(T::default(), n);
154 self.nulls.append_n_non_nulls(n);
155 }
156
157 unsafe fn append_nulls_unchecked(&mut self, n: usize) {
158 self.values.push_n(T::default(), n);
159 self.nulls.append_n_nulls(n);
160 }
161
162 fn append_scalar(&mut self, scalar: &Scalar) -> VortexResult<()> {
163 vortex_ensure!(
164 scalar.dtype() == self.dtype(),
165 "PrimitiveBuilder expected scalar with dtype {}, got {}",
166 self.dtype(),
167 scalar.dtype()
168 );
169
170 if let Some(pv) = scalar.as_primitive().pvalue() {
171 self.append_value(pv.cast::<T>()?)
172 } else {
173 self.append_null()
174 }
175
176 Ok(())
177 }
178
179 unsafe fn extend_from_array_unchecked(&mut self, array: &ArrayRef) {
180 let array = array.to_primitive();
181
182 debug_assert_eq!(
184 array.ptype(),
185 T::PTYPE,
186 "Cannot extend from array with different ptype"
187 );
188
189 self.values.extend_from_slice(array.as_slice::<T>());
190 self.nulls.append_validity_mask(
191 array
192 .as_ref()
193 .validity()
194 .vortex_expect("validity_mask")
195 .to_mask(
196 array.as_ref().len(),
197 &mut LEGACY_SESSION.create_execution_ctx(),
198 )
199 .vortex_expect("Failed to compute validity mask"),
200 );
201 }
202
203 fn reserve_exact(&mut self, additional: usize) {
204 self.values.reserve(additional);
205 self.nulls.reserve_exact(additional);
206 }
207
208 unsafe fn set_validity_unchecked(&mut self, validity: Mask) {
209 self.nulls = LazyBitBufferBuilder::new(validity.len());
210 self.nulls.append_validity_mask(validity);
211 }
212
213 fn finish(&mut self) -> ArrayRef {
214 self.finish_into_primitive().into_array()
215 }
216
217 fn finish_into_canonical(&mut self) -> Canonical {
218 Canonical::Primitive(self.finish_into_primitive())
219 }
220}
221
222pub struct UninitRange<'a, T> {
224 len: usize,
228
229 builder: &'a mut PrimitiveBuilder<T>,
234}
235
236impl<T> UninitRange<'_, T> {
237 #[inline]
239 pub fn len(&self) -> usize {
240 self.len
241 }
242
243 #[inline]
245 pub fn is_empty(&self) -> bool {
246 self.len == 0
247 }
248
249 #[inline]
255 pub fn set_value(&mut self, index: usize, value: T) {
256 assert!(index < self.len, "index out of bounds");
257 let spare = self.builder.values.spare_capacity_mut();
258 spare[index] = MaybeUninit::new(value);
259 }
260
261 pub unsafe fn append_mask(&mut self, mask: Mask) {
273 assert_eq!(
274 mask.len(),
275 self.len,
276 "Tried to append a mask to an `UninitRange` that was beyond the allowed range"
277 );
278
279 self.builder.nulls.append_validity_mask(mask);
284 }
285
286 pub fn set_validity_bit(&mut self, index: usize, v: bool) {
293 assert!(index < self.len, "set_bit index out of bounds");
294 let absolute_index = self.builder.values.len() + index;
297 self.builder.nulls.set_bit(absolute_index, v);
298 }
299
300 pub fn copy_from_slice(&mut self, local_offset: usize, src: &[T])
305 where
306 T: Copy,
307 {
308 debug_assert!(
309 local_offset + src.len() <= self.len,
310 "tried to copy a slice into a `UninitRange` past its boundary"
311 );
312
313 let uninit_src: &[MaybeUninit<T>] = unsafe { std::mem::transmute(src) };
315
316 let dst =
319 &mut self.builder.values.spare_capacity_mut()[local_offset..local_offset + src.len()];
320 dst.copy_from_slice(uninit_src);
321 }
322
323 pub unsafe fn slice_uninit_mut(&mut self, offset: usize, len: usize) -> &mut [MaybeUninit<T>] {
337 assert!(
338 offset + len <= self.len,
339 "slice_uninit_mut: offset {} + len {} exceeds range length {}",
340 offset,
341 len,
342 self.len
343 );
344 &mut self.builder.values.spare_capacity_mut()[offset..offset + len]
345 }
346
347 pub unsafe fn finish(self) {
361 let new_len = self.builder.values.len() + self.len;
363 unsafe { self.builder.values.set_len(new_len) };
364 }
365}
366
367#[cfg(test)]
368mod tests {
369 use vortex_error::VortexExpect;
370
371 use super::*;
372 use crate::assert_arrays_eq;
373
374 #[test]
379 fn test_multiple_uninit_ranges_correct_offsets() {
380 let mut builder = PrimitiveBuilder::<i32>::with_capacity(Nullability::NonNullable, 10);
381
382 let mut range1 = builder.uninit_range(3);
384 range1.copy_from_slice(0, &[1, 2, 3]);
385
386 unsafe {
388 range1.finish();
389 }
390
391 assert_eq!(builder.values(), &[1, 2, 3]);
393
394 let mut range2 = builder.uninit_range(2);
396
397 range2.copy_from_slice(0, &[4, 5]);
399
400 unsafe {
402 range2.finish();
403 }
404
405 assert_eq!(builder.values(), &[1, 2, 3, 4, 5]);
407
408 let array = builder.finish_into_primitive();
409 assert_arrays_eq!(array, PrimitiveArray::from_iter([1i32, 2, 3, 4, 5]));
410 }
411
412 #[test]
418 fn test_append_mask_on_uninit_range() {
419 let mut builder = PrimitiveBuilder::<i32>::with_capacity(Nullability::Nullable, 5);
420 let mut range = builder.uninit_range(3);
421
422 let mask = Mask::from_iter([true, false, true]);
424
425 unsafe {
427 range.append_mask(mask);
428 }
429
430 range.copy_from_slice(0, &[10, 20, 30]);
432
433 unsafe {
435 range.finish();
436 }
437
438 let array = builder.finish_into_primitive();
439 assert_eq!(array.len(), 3);
440 assert!(
442 !array
443 .execute_scalar(0, &mut LEGACY_SESSION.create_execution_ctx())
444 .unwrap()
445 .is_null()
446 );
447 assert!(
448 array
449 .execute_scalar(1, &mut LEGACY_SESSION.create_execution_ctx())
450 .unwrap()
451 .is_null()
452 );
453 assert!(
454 !array
455 .execute_scalar(2, &mut LEGACY_SESSION.create_execution_ctx())
456 .unwrap()
457 .is_null()
458 );
459 }
460
461 #[test]
465 #[should_panic(
466 expected = "Tried to append a mask to an `UninitRange` that was beyond the allowed range"
467 )]
468 fn test_append_mask_wrong_length_panics() {
469 let mut builder = PrimitiveBuilder::<i32>::with_capacity(Nullability::Nullable, 10);
470 let mut range = builder.uninit_range(5);
471
472 let wrong_mask = Mask::from_iter([true, false, true]);
474
475 unsafe {
477 range.append_mask(wrong_mask);
478 }
479 }
480
481 #[test]
485 fn test_copy_from_slice_with_offsets() {
486 let mut builder = PrimitiveBuilder::<i32>::with_capacity(Nullability::NonNullable, 10);
487 let mut range = builder.uninit_range(6);
488
489 range.copy_from_slice(0, &[1, 2]);
491 range.copy_from_slice(2, &[3, 4]);
492 range.copy_from_slice(4, &[5, 6]);
493
494 unsafe {
496 range.finish();
497 }
498
499 let array = builder.finish_into_primitive();
500 assert_arrays_eq!(array, PrimitiveArray::from_iter([1i32, 2, 3, 4, 5, 6]));
501 }
502
503 #[test]
509 fn test_set_bit_relative_indexing() {
510 let mut builder = PrimitiveBuilder::<i32>::with_capacity(Nullability::Nullable, 10);
511
512 builder.append_value(100);
514 builder.append_value(200);
515
516 let mut range = builder.uninit_range(3);
518
519 let initial_mask = Mask::from_iter([false, false, false]);
521 unsafe {
523 range.append_mask(initial_mask);
524 }
525
526 range.set_validity_bit(0, true); range.set_validity_bit(2, true); range.copy_from_slice(0, &[10, 20, 30]);
533
534 unsafe {
536 range.finish();
537 }
538
539 let array = builder.finish_into_primitive();
540
541 assert_eq!(array.len(), 5);
543 assert_eq!(array.as_slice::<i32>(), &[100, 200, 10, 20, 30]);
544
545 assert!(
547 !array
548 .execute_scalar(0, &mut LEGACY_SESSION.create_execution_ctx())
549 .unwrap()
550 .is_null()
551 ); assert!(
553 !array
554 .execute_scalar(1, &mut LEGACY_SESSION.create_execution_ctx())
555 .unwrap()
556 .is_null()
557 ); assert!(
561 !array
562 .execute_scalar(2, &mut LEGACY_SESSION.create_execution_ctx())
563 .unwrap()
564 .is_null()
565 ); assert!(
567 array
568 .execute_scalar(3, &mut LEGACY_SESSION.create_execution_ctx())
569 .unwrap()
570 .is_null()
571 ); assert!(
573 !array
574 .execute_scalar(4, &mut LEGACY_SESSION.create_execution_ctx())
575 .unwrap()
576 .is_null()
577 ); }
579
580 #[test]
582 #[should_panic(expected = "cannot create an uninit range of length 0")]
583 fn test_zero_length_uninit_range_panics() {
584 let mut builder = PrimitiveBuilder::<i32>::new(Nullability::NonNullable);
585 let _range = builder.uninit_range(0);
586 }
587
588 #[test]
590 #[should_panic(
591 expected = "uninit_range of len 10 exceeds builder with length 0 and capacity 6"
592 )]
593 fn test_uninit_range_exceeds_capacity_panics() {
594 let mut builder = PrimitiveBuilder::<i32>::with_capacity(Nullability::NonNullable, 5);
595 let _range = builder.uninit_range(10);
596 }
597
598 #[test]
602 #[cfg(debug_assertions)]
603 #[should_panic(expected = "tried to copy a slice into a `UninitRange` past its boundary")]
604 fn test_copy_from_slice_out_of_bounds() {
605 let mut builder = PrimitiveBuilder::<i32>::with_capacity(Nullability::NonNullable, 10);
606 let mut range = builder.uninit_range(3);
607
608 range.copy_from_slice(1, &[1, 2, 3]);
610 }
611
612 #[test]
616 fn test_finish_unsafe_contract() {
617 let mut builder = PrimitiveBuilder::<i32>::with_capacity(Nullability::Nullable, 5);
618 let mut range = builder.uninit_range(3);
619
620 let mask = Mask::from_iter([true, true, false]);
622 unsafe {
624 range.append_mask(mask);
625 }
626
627 range.copy_from_slice(0, &[10, 20, 30]);
629
630 unsafe {
632 range.finish();
633 }
634
635 let array = builder.finish_into_primitive();
636 assert_eq!(array.len(), 3);
637 assert_eq!(array.as_slice::<i32>(), &[10, 20, 30]);
638 }
639
640 #[test]
641 fn test_append_scalar() {
642 use crate::dtype::DType;
643 use crate::scalar::Scalar;
644
645 let mut builder = PrimitiveBuilder::<i32>::with_capacity(Nullability::Nullable, 10);
646
647 let scalar1 = Scalar::primitive(42i32, Nullability::Nullable);
649 builder.append_scalar(&scalar1).unwrap();
650
651 let scalar2 = Scalar::primitive(84i32, Nullability::Nullable);
653 builder.append_scalar(&scalar2).unwrap();
654
655 let null_scalar = Scalar::null(DType::Primitive(
657 crate::dtype::PType::I32,
658 Nullability::Nullable,
659 ));
660 builder.append_scalar(&null_scalar).unwrap();
661
662 let array = builder.finish_into_primitive();
663 assert_eq!(array.len(), 3);
664
665 let values = array.as_slice::<i32>();
667 assert_eq!(values[0], 42);
668 assert_eq!(values[1], 84);
669 assert!(
673 array
674 .validity()
675 .vortex_expect("primitive validity should be derivable")
676 .is_valid(0)
677 .unwrap()
678 );
679 assert!(
680 array
681 .validity()
682 .vortex_expect("primitive validity should be derivable")
683 .is_valid(1)
684 .unwrap()
685 );
686 assert!(
687 !array
688 .validity()
689 .vortex_expect("primitive validity should be derivable")
690 .is_valid(2)
691 .unwrap()
692 );
693
694 let mut builder = PrimitiveBuilder::<i32>::with_capacity(Nullability::NonNullable, 10);
696 let wrong_scalar = Scalar::from(true);
697 assert!(builder.append_scalar(&wrong_scalar).is_err());
698 }
699}