Add take-based rebuild strategy for ListView with small lists (vortex-data#6492)

- Split `naive_rebuild` into two strategies: `rebuild_with_take` (bulk
take) and `rebuild_list_by_list` (per-list slice + builder)
- Add a simple heuristic (`should_use_take`) that picks take when
average list size < 128, defaulting to list-by-list otherwise
- Take avoids per-list builder overhead (slice + extend) and is 10-100×
faster for small lists
 - Arrow exporter uses the same codepath

---------

Signed-off-by: Baris Palaska <barispalaska@gmail.com>

Author: palaska

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

@@ -106,12 +106,92 @@ impl ListViewArray {
         })
     }
 
-    // TODO(connor)[ListView]: We should benchmark if it is faster to use `take` on the elements
-    // instead of using a builder.
-    /// The inner function for `rebuild_zero_copy_to_list`, which rebuilds a `ListViewArray` piece
-    /// by piece.
+    /// Picks between [`rebuild_with_take`](Self::rebuild_with_take) and
+    /// [`rebuild_list_by_list`](Self::rebuild_list_by_list) based on element dtype and average
+    /// list size.
     fn naive_rebuild<O: IntegerPType, NewOffset: IntegerPType, S: IntegerPType>(
         &self,
+    ) -> VortexResult<ListViewArray> {
+        let sizes_canonical = self.sizes().to_primitive();
+        let total: u64 = sizes_canonical
+            .as_slice::<S>()
+            .iter()
+            .map(|s| (*s).as_() as u64)
+            .sum();
+        if Self::should_use_take(total, self.len()) {
+            self.rebuild_with_take::<O, NewOffset, S>()
+        } else {
+            self.rebuild_list_by_list::<O, NewOffset, S>()
+        }
+    }
+
+    /// Returns `true` when we are confident that `rebuild_with_take` will
+    /// outperform `rebuild_list_by_list`.
+    ///
+    /// Take is dramatically faster for small lists (often 10-100×) because it
+    /// avoids per-list builder overhead. LBL is the safer default for larger
+    /// lists since its sequential memcpy scales well. We only choose take when
+    /// the average list size is small enough that take clearly dominates.
+    fn should_use_take(total_output_elements: u64, num_lists: usize) -> bool {
+        if num_lists == 0 {
+            return true;
+        }
+        let avg = total_output_elements / num_lists as u64;
+        avg < 128
+    }
+
+    /// Rebuilds elements using a single bulk `take`: collect all element indices into a flat
+    /// `BufferMut<u64>`, perform a single `take`.
+    fn rebuild_with_take<O: IntegerPType, NewOffset: IntegerPType, S: IntegerPType>(
+        &self,
+    ) -> VortexResult<ListViewArray> {
+        let offsets_canonical = self.offsets().to_primitive();
+        let offsets_slice = offsets_canonical.as_slice::<O>();
+        let sizes_canonical = self.sizes().to_primitive();
+        let sizes_slice = sizes_canonical.as_slice::<S>();
+
+        let len = offsets_slice.len();
+
+        let mut new_offsets = BufferMut::<NewOffset>::with_capacity(len);
+        let mut new_sizes = BufferMut::<S>::with_capacity(len);
+        let mut take_indices = BufferMut::<u64>::with_capacity(self.elements().len());
+
+        let mut n_elements = NewOffset::zero();
+        for index in 0..len {
+            if !self.is_valid(index)? {
+                new_offsets.push(n_elements);
+                new_sizes.push(S::zero());
+                continue;
+            }
+
+            let offset = offsets_slice[index];
+            let size = sizes_slice[index];
+            let start = offset.as_();
+            let stop = start + size.as_();
+
+            new_offsets.push(n_elements);
+            new_sizes.push(size);
+            take_indices.extend(start as u64..stop as u64);
+            n_elements += num_traits::cast(size).vortex_expect("Cast failed");
+        }
+
+        let elements = self.elements().take(take_indices.into_array())?;
+        let offsets = new_offsets.into_array();
+        let sizes = new_sizes.into_array();
+
+        // SAFETY: same invariants as `rebuild_list_by_list` — offsets are sequential and
+        // non-overlapping, all (offset, size) pairs reference valid elements, and the validity
+        // array is preserved from the original.
+        Ok(unsafe {
+            ListViewArray::new_unchecked(elements, offsets, sizes, self.validity.clone())
+                .with_zero_copy_to_list(true)
+        })
+    }
+
+    /// Rebuilds elements list-by-list: canonicalize elements upfront, then for each list `slice`
+    /// the relevant range and `extend_from_array` into a typed builder.
+    fn rebuild_list_by_list<O: IntegerPType, NewOffset: IntegerPType, S: IntegerPType>(
+        &self,
     ) -> VortexResult<ListViewArray> {
         let element_dtype = self
             .dtype()
@@ -273,6 +353,7 @@ impl ListViewArray {
 }
 
 #[cfg(test)]
+#[allow(clippy::cast_possible_truncation)]
 mod tests {
     use vortex_buffer::BitBuffer;
     use vortex_error::VortexResult;
@@ -459,4 +540,19 @@ mod tests {
         );
         Ok(())
     }
+
+    // ── should_use_take heuristic tests ────────────────────────────────────
+
+    #[test]
+    fn heuristic_zero_lists_uses_take() {
+        assert!(ListViewArray::should_use_take(0, 0));
+    }
+
+    #[test]
+    fn heuristic_small_lists_use_take() {
+        // avg = 127 → take
+        assert!(ListViewArray::should_use_take(127_000, 1_000));
+        // avg = 128 → LBL
+        assert!(!ListViewArray::should_use_take(128_000, 1_000));
+    }
 }

