fix: array correctly update their validity

joseph-isaacs · joseph-isaacs · commit 3d8c1f9e71b9 · 2026-04-07T11:55:52.000+01:00
Signed-off-by: Joe Isaacs &lt;joe.isaacs@live.co.uk&gt;
diff --git a/encodings/fsst/src/array.rs b/encodings/fsst/src/array.rs
@@ -77,7 +77,7 @@ impl ArrayHash for FSSTData {
     fn array_hash<H: Hasher>(&self, state: &mut H, precision: Precision) {
         self.symbols.array_hash(state, precision);
         self.symbol_lengths.array_hash(state, precision);
-        self.codes_bytes.array_hash(state, precision);
+        self.codes.clone().into_array().array_hash(state, precision);
     }
 }
 
@@ -87,7 +87,11 @@ impl ArrayEq for FSSTData {
             && self
                 .symbol_lengths
                 .array_eq(&other.symbol_lengths, precision)
-            && self.codes_bytes.array_eq(&other.codes_bytes, precision)
+            && self
+                .codes
+                .clone()
+                .into_array()
+                .array_eq(&other.codes.clone().into_array(), precision)
     }
 }
 
@@ -118,7 +122,7 @@ impl VTable for FSST {
         match idx {
             0 => BufferHandle::new_host(array.symbols().clone().into_byte_buffer()),
             1 => BufferHandle::new_host(array.symbol_lengths().clone().into_byte_buffer()),
-            2 => array.codes_bytes.clone(),
+            2 => array.codes.bytes_handle().clone(),
             _ => vortex_panic!("FSSTArray buffer index {idx} out of bounds"),
         }
     }
@@ -136,7 +140,7 @@ impl VTable for FSST {
         Ok(Some(
             FSSTMetadata {
                 uncompressed_lengths_ptype: array.uncompressed_lengths().dtype().as_ptype().into(),
-                codes_offsets_ptype: array.codes().offsets().dtype().as_ptype().into(),
+                codes_offsets_ptype: array.codes.offsets().dtype().as_ptype().into(),
             }
             .encode_to_vec(),
         ))
@@ -315,7 +319,7 @@ pub(crate) const SLOT_NAMES: [&str; NUM_SLOTS] =
 pub struct FSSTData {
     symbols: Buffer<Symbol>,
     symbol_lengths: Buffer<u8>,
-    codes_bytes: BufferHandle,
+    codes: VarBinArray,
 
     /// Memoized compressor used for push-down of compute by compressing the RHS.
     compressor: Arc<LazyLock<Compressor, Box<dyn Fn() -> Compressor + Send>>>,
@@ -326,7 +330,7 @@ impl Debug for FSSTData {
         f.debug_struct("FSSTArray")
             .field("symbols", &self.symbols)
             .field("symbol_lengths", &self.symbol_lengths)
-            .field("codes_bytes", &self.codes_bytes)
+            .field("codes", &self.codes)
             .field("uncompressed_lengths", &"<outer slot>")
             .finish()
     }
@@ -365,8 +369,7 @@ impl FSST {
                 codes.len(),
             ),
         ];
-        let codes_bytes = codes.bytes_handle().clone();
-        let data = FSSTData::new_from_parts(symbols, symbol_lengths, codes_bytes);
+        let data = FSSTData::try_new(symbols, symbol_lengths, codes, &dtype)?;
         Ok(unsafe {
             Array::from_parts_unchecked(ArrayParts::new(FSST, dtype, len, data).with_slots(slots))
         })
@@ -390,8 +393,7 @@ impl FSST {
                 codes.len(),
             ),
         ];
-        let codes_bytes = codes.bytes_handle().clone();
-        let data = FSSTData::new_from_parts(symbols, symbol_lengths, codes_bytes);
+        let data = unsafe { FSSTData::new_unchecked(symbols, symbol_lengths, codes) };
         unsafe {
             Array::from_parts_unchecked(ArrayParts::new(FSST, dtype, len, data).with_slots(slots))
         }
