arrow_array/builder/
generic_bytes_view_builder.rs

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

arrow_array/builder/
generic_bytes_view_builder.rs
