vortex_array/arrays/varbinview/
compact.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4//! Defines a compaction operation for VarBinViewArrays that evicts unused buffers so they can
5//! be dropped.
6
7use std::ops::Range;
8
9use vortex_error::VortexExpect;
10use vortex_error::VortexResult;
11use vortex_mask::Mask;
12
13use crate::ExecutionCtx;
14use crate::arrays::VarBinViewArray;
15use crate::arrays::varbinview::Ref;
16use crate::builders::VarBinViewBuilder;
17
18const DEFAULT_COMPACTION_THRESHOLD: f64 = 0.5;
19const MIN_RETAINED_BYTES_PER_ROW_TO_CHECK_COMPACTION: u64 = 128;
20
21impl VarBinViewArray {
22    /// Returns a compacted copy of the input array, where all wasted space has been cleaned up. This
23    /// operation can be very expensive, in the worst case copying all existing string data into
24    /// a new allocation.
25    ///
26    /// After slicing/taking operations `VarBinViewArray`s can continue to hold references to buffers
27    /// that are no longer visible. We detect when there is wasted space in any of the buffers, and if
28    /// so, will aggressively compact all visible outlined string data into new buffers while keeping
29    /// well-utilized buffers unchanged.
30    pub fn compact_buffers(&self, ctx: &mut ExecutionCtx) -> VortexResult<VarBinViewArray> {
31        // If there is nothing to be gained by compaction, return the original array untouched.
32        if !self.should_compact(ctx)? {
33            return Ok(self.clone());
34        }
35
36        self.compact_with_threshold(DEFAULT_COMPACTION_THRESHOLD, ctx)
37    }
38
39    fn should_compact(&self, ctx: &mut ExecutionCtx) -> VortexResult<bool> {
40        let nbuffers = self.data_buffers().len();
41
42        // If the array is entirely inlined strings, do not attempt to compact.
43        if nbuffers == 0 {
44            return Ok(false);
45        }
46
47        // These will fail to write, so in most cases we want to compact this.
48        if nbuffers > u16::MAX as usize {
49            return Ok(true);
50        }
51
52        let buffer_total_bytes: u64 = self.buffers.iter().map(|buf| buf.len() as u64).sum();
53        if buffer_total_bytes == 0 {
54            return Ok(true);
55        }
56
57        let len = u64::try_from(self.len()).unwrap_or(u64::MAX);
58        if len > 0 && buffer_total_bytes / len <= MIN_RETAINED_BYTES_PER_ROW_TO_CHECK_COMPACTION {
59            return Ok(false);
60        }
61
62        let bytes_referenced: u64 = self.count_referenced_bytes(ctx)?;
63        Ok((bytes_referenced as f64 / buffer_total_bytes as f64) < DEFAULT_COMPACTION_THRESHOLD)
64    }
65
66    /// Iterates over all valid, non-inlined views, calling the provided
67    /// closure for each one.
68    #[allow(clippy::inline_always)]
69    #[inline(always)]
70    fn iter_valid_views<F>(&self, ctx: &mut ExecutionCtx, mut f: F) -> VortexResult<()>
71    where
72        F: FnMut(&Ref),
73    {
74        match self.as_ref().validity()?.execute_mask(self.len(), ctx)? {
75            Mask::AllTrue(_) => {
76                for &view in self.views().iter() {
77                    if !view.is_inlined() {
78                        f(view.as_view());
79                    }
80                }
81            }
82            Mask::AllFalse(_) => {}
83            Mask::Values(v) => {
84                for (&view, is_valid) in self.views().iter().zip(v.bit_buffer().iter()) {
85                    if is_valid && !view.is_inlined() {
86                        f(view.as_view());
87                    }
88                }
89            }
90        }
91        Ok(())
92    }
93
94    /// Count the number of bytes addressed by the views, not including null
95    /// values or any inlined strings.
96    fn count_referenced_bytes(&self, ctx: &mut ExecutionCtx) -> VortexResult<u64> {
97        let mut total = 0u64;
98        self.iter_valid_views(ctx, |view| total += view.size as u64)?;
99        Ok(total)
100    }
101
102    pub(crate) fn buffer_utilizations(
103        &self,
104        ctx: &mut ExecutionCtx,
105    ) -> VortexResult<Vec<BufferUtilization>> {
106        let mut utilizations: Vec<BufferUtilization> = self
107            .data_buffers()
108            .iter()
109            .map(|buf| {
110                let len = u32::try_from(buf.len()).vortex_expect("buffer sizes must fit in u32");
111                BufferUtilization::zero(len)
112            })
113            .collect();
114
115        self.iter_valid_views(ctx, |view| {
116            utilizations[view.buffer_index as usize].add(view.offset, view.size);
117        })?;
118
119        Ok(utilizations)
120    }
121
122    /// Returns a compacted copy of the input array using selective buffer compaction.
123    ///
124    /// This method analyzes each buffer's utilization and applies one of three strategies:
125    /// - **KeepFull** (zero-copy): Well-utilized buffers are kept unchanged
126    /// - **Slice** (zero-copy): Buffers with contiguous ranges of used data are sliced to that range
127    /// - **Rewrite**: Poorly-utilized buffers have their data copied to new compact buffers
128    ///
129    /// By preserving or slicing well-utilized buffers, compaction becomes zero-copy in many cases.
130    ///
131    /// # Arguments
132    ///
133    /// * `buffer_utilization_threshold` - Threshold in range [0, 1]. Buffers with utilization
134    ///   below this value will be compacted. Use 0.0 for no compaction, 1.0 for aggressive
135    ///   compaction of any buffer with wasted space.
136    pub fn compact_with_threshold(
137        &self,
138        buffer_utilization_threshold: f64, // [0, 1]
139        ctx: &mut ExecutionCtx,
140    ) -> VortexResult<VarBinViewArray> {
141        let mut builder = VarBinViewBuilder::with_compaction_in(
142            self.dtype().clone(),
143            self.len(),
144            buffer_utilization_threshold,
145            ctx.allocator().clone(),
146        );
147        builder.append_varbinview_array(self, ctx)?;
148        Ok(builder.finish_into_varbinview())
149    }
150}
151
152pub(crate) struct BufferUtilization {
153    len: u32,
154    used: u32,
155    min_offset: u32,
156    max_offset_end: u32,
157}
158
159impl BufferUtilization {
160    pub(crate) fn zero(len: u32) -> Self {
161        BufferUtilization {
162            len,
163            used: 0u32,
164            min_offset: u32::MAX,
165            max_offset_end: 0,
166        }
167    }
168
169    pub(crate) fn add(&mut self, offset: u32, size: u32) {
170        self.used += size;
171        self.min_offset = self.min_offset.min(offset);
172        self.max_offset_end = self.max_offset_end.max(offset + size);
173    }
174
175    pub fn overall_utilization(&self) -> f64 {
176        match self.len {
177            0 => 0.0,
178            len => self.used as f64 / len as f64,
179        }
180    }
181
182    pub fn range_utilization(&self) -> f64 {
183        match self.range_span() {
184            0 => 0.0,
185            span => self.used as f64 / span as f64,
186        }
187    }
188
189    pub fn range(&self) -> Range<u32> {
190        self.min_offset..self.max_offset_end
191    }
192
193    fn range_span(&self) -> u32 {
194        self.max_offset_end.saturating_sub(self.min_offset)
195    }
196}
197
198#[cfg(test)]
199mod tests {
200    use rstest::rstest;
201    use vortex_buffer::buffer;
202
203    use crate::IntoArray;
204    use crate::VortexSessionExecute;
205    use crate::array_session;
206    use crate::arrays::VarBinArray;
207    use crate::arrays::VarBinViewArray;
208    use crate::assert_arrays_eq;
209    use crate::dtype::DType;
210    use crate::dtype::Nullability;
211    #[test]
212    fn test_optimize_compacts_buffers() {
213        let mut ctx = array_session().create_execution_ctx();
214        // Create a VarBinViewArray with some long strings that will create multiple buffers
215        let excluded = "unused long string".repeat(100);
216        let original = VarBinViewArray::from_iter_nullable_str([
217            Some("short"),
218            Some("this is a longer string that will be stored in a buffer"),
219            Some(excluded.as_str()),
220            Some("another very long string that definitely needs a buffer to store it"),
221            Some("tiny"),
222        ]);
223
224        // Verify it has buffers
225        assert!(!original.data_buffers().is_empty());
226        let original_buffers = original.data_buffers().len();
227
228        // Take only the first and last elements (indices 0 and 4)
229        let indices = buffer![0u32, 4u32].into_array();
230        let taken = original.take(indices).unwrap();
231        let taken = taken.execute::<VarBinViewArray>(&mut ctx).unwrap();
232        // The taken array should still have the same number of buffers
233        assert_eq!(taken.data_buffers().len(), original_buffers);
234
235        // Now optimize the taken array
236        let optimized_array = taken.compact_buffers(&mut ctx).unwrap();
237
238        assert!(optimized_array.data_buffers().is_empty());
239
240        // Verify the data is still correct
241        assert_arrays_eq!(
242            optimized_array,
243            <VarBinArray as FromIterator<_>>::from_iter([Some("short"), Some("tiny")]),
244            &mut ctx
245        );
246    }
247
248    #[test]
249    fn test_optimize_with_long_strings() {
250        let mut ctx = array_session().create_execution_ctx();
251        // Create strings that are definitely longer than 12 bytes
252        let long_string_1 = "this is definitely a very long string that exceeds the inline limit";
253        let long_string_2 = "another extremely long string that also needs external buffer storage";
254        let long_string_3 = "yet another long string for testing buffer compaction functionality";
255
256        let original = VarBinViewArray::from_iter_str([
257            long_string_1,
258            long_string_2,
259            long_string_3,
260            "short1",
261            "short2",
262        ]);
263
264        // Take only the first and third long strings (indices 0 and 2)
265        let indices = buffer![0u32, 2u32].into_array();
266        let taken = original.take(indices).unwrap();
267        let taken_array = taken
268            .execute::<VarBinViewArray>(&mut array_session().create_execution_ctx())
269            .unwrap();
270
271        let optimized_array = taken_array.compact_with_threshold(1.0, &mut ctx).unwrap();
272
273        // The optimized array should have exactly 1 buffer (consolidated)
274        assert_eq!(optimized_array.data_buffers().len(), 1);
275        assert_eq!(
276            optimized_array.buffer(0).len(),
277            long_string_1.len() + long_string_3.len()
278        );
279
280        // Verify the data is still correct
281        assert_arrays_eq!(
282            optimized_array,
283            VarBinArray::from(vec![long_string_1, long_string_3]),
284            &mut ctx
285        );
286    }
287
288    #[test]
289    fn test_optimize_no_buffers() {
290        let mut ctx = array_session().create_execution_ctx();
291        // Create an array with only short strings (all inlined)
292        let original = VarBinViewArray::from_iter_str(["a", "bb", "ccc", "dddd"]);
293
294        // This should have no buffers
295        assert_eq!(original.data_buffers().len(), 0);
296
297        // Optimize should return the same array
298        let optimized_array = original.compact_buffers(&mut ctx).unwrap();
299
300        assert_eq!(optimized_array.data_buffers().len(), 0);
301
302        assert_arrays_eq!(optimized_array, original, &mut ctx);
303    }
304
305    #[test]
306    fn test_optimize_single_buffer() {
307        let mut ctx = array_session().create_execution_ctx();
308        // Create an array that naturally has only one buffer
309        let str1 = "this is a long string that goes into a buffer";
310        let str2 = "another long string in the same buffer";
311        let original = VarBinViewArray::from_iter_str([str1, str2]);
312
313        // Should have 1 compact buffer
314        assert_eq!(original.data_buffers().len(), 1);
315        assert_eq!(original.buffer(0).len(), str1.len() + str2.len());
316
317        // Optimize should return the same array (no change needed)
318        let optimized_array = original.compact_buffers(&mut ctx).unwrap();
319
320        assert_eq!(optimized_array.data_buffers().len(), 1);
321
322        assert_arrays_eq!(optimized_array, original, &mut ctx);
323    }
324
325    #[test]
326    fn test_selective_compaction_with_threshold_zero() {
327        let mut ctx = array_session().create_execution_ctx();
328        // threshold=0 should keep all buffers (no compaction)
329        let original = VarBinViewArray::from_iter_str([
330            "this is a longer string that will be stored in a buffer",
331            "another very long string that definitely needs a buffer to store it",
332        ]);
333
334        let original_buffers = original.data_buffers().len();
335        assert!(original_buffers > 0);
336
337        // Take only first element
338        let indices = buffer![0u32].into_array();
339        let taken = original.take(indices).unwrap();
340        let taken = taken
341            .execute::<VarBinViewArray>(&mut array_session().create_execution_ctx())
342            .unwrap();
343        // Compact with threshold=0 (should not compact)
344        let compacted = taken.compact_with_threshold(0.0, &mut ctx).unwrap();
345
346        // Should still have the same number of buffers as the taken array
347        assert_eq!(compacted.data_buffers().len(), taken.data_buffers().len());
348
349        // Verify correctness
350        assert_arrays_eq!(compacted, taken, &mut ctx);
351    }
352
353    #[test]
354    fn test_selective_compaction_with_high_threshold() {
355        let mut ctx = array_session().create_execution_ctx();
356        // threshold=1.0 should compact any buffer with waste
357        let original = VarBinViewArray::from_iter_str([
358            "this is a longer string that will be stored in a buffer",
359            "another very long string that definitely needs a buffer to store it",
360            "yet another long string",
361        ]);
362
363        // Take only first and last elements
364        let indices = buffer![0u32, 2u32].into_array();
365        let taken = original.take(indices).unwrap();
366        let taken = taken
367            .execute::<VarBinViewArray>(&mut array_session().create_execution_ctx())
368            .unwrap();
369
370        let original_buffers = taken.data_buffers().len();
371
372        // Compact with threshold=1.0 (aggressive compaction)
373        let compacted = taken.compact_with_threshold(1.0, &mut ctx).unwrap();
374
375        // Should have compacted buffers
376        assert!(compacted.data_buffers().len() <= original_buffers);
377        assert!(
378            compacted
379                .data_buffers()
380                .iter()
381                .map(|b| b.len())
382                .sum::<usize>()
383                < taken.data_buffers().iter().map(|b| b.len()).sum::<usize>()
384        );
385
386        // Verify correctness
387        assert_arrays_eq!(compacted, taken, &mut ctx);
388    }
389
390    #[test]
391    fn test_selective_compaction_preserves_well_utilized_buffers() {
392        let mut ctx = array_session().create_execution_ctx();
393        // Create an array with multiple strings in one buffer (well-utilized)
394        let str1 = "first long string that needs external buffer storage";
395        let str2 = "second long string also in buffer";
396        let str3 = "third long string in same buffer";
397
398        let original = VarBinViewArray::from_iter_str([str1, str2, str3]);
399
400        // All strings should be in one well-utilized buffer
401        assert_eq!(original.data_buffers().len(), 1);
402
403        // Compact with high threshold
404        let compacted = original.compact_with_threshold(0.8, &mut ctx).unwrap();
405
406        // Well-utilized buffer should be preserved
407        assert_eq!(compacted.data_buffers().len(), 1);
408        assert_eq!(compacted.buffer(0).as_ptr(), original.buffer(0).as_ptr());
409        assert_eq!(
410            compacted.buffer(0).as_slice(),
411            original.buffer(0).as_slice()
412        );
413
414        // Verify all data is correct
415        assert_arrays_eq!(compacted, original, &mut ctx);
416    }
417
418    #[test]
419    fn test_selective_compaction_with_mixed_utilization() {
420        let mut ctx = array_session().create_execution_ctx();
421        // Create array with some long strings
422        let strings: Vec<String> = (0..10)
423            .map(|i| {
424                format!(
425                    "this is a long string number {} that needs buffer storage",
426                    i
427                )
428            })
429            .collect();
430
431        let original = VarBinViewArray::from_iter_str(strings.iter().map(|s| s.as_str()));
432
433        // Take every other element to create mixed utilization
434        let indices_array = buffer![0u32, 2u32, 4u32, 6u32, 8u32].into_array();
435        let taken = original.take(indices_array).unwrap();
436        let taken = taken
437            .execute::<VarBinViewArray>(&mut array_session().create_execution_ctx())
438            .unwrap();
439
440        // Compact with moderate threshold
441        let compacted = taken.compact_with_threshold(0.7, &mut ctx).unwrap();
442
443        let expected = VarBinViewArray::from_iter(
444            [0, 2, 4, 6, 8].map(|i| Some(strings[i].as_str())),
445            DType::Utf8(Nullability::NonNullable),
446        );
447        assert_arrays_eq!(compacted, expected, &mut ctx);
448    }
449
450    #[test]
451    fn test_slice_strategy_with_contiguous_range() {
452        let mut ctx = array_session().create_execution_ctx();
453        // Create array with strings that will be in one buffer
454        let strings: Vec<String> = (0..20)
455            .map(|i| format!("this is a long string number {} for slice test", i))
456            .collect();
457
458        let original = VarBinViewArray::from_iter_str(strings.iter().map(|s| s.as_str()));
459
460        // Take only the first 5 elements - they should be in a contiguous range at the start
461        let indices_array = buffer![0u32, 1u32, 2u32, 3u32, 4u32].into_array();
462        let taken = original.take(indices_array).unwrap();
463        let taken = taken
464            .execute::<VarBinViewArray>(&mut array_session().create_execution_ctx())
465            .unwrap();
466        // Get buffer stats before compaction
467        let original_ptr = taken.buffer(0).as_ptr();
468
469        // Compact with a threshold that should trigger slicing
470        // The range utilization should be high even if overall utilization is low
471        let compacted = taken.compact_with_threshold(0.8, &mut ctx).unwrap();
472
473        // After compaction, we should still have buffers (sliced, not rewritten)
474        assert!(
475            !compacted.data_buffers().is_empty(),
476            "Should have buffers after slice compaction"
477        );
478
479        // Verify correctness
480        assert_arrays_eq!(&compacted, taken, &mut ctx);
481
482        assert_eq!(compacted.data_buffers().len(), 1);
483        assert_eq!(compacted.buffer(0).as_ptr(), original_ptr);
484        assert_eq!(
485            compacted.buffer(0).len(),
486            strings[..5].iter().map(String::len).sum::<usize>()
487        );
488    }
489
490    const LONG1: &str = "long string one!";
491    const LONG2: &str = "long string two!";
492    const SHORT: &str = "x";
493    const EXPECTED_BYTES: u64 = (LONG1.len() + LONG2.len()) as u64;
494
495    fn mixed_array() -> VarBinViewArray {
496        VarBinViewArray::from_iter_nullable_str([Some(LONG1), None, Some(LONG2), Some(SHORT)])
497    }
498
499    #[rstest]
500    #[case::non_nullable(VarBinViewArray::from_iter_str([LONG1, LONG2, SHORT]), EXPECTED_BYTES, &[1.0])]
501    #[case::all_valid(VarBinViewArray::from_iter_nullable_str([Some(LONG1), Some(LONG2), Some(SHORT)]), EXPECTED_BYTES, &[1.0])]
502    #[case::all_invalid(VarBinViewArray::from_iter_nullable_str([None::<&str>, None]), 0, &[])]
503    #[case::mixed_validity(mixed_array(), EXPECTED_BYTES, &[1.0])]
504    fn test_validity_code_paths(
505        #[case] arr: VarBinViewArray,
506        #[case] expected_bytes: u64,
507        #[case] expected_utils: &[f64],
508    ) {
509        let mut ctx = array_session().create_execution_ctx();
510        assert_eq!(
511            arr.count_referenced_bytes(&mut ctx).unwrap(),
512            expected_bytes
513        );
514        let utils: Vec<f64> = arr
515            .buffer_utilizations(&mut ctx)
516            .unwrap()
517            .iter()
518            .map(|u| u.overall_utilization())
519            .collect();
520        assert_eq!(utils, expected_utils);
521    }
522}
vortex_array/arrays/varbinview/compact.rs

vortex_array/arrays/varbinview/
compact.rs
