is_zero_copy_to_list ListView to List optimization

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

Author: connortsui20

encodings/sparse/src/canonical.rs

@@ -931,9 +931,16 @@ mod test {
         // List 3: [2] at offset 3, size 1
         let offsets = buffer![0u32, 1, 2, 3].into_array();
         let sizes = buffer![1u32, 1, 1, 1].into_array();
-        let lists = ListViewArray::try_new(elements, offsets, sizes, Validity::AllValid)
-            .unwrap()
-            .into_array();
+        let lists = unsafe {
+            ListViewArray::new_unchecked(
+                elements,
+                offsets,
+                sizes,
+                Validity::AllValid,
+                true, // Is zero-copy to list.
+            )
+        }
+        .into_array();
 
         let indices = buffer![0u8, 3u8, 4u8, 5u8].into_array();
         let fill_value = Scalar::null(lists.dtype().clone());
@@ -978,9 +985,16 @@ mod test {
         let elements = buffer![1i32, 2, 1, 2, 1, 2, 1, 2].into_array();
         let offsets = buffer![0u32, 1, 2, 3, 4, 5, 6, 7].into_array();
         let sizes = buffer![1u32, 1, 1, 1, 1, 1, 1, 1].into_array();
-        let lists = ListViewArray::try_new(elements, offsets, sizes, Validity::AllValid)
-            .unwrap()
-            .into_array();
+        let lists = unsafe {
+            ListViewArray::new_unchecked(
+                elements,
+                offsets,
+                sizes,
+                Validity::AllValid,
+                true, // Is zero-copy to list.
+            )
+        }
+        .into_array();
 
         // Slice to get lists 2..6, which are: [1], [2], [1], [2]
         let lists = lists.slice(2..6);
@@ -1022,9 +1036,16 @@ mod test {
         let elements = buffer![1i32, 2, 1, 2].into_array();
         let offsets = buffer![0u32, 1, 2, 3].into_array();
         let sizes = buffer![1u32, 1, 1, 1].into_array();
-        let lists = ListViewArray::try_new(elements, offsets, sizes, Validity::AllValid)
-            .unwrap()
-            .into_array();
+        let lists = unsafe {
+            ListViewArray::new_unchecked(
+                elements,
+                offsets,
+                sizes,
+                Validity::AllValid,
+                true, // Is zero-copy to list.
+            )
+        }
+        .into_array();
 
         let indices = buffer![0u8, 3u8, 4u8, 5u8].into_array();
         let fill_value = Scalar::from(Some(vec![5i32, 6, 7, 8]));
@@ -1384,8 +1405,15 @@ mod test {
         let offsets = buffer![0u32, 3, 5].into_array();
         let sizes = buffer![3u32, 2, 4].into_array();
 
-        let list_view =
-            ListViewArray::try_new(elements.clone(), offsets, sizes, Validity::AllValid).unwrap();
+        let list_view = unsafe {
+            ListViewArray::new_unchecked(
+                elements.clone(),
+                offsets,
+                sizes,
+                Validity::AllValid,
+                true, // Is zero-copy to list.
+            )
+        };
 
         let list_dtype = list_view.dtype().clone();
 
@@ -1460,9 +1488,16 @@ mod test {
         let offsets = buffer![0u32, 2, 5, 6, 8].into_array();
         let sizes = buffer![2u32, 3, 1, 2, 2].into_array();
 
-        let full_listview = ListViewArray::try_new(elements, offsets, sizes, Validity::AllValid)
-            .unwrap()
-            .into_array();
+        let full_listview = unsafe {
+            ListViewArray::new_unchecked(
+                elements,
+                offsets,
+                sizes,
+                Validity::AllValid,
+                true, // Is zero-copy to list.
+            )
+        }
+        .into_array();
 
         // Slice to get lists 1, 2, 3 (indices 1..4)
         // This gives us lists with elements:

fuzz/src/array/mask.rs

@@ -57,13 +57,18 @@ pub fn mask_canonical_array(canonical: Canonical, mask: &Mask) -> VortexResult<A
         }
         Canonical::List(array) => {
             let new_validity = apply_mask_to_validity(array.validity(), mask);
-            ListViewArray::try_new(
-                array.elements().clone(),
-                array.offsets().clone(),
-                array.sizes().clone(),
-                new_validity,
-            )
-            .vortex_unwrap()
+
+            // SAFETY: Since we are only masking the validity and everything else comes from an
+            // already valid `ListViewArray`, all of the invariants are still upheld.
+            unsafe {
+                ListViewArray::new_unchecked(
+                    array.elements().clone(),
+                    array.offsets().clone(),
+                    array.sizes().clone(),
+                    new_validity,
+                    array.is_zero_copy_to_list(),
+                )
+            }
             .into_array()
         }
         Canonical::FixedSizeList(array) => {
@@ -243,9 +248,15 @@ mod tests {
         let elements = PrimitiveArray::from_iter([1i32, 2, 3, 4, 5, 6]).into_array();
         let offsets = PrimitiveArray::from_iter([0i32, 2, 4]).into_array();
         let sizes = PrimitiveArray::from_iter([2i32, 2, 2]).into_array();
-        let array =
-            ListViewArray::try_new(elements, offsets, sizes, Nullability::NonNullable.into())
-                .unwrap();
+        let array = unsafe {
+            ListViewArray::new_unchecked(
+                elements,
+                offsets,
+                sizes,
+                Nullability::NonNullable.into(),
+                true, // Is zero-copy to list.
+            )
+        };
 
         let mask = Mask::from_iter([false, true, false]);
 

fuzz/src/array/slice.rs

@@ -74,7 +74,19 @@ pub fn slice_canonical_array(
             // Since the list view elements can be stored out of order, we cannot slice it.
             let elements = list_array.elements().clone();
 
-            ListViewArray::try_new(elements, offsets, sizes, validity).map(|a| a.into_array())
+            // SAFETY: If the array was already zero-copyable to list, slicing the offsets and sizes
+            // only causes there to be leading and trailing garbage data, which is still
+            // zero-copyable to a `ListArray`.
+            Ok(unsafe {
+                ListViewArray::new_unchecked(
+                    elements,
+                    offsets,
+                    sizes,
+                    validity,
+                    list_array.is_zero_copy_to_list(),
+                )
+            }
+            .into_array())
         }
         DType::FixedSizeList(..) => {
             let fsl_array = array.to_fixed_size_list();

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

@@ -5,7 +5,9 @@ use vortex_buffer::BufferMut;
 use vortex_dtype::{DType, Nullability, PType, StructFields};
 use vortex_error::VortexExpect;
 
-use crate::arrays::{ChunkedArray, ChunkedVTable, ListViewArray, PrimitiveArray, StructArray};
+use crate::arrays::{
+    ChunkedArray, ChunkedVTable, ListViewArray, ListViewRebuildMode, PrimitiveArray, StructArray,
+};
 use crate::builders::{ArrayBuilder, builder_with_capacity};
 use crate::compute::cast;
 use crate::validity::Validity;
@@ -121,6 +123,9 @@ fn swizzle_list_chunks(
 
     for chunk in chunks {
         let chunk_array = chunk.to_listview();
+        // 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());
@@ -158,12 +163,16 @@ fn swizzle_list_chunks(
     let offsets = PrimitiveArray::new(offsets.freeze(), Validity::NonNullable).into_array();
     let sizes = PrimitiveArray::new(sizes.freeze(), Validity::NonNullable).into_array();
 
+    // Since we made sure that all chunk were zero-copyable to list above, we know that the final
+    // output is also zero-copyable to a list.
+    let is_zctl = true;
+
     // SAFETY:
     // - `offsets` and `sizes` are non-nullable u64 arrays of the same length
     // - Each `offset[i] + size[i]` list view is within bounds of elements array because it came
     //   from valid chunks
     // - Validity came from the outer chunked array so it must have the same length
-    unsafe { ListViewArray::new_unchecked(chunked_elements, offsets, sizes, validity) }
+    unsafe { ListViewArray::new_unchecked(chunked_elements, offsets, sizes, validity, is_zctl) }
 }
 
 #[cfg(test)]

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

@@ -253,10 +253,13 @@ fn constant_canonical_list_array(scalar: &Scalar, len: usize) -> ListViewArray {
     debug_assert!(!offsets.dtype().is_nullable());
     debug_assert!(!sizes.dtype().is_nullable());
 
+    // Since all views are pointing to the same exact thing, it is not zero-copyable to `ListArray`.
+    let is_zctl = false;
+
     // SAFETY: All views point to the same range [0, list.len()) in the elements array.
     // The elements array contains `len` copies of the same value, offsets are all 0,
     // and sizes are all equal to the list length. The validity matches the scalar's nullability.
-    unsafe { ListViewArray::new_unchecked(elements, offsets, sizes, validity) }
+    unsafe { ListViewArray::new_unchecked(elements, offsets, sizes, validity, is_zctl) }
 }
 
 fn constant_canonical_fixed_size_list_array(

vortex-array/src/arrays/list/array.rs

@@ -4,7 +4,7 @@
 use std::sync::Arc;
 
 use num_traits::AsPrimitive;
-use vortex_dtype::{DType, match_each_integer_ptype, match_each_native_ptype};
+use vortex_dtype::{DType, NativePType, match_each_integer_ptype, match_each_native_ptype};
 use vortex_error::{VortexExpect, VortexResult, vortex_bail, vortex_ensure};
 
 use crate::arrays::{ListVTable, PrimitiveVTable};
@@ -191,8 +191,11 @@ impl ListArray {
 
                 vortex_ensure!(
                     max_offset
-                        <= P::try_from(elements.len())
-                            .vortex_expect("Offsets type must be able to fit elements length"),
+                        <= P::try_from(elements.len()).vortex_expect(&format!(
+                            "Offsets type {} must be able to fit elements length {}",
+                            <P as NativePType>::PTYPE,
+                            elements.len()
+                        )),
                     "Max offset {max_offset} is beyond the length of the elements array {}",
                     elements.len()
                 );
@@ -266,6 +269,11 @@ impl ListArray {
         &self.elements
     }
 
+    // TODO(connor)[ListView]: Create 2 functions `reset_offsets` and `recursive_reset_offsets`,
+    // where `reset_offsets` is infallible.
+    // Also, `reset_offsets` can be made more efficient by replacing `sub_scalar` with a match on
+    // the offset type and manual subtraction and fast path where `offsets[0] == 0`.
+
     /// Create a copy of this array by adjusting `offsets` to start at `0` and removing elements not
     /// referenced by the `offsets`.
     pub fn reset_offsets(&self, recurse: bool) -> VortexResult<Self> {