vortex-array/src/arrow/executor/list.rs

@@ -1,7 +1,6 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
-use std::any::type_name;
 use std::sync::Arc;
 
 use arrow_array::ArrayRef as ArrowArrayRef;
@@ -13,26 +12,23 @@ use vortex_buffer::BufferMut;
 use vortex_error::VortexExpect;
 use vortex_error::VortexResult;
 use vortex_error::vortex_ensure;
-use vortex_error::vortex_err;
 
 use crate::Array;
 use crate::ArrayRef;
 use crate::Canonical;
 use crate::ExecutionCtx;
-use crate::IntoArray;
 use crate::arrays::ListArray;
 use crate::arrays::ListVTable;
 use crate::arrays::ListViewArray;
 use crate::arrays::ListViewArrayParts;
+use crate::arrays::ListViewRebuildMode;
 use crate::arrays::ListViewVTable;
-use crate::arrays::PrimitiveArray;
 use crate::arrow::ArrowArrayExecutor;
 use crate::arrow::executor::validity::to_arrow_null_buffer;
 use crate::builtins::ArrayBuiltins;
 use crate::dtype::DType;
 use crate::dtype::NativePType;
 use crate::dtype::Nullability;
-use crate::validity::Validity;
 use crate::vtable::ValidityHelper;
 
 /// Convert a Vortex array into an Arrow GenericBinaryArray.
@@ -46,37 +42,30 @@ pub(super) fn to_arrow_list<O: OffsetSizeTrait + NativePType>(
         return list_to_list::<O>(array, elements_field, ctx);
     }
 
-    // If the Vortex array is a ListViewArray, we check for our magic cheap conversion flag.
+    // If the Vortex array is a ListViewArray, rebuild to ZCTL if needed and convert.
     let array = match array.try_into::<ListViewVTable>() {
         Ok(array) => {
-            if array.is_zero_copy_to_list() {
-                return list_view_zctl::<O>(array, elements_field, ctx);
+            let zctl = if array.is_zero_copy_to_list() {
+                array
             } else {
-                return list_view_to_list::<O>(array, elements_field, ctx);
-            }
+                array.rebuild(ListViewRebuildMode::MakeZeroCopyToList)?
+            };
+            return list_view_zctl::<O>(zctl, elements_field, ctx);
         }
         Err(a) => a,
     };
 
-    // TODO(ngates): we should do the slightly more expensive thing which is to verify ZCTL.
-    //  In other words, check that offsets + sizes are monotonically increasing.
-
-    // Otherwise, we execute the array to become a ListViewArray.
+    // Otherwise, we execute the array to become a ListViewArray, then rebuild to ZCTL.
+    // Note: arrow_cast::cast supports ListView → List (apache/arrow-rs#8735), but it
+    // unconditionally uses take. Our rebuild uses a heuristic that picks list-by-list
+    // for large lists, which avoids materializing a large index buffer.
     let list_view = array.execute::<ListViewArray>(ctx)?;
-    if list_view.is_zero_copy_to_list() {
-        list_view_zctl::<O>(list_view, elements_field, ctx)
+    let zctl = if list_view.is_zero_copy_to_list() {
+        list_view
     } else {
-        list_view_to_list::<O>(list_view, elements_field, ctx)
-    }
-
-    // FIXME(ngates): we need this PR from arrow-rs:
-    //  https://github.com/apache/arrow-rs/pull/8735
-    // let list_view = array.execute(session)?.into_arrow()?;
-    // match O::IS_LARGE {
-    //     true => arrow_cast::cast(&list_view, &DataType::LargeList(elements_field.clone())),
-    //     false => arrow_cast::cast(&list_view, &DataType::List(elements_field.clone())),
-    // }
-    // .map_err(VortexError::from)
+        list_view.rebuild(ListViewRebuildMode::MakeZeroCopyToList)?
+    };
+    list_view_zctl::<O>(zctl, elements_field, ctx)
 }
 
 /// Convert a Vortex VarBinArray into an Arrow GenericBinaryArray.
