feat: Add optimize ArrayOp with VBView implementation (#3825)

a10y · web-flow · commit ec9d96240c4b · 2025-07-10T20:57:48.000+01:00
Adds a new OperationVTable method `optimize` which will return a new
array with unused data garbage-collected. This is similar to the
arrow-rs `StringViewArray::gc` method. See doc comments in the PR for
more info.

This implements a simple optimize operation that will fully rebuild a
VarBinViewArray if there is any wasted space in the buffers.

Signed-off-by: Andrew Duffy &lt;andrew@a10y.dev&gt;
diff --git a/vortex-array/src/array/mod.rs b/vortex-array/src/array/mod.rs
@@ -62,6 +62,11 @@ pub trait Array: 'static + private::Sealed + Send + Sync + Debug + ArrayVisitor
     /// Fetch the scalar at the given index.
     fn scalar_at(&self, index: usize) -> VortexResult<Scalar>;
 
+    /// Return an optimized version of the same array.
+    ///
+    /// See [`OperationsVTable::optimize`] for more details.
+    fn optimize(&self) -> VortexResult<ArrayRef>;
+
     /// Returns whether the array is of the given encoding.
     fn is_encoding(&self, encoding: EncodingId) -> bool {
         self.encoding_id() == encoding
@@ -183,6 +188,10 @@ impl Array for Arc<dyn Array> {
         self.as_ref().scalar_at(index)
     }
 
+    fn optimize(&self) -> VortexResult<ArrayRef> {
+        self.as_ref().optimize()
+    }
+
     fn is_valid(&self, index: usize) -> VortexResult<bool> {
         self.as_ref().is_valid(index)
     }
@@ -430,6 +439,39 @@ impl<V: VTable> Array for ArrayAdapter<V> {
         Ok(scalar)
     }
 
+    fn optimize(&self) -> VortexResult<ArrayRef> {
+        let result = <V::OperationsVTable as OperationsVTable<V>>::optimize(&self.0)?.into_array();
+
+        #[cfg(debug_assertions)]
+        {
+            let nbytes = self.0.nbytes();
+            let result_nbytes = result.nbytes();
+            assert!(
+                result_nbytes <= nbytes,
+                "optimize() made the array larger: {} bytes -> {} bytes",
+                nbytes,
+                result_nbytes
+            );
+        }
+
+        assert_eq!(
+            self.dtype(),
+            result.dtype(),
+            "optimize() changed DType from {} to {}",
+            self.dtype(),
+            result.dtype()
+        );
+        assert_eq!(
+            result.len(),
+            self.len(),
+            "optimize() changed len from {} to {}",
+            self.len(),
+            result.len()
+        );
+
+        Ok(result)
+    }
+
     fn is_valid(&self, index: usize) -> VortexResult<bool> {
         if index >= self.len() {
             vortex_bail!(OutOfBounds: index, 0, self.len());
diff --git a/vortex-array/src/arrays/varbinview/ops.rs b/vortex-array/src/arrays/varbinview/ops.rs
@@ -1,10 +1,12 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
-use vortex_error::VortexResult;
+use vortex_error::{VortexResult, VortexUnwrap};
 use vortex_scalar::Scalar;
 
 use crate::arrays::{VarBinViewArray, VarBinViewVTable, varbin_scalar};
+use crate::builders::{ArrayBuilder, VarBinViewBuilder};
+use crate::validity::Validity;
 use crate::vtable::{OperationsVTable, ValidityHelper};
 use crate::{ArrayRef, IntoArray};
 
@@ -24,4 +26,220 @@ impl OperationsVTable<VarBinViewVTable> for VarBinViewVTable {
     fn scalar_at(array: &VarBinViewArray, index: usize) -> VortexResult<Scalar> {
         Ok(varbin_scalar(array.bytes_at(index), array.dtype()))
     }
+
+    fn optimize(array: &VarBinViewArray) -> VortexResult<VarBinViewArray> {
+        // If there is nothing to be gained by compaction, return the original array untouched.
+        if !should_compact(array) {
+            return Ok(array.clone());
+        }
+
+        // Compaction pathways, depend on the validity
+        match array.validity {
+            // The array contains no values, all buffers can be dropped.
+            Validity::AllInvalid => Ok(VarBinViewArray::try_new(
+                array.views().clone(),
+                vec![],
+                array.dtype().clone(),
+                array.validity().clone(),
+            )?),
+            // Non-null pathway
+            Validity::NonNullable | Validity::AllValid => rebuild_nonnull(array),
+            // Nullable pathway, requires null-checks for each value
+            Validity::Array(_) => rebuild_nullable(array),
+        }
+    }
+}
+
+fn should_compact(array: &VarBinViewArray) -> bool {
+    // If the array is entirely inlined strings, do not attempt to compact.
+    if array.nbuffers() == 0 {
+        return false;
+    }
+
+    let bytes_referenced: u64 = count_referenced_bytes(array);
+    let buffer_total_bytes: u64 = array.buffers.iter().map(|buf| buf.len() as u64).sum();
+
+    // If there is any wasted space, we want to repack.
+    // This is very aggressive.
+    bytes_referenced < buffer_total_bytes
+}
+
+// count the number of bytes addressed by the views, not including null
+// values or any inlined strings.
+fn count_referenced_bytes(array: &VarBinViewArray) -> u64 {
+    match array.validity() {
+        Validity::AllInvalid => 0u64,
+        _ => {
+            array
+                .views()
+                .iter()
+                .enumerate()
+                .map(|(idx, &view)| {
+                    if !array.is_valid(idx).vortex_unwrap() || view.is_inlined() {
+                        0u64
+                    } else {
+                        // SAFETY: in this branch the view is not inlined.
+                        unsafe { view._ref }.size as u64
+                    }
+                })
+                .sum()
+        }
+    }
+}
+
+// Nullable string array compaction pathway.
+// This requires a null check on every append.
+fn rebuild_nullable(array: &VarBinViewArray) -> VortexResult<VarBinViewArray> {
+    let mut builder = VarBinViewBuilder::with_capacity(array.dtype().clone(), array.len());
+    for i in 0..array.len() {
+        if !array.is_valid(i)? {
+            builder.append_null();
+        } else {
+            let bytes = array.bytes_at(i);
+            builder.append_value(bytes.as_slice());
+        }
+    }
+
+    Ok(builder.finish_into_varbinview())
+}
+
+// Compaction for string arrays that contain no null values. Saves a branch
+// for every string element.
+fn rebuild_nonnull(array: &VarBinViewArray) -> VortexResult<VarBinViewArray> {
+    let mut builder = VarBinViewBuilder::with_capacity(array.dtype().clone(), array.len());
+    for i in 0..array.len() {
+        builder.append_value(array.bytes_at(i).as_ref());
+    }
+    Ok(builder.finish_into_varbinview())
+}
+
+#[cfg(test)]
+mod tests {
+    use vortex_buffer::buffer;
+
+    use crate::IntoArray;
+    use crate::arrays::{VarBinViewArray, VarBinViewVTable};
+    use crate::compute::take;
+
+    #[test]
+    fn test_optimize_compacts_buffers() {
+        // Create a VarBinViewArray with some long strings that will create multiple buffers
+        let original = VarBinViewArray::from_iter_nullable_str([
+            Some("short"),
+            Some("this is a longer string that will be stored in a buffer"),
+            Some("medium length string"),
+            Some("another very long string that definitely needs a buffer to store it"),
+            Some("tiny"),
+        ]);
+
+        // Verify it has buffers
+        assert!(original.nbuffers() > 0);
+        let original_buffers = original.nbuffers();
+
+        // Take only the first and last elements (indices 0 and 4)
+        let indices = buffer![0u32, 4u32].into_array();
+        let taken = take(original.as_ref(), &indices).unwrap();
+        let taken_array = taken.as_::<VarBinViewVTable>();
+
+        // The taken array should still have the same number of buffers
+        assert_eq!(taken_array.nbuffers(), original_buffers);
+
+        // Now optimize the taken array
+        let optimized = taken_array.optimize().unwrap();
+        let optimized_array = optimized.as_::<VarBinViewVTable>();
+
+        // The optimized array should have compacted buffers
+        // Since both remaining strings are short, they should be inlined
+        // so we might have 0 buffers, or 1 buffer if any were not inlined
+        assert!(optimized_array.nbuffers() <= 1);
+
+        // Verify the data is still correct
+        assert_eq!(optimized_array.len(), 2);
+        assert_eq!(optimized_array.scalar_at(0).unwrap(), "short".into());
+        assert_eq!(optimized_array.scalar_at(1).unwrap(), "tiny".into());
+    }
+
+    #[test]
+    fn test_optimize_with_long_strings() {
+        // Create strings that are definitely longer than 12 bytes
+        let long_string_1 = "this is definitely a very long string that exceeds the inline limit";
+        let long_string_2 = "another extremely long string that also needs external buffer storage";
+        let long_string_3 = "yet another long string for testing buffer compaction functionality";
+
+        let original = VarBinViewArray::from_iter_str([
+            long_string_1,
+            long_string_2,
+            long_string_3,
+            "short1",
+            "short2",
+        ]);
+
+        // Take only the first and third long strings (indices 0 and 2)
+        let indices = buffer![0u32, 2u32].into_array();
+        let taken = take(original.as_ref(), &indices).unwrap();
+        let taken_array = taken.as_::<VarBinViewVTable>();
+
+        // Optimize the taken array
+        let optimized = taken_array.optimize().unwrap();
+        let optimized_array = optimized.as_::<VarBinViewVTable>();
+
+        // The optimized array should have exactly 1 buffer (consolidated)
+        assert_eq!(optimized_array.nbuffers(), 1);
+
+        // Verify the data is still correct
+        assert_eq!(optimized_array.len(), 2);
+        assert_eq!(optimized_array.scalar_at(0).unwrap(), long_string_1.into());
+        assert_eq!(optimized_array.scalar_at(1).unwrap(), long_string_3.into());
+    }
+
+    #[test]
+    fn test_optimize_no_buffers() {
+        // Create an array with only short strings (all inlined)
+        let original = VarBinViewArray::from_iter_str(["a", "bb", "ccc", "dddd"]);
+
+        // This should have no buffers
+        assert_eq!(original.nbuffers(), 0);
+
+        // Optimize should return the same array
+        let optimized = original.optimize().unwrap();
+        let optimized_array = optimized.as_::<VarBinViewVTable>();
+
+        assert_eq!(optimized_array.nbuffers(), 0);
+        assert_eq!(optimized_array.len(), 4);
+
+        // Verify all values are preserved
+        for i in 0..4 {
+            assert_eq!(
+                optimized_array.scalar_at(i).unwrap(),
+                original.scalar_at(i).unwrap()
+            );
+        }
+    }
+
+    #[test]
+    fn test_optimize_single_buffer() {
+        // Create an array that naturally has only one buffer
+        let str1 = "this is a long string that goes into a buffer";
+        let str2 = "another long string in the same buffer";
+        let original = VarBinViewArray::from_iter_str([str1, str2]);
+
+        // Should have 1 compact buffer
+        assert_eq!(original.nbuffers(), 1);
+        assert_eq!(original.buffer(0).len(), str1.len() + str2.len());
+
+        // Optimize should return the same array (no change needed)
+        let optimized = original.optimize().unwrap();
+        let optimized_array = optimized.as_::<VarBinViewVTable>();
+
+        assert_eq!(optimized_array.nbuffers(), 1);
+        assert_eq!(optimized_array.len(), 2);
+
+        // Verify all values are preserved
+        for i in 0..2 {
+            assert_eq!(
+                optimized_array.scalar_at(i).unwrap(),
+                original.scalar_at(i).unwrap()
+            );
+        }
+    }
 }
diff --git a/vortex-array/src/builders/varbinview.rs b/vortex-array/src/builders/varbinview.rs
@@ -120,6 +120,31 @@ impl VarBinViewBuilder {
         debug_assert_eq!(self.null_buffer_builder.len(), self.views_builder.len())
     }
 
+    pub fn finish_into_varbinview(&mut self) -> VarBinViewArray {
+        self.flush_in_progress();
+        let buffers = std::mem::take(&mut self.completed);
+
+        assert_eq!(
+            self.views_builder.len(),
+            self.null_buffer_builder.len(),
+            "View and validity length must match"
+        );
+
+        let validity = self
+            .null_buffer_builder
+            .finish_with_nullability(self.nullability);
+
+        VarBinViewArray::try_new(
+            std::mem::take(&mut self.views_builder).freeze(),
+            buffers,
+            std::mem::replace(&mut self.dtype, DType::Null),
+            validity,
+        )
+        .vortex_expect("VarBinViewArray components should be valid.")
+    }
+}
+
+impl VarBinViewBuilder {
     // Pushes a validity mask into the builder not affecting the views or buffers
     fn push_only_validity_mask(&mut self, validity_mask: Mask) {
         self.null_buffer_builder.append_validity_mask(validity_mask);
@@ -191,27 +216,7 @@ impl ArrayBuilder for VarBinViewBuilder {
     }
 
     fn finish(&mut self) -> ArrayRef {
-        self.flush_in_progress();
-        let buffers = std::mem::take(&mut self.completed);
-
-        assert_eq!(
-            self.views_builder.len(),
-            self.null_buffer_builder.len(),
-            "View and validity length must match"
-        );
-
-        let validity = self
-            .null_buffer_builder
-            .finish_with_nullability(self.nullability);
-
-        VarBinViewArray::try_new(
-            std::mem::take(&mut self.views_builder).freeze(),
-            buffers,
-            std::mem::replace(&mut self.dtype, DType::Null),
-            validity,
-        )
-        .vortex_expect("VarBinViewArray components should be valid.")
-        .into_array()
+        self.finish_into_varbinview().into_array()
     }
 }
 
diff --git a/vortex-array/src/vtable/operations.rs b/vortex-array/src/vtable/operations.rs
@@ -28,4 +28,25 @@ pub trait OperationsVTable<V: VTable> {
     /// Bounds-checking has already been performed by the time this function is called,
     /// and the index is guaranteed to be non-null.
     fn scalar_at(array: &V::Array, index: usize) -> VortexResult<Scalar>;
+
+    /// Return an optimized copy of an Array, with any unreferenced data blocks freed and
+    /// any extraneous information removed.
+    ///
+    /// Many simple contiguous array types do not benefit from this operation, but it is
+    /// especially useful for variable-length types which keep a variable number of buffers
+    /// that can be dereferenced by other internal data structures via simple zero-copy
+    /// filter and take operations. After multiple operations are applied to these arrays, it is
+    /// common for the majority of owned buffer data to no longer logically be referenced.
+    ///
+    /// This operation can be called to return a new copy of an array with the same encoding,
+    /// but with all unreferenced data unlinked.
+    ///
+    /// ## Default behavior
+    ///
+    /// For most arrays that do not contain variable buffer counts, such as the canonical
+    /// arrays, the default implementation will not attempt to perform compaction and instead
+    /// return the original array.
+    fn optimize(array: &V::Array) -> VortexResult<V::Array> {
+        Ok(array.clone())
+    }
 }
