fix

Signed-off-by: Joe Isaacs <joe.isaacs@live.co.uk>

Author: joseph-isaacs

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

@@ -8,6 +8,7 @@ use vortex_buffer::buffer;
 
 use crate::IntoArray;
 use crate::accessor::ArrayAccessor;
+use crate::arrays::Chunked;
 use crate::arrays::ChunkedArray;
 use crate::arrays::ListArray;
 use crate::arrays::PrimitiveArray;
@@ -20,6 +21,7 @@ use crate::dtype::Nullability;
 use crate::dtype::PType;
 use crate::dtype::PType::I32;
 use crate::validity::Validity;
+use crate::vtable::VTable;
 
 fn chunked_array() -> ChunkedArray {
     ChunkedArray::try_new(
@@ -194,3 +196,33 @@ pub fn pack_nested_lists() {
     assert_eq!(l1.scalar_at(0).unwrap(), canon_values.scalar_at(0).unwrap());
     assert_eq!(l2.scalar_at(0).unwrap(), canon_values.scalar_at(1).unwrap());
 }
+
+#[test]
+fn with_slots_updates_nchunks_len_and_offsets() {
+    let mut array = chunked_array();
+    let slots = vec![
+        Some(buffer![0u64, 4, 9].into_array()),
+        Some(buffer![10u64, 11, 12, 13].into_array()),
+        Some(buffer![14u64, 15, 16, 17, 18].into_array()),
+    ];
+    let expected_nchunks = slots.len() - 1;
+    let expected_len = array.len();
+
+    <Chunked as VTable>::with_slots(&mut array, slots).unwrap();
+
+    assert_eq!(array.nchunks(), expected_nchunks);
+    assert_eq!(array.len(), expected_len);
+    assert_eq!(array.chunk_offsets(), buffer![0u64, 4, 9]);
+    assert_arrays_eq!(
+        array.chunk(0).clone(),
+        PrimitiveArray::from_iter([10u64, 11, 12, 13])
+    );
+    assert_arrays_eq!(
+        array.chunk(1).clone(),
+        PrimitiveArray::from_iter([14u64, 15, 16, 17, 18])
+    );
+    assert_arrays_eq!(
+        array,
+        PrimitiveArray::from_iter([10u64, 11, 12, 13, 14, 15, 16, 17, 18])
+    );
+}

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

@@ -6,6 +6,7 @@ use std::hash::Hash;
 use itertools::Itertools;
 use vortex_error::VortexResult;
 use vortex_error::vortex_bail;
+use vortex_error::vortex_ensure;
 use vortex_error::vortex_err;
 use vortex_error::vortex_panic;
 use vortex_session::VortexSession;
@@ -19,6 +20,7 @@ use crate::IntoArray;
 use crate::Precision;
 use crate::ToCanonical;
 use crate::arrays::ChunkedArray;
+use crate::arrays::Primitive;
 use crate::arrays::PrimitiveArray;
 use crate::arrays::chunked::compute::kernel::PARENT_KERNELS;
 use crate::arrays::chunked::compute::rules::PARENT_RULES;
@@ -210,6 +212,46 @@ impl VTable for Chunked {
     }
 
     fn with_slots(array: &mut ChunkedArray, slots: Vec<Option<ArrayRef>>) -> VortexResult<()> {
+        vortex_ensure!(!slots.is_empty(), "Chunked array needs at least one slot");
+
+        let chunk_offsets = slots[0]
+            .as_ref()
+            .ok_or_else(|| vortex_err!("Chunked array chunk_offsets slot must be present"))?;
+        let chunk_offsets_dtype = DType::Primitive(PType::U64, Nullability::NonNullable);
+        vortex_ensure!(
+            chunk_offsets.dtype() == &chunk_offsets_dtype,
+            MismatchedTypes: &chunk_offsets_dtype,
+            chunk_offsets.dtype()
+        );
+
+        #[cfg(debug_assertions)]
+        {
+            let chunk_offsets = chunk_offsets.as_opt::<Primitive>().unwrap_or_else(|| {
+                vortex_panic!("Chunked array chunk_offsets slot must be primitive")
+            });
+            debug_assert_eq!(
+                chunk_offsets.len(),
+                slots.len(),
+                "Expected {} chunk offsets, found {}",
+                slots.len(),
+                chunk_offsets.len(),
+            );
+        }
+
+        #[cfg(debug_assertions)]
+        for (idx, chunk_slot) in slots[1..].iter().enumerate() {
+            let chunk = chunk_slot
+                .as_ref()
+                .unwrap_or_else(|| vortex_panic!("Chunked array chunk slot {idx} must be present"));
+            debug_assert!(
+                chunk.dtype() == array.dtype(),
+                "Chunked array chunk slot {} has dtype {:?}, expected {:?}",
+                idx,
+                chunk.dtype(),
+                array.dtype(),
+            );
+        }
+
         array.slots = slots;
         Ok(())
     }