diff --git a/encodings/alp/src/alp/compress.rs b/encodings/alp/src/alp/compress.rs index 521d2cf01f2..53459860af7 100644 --- a/encodings/alp/src/alp/compress.rs +++ b/encodings/alp/src/alp/compress.rs @@ -69,7 +69,9 @@ pub fn decompress(array: ALPArray) -> VortexResult { let ptype = array.dtype().try_into()?; let decoded = match_each_alp_float_ptype!(ptype, |$T| { PrimitiveArray::from_vec( - <$T>::decode_vec(encoded.into_maybe_null_slice(), array.exponents()), + // If encoded is uniquely owned (e.g. it was just decompressed by into_primitive), Rust + // will re-use the allocation + <$T>::decode_vec(encoded.into_maybe_null_vec(), array.exponents()), validity, ) }); diff --git a/encodings/alp/src/alp_rd/array.rs b/encodings/alp/src/alp_rd/array.rs index bd6adfa3d4d..0051bae5955 100644 --- a/encodings/alp/src/alp_rd/array.rs +++ b/encodings/alp/src/alp_rd/array.rs @@ -204,7 +204,7 @@ impl IntoCanonical for ALPRDArray { left_parts.maybe_null_slice::(), left_parts_dict, self.metadata().right_bit_width, - right_parts.maybe_null_slice::(), + right_parts.into_maybe_null_vec::(), self.left_parts_patches(), )?, self.logical_validity().into_validity(), @@ -215,7 +215,7 @@ impl IntoCanonical for ALPRDArray { left_parts.maybe_null_slice::(), left_parts_dict, self.metadata().right_bit_width, - right_parts.maybe_null_slice::(), + right_parts.into_maybe_null_vec::(), self.left_parts_patches(), )?, self.logical_validity().into_validity(), diff --git a/encodings/alp/src/alp_rd/mod.rs b/encodings/alp/src/alp_rd/mod.rs index f7f8c59072f..51ee8595a32 100644 --- a/encodings/alp/src/alp_rd/mod.rs +++ b/encodings/alp/src/alp_rd/mod.rs @@ -17,7 +17,7 @@ use vortex_array::aliases::hash_map::HashMap; use vortex_array::array::PrimitiveArray; use vortex_array::{ArrayDType, IntoArrayData}; use vortex_dtype::{DType, NativePType}; -use vortex_error::{vortex_bail, VortexExpect, VortexResult, VortexUnwrap}; +use vortex_error::{vortex_bail, vortex_err, VortexExpect, VortexResult, VortexUnwrap}; use vortex_fastlanes::bitpack_encode_unchecked; use crate::match_each_alp_float_ptype; @@ -254,6 +254,10 @@ impl RDEncoder { /// Decode a vector of ALP-RD encoded values back into their original floating point format. /// +/// `left_parts_patches` is taken as an owned vector because we always need an allocation of that +/// size for the result and Rust will re-use the allocation (which has compatible length and +/// bit-width).sr +/// /// # Panics /// /// The function panics if the provided `left_parts` and `right_parts` differ in length. @@ -261,7 +265,7 @@ pub fn alp_rd_decode( left_parts: &[u16], left_parts_dict: &[u16], right_bit_width: u8, - right_parts: &[T::UINT], + right_parts: Vec, left_parts_patches: Option, ) -> VortexResult> { if left_parts.len() != right_parts.len() { @@ -281,14 +285,16 @@ pub fn alp_rd_decode( let patched_left_parts = match left_parts_patches { Some(patches) => PrimitiveArray::from(left_parts_decoded) .patch(patches)? - .into_maybe_null_slice::(), + .try_into_maybe_null_vec::() + .map_err(|_| vortex_err!("could not zero copy patched left_parts_decoded")) + .vortex_expect("the buffer backing this PrimitiveArray is uniquely derived from a Vec uniquely owned by this method"), None => left_parts_decoded, }; // recombine the left-and-right parts, adjusting by the right_bit_width. Ok(patched_left_parts .into_iter() - .zip(right_parts.iter().copied()) + .zip_eq(right_parts) .map(|(left, right)| { let left = ::from_u16(left); T::from_bits((left << (right_bit_width as usize)) | right) diff --git a/encodings/datetime-parts/src/compress.rs b/encodings/datetime-parts/src/compress.rs index b088f616a72..fdb3c108fef 100644 --- a/encodings/datetime-parts/src/compress.rs +++ b/encodings/datetime-parts/src/compress.rs @@ -35,15 +35,28 @@ pub fn split_temporal(array: TemporalArray) -> VortexResult { }; let length = timestamps.len(); - let mut days = Vec::with_capacity(length); let mut seconds = Vec::with_capacity(length); let mut subsecond = Vec::with_capacity(length); - for &t in timestamps.maybe_null_slice::().iter() { - days.push(t / (86_400 * divisor)); - seconds.push((t % (86_400 * divisor)) / divisor); - subsecond.push((t % (86_400 * divisor)) % divisor); - } + let days = match timestamps.try_into_maybe_null_vec::() { + Ok(vector) => vector + .into_iter() + .map(|t| { + seconds.push((t % (86_400 * divisor)) / divisor); + subsecond.push((t % (86_400 * divisor)) % divisor); + t / (86_400 * divisor) + }) + .collect::>(), + Err(timestamps) => timestamps + .maybe_null_slice::() + .iter() + .map(|t| { + seconds.push((t % (86_400 * divisor)) / divisor); + subsecond.push((t % (86_400 * divisor)) % divisor); + t / (86_400 * divisor) + }) + .collect::>(), + }; Ok(TemporalParts { days: PrimitiveArray::from_vec(days, validity).into_array(), diff --git a/encodings/datetime-parts/src/compute/mod.rs b/encodings/datetime-parts/src/compute/mod.rs index d2d374d6d3d..dbf08840683 100644 --- a/encodings/datetime-parts/src/compute/mod.rs +++ b/encodings/datetime-parts/src/compute/mod.rs @@ -118,7 +118,7 @@ pub fn decode_to_temporal(array: &DateTimePartsArray) -> VortexResult = days_buf - .into_maybe_null_slice::() + .into_maybe_null_vec::() // attempt to reuse the i64 allocation for values .into_iter() .map(|d| d * 86_400 * divisor) .collect(); diff --git a/encodings/fastlanes/src/bitpacking/compute/filter.rs b/encodings/fastlanes/src/bitpacking/compute/filter.rs index 88a72562cdd..7706a3e2d1c 100644 --- a/encodings/fastlanes/src/bitpacking/compute/filter.rs +++ b/encodings/fastlanes/src/bitpacking/compute/filter.rs @@ -188,7 +188,7 @@ mod test { .unwrap() .into_primitive() .unwrap() - .into_maybe_null_slice::(); + .into_maybe_null_vec::(); assert_eq!(filtered, values); } diff --git a/encodings/fastlanes/src/delta/compress.rs b/encodings/fastlanes/src/delta/compress.rs index 3941d2edcac..e98ba8e1195 100644 --- a/encodings/fastlanes/src/delta/compress.rs +++ b/encodings/fastlanes/src/delta/compress.rs @@ -110,6 +110,9 @@ pub fn delta_decompress(array: DeltaArray) -> VortexResult { slice(decoded, array.offset(), array.offset() + array.len())?.into_primitive() } +// FIXME(DK): This method may benefit from taking deltas as an owned vector. Rust could re-use that +// allocation for `output` but proving that is a bit complex compared to +// `vector.into_iter().map().collect()`. fn decompress_primitive( bases: &[T], deltas: &[T], diff --git a/encodings/fastlanes/src/for/compress.rs b/encodings/fastlanes/src/for/compress.rs index e717ade2c62..0093840d635 100644 --- a/encodings/fastlanes/src/for/compress.rs +++ b/encodings/fastlanes/src/for/compress.rs @@ -27,7 +27,7 @@ pub fn for_compress(array: &PrimitiveArray) -> VortexResult { encoded_zero::<$T>(array.validity().to_logical(array.len()), nullability) .vortex_expect("Failed to encode all zeroes") } else { - compress_primitive::<$T>(&array, shift, $T::try_from(&min)?) + compress_primitive::<$T>(array, shift, $T::try_from(&min)?) .reinterpret_cast(array.ptype().to_unsigned()) .into_array() } @@ -120,7 +120,7 @@ pub fn decompress(array: FoRArray) -> VortexResult { encoded } else { PrimitiveArray::from_vec( - decompress_primitive(encoded.into_maybe_null_slice::<$T>(), min, shift), + decompress_primitive(encoded.into_maybe_null_vec::<$T>(), min, shift), validity, ) } @@ -128,6 +128,9 @@ pub fn decompress(array: FoRArray) -> VortexResult { })) } +/// Decompresses a frame-of-reference encoded vector given the shift and the bias (`min`). +/// +/// `values` is taken as an owned vector because we always need an allocation of that size. fn decompress_primitive( values: Vec, min: T, diff --git a/encodings/runend-bool/src/compute/mod.rs b/encodings/runend-bool/src/compute/mod.rs index 1f82f81d1ef..d86b041a2c0 100644 --- a/encodings/runend-bool/src/compute/mod.rs +++ b/encodings/runend-bool/src/compute/mod.rs @@ -1,11 +1,12 @@ mod invert; -use arrow_buffer::BooleanBuffer; -use vortex_array::array::BoolArray; +use arrow_buffer::{ArrowNativeType, BooleanBuffer}; +use num_traits::AsPrimitive; +use vortex_array::array::{BoolArray, PrimitiveArray}; use vortex_array::compute::{slice, ComputeVTable, InvertFn, ScalarAtFn, SliceFn, TakeFn}; use vortex_array::variants::PrimitiveArrayTrait; use vortex_array::{ArrayDType, ArrayData, ArrayLen, IntoArrayData, IntoArrayVariant}; -use vortex_dtype::match_each_integer_ptype; +use vortex_dtype::{match_each_integer_ptype, NativePType}; use vortex_error::{vortex_bail, VortexResult}; use vortex_scalar::Scalar; @@ -41,17 +42,7 @@ impl TakeFn for RunEndBoolEncoding { fn take(&self, array: &RunEndBoolArray, indices: &ArrayData) -> VortexResult { let primitive_indices = indices.clone().into_primitive()?; let physical_indices = match_each_integer_ptype!(primitive_indices.ptype(), |$P| { - primitive_indices - .into_maybe_null_slice::<$P>() - .into_iter() - .map(|idx| idx as usize) - .map(|idx| { - if idx >= array.len() { - vortex_bail!(OutOfBounds: idx, 0, array.len()) - } - array.find_physical_index(idx) - }) - .collect::>>()? + valid_physical_indices::<$P>(array, primitive_indices)? }); let start = array.start(); BoolArray::try_new( @@ -66,6 +57,36 @@ impl TakeFn for RunEndBoolEncoding { } } +fn valid_physical_indices>( + array: &RunEndBoolArray, + primitive_indices: PrimitiveArray, +) -> VortexResult> { + match primitive_indices.try_into_maybe_null_vec::

() { + // This allocation can be re-used when P's width matches usize's. + Ok(vector) => vector + .into_iter() + .map(|idx| idx.as_()) + .map(|idx| { + if idx >= array.len() { + vortex_bail!(OutOfBounds: idx, 0, array.len()) + } + array.find_physical_index(idx) + }) + .collect::>>(), + Err(primitive_indices) => primitive_indices + .maybe_null_slice::

() + .iter() + .map(|idx| (*idx).as_()) + .map(|idx| { + if idx >= array.len() { + vortex_bail!(OutOfBounds: idx, 0, array.len()) + } + array.find_physical_index(idx) + }) + .collect::>>(), + } +} + impl SliceFn for RunEndBoolEncoding { fn slice(&self, array: &RunEndBoolArray, start: usize, stop: usize) -> VortexResult { let new_length = stop - start; diff --git a/encodings/runend/src/compute/take.rs b/encodings/runend/src/compute/take.rs index 57d54acfcef..0a2e78fd683 100644 --- a/encodings/runend/src/compute/take.rs +++ b/encodings/runend/src/compute/take.rs @@ -13,10 +13,10 @@ impl TakeFn for RunEndEncoding { let primitive_indices = indices.clone().into_primitive()?; let usize_indices = match_each_integer_ptype!(primitive_indices.ptype(), |$P| { primitive_indices - .into_maybe_null_slice::<$P>() + .maybe_null_slice::<$P>() .into_iter() .map(|idx| { - let usize_idx = idx as usize; + let usize_idx = *idx as usize; if usize_idx >= array.len() { vortex_error::vortex_bail!(OutOfBounds: usize_idx, 0, array.len()); } diff --git a/encodings/zigzag/src/compress.rs b/encodings/zigzag/src/compress.rs index 301874335a0..4d418061e6b 100644 --- a/encodings/zigzag/src/compress.rs +++ b/encodings/zigzag/src/compress.rs @@ -11,10 +11,10 @@ use crate::ZigZagArray; pub fn zigzag_encode(parray: PrimitiveArray) -> VortexResult { let validity = parray.validity(); let encoded = match parray.ptype() { - PType::I8 => zigzag_encode_primitive::(parray.into_maybe_null_slice(), validity), - PType::I16 => zigzag_encode_primitive::(parray.into_maybe_null_slice(), validity), - PType::I32 => zigzag_encode_primitive::(parray.into_maybe_null_slice(), validity), - PType::I64 => zigzag_encode_primitive::(parray.into_maybe_null_slice(), validity), + PType::I8 => zigzag_encode_primitive::(parray.into_maybe_null_vec(), validity), + PType::I16 => zigzag_encode_primitive::(parray.into_maybe_null_vec(), validity), + PType::I32 => zigzag_encode_primitive::(parray.into_maybe_null_vec(), validity), + PType::I64 => zigzag_encode_primitive::(parray.into_maybe_null_vec(), validity), _ => vortex_bail!( "ZigZag can only encode signed integers, got {}", parray.ptype() @@ -36,10 +36,10 @@ where pub fn zigzag_decode(parray: PrimitiveArray) -> VortexResult { let validity = parray.validity(); let decoded = match parray.ptype() { - PType::U8 => zigzag_decode_primitive::(parray.into_maybe_null_slice(), validity), - PType::U16 => zigzag_decode_primitive::(parray.into_maybe_null_slice(), validity), - PType::U32 => zigzag_decode_primitive::(parray.into_maybe_null_slice(), validity), - PType::U64 => zigzag_decode_primitive::(parray.into_maybe_null_slice(), validity), + PType::U8 => zigzag_decode_primitive::(parray.into_maybe_null_vec(), validity), + PType::U16 => zigzag_decode_primitive::(parray.into_maybe_null_vec(), validity), + PType::U32 => zigzag_decode_primitive::(parray.into_maybe_null_vec(), validity), + PType::U64 => zigzag_decode_primitive::(parray.into_maybe_null_vec(), validity), _ => vortex_bail!( "ZigZag can only decode unsigned integers, got {}", parray.ptype() diff --git a/fuzz/src/slice.rs b/fuzz/src/slice.rs index 00b82029f27..64684cb127c 100644 --- a/fuzz/src/slice.rs +++ b/fuzz/src/slice.rs @@ -30,7 +30,7 @@ pub fn slice_canonical_array(array: &ArrayData, start: usize, stop: usize) -> Ar } DType::Primitive(p, _) => match_each_native_ptype!(p, |$P| { let primitive_array = array.clone().into_primitive().unwrap(); - let vec_values = primitive_array.into_maybe_null_slice::<$P>(); + let vec_values = primitive_array.maybe_null_slice::<$P>(); PrimitiveArray::from_vec(vec_values[start..stop].into(), validity).into_array() }), DType::Utf8(_) | DType::Binary(_) => { diff --git a/vortex-array/src/array/bool/patch.rs b/vortex-array/src/array/bool/patch.rs index 8c2e1a4dfbc..ebff7466ce3 100644 --- a/vortex-array/src/array/bool/patch.rs +++ b/vortex-array/src/array/bool/patch.rs @@ -10,8 +10,9 @@ use crate::{ArrayLen, IntoArrayVariant, ToArrayData}; impl BoolArray { pub fn patch(self, patches: Patches) -> VortexResult { let length = self.len(); - let indices = patches.indices().clone().into_primitive()?; - let values = patches.values().clone().into_bool()?; + let (_, indices, values) = patches.into_parts(); + let indices = indices.into_primitive()?; + let values = values.into_bool()?; let patched_validity = self.validity() @@ -20,11 +21,11 @@ impl BoolArray { let (mut own_values, bit_offset) = self.into_boolean_builder(); match_each_integer_ptype!(indices.ptype(), |$I| { for (idx, value) in indices - .into_maybe_null_slice::<$I>() - .into_iter() + .maybe_null_slice::<$I>() + .iter() .zip_eq(values.boolean_buffer().iter()) { - own_values.set_bit(idx as usize + bit_offset, value); + own_values.set_bit(*idx as usize + bit_offset, value); } }); diff --git a/vortex-array/src/array/chunked/compute/mod.rs b/vortex-array/src/array/chunked/compute/mod.rs index 242082ae032..db055873440 100644 --- a/vortex-array/src/array/chunked/compute/mod.rs +++ b/vortex-array/src/array/chunked/compute/mod.rs @@ -109,7 +109,7 @@ mod test { .unwrap() .into_primitive() .unwrap() - .into_maybe_null_slice::(), + .into_maybe_null_vec::(), vec![0u64, 1, 2, 3], ); } diff --git a/vortex-array/src/array/primitive/compute/take.rs b/vortex-array/src/array/primitive/compute/take.rs index 4c2df9a37f9..081f23cc41a 100644 --- a/vortex-array/src/array/primitive/compute/take.rs +++ b/vortex-array/src/array/primitive/compute/take.rs @@ -14,6 +14,9 @@ impl TakeFn for PrimitiveEncoding { let indices = indices.clone().into_primitive()?; let validity = array.validity().take(indices.as_ref())?; + // FIXME(DK): we could save an allocation and re-use memory if: we take the indices as + // owned, there are no other references to the underlying indices buffer, and the indices + // bit-width matches the array's bit-width. match_each_native_ptype!(array.ptype(), |$T| { match_each_integer_ptype!(indices.ptype(), |$I| { let values = take_primitive(array.maybe_null_slice::<$T>(), indices.maybe_null_slice::<$I>()); @@ -30,6 +33,9 @@ impl TakeFn for PrimitiveEncoding { let indices = indices.clone().into_primitive()?; let validity = unsafe { array.validity().take_unchecked(indices.as_ref())? }; + // FIXME(DK): we could save an allocation and re-use memory if: We take the indices as + // owned, there are no other references to the underlying indices buffer, and the indices + // bit-width matches the array's bit-width. 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.maybe_null_slice::<$I>()); @@ -39,8 +45,6 @@ impl TakeFn for PrimitiveEncoding { } } -// We pass a Vec in case we're T == u64. -// In which case, Rust should reuse the same Vec the result. fn take_primitive>( array: &[T], indices: &[I], @@ -48,8 +52,6 @@ fn take_primitive>( 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: &[I], diff --git a/vortex-array/src/array/primitive/mod.rs b/vortex-array/src/array/primitive/mod.rs index 0467653e911..e2eb58c14aa 100644 --- a/vortex-array/src/array/primitive/mod.rs +++ b/vortex-array/src/array/primitive/mod.rs @@ -115,8 +115,26 @@ impl PrimitiveArray { } /// Convert the array into a mutable vec of the given type. - /// If possible, this will be zero-copy. - pub fn into_maybe_null_slice(self) -> Vec { + /// + /// If there are no other references to the underlying buffer, no data is copied. + pub fn into_maybe_null_vec(self) -> Vec { + match self.try_into_maybe_null_vec::() { + Ok(vector) => vector, + Err(array) => { + let buffer = array.into_buffer(); + let (prefix, values, suffix) = unsafe { buffer.as_ref().align_to::() }; + assert!(prefix.is_empty() && suffix.is_empty()); + Vec::from(values) + } + } + } + + /// Try to convert the array into a mutable vec of the given type without copying. + /// + /// If the underlying buffer is not uniquely owned, which would necessitate a copy, then return + /// a primitive array with the same buffer and validity. + #[allow(clippy::panic_in_result_fn)] + pub fn try_into_maybe_null_vec(self) -> Result, Self> { assert_eq!( T::PTYPE, self.ptype(), @@ -124,11 +142,11 @@ impl PrimitiveArray { T::PTYPE, self.ptype(), ); - self.into_buffer().into_vec::().unwrap_or_else(|b| { - let (prefix, values, suffix) = unsafe { b.as_ref().align_to::() }; - assert!(prefix.is_empty() && suffix.is_empty()); - Vec::from(values) - }) + let ptype = self.ptype(); + let (buffer, validity) = self.into_parts(); + buffer + .into_vec::() + .map_err(|buffer| PrimitiveArray::new(buffer, ptype, validity)) } pub fn get_as_cast(&self, idx: usize) -> T { @@ -151,6 +169,12 @@ impl PrimitiveArray { PrimitiveArray::new(self.buffer().clone(), ptype, self.validity()) } + pub fn into_parts(self) -> (Buffer, Validity) { + let validity = self.validity(); + let buffer = self.into_buffer(); + (buffer, validity) + } + pub fn into_buffer(self) -> Buffer { self.into_array() .into_buffer() diff --git a/vortex-array/src/array/primitive/patch.rs b/vortex-array/src/array/primitive/patch.rs index 4b6dc0419f3..01e54194706 100644 --- a/vortex-array/src/array/primitive/patch.rs +++ b/vortex-array/src/array/primitive/patch.rs @@ -36,12 +36,12 @@ impl PrimitiveArray { T: NativePType + ArrowNativeType, I: NativePType + ArrowNativeType, { - let mut own_values = self.into_maybe_null_slice::(); + let mut own_values = self.into_maybe_null_vec::(); - let patch_indices = patch_indices.into_maybe_null_slice::(); - let patch_values = patch_values.into_maybe_null_slice::(); + let patch_indices = patch_indices.maybe_null_slice::(); + let patch_values = patch_values.maybe_null_slice::(); for (idx, value) in itertools::zip_eq(patch_indices, patch_values) { - own_values[idx.as_usize()] = value; + own_values[idx.as_usize()] = *value; } Ok(Self::from_vec(own_values, patched_validity)) } @@ -62,7 +62,7 @@ mod tests { sliced .into_primitive() .unwrap() - .into_maybe_null_slice::(), + .into_maybe_null_vec::(), vec![2u32; 6] ); } diff --git a/vortex-array/src/array/varbinview/mod.rs b/vortex-array/src/array/varbinview/mod.rs index 4aceaf70cf6..bf13d8231ba 100644 --- a/vortex-array/src/array/varbinview/mod.rs +++ b/vortex-array/src/array/varbinview/mod.rs @@ -477,7 +477,7 @@ impl VarBinViewArray { (view.len() + view_ref.offset()) as usize, )? .into_primitive()?; - Ok(data_buf.maybe_null_slice::().to_vec()) + Ok(data_buf.into_maybe_null_vec::()) } else { // Return access to the range of bytes around it. Ok(view.as_inlined().value().to_vec()) diff --git a/vortex-array/src/compute/filter.rs b/vortex-array/src/compute/filter.rs index 030f8195b58..8f3aa78c723 100644 --- a/vortex-array/src/compute/filter.rs +++ b/vortex-array/src/compute/filter.rs @@ -375,7 +375,7 @@ mod test { .unwrap() .into_primitive() .unwrap() - .into_maybe_null_slice::(), + .into_maybe_null_vec::(), vec![0i32, 1i32, 2i32] ); } diff --git a/vortex-array/src/patches.rs b/vortex-array/src/patches.rs index e67e6350855..676fe8c35fc 100644 --- a/vortex-array/src/patches.rs +++ b/vortex-array/src/patches.rs @@ -202,12 +202,13 @@ impl Patches { let flat_indices = self.indices().clone().into_primitive()?; match_each_integer_ptype!(flat_indices.ptype(), |$I| { - for (value_idx, coordinate) in flat_indices.into_maybe_null_slice::<$I>().into_iter().enumerate() { - if buffer.value(coordinate as usize) { + for (value_idx, coordinate) in flat_indices.maybe_null_slice::<$I>().iter().enumerate() { + let coordinate = *coordinate as usize; + if buffer.value(coordinate) { // We count the number of truthy values between this coordinate and the previous truthy one - let adjusted_coordinate = buffer.slice(last_inserted_index, (coordinate as usize) - last_inserted_index).count_set_bits() as u64; + let adjusted_coordinate = buffer.slice(last_inserted_index, (coordinate) - last_inserted_index).count_set_bits() as u64; coordinate_indices.push(adjusted_coordinate + coordinate_indices.last().copied().unwrap_or_default()); - last_inserted_index = coordinate as usize; + last_inserted_index = coordinate; value_indices.push(value_idx as u64); } } @@ -266,11 +267,17 @@ impl Patches { pub fn take_search(&self, take_indices: PrimitiveArray) -> VortexResult> { let new_length = take_indices.len(); let take_indices = match_each_integer_ptype!(take_indices.ptype(), |$P| { - take_indices - .into_maybe_null_slice::<$P>() - .into_iter() - .map(usize::try_from) - .collect::, _>>()? + match take_indices.try_into_maybe_null_vec::<$P>() { + Ok(vec) => vec + .into_iter() + .map(usize::try_from) + .collect::, _>>()?, + Err(take_indices) => take_indices + .maybe_null_slice::<$P>() + .iter() + .map(|x| usize::try_from(*x)) + .collect::, _>>()?, + } }); let (values_indices, new_indices): (Vec, Vec) =