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

vortex_array/arrays/varbinview/
compact.rs
