feat[vortex-array] chunked array has owned chunk_offsets child (#5426)

We should ensure that visitors are given ref to externally owned array

Signed-off-by: Joe Isaacs <[email protected]>

Author: joseph-isaacs

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

@@ -12,17 +12,19 @@ use vortex_buffer::{Buffer, BufferMut};
 use vortex_dtype::DType;
 use vortex_error::{VortexExpect as _, VortexResult, VortexUnwrap, vortex_bail};
 
+use crate::arrays::PrimitiveArray;
 use crate::iter::{ArrayIterator, ArrayIteratorAdapter};
 use crate::search_sorted::{SearchSorted, SearchSortedSide};
 use crate::stats::ArrayStats;
 use crate::stream::{ArrayStream, ArrayStreamAdapter};
+use crate::validity::Validity;
 use crate::{Array, ArrayRef, IntoArray};
 
 #[derive(Clone, Debug)]
 pub struct ChunkedArray {
     pub(super) dtype: DType,
     pub(super) len: usize,
-    pub(super) chunk_offsets: Buffer<u64>,
+    pub(super) chunk_offsets: PrimitiveArray,
     pub(super) chunks: Vec<ArrayRef>,
     pub(super) stats_set: ArrayStats,
 }
@@ -58,20 +60,22 @@ impl ChunkedArray {
 
         let nchunks = chunks.len();
 
-        let mut chunk_offsets = BufferMut::<u64>::with_capacity(nchunks + 1);
+        let mut chunk_offsets_buf = BufferMut::<u64>::with_capacity(nchunks + 1);
         // SAFETY: nchunks + 1
-        unsafe { chunk_offsets.push_unchecked(0) }
+        unsafe { chunk_offsets_buf.push_unchecked(0) }
         let mut curr_offset = 0;
         for c in &chunks {
             curr_offset += c.len() as u64;
             // SAFETY: nchunks + 1
-            unsafe { chunk_offsets.push_unchecked(curr_offset) }
+            unsafe { chunk_offsets_buf.push_unchecked(curr_offset) }
         }
 
+        let chunk_offsets = PrimitiveArray::new(chunk_offsets_buf.freeze(), Validity::NonNullable);
+
         Self {
             dtype,
             len: curr_offset.try_into().vortex_unwrap(),
-            chunk_offsets: chunk_offsets.freeze(),
+            chunk_offsets,
             chunks,
             stats_set: Default::default(),
         }
@@ -102,8 +106,8 @@ impl ChunkedArray {
     }
 
     #[inline]
-    pub fn chunk_offsets(&self) -> &Buffer<u64> {
-        &self.chunk_offsets
+    pub fn chunk_offsets(&self) -> Buffer<u64> {
+        self.chunk_offsets.buffer()
     }
 
     pub(crate) fn find_chunk_idx(&self, index: usize) -> (usize, usize) {

vortex-array/src/arrays/chunked/compute/filter.rs

@@ -81,10 +81,10 @@ pub(crate) fn chunk_filters(
     let mut chunk_filters = vec![ChunkFilter::None; array.nchunks()];
 
     for (slice_start, slice_end) in slices {
-        let (start_chunk, start_idx) = find_chunk_idx(slice_start, chunk_offsets);
+        let (start_chunk, start_idx) = find_chunk_idx(slice_start, &chunk_offsets);
         // NOTE: we adjust slice end back by one, in case it ends on a chunk boundary, we do not
         // want to index into the unused chunk.
-        let (end_chunk, end_idx) = find_chunk_idx(slice_end - 1, chunk_offsets);
+        let (end_chunk, end_idx) = find_chunk_idx(slice_end - 1, &chunk_offsets);
         // Adjust back to an exclusive range
         let end_idx = end_idx + 1;
 
@@ -143,7 +143,7 @@ fn filter_indices(
     let chunk_offsets = array.chunk_offsets();
 
     for set_index in indices {
-        let (chunk_id, index) = find_chunk_idx(set_index, chunk_offsets);
+        let (chunk_id, index) = find_chunk_idx(set_index, &chunk_offsets);
         if chunk_id != current_chunk_id {
             // Push the chunk we've accumulated.
             if !chunk_indices.is_empty() {

vortex-array/src/arrays/chunked/compute/mask.rs

@@ -47,7 +47,7 @@ fn mask_indices(
     let chunk_offsets = array.chunk_offsets();
 
     for &set_index in indices {
-        let (chunk_id, index) = find_chunk_idx(set_index, chunk_offsets);
+        let (chunk_id, index) = find_chunk_idx(set_index, &chunk_offsets);
         if chunk_id != current_chunk_id {
             let chunk = array.chunk(current_chunk_id);
             let masked_chunk = mask(chunk, &Mask::from_indices(chunk.len(), chunk_indices))?;

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

@@ -27,7 +27,7 @@ impl ArrayVTable<ChunkedVTable> for ChunkedVTable {
     fn array_hash<H: std::hash::Hasher>(array: &ChunkedArray, state: &mut H, precision: Precision) {
         array.dtype.hash(state);
         array.len.hash(state);
-        array.chunk_offsets.array_hash(state, precision);
+        array.chunk_offsets.as_ref().array_hash(state, precision);
         for chunk in &array.chunks {
             chunk.array_hash(state, precision);
         }
@@ -38,7 +38,8 @@ impl ArrayVTable<ChunkedVTable> for ChunkedVTable {
             && array.len == other.len
             && array
                 .chunk_offsets
-                .array_eq(&other.chunk_offsets, precision)
+                .as_ref()
+                .array_eq(other.chunk_offsets.as_ref(), precision)
             && array.chunks.len() == other.chunks.len()
             && array
                 .chunks

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

@@ -7,9 +7,10 @@ use vortex_dtype::{DType, Nullability, PType};
 use vortex_error::{VortexResult, vortex_bail, vortex_err};
 use vortex_vector::{Vector, VectorMut, VectorMutOps};
 
-use crate::arrays::ChunkedArray;
+use crate::arrays::{ChunkedArray, PrimitiveArray};
 use crate::execution::ExecutionCtx;
 use crate::serde::ArrayChildren;
+use crate::validity::Validity;
 use crate::vtable::{NotSupported, VTable};
 use crate::{ArrayOperator, EmptyMetadata, EncodingId, EncodingRef, ToCanonical, vtable};
 
@@ -71,18 +72,19 @@ impl VTable for ChunkedVTable {
         let nchunks = children.len() - 1;
 
         // The first child contains the row offsets of the chunks
-        let chunk_offsets = children
+        let chunk_offsets_array = children
             .get(
                 0,
                 &DType::Primitive(PType::U64, Nullability::NonNullable),
                 // 1 extra offset for the end of the last chunk
                 nchunks + 1,
             )?
-            .to_primitive()
-            .buffer::<u64>();
+            .to_primitive();
+
+        let chunk_offsets_buf = chunk_offsets_array.buffer::<u64>();
 
         // The remaining children contain the actual data of the chunks
-        let chunks = chunk_offsets
+        let chunks = chunk_offsets_buf
             .iter()
             .tuple_windows()
             .enumerate()
@@ -93,9 +95,22 @@ impl VTable for ChunkedVTable {
             })
             .try_collect()?;
 
-        // SAFETY: All chunks are deserialized with the same dtype that was serialized.
-        // Each chunk was validated during deserialization to match the expected dtype.
-        unsafe { Ok(ChunkedArray::new_unchecked(chunks, dtype.clone())) }
+        let chunk_offsets = PrimitiveArray::new(chunk_offsets_buf.clone(), Validity::NonNullable);
+
+        let total_len = chunk_offsets_buf
+            .last()
+            .ok_or_else(|| vortex_err!("chunk_offsets must not be empty"))?;
+        let len = usize::try_from(*total_len)
+            .map_err(|_| vortex_err!("total length {} exceeds usize range", total_len))?;
+
+        // Construct directly using the struct fields to avoid recomputing chunk_offsets
+        Ok(ChunkedArray {
+            dtype: dtype.clone(),
+            len,
+            chunk_offsets,
+            chunks,
+            stats_set: Default::default(),
+        })
     }
 
     fn execute(array: &Self::Array, ctx: &mut dyn ExecutionCtx) -> VortexResult<Vector> {

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

@@ -1,18 +1,15 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
-use crate::arrays::{ChunkedArray, ChunkedVTable, PrimitiveArray};
-use crate::validity::Validity;
+use crate::arrays::{ChunkedArray, ChunkedVTable};
 use crate::vtable::VisitorVTable;
 use crate::{ArrayBufferVisitor, ArrayChildVisitor};
 
 impl VisitorVTable<ChunkedVTable> for ChunkedVTable {
     fn visit_buffers(_array: &ChunkedArray, _visitor: &mut dyn ArrayBufferVisitor) {}
 
     fn visit_children(array: &ChunkedArray, visitor: &mut dyn ArrayChildVisitor) {
-        let chunk_offsets =
-            PrimitiveArray::new(array.chunk_offsets().clone(), Validity::NonNullable);
-        visitor.visit_child("chunk_offsets", chunk_offsets.as_ref());
+        visitor.visit_child("chunk_offsets", array.chunk_offsets.as_ref());
 
         for (idx, chunk) in array.chunks().iter().enumerate() {
             visitor.visit_child(format!("chunks[{idx}]").as_str(), chunk);