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        // Anything at most `MAX_INLINE_VIEW_LEN` bytes long is inlined; everything else
354        // needs a four byte prefix, which `first_chunk` gives us without any indexing.
355        let prefix = match v.first_chunk::<4>() {
356            Some(prefix) if length > MAX_INLINE_VIEW_LEN => u32::from_le_bytes(*prefix),
357            _ => {
358                let mut view_buffer = [0; 16];
359                view_buffer[0..4].copy_from_slice(&length.to_le_bytes());
360                view_buffer[4..4 + v.len()].copy_from_slice(v);
361                self.views_buffer.push(u128::from_le_bytes(view_buffer));
362                self.null_buffer_builder.append_non_null();
363                return Ok(());
364            }
365        };
366
367        // Deduplication if:
368        // (1) deduplication is enabled.
369        // (2) len > `MAX_INLINE_VIEW_LEN` and len <= `max_deduplication_len`
370        let can_deduplicate = self.string_tracker.is_some()
371            && self
372                .max_deduplication_len
373                .map(|max_length| length <= max_length)
374                .unwrap_or(true);
375        if can_deduplicate && let Some((mut ht, hasher)) = self.string_tracker.take() {
376            let hash_val = hasher.hash_one(v);
377            let hasher_fn = |v: &_| hasher.hash_one(v);
378
379            let entry = ht.entry(
380                hash_val,
381                |idx| {
382                    let stored_value = self.get_value(*idx);
383                    v == stored_value
384                },
385                hasher_fn,
386            );
387            match entry {
388                Entry::Occupied(occupied) => {
389                    // If the string already exists, we will directly use the view
390                    let idx = occupied.get();
391                    self.views_buffer.push(self.views_buffer[*idx]);
392                    self.null_buffer_builder.append_non_null();
393                    self.string_tracker = Some((ht, hasher));
394                    return Ok(());
395                }
396                Entry::Vacant(vacant) => {
397                    // o.w. we insert the (string hash -> view index)
398                    // the idx is current length of views_builder, as we are inserting a new view
399                    vacant.insert(self.views_buffer.len());
400                }
401            }
402            self.string_tracker = Some((ht, hasher));
403        }
404
405        let required_cap = self.in_progress.len() + v.len();
406        if self.in_progress.capacity() < required_cap {
407            self.flush_in_progress();
408            let to_reserve = v.len().max(self.block_size.next_size() as usize);
409            self.in_progress.reserve(to_reserve);
410        }
411
412        let offset = self.in_progress.len() as u32;
413        self.in_progress.extend_from_slice(v);
414
415        let buffer_index: u32 = self.completed.len().try_into().map_err(|_| {
416            ArrowError::InvalidArgumentError(format!(
417                "Buffer count {} exceeds u32::MAX",
418                self.completed.len()
419            ))
420        })?;
421
422        let view = ByteView {
423            length,
424            prefix,
425            buffer_index,
426            offset,
427        };
428        self.views_buffer.push(view.into());
429        self.null_buffer_builder.append_non_null();
430
431        Ok(())
432    }
433
434    /// Append an `Option` value into the builder
435    #[inline]
436    pub fn append_option(&mut self, value: Option<impl AsRef<T::Native>>) {
437        match value {
438            None => self.append_null(),
439            Some(v) => self.append_value(v),
440        }
441    }
442
443    /// Append the same value `n` times into the builder
444    ///
445    /// This is more efficient than calling [`Self::try_append_value`] `n` times,
446    /// especially when deduplication is enabled, as it only hashes the value once.
447    ///
448    /// # Errors
449    ///
450    /// Returns an error if
451    /// - String buffer count exceeds `u32::MAX`
452    /// - String length exceeds `u32::MAX`
453    ///
454    /// # Example
455    /// ```
456    /// # use arrow_array::builder::StringViewBuilder;
457    /// # use arrow_array::Array;
458    /// let mut builder = StringViewBuilder::new().with_deduplicate_strings();
459    ///
460    /// // Append "hello" 1000 times efficiently
461    /// builder.try_append_value_n("hello", 1000)?;
462    ///
463    /// let array = builder.finish();
464    /// assert_eq!(array.len(), 1000);
465    ///
466    /// // All values are "hello"
467    /// for value in array.iter() {
468    ///     assert_eq!(value, Some("hello"));
469    /// }
470    /// # Ok::<(), arrow_schema::ArrowError>(())
471    /// ```
472    #[inline]
473    pub fn try_append_value_n(
474        &mut self,
475        value: impl AsRef<T::Native>,
476        n: usize,
477    ) -> Result<(), ArrowError> {
478        if n == 0 {
479            return Ok(());
480        }
481        // Process value once (handles deduplication, buffer management, view creation)
482        self.try_append_value(value)?;
483        // Reuse the view (n-1) times
484        let view = *self.views_buffer.last().unwrap();
485        self.views_buffer.extend(std::iter::repeat_n(view, n - 1));
486        self.null_buffer_builder.append_n_non_nulls(n - 1);
487        Ok(())
488    }
489
490    /// Append a null value into the builder
491    #[inline]
492    pub fn append_null(&mut self) {
493        self.null_buffer_builder.append_null();
494        self.views_buffer.push(0);
495    }
496
497    /// Builds the [`GenericByteViewArray`] and reset this builder
498    pub fn finish(&mut self) -> GenericByteViewArray<T> {
499        self.flush_in_progress();
500        let completed = std::mem::take(&mut self.completed);
501        let nulls = self.null_buffer_builder.finish();
502        if let Some((ht, _)) = self.string_tracker.as_mut() {
503            ht.clear();
504        }
505        let views = std::mem::take(&mut self.views_buffer);
506        // SAFETY: valid by construction
507        unsafe { GenericByteViewArray::new_unchecked(views.into(), completed.into(), nulls) }
508    }
509
510    /// Builds the [`GenericByteViewArray`] without resetting the builder
511    pub fn finish_cloned(&self) -> GenericByteViewArray<T> {
512        let mut completed = self.completed.clone();
513        if !self.in_progress.is_empty() {
514            completed.push(Buffer::from_slice_ref(&self.in_progress));
515        }
516        let len = self.views_buffer.len();
517        let views = Buffer::from_slice_ref(self.views_buffer.as_slice());
518        let views = ScalarBuffer::new(views, 0, len);
519        let nulls = self.null_buffer_builder.finish_cloned();
520        // SAFETY: valid by construction
521        unsafe { GenericByteViewArray::new_unchecked(views, completed.into(), nulls) }
522    }
523
524    /// Returns the current null buffer as a slice
525    pub fn validity_slice(&self) -> Option<&[u8]> {
526        self.null_buffer_builder.as_slice()
527    }
528
529    /// Return the allocated size of this builder in bytes, useful for memory accounting.
530    pub fn allocated_size(&self) -> usize {
531        let views = self.views_buffer.capacity() * std::mem::size_of::<u128>();
532        let null = self.null_buffer_builder.allocated_size();
533        let buffer_size = self.completed.iter().map(|b| b.capacity()).sum::<usize>();
534        let in_progress = self.in_progress.capacity();
535        let tracker = match &self.string_tracker {
536            Some((ht, _)) => ht.capacity() * std::mem::size_of::<usize>(),
537            None => 0,
538        };
539        buffer_size + in_progress + tracker + views + null
540    }
541}
542
543impl<T: ByteViewType + ?Sized> Default for GenericByteViewBuilder<T> {
544    fn default() -> Self {
545        Self::new()
546    }
547}
548
549impl<T: ByteViewType + ?Sized> std::fmt::Debug for GenericByteViewBuilder<T> {
550    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
551        write!(f, "{}ViewBuilder", T::PREFIX)?;
552        f.debug_struct("")
553            .field("views_buffer", &self.views_buffer)
554            .field("in_progress", &self.in_progress)
555            .field("completed", &self.completed)
556            .field("null_buffer_builder", &self.null_buffer_builder)
557            .finish()
558    }
559}
560
561impl<T: ByteViewType + ?Sized> ArrayBuilder for GenericByteViewBuilder<T> {
562    fn len(&self) -> usize {
563        self.null_buffer_builder.len()
564    }
565
566    fn finish(&mut self) -> ArrayRef {
567        Arc::new(self.finish())
568    }
569
570    fn finish_cloned(&self) -> ArrayRef {
571        Arc::new(self.finish_cloned())
572    }
573
574    fn as_any(&self) -> &dyn Any {
575        self
576    }
577
578    fn as_any_mut(&mut self) -> &mut dyn Any {
579        self
580    }
581
582    fn into_box_any(self: Box<Self>) -> Box<dyn Any> {
583        self
584    }
585}
586
587impl<T: ByteViewType + ?Sized, V: AsRef<T::Native>> Extend<Option<V>>
588    for GenericByteViewBuilder<T>
589{
590    #[inline]
591    fn extend<I: IntoIterator<Item = Option<V>>>(&mut self, iter: I) {
592        for v in iter {
593            self.append_option(v)
594        }
595    }
596}
597
598/// Array builder for [`StringViewArray`][crate::StringViewArray]
599///
600/// Values can be appended using [`GenericByteViewBuilder::append_value`], and nulls with
601/// [`GenericByteViewBuilder::append_null`] as normal.
602///
603/// # Example
604/// ```
605/// # use arrow_array::builder::StringViewBuilder;
606/// # use arrow_array::StringViewArray;
607/// let mut builder = StringViewBuilder::new();
608/// builder.append_value("hello");
609/// builder.append_null();
610/// builder.append_value("world");
611/// let array = builder.finish();
612///
613/// let expected = vec![Some("hello"), None, Some("world")];
614/// let actual: Vec<_> = array.iter().collect();
615/// assert_eq!(expected, actual);
616/// ```
617pub type StringViewBuilder = GenericByteViewBuilder<StringViewType>;
618
619impl StringLikeArrayBuilder for StringViewBuilder {
620    fn type_name() -> &'static str {
621        std::any::type_name::<StringViewBuilder>()
622    }
623    fn with_capacity(capacity: usize) -> Self {
624        Self::with_capacity(capacity)
625    }
626    fn append_value(&mut self, value: &str) {
627        Self::append_value(self, value);
628    }
629    fn append_null(&mut self) {
630        Self::append_null(self);
631    }
632}
633
634///  Array builder for [`BinaryViewArray`][crate::BinaryViewArray]
635///
636/// Values can be appended using [`GenericByteViewBuilder::append_value`], and nulls with
637/// [`GenericByteViewBuilder::append_null`] as normal.
638///
639/// # Example
640/// ```
641/// # use arrow_array::builder::BinaryViewBuilder;
642/// use arrow_array::BinaryViewArray;
643/// let mut builder = BinaryViewBuilder::new();
644/// builder.append_value("hello");
645/// builder.append_null();
646/// builder.append_value("world");
647/// let array = builder.finish();
648///
649/// let expected: Vec<Option<&[u8]>> = vec![Some(b"hello"), None, Some(b"world")];
650/// let actual: Vec<_> = array.iter().collect();
651/// assert_eq!(expected, actual);
652/// ```
653///
654pub type BinaryViewBuilder = GenericByteViewBuilder<BinaryViewType>;
655
656impl BinaryLikeArrayBuilder for BinaryViewBuilder {
657    fn type_name() -> &'static str {
658        std::any::type_name::<BinaryViewBuilder>()
659    }
660    fn with_capacity(capacity: usize) -> Self {
661        Self::with_capacity(capacity)
662    }
663    fn append_value(&mut self, value: &[u8]) {
664        Self::append_value(self, value);
665    }
666    fn append_null(&mut self) {
667        Self::append_null(self);
668    }
669}
670
671/// Creates a view from a fixed length input (the compiler can generate
672/// specialized code for this)
673fn make_inlined_view<const LEN: usize>(data: &[u8]) -> u128 {
674    let mut view_buffer = [0; 16];
675    view_buffer[0..4].copy_from_slice(&(LEN as u32).to_le_bytes());
676    view_buffer[4..4 + LEN].copy_from_slice(&data[..LEN]);
677    u128::from_le_bytes(view_buffer)
678}
679
680/// Create a view based on the given data, block id and offset.
681///
682/// Note that the code below is carefully examined with x86_64 assembly code: <https://godbolt.org/z/685YPsd5G>
683/// The goal is to avoid calling into `ptr::copy_non_interleave`, which makes function call (i.e., not inlined),
684/// which slows down things.
685#[inline(never)]
686pub fn make_view(data: &[u8], block_id: u32, offset: u32) -> u128 {
687    let len = data.len();
688
689    // Generate specialized code for each potential small string length
690    // to improve performance
691    match len {
692        0 => make_inlined_view::<0>(data),
693        1 => make_inlined_view::<1>(data),
694        2 => make_inlined_view::<2>(data),
695        3 => make_inlined_view::<3>(data),
696        4 => make_inlined_view::<4>(data),
697        5 => make_inlined_view::<5>(data),
698        6 => make_inlined_view::<6>(data),
699        7 => make_inlined_view::<7>(data),
700        8 => make_inlined_view::<8>(data),
701        9 => make_inlined_view::<9>(data),
702        10 => make_inlined_view::<10>(data),
703        11 => make_inlined_view::<11>(data),
704        12 => make_inlined_view::<12>(data),
705        // When string is longer than 12 bytes, it can't be inlined, we create a ByteView instead.
706        _ => {
707            let view = ByteView {
708                length: len as u32,
709                // this arm only matches lengths above 12, so there are at least four bytes
710                prefix: u32::from_le_bytes([data[0], data[1], data[2], data[3]]),
711                buffer_index: block_id,
712                offset,
713            };
714            view.as_u128()
715        }
716    }
717}
718
719#[cfg(test)]
720mod tests {
721    use core::str;
722
723    use arrow_buffer::ArrowNativeType;
724
725    use super::*;
726
727    #[test]
728    fn test_string_max_deduplication_len() {
729        let value_1 = "short";
730        let value_2 = "not so similar string but long";
731        let value_3 = "1234567890123";
732
733        let max_deduplication_len = MAX_INLINE_VIEW_LEN * 2;
734
735        let mut builder = StringViewBuilder::new()
736            .with_deduplicate_strings()
737            .with_max_deduplication_len(max_deduplication_len);
738
739        assert!(value_1.len() < MAX_INLINE_VIEW_LEN.as_usize());
740        assert!(value_2.len() > max_deduplication_len.as_usize());
741        assert!(
742            value_3.len() > MAX_INLINE_VIEW_LEN.as_usize()
743                && value_3.len() < max_deduplication_len.as_usize()
744        );
745
746        // append value1 (short), expect it is inlined and not deduplicated
747        builder.append_value(value_1); // view 0
748        builder.append_value(value_1); // view 1
749        // append value2, expect second copy is not deduplicated as it exceeds max_deduplication_len
750        builder.append_value(value_2); // view 2
751        builder.append_value(value_2); // view 3
752        // append value3, expect second copy is deduplicated
753        builder.append_value(value_3); // view 4
754        builder.append_value(value_3); // view 5
755
756        let array = builder.finish();
757
758        // verify
759        let v2 = ByteView::from(array.views()[2]);
760        let v3 = ByteView::from(array.views()[3]);
761        assert_eq!(v2.buffer_index, v3.buffer_index); // stored in same buffer
762        assert_ne!(v2.offset, v3.offset); // different offsets --> not deduplicated
763
764        let v4 = ByteView::from(array.views()[4]);
765        let v5 = ByteView::from(array.views()[5]);
766        assert_eq!(v4.buffer_index, v5.buffer_index); // stored in same buffer
767        assert_eq!(v4.offset, v5.offset); // same offsets --> deduplicated
768    }
769
770    #[test]
771    fn test_string_view_deduplicate() {
772        let value_1 = "long string to test string view";
773        let value_2 = "not so similar string but long";
774
775        let mut builder = StringViewBuilder::new()
776            .with_deduplicate_strings()
777            .with_fixed_block_size(value_1.len() as u32 * 2); // so that we will have multiple buffers
778
779        let values = vec![
780            Some(value_1),
781            Some(value_2),
782            Some("short"),
783            Some(value_1),
784            None,
785            Some(value_2),
786            Some(value_1),
787        ];
788        builder.extend(values.clone());
789
790        let array = builder.finish_cloned();
791        array.to_data().validate_full().unwrap();
792        assert_eq!(array.data_buffers().len(), 1); // without duplication we would need 3 buffers.
793        let actual: Vec<_> = array.iter().collect();
794        assert_eq!(actual, values);
795
796        let view0 = array.views().first().unwrap();
797        let view3 = array.views().get(3).unwrap();
798        let view6 = array.views().get(6).unwrap();
799
800        assert_eq!(view0, view3);
801        assert_eq!(view0, view6);
802
803        assert_eq!(array.views().get(1), array.views().get(5));
804    }
805
806    #[test]
807    fn test_string_view_deduplicate_after_finish() {
808        let mut builder = StringViewBuilder::new().with_deduplicate_strings();
809
810        let value_1 = "long string to test string view";
811        let value_2 = "not so similar string but long";
812        builder.append_value(value_1);
813        let _array = builder.finish();
814        builder.append_value(value_2);
815        let _array = builder.finish();
816        builder.append_value(value_1);
817        let _array = builder.finish();
818    }
819
820    #[test]
821    fn test_string_view() {
822        let b1 = Buffer::from(b"world\xFFbananas\xF0\x9F\x98\x81");
823        let b2 = Buffer::from(b"cupcakes");
824        let b3 = Buffer::from(b"Many strings are here contained of great length and verbosity");
825
826        let mut v = StringViewBuilder::new();
827        assert_eq!(v.append_block(b1), 0);
828
829        v.append_value("This is a very long string that exceeds the inline length");
830        v.append_value("This is another very long string that exceeds the inline length");
831
832        assert_eq!(v.append_block(b2), 2);
833        assert_eq!(v.append_block(b3), 3);
834
835        // Test short strings
836        v.try_append_view(0, 0, 5).unwrap(); // world
837        v.try_append_view(0, 6, 7).unwrap(); // bananas
838        v.try_append_view(2, 3, 5).unwrap(); // cake
839        v.try_append_view(2, 0, 3).unwrap(); // cup
840        v.try_append_view(2, 0, 8).unwrap(); // cupcakes
841        v.try_append_view(0, 13, 4).unwrap(); // 😁
842        v.try_append_view(0, 13, 0).unwrap(); //
843
844        // Test longer strings
845        v.try_append_view(3, 0, 16).unwrap(); // Many strings are
846        v.try_append_view(1, 0, 19).unwrap(); // This is a very long
847        v.try_append_view(3, 13, 27).unwrap(); // here contained of great length
848
849        v.append_value("I do so like long strings");
850
851        let array = v.finish_cloned();
852        array.to_data().validate_full().unwrap();
853        assert_eq!(array.data_buffers().len(), 5);
854        let actual: Vec<_> = array.iter().flatten().collect();
855        assert_eq!(
856            actual,
857            &[
858                "This is a very long string that exceeds the inline length",
859                "This is another very long string that exceeds the inline length",
860                "world",
861                "bananas",
862                "cakes",
863                "cup",
864                "cupcakes",
865                "😁",
866                "",
867                "Many strings are",
868                "This is a very long",
869                "are here contained of great",
870                "I do so like long strings"
871            ]
872        );
873
874        let err = v.try_append_view(0, u32::MAX, 1).unwrap_err();
875        assert_eq!(
876            err.to_string(),
877            "Invalid argument error: Range 4294967295..4294967296 out of bounds for block of length 17"
878        );
879
880        let err = v.try_append_view(0, 1, u32::MAX).unwrap_err();
881        assert_eq!(
882            err.to_string(),
883            "Invalid argument error: Range 1..4294967296 out of bounds for block of length 17"
884        );
885
886        let err = v.try_append_view(0, 13, 2).unwrap_err();
887        assert_eq!(err.to_string(), "Invalid argument error: Invalid view data");
888
889        let err = v.try_append_view(0, 40, 0).unwrap_err();
890        assert_eq!(
891            err.to_string(),
892            "Invalid argument error: Range 40..40 out of bounds for block of length 17"
893        );
894
895        let err = v.try_append_view(5, 0, 0).unwrap_err();
896        assert_eq!(
897            err.to_string(),
898            "Invalid argument error: No block found with index 5"
899        );
900    }
901
902    #[test]
903    fn test_string_view_with_block_size_growth() {
904        let mut exp_builder = StringViewBuilder::new();
905        let mut fixed_builder = StringViewBuilder::new().with_fixed_block_size(STARTING_BLOCK_SIZE);
906
907        let long_string = str::from_utf8(&[b'a'; STARTING_BLOCK_SIZE as usize]).unwrap();
908
909        for i in 0..9 {
910            // 8k, 16k, 32k, 64k, 128k, 256k, 512k, 1M, 2M
911            for _ in 0..(2_u32.pow(i)) {
912                exp_builder.append_value(long_string);
913                fixed_builder.append_value(long_string);
914            }
915            exp_builder.flush_in_progress();
916            fixed_builder.flush_in_progress();
917
918            // Every step only add one buffer, but the buffer size is much larger
919            assert_eq!(exp_builder.completed.len(), i as usize + 1);
920            assert_eq!(
921                exp_builder.completed[i as usize].len(),
922                STARTING_BLOCK_SIZE as usize * 2_usize.pow(i)
923            );
924
925            // This step we added 2^i blocks, the sum of blocks should be 2^(i+1) - 1
926            assert_eq!(fixed_builder.completed.len(), 2_usize.pow(i + 1) - 1);
927
928            // Every buffer is fixed size
929            assert!(
930                fixed_builder
931                    .completed
932                    .iter()
933                    .all(|b| b.len() == STARTING_BLOCK_SIZE as usize)
934            );
935        }
936
937        // Add one more value, and the buffer stop growing.
938        exp_builder.append_value(long_string);
939        exp_builder.flush_in_progress();
940        assert_eq!(
941            exp_builder.completed.last().unwrap().capacity(),
942            MAX_BLOCK_SIZE as usize
943        );
944    }
945
946    #[test]
947    fn test_append_value_n() {
948        // Test with inline strings (<=12 bytes)
949        let mut builder = StringViewBuilder::new();
950
951        builder.try_append_value_n("hello", 100).unwrap();
952        builder.append_value("world");
953        builder.try_append_value_n("foo", 50).unwrap();
954
955        let array = builder.finish();
956        assert_eq!(array.len(), 151);
957        assert_eq!(array.null_count(), 0);
958
959        // Verify the values
960        for i in 0..100 {
961            assert_eq!(array.value(i), "hello");
962        }
963        assert_eq!(array.value(100), "world");
964        for i in 101..151 {
965            assert_eq!(array.value(i), "foo");
966        }
967
968        // All inline strings should have no data buffers
969        assert_eq!(array.data_buffers().len(), 0);
970    }
971
972    #[test]
973    fn test_append_value_n_with_deduplication() {
974        let long_string = "This is a very long string that exceeds the inline length";
975
976        // Test with deduplication enabled
977        let mut builder = StringViewBuilder::new().with_deduplicate_strings();
978
979        // First append the string once to add it to the hash map
980        builder.append_value(long_string);
981
982        // Then append_n the same string - should deduplicate and reuse the existing value
983        builder.try_append_value_n(long_string, 999).unwrap();
984
985        let array = builder.finish();
986        assert_eq!(array.len(), 1000);
987        assert_eq!(array.null_count(), 0);
988
989        // Verify all values are the same
990        for i in 0..1000 {
991            assert_eq!(array.value(i), long_string);
992        }
993
994        // With deduplication, should only have 1 data buffer containing the string once
995        assert_eq!(array.data_buffers().len(), 1);
996
997        // All views should be identical
998        let first_view = array.views()[0];
999        for view in array.views() {
1000            assert_eq!(*view, first_view);
1001        }
1002    }
1003
1004    #[test]
1005    fn test_append_value_n_zero() {
1006        let mut builder = StringViewBuilder::new();
1007
1008        builder.append_value("first");
1009        builder.try_append_value_n("should not appear", 0).unwrap();
1010        builder.append_value("second");
1011
1012        let array = builder.finish();
1013        assert_eq!(array.len(), 2);
1014        assert_eq!(array.value(0), "first");
1015        assert_eq!(array.value(1), "second");
1016    }
1017}
arrow_array/builder/generic_bytes_view_builder.rs

arrow_array/builder/
generic_bytes_view_builder.rs