@@ -426,8 +428,8 @@ impl FSSTData {
         codes: VarBinArray,
         _dtype: &DType,
     ) -> VortexResult<Self> {
-        let codes_bytes = codes.bytes_handle().clone();
-        Ok(Self::new_from_parts(symbols, symbol_lengths, codes_bytes))
+        // SAFETY: all components validated above
+        unsafe { Ok(Self::new_unchecked(symbols, symbol_lengths, codes)) }
     }
 
     pub fn validate(
@@ -436,11 +438,10 @@ impl FSSTData {
         len: usize,
         slots: &[Option<ArrayRef>],
     ) -> VortexResult<()> {
-        let codes = self.codes_from_slots(slots, dtype);
         Self::validate_parts(
             &self.symbols,
             &self.symbol_lengths,
-            &codes,
+            &self.codes,
             uncompressed_lengths_from_slots(slots),
             dtype,
             len,
@@ -491,26 +492,35 @@ impl FSSTData {
         Ok(())
     }
 
-    fn new_from_parts(
+    pub(crate) unsafe fn new_unchecked(
         symbols: Buffer<Symbol>,
         symbol_lengths: Buffer<u8>,
-        codes_bytes: BufferHandle,
+        codes: VarBinArray,
     ) -> Self {
         let symbols2 = symbols.clone();
         let symbol_lengths2 = symbol_lengths.clone();
         let compressor = Arc::new(LazyLock::new(Box::new(move || {
             Compressor::rebuild_from(symbols2.as_slice(), symbol_lengths2.as_slice())
         })
             as Box<dyn Fn() -> Compressor + Send>));
-
         Self {
             symbols,
             symbol_lengths,
-            codes_bytes,
+            codes,
             compressor,
         }
     }
 
+    /// Returns the number of elements in the array.
+    pub fn len(&self) -> usize {
+        self.codes.len()
+    }
+
+    /// Returns `true` if the array contains no elements.
+    pub fn is_empty(&self) -> bool {
+        self.codes.len() == 0
+    }
+
     /// Access the symbol table array
     pub fn symbols(&self) -> &Buffer<Symbol> {
         &self.symbols
@@ -521,9 +531,15 @@ impl FSSTData {
         &self.symbol_lengths
     }
 
-    /// Access the codes bytes buffer handle.
-    pub fn codes_bytes(&self) -> &BufferHandle {
-        &self.codes_bytes
+    /// Access the codes array
+    pub fn codes(&self) -> &VarBinArray {
+        &self.codes
+    }
+
+    /// Get the DType of the codes array
+    #[inline]
+    pub fn codes_dtype(&self) -> &DType {
+        self.codes.dtype()
     }
 
     /// Build a [`Decompressor`][fsst::Decompressor] that can be used to decompress values from
@@ -536,21 +552,6 @@ impl FSSTData {
     pub fn compressor(&self) -> &Compressor {
         self.compressor.as_ref()
     }
-
-    fn codes_from_slots(&self, slots: &[Option<ArrayRef>], dtype: &DType) -> VarBinArray {
-        let offsets = slots[CODES_OFFSETS_SLOT]
-            .clone()
-            .vortex_expect("FSSTArray codes_offsets slot must be present");
-        let validity = child_to_validity(&slots[CODES_VALIDITY_SLOT], dtype.nullability());
-        unsafe {
-            VarBinArray::new_unchecked_from_handle(
-                offsets,
-                self.codes_bytes.clone(),
-                DType::Binary(dtype.nullability()),
-                validity,
-            )
-        }
-    }
 }
 
 fn uncompressed_lengths_from_slots(slots: &[Option<ArrayRef>]) -> &ArrayRef {
@@ -567,32 +568,6 @@ pub trait FSSTArrayExt: TypedArrayRef<FSST> {
     fn uncompressed_lengths_dtype(&self) -> &DType {
         self.uncompressed_lengths().dtype()
     }
-
-    /// Reconstruct the codes VarBinArray on demand from stored bytes and slot components.
-    fn codes(&self) -> VarBinArray {
-        let slots = self.as_ref().slots();
-        let offsets = slots[CODES_OFFSETS_SLOT]
-            .clone()
-            .vortex_expect("FSSTArray codes_offsets slot must be present");
-        let validity = child_to_validity(
-            &slots[CODES_VALIDITY_SLOT],
-            self.as_ref().dtype().nullability(),
-        );
-        // SAFETY: components were validated at FSSTArray construction time.
-        unsafe {
-            VarBinArray::new_unchecked_from_handle(
-                offsets,
-                self.codes_bytes().clone(),
-                DType::Binary(self.as_ref().dtype().nullability()),
-                validity,
-            )
-        }
-    }
-
-    /// Get the DType of the codes array.
-    fn codes_dtype(&self) -> DType {
-        DType::Binary(self.as_ref().dtype().nullability())
-    }
 }
 
 impl<T: TypedArrayRef<FSST>> FSSTArrayExt for T {}
@@ -666,7 +641,7 @@ mod test {
             &compressor,
         );
 
-        let compressed_codes = fsst_array.codes();
+        let compressed_codes = fsst_array.codes().clone();
 
         // There were two buffers:
         // 1. The 8 byte symbols
diff --git a/encodings/fsst/src/compute/cast.rs b/encodings/fsst/src/compute/cast.rs
@@ -20,8 +20,9 @@ impl CastReduce for FSST {
             // Cast codes array to handle nullability
             let new_codes = array
                 .codes()
+                .clone()
                 .into_array()
-                .cast(array.codes_dtype().with_nullability(dtype.nullability()))?;
+                .cast(array.codes().dtype().with_nullability(dtype.nullability()))?;
 
             Ok(Some(
                 FSST::try_new(
diff --git a/encodings/fsst/src/compute/compare.rs b/encodings/fsst/src/compute/compare.rs
@@ -115,6 +115,7 @@ fn compare_fsst_constant(
 
     let rhs = ConstantArray::new(encoded_scalar, left.len());
     left.codes()
+        .clone()
         .into_array()
         .binary(rhs.into_array(), Operator::from(operator))
         .map(Some)
diff --git a/encodings/fsst/src/compute/filter.rs b/encodings/fsst/src/compute/filter.rs
@@ -21,7 +21,7 @@ impl FilterKernel for FSST {
         ctx: &mut ExecutionCtx,
     ) -> VortexResult<Option<ArrayRef>> {
         // Directly invoke VarBin's FilterKernel on the codes child.
-        let codes = array.codes();
+        let codes = array.codes().clone();
         let codes = codes.as_view();
         let filtered_codes_ref = <VarBin as FilterKernel>::filter(codes, mask, ctx)?
             .vortex_expect("VarBin filter kernel always returns Some");
diff --git a/encodings/fsst/src/compute/like.rs b/encodings/fsst/src/compute/like.rs
@@ -16,7 +16,6 @@ use vortex_array::scalar_fn::fns::like::LikeOptions;
 use vortex_error::VortexResult;
 
 use crate::FSST;
-use crate::array::FSSTArrayExt;
 use crate::dfa::FsstMatcher;
 use crate::dfa::dfa_scan_to_bitbuf;
 
diff --git a/encodings/fsst/src/ops.rs b/encodings/fsst/src/ops.rs
@@ -11,7 +11,6 @@ use vortex_error::VortexExpect;
 use vortex_error::VortexResult;
 
 use crate::FSST;
-use crate::array::FSSTArrayExt;
 
 impl OperationsVTable<FSST> for FSST {
     fn scalar_at(
diff --git a/vortex-array/src/array/vtable/validity.rs b/vortex-array/src/array/vtable/validity.rs
@@ -17,28 +17,6 @@ pub trait ValidityVTable<V: VTable> {
     fn validity(array: ArrayView<'_, V>) -> VortexResult<Validity>;
 }
 
-/// An implementation of the [`ValidityVTable`] for arrays that hold an unsliced validity
-/// and a slice into it.
-pub struct ValidityVTableFromValiditySliceHelper;
-
-pub trait ValiditySliceHelper {
-    fn unsliced_validity_and_slice(&self) -> (&Validity, usize, usize);
-
-    fn sliced_validity(&self) -> VortexResult<Validity> {
-        let (unsliced_validity, start, stop) = self.unsliced_validity_and_slice();
-        unsliced_validity.slice(start..stop)
-    }
-}
-
-impl<V: VTable> ValidityVTable<V> for ValidityVTableFromValiditySliceHelper
-where
-    V::ArrayData: ValiditySliceHelper,
-{
-    fn validity(array: ArrayView<'_, V>) -> VortexResult<Validity> {
-        array.data().sliced_validity()
-    }
-}
-
 /// An implementation of the [`ValidityVTable`] for arrays that delegate validity entirely
 /// to a child array.
 pub struct ValidityVTableFromChild;
