Fix: incorrect offset for nulls + don't chunk elements (#5423)

connortsui20 · web-flow · commit 71f0f0018035 · 2025-11-20T12:00:37.000-05:00
Potentially fixes? #5412 A continuation of #5417 # Notes There was an issue where the offset being appended for null lists was the previous offset instead of the previous offset + previous size which broke the invariant of no overlaps. The `validate_zctl` failed to catch this edge case. This then broke `rebuild_trim_elements` since it assumes that everything is in order, including the last offset. This PR fixes that bug. Additionally, this PR cleans up several things in the `ListView` implementation and improves the performance of a `naive_rebuild` and `list_from_list_view`. On develop on my AMD 7950X: ``` Timer precision: 10 ns listview_rebuild fastest │ slowest │ median │ mean │ samples │ iters ╰─ rebuild_naive 158.8 µs │ 175.3 µs │ 160.7 µs │ 162.6 µs │ 100 │ 100000 ``` On this branch: ``` Timer precision: 10 ns listview_rebuild fastest │ slowest │ median │ mean │ samples │ iters ╰─ rebuild_naive 110 µs │ 149.1 µs │ 110.7 µs │ 111.8 µs │ 100 │ 100000 ``` Signed-off-by: Connor Tsui <connor.tsui20@gmail.com>
diff --git a/encodings/zstd/benches/listview_rebuild.rs b/encodings/zstd/benches/listview_rebuild.rs
@@ -10,7 +10,7 @@ use vortex_array::validity::Validity;
 use vortex_buffer::Buffer;
 use vortex_zstd::ZstdArray;
 
