From 3555d87f6690c744186b02bb4734fc532aba2dc8 Mon Sep 17 00:00:00 2001 From: Dan King Date: Tue, 17 Dec 2024 11:25:45 -0500 Subject: [PATCH] chore: use PrimitiveArray::patch in bitpacking take (#1690) Resolves #213 (I think!). --- .../fastlanes/src/bitpacking/compute/take.rs | 38 ++++++------------- .../src/array/primitive/compute/take.rs | 14 +++---- vortex-array/src/array/primitive/mod.rs | 6 ++- vortex-array/src/array/primitive/patch.rs | 5 ++- 4 files changed, 26 insertions(+), 37 deletions(-) diff --git a/encodings/fastlanes/src/bitpacking/compute/take.rs b/encodings/fastlanes/src/bitpacking/compute/take.rs index 12911d2d88..67e6cc044b 100644 --- a/encodings/fastlanes/src/bitpacking/compute/take.rs +++ b/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 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 for BitPackedEncoding { fn take_primitive( array: &BitPackedArray, indices: &PrimitiveArray, -) -> VortexResult> { + taken_validity: Validity, +) -> VortexResult { if indices.is_empty() { - return Ok(vec![]); + return Ok(PrimitiveArray::from_vec(Vec::::new(), taken_validity)); } let offset = array.offset() as usize; @@ -104,29 +105,14 @@ fn take_primitive( } }); - 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::(); - - for (idx, v) in indices.maybe_null_slice::().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)] diff --git a/vortex-array/src/array/primitive/compute/take.rs b/vortex-array/src/array/primitive/compute/take.rs index d2570b40dc..4c2df9a37f 100644 --- a/vortex-array/src/array/primitive/compute/take.rs +++ b/vortex-array/src/array/primitive/compute/take.rs @@ -16,7 +16,7 @@ impl TakeFn 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 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 for PrimitiveEncoding { // In which case, Rust should reuse the same Vec the result. fn take_primitive>( array: &[T], - indices: Vec, + indices: &[I], ) -> Vec { - indices.into_iter().map(|idx| array[idx.as_()]).collect() + indices.iter().map(|idx| array[idx.as_()]).collect() } // We pass a Vec in case we're T == u64. // In which case, Rust should reuse the same Vec the result. unsafe fn take_primitive_unchecked>( array: &[T], - indices: Vec, + indices: &[I], ) -> Vec { 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]); } } diff --git a/vortex-array/src/array/primitive/mod.rs b/vortex-array/src/array/primitive/mod.rs index dfc1a9a323..b8dd2adfc6 100644 --- a/vortex-array/src/array/primitive/mod.rs +++ b/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(values: Vec, validity: Validity) -> Self { match_each_native_ptype!(T::PTYPE, |$P| { PrimitiveArray::new( - ArrowBuffer::from(MutableBuffer::from(unsafe { std::mem::transmute::, Vec<$P>>(values) })).into(), + Buffer::from( + ArrowBuffer::from(unsafe { std::mem::transmute::, Vec<$P>>(values) }) + ), T::PTYPE, validity, ) diff --git a/vortex-array/src/array/primitive/patch.rs b/vortex-array/src/array/primitive/patch.rs index 4620683d42..c37d94a370 100644 --- a/vortex-array/src/array/primitive/patch.rs +++ b/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 { - 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()