feat: Support patching bool arrays, patch primitive array validity and use patching when canonicalizing sparse arrays (#1218)

Co-authored-by: Will Manning <[email protected]>

Author: robert3005

encodings/alp/src/alp/compress.rs

@@ -83,9 +83,11 @@ fn patch_decoded(array: PrimitiveArray, patches: &Array) -> VortexResult<Primiti
         Sparse::ID => {
             match_each_alp_float_ptype!(array.ptype(), |$T| {
                 let typed_patches = SparseArray::try_from(patches).unwrap();
+                let primitive_values = typed_patches.values().into_primitive()?;
                 array.patch(
                     &typed_patches.resolved_indices(),
-                    typed_patches.values().into_primitive()?.maybe_null_slice::<$T>())
+                    primitive_values.maybe_null_slice::<$T>(),
+                    primitive_values.validity())
             })
         }
         _ => vortex_bail!(

encodings/fastlanes/src/bitpacking/compress.rs

@@ -194,9 +194,11 @@ fn patch_unpacked(array: PrimitiveArray, patches: &Array) -> VortexResult<Primit
         Sparse::ID => {
             match_each_integer_ptype!(array.ptype(), |$T| {
                 let typed_patches = SparseArray::try_from(patches).unwrap();
+                let primitive_values = typed_patches.values().into_primitive()?;
                 array.patch(
                     &typed_patches.resolved_indices(),
-                    typed_patches.values().into_primitive()?.maybe_null_slice::<$T>())
+                    primitive_values.maybe_null_slice::<$T>(),
+                    primitive_values.validity())
             })
         }
         _ => vortex_bail!(

vortex-array/src/array/bool/mod.rs

@@ -1,12 +1,13 @@
 use std::fmt::{Debug, Display};
 
 use arrow_buffer::bit_iterator::{BitIndexIterator, BitSliceIterator};
-use arrow_buffer::BooleanBuffer;
+use arrow_buffer::{BooleanBuffer, BooleanBufferBuilder, MutableBuffer};
 use itertools::Itertools;
+use num_traits::AsPrimitive;
 use serde::{Deserialize, Serialize};
 use vortex_buffer::Buffer;
 use vortex_dtype::DType;
-use vortex_error::{VortexExpect as _, VortexResult};
+use vortex_error::{vortex_bail, VortexExpect as _, VortexResult};
 
 use crate::array::visitor::{AcceptArrayVisitor, ArrayVisitor};
 use crate::encoding::ids;
@@ -34,12 +35,21 @@ impl Display for BoolMetadata {
 }
 
 impl BoolArray {
+    /// Access internal array buffer
     pub fn buffer(&self) -> &Buffer {
         self.as_ref()
             .buffer()
             .vortex_expect("Missing buffer in BoolArray")
     }
 
+    /// Convert array into its internal buffer
+    pub fn into_buffer(self) -> Buffer {
+        self.into_array()
+            .into_buffer()
+            .vortex_expect("BoolArray must have a buffer")
+    }
+
+    /// Get array values as an arrow [BooleanBuffer]
     pub fn boolean_buffer(&self) -> BooleanBuffer {
         BooleanBuffer::new(
             self.buffer().clone().into_arrow(),
@@ -48,6 +58,33 @@ impl BoolArray {
         )
     }
 
+    /// Get a mutable version of this array.
+    ///
+    /// If the caller holds the only reference to the underlying buffer the underlying buffer is returned
+    /// otherwise a copy is created.
+    ///
+    /// The second value of the tuple is a bit_offset of first value in first byte of the returned builder
+    pub fn into_boolean_builder(self) -> (BooleanBufferBuilder, usize) {
+        let first_byte_bit_offset = self.metadata().first_byte_bit_offset as usize;
+        let len = self.len();
+        let arrow_buffer = self.into_buffer().into_arrow();
+        let mutable_buf = if arrow_buffer.ptr_offset() == 0 {
+            arrow_buffer.into_mutable().unwrap_or_else(|b| {
+                let mut buf = MutableBuffer::with_capacity(b.len());
+                buf.extend_from_slice(b.as_slice());
+                buf
+            })
+        } else {
+            let mut buf = MutableBuffer::with_capacity(arrow_buffer.len());
+            buf.extend_from_slice(arrow_buffer.as_slice());
+            buf
+        };
+        (
+            BooleanBufferBuilder::new_from_buffer(mutable_buf, len + first_byte_bit_offset),
+            first_byte_bit_offset,
+        )
+    }
+
     pub fn validity(&self) -> Validity {
         self.metadata().validity.to_validity(|| {
             self.as_ref()
@@ -85,6 +122,34 @@ impl BoolArray {
         let buffer = BooleanBuffer::from(bools);
         Self::try_new(buffer, validity).vortex_expect("Failed to create BoolArray from vec")
     }
+
+    pub fn patch<P: AsPrimitive<usize>>(
+        self,
+        positions: &[P],
+        values: BoolArray,
+    ) -> VortexResult<Self> {
+        if positions.len() != values.len() {
+            vortex_bail!(
+                "Positions and values passed to patch had different lengths {} and {}",
+                positions.len(),
+                values.len()
+            );
+        }
+        if let Some(last_pos) = positions.last() {
+            if last_pos.as_() >= self.len() {
+                vortex_bail!(OutOfBounds: last_pos.as_(), 0, self.len())
+            }
+        }
+
+        let len = self.len();
+        let result_validity = self.validity().patch(len, positions, values.validity())?;
+        let (mut own_values, bit_offset) = self.into_boolean_builder();
+        for (idx, value) in positions.iter().zip_eq(values.boolean_buffer().iter()) {
+            own_values.set_bit(idx.as_() + bit_offset, value);
+        }
+
+        Self::try_new(own_values.finish().slice(bit_offset, len), result_validity)
+    }
 }
 
 impl ArrayTrait for BoolArray {}
@@ -165,13 +230,15 @@ impl AcceptArrayVisitor for BoolArray {
 
 #[cfg(test)]
 mod tests {
+    use arrow_buffer::BooleanBuffer;
     use itertools::Itertools;
 
     use crate::array::BoolArray;
+    use crate::compute::slice;
     use crate::compute::unary::scalar_at;
     use crate::validity::Validity;
     use crate::variants::BoolArrayTrait;
-    use crate::IntoArray;
+    use crate::{IntoArray, IntoArrayVariant};
 
     #[test]
     fn bool_array() {
@@ -238,4 +305,31 @@ mod tests {
         assert_eq!(0, arr.maybe_null_indices_iter().collect_vec().len());
         assert_eq!(0, arr.maybe_null_slices_iter().collect_vec().len());
     }
+
+    #[test]
+    fn patch_sliced_bools() {
+        let arr = BoolArray::from(BooleanBuffer::new_set(12));
+        let sliced = slice(arr, 4, 12).unwrap();
+        let (values, offset) = sliced.into_bool().unwrap().into_boolean_builder();
+        assert_eq!(offset, 4);
+        assert_eq!(values.as_slice(), &[255, 15]);
+    }
+
+    #[test]
+    fn patch_sliced_bools_offset() {
+        let arr = BoolArray::from(BooleanBuffer::new_set(15));
+        let sliced = slice(arr, 4, 15).unwrap();
+        let (values, offset) = sliced.into_bool().unwrap().into_boolean_builder();
+        assert_eq!(offset, 4);
+        assert_eq!(values.as_slice(), &[255, 127]);
+    }
+
+    #[test]
+    fn patch_sliced_bools_even() {
+        let arr = BoolArray::from(BooleanBuffer::new_set(31));
+        let sliced = slice(arr, 8, 24).unwrap();
+        let (values, offset) = sliced.into_bool().unwrap().into_boolean_builder();
+        assert_eq!(offset, 0);
+        assert_eq!(values.as_slice(), &[255, 255]);
+    }
 }

vortex-array/src/array/primitive/mod.rs

@@ -162,19 +162,34 @@ impl PrimitiveArray {
         self,
         positions: &[P],
         values: &[T],
+        values_validity: Validity,
     ) -> VortexResult<Self> {
+        if positions.len() != values.len() {
+            vortex_bail!(
+                "Positions and values passed to patch had different lengths {} and {}",
+                positions.len(),
+                values.len()
+            );
+        }
+        if let Some(last_pos) = positions.last() {
+            if last_pos.as_() >= self.len() {
+                vortex_bail!(OutOfBounds: last_pos.as_(), 0, self.len())
+            }
+        }
+
         if self.ptype() != T::PTYPE {
             vortex_bail!(MismatchedTypes: self.dtype(), T::PTYPE)
         }
 
-        let validity = self.validity();
-
-        let mut own_values = self.into_maybe_null_slice();
-        // TODO(robert): Also patch validity
-        for (idx, value) in positions.iter().zip_eq(values.iter()) {
-            own_values[(*idx).as_()] = *value;
+        let result_validity = self
+            .validity()
+            .patch(self.len(), positions, values_validity)?;
+        let mut own_values = self.into_maybe_null_slice::<T>();
+        for (idx, value) in positions.iter().zip_eq(values) {
+            own_values[idx.as_()] = *value;
         }
-        Ok(Self::from_vec(own_values, validity))
+
+        Ok(Self::from_vec(own_values, result_validity))
     }
 
     pub fn into_buffer(self) -> Buffer {
@@ -461,6 +476,8 @@ mod tests {
     use vortex_scalar::Scalar;
 
     use super::*;
+    use crate::compute::slice;
+    use crate::IntoArrayVariant;
 
     #[test]
     fn batched_iter() {
@@ -546,4 +563,17 @@ mod tests {
             assert_eq!(o.unwrap(), 3);
         }
     }
+
+    #[test]
+    fn patch_sliced() {
+        let input = PrimitiveArray::from_vec(vec![2u32; 10], Validity::AllValid);
+        let sliced = slice(input, 2, 8).unwrap();
+        assert_eq!(
+            sliced
+                .into_primitive()
+                .unwrap()
+                .into_maybe_null_slice::<u32>(),
+            vec![2u32; 6]
+        );
+    }
 }

vortex-array/src/array/sparse/canonical.rs

@@ -1,5 +1,5 @@
-use arrow_buffer::BooleanBufferBuilder;
-use vortex_dtype::{match_each_native_ptype, DType, NativePType};
+use arrow_buffer::{ArrowNativeType, BooleanBuffer};
+use vortex_dtype::{match_each_native_ptype, DType, NativePType, Nullability};
 use vortex_error::{VortexError, VortexResult};
 use vortex_scalar::ScalarValue;
 
@@ -12,93 +12,82 @@ use crate::{ArrayDType, Canonical, IntoArrayVariant, IntoCanonical};
 
 impl IntoCanonical for SparseArray {
     fn into_canonical(self) -> VortexResult<Canonical> {
-        // Resolve our indices into a vector of usize, applying the offset
+        // Resolve our indices into a vector of usize applying the offset
         let indices = self.resolved_indices();
 
         if matches!(self.dtype(), DType::Bool(_)) {
-            canonicalize_sparse_bools(
-                self.values().into_bool()?,
-                &indices,
-                self.len(),
-                self.fill_value(),
-            )
+            let values = self.values().into_bool()?;
+            canonicalize_sparse_bools(values, &indices, self.len(), self.fill_value())
         } else {
             let values = self.values().into_primitive()?;
             match_each_native_ptype!(values.ptype(), |$P| {
-                canonicalize_sparse_primitives(
-                    values.maybe_null_slice::<$P>(),
-                    values.validity(),
+                canonicalize_sparse_primitives::<$P>(
+                    values,
                     &indices,
                     self.len(),
-                    self.fill_value()
+                    self.fill_value(),
                 )
             })
         }
     }
 }
 
 fn canonicalize_sparse_bools(
-    values_array: BoolArray,
+    values: BoolArray,
     indices: &[usize],
     len: usize,
     fill_value: &ScalarValue,
 ) -> VortexResult<Canonical> {
-    let fill_bool = if fill_value.is_null() {
-        false
+    let (fill_bool, validity) = if fill_value.is_null() {
+        (false, Validity::AllInvalid)
     } else {
-        fill_value.try_into()?
+        (
+            fill_value.try_into()?,
+            if values.dtype().nullability() == Nullability::NonNullable {
+                Validity::NonNullable
+            } else {
+                Validity::AllValid
+            },
+        )
     };
 
-    let values_validity = values_array.validity();
-    let values = values_array.boolean_buffer();
-
-    // pre-fill both values and validity based on fill_value
-    // this optimizes performance for the common case where indices.len() < len / 2
-    let mut flat_bools = BooleanBufferBuilder::new(len);
-    flat_bools.append_n(len, fill_bool);
-    let mut validity_buffer = BooleanBufferBuilder::new(len);
-    validity_buffer.append_n(len, !fill_value.is_null());
-
-    // patch in the actual values and validity
-    for (i, idx) in indices.iter().enumerate() {
-        flat_bools.set_bit(*idx, values.value(i));
-        validity_buffer.set_bit(*idx, values_validity.is_valid(i));
-    }
-
-    BoolArray::try_new(
-        flat_bools.finish(),
-        Validity::from(validity_buffer.finish()),
-    )
-    .map(Canonical::Bool)
+    let bools = BoolArray::try_new(
+        if fill_bool {
+            BooleanBuffer::new_set(len)
+        } else {
+            BooleanBuffer::new_unset(len)
+        },
+        validity,
+    )?;
+    let patched = bools.patch(indices, values)?;
+    Ok(Canonical::Bool(patched))
 }
 
 fn canonicalize_sparse_primitives<
-    T: NativePType + for<'a> TryFrom<&'a ScalarValue, Error = VortexError>,
+    T: NativePType + for<'a> TryFrom<&'a ScalarValue, Error = VortexError> + ArrowNativeType,
 >(
-    values: &[T],
-    values_validity: Validity,
+    values: PrimitiveArray,
     indices: &[usize],
     len: usize,
     fill_value: &ScalarValue,
 ) -> VortexResult<Canonical> {
-    let primitive_fill = if fill_value.is_null() {
-        T::default()
+    let values_validity = values.validity();
+    let (primitive_fill, validity) = if fill_value.is_null() {
+        (T::default(), Validity::AllInvalid)
     } else {
-        fill_value.try_into()?
+        (
+            fill_value.try_into()?,
+            if values_validity == Validity::NonNullable {
+                Validity::NonNullable
+            } else {
+                Validity::AllValid
+            },
+        )
     };
-    let mut result = vec![primitive_fill; len];
-    let mut validity = BooleanBufferBuilder::new(len);
-    validity.append_n(len, !fill_value.is_null());
-
-    for (i, idx) in indices.iter().enumerate() {
-        result[*idx] = values[i];
-        validity.set_bit(*idx, values_validity.is_valid(i));
-    }
 
-    Ok(Canonical::Primitive(PrimitiveArray::from_vec(
-        result,
-        Validity::from(validity.finish()),
-    )))
+    let parray = PrimitiveArray::from_vec(vec![primitive_fill; len], validity);
+    let patched = parray.patch(indices, values.maybe_null_slice::<T>(), values_validity)?;
+    Ok(Canonical::Primitive(patched))
 }
 
 #[cfg(test)]