fix compressor bug

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

Author: connortsui20

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()
                 );

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

@@ -377,6 +377,13 @@ where
             vortex_err!("offset[{i}] ({offset_u64}) + size[{i}] ({size_u64}) would overflow u64")
         })?;
 
+        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"
+            );
+        }
+
         vortex_ensure!(
             end <= elements_len,
             "offset[{i}] + size[{i}] = {offset_u64} + {size_u64} = {end} \

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

@@ -21,20 +21,22 @@ pub fn list_view_from_list(list: ListArray) -> ListViewArray {
         return Canonical::empty(list.dtype()).into_listview();
     }
 
+    // TODO(connor): Should we reset offsets?
+
     let len = list.len();
 
     // Get the `offsets` array directly from the `ListArray` (preserving its type).
     let list_offsets = list.offsets().clone();
 
-    // We need to slice the `offsets` to remove the last element (`ListArray` has n+1 offsets).
-    let adjusted_offsets = list_offsets.slice(0..len);
-
     // Create sizes array by computing differences between consecutive offsets.
     // Use the same dtype as the offsets array to ensure compatibility.
     let sizes = match_each_integer_ptype!(list_offsets.dtype().as_ptype(), |O| {
         build_sizes_from_offsets::<O>(&list)
     });
 
+    // We need to slice the `offsets` to remove the last element (`ListArray` has `n + 1` offsets).
+    let adjusted_offsets = list_offsets.slice(0..len);
+
     // Since the data came from a valid `ListArray`, we know it is zero-copyable to a `ListArray`.
     let shape = ListViewShape::as_zero_copy_to_list();
 
@@ -58,8 +60,11 @@ fn build_sizes_from_offsets<O: IntegerPType>(list: &ListArray) -> ArrayRef {
 
     // Create `UninitRange` for direct memory access.
     let mut sizes_range = sizes_builder.uninit_range(len);
+
     let offsets = list.offsets().to_primitive();
     let offsets_slice = offsets.as_slice::<O>();
+    debug_assert_eq!(len + 1, offsets_slice.len());
+    debug_assert!(offsets_slice.is_sorted());
 
     // Compute sizes as the difference between consecutive offsets.
     for i in 0..len {

vortex-array/src/builders/listview.rs

@@ -261,12 +261,12 @@ impl<O: IntegerPType, S: IntegerPType> ArrayBuilder for ListViewBuilder<O, S> {
 
         self.nulls.append_validity_mask(array.validity_mask());
 
-        let curr_elements_len = self.elements_builder.len();
-
         // Bulk append the new elements (which should have no gaps or overlaps).
+        let old_elements_len = self.elements_builder.len();
         self.elements_builder
             .reserve_exact(listview.elements().len());
         self.elements_builder.extend_from_array(listview.elements());
+        let new_elements_len = self.elements_builder.len();
 
         // Reserve enough space for the new views.
         let extend_length = listview.len();
@@ -287,36 +287,13 @@ impl<O: IntegerPType, S: IntegerPType> ArrayBuilder for ListViewBuilder<O, S> {
         // This should be cheap because we didn't compress after rebuilding.
         let new_offsets = listview.offsets().to_primitive();
 
-        fn adjust_offsets<'a, O: IntegerPType, A: IntegerPType>(
-            mut uninit_range: UninitRange<'a, O>,
-            new_offsets: PrimitiveArray,
-            curr_elements_len: usize,
-        ) {
-            let new_offsets_slice = new_offsets.as_slice::<A>();
-            let curr_elements_len = O::from_usize(curr_elements_len).vortex_expect(
-                "the current elements 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");
-
-                let adjusted_new_offset = new_offset_correct_type + curr_elements_len;
-                uninit_range.set_value(i, adjusted_new_offset);
-            }
-
-            // 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_elements_len);
+            adjust_offsets::<O, A>(
+                uninit_range,
+                new_offsets,
+                old_elements_len,
+                new_elements_len,
+            );
         })
     }
 
@@ -341,6 +318,42 @@ impl<O: IntegerPType, S: IntegerPType> ArrayBuilder for ListViewBuilder<O, S> {
     }
 }
 
+/// Given new offsets, adds them to the `UninitRange` after adding the `old_elements_len` to each
+/// offset.
+fn adjust_offsets<'a, O: IntegerPType, A: IntegerPType>(
+    mut uninit_range: UninitRange<'a, O>,
+    new_offsets: PrimitiveArray,
+    old_elements_len: usize,
+    new_elements_len: usize,
+) {
+    let new_offsets_slice = new_offsets.as_slice::<A>();
+    let old_elements_len = O::from_usize(old_elements_len)
+        .vortex_expect("the old elements length did not fit into the offset type (impossible)");
+    let new_elements_len = O::from_usize(new_elements_len)
+        .vortex_expect("the current elements length did not fit into the offset type (impossible)");
+
+    for i in 0..uninit_range.len() {
+        let new_offset = O::from_usize(
+            new_offsets_slice[i]
+                .to_usize()
+                .vortex_expect("Offsets must always fit in usize"),
+        )
+        .vortex_expect("New offset somehow did not fit into the builder's offset type");
+
+        let adjusted_new_offset = new_offset + old_elements_len;
+        debug_assert!(
+            adjusted_new_offset <= new_elements_len,
+            "[{i}/{}]: {new_offset} + {old_elements_len} = {adjusted_new_offset} <= {new_elements_len} failed",
+            uninit_range.len()
+        );
+        uninit_range.set_value(i, adjusted_new_offset);
+    }
+
+    // SAFETY: We have set all the values in the range, and since `offsets` are
+    // non-nullable, we are done.
+    unsafe { uninit_range.finish() };
+}
+
 #[cfg(test)]
 mod tests {
     use std::sync::Arc;

vortex-btrblocks/src/lib.rs

@@ -403,10 +403,15 @@ impl BtrBlocksCompressor {
                 // list and list view.
                 let list_array = list_from_list_view(list_view_array);
 
+                // Reset the offsets to remove garbage data that might prevent us from narrowing our
+                // offsets.
+                let list_array = list_array.reset_offsets(true)?;
+
                 let compressed_elems = self.compress(list_array.elements())?;
 
-                // Note that since the type of our offsets are not encoded in our `DType`,
-                // we may narrow the widths.
+                // Note that since the type of our offsets are not encoded in our `DType`, and since
+                // we guarantee above that all elements are referenced by offsets, we may narrow the
+                // widths.
 
                 let compressed_offsets = IntCompressor::compress_no_dict(
                     &list_array.offsets().to_primitive().narrow()?,