diff --git a/vortex-array/src/arrays/decimal/array.rs b/vortex-array/src/arrays/decimal/array.rs index 1f4d6c0c277..0f1f3e91ea4 100644 --- a/vortex-array/src/arrays/decimal/array.rs +++ b/vortex-array/src/arrays/decimal/array.rs @@ -44,8 +44,6 @@ use crate::validity::Validity; /// The validity bitmap indicating which elements are non-null. pub(super) const VALIDITY_SLOT: usize = 0; -pub(super) const NUM_SLOTS: usize = 1; -pub(super) const SLOT_NAMES: [&str; NUM_SLOTS] = ["validity"]; /// A decimal array that stores fixed-precision decimal numbers with configurable scale. /// diff --git a/vortex-array/src/arrays/decimal/compute/take.rs b/vortex-array/src/arrays/decimal/compute/take.rs index d7bde288261..0e179bf911a 100644 --- a/vortex-array/src/arrays/decimal/compute/take.rs +++ b/vortex-array/src/arrays/decimal/compute/take.rs @@ -1,33 +1,32 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -use std::ptr; - -use itertools::Itertools as _; -use vortex_buffer::Buffer; -use vortex_buffer::BufferMut; use vortex_error::VortexResult; -use vortex_error::vortex_ensure; -use vortex_error::vortex_err; use crate::ArrayRef; use crate::Columnar; use crate::IntoArray; use crate::array::ArrayView; +use crate::arrays::ConstantArray; use crate::arrays::Decimal; use crate::arrays::DecimalArray; use crate::arrays::PiecewiseSequence; use crate::arrays::PrimitiveArray; use crate::arrays::dict::TakeExecute; +use crate::arrays::fixed_width::take::PreparedTakeIndices; +use crate::arrays::fixed_width::take::prepare_take_indices; +use crate::arrays::fixed_width::take::take_slices as take_slices_to_buffer; +use crate::arrays::fixed_width::take::take_slices_constant_length as take_slices_constant_length_to_buffer; +use crate::arrays::fixed_width::take::take_values; use crate::arrays::piecewise_sequence::constant_unsigned_usize; use crate::arrays::piecewise_sequence::maybe_contiguous_slices; -use crate::dtype::IntegerPType; use crate::dtype::NativeDecimalType; use crate::dtype::UnsignedPType; use crate::executor::ExecutionCtx; use crate::match_each_decimal_value_type; use crate::match_each_integer_ptype; use crate::match_each_unsigned_integer_ptype; +use crate::scalar::Scalar; use crate::validity::Validity; impl TakeExecute for Decimal { @@ -42,15 +41,27 @@ impl TakeExecute for Decimal { return Ok(Some(taken)); } - let indices = indices.clone().execute::(ctx)?; - let validity = array.validity()?.take(&indices.clone().into_array())?; + let PreparedTakeIndices::Values { + indices, + validity: indices_validity, + } = prepare_take_indices(indices, ctx)? + else { + return Ok(Some( + ConstantArray::new(Scalar::null(array.dtype().as_nullable()), indices.len()) + .into_array(), + )); + }; + let validity = array + .validity()? + .take(&indices.clone().into_array())? + .and(indices_validity)?; // TODO(joe): if the true count of take indices validity is low, only take array values with // valid indices. let decimal = match_each_decimal_value_type!(array.values_type(), |D| { match_each_integer_ptype!(indices.ptype(), |I| { let buffer = - take_to_buffer::(indices.as_slice::(), array.buffer::().as_slice()); + take_values::(array.buffer::().as_slice(), indices.as_slice::())?; // SAFETY: Take operation preserves decimal dtype and creates valid buffer. // Validity is computed correctly from the parent array and indices. unsafe { DecimalArray::new_unchecked(buffer, array.decimal_dtype(), validity) } @@ -108,10 +119,10 @@ where D: NativeDecimalType, { match_each_unsigned_integer_ptype!(starts.ptype(), |S| { - let values = take_slices_constant_length_to_buffer::( + let values = take_slices_constant_length_to_buffer::( + array.buffer::().as_slice(), starts.as_slice::(), length, - array.buffer::().as_slice(), output_len, )?; @@ -162,10 +173,10 @@ where S: UnsignedPType, { match_each_unsigned_integer_ptype!(lengths.ptype(), |L| { - let values = take_slices_to_buffer::( + let values = take_slices_to_buffer::( + array.buffer::().as_slice(), starts.as_slice::(), lengths.as_slice::(), - array.buffer::().as_slice(), output_len, )?; @@ -177,93 +188,6 @@ where }) } -fn take_to_buffer(indices: &[I], values: &[T]) -> Buffer { - indices.iter().map(|idx| values[idx.as_()]).collect() -} - -fn take_slices_constant_length_to_buffer( - starts: &[S], - length: usize, - values: &[T], - output_len: usize, -) -> VortexResult> -where - S: UnsignedPType, - T: NativeDecimalType, -{ - let computed_len = starts - .len() - .checked_mul(length) - .ok_or_else(|| vortex_err!("PiecewiseSequenceArray output length overflows usize"))?; - vortex_ensure!( - computed_len == output_len, - "PiecewiseSequenceArray expanded length {computed_len} does not match declared length {output_len}" - ); - - let mut result = BufferMut::::with_capacity(output_len); - let spare = &mut result.spare_capacity_mut()[..output_len]; - let mut cursor = 0usize; - for &start in starts { - let start = start.as_(); - let src = &values[start..][..length]; - // SAFETY: `src` and the checked `spare` range have equal lengths and cannot overlap. - unsafe { - ptr::copy_nonoverlapping( - src.as_ptr(), - spare[cursor..][..src.len()].as_mut_ptr().cast::(), - src.len(), - ); - } - cursor += src.len(); - } - // SAFETY: the loop initialized the prefix `0..cursor` of the spare capacity. - unsafe { result.set_len(cursor) }; - vortex_ensure!( - result.len() == output_len, - "PiecewiseSequenceArray expanded length {} does not match declared length {output_len}", - result.len() - ); - Ok(result.freeze()) -} - -fn take_slices_to_buffer( - starts: &[S], - lengths: &[L], - values: &[T], - output_len: usize, -) -> VortexResult> -where - S: UnsignedPType, - L: UnsignedPType, - T: NativeDecimalType, -{ - let mut result = BufferMut::::with_capacity(output_len); - let spare = &mut result.spare_capacity_mut()[..output_len]; - let mut cursor = 0usize; - for (&start, &length) in starts.iter().zip_eq(lengths) { - let start = start.as_(); - let length = length.as_(); - let src = &values[start..][..length]; - // SAFETY: `src` and the checked `spare` range have equal lengths and cannot overlap. - unsafe { - ptr::copy_nonoverlapping( - src.as_ptr(), - spare[cursor..][..src.len()].as_mut_ptr().cast::(), - src.len(), - ); - } - cursor += src.len(); - } - // SAFETY: the loop initialized the prefix `0..cursor` of the spare capacity. - unsafe { result.set_len(cursor) }; - vortex_ensure!( - result.len() == output_len, - "PiecewiseSequenceArray expanded length {} does not match declared length {output_len}", - result.len() - ); - Ok(result.freeze()) -} - #[cfg(test)] mod tests { use rstest::rstest; @@ -314,6 +238,27 @@ mod tests { assert_arrays_eq!(expected, taken, &mut ctx); } + #[test] + fn test_take_null_indices_ignore_physical_value() { + let mut ctx = array_session().create_execution_ctx(); + let ddtype = DecimalDType::new(19, 1); + let array = DecimalArray::new( + buffer![10i128, 11i128, 12i128, 13i128], + ddtype, + Validity::NonNullable, + ); + + let indices = PrimitiveArray::new( + buffer![u32::MAX, 2, 3], + Validity::from_iter([false, true, true]), + ) + .into_array(); + let taken = array.take(indices).unwrap(); + + let expected = DecimalArray::from_option_iter([None, Some(12i128), Some(13)], ddtype); + assert_arrays_eq!(expected, taken, &mut ctx); + } + #[rstest] #[case(DecimalArray::new( buffer![100i128, 200i128, 300i128, 400i128, 500i128], diff --git a/vortex-array/src/arrays/decimal/vtable/mod.rs b/vortex-array/src/arrays/decimal/vtable/mod.rs index ec157b28c66..4d6cbd8678b 100644 --- a/vortex-array/src/arrays/decimal/vtable/mod.rs +++ b/vortex-array/src/arrays/decimal/vtable/mod.rs @@ -8,7 +8,6 @@ use vortex_buffer::Alignment; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_ensure; -use vortex_error::vortex_panic; use vortex_session::VortexSession; use crate::ArrayParts; @@ -19,6 +18,7 @@ use crate::array::Array; use crate::array::ArrayView; use crate::array::VTable; use crate::arrays::decimal::DecimalData; +use crate::arrays::fixed_width::vtable as fixed_width; use crate::buffer::BufferHandle; use crate::builders::ArrayBuilder; use crate::builders::DecimalBuilder; @@ -27,7 +27,6 @@ use crate::dtype::DecimalType; use crate::dtype::NativeDecimalType; use crate::match_each_decimal_value_type; use crate::serde::ArrayChildren; -use crate::validity::Validity; mod kernel; mod operations; mod validity; @@ -38,7 +37,6 @@ use vortex_session::registry::CachedId; use crate::EqMode; use crate::array::ArrayId; -use crate::arrays::decimal::array::SLOT_NAMES; use crate::arrays::decimal::compute::rules::RULES; use crate::hash::ArrayEq; use crate::hash::ArrayHash; @@ -85,17 +83,11 @@ impl VTable for Decimal { } fn buffer(array: ArrayView<'_, Self>, idx: usize) -> BufferHandle { - match idx { - 0 => array.values.clone(), - _ => vortex_panic!("DecimalArray buffer index {idx} out of bounds"), - } + fixed_width::buffer("DecimalArray", &array.values, idx) } fn buffer_name(_array: ArrayView<'_, Self>, idx: usize) -> Option { - match idx { - 0 => Some("values".to_string()), - _ => None, - } + fixed_width::buffer_name(idx) } fn with_buffers( @@ -103,13 +95,8 @@ impl VTable for Decimal { array: ArrayView<'_, Self>, buffers: &[BufferHandle], ) -> VortexResult> { - vortex_ensure!( - buffers.len() == 1, - "Expected 1 buffer, got {}", - buffers.len() - ); let mut data = array.data().clone(); - data.values = buffers[0].clone(); + data.values = fixed_width::replacement_buffer(buffers)?; Ok( ArrayParts::new(self.clone(), array.dtype().clone(), array.len(), data) .with_slots(array.slots().iter().cloned().collect()), @@ -138,25 +125,7 @@ impl VTable for Decimal { let DType::Decimal(_, nullability) = dtype else { vortex_bail!("Expected decimal dtype, got {dtype:?}"); }; - vortex_ensure!( - data.len() == len, - InvalidArgument: - "DecimalArray length {} does not match outer length {}", - data.len(), - len - ); - let validity = crate::array::child_to_validity(slots[0].as_ref(), *nullability); - if let Some(validity_len) = validity.maybe_len() { - vortex_ensure!( - validity_len == len, - InvalidArgument: - "DecimalArray validity len {} does not match outer length {}", - validity_len, - len - ); - } - - Ok(()) + fixed_width::validate_layout("DecimalArray", data.len(), *nullability, len, slots) } fn deserialize( @@ -174,14 +143,7 @@ impl VTable for Decimal { } let values = buffers[0].clone(); - let validity = if children.is_empty() { - Validity::from(dtype.nullability()) - } else if children.len() == 1 { - let validity = children.get(0, &Validity::DTYPE, len)?; - Validity::Array(validity) - } else { - vortex_bail!("Expected 0 or 1 child, got {}", children.len()); - }; + let validity = fixed_width::deserialize_validity(dtype.nullability(), len, children)?; let Some(decimal_dtype) = dtype.as_decimal_opt() else { vortex_bail!("Expected Decimal dtype, got {:?}", dtype) @@ -201,7 +163,7 @@ impl VTable for Decimal { } fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { - SLOT_NAMES[idx].to_string() + fixed_width::slot_name(idx) } fn execute(array: Array, _ctx: &mut ExecutionCtx) -> VortexResult { diff --git a/vortex-array/src/arrays/fixed_width/mod.rs b/vortex-array/src/arrays/fixed_width/mod.rs new file mode 100644 index 00000000000..c1ba80f78a3 --- /dev/null +++ b/vortex-array/src/arrays/fixed_width/mod.rs @@ -0,0 +1,7 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Shared structural operations for fixed-width canonical arrays. + +pub(crate) mod take; +pub(crate) mod vtable; diff --git a/vortex-array/src/arrays/fixed_width/take.rs b/vortex-array/src/arrays/fixed_width/take.rs new file mode 100644 index 00000000000..5731b3e27cd --- /dev/null +++ b/vortex-array/src/arrays/fixed_width/take.rs @@ -0,0 +1,168 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::ptr; + +use itertools::Itertools as _; +use vortex_buffer::Buffer; +use vortex_buffer::BufferMut; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_ensure; +use vortex_error::vortex_err; +use vortex_mask::Mask; + +use crate::ArrayRef; +use crate::arrays::PrimitiveArray; +use crate::builtins::ArrayBuiltins; +use crate::dtype::DType; +use crate::dtype::IntegerPType; +use crate::dtype::UnsignedPType; +use crate::executor::ExecutionCtx; +use crate::scalar::Scalar; +use crate::validity::Validity; + +pub(crate) enum PreparedTakeIndices { + AllInvalid, + Values { + indices: PrimitiveArray, + validity: Validity, + }, +} + +pub(crate) fn prepare_take_indices( + indices: &ArrayRef, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let DType::Primitive(ptype, nullability) = indices.dtype() else { + vortex_bail!("Invalid indices dtype: {}", indices.dtype()) + }; + + let validity = indices.validity()?; + // Null index lanes are semantically ignored, but their physical values may be out of bounds. + // Redirect those lanes to zero for the cast/gather, then restore their validity in the caller. + let zeroed = match validity.execute_mask(indices.len(), ctx)? { + Mask::AllTrue(_) => indices.clone(), + Mask::AllFalse(_) => return Ok(PreparedTakeIndices::AllInvalid), + Mask::Values(_) => indices + .clone() + .fill_null(Scalar::from(0).cast(indices.dtype())?)?, + }; + + let indices = if ptype.is_unsigned_int() { + zeroed.execute::(ctx)? + } else { + zeroed + .cast(DType::Primitive(ptype.to_unsigned(), *nullability))? + .execute::(ctx)? + }; + + Ok(PreparedTakeIndices::Values { indices, validity }) +} + +#[inline(always)] +pub(crate) fn take_values( + values: &[T], + indices: &[I], +) -> VortexResult> { + // The explicit pointer loop keeps the source length in a register and avoids a capacity check + // for every output value. + let mut result = BufferMut::with_capacity(indices.len()); + let ptr = result.spare_capacity_mut().as_mut_ptr().cast::(); + + for (i, idx) in indices.iter().enumerate() { + let index = idx.as_(); + if index >= values.len() { + vortex_bail!(OutOfBounds: index, 0, values.len()); + } + // SAFETY: `indices.len()` elements were reserved and each output position is written once. + unsafe { ptr.add(i).write(*values.get_unchecked(index)) }; + } + + // SAFETY: the loop initialized every element in the reserved output range. + unsafe { result.set_len(indices.len()) }; + Ok(result.freeze()) +} + +pub(crate) fn take_slices( + values: &[T], + starts: &[S], + lengths: &[L], + output_len: usize, +) -> VortexResult> +where + T: Copy, + S: UnsignedPType, + L: UnsignedPType, +{ + let mut result = BufferMut::::with_capacity(output_len); + let spare = &mut result.spare_capacity_mut()[..output_len]; + let mut cursor = 0usize; + for (&start, &length) in starts.iter().zip_eq(lengths) { + let start = start.as_(); + let length = length.as_(); + let src = &values[start..][..length]; + // SAFETY: `src` and the checked `spare` range have equal lengths and cannot overlap. + unsafe { + ptr::copy_nonoverlapping( + src.as_ptr(), + spare[cursor..][..src.len()].as_mut_ptr().cast::(), + src.len(), + ); + } + cursor += src.len(); + } + // SAFETY: the loop initialized the prefix `0..cursor` of the spare capacity. + unsafe { result.set_len(cursor) }; + vortex_ensure!( + result.len() == output_len, + "PiecewiseSequenceArray expanded length {} does not match declared length {output_len}", + result.len() + ); + Ok(result.freeze()) +} + +pub(crate) fn take_slices_constant_length( + values: &[T], + starts: &[S], + length: usize, + output_len: usize, +) -> VortexResult> +where + T: Copy, + S: UnsignedPType, +{ + let computed_len = starts + .len() + .checked_mul(length) + .ok_or_else(|| vortex_err!("PiecewiseSequenceArray output length overflows usize"))?; + vortex_ensure!( + computed_len == output_len, + "PiecewiseSequenceArray expanded length {computed_len} does not match declared length {output_len}" + ); + + let mut result = BufferMut::::with_capacity(output_len); + let spare = &mut result.spare_capacity_mut()[..output_len]; + let mut cursor = 0usize; + for &start in starts { + let start = start.as_(); + let src = &values[start..][..length]; + // SAFETY: `src` and the checked `spare` range have equal lengths and cannot overlap. + unsafe { + ptr::copy_nonoverlapping( + src.as_ptr(), + spare[cursor..][..src.len()].as_mut_ptr().cast::(), + src.len(), + ); + } + cursor += src.len(); + } + // SAFETY: the loop initialized the prefix `0..cursor` of the spare capacity. + unsafe { result.set_len(cursor) }; + vortex_ensure!( + result.len() == output_len, + "PiecewiseSequenceArray expanded length {} does not match declared length {output_len}", + result.len() + ); + Ok(result.freeze()) +} diff --git a/vortex-array/src/arrays/fixed_width/vtable.rs b/vortex-array/src/arrays/fixed_width/vtable.rs new file mode 100644 index 00000000000..9f4a9bd3f0e --- /dev/null +++ b/vortex-array/src/arrays/fixed_width/vtable.rs @@ -0,0 +1,82 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_ensure; +use vortex_error::vortex_panic; + +use crate::ArrayRef; +use crate::array::child_to_validity; +use crate::buffer::BufferHandle; +use crate::dtype::Nullability; +use crate::serde::ArrayChildren; +use crate::validity::Validity; + +pub(crate) fn buffer(array_name: &str, values: &BufferHandle, idx: usize) -> BufferHandle { + match idx { + 0 => values.clone(), + _ => vortex_panic!("{array_name} buffer index {idx} out of bounds"), + } +} + +pub(crate) fn buffer_name(idx: usize) -> Option { + match idx { + 0 => Some("values".to_string()), + _ => None, + } +} + +pub(crate) fn replacement_buffer(buffers: &[BufferHandle]) -> VortexResult { + vortex_ensure!( + buffers.len() == 1, + "Expected 1 buffer, got {}", + buffers.len() + ); + Ok(buffers[0].clone()) +} + +pub(crate) fn validate_layout( + array_name: &str, + data_len: usize, + nullability: Nullability, + len: usize, + slots: &[Option], +) -> VortexResult<()> { + vortex_ensure!(slots.len() == 1, "{array_name} expects one validity slot"); + vortex_ensure!( + data_len == len, + InvalidArgument: + "{array_name} length {data_len} does not match outer length {len}" + ); + let validity = child_to_validity(slots[0].as_ref(), nullability); + if let Some(validity_len) = validity.maybe_len() { + vortex_ensure!( + validity_len == len, + InvalidArgument: + "{array_name} validity len {validity_len} does not match outer length {len}" + ); + } + Ok(()) +} + +pub(crate) fn deserialize_validity( + nullability: Nullability, + len: usize, + children: &dyn ArrayChildren, +) -> VortexResult { + if children.is_empty() { + Ok(Validity::from(nullability)) + } else if children.len() == 1 { + Ok(Validity::Array(children.get(0, &Validity::DTYPE, len)?)) + } else { + vortex_bail!("Expected 0 or 1 child, got {}", children.len()) + } +} + +pub(crate) fn slot_name(idx: usize) -> String { + match idx { + 0 => "validity".to_string(), + _ => vortex_panic!("Fixed-width slot index {idx} out of bounds"), + } +} diff --git a/vortex-array/src/arrays/mod.rs b/vortex-array/src/arrays/mod.rs index 001e20dee65..e0483d8b220 100644 --- a/vortex-array/src/arrays/mod.rs +++ b/vortex-array/src/arrays/mod.rs @@ -60,6 +60,8 @@ pub mod filter; pub use filter::Filter; pub use filter::FilterArray; +pub(crate) mod fixed_width; + pub mod fixed_size_list; pub use fixed_size_list::FixedSizeList; pub use fixed_size_list::FixedSizeListArray; diff --git a/vortex-array/src/arrays/primitive/array/mod.rs b/vortex-array/src/arrays/primitive/array/mod.rs index 867c7d0c783..f8c9549f2e5 100644 --- a/vortex-array/src/arrays/primitive/array/mod.rs +++ b/vortex-array/src/arrays/primitive/array/mod.rs @@ -50,8 +50,6 @@ use crate::builtins::ArrayBuiltins; /// The validity bitmap indicating which elements are non-null. pub(super) const VALIDITY_SLOT: usize = 0; -pub(super) const NUM_SLOTS: usize = 1; -pub(super) const SLOT_NAMES: [&str; NUM_SLOTS] = ["validity"]; /// A primitive array that stores [native types][crate::dtype::NativePType] in a contiguous buffer /// of memory, along with an optional validity child. diff --git a/vortex-array/src/arrays/primitive/compute/take/avx2.rs b/vortex-array/src/arrays/primitive/compute/take/avx2.rs index fb64e6748b0..62ab4efb38a 100644 --- a/vortex-array/src/arrays/primitive/compute/take/avx2.rs +++ b/vortex-array/src/arrays/primitive/compute/take/avx2.rs @@ -36,14 +36,15 @@ use std::convert::identity; use vortex_buffer::Alignment; use vortex_buffer::Buffer; use vortex_buffer::BufferMut; +use vortex_error::VortexExpect; use vortex_error::VortexResult; use crate::ArrayRef; use crate::IntoArray; use crate::array::ArrayView; use crate::arrays::PrimitiveArray; +use crate::arrays::fixed_width::take::take_values; use crate::arrays::primitive::compute::take::TakeImpl; -use crate::arrays::primitive::compute::take::take_primitive_scalar; use crate::arrays::primitive::vtable::Primitive; use crate::dtype::NativePType; use crate::dtype::UnsignedPType; @@ -152,7 +153,8 @@ unsafe fn take_avx2(buffer: &[V], indices: &[I]) -> B 4 => dispatch!(u32), 8 => dispatch!(u64), // 1/2-byte and >8-byte values have no AVX2 gather lane, so fall back to scalar. - _ => take_primitive_scalar(buffer, indices), + _ => take_values(buffer, indices) + .vortex_expect("AVX2 scalar fallback take index should be in bounds"), } } diff --git a/vortex-array/src/arrays/primitive/compute/take/mod.rs b/vortex-array/src/arrays/primitive/compute/take/mod.rs index 5d663f55967..8fd08d66129 100644 --- a/vortex-array/src/arrays/primitive/compute/take/mod.rs +++ b/vortex-array/src/arrays/primitive/compute/take/mod.rs @@ -4,17 +4,9 @@ #[cfg(any(target_arch = "x86_64", target_arch = "x86"))] mod avx2; -use std::ptr; use std::sync::LazyLock; -use itertools::Itertools as _; -use vortex_buffer::Buffer; -use vortex_buffer::BufferMut; use vortex_error::VortexResult; -use vortex_error::vortex_bail; -use vortex_error::vortex_ensure; -use vortex_error::vortex_err; -use vortex_mask::Mask; use crate::ArrayRef; use crate::Columnar; @@ -25,11 +17,13 @@ use crate::arrays::PiecewiseSequence; use crate::arrays::Primitive; use crate::arrays::PrimitiveArray; use crate::arrays::dict::TakeExecute; +use crate::arrays::fixed_width::take::PreparedTakeIndices; +use crate::arrays::fixed_width::take::prepare_take_indices; +use crate::arrays::fixed_width::take::take_slices as take_slices_to_buffer; +use crate::arrays::fixed_width::take::take_slices_constant_length as take_slices_constant_length_to_buffer; +use crate::arrays::fixed_width::take::take_values; use crate::arrays::piecewise_sequence::constant_unsigned_usize; use crate::arrays::piecewise_sequence::maybe_contiguous_slices; -use crate::builtins::ArrayBuiltins; -use crate::dtype::DType; -use crate::dtype::IntegerPType; use crate::dtype::NativePType; use crate::dtype::UnsignedPType; use crate::executor::ExecutionCtx; @@ -77,7 +71,7 @@ impl TakeImpl for TakeKernelScalar { ) -> VortexResult { match_each_native_ptype!(array.ptype(), |T| { match_each_integer_ptype!(indices.ptype(), |I| { - let values = take_primitive_scalar(array.as_slice::(), indices.as_slice::()); + let values = take_values(array.as_slice::(), indices.as_slice::())?; Ok(PrimitiveArray::new(values, validity).into_array()) }) }) @@ -96,34 +90,15 @@ impl TakeExecute for Primitive { return Ok(Some(taken)); } - let DType::Primitive(ptype, null) = indices.dtype() else { - vortex_bail!("Invalid indices dtype: {}", indices.dtype()) - }; - - let indices_validity = indices.validity()?; - // Null index lanes are semantically ignored, but their physical values may be out of - // bounds. Redirect those lanes to zero for the cast/gather, then restore the original index - // validity below. - let indices_nulls_zeroed = match indices_validity.execute_mask(indices.len(), ctx)? { - Mask::AllTrue(_) => indices.clone(), - Mask::AllFalse(_) => { - return Ok(Some( - ConstantArray::new(Scalar::null(array.dtype().as_nullable()), indices.len()) - .into_array(), - )); - } - Mask::Values(_) => indices - .clone() - .fill_null(Scalar::from(0).cast(indices.dtype())?)?, - }; - - let unsigned_indices = if ptype.is_unsigned_int() { - indices_nulls_zeroed.execute::(ctx)? - } else { - // This will fail if all values cannot be converted to unsigned - indices_nulls_zeroed - .cast(DType::Primitive(ptype.to_unsigned(), *null))? - .execute::(ctx)? + let PreparedTakeIndices::Values { + indices: unsigned_indices, + validity: indices_validity, + } = prepare_take_indices(indices, ctx)? + else { + return Ok(Some( + ConstantArray::new(Scalar::null(array.dtype().as_nullable()), indices.len()) + .into_array(), + )); }; let validity = array @@ -164,27 +139,6 @@ fn take_contiguous_ranges( Ok(Some(taken)) } -// Compiler may see this as unused based on enabled features -#[inline(always)] -fn take_primitive_scalar(buffer: &[T], indices: &[I]) -> Buffer { - // NB: The simpler `indices.iter().map(|idx| buffer[idx.as_()]).collect()` generates suboptimal - // assembly where the buffer length is repeatedly loaded from the stack on each iteration. - - let mut result = BufferMut::with_capacity(indices.len()); - let ptr = result.spare_capacity_mut().as_mut_ptr().cast::(); - - // This explicit loop with pointer writes keeps the length in a register and avoids per-element - // capacity checks from `push()`. - for (i, idx) in indices.iter().enumerate() { - // SAFETY: We reserved `indices.len()` capacity, so `ptr.add(i)` is valid. - unsafe { ptr.add(i).write(buffer[idx.as_()]) }; - } - - // SAFETY: We just wrote exactly `indices.len()` elements. - unsafe { result.set_len(indices.len()) }; - result.freeze() -} - fn take_slices( array: ArrayView<'_, Primitive>, starts: &PrimitiveArray, @@ -234,44 +188,6 @@ where }) } -fn take_slices_to_buffer( - source: &[T], - starts: &[S], - lengths: &[L], - output_len: usize, -) -> VortexResult> -where - T: Copy, - S: UnsignedPType, - L: UnsignedPType, -{ - let mut values = BufferMut::::with_capacity(output_len); - let spare = &mut values.spare_capacity_mut()[..output_len]; - let mut cursor = 0usize; - for (&start, &length) in starts.iter().zip_eq(lengths) { - let start = start.as_(); - let length = length.as_(); - let src = &source[start..][..length]; - // SAFETY: `src` and the checked `spare` range have equal lengths and cannot overlap. - unsafe { - ptr::copy_nonoverlapping( - src.as_ptr(), - spare[cursor..][..src.len()].as_mut_ptr().cast::(), - src.len(), - ); - } - cursor += src.len(); - } - // SAFETY: the loop initialized the prefix `0..cursor` of the spare capacity. - unsafe { values.set_len(cursor) }; - vortex_ensure!( - values.len() == output_len, - "PiecewiseSequenceArray expanded length {} does not match declared length {output_len}", - values.len() - ); - Ok(values.freeze()) -} - fn take_slices_constant_length( array: ArrayView<'_, Primitive>, starts: &PrimitiveArray, @@ -305,73 +221,30 @@ where }) } -fn take_slices_constant_length_to_buffer( - source: &[T], - starts: &[S], - length: usize, - output_len: usize, -) -> VortexResult> -where - T: Copy, - S: UnsignedPType, -{ - let computed_len = starts - .len() - .checked_mul(length) - .ok_or_else(|| vortex_err!("PiecewiseSequenceArray output length overflows usize"))?; - vortex_ensure!( - computed_len == output_len, - "PiecewiseSequenceArray expanded length {computed_len} does not match declared length {output_len}" - ); - - let mut values = BufferMut::::with_capacity(output_len); - let spare = &mut values.spare_capacity_mut()[..output_len]; - let mut cursor = 0usize; - for &start in starts { - let start = start.as_(); - let src = &source[start..][..length]; - // SAFETY: `src` and the checked `spare` range have equal lengths and cannot overlap. - unsafe { - ptr::copy_nonoverlapping( - src.as_ptr(), - spare[cursor..][..src.len()].as_mut_ptr().cast::(), - src.len(), - ); - } - cursor += src.len(); - } - // SAFETY: the loop initialized the prefix `0..cursor` of the spare capacity. - unsafe { values.set_len(cursor) }; - vortex_ensure!( - values.len() == output_len, - "PiecewiseSequenceArray expanded length {} does not match declared length {output_len}", - values.len() - ); - Ok(values.freeze()) -} - #[cfg(any(target_arch = "x86_64", target_arch = "x86"))] #[cfg(test)] mod test { use rstest::rstest; use vortex_buffer::buffer; use vortex_error::VortexExpect; + use vortex_error::VortexResult; use crate::IntoArray; use crate::VortexSessionExecute; use crate::array_session; use crate::arrays::BoolArray; use crate::arrays::PrimitiveArray; - use crate::arrays::primitive::compute::take::take_primitive_scalar; + use crate::arrays::fixed_width::take::take_values; use crate::compute::conformance::take::test_take_conformance; use crate::scalar::Scalar; use crate::validity::Validity; #[test] - fn test_take() { + fn test_take() -> VortexResult<()> { let a = vec![1i32, 2, 3, 4, 5]; - let result = take_primitive_scalar(&a, &[0, 0, 4, 2]); + let result = take_values(&a, &[0, 0, 4, 2])?; assert_eq!(result.as_slice(), &[1i32, 1, 5, 3]); + Ok(()) } #[test] diff --git a/vortex-array/src/arrays/primitive/vtable/mod.rs b/vortex-array/src/arrays/primitive/vtable/mod.rs index 50189bf0037..2cbd08a7874 100644 --- a/vortex-array/src/arrays/primitive/vtable/mod.rs +++ b/vortex-array/src/arrays/primitive/vtable/mod.rs @@ -4,7 +4,6 @@ use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_ensure; -use vortex_error::vortex_panic; use crate::ArrayParts; use crate::ArrayRef; @@ -13,6 +12,7 @@ use crate::ExecutionResult; use crate::array::Array; use crate::array::ArrayView; use crate::array::VTable; +use crate::arrays::fixed_width::vtable as fixed_width; use crate::arrays::primitive::PrimitiveData; use crate::buffer::BufferHandle; use crate::builders::ArrayBuilder; @@ -21,7 +21,6 @@ use crate::dtype::DType; use crate::dtype::PType; use crate::match_each_native_ptype; use crate::serde::ArrayChildren; -use crate::validity::Validity; mod kernel; mod operations; mod validity; @@ -34,7 +33,6 @@ use vortex_session::registry::CachedId; use crate::EqMode; use crate::array::ArrayId; -use crate::arrays::primitive::array::SLOT_NAMES; use crate::arrays::primitive::compute::rules::RULES; use crate::hash::ArrayEq; use crate::hash::ArrayHash; @@ -74,17 +72,11 @@ impl VTable for Primitive { } fn buffer(array: ArrayView<'_, Self>, idx: usize) -> BufferHandle { - match idx { - 0 => array.buffer_handle().clone(), - _ => vortex_panic!("PrimitiveArray buffer index {idx} out of bounds"), - } + fixed_width::buffer("PrimitiveArray", array.buffer_handle(), idx) } fn buffer_name(_array: ArrayView<'_, Self>, idx: usize) -> Option { - match idx { - 0 => Some("values".to_string()), - _ => None, - } + fixed_width::buffer_name(idx) } fn with_buffers( @@ -92,13 +84,8 @@ impl VTable for Primitive { array: ArrayView<'_, Self>, buffers: &[BufferHandle], ) -> VortexResult> { - vortex_ensure!( - buffers.len() == 1, - "Expected 1 buffer, got {}", - buffers.len() - ); let mut data = array.data().clone(); - data.buffer = buffers[0].clone(); + data.buffer = fixed_width::replacement_buffer(buffers)?; Ok( ArrayParts::new(self.clone(), array.dtype().clone(), array.len(), data) .with_slots(array.slots().iter().cloned().collect()), @@ -122,23 +109,7 @@ impl VTable for Primitive { let DType::Primitive(_, nullability) = dtype else { vortex_bail!("Expected primitive dtype, got {dtype:?}"); }; - vortex_ensure!( - data.len() == len, - "PrimitiveArray length {} does not match outer length {}", - data.len(), - len - ); - let validity = crate::array::child_to_validity(slots[0].as_ref(), *nullability); - if let Some(validity_len) = validity.maybe_len() { - vortex_ensure!( - validity_len == len, - "PrimitiveArray validity len {} does not match outer length {}", - validity_len, - len - ); - } - - Ok(()) + fixed_width::validate_layout("PrimitiveArray", data.len(), *nullability, len, slots) } fn deserialize( @@ -162,14 +133,7 @@ impl VTable for Primitive { } let buffer = buffers[0].clone(); - let validity = if children.is_empty() { - Validity::from(dtype.nullability()) - } else if children.len() == 1 { - let validity = children.get(0, &Validity::DTYPE, len)?; - Validity::Array(validity) - } else { - vortex_bail!("Expected 0 or 1 child, got {}", children.len()); - }; + let validity = fixed_width::deserialize_validity(dtype.nullability(), len, children)?; let ptype = PType::try_from(dtype)?; @@ -202,7 +166,7 @@ impl VTable for Primitive { } fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { - SLOT_NAMES[idx].to_string() + fixed_width::slot_name(idx) } fn execute(array: Array, _ctx: &mut ExecutionCtx) -> VortexResult {