arrow_array/builder/
generic_bytes_view_builder.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: See the NOTICE file distributed with this work for additional information regarding copyright ownership
3// SPDX-FileContributor: Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements
4
5use std::any::Any;
6use std::marker::PhantomData;
7use std::sync::Arc;
8
9use arrow_buffer::{Buffer, NullBufferBuilder, ScalarBuffer};
10use arrow_data::{ByteView, MAX_INLINE_VIEW_LEN};
11use arrow_schema::ArrowError;
12use hashbrown::HashTable;
13use hashbrown::hash_table::Entry;
14
15use crate::builder::{ArrayBuilder, BinaryLikeArrayBuilder, StringLikeArrayBuilder};
16use crate::types::bytes::ByteArrayNativeType;
17use crate::types::{BinaryViewType, ByteViewType, StringViewType};
18use crate::{Array, ArrayRef, GenericByteViewArray};
19
20const STARTING_BLOCK_SIZE: u32 = 8 * 1024; // 8KiB
21const MAX_BLOCK_SIZE: u32 = 2 * 1024 * 1024; // 2MiB
22
23enum BlockSizeGrowthStrategy {
24    Fixed { size: u32 },
25    Exponential { current_size: u32 },
26}
27
28impl BlockSizeGrowthStrategy {
29    fn next_size(&mut self) -> u32 {
30        match self {
31            Self::Fixed { size } => *size,
32            Self::Exponential { current_size } => {
33                if *current_size < MAX_BLOCK_SIZE {
34                    // we have fixed start/end block sizes, so we can't overflow
35                    *current_size = current_size.saturating_mul(2);
36                    *current_size
37                } else {
38                    MAX_BLOCK_SIZE
39                }
40            }
41        }
42    }
43}
44
45/// A builder for [`GenericByteViewArray`]
46///
47/// A [`GenericByteViewArray`] consists of a list of data blocks containing string data,
48/// and a list of views into those buffers.
49///
50/// See examples on [`StringViewBuilder`] and [`BinaryViewBuilder`]
51///
52/// This builder can be used in two ways
53///
54/// # Append Values
55///
56/// To avoid bump allocating, this builder allocates data in fixed size blocks, configurable
57/// using [`GenericByteViewBuilder::with_fixed_block_size`]. [`GenericByteViewBuilder::append_value`]
58/// writes values larger than [`MAX_INLINE_VIEW_LEN`] bytes to the current in-progress block, with values smaller
59/// than [`MAX_INLINE_VIEW_LEN`] bytes inlined into the views. If a value is appended that will not fit in the
60/// in-progress block, it will be closed, and a new block of sufficient size allocated
61///
62/// # Append Views
63///
64/// Some use-cases may wish to reuse an existing allocation containing string data, for example,
65/// when parsing data from a parquet data page. In such a case entire blocks can be appended
66/// using [`GenericByteViewBuilder::append_block`] and then views into this block appended
67/// using [`GenericByteViewBuilder::try_append_view`]
68pub struct GenericByteViewBuilder<T: ByteViewType + ?Sized> {
69    views_buffer: Vec<u128>,
70    null_buffer_builder: NullBufferBuilder,
71    completed: Vec<Buffer>,
72    in_progress: Vec<u8>,
73    block_size: BlockSizeGrowthStrategy,
74    /// Some if deduplicating strings
75    /// map `<string hash> -> <index to the views>`
76    string_tracker: Option<(HashTable<usize>, ahash::RandomState)>,
77    max_deduplication_len: Option<u32>,
78    phantom: PhantomData<T>,
79}
80
81impl<T: ByteViewType + ?Sized> GenericByteViewBuilder<T> {
82    /// Creates a new [`GenericByteViewBuilder`].
83    pub fn new() -> Self {
84        Self::with_capacity(1024)
85    }
86
87    /// Creates a new [`GenericByteViewBuilder`] with space for `capacity` string values.
88    pub fn with_capacity(capacity: usize) -> Self {
89        Self {
90            views_buffer: Vec::with_capacity(capacity),
91            null_buffer_builder: NullBufferBuilder::new(capacity),
92            completed: vec![],
93            in_progress: vec![],
94            block_size: BlockSizeGrowthStrategy::Exponential {
95                current_size: STARTING_BLOCK_SIZE,
96            },
97            string_tracker: None,
98            max_deduplication_len: None,
99            phantom: Default::default(),
100        }
101    }
102
103    /// Configure max deduplication length when deduplicating strings while building the array.
104    /// Default is None.
105    ///
106    /// When [`Self::with_deduplicate_strings`] is enabled, the builder attempts to deduplicate
107    /// any strings longer than 12 bytes. However, since it takes time proportional to the length
108    /// of the string to deduplicate, setting this option limits the CPU overhead for this option.  
109    pub fn with_max_deduplication_len(self, max_deduplication_len: u32) -> Self {
110        debug_assert!(
111            max_deduplication_len > 0,
112            "max_deduplication_len must be greater than 0"
113        );
114        Self {
115            max_deduplication_len: Some(max_deduplication_len),
116            ..self
117        }
118    }
119
120    /// Set a fixed buffer size for variable length strings
121    ///
122    /// The block size is the size of the buffer used to store values greater
123    /// than [`MAX_INLINE_VIEW_LEN`] bytes. The builder allocates new buffers when the current
124    /// buffer is full.
125    ///
126    /// By default the builder balances buffer size and buffer count by
127    /// growing buffer size exponentially from 8KB up to 2MB. The
128    /// first buffer allocated is 8KB, then 16KB, then 32KB, etc up to 2MB.
129    ///
130    /// If this method is used, any new buffers allocated are
131    /// exactly this size. This can be useful for advanced users
132    /// that want to control the memory usage and buffer count.
133    ///
134    /// See <https://github.com/apache/arrow-rs/issues/6094> for more details on the implications.
135    pub fn with_fixed_block_size(self, block_size: u32) -> Self {
136        debug_assert!(block_size > 0, "Block size must be greater than 0");
137        Self {
138            block_size: BlockSizeGrowthStrategy::Fixed { size: block_size },
139            ..self
140        }
141    }
142
143    /// Deduplicate strings while building the array
144    ///
145    /// This will potentially decrease the memory usage if the array have repeated strings
146    /// It will also increase the time to build the array as it needs to hash the strings
147    pub fn with_deduplicate_strings(self) -> Self {
148        Self {
149            string_tracker: Some((
150                HashTable::with_capacity(self.views_buffer.capacity()),
151                Default::default(),
152            )),
153            ..self
154        }
155    }
156
157    /// Append a new data block returning the new block offset
158    ///
159    /// Note: this will first flush any in-progress block
160    ///
161    /// This allows appending views from blocks added using [`Self::append_block`]. See
162    /// [`Self::append_value`] for appending individual values
163    ///
164    /// ```
165    /// # use arrow_array::builder::StringViewBuilder;
166    /// let mut builder = StringViewBuilder::new();
167    ///
168    /// let block = builder.append_block(b"helloworldbingobongo".into());
169    ///
170    /// builder.try_append_view(block, 0, 5).unwrap();
171    /// builder.try_append_view(block, 5, 5).unwrap();
172    /// builder.try_append_view(block, 10, 5).unwrap();
173    /// builder.try_append_view(block, 15, 5).unwrap();
174    /// builder.try_append_view(block, 0, 15).unwrap();
175    /// let array = builder.finish();
176    ///
177    /// let actual: Vec<_> = array.iter().flatten().collect();
178    /// let expected = &["hello", "world", "bingo", "bongo", "helloworldbingo"];
179    /// assert_eq!(actual, expected);
180    /// ```
181    ///
182    /// # Panics
183    ///
184    /// Panics if `buffer.len() >= u32::MAX`
185    pub fn append_block(&mut self, buffer: Buffer) -> u32 {
186        assert!(buffer.len() < u32::MAX as usize);
187
188        self.flush_in_progress();
189        let offset = self.completed.len();
190        self.push_completed(buffer);
191        offset as u32
192    }
193
194    /// Append a view of the given `block`, `offset` and `length`
195    ///
196    /// # Safety
197    /// (1) The block must have been added using [`Self::append_block`]
198    /// (2) The range `offset..offset+length` must be within the bounds of the block
199    /// (3) The data in the block must be valid of type `T`
200    pub unsafe fn append_view_unchecked(&mut self, block: u32, offset: u32, len: u32) {
201        let b = unsafe { self.completed.get_unchecked(block as usize) };
202        let start = offset as usize;
203        let end = start.saturating_add(len as usize);
204        let b = unsafe { b.get_unchecked(start..end) };
205
206        let view = make_view(b, block, offset);
207        self.views_buffer.push(view);
208        self.null_buffer_builder.append_non_null();
209    }
210
211    /// Appends an array to the builder.
212    /// This will flush any in-progress block and append the data buffers
213    /// and add the (adapted) views.
214    ///
215    /// # Panics
216    ///
217    /// Panics if the builder would hold more than `u32::MAX` data buffers
218    pub fn append_array(&mut self, array: &GenericByteViewArray<T>) {
219        self.flush_in_progress();
220        // keep original views if this array is the first to be added or if there are no data buffers (all inline views)
221        let keep_views = self.completed.is_empty() || array.data_buffers().is_empty();
222        let starting_buffer = self.completed.len() as u32;
223
224        self.completed.extend(array.data_buffers().iter().cloned());
225        assert!(
226            u32::try_from(self.completed.len()).is_ok(),
227            "Too many blocks"
228        );
229
230        if keep_views {
231            self.views_buffer.extend_from_slice(array.views());
232        } else {
233            self.views_buffer.extend(array.views().iter().map(|v| {
234                // A view of a string of 12 bytes or fewer stores the string
235                // inline and has no buffer index, so adding 0 leaves it
236                // unchanged.
237                let offset = if *v as u32 > MAX_INLINE_VIEW_LEN {
238                    starting_buffer
239                } else {
240                    0
241                };
242                // The buffer index is bits 64 to 95 of a view. Adding
243                // `starting_buffer` gives the string's buffer index in the
244                // builder, which fits in 32 bits, so the sum never carries
245                // into bit 96.
246                v + ((offset as u128) << 64)
247            }));
248        }
249
250        if let Some(null_buffer) = array.nulls() {
251            self.null_buffer_builder.append_buffer(null_buffer);
252        } else {
253            self.null_buffer_builder.append_n_non_nulls(array.len());
254        }
255    }
256
257    /// Try to append a view of the given `block`, `offset` and `length`
258    ///
259    /// See [`Self::append_block`]
260    pub fn try_append_view(&mut self, block: u32, offset: u32, len: u32) -> Result<(), ArrowError> {
261        let b = self.completed.get(block as usize).ok_or_else(|| {
262            ArrowError::InvalidArgumentError(format!("No block found with index {block}"))
263        })?;
264        let start = offset as usize;
265        let end = start.saturating_add(len as usize);
266
267        let b = b.get(start..end).ok_or_else(|| {
268            ArrowError::InvalidArgumentError(format!(
269                "Range {start}..{end} out of bounds for block of length {}",
270                b.len()
271            ))
272        })?;
273
274        if T::Native::from_bytes_checked(b).is_none() {
275            return Err(ArrowError::InvalidArgumentError(
276                "Invalid view data".to_string(),
277            ));
278        }
279
280        unsafe {
281            self.append_view_unchecked(block, offset, len);
282        }
283        Ok(())
284    }
285
286    /// Flushes the in progress block if any
287    #[inline]
288    fn flush_in_progress(&mut self) {
289        if !self.in_progress.is_empty() {
290            let f = Buffer::from_vec(std::mem::take(&mut self.in_progress));
291            self.push_completed(f)
292        }
293    }
294
295    /// Append a block to `self.completed`, checking for overflow
296    #[inline]
297    fn push_completed(&mut self, block: Buffer) {
298        assert!(block.len() < u32::MAX as usize, "Block too large");
299        assert!(self.completed.len() < u32::MAX as usize, "Too many blocks");
300        self.completed.push(block);
301    }
302
303    /// Returns the value at the given index
304    ///
305    /// Useful if we want to know what value has been inserted to the builder
306    ///
307    /// # Panics
308    ///
309    /// Panics if `index >= self.len()`
310    pub fn get_value(&self, index: usize) -> &[u8] {
311        let view = self.views_buffer.as_slice().get(index).unwrap();
312        let len = *view as u32;
313        if len <= MAX_INLINE_VIEW_LEN {
314            // # Safety
315            // The view is valid from the builder
316            unsafe { GenericByteViewArray::<T>::inline_value(view, len as usize) }
317        } else {
318            let view = ByteView::from(*view);
319            if view.buffer_index < self.completed.len() as u32 {
320                let block = &self.completed[view.buffer_index as usize];
321                &block[view.offset as usize..view.offset as usize + view.length as usize]
322            } else {
323                &self.in_progress[view.offset as usize..view.offset as usize + view.length as usize]
324            }
325        }
326    }
327
328    /// Appends a value into the builder
329    ///
330    /// # Panics
331    ///
332    /// Panics if
333    /// - String buffer count exceeds `u32::MAX`
334    /// - String length exceeds `u32::MAX`
335    #[inline]
336    pub fn append_value(&mut self, value: impl AsRef<T::Native>) {
337        self.try_append_value(value).unwrap()
338    }
339
340    /// Appends a value into the builder
341    ///
342    /// # Errors
343    ///
344    /// Returns an error if:
345    /// - String buffer count exceeds `u32::MAX`
346    /// - String length exceeds `u32::MAX`
347    #[inline]
348    pub fn try_append_value(&mut self, value: impl AsRef<T::Native>) -> Result<(), ArrowError> {
349        let v: &[u8] = value.as_ref().as_ref();
350        let length: u32 = v.len().try_into().map_err(|_| {
351            ArrowError::InvalidArgumentError(format!("String length {} exceeds u32::MAX", v.len()))
352        })?;
353
354        // Anything at most `MAX_INLINE_VIEW_LEN` bytes long is inlined; everything else
355        // needs a four byte prefix, which `first_chunk` gives us without any indexing.
356        let prefix = match v.first_chunk::<4>() {
357            Some(prefix) if length > MAX_INLINE_VIEW_LEN => u32::from_le_bytes(*prefix),
358            _ => {
359                let mut view_buffer = [0; 16];
360                view_buffer[0..4].copy_from_slice(&length.to_le_bytes());
361                view_buffer[4..4 + v.len()].copy_from_slice(v);
362                self.views_buffer.push(u128::from_le_bytes(view_buffer));
363                self.null_buffer_builder.append_non_null();
364                return Ok(());
365            }
366        };
367
368        // Deduplication if:
369        // (1) deduplication is enabled.
370        // (2) len > `MAX_INLINE_VIEW_LEN` and len <= `max_deduplication_len`
371        let can_deduplicate = self.string_tracker.is_some()
372            && self
373                .max_deduplication_len
374                .map(|max_length| length <= max_length)
375                .unwrap_or(true);
376        if can_deduplicate && let Some((mut ht, hasher)) = self.string_tracker.take() {
377            let hash_val = hasher.hash_one(v);
378            let hasher_fn = |v: &_| hasher.hash_one(v);
379
380            let entry = ht.entry(
381                hash_val,
382                |idx| {
383                    let stored_value = self.get_value(*idx);
384                    v == stored_value
385                },
386                hasher_fn,
387            );
388            match entry {
389                Entry::Occupied(occupied) => {
390                    // If the string already exists, we will directly use the view
391                    let idx = occupied.get();
392                    self.views_buffer.push(self.views_buffer[*idx]);
393                    self.null_buffer_builder.append_non_null();
394                    self.string_tracker = Some((ht, hasher));
395                    return Ok(());
396                }
397                Entry::Vacant(vacant) => {
398                    // o.w. we insert the (string hash -> view index)
399                    // the idx is current length of views_builder, as we are inserting a new view
400                    vacant.insert(self.views_buffer.len());
401                }
402            }
403            self.string_tracker = Some((ht, hasher));
404        }
405
406        let required_cap = self.in_progress.len() + v.len();
407        if self.in_progress.capacity() < required_cap {
408            self.flush_in_progress();
409            let to_reserve = v.len().max(self.block_size.next_size() as usize);
410            self.in_progress.reserve(to_reserve);
411        }
412
413        let offset = self.in_progress.len() as u32;
414        self.in_progress.extend_from_slice(v);
415
416        let buffer_index: u32 = self.completed.len().try_into().map_err(|_| {
417            ArrowError::InvalidArgumentError(format!(
418                "Buffer count {} exceeds u32::MAX",
419                self.completed.len()
420            ))
421        })?;
422
423        let view = ByteView {
424            length,
425            prefix,
426            buffer_index,
427            offset,
428        };
429        self.views_buffer.push(view.into());
430        self.null_buffer_builder.append_non_null();
431
432        Ok(())
433    }
434
435    /// Append an `Option` value into the builder
436    #[inline]
437    pub fn append_option(&mut self, value: Option<impl AsRef<T::Native>>) {
438        match value {
439            None => self.append_null(),
440            Some(v) => self.append_value(v),
441        }
442    }
443
444    /// Append the same value `n` times into the builder
445    ///
446    /// This is more efficient than calling [`Self::try_append_value`] `n` times,
447    /// especially when deduplication is enabled, as it only hashes the value once.
448    ///
449    /// # Errors
450    ///
451    /// Returns an error if
452    /// - String buffer count exceeds `u32::MAX`
453    /// - String length exceeds `u32::MAX`
454    ///
455    /// # Example
456    /// ```
457    /// # use arrow_array::builder::StringViewBuilder;
458    /// # use arrow_array::Array;
459    /// let mut builder = StringViewBuilder::new().with_deduplicate_strings();
460    ///
461    /// // Append "hello" 1000 times efficiently
462    /// builder.try_append_value_n("hello", 1000)?;
463    ///
464    /// let array = builder.finish();
465    /// assert_eq!(array.len(), 1000);
466    ///
467    /// // All values are "hello"
468    /// for value in array.iter() {
469    ///     assert_eq!(value, Some("hello"));
470    /// }
471    /// # Ok::<(), arrow_schema::ArrowError>(())
472    /// ```
473    #[inline]
474    pub fn try_append_value_n(
475        &mut self,
476        value: impl AsRef<T::Native>,
477        n: usize,
478    ) -> Result<(), ArrowError> {
479        if n == 0 {
480            return Ok(());
481        }
482        // Process value once (handles deduplication, buffer management, view creation)
483        self.try_append_value(value)?;
484        // Reuse the view (n-1) times
485        let view = *self.views_buffer.last().unwrap();
486        self.views_buffer.extend(std::iter::repeat_n(view, n - 1));
487        self.null_buffer_builder.append_n_non_nulls(n - 1);
488        Ok(())
489    }
490
491    /// Append a null value into the builder
492    #[inline]
493    pub fn append_null(&mut self) {
494        self.null_buffer_builder.append_null();
495        self.views_buffer.push(0);
496    }
497
498    /// Builds the [`GenericByteViewArray`] and reset this builder
499    pub fn finish(&mut self) -> GenericByteViewArray<T> {
500        self.flush_in_progress();
501        let completed = std::mem::take(&mut self.completed);
502        let nulls = self.null_buffer_builder.finish();
503        if let Some((ht, _)) = self.string_tracker.as_mut() {
504            ht.clear();
505        }
506        let views = std::mem::take(&mut self.views_buffer);
507        // SAFETY: valid by construction
508        unsafe { GenericByteViewArray::new_unchecked(views.into(), completed.into(), nulls) }
509    }
510
511    /// Builds the [`GenericByteViewArray`] without resetting the builder
512    pub fn finish_cloned(&self) -> GenericByteViewArray<T> {
513        let mut completed = self.completed.clone();
514        if !self.in_progress.is_empty() {
515            completed.push(Buffer::from_slice_ref(&self.in_progress));
516        }
517        let len = self.views_buffer.len();
518        let views = Buffer::from_slice_ref(self.views_buffer.as_slice());
519        let views = ScalarBuffer::new(views, 0, len);
520        let nulls = self.null_buffer_builder.finish_cloned();
521        // SAFETY: valid by construction
522        unsafe { GenericByteViewArray::new_unchecked(views, completed.into(), nulls) }
523    }
524
525    /// Returns the current null buffer as a slice
526    pub fn validity_slice(&self) -> Option<&[u8]> {
527        self.null_buffer_builder.as_slice()
528    }
529
530    /// Return the allocated size of this builder in bytes, useful for memory accounting.
531    pub fn allocated_size(&self) -> usize {
532        let views = self.views_buffer.capacity() * std::mem::size_of::<u128>();
533        let null = self.null_buffer_builder.allocated_size();
534        let buffer_size = self.completed.iter().map(|b| b.capacity()).sum::<usize>();
535        let in_progress = self.in_progress.capacity();
536        let tracker = match &self.string_tracker {
537            Some((ht, _)) => ht.capacity() * std::mem::size_of::<usize>(),
538            None => 0,
539        };
540        buffer_size + in_progress + tracker + views + null
541    }
542}
543
544impl<T: ByteViewType + ?Sized> Default for GenericByteViewBuilder<T> {
545    fn default() -> Self {
546        Self::new()
547    }
548}
549
550impl<T: ByteViewType + ?Sized> std::fmt::Debug for GenericByteViewBuilder<T> {
551    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
552        write!(f, "{}ViewBuilder", T::PREFIX)?;
553        f.debug_struct("")
554            .field("views_buffer", &self.views_buffer)
555            .field("in_progress", &self.in_progress)
556            .field("completed", &self.completed)
557            .field("null_buffer_builder", &self.null_buffer_builder)
558            .finish()
559    }
560}
561
562impl<T: ByteViewType + ?Sized> ArrayBuilder for GenericByteViewBuilder<T> {
563    fn len(&self) -> usize {
564        self.null_buffer_builder.len()
565    }
566
567    fn finish(&mut self) -> ArrayRef {
568        Arc::new(self.finish())
569    }
570
571    fn finish_cloned(&self) -> ArrayRef {
572        Arc::new(self.finish_cloned())
573    }
574
575    fn as_any(&self) -> &dyn Any {
576        self
577    }
578
579    fn as_any_mut(&mut self) -> &mut dyn Any {
580        self
581    }
582
583    fn into_box_any(self: Box<Self>) -> Box<dyn Any> {
584        self
585    }
586}
587
588impl<T: ByteViewType + ?Sized, V: AsRef<T::Native>> Extend<Option<V>>
589    for GenericByteViewBuilder<T>
590{
591    #[inline]
592    fn extend<I: IntoIterator<Item = Option<V>>>(&mut self, iter: I) {
593        for v in iter {
594            self.append_option(v)
595        }
596    }
597}
598
599/// Array builder for [`StringViewArray`][crate::StringViewArray]
600///
601/// Values can be appended using [`GenericByteViewBuilder::append_value`], and nulls with
602/// [`GenericByteViewBuilder::append_null`] as normal.
603///
604/// # Example
605/// ```
606/// # use arrow_array::builder::StringViewBuilder;
607/// # use arrow_array::StringViewArray;
608/// let mut builder = StringViewBuilder::new();
609/// builder.append_value("hello");
610/// builder.append_null();
611/// builder.append_value("world");
612/// let array = builder.finish();
613///
614/// let expected = vec![Some("hello"), None, Some("world")];
615/// let actual: Vec<_> = array.iter().collect();
616/// assert_eq!(expected, actual);
617/// ```
618pub type StringViewBuilder = GenericByteViewBuilder<StringViewType>;
619
620impl StringLikeArrayBuilder for StringViewBuilder {
621    fn type_name() -> &'static str {
622        std::any::type_name::<StringViewBuilder>()
623    }
624    fn with_capacity(capacity: usize) -> Self {
625        Self::with_capacity(capacity)
626    }
627    fn append_value(&mut self, value: &str) {
628        Self::append_value(self, value);
629    }
630    fn append_null(&mut self) {
631        Self::append_null(self);
632    }
633}
634
635///  Array builder for [`BinaryViewArray`][crate::BinaryViewArray]
636///
637/// Values can be appended using [`GenericByteViewBuilder::append_value`], and nulls with
638/// [`GenericByteViewBuilder::append_null`] as normal.
639///
640/// # Example
641/// ```
642/// # use arrow_array::builder::BinaryViewBuilder;
643/// use arrow_array::BinaryViewArray;
644/// let mut builder = BinaryViewBuilder::new();
645/// builder.append_value("hello");
646/// builder.append_null();
647/// builder.append_value("world");
648/// let array = builder.finish();
649///
650/// let expected: Vec<Option<&[u8]>> = vec![Some(b"hello"), None, Some(b"world")];
651/// let actual: Vec<_> = array.iter().collect();
652/// assert_eq!(expected, actual);
653/// ```
654///
655pub type BinaryViewBuilder = GenericByteViewBuilder<BinaryViewType>;
656
657impl BinaryLikeArrayBuilder for BinaryViewBuilder {
658    fn type_name() -> &'static str {
659        std::any::type_name::<BinaryViewBuilder>()
660    }
661    fn with_capacity(capacity: usize) -> Self {
662        Self::with_capacity(capacity)
663    }
664    fn append_value(&mut self, value: &[u8]) {
665        Self::append_value(self, value);
666    }
667    fn append_null(&mut self) {
668        Self::append_null(self);
669    }
670}
671
672/// Creates a view from a fixed length input (the compiler can generate
673/// specialized code for this)
674fn make_inlined_view<const LEN: usize>(data: &[u8]) -> u128 {
675    let mut view_buffer = [0; 16];
676    view_buffer[0..4].copy_from_slice(&(LEN as u32).to_le_bytes());
677    view_buffer[4..4 + LEN].copy_from_slice(&data[..LEN]);
678    u128::from_le_bytes(view_buffer)
679}
680
681/// Create a view based on the given data, block id and offset.
682///
683/// Note that the code below is carefully examined with x86_64 assembly code: <https://godbolt.org/z/685YPsd5G>
684/// The goal is to avoid calling into `ptr::copy_non_interleave`, which makes function call (i.e., not inlined),
685/// which slows down things.
686#[inline(never)]
687pub fn make_view(data: &[u8], block_id: u32, offset: u32) -> u128 {
688    let len = data.len();
689
690    // Generate specialized code for each potential small string length
691    // to improve performance
692    match len {
693        0 => make_inlined_view::<0>(data),
694        1 => make_inlined_view::<1>(data),
695        2 => make_inlined_view::<2>(data),
696        3 => make_inlined_view::<3>(data),
697        4 => make_inlined_view::<4>(data),
698        5 => make_inlined_view::<5>(data),
699        6 => make_inlined_view::<6>(data),
700        7 => make_inlined_view::<7>(data),
701        8 => make_inlined_view::<8>(data),
702        9 => make_inlined_view::<9>(data),
703        10 => make_inlined_view::<10>(data),
704        11 => make_inlined_view::<11>(data),
705        12 => make_inlined_view::<12>(data),
706        // When string is longer than 12 bytes, it can't be inlined, we create a ByteView instead.
707        _ => {
708            let view = ByteView {
709                length: len as u32,
710                // this arm only matches lengths above 12, so there are at least four bytes
711                prefix: u32::from_le_bytes([data[0], data[1], data[2], data[3]]),
712                buffer_index: block_id,
713                offset,
714            };
715            view.as_u128()
716        }
717    }
718}
719
720#[cfg(test)]
721mod tests {
722    use core::str;
723
724    use arrow_buffer::ArrowNativeType;
725
726    use super::*;
727
728    #[test]
729    fn test_string_max_deduplication_len() {
730        let value_1 = "short";
731        let value_2 = "not so similar string but long";
732        let value_3 = "1234567890123";
733
734        let max_deduplication_len = MAX_INLINE_VIEW_LEN * 2;
735
736        let mut builder = StringViewBuilder::new()
737            .with_deduplicate_strings()
738            .with_max_deduplication_len(max_deduplication_len);
739
740        assert!(value_1.len() < MAX_INLINE_VIEW_LEN.as_usize());
741        assert!(value_2.len() > max_deduplication_len.as_usize());
742        assert!(
743            value_3.len() > MAX_INLINE_VIEW_LEN.as_usize()
744                && value_3.len() < max_deduplication_len.as_usize()
745        );
746
747        // append value1 (short), expect it is inlined and not deduplicated
748        builder.append_value(value_1); // view 0
749        builder.append_value(value_1); // view 1
750        // append value2, expect second copy is not deduplicated as it exceeds max_deduplication_len
751        builder.append_value(value_2); // view 2
752        builder.append_value(value_2); // view 3
753        // append value3, expect second copy is deduplicated
754        builder.append_value(value_3); // view 4
755        builder.append_value(value_3); // view 5
756
757        let array = builder.finish();
758
759        // verify
760        let v2 = ByteView::from(array.views()[2]);
761        let v3 = ByteView::from(array.views()[3]);
762        assert_eq!(v2.buffer_index, v3.buffer_index); // stored in same buffer
763        assert_ne!(v2.offset, v3.offset); // different offsets --> not deduplicated
764
765        let v4 = ByteView::from(array.views()[4]);
766        let v5 = ByteView::from(array.views()[5]);
767        assert_eq!(v4.buffer_index, v5.buffer_index); // stored in same buffer
768        assert_eq!(v4.offset, v5.offset); // same offsets --> deduplicated
769    }
770
771    #[test]
772    fn test_string_view_deduplicate() {
773        let value_1 = "long string to test string view";
774        let value_2 = "not so similar string but long";
775
776        let mut builder = StringViewBuilder::new()
777            .with_deduplicate_strings()
778            .with_fixed_block_size(value_1.len() as u32 * 2); // so that we will have multiple buffers
779
780        let values = vec![
781            Some(value_1),
782            Some(value_2),
783            Some("short"),
784            Some(value_1),
785            None,
786            Some(value_2),
787            Some(value_1),
788        ];
789        builder.extend(values.clone());
790
791        let array = builder.finish_cloned();
792        array.to_data().validate_full().unwrap();
793        assert_eq!(array.data_buffers().len(), 1); // without duplication we would need 3 buffers.
794        let actual: Vec<_> = array.iter().collect();
795        assert_eq!(actual, values);
796
797        let view0 = array.views().first().unwrap();
798        let view3 = array.views().get(3).unwrap();
799        let view6 = array.views().get(6).unwrap();
800
801        assert_eq!(view0, view3);
802        assert_eq!(view0, view6);
803
804        assert_eq!(array.views().get(1), array.views().get(5));
805    }
806
807    #[test]
808    fn test_string_view_deduplicate_after_finish() {
809        let mut builder = StringViewBuilder::new().with_deduplicate_strings();
810
811        let value_1 = "long string to test string view";
812        let value_2 = "not so similar string but long";
813        builder.append_value(value_1);
814        let _array = builder.finish();
815        builder.append_value(value_2);
816        let _array = builder.finish();
817        builder.append_value(value_1);
818        let _array = builder.finish();
819    }
820
821    #[test]
822    fn test_string_view() {
823        let b1 = Buffer::from(b"world\xFFbananas\xF0\x9F\x98\x81");
824        let b2 = Buffer::from(b"cupcakes");
825        let b3 = Buffer::from(b"Many strings are here contained of great length and verbosity");
826
827        let mut v = StringViewBuilder::new();
828        assert_eq!(v.append_block(b1), 0);
829
830        v.append_value("This is a very long string that exceeds the inline length");
831        v.append_value("This is another very long string that exceeds the inline length");
832
833        assert_eq!(v.append_block(b2), 2);
834        assert_eq!(v.append_block(b3), 3);
835
836        // Test short strings
837        v.try_append_view(0, 0, 5).unwrap(); // world
838        v.try_append_view(0, 6, 7).unwrap(); // bananas
839        v.try_append_view(2, 3, 5).unwrap(); // cake
840        v.try_append_view(2, 0, 3).unwrap(); // cup
841        v.try_append_view(2, 0, 8).unwrap(); // cupcakes
842        v.try_append_view(0, 13, 4).unwrap(); // 😁
843        v.try_append_view(0, 13, 0).unwrap(); //
844
845        // Test longer strings
846        v.try_append_view(3, 0, 16).unwrap(); // Many strings are
847        v.try_append_view(1, 0, 19).unwrap(); // This is a very long
848        v.try_append_view(3, 13, 27).unwrap(); // here contained of great length
849
850        v.append_value("I do so like long strings");
851
852        let array = v.finish_cloned();
853        array.to_data().validate_full().unwrap();
854        assert_eq!(array.data_buffers().len(), 5);
855        let actual: Vec<_> = array.iter().flatten().collect();
856        assert_eq!(
857            actual,
858            &[
859                "This is a very long string that exceeds the inline length",
860                "This is another very long string that exceeds the inline length",
861                "world",
862                "bananas",
863                "cakes",
864                "cup",
865                "cupcakes",
866                "😁",
867                "",
868                "Many strings are",
869                "This is a very long",
870                "are here contained of great",
871                "I do so like long strings"
872            ]
873        );
874
875        let err = v.try_append_view(0, u32::MAX, 1).unwrap_err();
876        assert_eq!(
877            err.to_string(),
878            "Invalid argument error: Range 4294967295..4294967296 out of bounds for block of length 17"
879        );
880
881        let err = v.try_append_view(0, 1, u32::MAX).unwrap_err();
882        assert_eq!(
883            err.to_string(),
884            "Invalid argument error: Range 1..4294967296 out of bounds for block of length 17"
885        );
886
887        let err = v.try_append_view(0, 13, 2).unwrap_err();
888        assert_eq!(err.to_string(), "Invalid argument error: Invalid view data");
889
890        let err = v.try_append_view(0, 40, 0).unwrap_err();
891        assert_eq!(
892            err.to_string(),
893            "Invalid argument error: Range 40..40 out of bounds for block of length 17"
894        );
895
896        let err = v.try_append_view(5, 0, 0).unwrap_err();
897        assert_eq!(
898            err.to_string(),
899            "Invalid argument error: No block found with index 5"
900        );
901    }
902
903    #[test]
904    fn test_string_view_with_block_size_growth() {
905        let mut exp_builder = StringViewBuilder::new();
906        let mut fixed_builder = StringViewBuilder::new().with_fixed_block_size(STARTING_BLOCK_SIZE);
907
908        let long_string = str::from_utf8(&[b'a'; STARTING_BLOCK_SIZE as usize]).unwrap();
909
910        for i in 0..9 {
911            // 8k, 16k, 32k, 64k, 128k, 256k, 512k, 1M, 2M
912            for _ in 0..(2_u32.pow(i)) {
913                exp_builder.append_value(long_string);
914                fixed_builder.append_value(long_string);
915            }
916            exp_builder.flush_in_progress();
917            fixed_builder.flush_in_progress();
918
919            // Every step only add one buffer, but the buffer size is much larger
920            assert_eq!(exp_builder.completed.len(), i as usize + 1);
921            assert_eq!(
922                exp_builder.completed[i as usize].len(),
923                STARTING_BLOCK_SIZE as usize * 2_usize.pow(i)
924            );
925
926            // This step we added 2^i blocks, the sum of blocks should be 2^(i+1) - 1
927            assert_eq!(fixed_builder.completed.len(), 2_usize.pow(i + 1) - 1);
928
929            // Every buffer is fixed size
930            assert!(
931                fixed_builder
932                    .completed
933                    .iter()
934                    .all(|b| b.len() == STARTING_BLOCK_SIZE as usize)
935            );
936        }
937
938        // Add one more value, and the buffer stop growing.
939        exp_builder.append_value(long_string);
940        exp_builder.flush_in_progress();
941        assert_eq!(
942            exp_builder.completed.last().unwrap().capacity(),
943            MAX_BLOCK_SIZE as usize
944        );
945    }
946
947    #[test]
948    fn test_append_value_n() {
949        // Test with inline strings (<=12 bytes)
950        let mut builder = StringViewBuilder::new();
951
952        builder.try_append_value_n("hello", 100).unwrap();
953        builder.append_value("world");
954        builder.try_append_value_n("foo", 50).unwrap();
955
956        let array = builder.finish();
957        assert_eq!(array.len(), 151);
958        assert_eq!(array.null_count(), 0);
959
960        // Verify the values
961        for i in 0..100 {
962            assert_eq!(array.value(i), "hello");
963        }
964        assert_eq!(array.value(100), "world");
965        for i in 101..151 {
966            assert_eq!(array.value(i), "foo");
967        }
968
969        // All inline strings should have no data buffers
970        assert_eq!(array.data_buffers().len(), 0);
971    }
972
973    #[test]
974    fn test_append_value_n_with_deduplication() {
975        let long_string = "This is a very long string that exceeds the inline length";
976
977        // Test with deduplication enabled
978        let mut builder = StringViewBuilder::new().with_deduplicate_strings();
979
980        // First append the string once to add it to the hash map
981        builder.append_value(long_string);
982
983        // Then append_n the same string - should deduplicate and reuse the existing value
984        builder.try_append_value_n(long_string, 999).unwrap();
985
986        let array = builder.finish();
987        assert_eq!(array.len(), 1000);
988        assert_eq!(array.null_count(), 0);
989
990        // Verify all values are the same
991        for i in 0..1000 {
992            assert_eq!(array.value(i), long_string);
993        }
994
995        // With deduplication, should only have 1 data buffer containing the string once
996        assert_eq!(array.data_buffers().len(), 1);
997
998        // All views should be identical
999        let first_view = array.views()[0];
1000        for view in array.views() {
1001            assert_eq!(*view, first_view);
1002        }
1003    }
1004
1005    #[test]
1006    fn test_append_value_n_zero() {
1007        let mut builder = StringViewBuilder::new();
1008
1009        builder.append_value("first");
1010        builder.try_append_value_n("should not appear", 0).unwrap();
1011        builder.append_value("second");
1012
1013        let array = builder.finish();
1014        assert_eq!(array.len(), 2);
1015        assert_eq!(array.value(0), "first");
1016        assert_eq!(array.value(1), "second");
1017    }
1018}
arrow_array/builder/generic_bytes_view_builder.rs

arrow_array/builder/
generic_bytes_view_builder.rs
