arrow_select/coalesce/
byte_view.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: See the NOTICE file distributed with this work for additional information regarding copyright ownership
3// SPDX-FileContributor: Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements
4
5use crate::coalesce::{InProgressArray, check_filter_len};
6use crate::filter::{FilterPredicate, FilterSelection};
7use arrow_array::cast::AsArray;
8use arrow_array::types::ByteViewType;
9use arrow_array::{Array, ArrayRef, GenericByteViewArray};
10use arrow_buffer::{Buffer, NullBuffer, NullBufferBuilder};
11use arrow_data::{ByteView, MAX_INLINE_VIEW_LEN};
12use arrow_schema::ArrowError;
13use std::marker::PhantomData;
14use std::sync::Arc;
15
16/// InProgressArray for [`StringViewArray`] and [`BinaryViewArray`]
17///
18/// This structure buffers the views and data buffers as they are copied from
19/// the source array, and then produces a new array when `finish` is called. It
20/// also handles "garbage collection" by copying strings to a new buffer when
21/// the source buffer is sparse (i.e. uses at least 2x more than the memory it
22/// needs).
23///
24/// [`StringViewArray`]: arrow_array::StringViewArray
25/// [`BinaryViewArray`]: arrow_array::BinaryViewArray
26pub(crate) struct InProgressByteViewArray<B: ByteViewType> {
27    /// The source array and information
28    source: Option<Source>,
29    /// the target batch size (and thus size for views allocation)
30    batch_size: usize,
31    /// The in progress views (guaranteed to be valid)
32    views: Vec<u128>,
33    /// In progress nulls
34    nulls: NullBufferBuilder,
35    /// current buffer
36    current: Option<Vec<u8>>,
37    /// completed buffers (corresponding to content of views)
38    completed: Vec<Buffer>,
39    /// Allocates new buffers of increasing size as needed
40    buffer_source: BufferSource,
41    /// Phantom so we can use the same struct for both StringViewArray and
42    /// BinaryViewArray
43    _phantom: PhantomData<B>,
44    /// The size in bytes the [`Buffer`]s in [`Self::completed`] is taking
45    completed_buffers_size: usize,
46    /// The size in bytes from [`Self::source`] that it is being used in [`Self::completed`]
47    size_of_completed_buffers_from_current_source: usize,
48}
49
50struct Source {
51    /// The array to copy form
52    array: ArrayRef,
53    /// Should the strings from the source array be copied into new buffers?
54    need_gc: bool,
55    /// How many bytes were actually used in the source array's buffers?
56    ideal_buffer_size: usize,
57}
58
59// manually implement Debug because ByteViewType doesn't implement Debug
60impl<B: ByteViewType> std::fmt::Debug for InProgressByteViewArray<B> {
61    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
62        f.debug_struct("InProgressByteViewArray")
63            .field("batch_size", &self.batch_size)
64            .field("views", &self.views.len())
65            .field("nulls", &self.nulls)
66            .field("current", &self.current.as_ref().map(|_| "Some(...)"))
67            .field("completed", &self.completed.len())
68            .finish()
69    }
70}
71
72impl<B: ByteViewType> InProgressByteViewArray<B> {
73    pub(crate) fn new(batch_size: usize) -> Self {
74        let buffer_source = BufferSource::new();
75
76        Self {
77            batch_size,
78            source: None,
79            views: Vec::new(),                         // allocate in push
80            nulls: NullBufferBuilder::new(batch_size), // no allocation
81            current: None,
82            completed: vec![],
83            completed_buffers_size: 0,
84            size_of_completed_buffers_from_current_source: 0,
85            buffer_source,
86            _phantom: PhantomData,
87        }
88    }
89
90    /// Allocate space for output views and nulls if needed
91    ///
92    /// This is done on write (when we know it is necessary) rather than
93    /// eagerly to avoid allocations that are not used.
94    fn ensure_capacity(&mut self) {
95        if self.views.capacity() == 0 {
96            self.views.reserve(self.batch_size);
97        }
98    }
99
100    /// Finishes in progress buffer, if any
101    fn finish_current(&mut self) {
102        let Some(next_buffer) = self.current.take() else {
103            return;
104        };
105        let buffer: Buffer = next_buffer.into();
106
107        self.completed_buffers_size += buffer.capacity();
108        self.completed.push(buffer);
109    }
110
111    /// # Safety
112    ///
113    /// `filter.filter_len()` must not exceed `views.len()`.
114    unsafe fn append_views_by_filter(&mut self, views: &[u128], filter: &FilterPredicate) {
115        let selected_count = filter.count();
116        let current_len = self.views.len();
117        self.views.reserve(selected_count);
118
119        let mut written = 0;
120
121        // SAFETY: the caller guarantees that selected rows are in bounds,
122        // `FilterPredicate` yields `selected_count` rows, and capacity was
123        // reserved above.
124        unsafe {
125            let mut out = self.views.spare_capacity_mut().as_mut_ptr().cast::<u128>();
126
127            match filter.selection() {
128                FilterSelection::None => {}
129                FilterSelection::All { .. } => {
130                    std::ptr::copy_nonoverlapping(views.as_ptr(), out, selected_count);
131                    written = selected_count;
132                }
133                FilterSelection::Slices(slices) => {
134                    slices.for_each(|(start, end)| {
135                        let len = end - start;
136                        std::ptr::copy_nonoverlapping(views.as_ptr().add(start), out, len);
137                        out = out.add(len);
138                        written += len;
139                    });
140                }
141                FilterSelection::Indices(indices) => {
142                    indices.for_each(|idx| {
143                        out.write(*views.get_unchecked(idx));
144                        out = out.add(1);
145                        written += 1;
146                    });
147                }
148            }
149
150            self.views.set_len(current_len + written);
151        }
152
153        debug_assert_eq!(written, selected_count);
154    }
155
156    fn append_nulls_by_filter(
157        &mut self,
158        filter: &FilterPredicate,
159        source_nulls: Option<&NullBuffer>,
160    ) {
161        if let Some(nulls) = filter.filter_nulls(source_nulls) {
162            self.nulls.append_buffer(&nulls);
163        } else {
164            self.nulls.append_n_non_nulls(filter.count());
165        }
166    }
167
168    /// Append views to self.views, updating the buffer index if necessary
169    ///
170    /// # Safety
171    ///
172    /// [`B::validate`](ByteViewType::validate) must succeed for `views` and
173    /// `buffers`.
174    ///
175    /// Note: this function performs no unsafe operations itself, but it assumes
176    /// `views` are valid views for self.views and `buffers` for self.buffers
177    #[inline(never)]
178    unsafe fn append_views_and_update_buffer_index(
179        &mut self,
180        views: &[u128],
181        buffers: &[Buffer],
182        is_reused: bool,
183    ) {
184        if let Some(buffer) = self.current.take() {
185            let buffer: Buffer = buffer.into();
186            self.completed_buffers_size += buffer.capacity();
187            self.completed.push(buffer);
188        }
189
190        let buffers_size = buffers.iter().map(|b| b.capacity()).sum::<usize>();
191        if !is_reused {
192            self.completed_buffers_size += buffers_size;
193        } else if self.size_of_completed_buffers_from_current_source == 0 {
194            // Don't double count buffers size if already counted that
195            self.size_of_completed_buffers_from_current_source += buffers_size;
196        }
197
198        let starting_buffer = self.completed.len() as u32;
199        self.completed.extend_from_slice(buffers);
200        assert!(
201            u32::try_from(self.completed.len()).is_ok(),
202            "too many buffers"
203        );
204
205        if starting_buffer == 0 {
206            // If there are no buffers, we can just use the views as is
207            self.views.extend_from_slice(views);
208        } else {
209            // If there are buffers, we need to update the buffer index
210            let updated_views = views.iter().map(|v| {
211                // A view of a string of 12 bytes or fewer stores the string
212                // inline and has no buffer index, so adding 0 leaves it
213                // unchanged.
214                let offset = if *v as u32 > MAX_INLINE_VIEW_LEN {
215                    starting_buffer
216                } else {
217                    0
218                };
219                // The buffer index is bits 64 to 95 of a view. Adding
220                // `starting_buffer` gives the string's buffer index in the
221                // output, which fits in 32 bits, so the sum never carries into
222                // bit 96.
223                v + ((offset as u128) << 64)
224            });
225
226            self.views.extend(updated_views);
227        }
228    }
229
230    /// Append views to self.views, copying data from the buffers into
231    /// self.buffers and updating the buffer index as necessary.
232    ///
233    /// # Arguments
234    /// - `views` - the views to append
235    /// - `view_buffer_size` - the total number of bytes pointed to by all
236    ///   views (used to allocate new buffers if needed)
237    /// - `buffers` - the buffers the views point to
238    ///
239    /// # Safety
240    ///
241    /// [`B::validate`](ByteViewType::validate) must succeed for `views` and
242    /// `buffers`.
243    #[inline(never)]
244    unsafe fn append_views_and_copy_strings(
245        &mut self,
246        views: &[u128],
247        view_buffer_size: usize,
248        buffers: &[Buffer],
249    ) {
250        // Note: the calculations below are designed to avoid any reallocations
251        // of the current buffer, and to only allocate new buffers when
252        // necessary, which is critical for performance.
253
254        // If there is no current buffer, allocate a new one
255        let Some(current) = self.current.take() else {
256            let new_buffer = self.buffer_source.next_buffer(view_buffer_size);
257            // SAFETY: the caller guarantees the views and buffers are valid.
258            unsafe { self.append_views_and_copy_strings_inner(views, new_buffer, buffers) };
259            return;
260        };
261
262        // If there is a current buffer with enough space, append the views and
263        // copy the strings into the existing buffer.
264        let mut remaining_capacity = current.capacity() - current.len();
265        if view_buffer_size <= remaining_capacity {
266            // SAFETY: the caller guarantees the views and buffers are valid.
267            unsafe { self.append_views_and_copy_strings_inner(views, current, buffers) };
268            return;
269        }
270
271        // Here there is a current buffer, but it doesn't have enough space to
272        // hold all the strings. Copy as many views as we can into the current
273        // buffer and then allocate a new buffer for the remaining views
274        //
275        // TODO: should we copy the strings too at the same time?
276        let mut num_view_to_current = 0;
277        for view in views {
278            let b = ByteView::from(*view);
279            let str_len = b.length;
280            if remaining_capacity < str_len as usize {
281                break;
282            }
283            if str_len > MAX_INLINE_VIEW_LEN {
284                remaining_capacity -= str_len as usize;
285            }
286            num_view_to_current += 1;
287        }
288
289        let first_views = &views[0..num_view_to_current];
290        let string_bytes_to_copy = current.capacity() - current.len() - remaining_capacity;
291        let remaining_view_buffer_size = view_buffer_size - string_bytes_to_copy;
292
293        // SAFETY: `first_views` is a subset of the caller-validated views.
294        unsafe { self.append_views_and_copy_strings_inner(first_views, current, buffers) };
295        let completed: Buffer = self.current.take().expect("completed").into();
296        self.completed_buffers_size += completed.capacity();
297        self.completed.push(completed);
298
299        // Copy any remaining views into a new buffer
300        let remaining_views = &views[num_view_to_current..];
301        let new_buffer = self.buffer_source.next_buffer(remaining_view_buffer_size);
302        // SAFETY: `remaining_views` is a subset of the caller-validated views.
303        unsafe { self.append_views_and_copy_strings_inner(remaining_views, new_buffer, buffers) };
304    }
305
306    /// Append views to self.views, copying data from the buffers into
307    /// dst_buffer, which is then set as self.current
308    ///
309    /// # Panics:
310    /// If `self.current` is `Some`
311    ///
312    /// # Safety
313    ///
314    /// [`B::validate`](ByteViewType::validate) must succeed for `views` and
315    /// `buffers`.
316    ///
317    /// See `append_views_and_copy_strings` for more details
318    #[inline(never)]
319    unsafe fn append_views_and_copy_strings_inner(
320        &mut self,
321        views: &[u128],
322        mut dst_buffer: Vec<u8>,
323        buffers: &[Buffer],
324    ) {
325        assert!(self.current.is_none(), "current buffer should be None");
326
327        if views.is_empty() {
328            self.current = Some(dst_buffer);
329            return;
330        }
331
332        let new_buffer_index: u32 = self.completed.len().try_into().expect("too many buffers");
333
334        // In debug builds, check that the vector has enough capacity to copy
335        // the views into it without reallocating.
336        #[cfg(debug_assertions)]
337        {
338            let total_length: usize = views
339                .iter()
340                .filter_map(|v| {
341                    let b = ByteView::from(*v);
342                    if b.length > MAX_INLINE_VIEW_LEN {
343                        Some(b.length as usize)
344                    } else {
345                        None
346                    }
347                })
348                .sum();
349            debug_assert!(
350                dst_buffer.capacity() >= total_length,
351                "dst_buffer capacity {} is less than total length {}",
352                dst_buffer.capacity(),
353                total_length
354            );
355        }
356
357        // Copy the views, updating the buffer index and copying the data as needed
358        let new_views = views.iter().map(|v| {
359            let mut b: ByteView = ByteView::from(*v);
360            if b.length > MAX_INLINE_VIEW_LEN {
361                let buffer_index = b.buffer_index as usize;
362                let buffer_offset = b.offset as usize;
363                let str_len = b.length as usize;
364
365                // Update view to location in current
366                b.offset = dst_buffer.len() as u32;
367                b.buffer_index = new_buffer_index;
368
369                // safety: input views are validly constructed
370                let src = unsafe {
371                    buffers
372                        .get_unchecked(buffer_index)
373                        .get_unchecked(buffer_offset..buffer_offset + str_len)
374                };
375                dst_buffer.extend_from_slice(src);
376            }
377            b.as_u128()
378        });
379
380        self.views.extend(new_views);
381        self.current = Some(dst_buffer);
382    }
383}
384
385impl<B: ByteViewType> InProgressArray for InProgressByteViewArray<B> {
386    fn set_source(&mut self, source: Option<ArrayRef>) {
387        // If used values from source, add only the size that was used
388        self.completed_buffers_size += self.size_of_completed_buffers_from_current_source;
389        self.size_of_completed_buffers_from_current_source = 0;
390
391        self.source = source.map(|array| {
392            let s = array.as_byte_view::<B>();
393
394            let (need_gc, ideal_buffer_size) = if s.data_buffers().is_empty() {
395                (false, 0)
396            } else {
397                let ideal_buffer_size = s.total_buffer_bytes_used();
398                // copying strings is expensive, so only do it if the array is
399                // sparse (uses at least 2x the memory it needs)
400                let need_gc = if ideal_buffer_size != 0 {
401                    let mut is_sparse = false;
402                    let mut curr_buffer_size = 0;
403                    for buffer in s.data_buffers().iter() {
404                        curr_buffer_size += buffer.capacity();
405                        if curr_buffer_size > ideal_buffer_size * 2 {
406                            is_sparse = true;
407                            break;
408                        }
409                    }
410                    is_sparse
411                } else {
412                    false
413                };
414                (need_gc, ideal_buffer_size)
415            };
416
417            Source {
418                array,
419                need_gc,
420                ideal_buffer_size,
421            }
422        });
423    }
424
425    fn copy_rows(&mut self, offset: usize, len: usize) -> Result<(), ArrowError> {
426        self.ensure_capacity();
427        let source = self.source.take().ok_or_else(|| {
428            ArrowError::InvalidArgumentError(
429                "Internal Error: InProgressByteViewArray: source not set".to_string(),
430            )
431        })?;
432
433        // If creating StringViewArray output, ensure input was valid utf8 too
434        let s = source.array.as_byte_view::<B>();
435
436        // add any nulls, as necessary
437        if let Some(nulls) = s.nulls().as_ref() {
438            let nulls = nulls.slice(offset, len);
439            self.nulls.append_buffer(&nulls);
440        } else {
441            self.nulls.append_n_non_nulls(len);
442        }
443
444        let buffers = s.data_buffers();
445        let views = &s.views().as_ref()[offset..offset + len];
446
447        // If there are no data buffers in s (all inlined views), can append the
448        // views/nulls and done
449        if source.ideal_buffer_size == 0 {
450            self.views.extend_from_slice(views);
451            self.source = Some(source);
452            return Ok(());
453        }
454
455        // Copying the strings into a buffer can be time-consuming so
456        // only do it if the array is sparse
457        if source.need_gc {
458            // SAFETY: `views` and `buffers` come from the same valid byte-view array.
459            unsafe { self.append_views_and_copy_strings(views, source.ideal_buffer_size, buffers) };
460        } else {
461            // SAFETY: `views` and `buffers` come from the same valid byte-view array.
462            unsafe { self.append_views_and_update_buffer_index(views, buffers, true) };
463        }
464        self.source = Some(source);
465        Ok(())
466    }
467
468    fn copy_rows_by_filter(&mut self, filter: &FilterPredicate) -> Result<(), ArrowError> {
469        self.ensure_capacity();
470        let source = self.source.take().ok_or_else(|| {
471            ArrowError::InvalidArgumentError(
472                "Internal Error: InProgressByteViewArray: source not set".to_string(),
473            )
474        })?;
475
476        let s = source.array.as_byte_view::<B>();
477
478        if !s.data_buffers().is_empty() {
479            // Restore the source taken above before returning the guard error.
480            self.source = Some(source);
481            return Err(ArrowError::InvalidArgumentError(
482                "Internal Error: InProgressByteViewArray::copy_rows_by_filter requires inline views"
483                    .to_string(),
484            ));
485        }
486
487        if let Err(e) = check_filter_len(filter, s.len()) {
488            // Restore the source taken above before returning the error.
489            self.source = Some(source);
490            return Err(e);
491        }
492
493        self.append_nulls_by_filter(filter, s.nulls());
494        // SAFETY: `check_filter_len` verified the filter is no longer than
495        // `s.views()`.
496        unsafe { self.append_views_by_filter(s.views(), filter) };
497
498        self.source = Some(source);
499        Ok(())
500    }
501
502    fn copy_rows_by_filter_from(
503        &mut self,
504        source: ArrayRef,
505        filter: &FilterPredicate,
506    ) -> Result<(), ArrowError> {
507        let s = source.as_byte_view::<B>();
508        // The source views reference no external buffers, so they must all be
509        // inline and we can copy just the nulls and views.
510        if s.data_buffers().is_empty() {
511            check_filter_len(filter, s.len())?;
512            self.ensure_capacity();
513            self.append_nulls_by_filter(filter, s.nulls());
514            // SAFETY: `check_filter_len` verified the filter is no longer than
515            // `s.views()`.
516            unsafe { self.append_views_by_filter(s.views(), filter) };
517            return Ok(());
518        }
519
520        // The views reference external buffers, so match the filter kernel:
521        // filter the views/nulls, but reuse the source's data buffers rather
522        // than copying the referenced string data.
523        let filtered = filter.filter(source.as_ref())?;
524        let filtered = filtered.as_byte_view::<B>();
525
526        self.ensure_capacity();
527        if let Some(nulls) = filtered.nulls().as_ref() {
528            self.nulls.append_buffer(nulls);
529        } else {
530            self.nulls.append_n_non_nulls(filter.count());
531        }
532        // SAFETY: the views and buffers come from the same valid filtered array.
533        unsafe {
534            self.append_views_and_update_buffer_index(
535                filtered.views(),
536                filtered.data_buffers(),
537                false,
538            )
539        };
540        Ok(())
541    }
542
543    fn finish(&mut self) -> Result<ArrayRef, ArrowError> {
544        self.finish_current();
545        assert!(self.current.is_none());
546        let buffers = std::mem::take(&mut self.completed);
547        let views = std::mem::take(&mut self.views);
548        let nulls = self.nulls.finish();
549        self.nulls = NullBufferBuilder::new(self.batch_size);
550
551        // Not reusing anything since we took all complete
552        self.size_of_completed_buffers_from_current_source = 0;
553        self.completed_buffers_size = 0;
554
555        // Safety: we created valid views and buffers above and the
556        // input arrays had value data and nulls
557        let new_array = unsafe {
558            GenericByteViewArray::<B>::new_unchecked(views.into(), buffers.into(), nulls)
559        };
560        Ok(Arc::new(new_array))
561    }
562
563    fn size(&self) -> usize {
564        self.completed_buffers_size
565            + self.current.as_ref().map_or(0, |c| c.capacity())
566            + self.nulls.allocated_size()
567            + self.views.capacity() * size_of::<u128>()
568            + self
569                .source
570                .as_ref()
571                .map_or(0, |s| s.array.get_array_memory_size())
572    }
573}
574
575const STARTING_BLOCK_SIZE: usize = 4 * 1024; // (note the first size used is actually 8KiB)
576const MAX_BLOCK_SIZE: usize = 1024 * 1024; // 1MiB
577
578/// Manages allocating new buffers for `StringViewArray` in increasing sizes
579#[derive(Debug)]
580struct BufferSource {
581    current_size: usize,
582}
583
584impl BufferSource {
585    fn new() -> Self {
586        Self {
587            current_size: STARTING_BLOCK_SIZE,
588        }
589    }
590
591    /// Return a new buffer, with a capacity of at least `min_size`
592    fn next_buffer(&mut self, min_size: usize) -> Vec<u8> {
593        let size = self.next_size(min_size);
594        Vec::with_capacity(size)
595    }
596
597    fn next_size(&mut self, min_size: usize) -> usize {
598        if self.current_size < MAX_BLOCK_SIZE {
599            // If the current size is less than the max size, we can double it
600            // we have fixed start/end block sizes, so we can't overflow
601            self.current_size = self.current_size.saturating_mul(2);
602        }
603        if self.current_size >= min_size {
604            self.current_size
605        } else {
606            // increase next size until we hit min_size or max  size
607            while self.current_size <= min_size && self.current_size < MAX_BLOCK_SIZE {
608                self.current_size = self.current_size.saturating_mul(2);
609            }
610            self.current_size.max(min_size)
611        }
612    }
613}
614
615#[cfg(test)]
616mod tests {
617    use super::*;
618    use crate::filter::FilterBuilder;
619    use arrow_array::types::BinaryViewType;
620    use arrow_array::{BinaryViewArray, BooleanArray};
621
622    #[test]
623    fn test_buffer_source() {
624        let mut source = BufferSource::new();
625        assert_eq!(source.next_buffer(1000).capacity(), 8192);
626        assert_eq!(source.next_buffer(1000).capacity(), 16384);
627        assert_eq!(source.next_buffer(1000).capacity(), 32768);
628        assert_eq!(source.next_buffer(1000).capacity(), 65536);
629        assert_eq!(source.next_buffer(1000).capacity(), 131072);
630        assert_eq!(source.next_buffer(1000).capacity(), 262144);
631        assert_eq!(source.next_buffer(1000).capacity(), 524288);
632        assert_eq!(source.next_buffer(1000).capacity(), 1024 * 1024);
633        // clamped to max size
634        assert_eq!(source.next_buffer(1000).capacity(), 1024 * 1024);
635        // Can override with larger size request
636        assert_eq!(source.next_buffer(10_000_000).capacity(), 10_000_000);
637    }
638
639    #[test]
640    fn test_buffer_source_with_min_small() {
641        let mut source = BufferSource::new();
642        // First buffer should be 8kb
643        assert_eq!(source.next_buffer(5_600).capacity(), 8 * 1024);
644        // then 16kb
645        assert_eq!(source.next_buffer(5_600).capacity(), 16 * 1024);
646        // then 32kb
647        assert_eq!(source.next_buffer(5_600).capacity(), 32 * 1024);
648    }
649
650    #[test]
651    fn test_buffer_source_with_min_large() {
652        let mut source = BufferSource::new();
653        assert_eq!(source.next_buffer(500_000).capacity(), 512 * 1024);
654        assert_eq!(source.next_buffer(500_000).capacity(), 1024 * 1024);
655        // clamped to max size
656        assert_eq!(source.next_buffer(500_000).capacity(), 1024 * 1024);
657        // Can override with larger size request
658        assert_eq!(source.next_buffer(2_000_000).capacity(), 2_000_000);
659    }
660
661    #[test]
662    fn test_copy_rows_by_filter_rejects_non_inline_views() {
663        let values: Vec<Option<&[u8]>> = vec![Some(b"This value is longer than 12 bytes")];
664        let array = BinaryViewArray::from_iter(values);
665        assert!(!array.data_buffers().is_empty());
666
667        let mut in_progress = InProgressByteViewArray::<BinaryViewType>::new(1);
668        in_progress.set_source(Some(Arc::new(array)));
669
670        let filter = BooleanArray::from(vec![true]);
671        let predicate = FilterBuilder::new(&filter).build();
672        let err = in_progress.copy_rows_by_filter(&predicate).unwrap_err();
673
674        assert!(
675            err.to_string().contains("requires inline views"),
676            "unexpected error: {err}"
677        );
678    }
679
680    #[test]
681    fn test_copy_rows_by_filter_from_reuses_non_inline_buffers() {
682        let values = (0..32)
683            .map(|i| format!("This value is longer than 12 bytes: {i}").into_bytes())
684            .collect::<Vec<_>>();
685        let array = BinaryViewArray::from_iter(values.iter().map(|v| Some(v.as_slice())));
686        assert!(!array.data_buffers().is_empty());
687        let source_buffer = array.data_buffers()[0].as_ptr();
688
689        let filter = BooleanArray::from((0..32).map(|i| i == 3 || i == 29).collect::<Vec<_>>());
690        let predicate = FilterBuilder::new(&filter).build();
691
692        let mut in_progress = InProgressByteViewArray::<BinaryViewType>::new(32);
693        in_progress
694            .copy_rows_by_filter_from(Arc::new(array), &predicate)
695            .unwrap();
696        let output = in_progress.finish().unwrap();
697        let output = output.as_binary_view();
698
699        assert_eq!(output.len(), 2);
700        assert_eq!(output.value(0), values[3].as_slice());
701        assert_eq!(output.value(1), values[29].as_slice());
702        assert!(
703            output
704                .data_buffers()
705                .iter()
706                .any(|buffer| std::ptr::addr_eq(buffer.as_ptr(), source_buffer)),
707            "expected filtered output to reuse the source data buffer"
708        );
709    }
710
711    /// Build a compacted BinaryViewArray whose values all spill into external
712    /// data buffers. `gc()` makes the buffers dense so the coalescer reuses them
713    /// (`need_gc == false`) rather than copying/compacting them.
714    fn non_inline_array(n: usize) -> (BinaryViewArray, usize) {
715        let values = (0..n)
716            .map(|i| format!("This value is longer than 12 bytes: {i}").into_bytes())
717            .collect::<Vec<_>>();
718        let array = BinaryViewArray::from_iter(values.iter().map(|v| Some(v.as_slice()))).gc();
719        assert!(!array.data_buffers().is_empty());
720        let buffer_capacity = array.data_buffers().iter().map(|b| b.capacity()).sum();
721        (array, buffer_capacity)
722    }
723
724    #[test]
725    fn test_size_empty() {
726        let in_progress = InProgressByteViewArray::<BinaryViewType>::new(64);
727        assert_eq!(in_progress.size(), 0);
728    }
729
730    #[test]
731    fn test_size_reused_buffers_not_double_counted() {
732        let (array, buffer_capacity) = non_inline_array(64);
733        let source: ArrayRef = Arc::new(array);
734
735        let mut in_progress = InProgressByteViewArray::<BinaryViewType>::new(64);
736        in_progress.set_source(Some(Arc::clone(&source)));
737
738        in_progress.copy_rows(0, 60).unwrap();
739        in_progress.copy_rows(60, 4).unwrap();
740
741        // The reused buffers now live in `completed`, but while the source is
742        // still set they are counted via the source, not `completed_buffers_size`,
743        // to avoid double counting them.
744        assert_eq!(in_progress.completed_buffers_size, 0);
745        assert_eq!(
746            in_progress.size_of_completed_buffers_from_current_source,
747            buffer_capacity,
748        );
749
750        // Setting a new source commits the pending reused-buffer bytes, since the
751        // old source (and its double count) is dropped.
752        let other: ArrayRef = Arc::new(BinaryViewArray::from_iter(std::iter::once(Some(
753            b"short".as_slice(),
754        ))));
755        in_progress.set_source(Some(Arc::clone(&other)));
756        assert_eq!(in_progress.completed_buffers_size, buffer_capacity);
757        assert_eq!(in_progress.size_of_completed_buffers_from_current_source, 0);
758
759        // finish() releases everything.
760        in_progress.finish().unwrap();
761        assert_eq!(in_progress.completed_buffers_size, 0);
762        assert_eq!(in_progress.size_of_completed_buffers_from_current_source, 0);
763    }
764
765    #[test]
766    fn size_should_be_the_same_if_copying_multiple_time_from_same_source_or_once() {
767        let (array, _buffer_capacity) = non_inline_array(64);
768        let source: ArrayRef = Arc::new(array);
769
770        let in_progress_size_with_split = {
771            let mut in_progress = InProgressByteViewArray::<BinaryViewType>::new(64);
772            in_progress.set_source(Some(Arc::clone(&source)));
773
774            in_progress.copy_rows(0, 60).unwrap();
775            in_progress.copy_rows(60, 4).unwrap();
776            in_progress.set_source(None);
777
778            in_progress.size()
779        };
780
781        let in_progress_size_without_split = {
782            let mut in_progress = InProgressByteViewArray::<BinaryViewType>::new(64);
783            in_progress.set_source(Some(Arc::clone(&source)));
784
785            in_progress.copy_rows(0, 64).unwrap();
786            in_progress.set_source(None);
787
788            in_progress.size()
789        };
790
791        assert_eq!(in_progress_size_with_split, in_progress_size_without_split);
792    }
793}
arrow_select/coalesce/byte_view.rs

arrow_select/coalesce/
byte_view.rs
