chore: use PrimitiveArray::patch in bitpacking take (#1690)

Resolves #213 (I think!).

encodings/fastlanes/src/bitpacking/compute/take.rs

@@ -1,13 +1,13 @@
 use fastlanes::BitPacking;
 use vortex_array::array::PrimitiveArray;
-use vortex_array::compute::{take, try_cast, TakeFn};
+use vortex_array::compute::{take, TakeFn};
+use vortex_array::validity::Validity;
 use vortex_array::variants::PrimitiveArrayTrait;
 use vortex_array::{
     ArrayDType, ArrayData, ArrayLen, IntoArrayData, IntoArrayVariant, IntoCanonical, ToArrayData,
 };
 use vortex_dtype::{
-    match_each_integer_ptype, match_each_unsigned_integer_ptype, DType, NativePType, Nullability,
-    PType,
+    match_each_integer_ptype, match_each_unsigned_integer_ptype, NativePType, PType,
 };
 use vortex_error::{VortexExpect as _, VortexResult};
 
@@ -33,7 +33,7 @@ impl TakeFn<BitPackedArray> for BitPackedEncoding {
         let indices = indices.clone().into_primitive()?;
         let taken = match_each_unsigned_integer_ptype!(ptype, |$T| {
             match_each_integer_ptype!(indices.ptype(), |$I| {
-                PrimitiveArray::from_vec(take_primitive::<$T, $I>(array, &indices)?, taken_validity)
+                take_primitive::<$T, $I>(array, &indices, taken_validity)?
             })
         });
         Ok(taken.reinterpret_cast(ptype).into_array())
@@ -43,9 +43,10 @@ impl TakeFn<BitPackedArray> for BitPackedEncoding {
 fn take_primitive<T: NativePType + BitPacking, I: NativePType>(
     array: &BitPackedArray,
     indices: &PrimitiveArray,
-) -> VortexResult<Vec<T>> {
+    taken_validity: Validity,
+) -> VortexResult<PrimitiveArray> {
     if indices.is_empty() {
-        return Ok(vec![]);
+        return Ok(PrimitiveArray::from_vec(Vec::<T>::new(), taken_validity));
     }
 
     let offset = array.offset() as usize;
@@ -104,29 +105,14 @@ fn take_primitive<T: NativePType + BitPacking, I: NativePType>(
         }
     });
 
-    if let Some(patches) = array
-        .patches()
-        .map(|p| p.take(&indices.to_array()))
-        .transpose()?
-        .flatten()
-    {
-        let indices = try_cast(
-            patches.indices(),
-            &DType::Primitive(PType::U64, Nullability::NonNullable),
-        )?
-        .into_primitive()?;
-
-        // TODO(ngates): can patch values themselves have nulls, or do we ensure they're in our
-        //  validity bitmap?
-        let values = patches.values().clone().into_primitive()?;
-        let values_slice = values.maybe_null_slice::<T>();
-
-        for (idx, v) in indices.maybe_null_slice::<u64>().iter().zip(values_slice) {
-            output[*idx as usize] = *v;
+    let unpatched_taken = PrimitiveArray::from_vec(output, taken_validity);
+    if let Some(patches) = array.patches() {
+        if let Some(patches) = patches.take(&indices.to_array())? {
+            return unpatched_taken.patch(patches);
         }
     }
 
-    Ok(output)
+    Ok(unpatched_taken)
 }
 
 #[cfg(test)]

vortex-array/src/array/primitive/compute/take.rs

@@ -16,7 +16,7 @@ impl TakeFn<PrimitiveArray> for PrimitiveEncoding {
 
         match_each_native_ptype!(array.ptype(), |$T| {
             match_each_integer_ptype!(indices.ptype(), |$I| {
-                let values = take_primitive(array.maybe_null_slice::<$T>(), indices.into_maybe_null_slice::<$I>());
+                let values = take_primitive(array.maybe_null_slice::<$T>(), indices.maybe_null_slice::<$I>());
                 Ok(PrimitiveArray::from_vec(values, validity).into_array())
             })
         })
@@ -32,7 +32,7 @@ impl TakeFn<PrimitiveArray> for PrimitiveEncoding {
 
         match_each_native_ptype!(array.ptype(), |$T| {
             match_each_integer_ptype!(indices.ptype(), |$I| {
-                let values = take_primitive_unchecked(array.maybe_null_slice::<$T>(), indices.into_maybe_null_slice::<$I>());
+                let values = take_primitive_unchecked(array.maybe_null_slice::<$T>(), indices.maybe_null_slice::<$I>());
                 Ok(PrimitiveArray::from_vec(values, validity).into_array())
             })
         })
@@ -43,19 +43,19 @@ impl TakeFn<PrimitiveArray> for PrimitiveEncoding {
 // In which case, Rust should reuse the same Vec<u64> the result.
 fn take_primitive<T: NativePType, I: NativePType + AsPrimitive<usize>>(
     array: &[T],
-    indices: Vec<I>,
+    indices: &[I],
 ) -> Vec<T> {
-    indices.into_iter().map(|idx| array[idx.as_()]).collect()
+    indices.iter().map(|idx| array[idx.as_()]).collect()
 }
 
 // We pass a Vec<I> in case we're T == u64.
 // In which case, Rust should reuse the same Vec<u64> the result.
 unsafe fn take_primitive_unchecked<T: NativePType, I: NativePType + AsPrimitive<usize>>(
     array: &[T],
-    indices: Vec<I>,
+    indices: &[I],
 ) -> Vec<T> {
     indices
-        .into_iter()
+        .iter()
         .map(|idx| unsafe { *array.get_unchecked(idx.as_()) })
         .collect()
 }
@@ -67,7 +67,7 @@ mod test {
     #[test]
     fn test_take() {
         let a = vec![1i32, 2, 3, 4, 5];
-        let result = take_primitive(&a, vec![0, 0, 4, 2]);
+        let result = take_primitive(&a, &[0, 0, 4, 2]);
         assert_eq!(result, vec![1i32, 1, 5, 3]);
     }
 }

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

@@ -3,7 +3,7 @@ use std::ptr;
 use std::sync::Arc;
 mod accessor;
 
-use arrow_buffer::{ArrowNativeType, BooleanBufferBuilder, Buffer as ArrowBuffer, MutableBuffer};
+use arrow_buffer::{ArrowNativeType, BooleanBufferBuilder, Buffer as ArrowBuffer};
 use bytes::Bytes;
 use itertools::Itertools;
 use serde::{Deserialize, Serialize};
@@ -71,7 +71,9 @@ impl PrimitiveArray {
     pub fn from_vec<T: NativePType>(values: Vec<T>, validity: Validity) -> Self {
         match_each_native_ptype!(T::PTYPE, |$P| {
             PrimitiveArray::new(
-                ArrowBuffer::from(MutableBuffer::from(unsafe { std::mem::transmute::<Vec<T>, Vec<$P>>(values) })).into(),
+                Buffer::from(
+                    ArrowBuffer::from(unsafe { std::mem::transmute::<Vec<T>, Vec<$P>>(values) })
+                ),
                 T::PTYPE,
                 validity,
             )

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

@@ -10,8 +10,9 @@ use crate::{ArrayLen, IntoArrayVariant};
 impl PrimitiveArray {
     #[allow(clippy::cognitive_complexity)]
     pub fn patch(self, patches: Patches) -> VortexResult<Self> {
-        let indices = patches.indices().clone().into_primitive()?;
-        let values = patches.values().clone().into_primitive()?;
+        let (_, indices, values) = patches.into_parts();
+        let indices = indices.into_primitive()?;
+        let values = values.into_primitive()?;
 
         let patched_validity =
             self.validity()