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_eq;
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_eq!(
839 scalar.dtype(),
840 self.dtype(),
841 "VarBinViewBuilder received a scalar with the wrong dtype"
842 );
843
844 match self.dtype() {
845 DType::Utf8(_) => match scalar.as_utf8().value() {
846 Some(value) => self.append_value(value),
847 None => self.append_null(),
848 },
849 DType::Binary(_) => match scalar.as_binary().value() {
850 Some(value) => self.append_value(value),
851 None => self.append_null(),
852 },
853 _ => vortex_bail!(
854 "VarBinViewBuilder can only handle Utf8 or Binary scalars, got {:?}",
855 scalar.dtype()
856 ),
857 }
858
859 Ok(())
860 }
861
862 fn reserve_exact(&mut self, additional: usize) {
863 self.views_builder.reserve(additional);
864 self.nulls.reserve_exact(additional);
865 }
866
867 fn finish(&mut self) -> ArrayRef {
868 self.finish_into_varbinview().into_array()
869 }
870
871 fn finish_into_canonical(&mut self, _ctx: &mut ExecutionCtx) -> Canonical {
872 Canonical::VarBinView(self.finish_into_varbinview())
873 }
874}
875
876impl VarBinViewBuilder {
877 #[allow(clippy::inline_always)]
878 #[inline(always)]
879 fn push_view(
880 &mut self,
881 view: BinaryView,
882 adjustment: &RewritingViewAdjustment,
883 array: &VarBinViewArray,
884 idx: usize,
885 ) -> BinaryView {
886 if view.is_inlined() {
887 view
888 } else if let Some(adjusted) = adjustment.adjust_view(&view) {
889 adjusted
890 } else {
891 let bytes = array.bytes_at(idx);
892 let (new_buf_idx, new_offset) = self.append_value_to_buffer(&bytes);
893 BinaryView::make_view(bytes.as_slice(), new_buf_idx, new_offset)
894 }
895 }
896}
897
898#[inline]
902fn remap_view(view: BinaryView, mapping: &[u32]) -> BinaryView {
903 if view.is_inlined() {
904 view
905 } else {
906 let view_ref = view.as_view();
907 view_ref
908 .with_buffer_and_offset(mapping[view_ref.buffer_index as usize], view_ref.offset)
909 .into()
910 }
911}
912
913enum CompactedSlot {
915 Kept { index: u32 },
917 Sliced { index: u32, shift: u32 },
919 Rewrite { source: ByteBuffer },
922}
923
924fn unmeasured_utilizations(buffers: &[ByteBuffer]) -> Vec<BufferUtilization> {
926 buffers
927 .iter()
928 .map(|buffer| {
929 BufferUtilization::zero(u32::try_from(buffer.len()).unwrap_or(u32::MAX))
933 })
934 .collect()
935}
936
937fn measure_view(utilizations: &mut [BufferUtilization], view: &BinaryView) {
939 if !view.is_inlined() {
940 let view_ref = view.as_view();
941 utilizations[view_ref.buffer_index as usize].add(view_ref.offset, view_ref.size);
942 }
943}
944
945pub enum CompletedBuffers {
946 Default(Vec<ByteBuffer>),
947 Deduplicated(DeduplicatedBuffers),
948}
949
950impl Default for CompletedBuffers {
951 fn default() -> Self {
952 Self::Default(Vec::new())
953 }
954}
955
956#[expect(clippy::cast_possible_truncation)]
958impl CompletedBuffers {
959 fn len(&self) -> u32 {
960 match self {
961 Self::Default(buffers) => buffers.len() as u32,
962 Self::Deduplicated(buffers) => buffers.len(),
963 }
964 }
965
966 fn push(&mut self, block: ByteBuffer) -> u32 {
969 match self {
970 Self::Default(buffers) => {
971 assert!(buffers.len() < u32::MAX as usize, "Too many blocks");
972 buffers.push(block);
973 self.len() - 1
974 }
975 Self::Deduplicated(buffers) => buffers.push(block),
976 }
977 }
978
979 fn extend_from_slice_unchecked(&mut self, buffers: &[ByteBuffer]) {
981 for buffer in buffers {
982 self.push(buffer.clone());
983 }
984 }
985
986 fn extend_from_compaction(&mut self, buffers: BuffersWithOffsets) -> ViewAdjustment {
987 match (self, buffers) {
988 (
989 Self::Default(completed_buffers),
990 BuffersWithOffsets::AllKept { buffers, offsets },
991 ) => {
992 let buffer_offset = completed_buffers.len() as u32;
993 completed_buffers.extend_from_slice(&buffers);
994 ViewAdjustment::shift(buffer_offset, offsets)
995 }
996 (
997 Self::Default(completed_buffers),
998 BuffersWithOffsets::SomeCompacted { buffers, offsets },
999 ) => {
1000 let lookup = buffers
1001 .iter()
1002 .map(|maybe_buffer| {
1003 maybe_buffer.as_ref().map(|buffer| {
1004 completed_buffers.push(buffer.clone());
1005 completed_buffers.len() as u32 - 1
1006 })
1007 })
1008 .collect();
1009 ViewAdjustment::rewriting(lookup, offsets)
1010 }
1011
1012 (
1013 Self::Deduplicated(completed_buffers),
1014 BuffersWithOffsets::AllKept { buffers, offsets },
1015 ) => {
1016 let buffer_lookup = completed_buffers.extend_from_iter(buffers.iter().cloned());
1017 ViewAdjustment::lookup(buffer_lookup, offsets)
1018 }
1019 (
1020 Self::Deduplicated(completed_buffers),
1021 BuffersWithOffsets::SomeCompacted { buffers, offsets },
1022 ) => {
1023 let buffer_lookup = completed_buffers.extend_from_option_slice(&buffers);
1024 ViewAdjustment::rewriting(buffer_lookup, offsets)
1025 }
1026 }
1027 }
1028
1029 fn finish(self) -> Arc<[ByteBuffer]> {
1030 match self {
1031 Self::Default(buffers) => Arc::from(buffers),
1032 Self::Deduplicated(buffers) => buffers.finish(),
1033 }
1034 }
1035}
1036
1037#[derive(Default)]
1038pub struct DeduplicatedBuffers {
1039 buffers: Vec<ByteBuffer>,
1040 buffer_to_idx: HashMap<BufferId, u32>,
1041}
1042
1043impl DeduplicatedBuffers {
1044 #[expect(clippy::cast_possible_truncation)]
1046 fn len(&self) -> u32 {
1047 self.buffers.len() as u32
1048 }
1049
1050 pub(crate) fn push(&mut self, block: ByteBuffer) -> u32 {
1052 assert!(self.buffers.len() < u32::MAX as usize, "Too many blocks");
1053
1054 let initial_len = self.len();
1055 let id = BufferId::from(&block);
1056 match self.buffer_to_idx.entry(id) {
1057 Entry::Occupied(idx) => *idx.get(),
1058 Entry::Vacant(entry) => {
1059 let idx = initial_len;
1060 entry.insert(idx);
1061 self.buffers.push(block);
1062 idx
1063 }
1064 }
1065 }
1066
1067 pub(crate) fn extend_from_option_slice(
1068 &mut self,
1069 buffers: &[Option<ByteBuffer>],
1070 ) -> Vec<Option<u32>> {
1071 buffers
1072 .iter()
1073 .map(|buffer| buffer.as_ref().map(|buf| self.push(buf.clone())))
1074 .collect()
1075 }
1076
1077 pub(crate) fn extend_from_iter(
1078 &mut self,
1079 buffers: impl Iterator<Item = ByteBuffer>,
1080 ) -> Vec<u32> {
1081 buffers.map(|buffer| self.push(buffer)).collect()
1082 }
1083
1084 pub(crate) fn finish(self) -> Arc<[ByteBuffer]> {
1085 Arc::from(self.buffers)
1086 }
1087}
1088
1089#[derive(PartialEq, Eq, Hash)]
1090struct BufferId {
1091 ptr: usize,
1093 len: usize,
1094}
1095
1096impl BufferId {
1097 fn from(buffer: &ByteBuffer) -> Self {
1098 let slice = buffer.as_slice();
1099 Self {
1100 ptr: slice.as_ptr() as usize,
1101 len: slice.len(),
1102 }
1103 }
1104}
1105
1106#[derive(Debug, Clone)]
1107pub enum BufferGrowthStrategy {
1108 Fixed { size: u32 },
1110 Exponential { current_size: u32, max_size: u32 },
1112}
1113
1114impl Default for BufferGrowthStrategy {
1115 fn default() -> Self {
1116 Self::Exponential {
1117 current_size: 4 * 1024, max_size: 2 * 1024 * 1024, }
1120 }
1121}
1122
1123impl BufferGrowthStrategy {
1124 pub fn fixed(size: u32) -> Self {
1125 Self::Fixed { size }
1126 }
1127
1128 pub fn exponential(initial_size: u32, max_size: u32) -> Self {
1129 Self::Exponential {
1130 current_size: initial_size,
1131 max_size,
1132 }
1133 }
1134
1135 pub fn next_size(&mut self) -> u32 {
1137 match self {
1138 Self::Fixed { size } => *size,
1139 Self::Exponential {
1140 current_size,
1141 max_size,
1142 } => {
1143 let result = *current_size;
1144 if *current_size < *max_size {
1145 *current_size = current_size.saturating_mul(2).min(*max_size);
1146 }
1147 result
1148 }
1149 }
1150 }
1151}
1152
1153enum BuffersWithOffsets {
1154 AllKept {
1155 buffers: Arc<[ByteBuffer]>,
1156 offsets: Option<Vec<u32>>,
1157 },
1158 SomeCompacted {
1159 buffers: Vec<Option<ByteBuffer>>,
1160 offsets: Option<Vec<u32>>,
1161 },
1162}
1163
1164impl BuffersWithOffsets {
1165 pub fn from_array(
1166 array: &VarBinViewArray,
1167 compaction_threshold: f64,
1168 ctx: &mut ExecutionCtx,
1169 ) -> Self {
1170 if compaction_threshold == 0.0 {
1171 return Self::AllKept {
1172 buffers: Arc::from(
1173 array
1174 .data_buffers()
1175 .iter()
1176 .cloned()
1177 .map(|b| b.unwrap_host())
1178 .collect_vec(),
1179 ),
1180 offsets: None,
1181 };
1182 }
1183
1184 let buffer_utilizations = array
1185 .buffer_utilizations(ctx)
1186 .vortex_expect("buffer_utilizations in BuffersWithOffsets::from_array");
1187 let mut has_rewrite = false;
1188 let mut has_nonzero_offset = false;
1189 for utilization in buffer_utilizations.iter() {
1190 match compaction_strategy(utilization, compaction_threshold) {
1191 CompactionStrategy::KeepFull => continue,
1192 CompactionStrategy::Slice { .. } => has_nonzero_offset = true,
1193 CompactionStrategy::Rewrite => has_rewrite = true,
1194 }
1195 }
1196
1197 let buffers_with_offsets_iter = buffer_utilizations
1198 .iter()
1199 .zip(array.data_buffers().iter())
1200 .map(|(utilization, buffer)| {
1201 match compaction_strategy(utilization, compaction_threshold) {
1202 CompactionStrategy::KeepFull => (Some(buffer.as_host().clone()), 0),
1203 CompactionStrategy::Slice { start, end } => (
1204 Some(buffer.as_host().slice(start as usize..end as usize)),
1205 start,
1206 ),
1207 CompactionStrategy::Rewrite => (None, 0),
1208 }
1209 });
1210
1211 match (has_rewrite, has_nonzero_offset) {
1212 (false, false) => {
1214 let buffers: Vec<_> = buffers_with_offsets_iter
1215 .map(|(b, _)| b.vortex_expect("already checked for rewrite"))
1216 .collect();
1217 Self::AllKept {
1218 buffers: Arc::from(buffers),
1219 offsets: None,
1220 }
1221 }
1222 (true, false) => {
1224 let buffers: Vec<_> = buffers_with_offsets_iter.map(|(b, _)| b).collect();
1225 Self::SomeCompacted {
1226 buffers,
1227 offsets: None,
1228 }
1229 }
1230 (false, true) => {
1232 let (buffers, offsets): (Vec<_>, _) = buffers_with_offsets_iter
1233 .map(|(buffer, offset)| {
1234 (buffer.vortex_expect("already checked for rewrite"), offset)
1235 })
1236 .collect();
1237 Self::AllKept {
1238 buffers: Arc::from(buffers),
1239 offsets: Some(offsets),
1240 }
1241 }
1242 (true, true) => {
1244 let (buffers, offsets) = buffers_with_offsets_iter.collect();
1245 Self::SomeCompacted {
1246 buffers,
1247 offsets: Some(offsets),
1248 }
1249 }
1250 }
1251 }
1252}
1253
1254#[derive(Debug, Clone, Copy, PartialEq, Eq)]
1255enum CompactionStrategy {
1256 KeepFull,
1257 Slice {
1259 start: u32,
1260 end: u32,
1261 },
1262 Rewrite,
1264}
1265
1266fn compaction_strategy(
1267 buffer_utilization: &BufferUtilization,
1268 threshold: f64,
1269) -> CompactionStrategy {
1270 match buffer_utilization.overall_utilization() {
1271 0.0 => CompactionStrategy::Rewrite,
1273 utilised if utilised >= threshold => CompactionStrategy::KeepFull,
1274 _ if buffer_utilization.range_utilization() >= threshold => {
1275 let Range { start, end } = buffer_utilization.range();
1276 CompactionStrategy::Slice { start, end }
1277 }
1278 _ => CompactionStrategy::Rewrite,
1279 }
1280}
1281
1282enum ViewAdjustment {
1283 Precomputed(PrecomputedViewAdjustment),
1284 Rewriting(RewritingViewAdjustment),
1285}
1286
1287impl ViewAdjustment {
1288 fn shift(buffer_offset: u32, offsets: Option<Vec<u32>>) -> Self {
1289 Self::Precomputed(PrecomputedViewAdjustment::Shift {
1290 buffer_offset,
1291 offsets,
1292 })
1293 }
1294
1295 fn lookup(buffer_lookup: Vec<u32>, offsets: Option<Vec<u32>>) -> Self {
1296 Self::Precomputed(PrecomputedViewAdjustment::Lookup {
1297 buffer_lookup,
1298 offsets,
1299 })
1300 }
1301
1302 fn rewriting(buffer_lookup: Vec<Option<u32>>, offsets: Option<Vec<u32>>) -> Self {
1303 Self::Rewriting(RewritingViewAdjustment {
1304 buffer_lookup,
1305 offsets,
1306 })
1307 }
1308}
1309
1310enum PrecomputedViewAdjustment {
1312 Shift {
1313 buffer_offset: u32,
1314 offsets: Option<Vec<u32>>,
1315 },
1316 Lookup {
1317 buffer_lookup: Vec<u32>,
1318 offsets: Option<Vec<u32>>,
1319 },
1320}
1321
1322impl PrecomputedViewAdjustment {
1323 #[inline]
1324 fn adjust_view(&self, view: &BinaryView) -> BinaryView {
1325 if view.is_inlined() {
1326 return *view;
1327 }
1328 let view_ref = view.as_view();
1329 match self {
1330 Self::Shift {
1331 buffer_offset,
1332 offsets,
1333 } => {
1334 let b_idx = view_ref.buffer_index;
1335 let offset_shift = offsets
1336 .as_ref()
1337 .map(|o| o[b_idx as usize])
1338 .unwrap_or_default();
1339
1340 if view_ref.offset < offset_shift {
1343 return BinaryView::empty_view();
1344 }
1345
1346 view_ref
1347 .with_buffer_and_offset(b_idx + buffer_offset, view_ref.offset - offset_shift)
1348 }
1349 Self::Lookup {
1350 buffer_lookup,
1351 offsets,
1352 } => {
1353 let b_idx = view_ref.buffer_index;
1354 let buffer = buffer_lookup[b_idx as usize];
1355 let offset_shift = offsets
1356 .as_ref()
1357 .map(|o| o[b_idx as usize])
1358 .unwrap_or_default();
1359
1360 if view_ref.offset < offset_shift {
1363 return BinaryView::empty_view();
1364 }
1365
1366 view_ref.with_buffer_and_offset(buffer, view_ref.offset - offset_shift)
1367 }
1368 }
1369 .into()
1370 }
1371}
1372
1373struct RewritingViewAdjustment {
1374 buffer_lookup: Vec<Option<u32>>,
1375 offsets: Option<Vec<u32>>,
1376}
1377
1378impl RewritingViewAdjustment {
1379 #[inline]
1382 fn adjust_view(&self, view: &BinaryView) -> Option<BinaryView> {
1383 if view.is_inlined() {
1384 return Some(*view);
1385 }
1386
1387 let view_ref = view.as_view();
1388 self.buffer_lookup[view_ref.buffer_index as usize].map(|buffer| {
1389 let offset_shift = self
1390 .offsets
1391 .as_ref()
1392 .map(|o| o[view_ref.buffer_index as usize])
1393 .unwrap_or_default();
1394 view_ref
1395 .with_buffer_and_offset(buffer, view_ref.offset - offset_shift)
1396 .into()
1397 })
1398 }
1399}
1400
1401#[cfg(test)]
1402mod tests {
1403 use vortex_buffer::BufferAllocatorRef;
1404 use vortex_buffer::ByteBuffer;
1405 use vortex_error::VortexResult;
1406 use vortex_mask::Mask;
1407
1408 use crate::IntoArray;
1409 use crate::VortexSessionExecute;
1410 use crate::array_session;
1411 use crate::assert_arrays_eq;
1412 use crate::builders::ArrayBuilder;
1413 use crate::builders::VarBinViewBuilder;
1414 use crate::builders::varbinview::VarBinViewArray;
1415 use crate::dtype::DType;
1416 use crate::dtype::Nullability;
1417
1418 const LONG: &str = "a value that is far too long to inline";
1420
1421 #[test]
1424 fn test_append_buffer_with_lengths() {
1425 let mut ctx = array_session().create_execution_ctx();
1426 let mut builder = VarBinViewBuilder::with_capacity_in(
1427 DType::Utf8(Nullability::Nullable),
1428 8,
1429 BufferAllocatorRef::statically_allocated(),
1430 );
1431
1432 builder.append_value(LONG);
1434
1435 let heap = ByteBuffer::copy_from([LONG.as_bytes(), b"", b"tiny"].concat());
1436 let heap_ptr = heap.as_ptr();
1437 let lengths = [u32::try_from(LONG.len()).unwrap(), 0, 4];
1438 builder.append_buffer_with_lengths(heap, &lengths, &Mask::from_iter([true, false, true]));
1439
1440 builder.append_value("tail");
1441
1442 let actual = builder.finish_into_varbinview();
1443 assert_eq!(actual.data_buffers()[1].as_host().as_ptr(), heap_ptr);
1445
1446 let expected = <VarBinViewArray as FromIterator<_>>::from_iter([
1447 Some(LONG),
1448 Some(LONG),
1449 None,
1450 Some("tiny"),
1451 Some("tail"),
1452 ]);
1453 assert_arrays_eq!(actual, expected, &mut ctx);
1454 }
1455
1456 #[test]
1458 fn test_append_buffer_with_offsets() {
1459 let mut ctx = array_session().create_execution_ctx();
1460 let mut builder = VarBinViewBuilder::with_capacity_in(
1461 DType::Utf8(Nullability::Nullable),
1462 8,
1463 BufferAllocatorRef::statically_allocated(),
1464 );
1465
1466 let heap = ByteBuffer::copy_from(format!("..{LONG}tiny!!"));
1467 let long_len = u32::try_from(LONG.len()).unwrap();
1468 let offsets = [2u32, 2 + long_len, 2 + long_len, 2 + long_len + 4];
1469 builder.append_buffer_with_offsets(
1470 heap.clone(),
1471 &offsets,
1472 &Mask::from_iter([true, false, true]),
1473 );
1474
1475 let actual = builder.finish_into_varbinview();
1476 assert_eq!(actual.data_buffers()[0].as_host().as_ptr(), unsafe {
1479 heap.as_ptr().add(2)
1480 });
1481 assert_eq!(
1482 actual.data_buffers()[0].len(),
1483 LONG.len() + 4,
1484 "only the referenced range must be adopted"
1485 );
1486
1487 let expected =
1488 <VarBinViewArray as FromIterator<_>>::from_iter([Some(LONG), None, Some("tiny")]);
1489 assert_arrays_eq!(actual, expected, &mut ctx);
1490 }
1491
1492 #[test]
1495 fn test_append_buffer_with_lengths_compacts_underutilized_heap() {
1496 let mut ctx = array_session().create_execution_ctx();
1497 let mut builder = VarBinViewBuilder::with_compaction_in(
1498 DType::Utf8(Nullability::Nullable),
1499 4,
1500 1.0,
1501 BufferAllocatorRef::statically_allocated(),
1502 );
1503
1504 let heap = ByteBuffer::copy_from([b"short".as_slice(), LONG.as_bytes(), b"tiny"].concat());
1505 let lengths = [5u32, u32::try_from(LONG.len()).unwrap(), 4];
1506 builder.append_buffer_with_lengths(heap, &lengths, &Mask::new_true(3));
1507
1508 let actual = builder.finish_into_varbinview();
1509 assert_eq!(actual.data_buffers().len(), 1);
1510 assert_eq!(
1511 actual.data_buffers()[0].len(),
1512 LONG.len(),
1513 "the compact buffer must hold only the non-inlined values"
1514 );
1515
1516 let expected = <VarBinViewArray as FromIterator<_>>::from_iter([
1517 Some("short"),
1518 Some(LONG),
1519 Some("tiny"),
1520 ]);
1521 assert_arrays_eq!(actual, expected, &mut ctx);
1522 }
1523
1524 #[test]
1527 fn test_append_buffer_with_lengths_compaction_skips_null_bytes() {
1528 let mut ctx = array_session().create_execution_ctx();
1529 let mut builder = VarBinViewBuilder::with_compaction_in(
1530 DType::Utf8(Nullability::Nullable),
1531 2,
1532 1.0,
1533 BufferAllocatorRef::statically_allocated(),
1534 );
1535
1536 let heap = ByteBuffer::copy_from([LONG.as_bytes(), LONG.as_bytes()].concat());
1537 let lengths = [u32::try_from(LONG.len()).unwrap(); 2];
1538 builder.append_buffer_with_lengths(heap, &lengths, &Mask::from_iter([false, true]));
1539
1540 let actual = builder.finish_into_varbinview();
1541 assert_eq!(actual.data_buffers().len(), 1);
1542 assert_eq!(
1543 actual.data_buffers()[0].len(),
1544 LONG.len(),
1545 "the compact buffer must omit bytes belonging to null rows"
1546 );
1547
1548 let expected = <VarBinViewArray as FromIterator<_>>::from_iter([None::<&str>, Some(LONG)]);
1549 assert_arrays_eq!(actual, expected, &mut ctx);
1550 }
1551
1552 #[test]
1555 fn test_push_buffers_deduplicates() {
1556 let mut builder = VarBinViewBuilder::with_buffer_deduplication_in(
1557 DType::Utf8(Nullability::Nullable),
1558 8,
1559 BufferAllocatorRef::statically_allocated(),
1560 );
1561
1562 let first = ByteBuffer::copy_from(LONG);
1563 let second = ByteBuffer::copy_from("another value far too long to inline");
1564
1565 assert_eq!(
1566 builder.push_buffers([first.clone(), second.clone()]),
1567 [0, 1]
1568 );
1569 assert_eq!(builder.push_buffers([second, first]), [1, 0]);
1570 assert_eq!(builder.completed_block_count(), 2);
1571 }
1572
1573 #[test]
1576 fn test_append_views_gathered() {
1577 let mut ctx = array_session().create_execution_ctx();
1578 let dictionary = <VarBinViewArray as FromIterator<_>>::from_iter([
1579 Some("tiny"),
1580 Some(LONG),
1581 Some("small"),
1582 ]);
1583 let buffers = dictionary
1584 .data_buffers()
1585 .iter()
1586 .map(|buffer| buffer.as_host().clone())
1587 .collect::<Vec<_>>();
1588
1589 let mut builder = VarBinViewBuilder::with_capacity_in(
1590 DType::Utf8(Nullability::Nullable),
1591 8,
1592 BufferAllocatorRef::statically_allocated(),
1593 );
1594 builder.append_value(LONG);
1596
1597 let codes: [usize; 4] = [1, 0, usize::MAX, 2];
1598 let views = dictionary.views();
1599 builder.append_views_gathered(
1600 buffers,
1601 views,
1602 &Mask::from_iter([true, true, false, true]),
1603 |row| codes[row],
1605 );
1606
1607 let expected = <VarBinViewArray as FromIterator<_>>::from_iter([
1608 Some(LONG),
1609 Some(LONG),
1610 Some("tiny"),
1611 None,
1612 Some("small"),
1613 ]);
1614 assert_arrays_eq!(builder.finish_into_varbinview(), expected, &mut ctx);
1615 }
1616
1617 #[test]
1620 fn test_append_views_scattered() {
1621 use crate::arrays::varbinview::build_views::BinaryView;
1622
1623 let mut ctx = array_session().create_execution_ctx();
1624 let patch_values =
1625 <VarBinViewArray as FromIterator<_>>::from_iter([Some(LONG), Some("tiny")]);
1626 let mut buffers = patch_values
1627 .data_buffers()
1628 .iter()
1629 .map(|buffer| buffer.as_host().clone())
1630 .collect::<Vec<_>>();
1631
1632 let fill_bytes = ByteBuffer::copy_from("a fill value too long to inline");
1633 buffers.push(fill_bytes.clone());
1634 let fill = BinaryView::make_view(
1635 fill_bytes.as_slice(),
1636 u32::try_from(buffers.len() - 1).unwrap(),
1637 0,
1638 );
1639
1640 let mut builder = VarBinViewBuilder::with_capacity_in(
1641 DType::Utf8(Nullability::Nullable),
1642 8,
1643 BufferAllocatorRef::statically_allocated(),
1644 );
1645 let views = patch_values.views();
1646 builder.append_views_scattered(
1647 buffers,
1648 5,
1649 fill,
1650 [(1usize, views[0]), (3usize, views[1])].into_iter(),
1651 &Mask::from_iter([true, true, false, true, true]),
1652 );
1653
1654 let expected = <VarBinViewArray as FromIterator<_>>::from_iter([
1655 Some("a fill value too long to inline"),
1656 Some(LONG),
1657 None,
1658 Some("tiny"),
1659 Some("a fill value too long to inline"),
1660 ]);
1661 assert_arrays_eq!(builder.finish_into_varbinview(), expected, &mut ctx);
1662 }
1663
1664 #[test]
1668 fn test_append_views_gathered_compacts_unreferenced_values() {
1669 const DEAD: &str = "a dead value nobody gathers, far too long to inline";
1670 const OTHER: &str = "another value far too long to inline";
1671
1672 let mut ctx = array_session().create_execution_ctx();
1673 let dictionary =
1674 <VarBinViewArray as FromIterator<_>>::from_iter([Some(LONG), Some(DEAD), Some(OTHER)]);
1675 assert_eq!(dictionary.data_buffers().len(), 1);
1676 let buffers = dictionary
1677 .data_buffers()
1678 .iter()
1679 .map(|buffer| buffer.as_host().clone())
1680 .collect::<Vec<_>>();
1681
1682 let mut builder = VarBinViewBuilder::with_compaction_in(
1683 DType::Utf8(Nullability::Nullable),
1684 8,
1685 1.0,
1686 BufferAllocatorRef::statically_allocated(),
1687 );
1688 let codes: [usize; 5] = [2, 0, usize::MAX, 2, 0];
1689 builder.append_views_gathered(
1690 buffers,
1691 dictionary.views(),
1692 &Mask::from_iter([true, true, false, true, true]),
1693 |row| codes[row],
1694 );
1695
1696 let actual = builder.finish_into_varbinview();
1697 assert_eq!(actual.data_buffers().len(), 1);
1698 assert_eq!(
1699 actual.data_buffers()[0].len(),
1700 LONG.len() + OTHER.len(),
1701 "the rewritten buffer must hold each referenced value exactly once"
1702 );
1703
1704 let expected = <VarBinViewArray as FromIterator<_>>::from_iter([
1705 Some(OTHER),
1706 Some(LONG),
1707 None,
1708 Some(OTHER),
1709 Some(LONG),
1710 ]);
1711 assert_arrays_eq!(actual, expected, &mut ctx);
1712 }
1713
1714 #[test]
1717 fn test_append_views_gathered_slices_contiguous_range() {
1718 const TAIL: &str = "the referenced tail value, too long to inline";
1719
1720 let mut ctx = array_session().create_execution_ctx();
1721 let dictionary = <VarBinViewArray as FromIterator<_>>::from_iter([Some(LONG), Some(TAIL)]);
1722 let heap = dictionary.data_buffers()[0].as_host().clone();
1723
1724 let mut builder = VarBinViewBuilder::with_compaction_in(
1725 DType::Utf8(Nullability::Nullable),
1726 8,
1727 1.0,
1728 BufferAllocatorRef::statically_allocated(),
1729 );
1730 builder.append_views_gathered(
1731 [heap.clone()],
1732 dictionary.views(),
1733 &Mask::new_true(2),
1734 |_| 1,
1735 );
1736
1737 let actual = builder.finish_into_varbinview();
1738 assert_eq!(actual.data_buffers().len(), 1);
1739 assert_eq!(actual.data_buffers()[0].len(), TAIL.len());
1740 assert_eq!(actual.data_buffers()[0].as_host().as_ptr(), unsafe {
1742 heap.as_ptr().add(LONG.len())
1743 });
1744
1745 let expected = <VarBinViewArray as FromIterator<_>>::from_iter([Some(TAIL), Some(TAIL)]);
1746 assert_arrays_eq!(actual, expected, &mut ctx);
1747 }
1748
1749 #[test]
1751 fn test_append_views_gathered_keeps_utilized_buffer() {
1752 const SHORTER: &str = "a shorter long value";
1753
1754 let mut ctx = array_session().create_execution_ctx();
1755 let dictionary =
1756 <VarBinViewArray as FromIterator<_>>::from_iter([Some(LONG), Some(SHORTER)]);
1757 let heap = dictionary.data_buffers()[0].as_host().clone();
1758
1759 let mut builder = VarBinViewBuilder::with_compaction_in(
1760 DType::Utf8(Nullability::Nullable),
1761 8,
1762 0.5,
1763 BufferAllocatorRef::statically_allocated(),
1764 );
1765 builder.append_views_gathered(
1766 [heap.clone()],
1767 dictionary.views(),
1768 &Mask::new_true(1),
1769 |_| 0,
1770 );
1771
1772 let actual = builder.finish_into_varbinview();
1773 assert_eq!(actual.data_buffers().len(), 1);
1774 assert_eq!(actual.data_buffers()[0].as_host().as_ptr(), heap.as_ptr());
1775 assert_eq!(actual.data_buffers()[0].len(), heap.len());
1776
1777 let expected = <VarBinViewArray as FromIterator<_>>::from_iter([Some(LONG)]);
1778 assert_arrays_eq!(actual, expected, &mut ctx);
1779 }
1780
1781 #[test]
1784 fn test_append_views_gathered_all_null_drops_buffers() {
1785 let mut ctx = array_session().create_execution_ctx();
1786 let dictionary = <VarBinViewArray as FromIterator<_>>::from_iter([Some(LONG)]);
1787 let buffers = dictionary
1788 .data_buffers()
1789 .iter()
1790 .map(|buffer| buffer.as_host().clone())
1791 .collect::<Vec<_>>();
1792
1793 let mut builder = VarBinViewBuilder::with_compaction_in(
1794 DType::Utf8(Nullability::Nullable),
1795 8,
1796 1.0,
1797 BufferAllocatorRef::statically_allocated(),
1798 );
1799 builder.append_views_gathered(buffers, dictionary.views(), &Mask::new_false(3), |_| {
1800 usize::MAX
1801 });
1802
1803 let actual = builder.finish_into_varbinview();
1804 assert!(
1805 actual.data_buffers().is_empty(),
1806 "an all-null gather must not retain any buffers"
1807 );
1808
1809 let expected = <VarBinViewArray as FromIterator<_>>::from_iter([None::<&str>, None, None]);
1810 assert_arrays_eq!(actual, expected, &mut ctx);
1811 }
1812
1813 #[test]
1815 fn test_append_buffer_with_lengths_drops_fully_inlined_heap() {
1816 let mut ctx = array_session().create_execution_ctx();
1817 let mut builder = VarBinViewBuilder::with_compaction_in(
1818 DType::Utf8(Nullability::Nullable),
1819 4,
1820 1.0,
1821 BufferAllocatorRef::statically_allocated(),
1822 );
1823
1824 let heap = ByteBuffer::copy_from(b"shorttinysmall".as_slice());
1825 builder.append_buffer_with_lengths(heap, &[5u32, 4, 5], &Mask::new_true(3));
1826
1827 let actual = builder.finish_into_varbinview();
1828 assert!(
1829 actual.data_buffers().is_empty(),
1830 "a fully-inlined append must not retain any value bytes"
1831 );
1832
1833 let expected = <VarBinViewArray as FromIterator<_>>::from_iter([
1834 Some("short"),
1835 Some("tiny"),
1836 Some("small"),
1837 ]);
1838 assert_arrays_eq!(actual, expected, &mut ctx);
1839 }
1840
1841 #[test]
1842 fn test_utf8_builder() {
1843 let mut ctx = array_session().create_execution_ctx();
1844 let mut builder = VarBinViewBuilder::with_capacity_in(
1845 DType::Utf8(Nullability::Nullable),
1846 10,
1847 BufferAllocatorRef::statically_allocated(),
1848 );
1849
1850 builder.append_value("Hello");
1851 builder.append_null();
1852 builder.append_value("World");
1853
1854 builder.append_nulls(2);
1855
1856 builder.append_zeros(2);
1857 builder.append_value("test");
1858
1859 let actual = builder.finish();
1860 let expected = <VarBinViewArray as FromIterator<_>>::from_iter([
1861 Some("Hello"),
1862 None,
1863 Some("World"),
1864 None,
1865 None,
1866 Some(""),
1867 Some(""),
1868 Some("test"),
1869 ]);
1870 assert_arrays_eq!(actual, expected, &mut ctx);
1871 }
1872
1873 #[test]
1874 fn test_utf8_builder_with_extend() {
1875 let mut ctx = array_session().create_execution_ctx();
1876 let array = {
1877 let mut builder = VarBinViewBuilder::with_capacity_in(
1878 DType::Utf8(Nullability::Nullable),
1879 10,
1880 BufferAllocatorRef::statically_allocated(),
1881 );
1882 builder.append_null();
1883 builder.append_value("Hello2");
1884 builder.finish()
1885 };
1886 let mut builder = VarBinViewBuilder::with_capacity_in(
1887 DType::Utf8(Nullability::Nullable),
1888 10,
1889 BufferAllocatorRef::statically_allocated(),
1890 );
1891
1892 builder.append_value("Hello1");
1893 array.append_to_builder(&mut builder, &mut ctx).unwrap();
1894 builder.append_nulls(2);
1895 builder.append_value("Hello3");
1896
1897 let actual = builder.finish_into_canonical(&mut ctx);
1898 let expected = <VarBinViewArray as FromIterator<_>>::from_iter([
1899 Some("Hello1"),
1900 None,
1901 Some("Hello2"),
1902 None,
1903 None,
1904 Some("Hello3"),
1905 ]);
1906 assert_arrays_eq!(actual.into_array(), expected.into_array(), &mut ctx);
1907 }
1908
1909 #[test]
1910 fn test_buffer_deduplication() -> VortexResult<()> {
1911 let array = {
1912 let mut builder = VarBinViewBuilder::with_capacity_in(
1913 DType::Utf8(Nullability::Nullable),
1914 10,
1915 BufferAllocatorRef::statically_allocated(),
1916 );
1917 builder.append_value("This is a long string that should not be inlined");
1918 builder.append_value("short string");
1919 builder.finish_into_varbinview()
1920 };
1921
1922 assert_eq!(array.data_buffers().len(), 1);
1923 let mut builder = VarBinViewBuilder::with_buffer_deduplication_in(
1924 DType::Utf8(Nullability::Nullable),
1925 10,
1926 BufferAllocatorRef::statically_allocated(),
1927 );
1928
1929 let mut ctx = array_session().create_execution_ctx();
1930
1931 array.append_to_builder(&mut builder, &mut ctx)?;
1932 assert_eq!(builder.completed_block_count(), 1);
1933
1934 array
1935 .slice(1..2)?
1936 .append_to_builder(&mut builder, &mut ctx)?;
1937 array
1938 .slice(0..1)?
1939 .append_to_builder(&mut builder, &mut ctx)?;
1940 assert_eq!(builder.completed_block_count(), 1);
1941
1942 let array2 = {
1943 let mut builder = VarBinViewBuilder::with_capacity_in(
1944 DType::Utf8(Nullability::Nullable),
1945 10,
1946 BufferAllocatorRef::statically_allocated(),
1947 );
1948 builder.append_value("This is a long string that should not be inlined");
1949 builder.finish_into_varbinview()
1950 };
1951
1952 array2.append_to_builder(&mut builder, &mut ctx)?;
1953 assert_eq!(builder.completed_block_count(), 2);
1954
1955 array
1956 .slice(0..1)?
1957 .append_to_builder(&mut builder, &mut ctx)?;
1958 array2
1959 .slice(0..1)?
1960 .append_to_builder(&mut builder, &mut ctx)?;
1961 assert_eq!(builder.completed_block_count(), 2);
1962 Ok(())
1963 }
1964
1965 #[test]
1966 fn test_append_scalar() {
1967 let mut ctx = array_session().create_execution_ctx();
1968 use crate::scalar::Scalar;
1969
1970 let mut utf8_builder = VarBinViewBuilder::with_capacity_in(
1972 DType::Utf8(Nullability::Nullable),
1973 10,
1974 BufferAllocatorRef::statically_allocated(),
1975 );
1976
1977 let utf8_scalar1 = Scalar::utf8("hello", Nullability::Nullable);
1979 utf8_builder.append_scalar(&utf8_scalar1).unwrap();
1980
1981 let utf8_scalar2 = Scalar::utf8("world", Nullability::Nullable);
1983 utf8_builder.append_scalar(&utf8_scalar2).unwrap();
1984
1985 let null_scalar = Scalar::null(DType::Utf8(Nullability::Nullable));
1987 utf8_builder.append_scalar(&null_scalar).unwrap();
1988
1989 let array = utf8_builder.finish();
1990 let expected =
1991 <VarBinViewArray as FromIterator<_>>::from_iter([Some("hello"), Some("world"), None]);
1992 assert_arrays_eq!(&array, &expected, &mut ctx);
1993
1994 let mut binary_builder = VarBinViewBuilder::with_capacity_in(
1996 DType::Binary(Nullability::Nullable),
1997 10,
1998 BufferAllocatorRef::statically_allocated(),
1999 );
2000
2001 let binary_scalar = Scalar::binary(vec![1u8, 2, 3], Nullability::Nullable);
2002 binary_builder.append_scalar(&binary_scalar).unwrap();
2003
2004 let binary_null = Scalar::null(DType::Binary(Nullability::Nullable));
2005 binary_builder.append_scalar(&binary_null).unwrap();
2006
2007 let binary_array = binary_builder.finish();
2008 let expected =
2009 <VarBinViewArray as FromIterator<_>>::from_iter([Some(vec![1u8, 2, 3]), None]);
2010 assert_arrays_eq!(&binary_array, &expected, &mut ctx);
2011
2012 let mut builder = VarBinViewBuilder::with_capacity_in(
2014 DType::Utf8(Nullability::NonNullable),
2015 10,
2016 BufferAllocatorRef::statically_allocated(),
2017 );
2018 let wrong_scalar = Scalar::from(42i32);
2019 assert!(builder.append_scalar(&wrong_scalar).is_err());
2020 }
2021
2022 #[test]
2023 fn test_buffer_growth_strategies() {
2024 use super::BufferGrowthStrategy;
2025
2026 let mut strategy = BufferGrowthStrategy::fixed(1024);
2028
2029 assert_eq!(strategy.next_size(), 1024);
2031 assert_eq!(strategy.next_size(), 1024);
2032 assert_eq!(strategy.next_size(), 1024);
2033
2034 let mut strategy = BufferGrowthStrategy::exponential(1024, 8192);
2036
2037 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); }
2044
2045 #[test]
2046 fn test_large_value_allocation() {
2047 use super::BufferGrowthStrategy;
2048 use super::VarBinViewBuilder;
2049
2050 let mut builder = VarBinViewBuilder::new_in(
2051 DType::Binary(Nullability::Nullable),
2052 10,
2053 Default::default(),
2054 BufferGrowthStrategy::exponential(1024, 4096),
2055 0.0,
2056 BufferAllocatorRef::statically_allocated(),
2057 );
2058
2059 let large_value = vec![0u8; 8192];
2061
2062 builder.append_value(&large_value);
2064
2065 let array = builder.finish_into_varbinview();
2066 assert_eq!(array.len(), 1);
2067
2068 let retrieved = array
2070 .execute_scalar(0, &mut array_session().create_execution_ctx())
2071 .unwrap()
2072 .as_binary()
2073 .value()
2074 .cloned()
2075 .unwrap();
2076 assert_eq!(retrieved.len(), 8192);
2077 assert_eq!(retrieved.as_slice(), &large_value);
2078 }
2079}