1use std::any::Any;
5use std::ops::Range;
6use std::sync::Arc;
7
8use itertools::Itertools;
9use num_traits::AsPrimitive;
10use vortex_buffer::Alignment;
11use vortex_buffer::Buffer;
12use vortex_buffer::BufferAllocatorRef;
13use vortex_buffer::BufferMut;
14use vortex_buffer::ByteBuffer;
15use vortex_buffer::ByteBufferMut;
16use vortex_error::VortexExpect;
17use vortex_error::VortexResult;
18use vortex_error::vortex_bail;
19use vortex_error::vortex_ensure;
20use vortex_mask::AllOr;
21use vortex_mask::Mask;
22use vortex_utils::aliases::hash_map::Entry;
23use vortex_utils::aliases::hash_map::HashMap;
24
25use crate::ArrayRef;
26use crate::ExecutionCtx;
27use crate::IntoArray;
28use crate::arrays::VarBinViewArray;
29use crate::arrays::varbinview::VarBinViewArrayExt;
30use crate::arrays::varbinview::build_views::BinaryView;
31use crate::arrays::varbinview::build_views::MAX_BUFFER_LEN;
32use crate::arrays::varbinview::build_views::extend_views;
33use crate::arrays::varbinview::compact::BufferUtilization;
34use crate::builders::ArrayBuilder;
35use crate::builders::LazyBitBufferBuilder;
36use crate::canonical::Canonical;
37use crate::dtype::DType;
38use crate::dtype::NativePType;
39use crate::scalar::Scalar;
40
41pub struct VarBinViewBuilder {
43 dtype: DType,
44 views_builder: BufferMut<BinaryView>,
45 nulls: LazyBitBufferBuilder,
46 completed: CompletedBuffers,
47 in_progress: Option<ByteBufferMut>,
48 growth_strategy: BufferGrowthStrategy,
49 compaction_threshold: f64,
50 allocator: BufferAllocatorRef,
51}
52
53impl VarBinViewBuilder {
54 #[deprecated(note = "use `with_capacity_in` with an explicit allocator")]
56 pub fn with_capacity(dtype: DType, capacity: usize) -> Self {
57 Self::with_capacity_in(dtype, capacity, BufferAllocatorRef::statically_allocated())
58 }
59
60 pub fn with_capacity_in(dtype: DType, capacity: usize, allocator: BufferAllocatorRef) -> Self {
62 Self::new_in(
63 dtype,
64 capacity,
65 Default::default(),
66 Default::default(),
67 0.0,
68 allocator,
69 )
70 }
71
72 #[deprecated(note = "use `with_buffer_deduplication_in` with an explicit allocator")]
74 pub fn with_buffer_deduplication(dtype: DType, capacity: usize) -> Self {
75 Self::with_buffer_deduplication_in(
76 dtype,
77 capacity,
78 BufferAllocatorRef::statically_allocated(),
79 )
80 }
81
82 pub fn with_buffer_deduplication_in(
84 dtype: DType,
85 capacity: usize,
86 allocator: BufferAllocatorRef,
87 ) -> Self {
88 Self::new_in(
89 dtype,
90 capacity,
91 CompletedBuffers::Deduplicated(Default::default()),
92 Default::default(),
93 0.0,
94 allocator,
95 )
96 }
97
98 #[deprecated(note = "use `with_compaction_in` with an explicit allocator")]
100 pub fn with_compaction(dtype: DType, capacity: usize, compaction_threshold: f64) -> Self {
101 Self::with_compaction_in(
102 dtype,
103 capacity,
104 compaction_threshold,
105 BufferAllocatorRef::statically_allocated(),
106 )
107 }
108
109 pub fn with_compaction_in(
111 dtype: DType,
112 capacity: usize,
113 compaction_threshold: f64,
114 allocator: BufferAllocatorRef,
115 ) -> Self {
116 Self::new_in(
117 dtype,
118 capacity,
119 Default::default(),
120 Default::default(),
121 compaction_threshold,
122 allocator,
123 )
124 }
125
126 #[deprecated(note = "use `new_in` with an explicit allocator")]
128 pub fn new(
129 dtype: DType,
130 capacity: usize,
131 completed: CompletedBuffers,
132 growth_strategy: BufferGrowthStrategy,
133 compaction_threshold: f64,
134 ) -> Self {
135 Self::new_in(
136 dtype,
137 capacity,
138 completed,
139 growth_strategy,
140 compaction_threshold,
141 BufferAllocatorRef::statically_allocated(),
142 )
143 }
144
145 pub fn new_in(
147 dtype: DType,
148 capacity: usize,
149 completed: CompletedBuffers,
150 growth_strategy: BufferGrowthStrategy,
151 compaction_threshold: f64,
152 allocator: BufferAllocatorRef,
153 ) -> Self {
154 assert!(
155 matches!(dtype, DType::Utf8(_) | DType::Binary(_)),
156 "VarBinViewBuilder DType must be Utf8 or Binary."
157 );
158 Self {
159 views_builder: BufferMut::with_capacity_preferred_aligned_in(
160 capacity,
161 Alignment::of::<BinaryView>(),
162 None,
163 allocator.clone(),
164 ),
165 nulls: LazyBitBufferBuilder::new(capacity, allocator.clone()),
166 completed,
167 in_progress: None,
168 dtype,
169 growth_strategy,
170 compaction_threshold,
171 allocator,
172 }
173 }
174
175 fn append_value_view(&mut self, value: &[u8]) {
176 let length =
177 u32::try_from(value.len()).vortex_expect("cannot have a single string >2^32 in length");
178 if length <= 12 {
179 self.views_builder.push(BinaryView::make_view(value, 0, 0));
180 return;
181 }
182
183 let (buffer_idx, offset) = self.append_value_to_buffer(value);
184 let view = BinaryView::make_view(value, buffer_idx, offset);
185 self.views_builder.push(view);
186 }
187
188 pub fn append_value<S: AsRef<[u8]>>(&mut self, value: S) {
190 self.append_value_view(value.as_ref());
191 self.nulls.append_non_null();
192 }
193
194 pub fn append_n_values<S: AsRef<[u8]>>(&mut self, value: S, n: usize) {
196 if n == 0 {
197 return;
198 }
199 let bytes = value.as_ref();
200 let view = if bytes.len() <= BinaryView::MAX_INLINED_SIZE {
201 BinaryView::make_view(bytes, 0, 0)
202 } else {
203 let (buffer_idx, offset) = self.append_value_to_buffer(bytes);
204 BinaryView::make_view(bytes, buffer_idx, offset)
205 };
206 self.views_builder.push_n(view, n);
207 self.nulls.append_n_non_nulls(n);
208 }
209
210 fn flush_in_progress(&mut self) {
211 let Some(block) = self.in_progress.take() else {
212 return;
213 };
214
215 assert!(block.len() < u32::MAX as usize, "Block too large");
216
217 let initial_len = self.completed.len();
218 self.completed.push(block.freeze());
219 assert_eq!(
220 self.completed.len(),
221 initial_len + 1,
222 "Invalid state, just completed block already exists"
223 );
224 }
225
226 fn init_in_progress(&mut self, min_len: usize) {
227 let next_buffer_size = self.growth_strategy.next_size() as usize;
228 let to_reserve = next_buffer_size.max(min_len);
229 self.in_progress = Some(ByteBufferMut::with_capacity_preferred_aligned_in(
230 to_reserve,
231 Alignment::of::<u8>(),
232 None,
233 self.allocator.clone(),
234 ));
235 }
236
237 fn append_value_to_buffer(&mut self, value: &[u8]) -> (u32, u32) {
239 assert!(
240 value.len() > BinaryView::MAX_INLINED_SIZE,
241 "must inline small strings"
242 );
243
244 if let Some(in_progress) = &mut self.in_progress {
245 let required_cap = in_progress.len() + value.len();
246 if in_progress.capacity() < required_cap {
247 self.flush_in_progress();
248 self.init_in_progress(value.len());
249 }
250 } else {
251 self.init_in_progress(value.len())
252 };
253
254 let in_progress = self
255 .in_progress
256 .as_mut()
257 .vortex_expect("in_progress just set");
258
259 let buffer_idx = self.completed.len();
260 let offset = u32::try_from(in_progress.len()).vortex_expect("too many buffers");
261 in_progress.extend_from_slice(value);
262
263 (buffer_idx, offset)
264 }
265
266 pub fn completed_block_count(&self) -> u32 {
267 self.completed.len()
268 }
269
270 fn compacts_buffers(&self) -> bool {
275 self.compaction_threshold > 0.0
276 }
277
278 pub fn append_views_built_at(
298 &mut self,
299 validity: &Mask,
300 build: impl FnOnce(u32) -> VortexResult<(Vec<ByteBuffer>, Buffer<BinaryView>)>,
301 ) -> VortexResult<()> {
302 self.flush_in_progress();
303
304 let start_index = self.completed.len();
305 let (buffers, views) = build(start_index)?;
306 assert_eq!(
307 views.len(),
308 validity.len(),
309 "Must build one view per validity entry"
310 );
311
312 let expected_completed_len = start_index as usize + buffers.len();
313 self.completed.extend_from_slice_unchecked(&buffers);
314 assert_eq!(
315 self.completed.len() as usize,
316 expected_completed_len,
317 "Some buffers already exist",
318 );
319 self.views_builder.extend_trusted(views.iter().copied());
320 self.nulls.append_validity_mask(validity);
321
322 debug_assert_eq!(self.nulls.len(), self.views_builder.len());
323 Ok(())
324 }
325
326 fn push_buffers(&mut self, buffers: impl IntoIterator<Item = ByteBuffer>) -> Vec<u32> {
342 self.flush_in_progress();
343 buffers
344 .into_iter()
345 .map(|buffer| self.completed.push(buffer))
346 .collect()
347 }
348
349 fn push_buffers_compacted(
355 &mut self,
356 buffers: Vec<ByteBuffer>,
357 utilizations: &[BufferUtilization],
358 ) -> Vec<CompactedSlot> {
359 self.flush_in_progress();
360 buffers
361 .into_iter()
362 .zip(utilizations)
363 .map(|(buffer, utilization)| {
364 match compaction_strategy(utilization, self.compaction_threshold) {
365 CompactionStrategy::KeepFull => CompactedSlot::Kept {
366 index: self.completed.push(buffer),
367 },
368 CompactionStrategy::Slice { start, end } => CompactedSlot::Sliced {
369 index: self
370 .completed
371 .push(buffer.slice(start as usize..end as usize)),
372 shift: start,
373 },
374 CompactionStrategy::Rewrite => CompactedSlot::Rewrite { source: buffer },
375 }
376 })
377 .collect()
378 }
379
380 fn remap_view_compacted(&mut self, view: BinaryView, slots: &[CompactedSlot]) -> BinaryView {
384 if view.is_inlined() {
385 return view;
386 }
387 let view_ref = view.as_view();
388 match &slots[view_ref.buffer_index as usize] {
389 CompactedSlot::Kept { index } => view_ref
390 .with_buffer_and_offset(*index, view_ref.offset)
391 .into(),
392 CompactedSlot::Sliced { index, shift } => view_ref
393 .with_buffer_and_offset(*index, view_ref.offset - shift)
394 .into(),
395 CompactedSlot::Rewrite { source } => {
396 let bytes = &source.as_slice()[view_ref.as_range()];
397 let (buffer_idx, offset) = self.append_value_to_buffer(bytes);
398 BinaryView::make_view(bytes, buffer_idx, offset)
399 }
400 }
401 }
402
403 pub fn append_views_gathered(
422 &mut self,
423 buffers: impl IntoIterator<Item = ByteBuffer>,
424 views: &[BinaryView],
425 validity: &Mask,
426 view_at: impl Fn(usize) -> usize,
427 ) {
428 if self.compacts_buffers() {
429 return self.append_views_gathered_compacted(
430 buffers.into_iter().collect(),
431 views,
432 validity,
433 view_at,
434 );
435 }
436
437 let mapping = self.push_buffers(buffers);
438
439 self.views_builder.reserve(validity.len());
440 match validity {
441 Mask::AllTrue(len) => self
442 .views_builder
443 .extend_trusted((0..*len).map(|row| remap_view(views[view_at(row)], &mapping))),
444 Mask::AllFalse(len) => self.views_builder.push_n(BinaryView::empty_view(), *len),
445 Mask::Values(values) => {
446 let bits = values.bit_buffer();
447 let start = self.views_builder.len();
448 self.views_builder
449 .push_n(BinaryView::empty_view(), bits.len());
450 let gathered = &mut self.views_builder[start..];
451 bits.for_each_set_index(|row| {
452 gathered[row] = remap_view(views[view_at(row)], &mapping);
453 });
454 }
455 }
456
457 self.nulls.append_validity_mask(validity);
458 debug_assert_eq!(self.nulls.len(), self.views_builder.len());
459 }
460
461 fn append_views_gathered_compacted(
467 &mut self,
468 buffers: Vec<ByteBuffer>,
469 views: &[BinaryView],
470 validity: &Mask,
471 view_at: impl Fn(usize) -> usize,
472 ) {
473 let mut referenced = vec![false; views.len()];
474 match validity {
475 Mask::AllTrue(len) => (0..*len).for_each(|row| referenced[view_at(row)] = true),
476 Mask::AllFalse(_) => {}
477 Mask::Values(values) => values
478 .bit_buffer()
479 .for_each_set_index(|row| referenced[view_at(row)] = true),
480 }
481
482 let mut utilizations = unmeasured_utilizations(&buffers);
483 for (view, _) in views.iter().zip(&referenced).filter(|(_, used)| **used) {
484 measure_view(&mut utilizations, view);
485 }
486 let slots = self.push_buffers_compacted(buffers, &utilizations);
487
488 let mut remapped = vec![BinaryView::empty_view(); views.len()];
489 for (idx, view) in views.iter().enumerate() {
490 if referenced[idx] {
491 remapped[idx] = self.remap_view_compacted(*view, &slots);
492 }
493 }
494
495 self.views_builder.reserve(validity.len());
496 match validity {
497 Mask::AllTrue(len) => self
498 .views_builder
499 .extend_trusted((0..*len).map(|row| remapped[view_at(row)])),
500 Mask::AllFalse(len) => self.views_builder.push_n(BinaryView::empty_view(), *len),
501 Mask::Values(values) => {
502 let bits = values.bit_buffer();
503 let start = self.views_builder.len();
504 self.views_builder
505 .push_n(BinaryView::empty_view(), bits.len());
506 let gathered = &mut self.views_builder[start..];
507 bits.for_each_set_index(|row| gathered[row] = remapped[view_at(row)]);
508 }
509 }
510
511 self.nulls.append_validity_mask(validity);
512 debug_assert_eq!(self.nulls.len(), self.views_builder.len());
513 }
514
515 pub fn append_views_scattered(
532 &mut self,
533 buffers: impl IntoIterator<Item = ByteBuffer>,
534 len: usize,
535 fill: BinaryView,
536 patches: impl Iterator<Item = (usize, BinaryView)>,
537 validity: &Mask,
538 ) {
539 assert_eq!(validity.len(), len, "Must have one validity entry per row");
540 let mapping = self.push_buffers(buffers);
541
542 let start = self.views_builder.len();
543 self.views_builder.push_n(remap_view(fill, &mapping), len);
544 let scattered = &mut self.views_builder[start..];
545 for (row, view) in patches {
546 scattered[row] = remap_view(view, &mapping);
547 }
548
549 self.nulls.append_validity_mask(validity);
550 debug_assert_eq!(self.nulls.len(), self.views_builder.len());
551 }
552
553 pub fn append_buffer_with_lengths<P: NativePType + AsPrimitive<usize>>(
571 &mut self,
572 bytes: ByteBuffer,
573 lengths: &[P],
574 validity: &Mask,
575 ) {
576 assert_eq!(
577 lengths.len(),
578 validity.len(),
579 "Must have one length per validity entry"
580 );
581 self.append_buffer_views(&bytes, lengths.len(), validity, |i| lengths[i].as_());
582 }
583
584 pub fn append_buffer_with_offsets<P: NativePType + AsPrimitive<usize>>(
597 &mut self,
598 bytes: ByteBuffer,
599 offsets: &[P],
600 validity: &Mask,
601 ) {
602 assert_eq!(
603 offsets.len(),
604 validity.len() + 1,
605 "Must have one more offset than validity entries"
606 );
607 let first: usize = offsets[0].as_();
608 let last: usize = offsets[offsets.len() - 1].as_();
609 let bytes = bytes.slice(first..last);
610 self.append_buffer_views(&bytes, validity.len(), validity, |i| {
613 AsPrimitive::<usize>::as_(offsets[i + 1])
614 .wrapping_sub(AsPrimitive::<usize>::as_(offsets[i]))
615 });
616 }
617
618 fn append_buffer_views(
621 &mut self,
622 bytes: &ByteBuffer,
623 count: usize,
624 validity: &Mask,
625 len_at: impl Fn(usize) -> usize,
626 ) {
627 self.flush_in_progress();
628
629 if self.compacts_buffers() && bytes.len() <= MAX_BUFFER_LEN {
635 let referenced: usize = match validity.bit_buffer() {
636 AllOr::All => (0..count)
637 .map(&len_at)
638 .filter(|len| *len > BinaryView::MAX_INLINED_SIZE)
639 .sum(),
640 AllOr::None => 0,
641 AllOr::Some(b) => {
642 let mut sum = 0;
643 b.for_each_set_index(|idx| {
644 let len = len_at(idx);
645 if len > BinaryView::MAX_INLINED_SIZE {
646 sum += len;
647 }
648 });
649 sum
650 }
651 };
652 #[expect(clippy::cast_precision_loss)]
653 if (referenced as f64) < self.compaction_threshold * (bytes.len() as f64) {
654 return self
655 .append_buffer_views_rewritten(bytes, count, validity, len_at, referenced);
656 }
657 }
658
659 let start_index = self.completed.len();
660 let segments = extend_views(
661 &mut self.views_builder,
662 start_index,
663 MAX_BUFFER_LEN,
664 bytes,
665 count,
666 len_at,
667 );
668
669 let expected_completed_len = start_index as usize + segments.len();
670 self.completed.extend_from_slice_unchecked(&segments);
671 assert_eq!(
672 self.completed.len() as usize,
673 expected_completed_len,
674 "Some buffers already exist",
675 );
676
677 self.nulls.append_validity_mask(validity);
678 debug_assert_eq!(self.nulls.len(), self.views_builder.len());
679 }
680
681 fn append_buffer_views_rewritten(
685 &mut self,
686 bytes: &ByteBuffer,
687 count: usize,
688 validity: &Mask,
689 len_at: impl Fn(usize) -> usize,
690 referenced: usize,
691 ) {
692 let buf_index = self.completed.len();
693 let mut compact = ByteBufferMut::with_capacity_in(referenced, self.allocator.clone());
694 self.views_builder.reserve(count);
695
696 let data = bytes.as_slice();
697 let mut offset = 0usize;
698 for (i, is_valid) in validity.iter().enumerate() {
699 let len = len_at(i);
700 let value = &data[offset..offset + len];
701 let view = if !is_valid {
702 BinaryView::empty_view()
703 } else if len > BinaryView::MAX_INLINED_SIZE {
704 #[expect(clippy::cast_possible_truncation)]
707 let view = BinaryView::make_view(value, buf_index, compact.len() as u32);
708 compact.extend_from_slice(value);
709 view
710 } else {
711 BinaryView::make_view(value, buf_index, 0)
712 };
713 self.views_builder.push(view);
714 offset += len;
715 }
716 assert_eq!(
717 offset,
718 data.len(),
719 "value lengths must describe the byte heap exactly"
720 );
721
722 if !compact.is_empty() {
723 let pushed_index = self.completed.push(compact.freeze());
724 assert_eq!(pushed_index, buf_index, "Buffer already exists");
725 }
726
727 self.nulls.append_validity_mask(validity);
728 debug_assert_eq!(self.nulls.len(), self.views_builder.len());
729 }
730
731 pub fn finish_into_varbinview(&mut self) -> VarBinViewArray {
733 self.flush_in_progress();
734 let buffers = std::mem::take(&mut self.completed);
735
736 assert_eq!(
737 self.views_builder.len(),
738 self.nulls.len(),
739 "View and validity length must match"
740 );
741
742 let validity = self.nulls.finish_with_nullability(self.dtype.nullability());
743
744 let views = std::mem::replace(
745 &mut self.views_builder,
746 BufferMut::empty_aligned_in(Alignment::of::<BinaryView>(), self.allocator.clone()),
747 )
748 .freeze();
749
750 unsafe {
752 VarBinViewArray::new_unchecked(views, buffers.finish(), self.dtype.clone(), validity)
753 }
754 }
755
756 pub(crate) fn append_varbinview_array(
757 &mut self,
758 array: &VarBinViewArray,
759 ctx: &mut ExecutionCtx,
760 ) -> VortexResult<()> {
761 self.flush_in_progress();
762
763 let mask = array.varbinview_validity().execute_mask(array.len(), ctx)?;
764
765 self.nulls.append_validity_mask(&mask);
766
767 let view_adjustment =
768 self.completed
769 .extend_from_compaction(BuffersWithOffsets::from_array(
770 array,
771 self.compaction_threshold,
772 ctx,
773 ));
774
775 match view_adjustment {
776 ViewAdjustment::Precomputed(adjustment) => self.views_builder.extend_trusted(
777 array
778 .views()
779 .iter()
780 .map(|view| adjustment.adjust_view(view)),
781 ),
782 ViewAdjustment::Rewriting(adjustment) => match mask {
783 Mask::AllTrue(_) => {
784 for (idx, &view) in array.views().iter().enumerate() {
785 let new_view = self.push_view(view, &adjustment, array, idx);
786 self.views_builder.push(new_view);
787 }
788 }
789 Mask::AllFalse(_) => {
790 self.views_builder
791 .push_n(BinaryView::empty_view(), array.len());
792 }
793 Mask::Values(v) => {
794 let views = array.views();
795 let start = self.views_builder.len();
796 self.views_builder
797 .push_n(BinaryView::empty_view(), array.len());
798 v.bit_buffer().for_each_set_index(|idx| {
799 let new_view = self.push_view(views[idx], &adjustment, array, idx);
800 self.views_builder[start + idx] = new_view;
801 });
802 }
803 },
804 }
805
806 Ok(())
807 }
808}
809
810impl ArrayBuilder for VarBinViewBuilder {
811 fn as_any(&self) -> &dyn Any {
812 self
813 }
814
815 fn as_any_mut(&mut self) -> &mut dyn Any {
816 self
817 }
818
819 fn dtype(&self) -> &DType {
820 &self.dtype
821 }
822
823 fn len(&self) -> usize {
824 self.nulls.len()
825 }
826
827 fn append_zeros(&mut self, n: usize) {
828 self.views_builder.push_n(BinaryView::empty_view(), n);
829 self.nulls.append_n_non_nulls(n);
830 }
831
832 unsafe fn append_nulls_unchecked(&mut self, n: usize) {
833 self.views_builder.push_n(BinaryView::empty_view(), n);
834 self.nulls.append_n_nulls(n);
835 }
836
837 fn append_scalar(&mut self, scalar: &Scalar) -> VortexResult<()> {
838 vortex_ensure!(
839 scalar.dtype() == self.dtype(),
840 "VarBinViewBuilder expected scalar with dtype {}, got {}",
841 self.dtype(),
842 scalar.dtype()
843 );
844
845 match self.dtype() {
846 DType::Utf8(_) => match scalar.as_utf8().value() {
847 Some(value) => self.append_value(value),
848 None => self.append_null(),
849 },
850 DType::Binary(_) => match scalar.as_binary().value() {
851 Some(value) => self.append_value(value),
852 None => self.append_null(),
853 },
854 _ => vortex_bail!(
855 "VarBinViewBuilder can only handle Utf8 or Binary scalars, got {:?}",
856 scalar.dtype()
857 ),
858 }
859
860 Ok(())
861 }
862
863 fn reserve_exact(&mut self, additional: usize) {
864 self.views_builder.reserve(additional);
865 self.nulls.reserve_exact(additional);
866 }
867
868 fn finish(&mut self) -> ArrayRef {
869 self.finish_into_varbinview().into_array()
870 }
871
872 fn finish_into_canonical(&mut self, _ctx: &mut ExecutionCtx) -> Canonical {
873 Canonical::VarBinView(self.finish_into_varbinview())
874 }
875}
876
877impl VarBinViewBuilder {
878 #[allow(clippy::inline_always)]
879 #[inline(always)]
880 fn push_view(
881 &mut self,
882 view: BinaryView,
883 adjustment: &RewritingViewAdjustment,
884 array: &VarBinViewArray,
885 idx: usize,
886 ) -> BinaryView {
887 if view.is_inlined() {
888 view
889 } else if let Some(adjusted) = adjustment.adjust_view(&view) {
890 adjusted
891 } else {
892 let bytes = array.bytes_at(idx);
893 let (new_buf_idx, new_offset) = self.append_value_to_buffer(&bytes);
894 BinaryView::make_view(bytes.as_slice(), new_buf_idx, new_offset)
895 }
896 }
897}
898
899#[inline]
903fn remap_view(view: BinaryView, mapping: &[u32]) -> BinaryView {
904 if view.is_inlined() {
905 view
906 } else {
907 let view_ref = view.as_view();
908 view_ref
909 .with_buffer_and_offset(mapping[view_ref.buffer_index as usize], view_ref.offset)
910 .into()
911 }
912}
913
914enum CompactedSlot {
916 Kept { index: u32 },
918 Sliced { index: u32, shift: u32 },
920 Rewrite { source: ByteBuffer },
923}
924
925fn unmeasured_utilizations(buffers: &[ByteBuffer]) -> Vec<BufferUtilization> {
927 buffers
928 .iter()
929 .map(|buffer| {
930 BufferUtilization::zero(u32::try_from(buffer.len()).unwrap_or(u32::MAX))
934 })
935 .collect()
936}
937
938fn measure_view(utilizations: &mut [BufferUtilization], view: &BinaryView) {
940 if !view.is_inlined() {
941 let view_ref = view.as_view();
942 utilizations[view_ref.buffer_index as usize].add(view_ref.offset, view_ref.size);
943 }
944}
945
946pub enum CompletedBuffers {
947 Default(Vec<ByteBuffer>),
948 Deduplicated(DeduplicatedBuffers),
949}
950
951impl Default for CompletedBuffers {
952 fn default() -> Self {
953 Self::Default(Vec::new())
954 }
955}
956
957#[expect(clippy::cast_possible_truncation)]
959impl CompletedBuffers {
960 fn len(&self) -> u32 {
961 match self {
962 Self::Default(buffers) => buffers.len() as u32,
963 Self::Deduplicated(buffers) => buffers.len(),
964 }
965 }
966
967 fn push(&mut self, block: ByteBuffer) -> u32 {
970 match self {
971 Self::Default(buffers) => {
972 assert!(buffers.len() < u32::MAX as usize, "Too many blocks");
973 buffers.push(block);
974 self.len() - 1
975 }
976 Self::Deduplicated(buffers) => buffers.push(block),
977 }
978 }
979
980 fn extend_from_slice_unchecked(&mut self, buffers: &[ByteBuffer]) {
982 for buffer in buffers {
983 self.push(buffer.clone());
984 }
985 }
986
987 fn extend_from_compaction(&mut self, buffers: BuffersWithOffsets) -> ViewAdjustment {
988 match (self, buffers) {
989 (
990 Self::Default(completed_buffers),
991 BuffersWithOffsets::AllKept { buffers, offsets },
992 ) => {
993 let buffer_offset = completed_buffers.len() as u32;
994 completed_buffers.extend_from_slice(&buffers);
995 ViewAdjustment::shift(buffer_offset, offsets)
996 }
997 (
998 Self::Default(completed_buffers),
999 BuffersWithOffsets::SomeCompacted { buffers, offsets },
1000 ) => {
1001 let lookup = buffers
1002 .iter()
1003 .map(|maybe_buffer| {
1004 maybe_buffer.as_ref().map(|buffer| {
1005 completed_buffers.push(buffer.clone());
1006 completed_buffers.len() as u32 - 1
1007 })
1008 })
1009 .collect();
1010 ViewAdjustment::rewriting(lookup, offsets)
1011 }
1012
1013 (
1014 Self::Deduplicated(completed_buffers),
1015 BuffersWithOffsets::AllKept { buffers, offsets },
1016 ) => {
1017 let buffer_lookup = completed_buffers.extend_from_iter(buffers.iter().cloned());
1018 ViewAdjustment::lookup(buffer_lookup, offsets)
1019 }
1020 (
1021 Self::Deduplicated(completed_buffers),
1022 BuffersWithOffsets::SomeCompacted { buffers, offsets },
1023 ) => {
1024 let buffer_lookup = completed_buffers.extend_from_option_slice(&buffers);
1025 ViewAdjustment::rewriting(buffer_lookup, offsets)
1026 }
1027 }
1028 }
1029
1030 fn finish(self) -> Arc<[ByteBuffer]> {
1031 match self {
1032 Self::Default(buffers) => Arc::from(buffers),
1033 Self::Deduplicated(buffers) => buffers.finish(),
1034 }
1035 }
1036}
1037
1038#[derive(Default)]
1039pub struct DeduplicatedBuffers {
1040 buffers: Vec<ByteBuffer>,
1041 buffer_to_idx: HashMap<BufferId, u32>,
1042}
1043
1044impl DeduplicatedBuffers {
1045 #[expect(clippy::cast_possible_truncation)]
1047 fn len(&self) -> u32 {
1048 self.buffers.len() as u32
1049 }
1050
1051 pub(crate) fn push(&mut self, block: ByteBuffer) -> u32 {
1053 assert!(self.buffers.len() < u32::MAX as usize, "Too many blocks");
1054
1055 let initial_len = self.len();
1056 let id = BufferId::from(&block);
1057 match self.buffer_to_idx.entry(id) {
1058 Entry::Occupied(idx) => *idx.get(),
1059 Entry::Vacant(entry) => {
1060 let idx = initial_len;
1061 entry.insert(idx);
1062 self.buffers.push(block);
1063 idx
1064 }
1065 }
1066 }
1067
1068 pub(crate) fn extend_from_option_slice(
1069 &mut self,
1070 buffers: &[Option<ByteBuffer>],
1071 ) -> Vec<Option<u32>> {
1072 buffers
1073 .iter()
1074 .map(|buffer| buffer.as_ref().map(|buf| self.push(buf.clone())))
1075 .collect()
1076 }
1077
1078 pub(crate) fn extend_from_iter(
1079 &mut self,
1080 buffers: impl Iterator<Item = ByteBuffer>,
1081 ) -> Vec<u32> {
1082 buffers.map(|buffer| self.push(buffer)).collect()
1083 }
1084
1085 pub(crate) fn finish(self) -> Arc<[ByteBuffer]> {
1086 Arc::from(self.buffers)
1087 }
1088}
1089
1090#[derive(PartialEq, Eq, Hash)]
1091struct BufferId {
1092 ptr: usize,
1094 len: usize,
1095}
1096
1097impl BufferId {
1098 fn from(buffer: &ByteBuffer) -> Self {
1099 let slice = buffer.as_slice();
1100 Self {
1101 ptr: slice.as_ptr() as usize,
1102 len: slice.len(),
1103 }
1104 }
1105}
1106
1107#[derive(Debug, Clone)]
1108pub enum BufferGrowthStrategy {
1109 Fixed { size: u32 },
1111 Exponential { current_size: u32, max_size: u32 },
1113}
1114
1115impl Default for BufferGrowthStrategy {
1116 fn default() -> Self {
1117 Self::Exponential {
1118 current_size: 4 * 1024, max_size: 2 * 1024 * 1024, }
1121 }
1122}
1123
1124impl BufferGrowthStrategy {
1125 pub fn fixed(size: u32) -> Self {
1126 Self::Fixed { size }
1127 }
1128
1129 pub fn exponential(initial_size: u32, max_size: u32) -> Self {
1130 Self::Exponential {
1131 current_size: initial_size,
1132 max_size,
1133 }
1134 }
1135
1136 pub fn next_size(&mut self) -> u32 {
1138 match self {
1139 Self::Fixed { size } => *size,
1140 Self::Exponential {
1141 current_size,
1142 max_size,
1143 } => {
1144 let result = *current_size;
1145 if *current_size < *max_size {
1146 *current_size = current_size.saturating_mul(2).min(*max_size);
1147 }
1148 result
1149 }
1150 }
1151 }
1152}
1153
1154enum BuffersWithOffsets {
1155 AllKept {
1156 buffers: Arc<[ByteBuffer]>,
1157 offsets: Option<Vec<u32>>,
1158 },
1159 SomeCompacted {
1160 buffers: Vec<Option<ByteBuffer>>,
1161 offsets: Option<Vec<u32>>,
1162 },
1163}
1164
1165impl BuffersWithOffsets {
1166 pub fn from_array(
1167 array: &VarBinViewArray,
1168 compaction_threshold: f64,
1169 ctx: &mut ExecutionCtx,
1170 ) -> Self {
1171 if compaction_threshold == 0.0 {
1172 return Self::AllKept {
1173 buffers: Arc::from(
1174 array
1175 .data_buffers()
1176 .iter()
1177 .cloned()
1178 .map(|b| b.unwrap_host())
1179 .collect_vec(),
1180 ),
1181 offsets: None,
1182 };
1183 }
1184
1185 let buffer_utilizations = array
1186 .buffer_utilizations(ctx)
1187 .vortex_expect("buffer_utilizations in BuffersWithOffsets::from_array");
1188 let mut has_rewrite = false;
1189 let mut has_nonzero_offset = false;
1190 for utilization in buffer_utilizations.iter() {
1191 match compaction_strategy(utilization, compaction_threshold) {
1192 CompactionStrategy::KeepFull => continue,
1193 CompactionStrategy::Slice { .. } => has_nonzero_offset = true,
1194 CompactionStrategy::Rewrite => has_rewrite = true,
1195 }
1196 }
1197
1198 let buffers_with_offsets_iter = buffer_utilizations
1199 .iter()
1200 .zip(array.data_buffers().iter())
1201 .map(|(utilization, buffer)| {
1202 match compaction_strategy(utilization, compaction_threshold) {
1203 CompactionStrategy::KeepFull => (Some(buffer.as_host().clone()), 0),
1204 CompactionStrategy::Slice { start, end } => (
1205 Some(buffer.as_host().slice(start as usize..end as usize)),
1206 start,
1207 ),
1208 CompactionStrategy::Rewrite => (None, 0),
1209 }
1210 });
1211
1212 match (has_rewrite, has_nonzero_offset) {
1213 (false, false) => {
1215 let buffers: Vec<_> = buffers_with_offsets_iter
1216 .map(|(b, _)| b.vortex_expect("already checked for rewrite"))
1217 .collect();
1218 Self::AllKept {
1219 buffers: Arc::from(buffers),
1220 offsets: None,
1221 }
1222 }
1223 (true, false) => {
1225 let buffers: Vec<_> = buffers_with_offsets_iter.map(|(b, _)| b).collect();
1226 Self::SomeCompacted {
1227 buffers,
1228 offsets: None,
1229 }
1230 }
1231 (false, true) => {
1233 let (buffers, offsets): (Vec<_>, _) = buffers_with_offsets_iter
1234 .map(|(buffer, offset)| {
1235 (buffer.vortex_expect("already checked for rewrite"), offset)
1236 })
1237 .collect();
1238 Self::AllKept {
1239 buffers: Arc::from(buffers),
1240 offsets: Some(offsets),
1241 }
1242 }
1243 (true, true) => {
1245 let (buffers, offsets) = buffers_with_offsets_iter.collect();
1246 Self::SomeCompacted {
1247 buffers,
1248 offsets: Some(offsets),
1249 }
1250 }
1251 }
1252 }
1253}
1254
1255#[derive(Debug, Clone, Copy, PartialEq, Eq)]
1256enum CompactionStrategy {
1257 KeepFull,
1258 Slice {
1260 start: u32,
1261 end: u32,
1262 },
1263 Rewrite,
1265}
1266
1267fn compaction_strategy(
1268 buffer_utilization: &BufferUtilization,
1269 threshold: f64,
1270) -> CompactionStrategy {
1271 match buffer_utilization.overall_utilization() {
1272 0.0 => CompactionStrategy::Rewrite,
1274 utilised if utilised >= threshold => CompactionStrategy::KeepFull,
1275 _ if buffer_utilization.range_utilization() >= threshold => {
1276 let Range { start, end } = buffer_utilization.range();
1277 CompactionStrategy::Slice { start, end }
1278 }
1279 _ => CompactionStrategy::Rewrite,
1280 }
1281}
1282
1283enum ViewAdjustment {
1284 Precomputed(PrecomputedViewAdjustment),
1285 Rewriting(RewritingViewAdjustment),
1286}
1287
1288impl ViewAdjustment {
1289 fn shift(buffer_offset: u32, offsets: Option<Vec<u32>>) -> Self {
1290 Self::Precomputed(PrecomputedViewAdjustment::Shift {
1291 buffer_offset,
1292 offsets,
1293 })
1294 }
1295
1296 fn lookup(buffer_lookup: Vec<u32>, offsets: Option<Vec<u32>>) -> Self {
1297 Self::Precomputed(PrecomputedViewAdjustment::Lookup {
1298 buffer_lookup,
1299 offsets,
1300 })
1301 }
1302
1303 fn rewriting(buffer_lookup: Vec<Option<u32>>, offsets: Option<Vec<u32>>) -> Self {
1304 Self::Rewriting(RewritingViewAdjustment {
1305 buffer_lookup,
1306 offsets,
1307 })
1308 }
1309}
1310
1311enum PrecomputedViewAdjustment {
1313 Shift {
1314 buffer_offset: u32,
1315 offsets: Option<Vec<u32>>,
1316 },
1317 Lookup {
1318 buffer_lookup: Vec<u32>,
1319 offsets: Option<Vec<u32>>,
1320 },
1321}
1322
1323impl PrecomputedViewAdjustment {
1324 #[inline]
1325 fn adjust_view(&self, view: &BinaryView) -> BinaryView {
1326 if view.is_inlined() {
1327 return *view;
1328 }
1329 let view_ref = view.as_view();
1330 match self {
1331 Self::Shift {
1332 buffer_offset,
1333 offsets,
1334 } => {
1335 let b_idx = view_ref.buffer_index;
1336 let offset_shift = offsets
1337 .as_ref()
1338 .map(|o| o[b_idx as usize])
1339 .unwrap_or_default();
1340
1341 if view_ref.offset < offset_shift {
1344 return BinaryView::empty_view();
1345 }
1346
1347 view_ref
1348 .with_buffer_and_offset(b_idx + buffer_offset, view_ref.offset - offset_shift)
1349 }
1350 Self::Lookup {
1351 buffer_lookup,
1352 offsets,
1353 } => {
1354 let b_idx = view_ref.buffer_index;
1355 let buffer = buffer_lookup[b_idx as usize];
1356 let offset_shift = offsets
1357 .as_ref()
1358 .map(|o| o[b_idx as usize])
1359 .unwrap_or_default();
1360
1361 if view_ref.offset < offset_shift {
1364 return BinaryView::empty_view();
1365 }
1366
1367 view_ref.with_buffer_and_offset(buffer, view_ref.offset - offset_shift)
1368 }
1369 }
1370 .into()
1371 }
1372}
1373
1374struct RewritingViewAdjustment {
1375 buffer_lookup: Vec<Option<u32>>,
1376 offsets: Option<Vec<u32>>,
1377}
1378
1379impl RewritingViewAdjustment {
1380 #[inline]
1383 fn adjust_view(&self, view: &BinaryView) -> Option<BinaryView> {
1384 if view.is_inlined() {
1385 return Some(*view);
1386 }
1387
1388 let view_ref = view.as_view();
1389 self.buffer_lookup[view_ref.buffer_index as usize].map(|buffer| {
1390 let offset_shift = self
1391 .offsets
1392 .as_ref()
1393 .map(|o| o[view_ref.buffer_index as usize])
1394 .unwrap_or_default();
1395 view_ref
1396 .with_buffer_and_offset(buffer, view_ref.offset - offset_shift)
1397 .into()
1398 })
1399 }
1400}
1401
1402#[cfg(test)]
1403mod tests {
1404 use vortex_buffer::BufferAllocatorRef;
1405 use vortex_buffer::ByteBuffer;
1406 use vortex_error::VortexResult;
1407 use vortex_mask::Mask;
1408
1409 use crate::IntoArray;
1410 use crate::VortexSessionExecute;
1411 use crate::array_session;
1412 use crate::assert_arrays_eq;
1413 use crate::builders::ArrayBuilder;
1414 use crate::builders::VarBinViewBuilder;
1415 use crate::builders::varbinview::VarBinViewArray;
1416 use crate::dtype::DType;
1417 use crate::dtype::Nullability;
1418
1419 const LONG: &str = "a value that is far too long to inline";
1421
1422 #[test]
1425 fn test_append_buffer_with_lengths() {
1426 let mut ctx = array_session().create_execution_ctx();
1427 let mut builder = VarBinViewBuilder::with_capacity_in(
1428 DType::Utf8(Nullability::Nullable),
1429 8,
1430 BufferAllocatorRef::statically_allocated(),
1431 );
1432
1433 builder.append_value(LONG);
1435
1436 let heap = ByteBuffer::copy_from([LONG.as_bytes(), b"", b"tiny"].concat());
1437 let heap_ptr = heap.as_ptr();
1438 let lengths = [u32::try_from(LONG.len()).unwrap(), 0, 4];
1439 builder.append_buffer_with_lengths(heap, &lengths, &Mask::from_iter([true, false, true]));
1440
1441 builder.append_value("tail");
1442
1443 let actual = builder.finish_into_varbinview();
1444 assert_eq!(actual.data_buffers()[1].as_host().as_ptr(), heap_ptr);
1446
1447 let expected = <VarBinViewArray as FromIterator<_>>::from_iter([
1448 Some(LONG),
1449 Some(LONG),
1450 None,
1451 Some("tiny"),
1452 Some("tail"),
1453 ]);
1454 assert_arrays_eq!(actual, expected, &mut ctx);
1455 }
1456
1457 #[test]
1459 fn test_append_buffer_with_offsets() {
1460 let mut ctx = array_session().create_execution_ctx();
1461 let mut builder = VarBinViewBuilder::with_capacity_in(
1462 DType::Utf8(Nullability::Nullable),
1463 8,
1464 BufferAllocatorRef::statically_allocated(),
1465 );
1466
1467 let heap = ByteBuffer::copy_from(format!("..{LONG}tiny!!"));
1468 let long_len = u32::try_from(LONG.len()).unwrap();
1469 let offsets = [2u32, 2 + long_len, 2 + long_len, 2 + long_len + 4];
1470 builder.append_buffer_with_offsets(
1471 heap.clone(),
1472 &offsets,
1473 &Mask::from_iter([true, false, true]),
1474 );
1475
1476 let actual = builder.finish_into_varbinview();
1477 assert_eq!(actual.data_buffers()[0].as_host().as_ptr(), unsafe {
1480 heap.as_ptr().add(2)
1481 });
1482 assert_eq!(
1483 actual.data_buffers()[0].len(),
1484 LONG.len() + 4,
1485 "only the referenced range must be adopted"
1486 );
1487
1488 let expected =
1489 <VarBinViewArray as FromIterator<_>>::from_iter([Some(LONG), None, Some("tiny")]);
1490 assert_arrays_eq!(actual, expected, &mut ctx);
1491 }
1492
1493 #[test]
1496 fn test_append_buffer_with_lengths_compacts_underutilized_heap() {
1497 let mut ctx = array_session().create_execution_ctx();
1498 let mut builder = VarBinViewBuilder::with_compaction_in(
1499 DType::Utf8(Nullability::Nullable),
1500 4,
1501 1.0,
1502 BufferAllocatorRef::statically_allocated(),
1503 );
1504
1505 let heap = ByteBuffer::copy_from([b"short".as_slice(), LONG.as_bytes(), b"tiny"].concat());
1506 let lengths = [5u32, u32::try_from(LONG.len()).unwrap(), 4];
1507 builder.append_buffer_with_lengths(heap, &lengths, &Mask::new_true(3));
1508
1509 let actual = builder.finish_into_varbinview();
1510 assert_eq!(actual.data_buffers().len(), 1);
1511 assert_eq!(
1512 actual.data_buffers()[0].len(),
1513 LONG.len(),
1514 "the compact buffer must hold only the non-inlined values"
1515 );
1516
1517 let expected = <VarBinViewArray as FromIterator<_>>::from_iter([
1518 Some("short"),
1519 Some(LONG),
1520 Some("tiny"),
1521 ]);
1522 assert_arrays_eq!(actual, expected, &mut ctx);
1523 }
1524
1525 #[test]
1528 fn test_append_buffer_with_lengths_compaction_skips_null_bytes() {
1529 let mut ctx = array_session().create_execution_ctx();
1530 let mut builder = VarBinViewBuilder::with_compaction_in(
1531 DType::Utf8(Nullability::Nullable),
1532 2,
1533 1.0,
1534 BufferAllocatorRef::statically_allocated(),
1535 );
1536
1537 let heap = ByteBuffer::copy_from([LONG.as_bytes(), LONG.as_bytes()].concat());
1538 let lengths = [u32::try_from(LONG.len()).unwrap(); 2];
1539 builder.append_buffer_with_lengths(heap, &lengths, &Mask::from_iter([false, true]));
1540
1541 let actual = builder.finish_into_varbinview();
1542 assert_eq!(actual.data_buffers().len(), 1);
1543 assert_eq!(
1544 actual.data_buffers()[0].len(),
1545 LONG.len(),
1546 "the compact buffer must omit bytes belonging to null rows"
1547 );
1548
1549 let expected = <VarBinViewArray as FromIterator<_>>::from_iter([None::<&str>, Some(LONG)]);
1550 assert_arrays_eq!(actual, expected, &mut ctx);
1551 }
1552
1553 #[test]
1556 fn test_push_buffers_deduplicates() {
1557 let mut builder = VarBinViewBuilder::with_buffer_deduplication_in(
1558 DType::Utf8(Nullability::Nullable),
1559 8,
1560 BufferAllocatorRef::statically_allocated(),
1561 );
1562
1563 let first = ByteBuffer::copy_from(LONG);
1564 let second = ByteBuffer::copy_from("another value far too long to inline");
1565
1566 assert_eq!(
1567 builder.push_buffers([first.clone(), second.clone()]),
1568 [0, 1]
1569 );
1570 assert_eq!(builder.push_buffers([second, first]), [1, 0]);
1571 assert_eq!(builder.completed_block_count(), 2);
1572 }
1573
1574 #[test]
1577 fn test_append_views_gathered() {
1578 let mut ctx = array_session().create_execution_ctx();
1579 let dictionary = <VarBinViewArray as FromIterator<_>>::from_iter([
1580 Some("tiny"),
1581 Some(LONG),
1582 Some("small"),
1583 ]);
1584 let buffers = dictionary
1585 .data_buffers()
1586 .iter()
1587 .map(|buffer| buffer.as_host().clone())
1588 .collect::<Vec<_>>();
1589
1590 let mut builder = VarBinViewBuilder::with_capacity_in(
1591 DType::Utf8(Nullability::Nullable),
1592 8,
1593 BufferAllocatorRef::statically_allocated(),
1594 );
1595 builder.append_value(LONG);
1597
1598 let codes: [usize; 4] = [1, 0, usize::MAX, 2];
1599 let views = dictionary.views();
1600 builder.append_views_gathered(
1601 buffers,
1602 views,
1603 &Mask::from_iter([true, true, false, true]),
1604 |row| codes[row],
1606 );
1607
1608 let expected = <VarBinViewArray as FromIterator<_>>::from_iter([
1609 Some(LONG),
1610 Some(LONG),
1611 Some("tiny"),
1612 None,
1613 Some("small"),
1614 ]);
1615 assert_arrays_eq!(builder.finish_into_varbinview(), expected, &mut ctx);
1616 }
1617
1618 #[test]
1621 fn test_append_views_scattered() {
1622 use crate::arrays::varbinview::build_views::BinaryView;
1623
1624 let mut ctx = array_session().create_execution_ctx();
1625 let patch_values =
1626 <VarBinViewArray as FromIterator<_>>::from_iter([Some(LONG), Some("tiny")]);
1627 let mut buffers = patch_values
1628 .data_buffers()
1629 .iter()
1630 .map(|buffer| buffer.as_host().clone())
1631 .collect::<Vec<_>>();
1632
1633 let fill_bytes = ByteBuffer::copy_from("a fill value too long to inline");
1634 buffers.push(fill_bytes.clone());
1635 let fill = BinaryView::make_view(
1636 fill_bytes.as_slice(),
1637 u32::try_from(buffers.len() - 1).unwrap(),
1638 0,
1639 );
1640
1641 let mut builder = VarBinViewBuilder::with_capacity_in(
1642 DType::Utf8(Nullability::Nullable),
1643 8,
1644 BufferAllocatorRef::statically_allocated(),
1645 );
1646 let views = patch_values.views();
1647 builder.append_views_scattered(
1648 buffers,
1649 5,
1650 fill,
1651 [(1usize, views[0]), (3usize, views[1])].into_iter(),
1652 &Mask::from_iter([true, true, false, true, true]),
1653 );
1654
1655 let expected = <VarBinViewArray as FromIterator<_>>::from_iter([
1656 Some("a fill value too long to inline"),
1657 Some(LONG),
1658 None,
1659 Some("tiny"),
1660 Some("a fill value too long to inline"),
1661 ]);
1662 assert_arrays_eq!(builder.finish_into_varbinview(), expected, &mut ctx);
1663 }
1664
1665 #[test]
1669 fn test_append_views_gathered_compacts_unreferenced_values() {
1670 const DEAD: &str = "a dead value nobody gathers, far too long to inline";
1671 const OTHER: &str = "another value far too long to inline";
1672
1673 let mut ctx = array_session().create_execution_ctx();
1674 let dictionary =
1675 <VarBinViewArray as FromIterator<_>>::from_iter([Some(LONG), Some(DEAD), Some(OTHER)]);
1676 assert_eq!(dictionary.data_buffers().len(), 1);
1677 let buffers = dictionary
1678 .data_buffers()
1679 .iter()
1680 .map(|buffer| buffer.as_host().clone())
1681 .collect::<Vec<_>>();
1682
1683 let mut builder = VarBinViewBuilder::with_compaction_in(
1684 DType::Utf8(Nullability::Nullable),
1685 8,
1686 1.0,
1687 BufferAllocatorRef::statically_allocated(),
1688 );
1689 let codes: [usize; 5] = [2, 0, usize::MAX, 2, 0];
1690 builder.append_views_gathered(
1691 buffers,
1692 dictionary.views(),
1693 &Mask::from_iter([true, true, false, true, true]),
1694 |row| codes[row],
1695 );
1696
1697 let actual = builder.finish_into_varbinview();
1698 assert_eq!(actual.data_buffers().len(), 1);
1699 assert_eq!(
1700 actual.data_buffers()[0].len(),
1701 LONG.len() + OTHER.len(),
1702 "the rewritten buffer must hold each referenced value exactly once"
1703 );
1704
1705 let expected = <VarBinViewArray as FromIterator<_>>::from_iter([
1706 Some(OTHER),
1707 Some(LONG),
1708 None,
1709 Some(OTHER),
1710 Some(LONG),
1711 ]);
1712 assert_arrays_eq!(actual, expected, &mut ctx);
1713 }
1714
1715 #[test]
1718 fn test_append_views_gathered_slices_contiguous_range() {
1719 const TAIL: &str = "the referenced tail value, too long to inline";
1720
1721 let mut ctx = array_session().create_execution_ctx();
1722 let dictionary = <VarBinViewArray as FromIterator<_>>::from_iter([Some(LONG), Some(TAIL)]);
1723 let heap = dictionary.data_buffers()[0].as_host().clone();
1724
1725 let mut builder = VarBinViewBuilder::with_compaction_in(
1726 DType::Utf8(Nullability::Nullable),
1727 8,
1728 1.0,
1729 BufferAllocatorRef::statically_allocated(),
1730 );
1731 builder.append_views_gathered(
1732 [heap.clone()],
1733 dictionary.views(),
1734 &Mask::new_true(2),
1735 |_| 1,
1736 );
1737
1738 let actual = builder.finish_into_varbinview();
1739 assert_eq!(actual.data_buffers().len(), 1);
1740 assert_eq!(actual.data_buffers()[0].len(), TAIL.len());
1741 assert_eq!(actual.data_buffers()[0].as_host().as_ptr(), unsafe {
1743 heap.as_ptr().add(LONG.len())
1744 });
1745
1746 let expected = <VarBinViewArray as FromIterator<_>>::from_iter([Some(TAIL), Some(TAIL)]);
1747 assert_arrays_eq!(actual, expected, &mut ctx);
1748 }
1749
1750 #[test]
1752 fn test_append_views_gathered_keeps_utilized_buffer() {
1753 const SHORTER: &str = "a shorter long value";
1754
1755 let mut ctx = array_session().create_execution_ctx();
1756 let dictionary =
1757 <VarBinViewArray as FromIterator<_>>::from_iter([Some(LONG), Some(SHORTER)]);
1758 let heap = dictionary.data_buffers()[0].as_host().clone();
1759
1760 let mut builder = VarBinViewBuilder::with_compaction_in(
1761 DType::Utf8(Nullability::Nullable),
1762 8,
1763 0.5,
1764 BufferAllocatorRef::statically_allocated(),
1765 );
1766 builder.append_views_gathered(
1767 [heap.clone()],
1768 dictionary.views(),
1769 &Mask::new_true(1),
1770 |_| 0,
1771 );
1772
1773 let actual = builder.finish_into_varbinview();
1774 assert_eq!(actual.data_buffers().len(), 1);
1775 assert_eq!(actual.data_buffers()[0].as_host().as_ptr(), heap.as_ptr());
1776 assert_eq!(actual.data_buffers()[0].len(), heap.len());
1777
1778 let expected = <VarBinViewArray as FromIterator<_>>::from_iter([Some(LONG)]);
1779 assert_arrays_eq!(actual, expected, &mut ctx);
1780 }
1781
1782 #[test]
1785 fn test_append_views_gathered_all_null_drops_buffers() {
1786 let mut ctx = array_session().create_execution_ctx();
1787 let dictionary = <VarBinViewArray as FromIterator<_>>::from_iter([Some(LONG)]);
1788 let buffers = dictionary
1789 .data_buffers()
1790 .iter()
1791 .map(|buffer| buffer.as_host().clone())
1792 .collect::<Vec<_>>();
1793
1794 let mut builder = VarBinViewBuilder::with_compaction_in(
1795 DType::Utf8(Nullability::Nullable),
1796 8,
1797 1.0,
1798 BufferAllocatorRef::statically_allocated(),
1799 );
1800 builder.append_views_gathered(buffers, dictionary.views(), &Mask::new_false(3), |_| {
1801 usize::MAX
1802 });
1803
1804 let actual = builder.finish_into_varbinview();
1805 assert!(
1806 actual.data_buffers().is_empty(),
1807 "an all-null gather must not retain any buffers"
1808 );
1809
1810 let expected = <VarBinViewArray as FromIterator<_>>::from_iter([None::<&str>, None, None]);
1811 assert_arrays_eq!(actual, expected, &mut ctx);
1812 }
1813
1814 #[test]
1816 fn test_append_buffer_with_lengths_drops_fully_inlined_heap() {
1817 let mut ctx = array_session().create_execution_ctx();
1818 let mut builder = VarBinViewBuilder::with_compaction_in(
1819 DType::Utf8(Nullability::Nullable),
1820 4,
1821 1.0,
1822 BufferAllocatorRef::statically_allocated(),
1823 );
1824
1825 let heap = ByteBuffer::copy_from(b"shorttinysmall".as_slice());
1826 builder.append_buffer_with_lengths(heap, &[5u32, 4, 5], &Mask::new_true(3));
1827
1828 let actual = builder.finish_into_varbinview();
1829 assert!(
1830 actual.data_buffers().is_empty(),
1831 "a fully-inlined append must not retain any value bytes"
1832 );
1833
1834 let expected = <VarBinViewArray as FromIterator<_>>::from_iter([
1835 Some("short"),
1836 Some("tiny"),
1837 Some("small"),
1838 ]);
1839 assert_arrays_eq!(actual, expected, &mut ctx);
1840 }
1841
1842 #[test]
1843 fn test_utf8_builder() {
1844 let mut ctx = array_session().create_execution_ctx();
1845 let mut builder = VarBinViewBuilder::with_capacity_in(
1846 DType::Utf8(Nullability::Nullable),
1847 10,
1848 BufferAllocatorRef::statically_allocated(),
1849 );
1850
1851 builder.append_value("Hello");
1852 builder.append_null();
1853 builder.append_value("World");
1854
1855 builder.append_nulls(2);
1856
1857 builder.append_zeros(2);
1858 builder.append_value("test");
1859
1860 let actual = builder.finish();
1861 let expected = <VarBinViewArray as FromIterator<_>>::from_iter([
1862 Some("Hello"),
1863 None,
1864 Some("World"),
1865 None,
1866 None,
1867 Some(""),
1868 Some(""),
1869 Some("test"),
1870 ]);
1871 assert_arrays_eq!(actual, expected, &mut ctx);
1872 }
1873
1874 #[test]
1875 fn test_utf8_builder_with_extend() {
1876 let mut ctx = array_session().create_execution_ctx();
1877 let array = {
1878 let mut builder = VarBinViewBuilder::with_capacity_in(
1879 DType::Utf8(Nullability::Nullable),
1880 10,
1881 BufferAllocatorRef::statically_allocated(),
1882 );
1883 builder.append_null();
1884 builder.append_value("Hello2");
1885 builder.finish()
1886 };
1887 let mut builder = VarBinViewBuilder::with_capacity_in(
1888 DType::Utf8(Nullability::Nullable),
1889 10,
1890 BufferAllocatorRef::statically_allocated(),
1891 );
1892
1893 builder.append_value("Hello1");
1894 array.append_to_builder(&mut builder, &mut ctx).unwrap();
1895 builder.append_nulls(2);
1896 builder.append_value("Hello3");
1897
1898 let actual = builder.finish_into_canonical(&mut ctx);
1899 let expected = <VarBinViewArray as FromIterator<_>>::from_iter([
1900 Some("Hello1"),
1901 None,
1902 Some("Hello2"),
1903 None,
1904 None,
1905 Some("Hello3"),
1906 ]);
1907 assert_arrays_eq!(actual.into_array(), expected.into_array(), &mut ctx);
1908 }
1909
1910 #[test]
1911 fn test_buffer_deduplication() -> VortexResult<()> {
1912 let array = {
1913 let mut builder = VarBinViewBuilder::with_capacity_in(
1914 DType::Utf8(Nullability::Nullable),
1915 10,
1916 BufferAllocatorRef::statically_allocated(),
1917 );
1918 builder.append_value("This is a long string that should not be inlined");
1919 builder.append_value("short string");
1920 builder.finish_into_varbinview()
1921 };
1922
1923 assert_eq!(array.data_buffers().len(), 1);
1924 let mut builder = VarBinViewBuilder::with_buffer_deduplication_in(
1925 DType::Utf8(Nullability::Nullable),
1926 10,
1927 BufferAllocatorRef::statically_allocated(),
1928 );
1929
1930 let mut ctx = array_session().create_execution_ctx();
1931
1932 array.append_to_builder(&mut builder, &mut ctx)?;
1933 assert_eq!(builder.completed_block_count(), 1);
1934
1935 array
1936 .slice(1..2)?
1937 .append_to_builder(&mut builder, &mut ctx)?;
1938 array
1939 .slice(0..1)?
1940 .append_to_builder(&mut builder, &mut ctx)?;
1941 assert_eq!(builder.completed_block_count(), 1);
1942
1943 let array2 = {
1944 let mut builder = VarBinViewBuilder::with_capacity_in(
1945 DType::Utf8(Nullability::Nullable),
1946 10,
1947 BufferAllocatorRef::statically_allocated(),
1948 );
1949 builder.append_value("This is a long string that should not be inlined");
1950 builder.finish_into_varbinview()
1951 };
1952
1953 array2.append_to_builder(&mut builder, &mut ctx)?;
1954 assert_eq!(builder.completed_block_count(), 2);
1955
1956 array
1957 .slice(0..1)?
1958 .append_to_builder(&mut builder, &mut ctx)?;
1959 array2
1960 .slice(0..1)?
1961 .append_to_builder(&mut builder, &mut ctx)?;
1962 assert_eq!(builder.completed_block_count(), 2);
1963 Ok(())
1964 }
1965
1966 #[test]
1967 fn test_append_scalar() {
1968 let mut ctx = array_session().create_execution_ctx();
1969 use crate::scalar::Scalar;
1970
1971 let mut utf8_builder = VarBinViewBuilder::with_capacity_in(
1973 DType::Utf8(Nullability::Nullable),
1974 10,
1975 BufferAllocatorRef::statically_allocated(),
1976 );
1977
1978 let utf8_scalar1 = Scalar::utf8("hello", Nullability::Nullable);
1980 utf8_builder.append_scalar(&utf8_scalar1).unwrap();
1981
1982 let utf8_scalar2 = Scalar::utf8("world", Nullability::Nullable);
1984 utf8_builder.append_scalar(&utf8_scalar2).unwrap();
1985
1986 let null_scalar = Scalar::null(DType::Utf8(Nullability::Nullable));
1988 utf8_builder.append_scalar(&null_scalar).unwrap();
1989
1990 let array = utf8_builder.finish();
1991 let expected =
1992 <VarBinViewArray as FromIterator<_>>::from_iter([Some("hello"), Some("world"), None]);
1993 assert_arrays_eq!(&array, &expected, &mut ctx);
1994
1995 let mut binary_builder = VarBinViewBuilder::with_capacity_in(
1997 DType::Binary(Nullability::Nullable),
1998 10,
1999 BufferAllocatorRef::statically_allocated(),
2000 );
2001
2002 let binary_scalar = Scalar::binary(vec![1u8, 2, 3], Nullability::Nullable);
2003 binary_builder.append_scalar(&binary_scalar).unwrap();
2004
2005 let binary_null = Scalar::null(DType::Binary(Nullability::Nullable));
2006 binary_builder.append_scalar(&binary_null).unwrap();
2007
2008 let binary_array = binary_builder.finish();
2009 let expected =
2010 <VarBinViewArray as FromIterator<_>>::from_iter([Some(vec![1u8, 2, 3]), None]);
2011 assert_arrays_eq!(&binary_array, &expected, &mut ctx);
2012
2013 let mut builder = VarBinViewBuilder::with_capacity_in(
2015 DType::Utf8(Nullability::NonNullable),
2016 10,
2017 BufferAllocatorRef::statically_allocated(),
2018 );
2019 let wrong_scalar = Scalar::from(42i32);
2020 assert!(builder.append_scalar(&wrong_scalar).is_err());
2021 }
2022
2023 #[test]
2024 fn test_buffer_growth_strategies() {
2025 use super::BufferGrowthStrategy;
2026
2027 let mut strategy = BufferGrowthStrategy::fixed(1024);
2029
2030 assert_eq!(strategy.next_size(), 1024);
2032 assert_eq!(strategy.next_size(), 1024);
2033 assert_eq!(strategy.next_size(), 1024);
2034
2035 let mut strategy = BufferGrowthStrategy::exponential(1024, 8192);
2037
2038 assert_eq!(strategy.next_size(), 1024); assert_eq!(strategy.next_size(), 2048); assert_eq!(strategy.next_size(), 4096); assert_eq!(strategy.next_size(), 8192); assert_eq!(strategy.next_size(), 8192); }
2045
2046 #[test]
2047 fn test_large_value_allocation() {
2048 use super::BufferGrowthStrategy;
2049 use super::VarBinViewBuilder;
2050
2051 let mut builder = VarBinViewBuilder::new_in(
2052 DType::Binary(Nullability::Nullable),
2053 10,
2054 Default::default(),
2055 BufferGrowthStrategy::exponential(1024, 4096),
2056 0.0,
2057 BufferAllocatorRef::statically_allocated(),
2058 );
2059
2060 let large_value = vec![0u8; 8192];
2062
2063 builder.append_value(&large_value);
2065
2066 let array = builder.finish_into_varbinview();
2067 assert_eq!(array.len(), 1);
2068
2069 let retrieved = array
2071 .execute_scalar(0, &mut array_session().create_execution_ctx())
2072 .unwrap()
2073 .as_binary()
2074 .value()
2075 .cloned()
2076 .unwrap();
2077 assert_eq!(retrieved.len(), 8192);
2078 assert_eq!(retrieved.as_slice(), &large_value);
2079 }
2080}