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::BufferMut;
13use vortex_buffer::ByteBuffer;
14use vortex_buffer::ByteBufferMut;
15use vortex_error::VortexExpect;
16use vortex_error::VortexResult;
17use vortex_error::vortex_bail;
18use vortex_error::vortex_ensure;
19use vortex_mask::AllOr;
20use vortex_mask::Mask;
21use vortex_utils::aliases::hash_map::Entry;
22use vortex_utils::aliases::hash_map::HashMap;
23
24use crate::ArrayRef;
25use crate::ExecutionCtx;
26use crate::IntoArray;
27use crate::arrays::VarBinViewArray;
28use crate::arrays::varbinview::VarBinViewArrayExt;
29use crate::arrays::varbinview::build_views::BinaryView;
30use crate::arrays::varbinview::build_views::MAX_BUFFER_LEN;
31use crate::arrays::varbinview::build_views::extend_views;
32use crate::arrays::varbinview::compact::BufferUtilization;
33use crate::builders::ArrayBuilder;
34use crate::builders::LazyBitBufferBuilder;
35use crate::canonical::Canonical;
36use crate::dtype::DType;
37use crate::dtype::NativePType;
38use crate::scalar::Scalar;
39
40pub struct VarBinViewBuilder {
42 dtype: DType,
43 views_builder: BufferMut<BinaryView>,
44 nulls: LazyBitBufferBuilder,
45 completed: CompletedBuffers,
46 in_progress: Option<ByteBufferMut>,
47 growth_strategy: BufferGrowthStrategy,
48 compaction_threshold: f64,
49}
50
51impl VarBinViewBuilder {
52 pub fn with_capacity(dtype: DType, capacity: usize) -> Self {
53 Self::new(dtype, capacity, Default::default(), Default::default(), 0.0)
54 }
55
56 pub fn with_buffer_deduplication(dtype: DType, capacity: usize) -> Self {
57 Self::new(
58 dtype,
59 capacity,
60 CompletedBuffers::Deduplicated(Default::default()),
61 Default::default(),
62 0.0,
63 )
64 }
65
66 pub fn with_compaction(dtype: DType, capacity: usize, compaction_threshold: f64) -> Self {
67 Self::new(
68 dtype,
69 capacity,
70 Default::default(),
71 Default::default(),
72 compaction_threshold,
73 )
74 }
75
76 pub fn new(
77 dtype: DType,
78 capacity: usize,
79 completed: CompletedBuffers,
80 growth_strategy: BufferGrowthStrategy,
81 compaction_threshold: f64,
82 ) -> Self {
83 assert!(
84 matches!(dtype, DType::Utf8(_) | DType::Binary(_)),
85 "VarBinViewBuilder DType must be Utf8 or Binary."
86 );
87 Self {
88 views_builder: BufferMut::with_capacity_preferred_aligned(
89 capacity,
90 Alignment::of::<BinaryView>(),
91 None,
92 ),
93 nulls: LazyBitBufferBuilder::new(capacity),
94 completed,
95 in_progress: None,
96 dtype,
97 growth_strategy,
98 compaction_threshold,
99 }
100 }
101
102 fn append_value_view(&mut self, value: &[u8]) {
103 let length =
104 u32::try_from(value.len()).vortex_expect("cannot have a single string >2^32 in length");
105 if length <= 12 {
106 self.views_builder.push(BinaryView::make_view(value, 0, 0));
107 return;
108 }
109
110 let (buffer_idx, offset) = self.append_value_to_buffer(value);
111 let view = BinaryView::make_view(value, buffer_idx, offset);
112 self.views_builder.push(view);
113 }
114
115 pub fn append_value<S: AsRef<[u8]>>(&mut self, value: S) {
117 self.append_value_view(value.as_ref());
118 self.nulls.append_non_null();
119 }
120
121 pub fn append_n_values<S: AsRef<[u8]>>(&mut self, value: S, n: usize) {
123 if n == 0 {
124 return;
125 }
126 let bytes = value.as_ref();
127 let view = if bytes.len() <= BinaryView::MAX_INLINED_SIZE {
128 BinaryView::make_view(bytes, 0, 0)
129 } else {
130 let (buffer_idx, offset) = self.append_value_to_buffer(bytes);
131 BinaryView::make_view(bytes, buffer_idx, offset)
132 };
133 self.views_builder.push_n(view, n);
134 self.nulls.append_n_non_nulls(n);
135 }
136
137 fn flush_in_progress(&mut self) {
138 let Some(block) = self.in_progress.take() else {
139 return;
140 };
141
142 assert!(block.len() < u32::MAX as usize, "Block too large");
143
144 let initial_len = self.completed.len();
145 self.completed.push(block.freeze());
146 assert_eq!(
147 self.completed.len(),
148 initial_len + 1,
149 "Invalid state, just completed block already exists"
150 );
151 }
152
153 fn init_in_progress(&mut self, min_len: usize) {
154 let next_buffer_size = self.growth_strategy.next_size() as usize;
155 let to_reserve = next_buffer_size.max(min_len);
156 self.in_progress = Some(ByteBufferMut::with_capacity_preferred_aligned(
157 to_reserve,
158 Alignment::of::<u8>(),
159 None,
160 ));
161 }
162
163 fn append_value_to_buffer(&mut self, value: &[u8]) -> (u32, u32) {
165 assert!(
166 value.len() > BinaryView::MAX_INLINED_SIZE,
167 "must inline small strings"
168 );
169
170 if let Some(in_progress) = &mut self.in_progress {
171 let required_cap = in_progress.len() + value.len();
172 if in_progress.capacity() < required_cap {
173 self.flush_in_progress();
174 self.init_in_progress(value.len());
175 }
176 } else {
177 self.init_in_progress(value.len())
178 };
179
180 let in_progress = self
181 .in_progress
182 .as_mut()
183 .vortex_expect("in_progress just set");
184
185 let buffer_idx = self.completed.len();
186 let offset = u32::try_from(in_progress.len()).vortex_expect("too many buffers");
187 in_progress.extend_from_slice(value);
188
189 (buffer_idx, offset)
190 }
191
192 pub fn completed_block_count(&self) -> u32 {
193 self.completed.len()
194 }
195
196 fn compacts_buffers(&self) -> bool {
201 self.compaction_threshold > 0.0
202 }
203
204 pub fn append_views_built_at(
224 &mut self,
225 validity: &Mask,
226 build: impl FnOnce(u32) -> VortexResult<(Vec<ByteBuffer>, Buffer<BinaryView>)>,
227 ) -> VortexResult<()> {
228 self.flush_in_progress();
229
230 let start_index = self.completed.len();
231 let (buffers, views) = build(start_index)?;
232 assert_eq!(
233 views.len(),
234 validity.len(),
235 "Must build one view per validity entry"
236 );
237
238 let expected_completed_len = start_index as usize + buffers.len();
239 self.completed.extend_from_slice_unchecked(&buffers);
240 assert_eq!(
241 self.completed.len() as usize,
242 expected_completed_len,
243 "Some buffers already exist",
244 );
245 self.views_builder.extend_trusted(views.iter().copied());
246 self.nulls.append_validity_mask(validity);
247
248 debug_assert_eq!(self.nulls.len(), self.views_builder.len());
249 Ok(())
250 }
251
252 fn push_buffers(&mut self, buffers: impl IntoIterator<Item = ByteBuffer>) -> Vec<u32> {
268 self.flush_in_progress();
269 buffers
270 .into_iter()
271 .map(|buffer| self.completed.push(buffer))
272 .collect()
273 }
274
275 fn push_buffers_compacted(
281 &mut self,
282 buffers: Vec<ByteBuffer>,
283 utilizations: &[BufferUtilization],
284 ) -> Vec<CompactedSlot> {
285 self.flush_in_progress();
286 buffers
287 .into_iter()
288 .zip(utilizations)
289 .map(|(buffer, utilization)| {
290 match compaction_strategy(utilization, self.compaction_threshold) {
291 CompactionStrategy::KeepFull => CompactedSlot::Kept {
292 index: self.completed.push(buffer),
293 },
294 CompactionStrategy::Slice { start, end } => CompactedSlot::Sliced {
295 index: self
296 .completed
297 .push(buffer.slice(start as usize..end as usize)),
298 shift: start,
299 },
300 CompactionStrategy::Rewrite => CompactedSlot::Rewrite { source: buffer },
301 }
302 })
303 .collect()
304 }
305
306 fn remap_view_compacted(&mut self, view: BinaryView, slots: &[CompactedSlot]) -> BinaryView {
310 if view.is_inlined() {
311 return view;
312 }
313 let view_ref = view.as_view();
314 match &slots[view_ref.buffer_index as usize] {
315 CompactedSlot::Kept { index } => view_ref
316 .with_buffer_and_offset(*index, view_ref.offset)
317 .into(),
318 CompactedSlot::Sliced { index, shift } => view_ref
319 .with_buffer_and_offset(*index, view_ref.offset - shift)
320 .into(),
321 CompactedSlot::Rewrite { source } => {
322 let bytes = &source.as_slice()[view_ref.as_range()];
323 let (buffer_idx, offset) = self.append_value_to_buffer(bytes);
324 BinaryView::make_view(bytes, buffer_idx, offset)
325 }
326 }
327 }
328
329 pub fn append_views_gathered(
348 &mut self,
349 buffers: impl IntoIterator<Item = ByteBuffer>,
350 views: &[BinaryView],
351 validity: &Mask,
352 view_at: impl Fn(usize) -> usize,
353 ) {
354 if self.compacts_buffers() {
355 return self.append_views_gathered_compacted(
356 buffers.into_iter().collect(),
357 views,
358 validity,
359 view_at,
360 );
361 }
362
363 let mapping = self.push_buffers(buffers);
364
365 self.views_builder.reserve(validity.len());
366 match validity {
367 Mask::AllTrue(len) => self
368 .views_builder
369 .extend_trusted((0..*len).map(|row| remap_view(views[view_at(row)], &mapping))),
370 Mask::AllFalse(len) => self.views_builder.push_n(BinaryView::empty_view(), *len),
371 Mask::Values(values) => {
372 let bits = values.bit_buffer();
373 let start = self.views_builder.len();
374 self.views_builder
375 .push_n(BinaryView::empty_view(), bits.len());
376 let gathered = &mut self.views_builder[start..];
377 bits.for_each_set_index(|row| {
378 gathered[row] = remap_view(views[view_at(row)], &mapping);
379 });
380 }
381 }
382
383 self.nulls.append_validity_mask(validity);
384 debug_assert_eq!(self.nulls.len(), self.views_builder.len());
385 }
386
387 fn append_views_gathered_compacted(
393 &mut self,
394 buffers: Vec<ByteBuffer>,
395 views: &[BinaryView],
396 validity: &Mask,
397 view_at: impl Fn(usize) -> usize,
398 ) {
399 let mut referenced = vec![false; views.len()];
400 match validity {
401 Mask::AllTrue(len) => (0..*len).for_each(|row| referenced[view_at(row)] = true),
402 Mask::AllFalse(_) => {}
403 Mask::Values(values) => values
404 .bit_buffer()
405 .for_each_set_index(|row| referenced[view_at(row)] = true),
406 }
407
408 let mut utilizations = unmeasured_utilizations(&buffers);
409 for (view, _) in views.iter().zip(&referenced).filter(|(_, used)| **used) {
410 measure_view(&mut utilizations, view);
411 }
412 let slots = self.push_buffers_compacted(buffers, &utilizations);
413
414 let mut remapped = vec![BinaryView::empty_view(); views.len()];
415 for (idx, view) in views.iter().enumerate() {
416 if referenced[idx] {
417 remapped[idx] = self.remap_view_compacted(*view, &slots);
418 }
419 }
420
421 self.views_builder.reserve(validity.len());
422 match validity {
423 Mask::AllTrue(len) => self
424 .views_builder
425 .extend_trusted((0..*len).map(|row| remapped[view_at(row)])),
426 Mask::AllFalse(len) => self.views_builder.push_n(BinaryView::empty_view(), *len),
427 Mask::Values(values) => {
428 let bits = values.bit_buffer();
429 let start = self.views_builder.len();
430 self.views_builder
431 .push_n(BinaryView::empty_view(), bits.len());
432 let gathered = &mut self.views_builder[start..];
433 bits.for_each_set_index(|row| gathered[row] = remapped[view_at(row)]);
434 }
435 }
436
437 self.nulls.append_validity_mask(validity);
438 debug_assert_eq!(self.nulls.len(), self.views_builder.len());
439 }
440
441 pub fn append_views_scattered(
458 &mut self,
459 buffers: impl IntoIterator<Item = ByteBuffer>,
460 len: usize,
461 fill: BinaryView,
462 patches: impl Iterator<Item = (usize, BinaryView)>,
463 validity: &Mask,
464 ) {
465 assert_eq!(validity.len(), len, "Must have one validity entry per row");
466 let mapping = self.push_buffers(buffers);
467
468 let start = self.views_builder.len();
469 self.views_builder.push_n(remap_view(fill, &mapping), len);
470 let scattered = &mut self.views_builder[start..];
471 for (row, view) in patches {
472 scattered[row] = remap_view(view, &mapping);
473 }
474
475 self.nulls.append_validity_mask(validity);
476 debug_assert_eq!(self.nulls.len(), self.views_builder.len());
477 }
478
479 pub fn append_buffer_with_lengths<P: NativePType + AsPrimitive<usize>>(
497 &mut self,
498 bytes: ByteBuffer,
499 lengths: &[P],
500 validity: &Mask,
501 ) {
502 assert_eq!(
503 lengths.len(),
504 validity.len(),
505 "Must have one length per validity entry"
506 );
507 self.append_buffer_views(&bytes, lengths.len(), validity, |i| lengths[i].as_());
508 }
509
510 pub fn append_buffer_with_offsets<P: NativePType + AsPrimitive<usize>>(
523 &mut self,
524 bytes: ByteBuffer,
525 offsets: &[P],
526 validity: &Mask,
527 ) {
528 assert_eq!(
529 offsets.len(),
530 validity.len() + 1,
531 "Must have one more offset than validity entries"
532 );
533 let first: usize = offsets[0].as_();
534 let last: usize = offsets[offsets.len() - 1].as_();
535 let bytes = bytes.slice(first..last);
536 self.append_buffer_views(&bytes, validity.len(), validity, |i| {
539 AsPrimitive::<usize>::as_(offsets[i + 1])
540 .wrapping_sub(AsPrimitive::<usize>::as_(offsets[i]))
541 });
542 }
543
544 fn append_buffer_views(
547 &mut self,
548 bytes: &ByteBuffer,
549 count: usize,
550 validity: &Mask,
551 len_at: impl Fn(usize) -> usize,
552 ) {
553 self.flush_in_progress();
554
555 if self.compacts_buffers() && bytes.len() <= MAX_BUFFER_LEN {
561 let referenced: usize = match validity.bit_buffer() {
562 AllOr::All => (0..count)
563 .map(&len_at)
564 .filter(|len| *len > BinaryView::MAX_INLINED_SIZE)
565 .sum(),
566 AllOr::None => 0,
567 AllOr::Some(b) => {
568 let mut sum = 0;
569 b.for_each_set_index(|idx| {
570 let len = len_at(idx);
571 if len > BinaryView::MAX_INLINED_SIZE {
572 sum += len;
573 }
574 });
575 sum
576 }
577 };
578 #[expect(clippy::cast_precision_loss)]
579 if (referenced as f64) < self.compaction_threshold * (bytes.len() as f64) {
580 return self
581 .append_buffer_views_rewritten(bytes, count, validity, len_at, referenced);
582 }
583 }
584
585 let start_index = self.completed.len();
586 let segments = extend_views(
587 &mut self.views_builder,
588 start_index,
589 MAX_BUFFER_LEN,
590 bytes,
591 count,
592 len_at,
593 );
594
595 let expected_completed_len = start_index as usize + segments.len();
596 self.completed.extend_from_slice_unchecked(&segments);
597 assert_eq!(
598 self.completed.len() as usize,
599 expected_completed_len,
600 "Some buffers already exist",
601 );
602
603 self.nulls.append_validity_mask(validity);
604 debug_assert_eq!(self.nulls.len(), self.views_builder.len());
605 }
606
607 fn append_buffer_views_rewritten(
611 &mut self,
612 bytes: &ByteBuffer,
613 count: usize,
614 validity: &Mask,
615 len_at: impl Fn(usize) -> usize,
616 referenced: usize,
617 ) {
618 let buf_index = self.completed.len();
619 let mut compact = ByteBufferMut::with_capacity(referenced);
620 self.views_builder.reserve(count);
621
622 let data = bytes.as_slice();
623 let mut offset = 0usize;
624 for (i, is_valid) in validity.iter().enumerate() {
625 let len = len_at(i);
626 let value = &data[offset..offset + len];
627 let view = if !is_valid {
628 BinaryView::empty_view()
629 } else if len > BinaryView::MAX_INLINED_SIZE {
630 #[expect(clippy::cast_possible_truncation)]
633 let view = BinaryView::make_view(value, buf_index, compact.len() as u32);
634 compact.extend_from_slice(value);
635 view
636 } else {
637 BinaryView::make_view(value, buf_index, 0)
638 };
639 self.views_builder.push(view);
640 offset += len;
641 }
642 assert_eq!(
643 offset,
644 data.len(),
645 "value lengths must describe the byte heap exactly"
646 );
647
648 if !compact.is_empty() {
649 let pushed_index = self.completed.push(compact.freeze());
650 assert_eq!(pushed_index, buf_index, "Buffer already exists");
651 }
652
653 self.nulls.append_validity_mask(validity);
654 debug_assert_eq!(self.nulls.len(), self.views_builder.len());
655 }
656
657 pub fn finish_into_varbinview(&mut self) -> VarBinViewArray {
659 self.flush_in_progress();
660 let buffers = std::mem::take(&mut self.completed);
661
662 assert_eq!(
663 self.views_builder.len(),
664 self.nulls.len(),
665 "View and validity length must match"
666 );
667
668 let validity = self.nulls.finish_with_nullability(self.dtype.nullability());
669
670 unsafe {
672 VarBinViewArray::new_unchecked(
673 std::mem::take(&mut self.views_builder).freeze(),
674 buffers.finish(),
675 self.dtype.clone(),
676 validity,
677 )
678 }
679 }
680
681 pub(crate) fn append_varbinview_array(
682 &mut self,
683 array: &VarBinViewArray,
684 ctx: &mut ExecutionCtx,
685 ) -> VortexResult<()> {
686 self.flush_in_progress();
687
688 let mask = array.varbinview_validity().execute_mask(array.len(), ctx)?;
689
690 self.nulls.append_validity_mask(&mask);
691
692 let view_adjustment =
693 self.completed
694 .extend_from_compaction(BuffersWithOffsets::from_array(
695 array,
696 self.compaction_threshold,
697 ctx,
698 ));
699
700 match view_adjustment {
701 ViewAdjustment::Precomputed(adjustment) => self.views_builder.extend_trusted(
702 array
703 .views()
704 .iter()
705 .map(|view| adjustment.adjust_view(view)),
706 ),
707 ViewAdjustment::Rewriting(adjustment) => match mask {
708 Mask::AllTrue(_) => {
709 for (idx, &view) in array.views().iter().enumerate() {
710 let new_view = self.push_view(view, &adjustment, array, idx);
711 self.views_builder.push(new_view);
712 }
713 }
714 Mask::AllFalse(_) => {
715 self.views_builder
716 .push_n(BinaryView::empty_view(), array.len());
717 }
718 Mask::Values(v) => {
719 let views = array.views();
720 let start = self.views_builder.len();
721 self.views_builder
722 .push_n(BinaryView::empty_view(), array.len());
723 v.bit_buffer().for_each_set_index(|idx| {
724 let new_view = self.push_view(views[idx], &adjustment, array, idx);
725 self.views_builder[start + idx] = new_view;
726 });
727 }
728 },
729 }
730
731 Ok(())
732 }
733}
734
735impl ArrayBuilder for VarBinViewBuilder {
736 fn as_any(&self) -> &dyn Any {
737 self
738 }
739
740 fn as_any_mut(&mut self) -> &mut dyn Any {
741 self
742 }
743
744 fn dtype(&self) -> &DType {
745 &self.dtype
746 }
747
748 fn len(&self) -> usize {
749 self.nulls.len()
750 }
751
752 fn append_zeros(&mut self, n: usize) {
753 self.views_builder.push_n(BinaryView::empty_view(), n);
754 self.nulls.append_n_non_nulls(n);
755 }
756
757 unsafe fn append_nulls_unchecked(&mut self, n: usize) {
758 self.views_builder.push_n(BinaryView::empty_view(), n);
759 self.nulls.append_n_nulls(n);
760 }
761
762 fn append_scalar(&mut self, scalar: &Scalar) -> VortexResult<()> {
763 vortex_ensure!(
764 scalar.dtype() == self.dtype(),
765 "VarBinViewBuilder expected scalar with dtype {}, got {}",
766 self.dtype(),
767 scalar.dtype()
768 );
769
770 match self.dtype() {
771 DType::Utf8(_) => match scalar.as_utf8().value() {
772 Some(value) => self.append_value(value),
773 None => self.append_null(),
774 },
775 DType::Binary(_) => match scalar.as_binary().value() {
776 Some(value) => self.append_value(value),
777 None => self.append_null(),
778 },
779 _ => vortex_bail!(
780 "VarBinViewBuilder can only handle Utf8 or Binary scalars, got {:?}",
781 scalar.dtype()
782 ),
783 }
784
785 Ok(())
786 }
787
788 fn reserve_exact(&mut self, additional: usize) {
789 self.views_builder.reserve(additional);
790 self.nulls.reserve_exact(additional);
791 }
792
793 unsafe fn set_validity_unchecked(&mut self, validity: Mask) {
794 self.nulls = LazyBitBufferBuilder::from_validity_mask(validity);
795 }
796
797 fn finish(&mut self) -> ArrayRef {
798 self.finish_into_varbinview().into_array()
799 }
800
801 fn finish_into_canonical(&mut self, _ctx: &mut ExecutionCtx) -> Canonical {
802 Canonical::VarBinView(self.finish_into_varbinview())
803 }
804}
805
806impl VarBinViewBuilder {
807 #[inline]
808 fn push_view(
809 &mut self,
810 view: BinaryView,
811 adjustment: &RewritingViewAdjustment,
812 array: &VarBinViewArray,
813 idx: usize,
814 ) -> BinaryView {
815 if view.is_inlined() {
816 view
817 } else if let Some(adjusted) = adjustment.adjust_view(&view) {
818 adjusted
819 } else {
820 let bytes = array.bytes_at(idx);
821 let (new_buf_idx, new_offset) = self.append_value_to_buffer(&bytes);
822 BinaryView::make_view(bytes.as_slice(), new_buf_idx, new_offset)
823 }
824 }
825}
826
827#[inline]
831fn remap_view(view: BinaryView, mapping: &[u32]) -> BinaryView {
832 if view.is_inlined() {
833 view
834 } else {
835 let view_ref = view.as_view();
836 view_ref
837 .with_buffer_and_offset(mapping[view_ref.buffer_index as usize], view_ref.offset)
838 .into()
839 }
840}
841
842enum CompactedSlot {
844 Kept { index: u32 },
846 Sliced { index: u32, shift: u32 },
848 Rewrite { source: ByteBuffer },
851}
852
853fn unmeasured_utilizations(buffers: &[ByteBuffer]) -> Vec<BufferUtilization> {
855 buffers
856 .iter()
857 .map(|buffer| {
858 BufferUtilization::zero(u32::try_from(buffer.len()).unwrap_or(u32::MAX))
862 })
863 .collect()
864}
865
866fn measure_view(utilizations: &mut [BufferUtilization], view: &BinaryView) {
868 if !view.is_inlined() {
869 let view_ref = view.as_view();
870 utilizations[view_ref.buffer_index as usize].add(view_ref.offset, view_ref.size);
871 }
872}
873
874pub enum CompletedBuffers {
875 Default(Vec<ByteBuffer>),
876 Deduplicated(DeduplicatedBuffers),
877}
878
879impl Default for CompletedBuffers {
880 fn default() -> Self {
881 Self::Default(Vec::new())
882 }
883}
884
885#[expect(clippy::cast_possible_truncation)]
887impl CompletedBuffers {
888 fn len(&self) -> u32 {
889 match self {
890 Self::Default(buffers) => buffers.len() as u32,
891 Self::Deduplicated(buffers) => buffers.len(),
892 }
893 }
894
895 fn push(&mut self, block: ByteBuffer) -> u32 {
898 match self {
899 Self::Default(buffers) => {
900 assert!(buffers.len() < u32::MAX as usize, "Too many blocks");
901 buffers.push(block);
902 self.len() - 1
903 }
904 Self::Deduplicated(buffers) => buffers.push(block),
905 }
906 }
907
908 fn extend_from_slice_unchecked(&mut self, buffers: &[ByteBuffer]) {
910 for buffer in buffers {
911 self.push(buffer.clone());
912 }
913 }
914
915 fn extend_from_compaction(&mut self, buffers: BuffersWithOffsets) -> ViewAdjustment {
916 match (self, buffers) {
917 (
918 Self::Default(completed_buffers),
919 BuffersWithOffsets::AllKept { buffers, offsets },
920 ) => {
921 let buffer_offset = completed_buffers.len() as u32;
922 completed_buffers.extend_from_slice(&buffers);
923 ViewAdjustment::shift(buffer_offset, offsets)
924 }
925 (
926 Self::Default(completed_buffers),
927 BuffersWithOffsets::SomeCompacted { buffers, offsets },
928 ) => {
929 let lookup = buffers
930 .iter()
931 .map(|maybe_buffer| {
932 maybe_buffer.as_ref().map(|buffer| {
933 completed_buffers.push(buffer.clone());
934 completed_buffers.len() as u32 - 1
935 })
936 })
937 .collect();
938 ViewAdjustment::rewriting(lookup, offsets)
939 }
940
941 (
942 Self::Deduplicated(completed_buffers),
943 BuffersWithOffsets::AllKept { buffers, offsets },
944 ) => {
945 let buffer_lookup = completed_buffers.extend_from_iter(buffers.iter().cloned());
946 ViewAdjustment::lookup(buffer_lookup, offsets)
947 }
948 (
949 Self::Deduplicated(completed_buffers),
950 BuffersWithOffsets::SomeCompacted { buffers, offsets },
951 ) => {
952 let buffer_lookup = completed_buffers.extend_from_option_slice(&buffers);
953 ViewAdjustment::rewriting(buffer_lookup, offsets)
954 }
955 }
956 }
957
958 fn finish(self) -> Arc<[ByteBuffer]> {
959 match self {
960 Self::Default(buffers) => Arc::from(buffers),
961 Self::Deduplicated(buffers) => buffers.finish(),
962 }
963 }
964}
965
966#[derive(Default)]
967pub struct DeduplicatedBuffers {
968 buffers: Vec<ByteBuffer>,
969 buffer_to_idx: HashMap<BufferId, u32>,
970}
971
972impl DeduplicatedBuffers {
973 #[expect(clippy::cast_possible_truncation)]
975 fn len(&self) -> u32 {
976 self.buffers.len() as u32
977 }
978
979 pub(crate) fn push(&mut self, block: ByteBuffer) -> u32 {
981 assert!(self.buffers.len() < u32::MAX as usize, "Too many blocks");
982
983 let initial_len = self.len();
984 let id = BufferId::from(&block);
985 match self.buffer_to_idx.entry(id) {
986 Entry::Occupied(idx) => *idx.get(),
987 Entry::Vacant(entry) => {
988 let idx = initial_len;
989 entry.insert(idx);
990 self.buffers.push(block);
991 idx
992 }
993 }
994 }
995
996 pub(crate) fn extend_from_option_slice(
997 &mut self,
998 buffers: &[Option<ByteBuffer>],
999 ) -> Vec<Option<u32>> {
1000 buffers
1001 .iter()
1002 .map(|buffer| buffer.as_ref().map(|buf| self.push(buf.clone())))
1003 .collect()
1004 }
1005
1006 pub(crate) fn extend_from_iter(
1007 &mut self,
1008 buffers: impl Iterator<Item = ByteBuffer>,
1009 ) -> Vec<u32> {
1010 buffers.map(|buffer| self.push(buffer)).collect()
1011 }
1012
1013 pub(crate) fn finish(self) -> Arc<[ByteBuffer]> {
1014 Arc::from(self.buffers)
1015 }
1016}
1017
1018#[derive(PartialEq, Eq, Hash)]
1019struct BufferId {
1020 ptr: usize,
1022 len: usize,
1023}
1024
1025impl BufferId {
1026 fn from(buffer: &ByteBuffer) -> Self {
1027 let slice = buffer.as_slice();
1028 Self {
1029 ptr: slice.as_ptr() as usize,
1030 len: slice.len(),
1031 }
1032 }
1033}
1034
1035#[derive(Debug, Clone)]
1036pub enum BufferGrowthStrategy {
1037 Fixed { size: u32 },
1039 Exponential { current_size: u32, max_size: u32 },
1041}
1042
1043impl Default for BufferGrowthStrategy {
1044 fn default() -> Self {
1045 Self::Exponential {
1046 current_size: 4 * 1024, max_size: 2 * 1024 * 1024, }
1049 }
1050}
1051
1052impl BufferGrowthStrategy {
1053 pub fn fixed(size: u32) -> Self {
1054 Self::Fixed { size }
1055 }
1056
1057 pub fn exponential(initial_size: u32, max_size: u32) -> Self {
1058 Self::Exponential {
1059 current_size: initial_size,
1060 max_size,
1061 }
1062 }
1063
1064 pub fn next_size(&mut self) -> u32 {
1066 match self {
1067 Self::Fixed { size } => *size,
1068 Self::Exponential {
1069 current_size,
1070 max_size,
1071 } => {
1072 let result = *current_size;
1073 if *current_size < *max_size {
1074 *current_size = current_size.saturating_mul(2).min(*max_size);
1075 }
1076 result
1077 }
1078 }
1079 }
1080}
1081
1082enum BuffersWithOffsets {
1083 AllKept {
1084 buffers: Arc<[ByteBuffer]>,
1085 offsets: Option<Vec<u32>>,
1086 },
1087 SomeCompacted {
1088 buffers: Vec<Option<ByteBuffer>>,
1089 offsets: Option<Vec<u32>>,
1090 },
1091}
1092
1093impl BuffersWithOffsets {
1094 pub fn from_array(
1095 array: &VarBinViewArray,
1096 compaction_threshold: f64,
1097 ctx: &mut ExecutionCtx,
1098 ) -> Self {
1099 if compaction_threshold == 0.0 {
1100 return Self::AllKept {
1101 buffers: Arc::from(
1102 array
1103 .data_buffers()
1104 .iter()
1105 .cloned()
1106 .map(|b| b.unwrap_host())
1107 .collect_vec(),
1108 ),
1109 offsets: None,
1110 };
1111 }
1112
1113 let buffer_utilizations = array
1114 .buffer_utilizations(ctx)
1115 .vortex_expect("buffer_utilizations in BuffersWithOffsets::from_array");
1116 let mut has_rewrite = false;
1117 let mut has_nonzero_offset = false;
1118 for utilization in buffer_utilizations.iter() {
1119 match compaction_strategy(utilization, compaction_threshold) {
1120 CompactionStrategy::KeepFull => continue,
1121 CompactionStrategy::Slice { .. } => has_nonzero_offset = true,
1122 CompactionStrategy::Rewrite => has_rewrite = true,
1123 }
1124 }
1125
1126 let buffers_with_offsets_iter = buffer_utilizations
1127 .iter()
1128 .zip(array.data_buffers().iter())
1129 .map(|(utilization, buffer)| {
1130 match compaction_strategy(utilization, compaction_threshold) {
1131 CompactionStrategy::KeepFull => (Some(buffer.as_host().clone()), 0),
1132 CompactionStrategy::Slice { start, end } => (
1133 Some(buffer.as_host().slice(start as usize..end as usize)),
1134 start,
1135 ),
1136 CompactionStrategy::Rewrite => (None, 0),
1137 }
1138 });
1139
1140 match (has_rewrite, has_nonzero_offset) {
1141 (false, false) => {
1143 let buffers: Vec<_> = buffers_with_offsets_iter
1144 .map(|(b, _)| b.vortex_expect("already checked for rewrite"))
1145 .collect();
1146 Self::AllKept {
1147 buffers: Arc::from(buffers),
1148 offsets: None,
1149 }
1150 }
1151 (true, false) => {
1153 let buffers: Vec<_> = buffers_with_offsets_iter.map(|(b, _)| b).collect();
1154 Self::SomeCompacted {
1155 buffers,
1156 offsets: None,
1157 }
1158 }
1159 (false, true) => {
1161 let (buffers, offsets): (Vec<_>, _) = buffers_with_offsets_iter
1162 .map(|(buffer, offset)| {
1163 (buffer.vortex_expect("already checked for rewrite"), offset)
1164 })
1165 .collect();
1166 Self::AllKept {
1167 buffers: Arc::from(buffers),
1168 offsets: Some(offsets),
1169 }
1170 }
1171 (true, true) => {
1173 let (buffers, offsets) = buffers_with_offsets_iter.collect();
1174 Self::SomeCompacted {
1175 buffers,
1176 offsets: Some(offsets),
1177 }
1178 }
1179 }
1180 }
1181}
1182
1183#[derive(Debug, Clone, Copy, PartialEq, Eq)]
1184enum CompactionStrategy {
1185 KeepFull,
1186 Slice {
1188 start: u32,
1189 end: u32,
1190 },
1191 Rewrite,
1193}
1194
1195fn compaction_strategy(
1196 buffer_utilization: &BufferUtilization,
1197 threshold: f64,
1198) -> CompactionStrategy {
1199 match buffer_utilization.overall_utilization() {
1200 0.0 => CompactionStrategy::Rewrite,
1202 utilised if utilised >= threshold => CompactionStrategy::KeepFull,
1203 _ if buffer_utilization.range_utilization() >= threshold => {
1204 let Range { start, end } = buffer_utilization.range();
1205 CompactionStrategy::Slice { start, end }
1206 }
1207 _ => CompactionStrategy::Rewrite,
1208 }
1209}
1210
1211enum ViewAdjustment {
1212 Precomputed(PrecomputedViewAdjustment),
1213 Rewriting(RewritingViewAdjustment),
1214}
1215
1216impl ViewAdjustment {
1217 fn shift(buffer_offset: u32, offsets: Option<Vec<u32>>) -> Self {
1218 Self::Precomputed(PrecomputedViewAdjustment::Shift {
1219 buffer_offset,
1220 offsets,
1221 })
1222 }
1223
1224 fn lookup(buffer_lookup: Vec<u32>, offsets: Option<Vec<u32>>) -> Self {
1225 Self::Precomputed(PrecomputedViewAdjustment::Lookup {
1226 buffer_lookup,
1227 offsets,
1228 })
1229 }
1230
1231 fn rewriting(buffer_lookup: Vec<Option<u32>>, offsets: Option<Vec<u32>>) -> Self {
1232 Self::Rewriting(RewritingViewAdjustment {
1233 buffer_lookup,
1234 offsets,
1235 })
1236 }
1237}
1238
1239enum PrecomputedViewAdjustment {
1241 Shift {
1242 buffer_offset: u32,
1243 offsets: Option<Vec<u32>>,
1244 },
1245 Lookup {
1246 buffer_lookup: Vec<u32>,
1247 offsets: Option<Vec<u32>>,
1248 },
1249}
1250
1251impl PrecomputedViewAdjustment {
1252 #[inline]
1253 fn adjust_view(&self, view: &BinaryView) -> BinaryView {
1254 if view.is_inlined() {
1255 return *view;
1256 }
1257 let view_ref = view.as_view();
1258 match self {
1259 Self::Shift {
1260 buffer_offset,
1261 offsets,
1262 } => {
1263 let b_idx = view_ref.buffer_index;
1264 let offset_shift = offsets
1265 .as_ref()
1266 .map(|o| o[b_idx as usize])
1267 .unwrap_or_default();
1268
1269 if view_ref.offset < offset_shift {
1272 return BinaryView::empty_view();
1273 }
1274
1275 view_ref
1276 .with_buffer_and_offset(b_idx + buffer_offset, view_ref.offset - offset_shift)
1277 }
1278 Self::Lookup {
1279 buffer_lookup,
1280 offsets,
1281 } => {
1282 let b_idx = view_ref.buffer_index;
1283 let buffer = buffer_lookup[b_idx as usize];
1284 let offset_shift = offsets
1285 .as_ref()
1286 .map(|o| o[b_idx as usize])
1287 .unwrap_or_default();
1288
1289 if view_ref.offset < offset_shift {
1292 return BinaryView::empty_view();
1293 }
1294
1295 view_ref.with_buffer_and_offset(buffer, view_ref.offset - offset_shift)
1296 }
1297 }
1298 .into()
1299 }
1300}
1301
1302struct RewritingViewAdjustment {
1303 buffer_lookup: Vec<Option<u32>>,
1304 offsets: Option<Vec<u32>>,
1305}
1306
1307impl RewritingViewAdjustment {
1308 #[inline]
1311 fn adjust_view(&self, view: &BinaryView) -> Option<BinaryView> {
1312 if view.is_inlined() {
1313 return Some(*view);
1314 }
1315
1316 let view_ref = view.as_view();
1317 self.buffer_lookup[view_ref.buffer_index as usize].map(|buffer| {
1318 let offset_shift = self
1319 .offsets
1320 .as_ref()
1321 .map(|o| o[view_ref.buffer_index as usize])
1322 .unwrap_or_default();
1323 view_ref
1324 .with_buffer_and_offset(buffer, view_ref.offset - offset_shift)
1325 .into()
1326 })
1327 }
1328}
1329
1330#[cfg(test)]
1331mod tests {
1332 use vortex_buffer::ByteBuffer;
1333 use vortex_error::VortexResult;
1334 use vortex_mask::Mask;
1335
1336 use crate::IntoArray;
1337 use crate::VortexSessionExecute;
1338 use crate::array_session;
1339 use crate::assert_arrays_eq;
1340 use crate::builders::ArrayBuilder;
1341 use crate::builders::VarBinViewBuilder;
1342 use crate::builders::varbinview::VarBinViewArray;
1343 use crate::dtype::DType;
1344 use crate::dtype::Nullability;
1345
1346 const LONG: &str = "a value that is far too long to inline";
1348
1349 #[test]
1352 fn test_append_buffer_with_lengths() {
1353 let mut ctx = array_session().create_execution_ctx();
1354 let mut builder = VarBinViewBuilder::with_capacity(DType::Utf8(Nullability::Nullable), 8);
1355
1356 builder.append_value(LONG);
1358
1359 let heap = ByteBuffer::copy_from([LONG.as_bytes(), b"", b"tiny"].concat());
1360 let heap_ptr = heap.as_ptr();
1361 let lengths = [u32::try_from(LONG.len()).unwrap(), 0, 4];
1362 builder.append_buffer_with_lengths(heap, &lengths, &Mask::from_iter([true, false, true]));
1363
1364 builder.append_value("tail");
1365
1366 let actual = builder.finish_into_varbinview();
1367 assert_eq!(actual.data_buffers()[1].as_host().as_ptr(), heap_ptr);
1369
1370 let expected = <VarBinViewArray as FromIterator<_>>::from_iter([
1371 Some(LONG),
1372 Some(LONG),
1373 None,
1374 Some("tiny"),
1375 Some("tail"),
1376 ]);
1377 assert_arrays_eq!(actual, expected, &mut ctx);
1378 }
1379
1380 #[test]
1382 fn test_append_buffer_with_offsets() {
1383 let mut ctx = array_session().create_execution_ctx();
1384 let mut builder = VarBinViewBuilder::with_capacity(DType::Utf8(Nullability::Nullable), 8);
1385
1386 let heap = ByteBuffer::copy_from(format!("..{LONG}tiny!!"));
1387 let long_len = u32::try_from(LONG.len()).unwrap();
1388 let offsets = [2u32, 2 + long_len, 2 + long_len, 2 + long_len + 4];
1389 builder.append_buffer_with_offsets(
1390 heap.clone(),
1391 &offsets,
1392 &Mask::from_iter([true, false, true]),
1393 );
1394
1395 let actual = builder.finish_into_varbinview();
1396 assert_eq!(actual.data_buffers()[0].as_host().as_ptr(), unsafe {
1399 heap.as_ptr().add(2)
1400 });
1401 assert_eq!(
1402 actual.data_buffers()[0].len(),
1403 LONG.len() + 4,
1404 "only the referenced range must be adopted"
1405 );
1406
1407 let expected =
1408 <VarBinViewArray as FromIterator<_>>::from_iter([Some(LONG), None, Some("tiny")]);
1409 assert_arrays_eq!(actual, expected, &mut ctx);
1410 }
1411
1412 #[test]
1415 fn test_append_buffer_with_lengths_compacts_underutilized_heap() {
1416 let mut ctx = array_session().create_execution_ctx();
1417 let mut builder =
1418 VarBinViewBuilder::with_compaction(DType::Utf8(Nullability::Nullable), 4, 1.0);
1419
1420 let heap = ByteBuffer::copy_from([b"short".as_slice(), LONG.as_bytes(), b"tiny"].concat());
1421 let lengths = [5u32, u32::try_from(LONG.len()).unwrap(), 4];
1422 builder.append_buffer_with_lengths(heap, &lengths, &Mask::new_true(3));
1423
1424 let actual = builder.finish_into_varbinview();
1425 assert_eq!(actual.data_buffers().len(), 1);
1426 assert_eq!(
1427 actual.data_buffers()[0].len(),
1428 LONG.len(),
1429 "the compact buffer must hold only the non-inlined values"
1430 );
1431
1432 let expected = <VarBinViewArray as FromIterator<_>>::from_iter([
1433 Some("short"),
1434 Some(LONG),
1435 Some("tiny"),
1436 ]);
1437 assert_arrays_eq!(actual, expected, &mut ctx);
1438 }
1439
1440 #[test]
1443 fn test_append_buffer_with_lengths_compaction_skips_null_bytes() {
1444 let mut ctx = array_session().create_execution_ctx();
1445 let mut builder =
1446 VarBinViewBuilder::with_compaction(DType::Utf8(Nullability::Nullable), 2, 1.0);
1447
1448 let heap = ByteBuffer::copy_from([LONG.as_bytes(), LONG.as_bytes()].concat());
1449 let lengths = [u32::try_from(LONG.len()).unwrap(); 2];
1450 builder.append_buffer_with_lengths(heap, &lengths, &Mask::from_iter([false, true]));
1451
1452 let actual = builder.finish_into_varbinview();
1453 assert_eq!(actual.data_buffers().len(), 1);
1454 assert_eq!(
1455 actual.data_buffers()[0].len(),
1456 LONG.len(),
1457 "the compact buffer must omit bytes belonging to null rows"
1458 );
1459
1460 let expected = <VarBinViewArray as FromIterator<_>>::from_iter([None::<&str>, Some(LONG)]);
1461 assert_arrays_eq!(actual, expected, &mut ctx);
1462 }
1463
1464 #[test]
1467 fn test_push_buffers_deduplicates() {
1468 let mut builder =
1469 VarBinViewBuilder::with_buffer_deduplication(DType::Utf8(Nullability::Nullable), 8);
1470
1471 let first = ByteBuffer::copy_from(LONG);
1472 let second = ByteBuffer::copy_from("another value far too long to inline");
1473
1474 assert_eq!(
1475 builder.push_buffers([first.clone(), second.clone()]),
1476 [0, 1]
1477 );
1478 assert_eq!(builder.push_buffers([second, first]), [1, 0]);
1479 assert_eq!(builder.completed_block_count(), 2);
1480 }
1481
1482 #[test]
1485 fn test_append_views_gathered() {
1486 let mut ctx = array_session().create_execution_ctx();
1487 let dictionary = <VarBinViewArray as FromIterator<_>>::from_iter([
1488 Some("tiny"),
1489 Some(LONG),
1490 Some("small"),
1491 ]);
1492 let buffers = dictionary
1493 .data_buffers()
1494 .iter()
1495 .map(|buffer| buffer.as_host().clone())
1496 .collect::<Vec<_>>();
1497
1498 let mut builder = VarBinViewBuilder::with_capacity(DType::Utf8(Nullability::Nullable), 8);
1499 builder.append_value(LONG);
1501
1502 let codes: [usize; 4] = [1, 0, usize::MAX, 2];
1503 let views = dictionary.views();
1504 builder.append_views_gathered(
1505 buffers,
1506 views,
1507 &Mask::from_iter([true, true, false, true]),
1508 |row| codes[row],
1510 );
1511
1512 let expected = <VarBinViewArray as FromIterator<_>>::from_iter([
1513 Some(LONG),
1514 Some(LONG),
1515 Some("tiny"),
1516 None,
1517 Some("small"),
1518 ]);
1519 assert_arrays_eq!(builder.finish_into_varbinview(), expected, &mut ctx);
1520 }
1521
1522 #[test]
1525 fn test_append_views_scattered() {
1526 use crate::arrays::varbinview::build_views::BinaryView;
1527
1528 let mut ctx = array_session().create_execution_ctx();
1529 let patch_values =
1530 <VarBinViewArray as FromIterator<_>>::from_iter([Some(LONG), Some("tiny")]);
1531 let mut buffers = patch_values
1532 .data_buffers()
1533 .iter()
1534 .map(|buffer| buffer.as_host().clone())
1535 .collect::<Vec<_>>();
1536
1537 let fill_bytes = ByteBuffer::copy_from("a fill value too long to inline");
1538 buffers.push(fill_bytes.clone());
1539 let fill = BinaryView::make_view(
1540 fill_bytes.as_slice(),
1541 u32::try_from(buffers.len() - 1).unwrap(),
1542 0,
1543 );
1544
1545 let mut builder = VarBinViewBuilder::with_capacity(DType::Utf8(Nullability::Nullable), 8);
1546 let views = patch_values.views();
1547 builder.append_views_scattered(
1548 buffers,
1549 5,
1550 fill,
1551 [(1usize, views[0]), (3usize, views[1])].into_iter(),
1552 &Mask::from_iter([true, true, false, true, true]),
1553 );
1554
1555 let expected = <VarBinViewArray as FromIterator<_>>::from_iter([
1556 Some("a fill value too long to inline"),
1557 Some(LONG),
1558 None,
1559 Some("tiny"),
1560 Some("a fill value too long to inline"),
1561 ]);
1562 assert_arrays_eq!(builder.finish_into_varbinview(), expected, &mut ctx);
1563 }
1564
1565 #[test]
1569 fn test_append_views_gathered_compacts_unreferenced_values() {
1570 const DEAD: &str = "a dead value nobody gathers, far too long to inline";
1571 const OTHER: &str = "another value far too long to inline";
1572
1573 let mut ctx = array_session().create_execution_ctx();
1574 let dictionary =
1575 <VarBinViewArray as FromIterator<_>>::from_iter([Some(LONG), Some(DEAD), Some(OTHER)]);
1576 assert_eq!(dictionary.data_buffers().len(), 1);
1577 let buffers = dictionary
1578 .data_buffers()
1579 .iter()
1580 .map(|buffer| buffer.as_host().clone())
1581 .collect::<Vec<_>>();
1582
1583 let mut builder =
1584 VarBinViewBuilder::with_compaction(DType::Utf8(Nullability::Nullable), 8, 1.0);
1585 let codes: [usize; 5] = [2, 0, usize::MAX, 2, 0];
1586 builder.append_views_gathered(
1587 buffers,
1588 dictionary.views(),
1589 &Mask::from_iter([true, true, false, true, true]),
1590 |row| codes[row],
1591 );
1592
1593 let actual = builder.finish_into_varbinview();
1594 assert_eq!(actual.data_buffers().len(), 1);
1595 assert_eq!(
1596 actual.data_buffers()[0].len(),
1597 LONG.len() + OTHER.len(),
1598 "the rewritten buffer must hold each referenced value exactly once"
1599 );
1600
1601 let expected = <VarBinViewArray as FromIterator<_>>::from_iter([
1602 Some(OTHER),
1603 Some(LONG),
1604 None,
1605 Some(OTHER),
1606 Some(LONG),
1607 ]);
1608 assert_arrays_eq!(actual, expected, &mut ctx);
1609 }
1610
1611 #[test]
1614 fn test_append_views_gathered_slices_contiguous_range() {
1615 const TAIL: &str = "the referenced tail value, too long to inline";
1616
1617 let mut ctx = array_session().create_execution_ctx();
1618 let dictionary = <VarBinViewArray as FromIterator<_>>::from_iter([Some(LONG), Some(TAIL)]);
1619 let heap = dictionary.data_buffers()[0].as_host().clone();
1620
1621 let mut builder =
1622 VarBinViewBuilder::with_compaction(DType::Utf8(Nullability::Nullable), 8, 1.0);
1623 builder.append_views_gathered(
1624 [heap.clone()],
1625 dictionary.views(),
1626 &Mask::new_true(2),
1627 |_| 1,
1628 );
1629
1630 let actual = builder.finish_into_varbinview();
1631 assert_eq!(actual.data_buffers().len(), 1);
1632 assert_eq!(actual.data_buffers()[0].len(), TAIL.len());
1633 assert_eq!(actual.data_buffers()[0].as_host().as_ptr(), unsafe {
1635 heap.as_ptr().add(LONG.len())
1636 });
1637
1638 let expected = <VarBinViewArray as FromIterator<_>>::from_iter([Some(TAIL), Some(TAIL)]);
1639 assert_arrays_eq!(actual, expected, &mut ctx);
1640 }
1641
1642 #[test]
1644 fn test_append_views_gathered_keeps_utilized_buffer() {
1645 const SHORTER: &str = "a shorter long value";
1646
1647 let mut ctx = array_session().create_execution_ctx();
1648 let dictionary =
1649 <VarBinViewArray as FromIterator<_>>::from_iter([Some(LONG), Some(SHORTER)]);
1650 let heap = dictionary.data_buffers()[0].as_host().clone();
1651
1652 let mut builder =
1653 VarBinViewBuilder::with_compaction(DType::Utf8(Nullability::Nullable), 8, 0.5);
1654 builder.append_views_gathered(
1655 [heap.clone()],
1656 dictionary.views(),
1657 &Mask::new_true(1),
1658 |_| 0,
1659 );
1660
1661 let actual = builder.finish_into_varbinview();
1662 assert_eq!(actual.data_buffers().len(), 1);
1663 assert_eq!(actual.data_buffers()[0].as_host().as_ptr(), heap.as_ptr());
1664 assert_eq!(actual.data_buffers()[0].len(), heap.len());
1665
1666 let expected = <VarBinViewArray as FromIterator<_>>::from_iter([Some(LONG)]);
1667 assert_arrays_eq!(actual, expected, &mut ctx);
1668 }
1669
1670 #[test]
1673 fn test_append_views_gathered_all_null_drops_buffers() {
1674 let mut ctx = array_session().create_execution_ctx();
1675 let dictionary = <VarBinViewArray as FromIterator<_>>::from_iter([Some(LONG)]);
1676 let buffers = dictionary
1677 .data_buffers()
1678 .iter()
1679 .map(|buffer| buffer.as_host().clone())
1680 .collect::<Vec<_>>();
1681
1682 let mut builder =
1683 VarBinViewBuilder::with_compaction(DType::Utf8(Nullability::Nullable), 8, 1.0);
1684 builder.append_views_gathered(buffers, dictionary.views(), &Mask::new_false(3), |_| {
1685 usize::MAX
1686 });
1687
1688 let actual = builder.finish_into_varbinview();
1689 assert!(
1690 actual.data_buffers().is_empty(),
1691 "an all-null gather must not retain any buffers"
1692 );
1693
1694 let expected = <VarBinViewArray as FromIterator<_>>::from_iter([None::<&str>, None, None]);
1695 assert_arrays_eq!(actual, expected, &mut ctx);
1696 }
1697
1698 #[test]
1700 fn test_append_buffer_with_lengths_drops_fully_inlined_heap() {
1701 let mut ctx = array_session().create_execution_ctx();
1702 let mut builder =
1703 VarBinViewBuilder::with_compaction(DType::Utf8(Nullability::Nullable), 4, 1.0);
1704
1705 let heap = ByteBuffer::copy_from(b"shorttinysmall".as_slice());
1706 builder.append_buffer_with_lengths(heap, &[5u32, 4, 5], &Mask::new_true(3));
1707
1708 let actual = builder.finish_into_varbinview();
1709 assert!(
1710 actual.data_buffers().is_empty(),
1711 "a fully-inlined append must not retain any value bytes"
1712 );
1713
1714 let expected = <VarBinViewArray as FromIterator<_>>::from_iter([
1715 Some("short"),
1716 Some("tiny"),
1717 Some("small"),
1718 ]);
1719 assert_arrays_eq!(actual, expected, &mut ctx);
1720 }
1721
1722 #[test]
1723 fn test_utf8_builder() {
1724 let mut ctx = array_session().create_execution_ctx();
1725 let mut builder = VarBinViewBuilder::with_capacity(DType::Utf8(Nullability::Nullable), 10);
1726
1727 builder.append_value("Hello");
1728 builder.append_null();
1729 builder.append_value("World");
1730
1731 builder.append_nulls(2);
1732
1733 builder.append_zeros(2);
1734 builder.append_value("test");
1735
1736 let actual = builder.finish();
1737 let expected = <VarBinViewArray as FromIterator<_>>::from_iter([
1738 Some("Hello"),
1739 None,
1740 Some("World"),
1741 None,
1742 None,
1743 Some(""),
1744 Some(""),
1745 Some("test"),
1746 ]);
1747 assert_arrays_eq!(actual, expected, &mut ctx);
1748 }
1749
1750 #[test]
1751 fn test_utf8_builder_with_extend() {
1752 let mut ctx = array_session().create_execution_ctx();
1753 let array = {
1754 let mut builder =
1755 VarBinViewBuilder::with_capacity(DType::Utf8(Nullability::Nullable), 10);
1756 builder.append_null();
1757 builder.append_value("Hello2");
1758 builder.finish()
1759 };
1760 let mut builder = VarBinViewBuilder::with_capacity(DType::Utf8(Nullability::Nullable), 10);
1761
1762 builder.append_value("Hello1");
1763 array.append_to_builder(&mut builder, &mut ctx).unwrap();
1764 builder.append_nulls(2);
1765 builder.append_value("Hello3");
1766
1767 let actual = builder.finish_into_canonical(&mut ctx);
1768 let expected = <VarBinViewArray as FromIterator<_>>::from_iter([
1769 Some("Hello1"),
1770 None,
1771 Some("Hello2"),
1772 None,
1773 None,
1774 Some("Hello3"),
1775 ]);
1776 assert_arrays_eq!(actual.into_array(), expected.into_array(), &mut ctx);
1777 }
1778
1779 #[test]
1780 fn test_buffer_deduplication() -> VortexResult<()> {
1781 let array = {
1782 let mut builder =
1783 VarBinViewBuilder::with_capacity(DType::Utf8(Nullability::Nullable), 10);
1784 builder.append_value("This is a long string that should not be inlined");
1785 builder.append_value("short string");
1786 builder.finish_into_varbinview()
1787 };
1788
1789 assert_eq!(array.data_buffers().len(), 1);
1790 let mut builder =
1791 VarBinViewBuilder::with_buffer_deduplication(DType::Utf8(Nullability::Nullable), 10);
1792
1793 let mut ctx = array_session().create_execution_ctx();
1794
1795 array.append_to_builder(&mut builder, &mut ctx)?;
1796 assert_eq!(builder.completed_block_count(), 1);
1797
1798 array
1799 .slice(1..2)?
1800 .append_to_builder(&mut builder, &mut ctx)?;
1801 array
1802 .slice(0..1)?
1803 .append_to_builder(&mut builder, &mut ctx)?;
1804 assert_eq!(builder.completed_block_count(), 1);
1805
1806 let array2 = {
1807 let mut builder =
1808 VarBinViewBuilder::with_capacity(DType::Utf8(Nullability::Nullable), 10);
1809 builder.append_value("This is a long string that should not be inlined");
1810 builder.finish_into_varbinview()
1811 };
1812
1813 array2.append_to_builder(&mut builder, &mut ctx)?;
1814 assert_eq!(builder.completed_block_count(), 2);
1815
1816 array
1817 .slice(0..1)?
1818 .append_to_builder(&mut builder, &mut ctx)?;
1819 array2
1820 .slice(0..1)?
1821 .append_to_builder(&mut builder, &mut ctx)?;
1822 assert_eq!(builder.completed_block_count(), 2);
1823 Ok(())
1824 }
1825
1826 #[test]
1827 fn test_append_scalar() {
1828 let mut ctx = array_session().create_execution_ctx();
1829 use crate::scalar::Scalar;
1830
1831 let mut utf8_builder =
1833 VarBinViewBuilder::with_capacity(DType::Utf8(Nullability::Nullable), 10);
1834
1835 let utf8_scalar1 = Scalar::utf8("hello", Nullability::Nullable);
1837 utf8_builder.append_scalar(&utf8_scalar1).unwrap();
1838
1839 let utf8_scalar2 = Scalar::utf8("world", Nullability::Nullable);
1841 utf8_builder.append_scalar(&utf8_scalar2).unwrap();
1842
1843 let null_scalar = Scalar::null(DType::Utf8(Nullability::Nullable));
1845 utf8_builder.append_scalar(&null_scalar).unwrap();
1846
1847 let array = utf8_builder.finish();
1848 let expected =
1849 <VarBinViewArray as FromIterator<_>>::from_iter([Some("hello"), Some("world"), None]);
1850 assert_arrays_eq!(&array, &expected, &mut ctx);
1851
1852 let mut binary_builder =
1854 VarBinViewBuilder::with_capacity(DType::Binary(Nullability::Nullable), 10);
1855
1856 let binary_scalar = Scalar::binary(vec![1u8, 2, 3], Nullability::Nullable);
1857 binary_builder.append_scalar(&binary_scalar).unwrap();
1858
1859 let binary_null = Scalar::null(DType::Binary(Nullability::Nullable));
1860 binary_builder.append_scalar(&binary_null).unwrap();
1861
1862 let binary_array = binary_builder.finish();
1863 let expected =
1864 <VarBinViewArray as FromIterator<_>>::from_iter([Some(vec![1u8, 2, 3]), None]);
1865 assert_arrays_eq!(&binary_array, &expected, &mut ctx);
1866
1867 let mut builder =
1869 VarBinViewBuilder::with_capacity(DType::Utf8(Nullability::NonNullable), 10);
1870 let wrong_scalar = Scalar::from(42i32);
1871 assert!(builder.append_scalar(&wrong_scalar).is_err());
1872 }
1873
1874 #[test]
1875 fn test_buffer_growth_strategies() {
1876 use super::BufferGrowthStrategy;
1877
1878 let mut strategy = BufferGrowthStrategy::fixed(1024);
1880
1881 assert_eq!(strategy.next_size(), 1024);
1883 assert_eq!(strategy.next_size(), 1024);
1884 assert_eq!(strategy.next_size(), 1024);
1885
1886 let mut strategy = BufferGrowthStrategy::exponential(1024, 8192);
1888
1889 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); }
1896
1897 #[test]
1898 fn test_large_value_allocation() {
1899 use super::BufferGrowthStrategy;
1900 use super::VarBinViewBuilder;
1901
1902 let mut builder = VarBinViewBuilder::new(
1903 DType::Binary(Nullability::Nullable),
1904 10,
1905 Default::default(),
1906 BufferGrowthStrategy::exponential(1024, 4096),
1907 0.0,
1908 );
1909
1910 let large_value = vec![0u8; 8192];
1912
1913 builder.append_value(&large_value);
1915
1916 let array = builder.finish_into_varbinview();
1917 assert_eq!(array.len(), 1);
1918
1919 let retrieved = array
1921 .execute_scalar(0, &mut array_session().create_execution_ctx())
1922 .unwrap()
1923 .as_binary()
1924 .value()
1925 .cloned()
1926 .unwrap();
1927 assert_eq!(retrieved.len(), 8192);
1928 assert_eq!(retrieved.as_slice(), &large_value);
1929 }
1930}