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    pub fn append_block(&mut self, buffer: Buffer) -> u32 {
195        assert!(buffer.len() < u32::MAX as usize);
196
197        self.flush_in_progress();
198        let offset = self.completed.len();
199        self.push_completed(buffer);
200        offset as u32
201    }
202
203    /// Append a view of the given `block`, `offset` and `length`
204    ///
205    /// # Safety
206    /// (1) The block must have been added using [`Self::append_block`]
207    /// (2) The range `offset..offset+length` must be within the bounds of the block
208    /// (3) The data in the block must be valid of type `T`
209    pub unsafe fn append_view_unchecked(&mut self, block: u32, offset: u32, len: u32) {
210        let b = unsafe { self.completed.get_unchecked(block as usize) };
211        let start = offset as usize;
212        let end = start.saturating_add(len as usize);
213        let b = unsafe { b.get_unchecked(start..end) };
214
215        let view = make_view(b, block, offset);
216        self.views_buffer.push(view);
217        self.null_buffer_builder.append_non_null();
218    }
219
220    /// Appends an array to the builder.
221    /// This will flush any in-progress block and append the data buffers
222    /// and add the (adapted) views.
223    pub fn append_array(&mut self, array: &GenericByteViewArray<T>) {
224        self.flush_in_progress();
225        // keep original views if this array is the first to be added or if there are no data buffers (all inline views)
226        let keep_views = self.completed.is_empty() || array.data_buffers().is_empty();
227        let starting_buffer = self.completed.len() as u32;
228
229        self.completed.extend(array.data_buffers().iter().cloned());
230
231        if keep_views {
232            self.views_buffer.extend_from_slice(array.views());
233        } else {
234            self.views_buffer.extend(array.views().iter().map(|v| {
235                let mut byte_view = ByteView::from(*v);
236                if byte_view.length > MAX_INLINE_VIEW_LEN {
237                    // Small views (<=12 bytes) are inlined, so only need to update large views
238                    byte_view.buffer_index += starting_buffer;
239                };
240
241                byte_view.as_u128()
242            }));
243        }
244
245        if let Some(null_buffer) = array.nulls() {
246            self.null_buffer_builder.append_buffer(null_buffer);
247        } else {
248            self.null_buffer_builder.append_n_non_nulls(array.len());
249        }
250    }
251
252    /// Try to append a view of the given `block`, `offset` and `length`
253    ///
254    /// See [`Self::append_block`]
255    pub fn try_append_view(&mut self, block: u32, offset: u32, len: u32) -> Result<(), ArrowError> {
256        let b = self.completed.get(block as usize).ok_or_else(|| {
257            ArrowError::InvalidArgumentError(format!("No block found with index {block}"))
258        })?;
259        let start = offset as usize;
260        let end = start.saturating_add(len as usize);
261
262        let b = b.get(start..end).ok_or_else(|| {
263            ArrowError::InvalidArgumentError(format!(
264                "Range {start}..{end} out of bounds for block of length {}",
265                b.len()
266            ))
267        })?;
268
269        if T::Native::from_bytes_checked(b).is_none() {
270            return Err(ArrowError::InvalidArgumentError(
271                "Invalid view data".to_string(),
272            ));
273        }
274
275        unsafe {
276            self.append_view_unchecked(block, offset, len);
277        }
278        Ok(())
279    }
280
281    /// Flushes the in progress block if any
282    #[inline]
283    fn flush_in_progress(&mut self) {
284        if !self.in_progress.is_empty() {
285            let f = Buffer::from_vec(std::mem::take(&mut self.in_progress));
286            self.push_completed(f)
287        }
288    }
289
290    /// Append a block to `self.completed`, checking for overflow
291    #[inline]
292    fn push_completed(&mut self, block: Buffer) {
293        assert!(block.len() < u32::MAX as usize, "Block too large");
294        assert!(self.completed.len() < u32::MAX as usize, "Too many blocks");
295        self.completed.push(block);
296    }
297
298    /// Returns the value at the given index
299    /// Useful if we want to know what value has been inserted to the builder
300    /// The index has to be smaller than `self.len()`, otherwise it will panic
301    pub fn get_value(&self, index: usize) -> &[u8] {
302        let view = self.views_buffer.as_slice().get(index).unwrap();
303        let len = *view as u32;
304        if len <= MAX_INLINE_VIEW_LEN {
305            // # Safety
306            // The view is valid from the builder
307            unsafe { GenericByteViewArray::<T>::inline_value(view, len as usize) }
308        } else {
309            let view = ByteView::from(*view);
310            if view.buffer_index < self.completed.len() as u32 {
311                let block = &self.completed[view.buffer_index as usize];
312                &block[view.offset as usize..view.offset as usize + view.length as usize]
313            } else {
314                &self.in_progress[view.offset as usize..view.offset as usize + view.length as usize]
315            }
316        }
317    }
318
319    /// Appends a value into the builder
320    ///
321    /// # Panics
322    ///
323    /// Panics if
324    /// - String buffer count exceeds `u32::MAX`
325    /// - String length exceeds `u32::MAX`
326    #[inline]
327    pub fn append_value(&mut self, value: impl AsRef<T::Native>) {
328        self.try_append_value(value).unwrap()
329    }
330
331    /// Appends a value into the builder
332    ///
333    /// # Errors
334    ///
335    /// Returns an error if:
336    /// - String buffer count exceeds `u32::MAX`
337    /// - String length exceeds `u32::MAX`
338    #[inline]
339    pub fn try_append_value(&mut self, value: impl AsRef<T::Native>) -> Result<(), ArrowError> {
340        let v: &[u8] = value.as_ref().as_ref();
341        let length: u32 = v.len().try_into().map_err(|_| {
342            ArrowError::InvalidArgumentError(format!("String length {} exceeds u32::MAX", v.len()))
343        })?;
344
345        if length <= MAX_INLINE_VIEW_LEN {
346            let mut view_buffer = [0; 16];
347            view_buffer[0..4].copy_from_slice(&length.to_le_bytes());
348            view_buffer[4..4 + v.len()].copy_from_slice(v);
349            self.views_buffer.push(u128::from_le_bytes(view_buffer));
350            self.null_buffer_builder.append_non_null();
351            return Ok(());
352        }
353
354        // Deduplication if:
355        // (1) deduplication is enabled.
356        // (2) len > `MAX_INLINE_VIEW_LEN` and len <= `max_deduplication_len`
357        let can_deduplicate = self.string_tracker.is_some()
358            && self
359                .max_deduplication_len
360                .map(|max_length| length <= max_length)
361                .unwrap_or(true);
362        if can_deduplicate {
363            if let Some((mut ht, hasher)) = self.string_tracker.take() {
364                let hash_val = hasher.hash_one(v);
365                let hasher_fn = |v: &_| hasher.hash_one(v);
366
367                let entry = ht.entry(
368                    hash_val,
369                    |idx| {
370                        let stored_value = self.get_value(*idx);
371                        v == stored_value
372                    },
373                    hasher_fn,
374                );
375                match entry {
376                    Entry::Occupied(occupied) => {
377                        // If the string already exists, we will directly use the view
378                        let idx = occupied.get();
379                        self.views_buffer.push(self.views_buffer[*idx]);
380                        self.null_buffer_builder.append_non_null();
381                        self.string_tracker = Some((ht, hasher));
382                        return Ok(());
383                    }
384                    Entry::Vacant(vacant) => {
385                        // o.w. we insert the (string hash -> view index)
386                        // the idx is current length of views_builder, as we are inserting a new view
387                        vacant.insert(self.views_buffer.len());
388                    }
389                }
390                self.string_tracker = Some((ht, hasher));
391            }
392        }
393
394        let required_cap = self.in_progress.len() + v.len();
395        if self.in_progress.capacity() < required_cap {
396            self.flush_in_progress();
397            let to_reserve = v.len().max(self.block_size.next_size() as usize);
398            self.in_progress.reserve(to_reserve);
399        };
400
401        let offset = self.in_progress.len() as u32;
402        self.in_progress.extend_from_slice(v);
403
404        let buffer_index: u32 = self.completed.len().try_into().map_err(|_| {
405            ArrowError::InvalidArgumentError(format!(
406                "Buffer count {} exceeds u32::MAX",
407                self.completed.len()
408            ))
409        })?;
410
411        let view = ByteView {
412            length,
413            // This won't panic as we checked the length of prefix earlier.
414            prefix: u32::from_le_bytes(v[0..4].try_into().unwrap()),
415            buffer_index,
416            offset,
417        };
418        self.views_buffer.push(view.into());
419        self.null_buffer_builder.append_non_null();
420
421        Ok(())
422    }
423
424    /// Append an `Option` value into the builder
425    #[inline]
426    pub fn append_option(&mut self, value: Option<impl AsRef<T::Native>>) {
427        match value {
428            None => self.append_null(),
429            Some(v) => self.append_value(v),
430        };
431    }
432
433    /// Append a null value into the builder
434    #[inline]
435    pub fn append_null(&mut self) {
436        self.null_buffer_builder.append_null();
437        self.views_buffer.push(0);
438    }
439
440    /// Builds the [`GenericByteViewArray`] and reset this builder
441    pub fn finish(&mut self) -> GenericByteViewArray<T> {
442        self.flush_in_progress();
443        let completed = std::mem::take(&mut self.completed);
444        let nulls = self.null_buffer_builder.finish();
445        if let Some((ht, _)) = self.string_tracker.as_mut() {
446            ht.clear();
447        }
448        let views = std::mem::take(&mut self.views_buffer);
449        // SAFETY: valid by construction
450        unsafe { GenericByteViewArray::new_unchecked(views.into(), completed, nulls) }
451    }
452
453    /// Builds the [`GenericByteViewArray`] without resetting the builder
454    pub fn finish_cloned(&self) -> GenericByteViewArray<T> {
455        let mut completed = self.completed.clone();
456        if !self.in_progress.is_empty() {
457            completed.push(Buffer::from_slice_ref(&self.in_progress));
458        }
459        let len = self.views_buffer.len();
460        let views = Buffer::from_slice_ref(self.views_buffer.as_slice());
461        let views = ScalarBuffer::new(views, 0, len);
462        let nulls = self.null_buffer_builder.finish_cloned();
463        // SAFETY: valid by construction
464        unsafe { GenericByteViewArray::new_unchecked(views, completed, nulls) }
465    }
466
467    /// Returns the current null buffer as a slice
468    pub fn validity_slice(&self) -> Option<&[u8]> {
469        self.null_buffer_builder.as_slice()
470    }
471
472    /// Return the allocated size of this builder in bytes, useful for memory accounting.
473    pub fn allocated_size(&self) -> usize {
474        let views = self.views_buffer.capacity() * std::mem::size_of::<u128>();
475        let null = self.null_buffer_builder.allocated_size();
476        let buffer_size = self.completed.iter().map(|b| b.capacity()).sum::<usize>();
477        let in_progress = self.in_progress.capacity();
478        let tracker = match &self.string_tracker {
479            Some((ht, _)) => ht.capacity() * std::mem::size_of::<usize>(),
480            None => 0,
481        };
482        buffer_size + in_progress + tracker + views + null
483    }
484}
485
486impl<T: ByteViewType + ?Sized> Default for GenericByteViewBuilder<T> {
487    fn default() -> Self {
488        Self::new()
489    }
490}
491
492impl<T: ByteViewType + ?Sized> std::fmt::Debug for GenericByteViewBuilder<T> {
493    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
494        write!(f, "{}ViewBuilder", T::PREFIX)?;
495        f.debug_struct("")
496            .field("views_buffer", &self.views_buffer)
497            .field("in_progress", &self.in_progress)
498            .field("completed", &self.completed)
499            .field("null_buffer_builder", &self.null_buffer_builder)
500            .finish()
501    }
502}
503
504impl<T: ByteViewType + ?Sized> ArrayBuilder for GenericByteViewBuilder<T> {
505    fn len(&self) -> usize {
506        self.null_buffer_builder.len()
507    }
508
509    fn finish(&mut self) -> ArrayRef {
510        Arc::new(self.finish())
511    }
512
513    fn finish_cloned(&self) -> ArrayRef {
514        Arc::new(self.finish_cloned())
515    }
516
517    fn as_any(&self) -> &dyn Any {
518        self
519    }
520
521    fn as_any_mut(&mut self) -> &mut dyn Any {
522        self
523    }
524
525    fn into_box_any(self: Box<Self>) -> Box<dyn Any> {
526        self
527    }
528}
529
530impl<T: ByteViewType + ?Sized, V: AsRef<T::Native>> Extend<Option<V>>
531    for GenericByteViewBuilder<T>
532{
533    #[inline]
534    fn extend<I: IntoIterator<Item = Option<V>>>(&mut self, iter: I) {
535        for v in iter {
536            self.append_option(v)
537        }
538    }
539}
540
541/// Array builder for [`StringViewArray`][crate::StringViewArray]
542///
543/// Values can be appended using [`GenericByteViewBuilder::append_value`], and nulls with
544/// [`GenericByteViewBuilder::append_null`] as normal.
545///
546/// # Example
547/// ```
548/// # use arrow_array::builder::StringViewBuilder;
549/// # use arrow_array::StringViewArray;
550/// let mut builder = StringViewBuilder::new();
551/// builder.append_value("hello");
552/// builder.append_null();
553/// builder.append_value("world");
554/// let array = builder.finish();
555///
556/// let expected = vec![Some("hello"), None, Some("world")];
557/// let actual: Vec<_> = array.iter().collect();
558/// assert_eq!(expected, actual);
559/// ```
560pub type StringViewBuilder = GenericByteViewBuilder<StringViewType>;
561
562impl StringLikeArrayBuilder for StringViewBuilder {
563    fn type_name() -> &'static str {
564        std::any::type_name::<StringViewBuilder>()
565    }
566    fn with_capacity(capacity: usize) -> Self {
567        Self::with_capacity(capacity)
568    }
569    fn append_value(&mut self, value: &str) {
570        Self::append_value(self, value);
571    }
572    fn append_null(&mut self) {
573        Self::append_null(self);
574    }
575}
576
577///  Array builder for [`BinaryViewArray`][crate::BinaryViewArray]
578///
579/// Values can be appended using [`GenericByteViewBuilder::append_value`], and nulls with
580/// [`GenericByteViewBuilder::append_null`] as normal.
581///
582/// # Example
583/// ```
584/// # use arrow_array::builder::BinaryViewBuilder;
585/// use arrow_array::BinaryViewArray;
586/// let mut builder = BinaryViewBuilder::new();
587/// builder.append_value("hello");
588/// builder.append_null();
589/// builder.append_value("world");
590/// let array = builder.finish();
591///
592/// let expected: Vec<Option<&[u8]>> = vec![Some(b"hello"), None, Some(b"world")];
593/// let actual: Vec<_> = array.iter().collect();
594/// assert_eq!(expected, actual);
595/// ```
596///
597pub type BinaryViewBuilder = GenericByteViewBuilder<BinaryViewType>;
598
599impl BinaryLikeArrayBuilder for BinaryViewBuilder {
600    fn type_name() -> &'static str {
601        std::any::type_name::<BinaryViewBuilder>()
602    }
603    fn with_capacity(capacity: usize) -> Self {
604        Self::with_capacity(capacity)
605    }
606    fn append_value(&mut self, value: &[u8]) {
607        Self::append_value(self, value);
608    }
609    fn append_null(&mut self) {
610        Self::append_null(self);
611    }
612}
613
614/// Creates a view from a fixed length input (the compiler can generate
615/// specialized code for this)
616fn make_inlined_view<const LEN: usize>(data: &[u8]) -> u128 {
617    let mut view_buffer = [0; 16];
618    view_buffer[0..4].copy_from_slice(&(LEN as u32).to_le_bytes());
619    view_buffer[4..4 + LEN].copy_from_slice(&data[..LEN]);
620    u128::from_le_bytes(view_buffer)
621}
622
623/// Create a view based on the given data, block id and offset.
624///
625/// Note that the code below is carefully examined with x86_64 assembly code: <https://godbolt.org/z/685YPsd5G>
626/// The goal is to avoid calling into `ptr::copy_non_interleave`, which makes function call (i.e., not inlined),
627/// which slows down things.
628#[inline(never)]
629pub fn make_view(data: &[u8], block_id: u32, offset: u32) -> u128 {
630    let len = data.len();
631
632    // Generate specialized code for each potential small string length
633    // to improve performance
634    match len {
635        0 => make_inlined_view::<0>(data),
636        1 => make_inlined_view::<1>(data),
637        2 => make_inlined_view::<2>(data),
638        3 => make_inlined_view::<3>(data),
639        4 => make_inlined_view::<4>(data),
640        5 => make_inlined_view::<5>(data),
641        6 => make_inlined_view::<6>(data),
642        7 => make_inlined_view::<7>(data),
643        8 => make_inlined_view::<8>(data),
644        9 => make_inlined_view::<9>(data),
645        10 => make_inlined_view::<10>(data),
646        11 => make_inlined_view::<11>(data),
647        12 => make_inlined_view::<12>(data),
648        // When string is longer than 12 bytes, it can't be inlined, we create a ByteView instead.
649        _ => {
650            let view = ByteView {
651                length: len as u32,
652                prefix: u32::from_le_bytes(data[0..4].try_into().unwrap()),
653                buffer_index: block_id,
654                offset,
655            };
656            view.as_u128()
657        }
658    }
659}
660
661#[cfg(test)]
662mod tests {
663    use core::str;
664
665    use arrow_buffer::ArrowNativeType;
666
667    use super::*;
668
669    #[test]
670    fn test_string_max_deduplication_len() {
671        let value_1 = "short";
672        let value_2 = "not so similar string but long";
673        let value_3 = "1234567890123";
674
675        let max_deduplication_len = MAX_INLINE_VIEW_LEN * 2;
676
677        let mut builder = StringViewBuilder::new()
678            .with_deduplicate_strings()
679            .with_max_deduplication_len(max_deduplication_len);
680
681        assert!(value_1.len() < MAX_INLINE_VIEW_LEN.as_usize());
682        assert!(value_2.len() > max_deduplication_len.as_usize());
683        assert!(
684            value_3.len() > MAX_INLINE_VIEW_LEN.as_usize()
685                && value_3.len() < max_deduplication_len.as_usize()
686        );
687
688        // append value1 (short), expect it is inlined and not deduplicated
689        builder.append_value(value_1); // view 0
690        builder.append_value(value_1); // view 1
691        // append value2, expect second copy is not deduplicated as it exceeds max_deduplication_len
692        builder.append_value(value_2); // view 2
693        builder.append_value(value_2); // view 3
694        // append value3, expect second copy is deduplicated
695        builder.append_value(value_3); // view 4
696        builder.append_value(value_3); // view 5
697
698        let array = builder.finish();
699
700        // verify
701        let v2 = ByteView::from(array.views()[2]);
702        let v3 = ByteView::from(array.views()[3]);
703        assert_eq!(v2.buffer_index, v3.buffer_index); // stored in same buffer
704        assert_ne!(v2.offset, v3.offset); // different offsets --> not deduplicated
705
706        let v4 = ByteView::from(array.views()[4]);
707        let v5 = ByteView::from(array.views()[5]);
708        assert_eq!(v4.buffer_index, v5.buffer_index); // stored in same buffer
709        assert_eq!(v4.offset, v5.offset); // same offsets --> deduplicated
710    }
711
712    #[test]
713    fn test_string_view_deduplicate() {
714        let value_1 = "long string to test string view";
715        let value_2 = "not so similar string but long";
716
717        let mut builder = StringViewBuilder::new()
718            .with_deduplicate_strings()
719            .with_fixed_block_size(value_1.len() as u32 * 2); // so that we will have multiple buffers
720
721        let values = vec![
722            Some(value_1),
723            Some(value_2),
724            Some("short"),
725            Some(value_1),
726            None,
727            Some(value_2),
728            Some(value_1),
729        ];
730        builder.extend(values.clone());
731
732        let array = builder.finish_cloned();
733        array.to_data().validate_full().unwrap();
734        assert_eq!(array.data_buffers().len(), 1); // without duplication we would need 3 buffers.
735        let actual: Vec<_> = array.iter().collect();
736        assert_eq!(actual, values);
737
738        let view0 = array.views().first().unwrap();
739        let view3 = array.views().get(3).unwrap();
740        let view6 = array.views().get(6).unwrap();
741
742        assert_eq!(view0, view3);
743        assert_eq!(view0, view6);
744
745        assert_eq!(array.views().get(1), array.views().get(5));
746    }
747
748    #[test]
749    fn test_string_view_deduplicate_after_finish() {
750        let mut builder = StringViewBuilder::new().with_deduplicate_strings();
751
752        let value_1 = "long string to test string view";
753        let value_2 = "not so similar string but long";
754        builder.append_value(value_1);
755        let _array = builder.finish();
756        builder.append_value(value_2);
757        let _array = builder.finish();
758        builder.append_value(value_1);
759        let _array = builder.finish();
760    }
761
762    #[test]
763    fn test_string_view() {
764        let b1 = Buffer::from(b"world\xFFbananas\xF0\x9F\x98\x81");
765        let b2 = Buffer::from(b"cupcakes");
766        let b3 = Buffer::from(b"Many strings are here contained of great length and verbosity");
767
768        let mut v = StringViewBuilder::new();
769        assert_eq!(v.append_block(b1), 0);
770
771        v.append_value("This is a very long string that exceeds the inline length");
772        v.append_value("This is another very long string that exceeds the inline length");
773
774        assert_eq!(v.append_block(b2), 2);
775        assert_eq!(v.append_block(b3), 3);
776
777        // Test short strings
778        v.try_append_view(0, 0, 5).unwrap(); // world
779        v.try_append_view(0, 6, 7).unwrap(); // bananas
780        v.try_append_view(2, 3, 5).unwrap(); // cake
781        v.try_append_view(2, 0, 3).unwrap(); // cup
782        v.try_append_view(2, 0, 8).unwrap(); // cupcakes
783        v.try_append_view(0, 13, 4).unwrap(); // 😁
784        v.try_append_view(0, 13, 0).unwrap(); //
785
786        // Test longer strings
787        v.try_append_view(3, 0, 16).unwrap(); // Many strings are
788        v.try_append_view(1, 0, 19).unwrap(); // This is a very long
789        v.try_append_view(3, 13, 27).unwrap(); // here contained of great length
790
791        v.append_value("I do so like long strings");
792
793        let array = v.finish_cloned();
794        array.to_data().validate_full().unwrap();
795        assert_eq!(array.data_buffers().len(), 5);
796        let actual: Vec<_> = array.iter().flatten().collect();
797        assert_eq!(
798            actual,
799            &[
800                "This is a very long string that exceeds the inline length",
801                "This is another very long string that exceeds the inline length",
802                "world",
803                "bananas",
804                "cakes",
805                "cup",
806                "cupcakes",
807                "😁",
808                "",
809                "Many strings are",
810                "This is a very long",
811                "are here contained of great",
812                "I do so like long strings"
813            ]
814        );
815
816        let err = v.try_append_view(0, u32::MAX, 1).unwrap_err();
817        assert_eq!(
818            err.to_string(),
819            "Invalid argument error: Range 4294967295..4294967296 out of bounds for block of length 17"
820        );
821
822        let err = v.try_append_view(0, 1, u32::MAX).unwrap_err();
823        assert_eq!(
824            err.to_string(),
825            "Invalid argument error: Range 1..4294967296 out of bounds for block of length 17"
826        );
827
828        let err = v.try_append_view(0, 13, 2).unwrap_err();
829        assert_eq!(err.to_string(), "Invalid argument error: Invalid view data");
830
831        let err = v.try_append_view(0, 40, 0).unwrap_err();
832        assert_eq!(
833            err.to_string(),
834            "Invalid argument error: Range 40..40 out of bounds for block of length 17"
835        );
836
837        let err = v.try_append_view(5, 0, 0).unwrap_err();
838        assert_eq!(
839            err.to_string(),
840            "Invalid argument error: No block found with index 5"
841        );
842    }
843
844    #[test]
845    fn test_string_view_with_block_size_growth() {
846        let mut exp_builder = StringViewBuilder::new();
847        let mut fixed_builder = StringViewBuilder::new().with_fixed_block_size(STARTING_BLOCK_SIZE);
848
849        let long_string = str::from_utf8(&[b'a'; STARTING_BLOCK_SIZE as usize]).unwrap();
850
851        for i in 0..9 {
852            // 8k, 16k, 32k, 64k, 128k, 256k, 512k, 1M, 2M
853            for _ in 0..(2_u32.pow(i)) {
854                exp_builder.append_value(long_string);
855                fixed_builder.append_value(long_string);
856            }
857            exp_builder.flush_in_progress();
858            fixed_builder.flush_in_progress();
859
860            // Every step only add one buffer, but the buffer size is much larger
861            assert_eq!(exp_builder.completed.len(), i as usize + 1);
862            assert_eq!(
863                exp_builder.completed[i as usize].len(),
864                STARTING_BLOCK_SIZE as usize * 2_usize.pow(i)
865            );
866
867            // This step we added 2^i blocks, the sum of blocks should be 2^(i+1) - 1
868            assert_eq!(fixed_builder.completed.len(), 2_usize.pow(i + 1) - 1);
869
870            // Every buffer is fixed size
871            assert!(
872                fixed_builder
873                    .completed
874                    .iter()
875                    .all(|b| b.len() == STARTING_BLOCK_SIZE as usize)
876            );
877        }
878
879        // Add one more value, and the buffer stop growing.
880        exp_builder.append_value(long_string);
881        exp_builder.flush_in_progress();
882        assert_eq!(
883            exp_builder.completed.last().unwrap().capacity(),
884            MAX_BLOCK_SIZE as usize
885        );
886    }
887}
arrow_array/builder/generic_bytes_view_builder.rs

arrow_array/builder/
generic_bytes_view_builder.rs