@@ -190,70 +179,6 @@ fn list_view_zctl<O: OffsetSizeTrait + NativePType>(
     )))
 }
 
-fn list_view_to_list<O: OffsetSizeTrait + NativePType>(
-    array: ListViewArray,
-    elements_field: &FieldRef,
-    ctx: &mut ExecutionCtx,
-) -> VortexResult<ArrowArrayRef> {
-    let ListViewArrayParts {
-        elements,
-        offsets,
-        sizes,
-        validity,
-        ..
-    } = array.into_parts();
-
-    let offsets = offsets
-        .cast(DType::Primitive(O::PTYPE, Nullability::NonNullable))?
-        .execute::<Canonical>(ctx)?
-        .into_primitive()
-        .to_buffer::<O>();
-    let sizes = sizes
-        .cast(DType::Primitive(O::PTYPE, Nullability::NonNullable))?
-        .execute::<Canonical>(ctx)?
-        .into_primitive()
-        .to_buffer::<O>();
-
-    // We create a new offsets buffer for the final list array.
-    // And we also create an `indices` buffer for taking the elements.
-    let mut new_offsets = BufferMut::<O>::with_capacity(offsets.len() + 1);
-    let mut take_indices = BufferMut::<u32>::with_capacity(elements.len());
-
-    // Add the offset for the first subarray
-    new_offsets.push(O::zero());
-    for (offset, size) in offsets.iter().zip(sizes.iter()) {
-        let offset = offset.as_usize();
-        let size = size.as_usize();
-        let end = offset + size;
-        for j in offset..end {
-            take_indices.push(u32::try_from(j).map_err(|_| {
-                vortex_err!("List array too large for {} indices", type_name::<O>())
-            })?);
-        }
-        new_offsets.push(O::usize_as(take_indices.len()));
-    }
-    assert_eq!(new_offsets.len(), offsets.len() + 1);
-
-    // Now we can "take" the elements using the computed indices.
-    let elements =
-        elements.take(PrimitiveArray::new(take_indices, Validity::NonNullable).into_array())?;
-
-    let elements = elements.execute_arrow(Some(elements_field.data_type()), ctx)?;
-    vortex_ensure!(
-        elements_field.is_nullable() || elements.null_count() == 0,
-        "Cannot convert to non-nullable Arrow array with null elements"
-    );
-
-    let null_buffer = to_arrow_null_buffer(validity, sizes.len(), ctx)?;
-
-    Ok(Arc::new(GenericListArray::<O>::new(
-        elements_field.clone(),
-        new_offsets.freeze().into_arrow_offset_buffer(),
-        elements,
-        null_buffer,
-    )))
-}
-
 #[cfg(test)]
 mod tests {
     use std::sync::Arc;
@@ -364,6 +289,45 @@ mod tests {
         Ok(())
     }
 
+    #[test]
+    fn test_to_arrow_list_non_zctl() -> VortexResult<()> {
+        // Overlapping lists are NOT zero-copy-to-list, so this exercises the rebuild path.
+        // Elements: [1, 2, 3, 4], List 0: [1,2,3], List 1: [2,3,4] (overlap at indices 1-2)
+        let elements = PrimitiveArray::new(buffer![1i32, 2, 3, 4], Validity::NonNullable);
+        let offsets = PrimitiveArray::new(buffer![0i32, 1], Validity::NonNullable);
+        let sizes = PrimitiveArray::new(buffer![3i32, 3], Validity::NonNullable);
+
+        let list_array = ListViewArray::new(
+            elements.into_array(),
+            offsets.into_array(),
+            sizes.into_array(),
+            Validity::NonNullable,
+        );
+        assert!(!list_array.is_zero_copy_to_list());
+
+        let field = Field::new("item", DataType::Int32, false);
+        let arrow_dt = DataType::List(field.into());
+        let arrow_array = list_array.into_array().into_arrow(&arrow_dt)?;
+
+        let list = arrow_array
+            .as_any()
+            .downcast_ref::<GenericListArray<i32>>()
+            .unwrap();
+
+        assert_eq!(list.len(), 2);
+
+        let first = list.value(0);
+        assert_eq!(first.len(), 3);
+        let first_vals = first.as_any().downcast_ref::<Int32Array>().unwrap();
+        assert_eq!(first_vals.values(), &[1, 2, 3]);
+
+        let second = list.value(1);
+        assert_eq!(second.len(), 3);
+        let second_vals = second.as_any().downcast_ref::<Int32Array>().unwrap();
+        assert_eq!(second_vals.values(), &[2, 3, 4]);
+        Ok(())
+    }
+
     #[test]
     fn test_to_arrow_list_empty_zctl() -> VortexResult<()> {
         let dtype = DType::List(