implement more efficient `extend_from_array_unchecked (ugly)

Signed-off-by: Connor Tsui <[email protected]>

Author: connortsui20

vortex-array/src/arrays/chunked/vtable/canonical.rs

@@ -124,9 +124,9 @@ fn swizzle_list_chunks(
 
     for chunk in chunks {
         let chunk_array = chunk.to_listview();
-        // By rebuilding as zero-copy to `List`, we make the final output `ListView` also
-        // zero-copyable to `List`.
-        let chunk_array = chunk_array.rebuild(ListViewRebuildMode::MakeZeroCopyToList);
+        // By rebuilding as zero-copy to `List` and trimming all elements (to prevent gaps), we make
+        // the final output `ListView` also zero-copyable to `List`.
+        let chunk_array = chunk_array.rebuild(ListViewRebuildMode::MakeExact);
 
         // Add the `elements` of the current array as a new chunk.
         list_elements_chunks.push(chunk_array.elements().clone());

vortex-array/src/arrays/listview/conversion.rs

@@ -281,7 +281,8 @@ mod tests {
     };
     use super::recursive_list_from_list_view;
     use crate::arrays::{
-        list_from_list_view, list_view_from_list, BoolArray, FixedSizeListArray, ListArray, ListViewArray, ListViewShape, PrimitiveArray, StructArray
+        BoolArray, FixedSizeListArray, ListArray, ListViewArray, ListViewShape, PrimitiveArray,
+        StructArray, list_from_list_view, list_view_from_list,
     };
     use crate::validity::Validity;
     use crate::vtable::ValidityHelper;

vortex-array/src/arrays/listview/rebuild.rs

@@ -26,6 +26,11 @@ pub enum ListViewRebuildMode {
     /// the [`ListViewArray`].
     TrimElements,
 
+    /// Equivalent to `MakeZeroCopyToList` plus `TrimElements`.
+    ///
+    /// TODO more docs.
+    MakeExact,
+
     /// Removes any unused data from the underlying `elements` array.
     ///
     /// This mode will rebuild the `elements` array in the process, but it will keep any overlapping
@@ -63,6 +68,9 @@ impl ListViewArray {
         match mode {
             ListViewRebuildMode::MakeZeroCopyToList => self.rebuild_zero_copy_to_list(),
             ListViewRebuildMode::TrimElements => self.rebuild_trim_elements(),
+            ListViewRebuildMode::MakeExact => {
+                self.rebuild_zero_copy_to_list().rebuild_trim_elements()
+            }
             ListViewRebuildMode::RemoveGaps => self.rebuild_remove_gaps::<false>(),
             ListViewRebuildMode::RemoveNulls => self.rebuild_remove_gaps::<true>(),
             ListViewRebuildMode::OverlapCompression => unimplemented!("Does P=NP?"),
@@ -82,27 +90,38 @@ impl ListViewArray {
             return self.clone();
         }
 
+        let offsets_ptype = self.offsets().dtype().as_ptype();
+        let sizes_ptype = self.sizes().dtype().as_ptype();
+
+        match_each_integer_ptype!(offsets_ptype, |O| {
+            match_each_integer_ptype!(sizes_ptype, |S| { self.naive_rebuild::<O, S>() })
+        })
+    }
+
+    /// The inner function for `rebuild_zero_copy_to_list`, which naively rebuilds a `ListViewArray`
+    /// via `append_scalar`.
+    fn naive_rebuild<O: IntegerPType, S: IntegerPType>(&self) -> ListViewArray {
         let element_dtype = self
             .dtype()
             .as_list_element_opt()
             .vortex_expect("somehow had a canonical list that was not a list");
 
-        let offsets_ptype = self.offsets().dtype().as_ptype();
-        let sizes_ptype = self.sizes().dtype().as_ptype();
+        let mut builder = ListViewBuilder::<O, S>::with_capacity(
+            element_dtype.clone(),
+            self.dtype().nullability(),
+            self.elements().len(),
+            self.len(),
+        );
 
-        match_each_integer_ptype!(offsets_ptype, |O| {
-            match_each_integer_ptype!(sizes_ptype, |S| {
-                let mut builder = ListViewBuilder::<O, S>::with_capacity(
-                    element_dtype.clone(),
-                    self.dtype().nullability(),
-                    self.elements().len(),
-                    self.len(),
-                );
-
-                builder.extend_from_array(self.as_ref());
-                builder.finish_into_listview()
-            })
-        })
+        for i in 0..self.len() {
+            let list = self.scalar_at(i);
+
+            builder
+                .append_scalar(&list)
+                .vortex_expect("was unable to extend the `ListViewBuilder`")
+        }
+
+        builder.finish_into_listview()
     }
 
     /// Rebuilds a [`ListViewArray`] by trimming any unused / unreferenced leading and trailing
@@ -658,9 +677,8 @@ mod tests {
         let elements = PrimitiveArray::from_iter(vec![1i32, 2, 3, 4, 5]).into_array();
         let offsets = PrimitiveArray::from_iter(vec![0u32, 2, 4]).into_array();
         let sizes = PrimitiveArray::from_iter(vec![2u32, 2, 1]).into_array();
-        let validity = Validity::Array(
-            BoolArray::from(BitBuffer::from(vec![true, false, true])).into_array(),
-        );
+        let validity =
+            Validity::Array(BoolArray::from(BitBuffer::from(vec![true, false, true])).into_array());
 
         let listview = ListViewArray::try_new(
             elements,

vortex-array/src/builders/listview.rs

@@ -10,18 +10,18 @@
 
 use std::sync::Arc;
 
-use vortex_dtype::{DType, IntegerPType, Nullability};
+use vortex_dtype::{DType, IntegerPType, Nullability, match_each_integer_ptype};
 use vortex_error::{VortexExpect, VortexResult, vortex_ensure, vortex_panic};
 use vortex_mask::Mask;
 use vortex_scalar::{ListScalar, Scalar};
 
 use crate::array::{Array, ArrayRef, IntoArray};
-use crate::arrays::{ListViewArray, ListViewShape};
+use crate::arrays::{ListViewArray, ListViewRebuildMode, ListViewShape, PrimitiveArray};
 use crate::builders::lazy_null_builder::LazyBitBufferBuilder;
 use crate::builders::{
-    ArrayBuilder, DEFAULT_BUILDER_CAPACITY, PrimitiveBuilder, builder_with_capacity,
+    ArrayBuilder, DEFAULT_BUILDER_CAPACITY, PrimitiveBuilder, UninitRange, builder_with_capacity,
 };
-use crate::{Canonical, ToCanonical};
+use crate::{Canonical, ToCanonical, compute};
 
 /// A builder for creating [`ListViewArray`] instances, parameterized by the [`IntegerPType`] of
 /// the `offsets` and the `sizes` builders.
@@ -236,28 +236,74 @@ impl<O: IntegerPType, S: IntegerPType> ArrayBuilder for ListViewBuilder<O, S> {
     }
 
     unsafe fn extend_from_array_unchecked(&mut self, array: &dyn Array) {
-        let listview_array = array.to_listview();
-        if listview_array.is_empty() {
+        let listview = array.to_listview();
+        if listview.is_empty() {
             return;
         }
 
-        // TODO NOW: Rebuild the array we are extending by to be zero-copy to list and then be
-        // smarter about this.
-        
-        // TODO(connor)[ListView]: We could potentially concatenate the new elements on top of the
-        // existing elements and recalculate offsets (and then use `UninitRange`). However, that
-        // would mean we lose the guarantee that the output `ListViewArray` does not look like a
-        // `ListArray` (because the incoming array could have garbage data).
-
-        // We assume the worst case scenario, where the list view array is stored completely out of
-        // order, with many out-of-order offsets, and lots of garbage data. Thus, we simply iterate
-        // over all of the lists in the array and copy the data into this builder.
-        for i in 0..listview_array.len() {
-            let list = listview_array.scalar_at(i);
-
-            self.append_scalar(&list)
-                .vortex_expect("was unable to extend the `ListViewBuilder`")
+        // If we do not have the guarantee that the array is zero-copyable to a list, then we have
+        // to manually append each scalar.
+        if !listview.shape().is_zero_copy_to_list() {
+            for i in 0..listview.len() {
+                let list = listview.scalar_at(i);
+
+                self.append_scalar(&list)
+                    .vortex_expect("was unable to extend the `ListViewBuilder`")
+            }
         }
+
+        // Otherwise, after removing any leading and trailing elements, we can simply bulk append
+        // the entire array.
+        let listview = listview.rebuild(ListViewRebuildMode::TrimElements);
+
+        self.nulls.append_validity_mask(array.validity_mask());
+        self.elements_builder.extend_from_array(listview.elements());
+
+        self.sizes_builder.extend_from_array(
+            compute::cast(listview.sizes(), self.sizes_builder.dtype())
+                .vortex_expect(
+                    "was somehow unable to cast the new sizes to the type of the builder sizes",
+                )
+                .as_ref(),
+        );
+
+        // Adjust all of the offsets.
+        let new_offsets = listview.offsets().to_primitive(); // This should be cheap.
+        let num_new_offsets = new_offsets.len();
+
+        let curr_builder_len = self.len();
+        self.offsets_builder.reserve_exact(num_new_offsets);
+
+        let uninit_range = self.offsets_builder.uninit_range(num_new_offsets);
+
+        fn adjust_offsets<'a, O: IntegerPType, A: IntegerPType>(
+            mut uninit_range: UninitRange<'a, O>,
+            new_offsets: PrimitiveArray,
+            curr_builder_len: usize,
+        ) {
+            let new_offsets_slice = new_offsets.as_slice::<A>();
+            let curr_builder_len = O::from_usize(curr_builder_len)
+                .vortex_expect("the builder length did not fit into the offset type (impossible)");
+
+            for i in 0..uninit_range.len() {
+                let new_offset = new_offsets_slice[i];
+                let new_offset_correct_type = O::from_usize(
+                    new_offset
+                        .to_usize()
+                        .vortex_expect("Offsets must always fit in usize"),
+                )
+                .vortex_expect("New offset somehow did not fit into the builder's offset type");
+
+                uninit_range.set_value(i, new_offset_correct_type + curr_builder_len);
+            }
+            // SAFETY: We have set all the values in the range, and since `offsets` are
+            // non-nullable, we are done.
+            unsafe { uninit_range.finish() };
+        }
+
+        match_each_integer_ptype!(new_offsets.ptype(), |A| {
+            adjust_offsets::<O, A>(uninit_range, new_offsets, curr_builder_len);
+        })
     }
 
     fn reserve_exact(&mut self, capacity: usize) {
@@ -481,19 +527,15 @@ mod tests {
             .unwrap();
 
         // Extend from the ListArray.
-        unsafe {
-            builder.extend_from_array_unchecked(&source.into_array());
-        }
+        builder.extend_from_array(&source.into_array());
 
         // Extend from empty array (should be no-op).
         let empty_source = ListArray::from_iter_opt_slow::<u32, _, Vec<i32>>(
             std::iter::empty::<Option<Vec<i32>>>(),
             Arc::new(I32.into()),
         )
         .unwrap();
-        unsafe {
-            builder.extend_from_array_unchecked(&empty_source.into_array());
-        }
+        builder.extend_from_array(&empty_source.into_array());
 
         let listview = builder.finish_into_listview();
         assert_eq!(listview.len(), 4);

vortex-array/src/compute/list_contains.rs

@@ -53,7 +53,7 @@ pub(crate) fn warm_up_vtable() -> usize {
 ///
 /// ```rust
 /// # use vortex_array::{Array, IntoArray, ToCanonical};
-/// # use vortex_array::arrays::{ConstantArray, ListViewArray, VarBinArray};
+/// # use vortex_array::arrays::{ConstantArray, ListViewArray, ListViewShape, VarBinArray};
 /// # use vortex_array::compute;
 /// # use vortex_array::validity::Validity;
 /// # use vortex_buffer::{buffer, bitbuffer};
@@ -65,7 +65,13 @@ pub(crate) fn warm_up_vtable() -> usize {
 /// let offsets = buffer![0u32, 1, 3].into_array();
 /// let sizes = buffer![1u32, 2, 2].into_array();
 /// let list_array =
-///     ListViewArray::try_new(elements, offsets, sizes, Validity::NonNullable).unwrap();
+///     ListViewArray::try_new(
+///         elements,
+///         offsets,
+///         sizes,
+///         Validity::NonNullable,
+///         ListViewShape::as_zero_copy_to_list(),
+///     ).unwrap();
 ///
 /// let matches = compute::list_contains(
 ///     list_array.as_ref(),