-#[divan::bench]
+#[divan::bench(sample_size = 1000)]
 fn rebuild_naive(bencher: Bencher) {
     let dudes = VarBinViewArray::from_iter_str(["Washington", "Adams", "Jefferson", "Madison"])
         .into_array();
diff --git a/vortex-array/src/arrays/listview/array.rs b/vortex-array/src/arrays/listview/array.rs
@@ -99,8 +99,9 @@ pub struct ListViewArray {
     ///
     /// We use this information to help us more efficiently rebuild / compact our data.
     ///
-    /// When this flag is true (indicating sorted offsets with no gaps and no overlaps), conversions
-    /// can bypass the very expensive rebuild process (which just calls `append_scalar` in a loop).
+    /// When this flag is true (indicating sorted offsets with no gaps and no overlaps and all
+    /// `offsets[i] + sizes[i]` are in order), conversions can bypass the very expensive rebuild
+    /// process which must rebuild the array from scratch.
     is_zero_copy_to_list: bool,
 
     /// The validity / null map of the array.
@@ -264,6 +265,7 @@ impl ListViewArray {
     /// actually zero-copyable to a [`ListArray`]. This means:
     ///
     /// - Offsets must be sorted (but not strictly sorted, zero-length lists are allowed).
+    /// - `offsets[i] + sizes[i] == offsets[i + 1]` for all `i`.
     /// - No gaps in elements between first and last referenced elements.
     /// - No overlapping list views (each element referenced at most once).
     ///
@@ -425,7 +427,8 @@ where
         if offset_u64 == elements_len {
             vortex_ensure!(
                 size_u64 == 0,
-                "views to the end of the elements array (length {elements_len}) must have size 0"
+                "views to the end of the elements array (length {elements_len}) must have size 0 \
+                    (had size {size_u64})"
             );
         }
 
@@ -454,6 +457,51 @@ fn validate_zctl(
         vortex_bail!("offsets must report is_sorted statistic");
     }
 
+    // Validate that offset[i] + size[i] <= offset[i+1] for all items
+    // This ensures views are non-overlapping and properly ordered for zero-copy-to-list
+    fn validate_monotonic_ends<O: IntegerPType, S: IntegerPType>(
+        offsets_slice: &[O],
+        sizes_slice: &[S],
+        len: usize,
+    ) -> VortexResult<()> {
+        let mut max_end = 0usize;
+
+        for i in 0..len {
+            let offset = offsets_slice[i].to_usize().unwrap_or(usize::MAX);
+            let size = sizes_slice[i].to_usize().unwrap_or(usize::MAX);
+
+            // Check that this view starts at or after the previous view ended
+            vortex_ensure!(
+                offset >= max_end,
+                "Zero-copy-to-list requires views to be non-overlapping and ordered: \
+                 view[{}] starts at {} but previous views extend to {}",
+                i,
+                offset,
+                max_end
+            );
+
+            // Update max_end for the next iteration
+            let end = offset.saturating_add(size);
+            max_end = max_end.max(end);
+        }
+
+        Ok(())
+    }
+
+    let offsets_dtype = offsets_primitive.dtype();
+    let sizes_dtype = sizes_primitive.dtype();
+    let len = offsets_primitive.len();
+
+    // Check that offset + size values are monotonic (no overlaps)
+    match_each_integer_ptype!(offsets_dtype.as_ptype(), |O| {
+        match_each_integer_ptype!(sizes_dtype.as_ptype(), |S| {
+            let offsets_slice = offsets_primitive.as_slice::<O>();
+            let sizes_slice = sizes_primitive.as_slice::<S>();
+
+            validate_monotonic_ends(offsets_slice, sizes_slice, len)?;
+        })
+    });
+
     // TODO(connor)[ListView]: Making this allocation is expensive, but the more efficient
     // implementation would be even more complicated than this. We could use a bit buffer denoting
     // if positions in `elements` are used, and then additionally store a separate flag that tells
diff --git a/vortex-array/src/arrays/listview/conversion.rs b/vortex-array/src/arrays/listview/conversion.rs
@@ -6,8 +6,10 @@ use std::sync::Arc;
 use vortex_dtype::{IntegerPType, Nullability, match_each_integer_ptype};
 use vortex_error::VortexExpect;
 
-use crate::arrays::{ExtensionArray, FixedSizeListArray, ListArray, ListViewArray, StructArray};
-use crate::builders::{ArrayBuilder, ListBuilder, PrimitiveBuilder};
+use crate::arrays::{
+    ExtensionArray, FixedSizeListArray, ListArray, ListViewArray, ListViewRebuildMode, StructArray,
+};
+use crate::builders::PrimitiveBuilder;
 use crate::vtable::ValidityHelper;
 use crate::{Array, ArrayRef, Canonical, IntoArray, ToCanonical};
 
@@ -92,56 +94,26 @@ fn build_sizes_from_offsets<O: IntegerPType>(list: &ListArray) -> ArrayRef {
 /// Otherwise, this function fall back to the (very) expensive path and will rebuild the
 /// [`ListArray`] from scratch.
 pub fn list_from_list_view(list_view: ListViewArray) -> ListArray {
-    // Fast path if the array is zero-copyable to a `ListArray`.
-    if list_view.is_zero_copy_to_list() {
-        let list_offsets = match_each_integer_ptype!(list_view.offsets().dtype().as_ptype(), |O| {
-            // SAFETY: We checked that the array is zero-copyable to `ListArray`, so the safety
-            // contract is upheld.
-            unsafe { build_list_offsets_from_list_view::<O>(&list_view) }
-        });
-
-        // SAFETY: Because the shape of the `ListViewArray` is zero-copyable to a `ListArray`, we
-        // can simply reuse all of the data (besides the offsets).
-        let new_array = unsafe {
-            ListArray::new_unchecked(
-                list_view.elements().clone(),
-                list_offsets,
-                list_view.validity().clone(),
-            )
-        };
-
-        let new_array = new_array
-            .reset_offsets(false)
-            .vortex_expect("TODO(connor)[ListView]: This can't fail");
+    // Rebuild as zero-copyable to list array and also trim all leading and trailing elements.
+    let zctl_array = list_view.rebuild(ListViewRebuildMode::MakeExact);
+    debug_assert!(zctl_array.is_zero_copy_to_list());
+
+    let list_offsets = match_each_integer_ptype!(zctl_array.offsets().dtype().as_ptype(), |O| {
+        // SAFETY: We just made the array zero-copyable to `ListArray`, so the safety contract is
+        // upheld.
+        unsafe { build_list_offsets_from_list_view::<O>(&zctl_array) }
+    });
 
-        return new_array;
+    // SAFETY: Because the shape of the `ListViewArray` is zero-copyable to a `ListArray`, we
+    // can simply reuse all of the data (besides the offsets). We also trim all of the elements to
+    // make it easier for the caller to use the `ListArray`.
+    unsafe {
+        ListArray::new_unchecked(
+            zctl_array.elements().clone(),
+            list_offsets,
+            zctl_array.validity().clone(),
+        )
     }
-
-    let elements_dtype = list_view
-        .dtype()
-        .as_list_element_opt()
-        .vortex_expect("`DType` of `ListView` was somehow not a `List`");
-    let nullability = list_view.dtype().nullability();
-    let len = list_view.len();
-
-    match_each_integer_ptype!(list_view.offsets().dtype().as_ptype(), |O| {
-        let mut builder = ListBuilder::<O>::with_capacity(
-            elements_dtype.clone(),
-            nullability,
-            list_view.elements().len(),
-            len,
-        );
-
-        for i in 0..len {
-            builder
-                .append_scalar(&list_view.scalar_at(i))
-                .vortex_expect(
-                    "The `ListView` scalars are `ListScalar`, which the `ListBuilder` must accept",
-                )
-        }
-
-        builder.finish_into_list()
-    })
 }
 
 // TODO(connor)[ListView]: We can optimize this by always keeping extra memory in `ListViewArray`
@@ -165,6 +137,7 @@ unsafe fn build_list_offsets_from_list_view<O: IntegerPType>(
 
     let offsets = list_view.offsets().to_primitive();
     let offsets_slice = offsets.as_slice::<O>();
+    debug_assert!(offsets_slice.is_sorted());
 
     // Copy the existing n offsets.
     offsets_range.copy_from_slice(0, offsets_slice);
diff --git a/vortex-array/src/arrays/listview/rebuild.rs b/vortex-array/src/arrays/listview/rebuild.rs
@@ -7,7 +7,8 @@ use vortex_dtype::{IntegerPType, Nullability, match_each_integer_ptype};
 use vortex_error::VortexExpect;
 use vortex_scalar::Scalar;
 
-use crate::arrays::{ChunkedArray, ListViewArray};
+use crate::arrays::ListViewArray;
+use crate::builders::builder_with_capacity;
 use crate::vtable::ValidityHelper;
 use crate::{Array, IntoArray, ToCanonical, compute};
 
@@ -74,6 +75,13 @@ impl ListViewArray {
         let offsets_ptype = self.offsets().dtype().as_ptype();
         let sizes_ptype = self.sizes().dtype().as_ptype();
 
+        // One of the main purposes behind adding this "zero-copyable to `ListArray`" optimization
+        // is that we want to pass data to systems that expect Arrow data.
+        // The arrow specification only allows for `i32` and `i64` offset and sizes types, so in
+        // order to also make `ListView` zero-copyable to **Arrow**'s `ListArray` (not just Vortex's
+        // `ListArray`), we rebuild the offsets as 32-bit or 64-bit integer types.
+        // TODO(connor)[ListView]: This is true for `sizes` as well, we should do this conversion
+        // for sizes as well.
         match_each_integer_ptype!(sizes_ptype, |S| {
             match offsets_ptype {
                 PType::U8 => self.naive_rebuild::<u8, u32, S>(),
@@ -89,8 +97,10 @@ impl ListViewArray {
         })
     }
 
-    /// The inner function for `rebuild_zero_copy_to_list`, which naively rebuilds a `ListViewArray`
-    /// via `append_scalar`.
+    // 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.
     fn naive_rebuild<O: IntegerPType, NewOffset: IntegerPType, S: IntegerPType>(
         &self,
     ) -> ListViewArray {
@@ -99,59 +109,72 @@ impl ListViewArray {
             .as_list_element_opt()
             .vortex_expect("somehow had a canonical list that was not a list");
 
-        // Upfront canonicalize the list elements, we're going to be doing a lot of
-        // slicing with them.
-        let elements_canonical = self.elements().to_canonical().into_array();
         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 offsets_canonical = offsets_canonical.as_slice::<O>();
-        let sizes_canonical = sizes_canonical.as_slice::<S>();
+        let len = offsets_slice.len();
 
-        let mut offsets = BufferMut::<NewOffset>::with_capacity(self.len());
-        let mut sizes = BufferMut::<S>::with_capacity(self.len());
+        let mut new_offsets = BufferMut::<NewOffset>::with_capacity(len);
+        // TODO(connor)[ListView]: Do we really need to do this?
+        // The only reason we need to rebuild the sizes here is that the validity may indicate that
+        // a list is null even though it has a non-zero size. This rebuild will set the size of all
+        // null lists to 0.
+        let mut new_sizes = BufferMut::<S>::with_capacity(len);
 
-        let mut chunks = Vec::with_capacity(self.len());
+        // Canonicalize the elements up front as we will be slicing the elements quite a lot.
+        let elements_canonical = self.elements().to_canonical().into_array();
 
-        let mut n_elements = NewOffset::zero();
+        // Note that we do not know what the exact capacity should be of the new elements since
+        // there could be overlaps in the existing `ListViewArray`.
+        let mut new_elements_builder =
+            builder_with_capacity(element_dtype.as_ref(), self.elements().len());
 
-        for index in 0..self.len() {
+        let mut n_elements = NewOffset::zero();
+        for index in 0..len {
             if !self.is_valid(index) {
-                offsets.push(offsets.last().copied().unwrap_or_default());
-                sizes.push(S::zero());
+                // For NULL lists, place them after the previous item's data to maintain the
+                // no-overlap invariant for zero-copy to `ListArray` arrays.
+                new_offsets.push(n_elements);
+                new_sizes.push(S::zero());
                 continue;
             }
 
-            let offset = offsets_canonical[index];
-            let size = sizes_canonical[index];
+            let offset = offsets_slice[index];
+            let size = sizes_slice[index];
 
             let start = offset.as_();
             let stop = start + size.as_();
 
-            chunks.push(elements_canonical.slice(start..stop));
-            offsets.push(n_elements);
-            sizes.push(size);
+            new_offsets.push(n_elements);
+            new_sizes.push(size);
+            new_elements_builder.extend_from_array(&elements_canonical.slice(start..stop));
 
-            n_elements += num_traits::cast(size).vortex_expect("cast");
+            n_elements += num_traits::cast(size).vortex_expect("Cast failed");
         }
 
-        let offsets = offsets.into_array();
-        let sizes = sizes.into_array();
+        let offsets = new_offsets.into_array();
+        let sizes = new_sizes.into_array();
+        let elements = new_elements_builder.finish();
 
-        // SAFETY: all chunks were sliced from the same array so have same DType.
-        let elements =
-            unsafe { ChunkedArray::new_unchecked(chunks, element_dtype.as_ref().clone()) };
+        debug_assert_eq!(
+            n_elements.as_(),
+            elements.len(),
+            "The accumulated elements somehow had the wrong length"
+        );
 
-        // SAFETY: elements are contiguous, offsets and sizes hand-built to be zero copy
-        //  to list.
+        // SAFETY:
+        // - All offsets are sequential and non-overlapping (`n_elements` tracks running total).
+        // - Each `offset[i] + size[i]` equals `offset[i+1]` for all valid indices (including null
+        //   lists).
+        // - All elements referenced by (offset, size) pairs exist within the new `elements` array.
+        // - The validity array is preserved from the original array unchanged
+        // - The array satisfies the zero-copy-to-list property by having sorted offsets, no gaps,
+        //   and no overlaps.
         unsafe {
-            ListViewArray::new_unchecked(
-                elements.to_canonical().into_array(),
-                offsets,
-                sizes,
-                self.validity.clone(),
-            )
-            .with_zero_copy_to_list(true)
+            ListViewArray::new_unchecked(elements, offsets, sizes, self.validity.clone())
+                .with_zero_copy_to_list(true)
         }
     }
 
@@ -342,4 +365,53 @@ mod tests {
         let all_elements = trimmed.elements().to_primitive();
         assert_eq!(all_elements.scalar_at(2), 97i32.into());
     }
+
+    #[test]
+    fn test_rebuild_with_trailing_nulls_regression() {
+        // Regression test for issue #5412
+        // Tests that zero-copy-to-list arrays with trailing NULLs correctly calculate
+        // offsets for NULL items to maintain no-overlap invariant
+
+        // Create a ListViewArray with trailing NULLs
+        let elements = PrimitiveArray::from_iter(vec![1i32, 2, 3, 4]).into_array();
+        let offsets = PrimitiveArray::from_iter(vec![0u32, 2, 0, 0]).into_array();
+        let sizes = PrimitiveArray::from_iter(vec![2u32, 2, 0, 0]).into_array();
+        let validity = Validity::from_iter(vec![true, true, false, false]);
+
+        let listview = ListViewArray::new(elements, offsets, sizes, validity);
+
+        // First rebuild to make it zero-copy-to-list
+        let rebuilt = listview.rebuild(ListViewRebuildMode::MakeZeroCopyToList);
+        assert!(rebuilt.is_zero_copy_to_list());
+
+        // Verify NULL items have correct offsets (should not reuse previous offsets)
+        // After rebuild: offsets should be [0, 2, 4, 4] for zero-copy-to-list
+        assert_eq!(rebuilt.offset_at(0), 0);
+        assert_eq!(rebuilt.offset_at(1), 2);
+        assert_eq!(rebuilt.offset_at(2), 4); // NULL should be at position 4
+        assert_eq!(rebuilt.offset_at(3), 4); // Second NULL also at position 4
+
+        // All sizes should be correct
+        assert_eq!(rebuilt.size_at(0), 2);
+        assert_eq!(rebuilt.size_at(1), 2);
+        assert_eq!(rebuilt.size_at(2), 0); // NULL has size 0
+        assert_eq!(rebuilt.size_at(3), 0); // NULL has size 0
+
+        // Now rebuild with MakeExact (which calls naive_rebuild then trim_elements)
+        // This should not panic (issue #5412)
+        let exact = rebuilt.rebuild(ListViewRebuildMode::MakeExact);
+
+        // Verify the result is still valid
+        assert!(exact.is_valid(0));
+        assert!(exact.is_valid(1));
+        assert!(!exact.is_valid(2));
+        assert!(!exact.is_valid(3));
+
+        // Verify data is preserved
+        let list0 = exact.list_elements_at(0).to_primitive();
+        assert_eq!(list0.as_slice::<i32>(), &[1, 2]);
+
+        let list1 = exact.list_elements_at(1).to_primitive();
+        assert_eq!(list1.as_slice::<i32>(), &[3, 4]);
+    }
 }
diff --git a/vortex-array/src/arrays/listview/tests/basic.rs b/vortex-array/src/arrays/listview/tests/basic.rs
diff --git a/vortex-array/src/arrow/compute/to_arrow/canonical.rs b/vortex-array/src/arrow/compute/to_arrow/canonical.rs
