Skip to main content

vortex_array/builders/
varbinview.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use 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
40/// The builder for building a [`VarBinViewArray`].
41pub 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    /// Appends a value to the builder.
116    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    /// Appends `n` copies of `value` as non-null entries.
122    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    /// append a non inlined value to self.in_progress.
164    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    /// Whether this builder compacts the data buffers it is handed. The lengths-driven and
197    /// gather appends use this to gate their utilization measurement; the buffer-adopting escape
198    /// hatches ([`append_views_built_at`](Self::append_views_built_at) and
199    /// [`push_buffers`](Self::push_buffers)) always bypass it.
200    fn compacts_buffers(&self) -> bool {
201        self.compaction_threshold > 0.0
202    }
203
204    /// Adopts the buffers and views that `build` produces against the index its first buffer
205    /// will land at.
206    ///
207    /// The builder flushes its staged bytes, then hands `build` the index the next data buffer
208    /// will occupy; `build` returns data buffers — which land contiguously from that index — and
209    /// one view per entry of `validity`, already referencing them. This is the escape hatch for
210    /// an encoding that only discovers its views while walking its own byte format (e.g.
211    /// length-prefixed frames), where the lengths-driven appends cannot apply; keeping the
212    /// numbering inside this call is what makes the views come out right without a rebase pass.
213    ///
214    /// # Warning
215    ///
216    /// This method does not check utilization of the returned buffers. `build` must return
217    /// buffers that are fully utilized by its views.
218    ///
219    /// # Panics
220    ///
221    /// Panics if `build` returns a different view count than `validity.len()`, or if this
222    /// builder deduplicates buffers and already holds one of the returned buffers.
223    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    /// Adopts `buffers` as completed data buffers, returning the index each landed at.
253    ///
254    /// Views appended afterwards (e.g. via [`append_views_gathered`](Self::append_views_gathered)
255    /// or [`append_views_scattered`](Self::append_views_scattered)) reference values through the
256    /// returned indices. A deduplicating builder returns the existing index for a buffer it
257    /// already holds, so repeated appends over shared storage — chunks gathered through one
258    /// dictionary, slices of one array — adopt it once.
259    ///
260    /// # Warning
261    ///
262    /// Buffers are taken as they are, without utilization measurement; like
263    /// [`append_views_built_at`](Self::append_views_built_at), this bypasses any compaction the
264    /// builder was configured for. Scattered appends deliberately have these semantics;
265    /// compacting paths that measure utilization use
266    /// [`push_buffers_compacted`](Self::push_buffers_compacted) instead.
267    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    /// The compacting counterpart of [`push_buffers`](Self::push_buffers): adopts each buffer
276    /// according to the [`CompactionStrategy`] its measured utilization earns — whole, sliced to
277    /// its referenced range, or not at all. Views into the buffers must then be rebased through
278    /// [`remap_view_compacted`](Self::remap_view_compacted), which routes views into unadopted
279    /// buffers through the builder's own storage.
280    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    /// [`remap_view`] against [`CompactedSlot`]s: views into kept or sliced buffers are rebased
307    /// onto the index those landed at, and a view into a rewritten buffer has its bytes copied
308    /// into the builder's own storage. Inlined views pass through unchanged.
309    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    /// Appends `validity.len()` values, gathering each valid row's view from `views` through the
330    /// index `view_at` returns for it.
331    ///
332    /// `buffers` — the data buffers the views reference, in the numbering the views use — are
333    /// adopted through the builder's buffer storage, and every gathered view is rebased onto
334    /// the indices they land at as it is written, so the whole append is one view per row with no
335    /// byte copy and no intermediate array. `view_at` is only called for valid rows; null rows get
336    /// an empty view.
337    ///
338    /// When the builder is configured to compact buffers, utilization is measured from the views
339    /// the gather actually references — duplicate indices count once — and each buffer is
340    /// adopted, sliced, or rewritten accordingly, so callers never need a
341    /// canonicalize-and-compact fallback.
342    ///
343    /// # Panics
344    ///
345    /// Panics if `view_at` returns an index out of bounds of `views`, or if a gathered view
346    /// references a buffer index outside `buffers`.
347    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    /// The compacting arm of [`append_views_gathered`](Self::append_views_gathered): marks the
388    /// source views the gather references, measures each buffer's utilization from just those —
389    /// so entries gathered by many rows count once — and adopts the buffers through
390    /// [`push_buffers_compacted`](Self::push_buffers_compacted). Each referenced source view is
391    /// remapped once, so rows sharing an entry of a rewritten buffer share one copy of its bytes.
392    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    /// Appends `len` rows that are `fill` everywhere except at the given patch rows.
442    ///
443    /// `buffers` — the data buffers `fill` and the patch views reference, in the numbering they
444    /// use — are adopted through the builder's buffer storage and the views are rebased onto
445    /// the indices they land at. `patches` yields `(row, view)` pairs with rows below `len`; the
446    /// fill rows cost one bulk view fill and each patch one view write, with no byte copy.
447    /// Validity is appended exactly as given — invalid rows keep whichever view they got.
448    ///
449    /// This append deliberately adopts the supplied buffers as-is, even when the builder is
450    /// configured for compaction. Callers should use it only when retaining those buffers is
451    /// acceptable.
452    ///
453    /// # Panics
454    ///
455    /// Panics if `validity` is not `len` long, a patch row is out of bounds, or a view references
456    /// a buffer index outside `buffers`.
457    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    /// Appends values laid end-to-end in `bytes`, one per entry of `lengths`.
480    ///
481    /// The builder adopts `bytes` as a data buffer without copying it (splitting it only past the
482    /// `u32` view-offset limit) and builds the views directly into its own storage, so the whole
483    /// append costs one view per row. Every length is consumed, including those for null rows, so
484    /// the lengths must describe `bytes` exactly; bytes belonging to null rows are not retained
485    /// when compaction rewrites the buffer.
486    ///
487    /// When the builder is configured to compact buffers, the utilization is measured from the
488    /// lengths alone — only values too long to inline reference the buffer — and a heap below
489    /// the threshold is rewritten to just those values instead of adopted, so callers never need
490    /// a canonicalize-and-compact fallback.
491    ///
492    /// # Panics
493    ///
494    /// Panics if `lengths` and `validity` disagree in length, if the lengths do not describe
495    /// `bytes` exactly, or if this builder deduplicates buffers and already holds `bytes`.
496    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    /// [`append_buffer_with_lengths`](Self::append_buffer_with_lengths) for values described by
511    /// an offsets buffer instead of lengths.
512    ///
513    /// `offsets` are absolute positions into `bytes` — the layout a
514    /// [`VarBinArray`](crate::arrays::VarBinArray) stores — so there is one more offset than there
515    /// are values, and only the `offsets[0]..offsets[last]` range of `bytes` is adopted, again
516    /// without copying.
517    ///
518    /// # Panics
519    ///
520    /// Panics if `offsets` does not hold exactly one more entry than `validity`, or if the offsets
521    /// are not monotonically non-decreasing positions within `bytes`.
522    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        // Wrapping keeps corrupt non-monotonic offsets from panicking on the subtraction itself;
537        // the wrapped length then fails the in-bounds checks of the view-building loop.
538        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    /// Shared tail of the bulk buffer appends: builds the views straight into the builder's views
545    /// storage, then adopts the buffer segments and the validity.
546    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        // A compacting builder measures utilization before adopting the buffer. Only values too
556        // long to inline reference the heap, so the measurement is one pass over the lengths and
557        // never touches the bytes. Heaps past the single-buffer limit are adopted as they are:
558        // they roll over into multiple buffers, and per-segment accounting is not worth the rare
559        // >2GiB case.
560        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    /// The under-utilized arm of [`append_buffer_views`](Self::append_buffer_views): copies only
608    /// the values that actually reference the heap into a compact buffer, sized `referenced`,
609    /// instead of adopting the whole heap. A fully-inlined append pushes no buffer at all.
610    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                // In `u32` range: `referenced <= bytes.len() <= MAX_BUFFER_LEN` (checked by the
631                // caller), and `compact` never grows past `referenced`.
632                #[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    /// Finishes the builder directly into a [`VarBinViewArray`].
658    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        // SAFETY: the builder methods check safety at each step.
671        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/// Rebases a view built against a local buffer numbering onto the builder indices those buffers
828/// landed at, i.e. `mapping[i]` is where the caller's buffer `i` went. Inlined views carry no
829/// buffer reference and pass through unchanged.
830#[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
842/// Where a caller's buffer went under `VarBinViewBuilder::push_buffers_compacted`.
843enum CompactedSlot {
844    /// Adopted whole at this index.
845    Kept { index: u32 },
846    /// Adopted at this index as the slice starting `shift` bytes in, so view offsets shift down.
847    Sliced { index: u32, shift: u32 },
848    /// Not adopted: each referenced view's bytes are copied out of the source buffer into the
849    /// builder's own storage.
850    Rewrite { source: ByteBuffer },
851}
852
853/// One unmeasured [`BufferUtilization`] per buffer, ready for [`measure_view`] passes.
854fn unmeasured_utilizations(buffers: &[ByteBuffer]) -> Vec<BufferUtilization> {
855    buffers
856        .iter()
857        .map(|buffer| {
858            // Views address at most `u32` offsets, so measuring an oversized buffer against the
859            // saturated length under-reports utilization, which can only compact harder — and its
860            // unaddressable tail is dead weight worth compacting anyway.
861            BufferUtilization::zero(u32::try_from(buffer.len()).unwrap_or(u32::MAX))
862        })
863        .collect()
864}
865
866/// Counts `view`'s bytes against the buffer it references; inlined views reference none.
867fn 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// Self::push enforces len < u32::max
886#[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    /// Push a new block, returning the index it landed at (or, when deduplicating, the index of
896    /// the identical block already held).
897    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    /// Does not compact buffers, bypasses utilization checks.
909    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    // Self::push enforces len < u32::max
974    #[expect(clippy::cast_possible_truncation)]
975    fn len(&self) -> u32 {
976        self.buffers.len() as u32
977    }
978
979    /// Push a new block if not seen before. Returns the idx of the block.
980    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    // *const u8 stored as usize for `Send`
1021    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    /// Use a fixed buffer size for all allocations.
1038    Fixed { size: u32 },
1039    /// Use exponential growth starting from initial_size, doubling until max_size.
1040    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,    // 4KB starting size
1047            max_size: 2 * 1024 * 1024, // 2MB max size
1048        }
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    /// Returns the next buffer size to allocate and updates internal state.
1065    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            // keep all buffers
1142            (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            // rewrite, all zero offsets
1152            (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            // keep all buffers, but some have offsets
1160            (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            // rewrite and some have offsets
1172            (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 the buffer to [start, end) range
1187    Slice {
1188        start: u32,
1189        end: u32,
1190    },
1191    /// Rewrite data into new compacted buffer
1192    Rewrite,
1193}
1194
1195fn compaction_strategy(
1196    buffer_utilization: &BufferUtilization,
1197    threshold: f64,
1198) -> CompactionStrategy {
1199    match buffer_utilization.overall_utilization() {
1200        // rewrite empty or not used buffers TODO(os): maybe keep them
1201        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
1239// Care when adding new variants or fields in this enum, it will mess with inlining if it gets too big
1240enum 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 offset < offset_shift, this view was invalid and wasn't counted in buffer_utilizations.
1270                // Return an empty view to match how invalid views are handled in the Rewriting path.
1271                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 offset < offset_shift, this view was invalid and wasn't counted in buffer_utilizations.
1290                // Return an empty view to match how invalid views are handled in the Rewriting path.
1291                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    /// Can return None if this view can't be adjusted, because there is no precomputed lookup
1309    /// for the current buffer.
1310    #[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    /// A long-enough value that a view over it must reference a data buffer.
1347    const LONG: &str = "a value that is far too long to inline";
1348
1349    /// The heap is adopted zero-copy as a data buffer, the views are built against it in place,
1350    /// and the append composes with staged in-progress bytes on either side.
1351    #[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        // Stages an in-progress buffer the bulk append has to flush first.
1357        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        // The adopted heap sits after the flushed in-progress buffer, untouched.
1368        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    /// Offsets need not start at zero: only the referenced range of the heap is adopted.
1381    #[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        // Zero-copy adoption of just the `offsets[0]..offsets[last]` range.
1397        // SAFETY: offset 2 is in bounds of the heap.
1398        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    /// A compacting builder must measure the heap it is handed: values short enough to inline
1413    /// never reference it, so an under-utilized heap is rewritten to just the referencing values.
1414    #[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    /// Rewriting an under-utilized heap must consume null values' spans without retaining their
1441    /// bytes or producing views that reference them.
1442    #[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    /// `push_buffers` returns where each buffer landed, and a deduplicating builder maps a
1465    /// re-pushed buffer back to its existing index instead of holding it twice.
1466    #[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    /// Gathered views are rebased onto wherever the pushed buffers landed, and null rows never
1483    /// resolve their index.
1484    #[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        // Stages an in-progress buffer that the gather has to flush ahead of its own buffers.
1500        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            // The null row's code is garbage; the builder must not look it up.
1509            |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    /// Scattered patches overwrite the fill view at their rows, and both are rebased onto the
1523    /// adopted buffers.
1524    #[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    /// A compacting builder must measure the buffers a gather hands it: with the middle
1566    /// dictionary value never referenced, adopting the heap whole would keep its bytes alive.
1567    /// Duplicate codes must share one rewritten copy, and null rows must not resolve their code.
1568    #[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    /// A gather that references only a contiguous tail of the heap must adopt just that slice,
1612    /// zero-copy.
1613    #[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        // SAFETY: LONG.len() is in bounds of the heap.
1634        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    /// A buffer whose referenced share clears the threshold is adopted whole, zero-copy.
1643    #[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    /// An all-null gather references nothing, so a compacting builder must not retain any of the
1671    /// buffers — and must never resolve a code.
1672    #[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    /// A fully-inlined heap has zero utilization; a compacting builder must not retain it at all.
1699    #[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        // Test with Utf8 builder.
1832        let mut utf8_builder =
1833            VarBinViewBuilder::with_capacity(DType::Utf8(Nullability::Nullable), 10);
1834
1835        // Test appending a valid utf8 value.
1836        let utf8_scalar1 = Scalar::utf8("hello", Nullability::Nullable);
1837        utf8_builder.append_scalar(&utf8_scalar1).unwrap();
1838
1839        // Test appending another value.
1840        let utf8_scalar2 = Scalar::utf8("world", Nullability::Nullable);
1841        utf8_builder.append_scalar(&utf8_scalar2).unwrap();
1842
1843        // Test appending null value.
1844        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        // Test with Binary builder.
1853        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        // Test wrong dtype error.
1868        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        // Test Fixed strategy
1879        let mut strategy = BufferGrowthStrategy::fixed(1024);
1880
1881        // Should always return the fixed size
1882        assert_eq!(strategy.next_size(), 1024);
1883        assert_eq!(strategy.next_size(), 1024);
1884        assert_eq!(strategy.next_size(), 1024);
1885
1886        // Test Exponential strategy
1887        let mut strategy = BufferGrowthStrategy::exponential(1024, 8192);
1888
1889        // Should double each time until hitting max_size
1890        assert_eq!(strategy.next_size(), 1024); // First: 1024
1891        assert_eq!(strategy.next_size(), 2048); // Second: 2048
1892        assert_eq!(strategy.next_size(), 4096); // Third: 4096
1893        assert_eq!(strategy.next_size(), 8192); // Fourth: 8192 (max)
1894        assert_eq!(strategy.next_size(), 8192); // Fifth: 8192 (capped)
1895    }
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        // Create a value larger than max_size
1911        let large_value = vec![0u8; 8192];
1912
1913        // Should successfully append the large value
1914        builder.append_value(&large_value);
1915
1916        let array = builder.finish_into_varbinview();
1917        assert_eq!(array.len(), 1);
1918
1919        // Verify the value was stored correctly
1920        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}