Add ensure_capacity to array builders and fix the case where list element builders didn't allocate enough capacity for all elements (#2698)

Author: robert3005

encodings/fastlanes/src/bitpacking/compress.rs

@@ -213,7 +213,7 @@ pub fn unpack_primitive<T: NativePType, UnsignedT: NativePType + BitPacking>(
 ) -> VortexResult<PrimitiveArray> {
     let n = array.len();
     let mut builder = PrimitiveBuilder::with_capacity(array.dtype().nullability(), array.len());
-    assert!(size_of::<T>() == size_of::<UnsignedT>());
+    assert_eq!(size_of::<T>(), size_of::<UnsignedT>());
     unpack_into::<T, UnsignedT, _, _>(
         array,
         &mut builder,
@@ -224,7 +224,7 @@ pub fn unpack_primitive<T: NativePType, UnsignedT: NativePType + BitPacking>(
         // &mut [UnsignedT] is therefore safe.
         |x| unsafe { std::mem::transmute(x) },
     )?;
-    assert!(builder.len() == n, "{} {}", builder.len(), n);
+    assert_eq!(builder.len(), n);
     Ok(builder.finish_into_primitive())
 }
 

encodings/fastlanes/src/bitpacking/mod.rs

@@ -254,10 +254,10 @@ impl ArrayCanonicalImpl for BitPackedArray {
                     .as_any_mut()
                     .downcast_mut()
                     .vortex_expect("bit packed array must canonicalize into a primitive array"),
-                // SAFETY: UnsignedT is the unsigned verison of T, reinterpreting &[UnsignedT] to
+                // SAFETY: UnsignedT is the unsigned version of T, reinterpreting &[UnsignedT] to
                 // &[T] is therefore safe.
                 |x| unsafe { std::mem::transmute(x) },
-                // SAFETY: UnsignedT is the unsigned verison of T, reinterpreting &mut [T] to
+                // SAFETY: UnsignedT is the unsigned version of T, reinterpreting &mut [T] to
                 // &mut [UnsignedT] is therefore safe.
                 |x| unsafe { std::mem::transmute(x) },
             )

vortex-array/src/array/canonical.rs

@@ -21,7 +21,6 @@ pub trait ArrayCanonicalImpl {
     /// - The length of the builder is incremented by the length of the input array.
     fn _append_to_builder(&self, builder: &mut dyn ArrayBuilder) -> VortexResult<()> {
         let canonical = self._to_canonical()?;
-        // debug!("default impl append_to_builder {}", canonical.encoding());
         builder.extend_from_array(canonical.as_ref())
     }
 }

vortex-array/src/builders/bool.rs

@@ -86,6 +86,13 @@ impl ArrayBuilder for BoolBuilder {
         Ok(())
     }
 
+    fn ensure_capacity(&mut self, capacity: usize) {
+        if capacity > self.inner.capacity() {
+            self.nulls.ensure_capacity(capacity);
+            self.inner.reserve(capacity - self.inner.capacity());
+        }
+    }
+
     fn set_validity(&mut self, validity: Mask) {
         self.nulls = LazyNullBufferBuilder::new(validity.len());
         self.nulls.append_validity_mask(validity);

vortex-array/src/builders/extension.rs

@@ -83,6 +83,10 @@ impl ArrayBuilder for ExtensionBuilder {
         array.storage().append_to_builder(self.storage.as_mut())
     }
 
+    fn ensure_capacity(&mut self, capacity: usize) {
+        self.storage.ensure_capacity(capacity)
+    }
+
     fn set_validity(&mut self, validity: Mask) {
         self.storage.set_validity(validity);
     }

vortex-array/src/builders/lazy_validity_builder.rs

@@ -135,6 +135,20 @@ impl LazyNullBufferBuilder {
         }
     }
 
+    pub fn ensure_capacity(&mut self, capacity: usize) {
+        if self.inner.is_none() {
+            self.capacity = capacity;
+        } else {
+            let inner = self
+                .inner
+                .as_mut()
+                .vortex_expect("buffer just materialized");
+            if capacity < inner.capacity() {
+                inner.reserve(capacity - inner.len());
+            }
+        }
+    }
+
     #[inline]
     fn materialize_if_needed(&mut self) {
         if self.inner.is_none() {

vortex-array/src/builders/list.rs

@@ -11,7 +11,7 @@ use crate::arrays::{ConstantArray, ListArray, OffsetPType};
 use crate::builders::lazy_validity_builder::LazyNullBufferBuilder;
 use crate::builders::{ArrayBuilder, ArrayBuilderExt, PrimitiveBuilder, builder_with_capacity};
 use crate::compute::{binary_numeric, slice, try_cast};
-use crate::{Array, ArrayRef};
+use crate::{Array, ArrayRef, ToCanonical};
 
 pub struct ListBuilder<O: NativePType> {
     value_builder: Box<dyn ArrayBuilder>,
@@ -118,7 +118,7 @@ impl<O: OffsetPType> ArrayBuilder for ListBuilder<O> {
     fn extend_from_array(&mut self, array: &dyn Array) -> VortexResult<()> {
         self.nulls.append_validity_mask(array.validity_mask()?);
 
-        let list = array.to_canonical()?.into_list()?;
+        let list = array.to_list()?;
 
         let cursor_usize = self.value_builder.len();
         let cursor = O::from_usize(cursor_usize).ok_or_else(|| {
@@ -141,16 +141,20 @@ impl<O: OffsetPType> ArrayBuilder for ListBuilder<O> {
 
         if !list.is_empty() {
             let last_used_index = self.index_builder.values().last().vortex_expect("there must be at least one index because we just extended a non-zero list of offsets");
-            self.value_builder.extend_from_array(&slice(
-                list.elements(),
-                0,
-                last_used_index.as_() - cursor_usize,
-            )?)?;
+            let sliced_values = slice(list.elements(), 0, last_used_index.as_() - cursor_usize)?;
+            self.value_builder.ensure_capacity(sliced_values.len());
+            self.value_builder.extend_from_array(&sliced_values)?;
         }
 
         Ok(())
     }
 
+    fn ensure_capacity(&mut self, capacity: usize) {
+        self.index_builder.ensure_capacity(capacity);
+        self.value_builder.ensure_capacity(capacity);
+        self.nulls.ensure_capacity(capacity);
+    }
+
     fn set_validity(&mut self, validity: Mask) {
         self.nulls = LazyNullBufferBuilder::new(validity.len());
         self.nulls.append_validity_mask(validity);

vortex-array/src/builders/mod.rs

@@ -58,6 +58,9 @@ pub trait ArrayBuilder: Send {
     /// Extends the array with the provided array, canonicalizing if necessary.
     fn extend_from_array(&mut self, array: &dyn Array) -> VortexResult<()>;
 
+    /// Ensure that the builder can hold at least `capacity` number of items
+    fn ensure_capacity(&mut self, capacity: usize);
+
     /// Override builders validity with the one provided
     fn set_validity(&mut self, validity: Mask);
 

vortex-array/src/builders/null.rs

@@ -55,6 +55,8 @@ impl ArrayBuilder for NullBuilder {
         Ok(())
     }
 
+    fn ensure_capacity(&mut self, _capacity: usize) {}
+
     fn set_validity(&mut self, validity: Mask) {
         self.length = validity.len();
     }

vortex-array/src/builders/primitive.rs

@@ -12,7 +12,7 @@ use crate::builders::ArrayBuilder;
 use crate::builders::lazy_validity_builder::LazyNullBufferBuilder;
 use crate::validity::Validity;
 use crate::variants::PrimitiveArrayTrait;
-use crate::{Array, ArrayRef};
+use crate::{Array, ArrayRef, ToCanonical};
 
 /// Builder for [`PrimitiveArray`].
 pub struct PrimitiveBuilder<T> {
@@ -89,7 +89,8 @@ impl<T: NativePType> PrimitiveBuilder<T> {
         let offset = self.values.len();
         assert!(
             offset + len <= self.values.capacity(),
-            "uninit_range of len {len} exceeds builder capacity"
+            "uninit_range of len {len} exceeds builder capacity {}",
+            self.values.capacity()
         );
 
         UninitRange {
@@ -166,7 +167,7 @@ impl<T: NativePType> ArrayBuilder for PrimitiveBuilder<T> {
     }
 
     fn extend_from_array(&mut self, array: &dyn Array) -> VortexResult<()> {
-        let array = array.to_canonical()?.into_primitive()?;
+        let array = array.to_primitive()?;
         if array.ptype() != T::PTYPE {
             vortex_bail!("Cannot extend from array with different ptype");
         }
@@ -178,6 +179,13 @@ impl<T: NativePType> ArrayBuilder for PrimitiveBuilder<T> {
         Ok(())
     }
 
+    fn ensure_capacity(&mut self, capacity: usize) {
+        if capacity > self.values.capacity() {
+            self.values.reserve(capacity - self.values.len());
+            self.nulls.ensure_capacity(capacity);
+        }
+    }
+
     fn set_validity(&mut self, validity: Mask) {
         self.nulls = LazyNullBufferBuilder::new(validity.len());
         self.nulls.append_validity_mask(validity);