diff --git a/vortex-array/src/arrays/decimal/compute/fixed_width.rs b/vortex-array/src/arrays/decimal/compute/fixed_width.rs index 8ec0987e0f8..6ab794431c2 100644 --- a/vortex-array/src/arrays/decimal/compute/fixed_width.rs +++ b/vortex-array/src/arrays/decimal/compute/fixed_width.rs @@ -1,6 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_buffer::Buffer; use vortex_buffer::ByteBuffer; use vortex_error::VortexResult; @@ -16,8 +17,10 @@ impl FixedWidthArray for Decimal { array.values_type().byte_width() } - fn values(array: ArrayView<'_, Self>) -> ByteBuffer { - array.buffer_handle().to_host_sync() + fn values(array: ArrayView<'_, Self>) -> Buffer { + let values = array.buffer_handle().to_host_sync(); + let alignment = values.alignment(); + Buffer::from_byte_buffer_aligned(values, alignment) } fn with_values( diff --git a/vortex-array/src/arrays/fixed_width/array.rs b/vortex-array/src/arrays/fixed_width/array.rs index 96ee6e6255f..c760b14a397 100644 --- a/vortex-array/src/arrays/fixed_width/array.rs +++ b/vortex-array/src/arrays/fixed_width/array.rs @@ -1,6 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_buffer::Buffer; use vortex_buffer::ByteBuffer; use vortex_error::VortexResult; use vortex_error::vortex_ensure_eq; @@ -22,10 +23,11 @@ pub(crate) trait FixedWidthArray: VTable { /// Returns the number of bytes each record occupies. fn byte_width(array: ArrayView<'_, Self>) -> usize; - /// Returns the records of `array` as a single host-resident byte buffer. + /// Returns the storage of `array` as a host-resident buffer of `T`, preserving its alignment. /// /// The returned buffer must contain exactly `array.len() * byte_width` bytes. - fn values(array: ArrayView<'_, Self>) -> ByteBuffer; + /// Its byte length and alignment must be compatible with `T`. + fn values(array: ArrayView<'_, Self>) -> Buffer; /// Rebuilds an array of this encoding from a records buffer, preserving the logical type of /// `array`. diff --git a/vortex-array/src/arrays/fixed_width/filter.rs b/vortex-array/src/arrays/fixed_width/filter.rs index 7d7eec9d635..c5b6d1d3466 100644 --- a/vortex-array/src/arrays/fixed_width/filter.rs +++ b/vortex-array/src/arrays/fixed_width/filter.rs @@ -1,19 +1,19 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -use vortex_buffer::Buffer; -use vortex_buffer::BufferMut; use vortex_buffer::ByteBuffer; use vortex_error::VortexExpect; +use vortex_error::vortex_panic; use vortex_mask::MaskValues; use vortex_mask::MaskValuesRef; use super::FixedWidthArray; -use super::match_each_record_width; use super::with_values; use crate::array::Array; +use crate::array::ArrayView; use crate::arrays::filter::filter_buffer; use crate::arrays::filter::filter_validity; +use crate::dtype::i256; #[cfg(test)] #[expect(clippy::cast_possible_truncation)] @@ -21,7 +21,15 @@ mod tests; pub(crate) fn filter(array: &Array, mask: &MaskValuesRef) -> Array { let array = array.as_view(); - let values = filter_records(V::values(array), V::byte_width(array), mask.as_ref()); + let values = match V::byte_width(array) { + 1 => filter_records::(array, mask.as_ref()), + 2 => filter_records::(array, mask.as_ref()), + 4 => filter_records::(array, mask.as_ref()), + 8 => filter_records::(array, mask.as_ref()), + 16 => filter_records::(array, mask.as_ref()), + 32 => filter_records::(array, mask.as_ref()), + byte_width => vortex_panic!("Unsupported fixed-width byte width: {byte_width}"), + }; let validity = filter_validity( array .validity() @@ -32,39 +40,15 @@ pub(crate) fn filter(array: &Array, mask: &MaskValuesRef) .vortex_expect("filtering fixed-width values preserves array invariants") } -fn filter_records(values: ByteBuffer, byte_width: usize, mask: &MaskValues) -> ByteBuffer { +fn filter_records( + array: ArrayView<'_, V>, + mask: &MaskValues, +) -> ByteBuffer { + let values = V::values::(array); let alignment = values.alignment(); - - match_each_record_width!( - byte_width, - |W| { - let records = Buffer::<[u8; W]>::from_byte_buffer(values); - // `filter_buffer` picks between in-place compaction, cached indices/slices, - // byte-compress, and bitmap iteration based on record width and mask density. - let filtered = filter_buffer(records, mask); - filtered.into_byte_buffer().aligned(alignment) - }, - _ => { - match values.try_into_mut() { - Ok(mut values) => { - let mut destination = 0; - mask.bit_buffer().for_each_set_index(|index| { - let source = index * byte_width; - values.copy_within(source..source + byte_width, destination); - destination += byte_width; - }); - values.truncate(destination); - values.freeze().into_byte_buffer().aligned(alignment) - } - Err(values) => { - let mut filtered = BufferMut::with_capacity(mask.true_count() * byte_width); - mask.bit_buffer().for_each_set_index(|index| { - let start = index * byte_width; - filtered.extend_from_slice(&values[start..start + byte_width]); - }); - filtered.freeze().into_byte_buffer().aligned(alignment) - } - } - } - ) + // `filter_buffer` picks between in-place compaction, cached indices/slices, + // byte-compress, and bitmap iteration based on record width and mask density. + filter_buffer(values, mask) + .into_byte_buffer() + .aligned(alignment) } diff --git a/vortex-array/src/arrays/fixed_width/filter/tests.rs b/vortex-array/src/arrays/fixed_width/filter/tests.rs index cf01e5d39e5..0fe6c7eb7ed 100644 --- a/vortex-array/src/arrays/fixed_width/filter/tests.rs +++ b/vortex-array/src/arrays/fixed_width/filter/tests.rs @@ -2,17 +2,21 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use rstest::rstest; +use vortex_buffer::Alignment; use vortex_buffer::Buffer; use vortex_buffer::buffer; +use vortex_error::VortexResult; use vortex_mask::Mask; -use super::filter_records; use crate::ArrayRef; use crate::IntoArray; use crate::VortexSessionExecute; +use crate::array::Array; use crate::array_session; use crate::arrays::DecimalArray; use crate::arrays::PrimitiveArray; +use crate::arrays::fixed_width::FixedWidthArray; +use crate::arrays::fixed_width::with_values; use crate::compute::conformance::filter::LARGE_SIZE; use crate::compute::conformance::filter::MEDIUM_SIZE; use crate::compute::conformance::filter::test_filter_conformance; @@ -20,23 +24,39 @@ use crate::dtype::DecimalDType; use crate::dtype::i256; use crate::validity::Validity; -#[test] -fn filter_fallback_width_records() { +#[rstest] +#[case::u8(PrimitiveArray::from_iter([10u8, 20, 30, 40]))] +#[case::u16(PrimitiveArray::from_iter([10u16, 20, 30, 40]))] +#[case::u32(PrimitiveArray::from_iter([10u32, 20, 30, 40]))] +#[case::u64(PrimitiveArray::from_iter([10u64, 20, 30, 40]))] +#[case::u128(DecimalArray::new( + buffer![10i128, -20, -30, -40], + DecimalDType::new(19, 0), + Validity::NonNullable, +))] +#[case::i256(DecimalArray::new( + Buffer::from_iter([10, -20, -30, -40].map(|value| i256::from_parts(1, value))), + DecimalDType::new(76, 0), + Validity::NonNullable, +))] +fn filter_typed_records(#[case] array: Array) -> VortexResult<()> { let Mask::Values(mask) = Mask::from_iter([true, false, true, false]) else { panic!("a mixed mask must have mask values"); }; - let expected = [0u8, 1, 2, 6, 7, 8]; - - // A uniquely owned buffer takes the in-place `copy_within` path. - let owned = Buffer::from_iter(0u8..12); - let filtered = filter_records(owned, 3, &mask); - assert_eq!(filtered.as_slice(), &expected); - - // Retaining a second reference forces the copying path instead. - let shared = Buffer::from_iter(0u8..12); - let _retained = shared.clone(); - let filtered = filter_records(shared, 3, &mask); - assert_eq!(filtered.as_slice(), &expected); + let byte_width = V::byte_width(array.as_view()); + let values = V::values::(array.as_view()).aligned(Alignment::new(64)); + let expected = values[..byte_width] + .iter() + .chain(&values[2 * byte_width..3 * byte_width]) + .copied() + .collect::>(); + let array = with_values(array.as_view(), values, array.len(), Validity::NonNullable)?; + let alignment = V::values::(array.as_view()).alignment(); + let filtered = super::filter(&array, &mask).into_array(); + let buffers = filtered.buffers(); + assert_eq!(buffers[0].as_slice(), expected); + assert_eq!(buffers[0].alignment(), alignment); + Ok(()) } #[rstest] diff --git a/vortex-array/src/arrays/fixed_width/mod.rs b/vortex-array/src/arrays/fixed_width/mod.rs index ec19f05413c..c0f39120fcc 100644 --- a/vortex-array/src/arrays/fixed_width/mod.rs +++ b/vortex-array/src/arrays/fixed_width/mod.rs @@ -10,38 +10,3 @@ pub(crate) mod vtable; pub(crate) use self::array::FixedWidthArray; pub(crate) use self::array::with_values; - -/// Dispatches a runtime byte width to a compile-time `const $W: usize` for every record width -/// with a dedicated fixed-width kernel, falling back to `$fallback` for any other width. -macro_rules! match_each_record_width { - ($byte_width:expr, | $W:ident | $body:block,_ => $fallback:block) => { - match $byte_width { - 1 => { - const $W: usize = 1; - $body - } - 2 => { - const $W: usize = 2; - $body - } - 4 => { - const $W: usize = 4; - $body - } - 8 => { - const $W: usize = 8; - $body - } - 16 => { - const $W: usize = 16; - $body - } - 32 => { - const $W: usize = 32; - $body - } - _ => $fallback, - } - }; -} -pub(crate) use match_each_record_width; diff --git a/vortex-array/src/arrays/fixed_width/take/mod.rs b/vortex-array/src/arrays/fixed_width/take/mod.rs index d6e2f744bdd..20a4bfdbdc0 100644 --- a/vortex-array/src/arrays/fixed_width/take/mod.rs +++ b/vortex-array/src/arrays/fixed_width/take/mod.rs @@ -15,9 +15,10 @@ use std::sync::LazyLock; use vortex_buffer::Buffer; use vortex_error::VortexResult; use vortex_error::vortex_bail; +use vortex_error::vortex_ensure; use vortex_mask::Mask; -use self::records::take_byte_records; +use self::records::take_records; use self::scalar::take_values_scalar; use self::slices::take_slices; use self::slices::take_slices_constant_length; @@ -38,6 +39,7 @@ use crate::builtins::ArrayBuiltins; use crate::dtype::DType; use crate::dtype::UnsignedPType; use crate::dtype::half::f16; +use crate::dtype::i256; use crate::match_each_unsigned_integer_ptype; use crate::scalar::Scalar; @@ -71,7 +73,9 @@ macro_rules! impl_fixed_width_take_value { }; } -impl_fixed_width_take_value!(u8, u16, u32, u64, i8, i16, i32, i64, f16, f32, f64,); +impl_fixed_width_take_value!( + u8, u16, u32, u64, u128, i8, i16, i32, i64, i256, f16, f32, f64, +); // SAFETY: Byte arrays have no padding and every byte is initialized. unsafe impl FixedWidthTakeValue for [u8; N] {} @@ -96,10 +100,19 @@ pub(crate) fn take( indices: &ArrayRef, ctx: &mut ExecutionCtx, ) -> VortexResult> { - if let Some(piecewise_indices) = indices.as_opt::() - && let Some(taken) = take_contiguous_ranges(array, piecewise_indices, indices, ctx)? - { - return Ok(Some(taken)); + if let Some(piecewise_indices) = indices.as_opt::() { + let taken = match V::byte_width(array) { + 1 => take_contiguous_ranges::(array, piecewise_indices, indices, ctx)?, + 2 => take_contiguous_ranges::(array, piecewise_indices, indices, ctx)?, + 4 => take_contiguous_ranges::(array, piecewise_indices, indices, ctx)?, + 8 => take_contiguous_ranges::(array, piecewise_indices, indices, ctx)?, + 16 => take_contiguous_ranges::(array, piecewise_indices, indices, ctx)?, + 32 => take_contiguous_ranges::(array, piecewise_indices, indices, ctx)?, + _ => None, + }; + if taken.is_some() { + return Ok(taken); + } } let DType::Primitive(ptype, nullability) = indices.dtype() else { @@ -135,21 +148,23 @@ pub(crate) fn take( .take(&indices.clone().into_array())? .and(indices_validity)?; - let source = V::values(array); - let values = match_each_unsigned_integer_ptype!(indices.ptype(), |I| { - take_byte_records( - &source, - V::byte_width(array), - array.len(), - indices.as_slice::(), - ) - })?; + let values = match V::byte_width(array) { + 1 => take_records::(array, &indices)?, + 2 => take_records::(array, &indices)?, + 4 => take_records::(array, &indices)?, + 8 => take_records::(array, &indices)?, + 16 => take_records::(array, &indices)?, + 32 => take_records::(array, &indices)?, + _ => return Ok(None), + }; Ok(Some( with_values(array, values, indices.len(), validity)?.into_array(), )) } -fn take_contiguous_ranges( +// Avoid duplicating the starts/lengths dispatch in every record-width arm of `take`. +#[inline(never)] +fn take_contiguous_ranges( array: ArrayView<'_, V>, indices: ArrayView<'_, PiecewiseSequence>, indices_ref: &ArrayRef, @@ -159,21 +174,17 @@ fn take_contiguous_ranges( return Ok(None); }; - let values = V::values(array); - let byte_width = V::byte_width(array); + let values = V::values::(array); + vortex_ensure!( + values.len() == array.len(), + "Fixed-width values buffer length does not match record count" + ); let output_len = indices_ref.len(); let taken = match lengths { Columnar::Constant(lengths) => { let length = constant_unsigned_usize(&lengths); match_each_unsigned_integer_ptype!(starts.ptype(), |S| { - take_slices_constant_length( - &values, - byte_width, - array.len(), - starts.as_slice::(), - length, - output_len, - ) + take_slices_constant_length(&values, starts.as_slice::(), length, output_len) }) } Columnar::Canonical(lengths) => { @@ -182,8 +193,6 @@ fn take_contiguous_ranges( match_each_unsigned_integer_ptype!(lengths.ptype(), |L| { take_slices( &values, - byte_width, - array.len(), starts.as_slice::(), lengths.as_slice::(), output_len, @@ -194,6 +203,6 @@ fn take_contiguous_ranges( }?; let validity = array.validity()?.take(indices_ref)?; Ok(Some( - with_values(array, taken, output_len, validity)?.into_array(), + with_values(array, taken.into_byte_buffer(), output_len, validity)?.into_array(), )) } diff --git a/vortex-array/src/arrays/fixed_width/take/records.rs b/vortex-array/src/arrays/fixed_width/take/records.rs index 6cb47957991..e18c72e1113 100644 --- a/vortex-array/src/arrays/fixed_width/take/records.rs +++ b/vortex-array/src/arrays/fixed_width/take/records.rs @@ -1,49 +1,30 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -use vortex_buffer::Buffer; -use vortex_buffer::BufferMut; use vortex_buffer::ByteBuffer; use vortex_error::VortexResult; -use vortex_error::vortex_err; +use vortex_error::vortex_ensure; +use super::FixedWidthTakeValue; use super::take_values; -use crate::arrays::fixed_width::match_each_record_width; -use crate::dtype::UnsignedPType; +use crate::array::ArrayView; +use crate::arrays::PrimitiveArray; +use crate::arrays::fixed_width::FixedWidthArray; +use crate::match_each_unsigned_integer_ptype; -pub(super) fn take_byte_records( - values: &ByteBuffer, - byte_width: usize, - record_count: usize, - indices: &[I], +// Avoid duplicating the indices dispatch in every record-width arm of `take`. +#[inline(never)] +pub(super) fn take_records( + array: ArrayView<'_, V>, + indices: &PrimitiveArray, ) -> VortexResult { - let alignment = values.alignment(); - - match_each_record_width!( - byte_width, - |W| { - let records = Buffer::<[u8; W]>::from_byte_buffer(values.clone()); - debug_assert_eq!(records.len(), record_count); - Ok(take_values(records.as_slice(), indices) - .into_byte_buffer() - .aligned(alignment)) - }, - _ => { - let output_len = indices - .len() - .checked_mul(byte_width) - .ok_or_else(|| vortex_err!("Fixed-width take output length overflows usize"))?; - let mut result = BufferMut::::with_capacity(output_len); - for index in indices { - let index = index.as_(); - assert!( - index < record_count, - "take index {index} out of bounds for length {record_count}" - ); - let start = index * byte_width; - result.extend_from_slice(&values[start..start + byte_width]); - } - Ok(result.freeze().into_byte_buffer().aligned(alignment)) - } - ) + let values = V::values::(array); + vortex_ensure!( + values.len() == array.len(), + "Fixed-width values buffer length does not match record count" + ); + let taken = match_each_unsigned_integer_ptype!(indices.ptype(), |I| { + take_values(values.as_slice(), indices.as_slice::()) + }); + Ok(taken.into_byte_buffer().aligned(values.alignment())) } diff --git a/vortex-array/src/arrays/fixed_width/take/slices.rs b/vortex-array/src/arrays/fixed_width/take/slices.rs index 28efe82b625..190f32b6069 100644 --- a/vortex-array/src/arrays/fixed_width/take/slices.rs +++ b/vortex-array/src/arrays/fixed_width/take/slices.rs @@ -2,37 +2,33 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use itertools::Itertools as _; +use vortex_buffer::Buffer; use vortex_buffer::BufferMut; -use vortex_buffer::ByteBuffer; use vortex_error::VortexResult; use vortex_error::vortex_ensure; use vortex_error::vortex_err; use crate::dtype::UnsignedPType; -pub(super) fn take_slices( - values: &ByteBuffer, - byte_width: usize, - record_count: usize, +pub(super) fn take_slices( + values: &Buffer, starts: &[S], lengths: &[L], output_len: usize, -) -> VortexResult { +) -> VortexResult> { let slices = starts .iter() .zip_eq(lengths) .map(|(&start, &length)| (start.as_(), length.as_())); - copy_slices(values, byte_width, record_count, slices, output_len) + copy_slices(values, slices, output_len) } -pub(super) fn take_slices_constant_length( - values: &ByteBuffer, - byte_width: usize, - record_count: usize, +pub(super) fn take_slices_constant_length( + values: &Buffer, starts: &[S], length: usize, output_len: usize, -) -> VortexResult { +) -> VortexResult> { let computed_len = starts .len() .checked_mul(length) @@ -43,34 +39,25 @@ pub(super) fn take_slices_constant_length( ); copy_slices( values, - byte_width, - record_count, starts.iter().map(|start| (start.as_(), length)), output_len, ) } -fn copy_slices( - values: &ByteBuffer, - byte_width: usize, - record_count: usize, +// Keeping this kernel separate improves the take_fsl benchmark's small-range copies. +#[inline(never)] +fn copy_slices( + values: &Buffer, slices: impl IntoIterator, output_len: usize, -) -> VortexResult { - let input_byte_len = record_count - .checked_mul(byte_width) - .ok_or_else(|| vortex_err!("Fixed-width values buffer length overflows usize"))?; - vortex_ensure!( - values.len() == input_byte_len, - "Fixed-width values buffer length does not match record count" - ); - - let output_byte_len = output_len - .checked_mul(byte_width) +) -> VortexResult> { + output_len + .checked_mul(size_of::()) .ok_or_else(|| vortex_err!("PiecewiseSequenceArray output length overflows usize"))?; - let mut result = BufferMut::::with_capacity_aligned(output_byte_len, values.alignment()); - let spare = &mut result.spare_capacity_mut()[..output_byte_len]; + let mut result = BufferMut::::with_capacity_aligned(output_len, values.alignment()); + let spare = &mut result.spare_capacity_mut()[..output_len]; let mut cursor = 0usize; + let record_count = values.len(); for (start, length) in slices { let end = start @@ -80,21 +67,17 @@ fn copy_slices( end <= record_count, "PiecewiseSequenceArray slice {start}..{end} exceeds array length {record_count}" ); - // These multiplications cannot overflow because `end <= record_count` and the complete - // values buffer length was checked above. - let byte_start = start * byte_width; - let byte_length = length * byte_width; - let source = &values[byte_start..][..byte_length]; - spare[cursor..][..source.len()].write_copy_of_slice(source); - cursor += source.len(); + let source = &values[start..end]; + spare[cursor..][..length].write_copy_of_slice(source); + cursor += length; } // SAFETY: The loop initialized the prefix `0..cursor` of the spare capacity. unsafe { result.set_len(cursor) }; vortex_ensure!( - result.len() == output_byte_len, - "PiecewiseSequenceArray expanded length {} does not match declared length {output_byte_len}", + result.len() == output_len, + "PiecewiseSequenceArray expanded length {} does not match declared length {output_len}", result.len() ); - Ok(result.freeze().into_byte_buffer()) + Ok(result.freeze()) } diff --git a/vortex-array/src/arrays/fixed_width/take/tests.rs b/vortex-array/src/arrays/fixed_width/take/tests.rs index ff167d27499..8cdd6efed46 100644 --- a/vortex-array/src/arrays/fixed_width/take/tests.rs +++ b/vortex-array/src/arrays/fixed_width/take/tests.rs @@ -1,27 +1,33 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use std::fmt::Debug; + use rstest::rstest; +use vortex_buffer::Alignment; use vortex_buffer::Buffer; use vortex_buffer::buffer; use vortex_error::VortexResult; -use super::records::take_byte_records; use super::slices::take_slices; use super::slices::take_slices_constant_length; use super::take_values; use crate::ArrayRef; use crate::IntoArray; use crate::VortexSessionExecute; +use crate::array::Array; use crate::array_session; use crate::arrays::BoolArray; use crate::arrays::ConstantArray; use crate::arrays::DecimalArray; use crate::arrays::PiecewiseSequenceArray; use crate::arrays::PrimitiveArray; +use crate::arrays::fixed_width::FixedWidthArray; +use crate::arrays::fixed_width::with_values; use crate::assert_arrays_eq; use crate::compute::conformance::take::test_take_conformance; use crate::dtype::DecimalDType; +use crate::dtype::half::f16; use crate::dtype::i256; use crate::validity::Validity; @@ -39,53 +45,83 @@ fn take_eight_byte_values() { } #[rstest] -#[case(1)] -#[case(2)] -#[case(4)] -#[case(8)] -#[case(16)] -#[case(32)] -#[case::fallback(3)] -#[case::fallback_wide(12)] -fn take_runtime_width_records(#[case] byte_width: usize) -> VortexResult<()> { - let values = Buffer::from_iter((0u8..).take(3 * byte_width)); +#[case::u8(PrimitiveArray::from_iter([10u8, 20, 30]))] +#[case::u16(PrimitiveArray::from_iter([10u16, 20, 30]))] +#[case::u32(PrimitiveArray::from_iter([10u32, 20, 30]))] +#[case::u64(PrimitiveArray::from_iter([10u64, 20, 30]))] +#[case::u128(DecimalArray::new( + buffer![10i128, -20, -30], + DecimalDType::new(19, 0), + Validity::NonNullable, +))] +#[case::i256(DecimalArray::new( + Buffer::from_iter([10, -20, -30].map(|value| i256::from_parts(1, value))), + DecimalDType::new(76, 0), + Validity::NonNullable, +))] +fn take_typed_records( + #[case] array: Array, + #[values( + PrimitiveArray::from_iter([2u8, 0]), + PrimitiveArray::from_iter([2u16, 0]), + PrimitiveArray::from_iter([2u32, 0]), + PrimitiveArray::from_iter([2u64, 0]) + )] + indices: PrimitiveArray, +) -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let byte_width = V::byte_width(array.as_view()); + let values = V::values::(array.as_view()).aligned(Alignment::new(64)); let expected = values[2 * byte_width..3 * byte_width] .iter() .chain(&values[..byte_width]) .copied() .collect::>(); - let taken = take_byte_records(&values.into_byte_buffer(), byte_width, 3, &[2u32, 0])?; - assert_eq!(taken.as_slice(), expected); + let array = with_values(array.as_view(), values, array.len(), Validity::NonNullable)?; + let alignment = V::values::(array.as_view()).alignment(); + let taken = super::take(array.as_view(), &indices.into_array(), &mut ctx)? + .expect("native-width take must be supported"); + let buffers = taken.buffers(); + assert_eq!(buffers[0].as_slice(), expected); + assert_eq!(buffers[0].alignment(), alignment); Ok(()) } #[test] -#[should_panic(expected = "take index 3 out of bounds for length 3")] -fn fallback_take_rejects_out_of_bounds_index() { - let values = Buffer::from_iter((0u8..).take(9)).into_byte_buffer(); - drop(take_byte_records(&values, 3, 3, &[3u32])); +#[should_panic(expected = "out of bounds")] +fn typed_take_rejects_out_of_bounds_index() { + let mut ctx = array_session().create_execution_ctx(); + let array = PrimitiveArray::from_iter([10u32, 20, 30]); + let indices = PrimitiveArray::from_iter([3u32]).into_array(); + drop(super::take(array.as_view(), &indices, &mut ctx)); } -#[test] -fn take_variable_length_slices() -> VortexResult<()> { - let values = buffer![10u8, 11, 12, 13, 14].into_byte_buffer(); - let taken = take_slices(&values, 1, 5, &[1u32, 3], &[2u32, 1], 3)?; - assert_eq!(taken.as_slice(), &[11, 12, 13]); +#[rstest] +#[case::u8(buffer![10u8, 11, 12, 13, 14])] +#[case::u16(buffer![10u16, 11, 12, 13, 14])] +#[case::u32(buffer![10u32, 11, 12, 13, 14])] +#[case::u64(buffer![10u64, 11, 12, 13, 14])] +#[case::u128(buffer![10u128, 11, 12, 13, 14])] +#[case::i256(Buffer::from_iter((10..15).map(i256::from_i128)))] +#[case::three_bytes(buffer![[10u8; 3], [11; 3], [12; 3], [13; 3], [14; 3]])] +#[case::twelve_bytes(buffer![[10u8; 12], [11; 12], [12; 12], [13; 12], [14; 12]])] +fn take_typed_slices(#[case] values: Buffer) -> VortexResult<()> { + let values = values.aligned(Alignment::new(64)); + let taken = take_slices(&values, &[1u32, 3], &[2u32, 1], 3)?; + assert_eq!(taken.as_slice(), &values[1..4]); + assert_eq!(taken.alignment(), values.alignment()); + + let taken = take_slices_constant_length(&values, &[0u32, 3], 2, 4)?; + let expected = [&values[..2], &values[3..]].concat(); + assert_eq!(taken.as_slice(), expected); + assert_eq!(taken.alignment(), values.alignment()); Ok(()) } #[test] fn variable_length_slices_validate_output_length() { - let values = buffer![10u8, 11, 12, 13].into_byte_buffer(); - assert!(take_slices(&values, 1, 4, &[0u32, 2], &[1u32, 1], 3).is_err()); -} - -#[test] -fn take_constant_length_slices() -> VortexResult<()> { - let values = buffer![10u8, 11, 12, 13, 14].into_byte_buffer(); - let taken = take_slices_constant_length(&values, 1, 5, &[0u32, 3], 2, 4)?; - assert_eq!(taken.as_slice(), &[10, 11, 13, 14]); - Ok(()) + let values = buffer![10u8, 11, 12, 13]; + assert!(take_slices(&values, &[0u32, 2], &[1u32, 1], 3).is_err()); } #[test] @@ -131,38 +167,39 @@ fn null_index_skips_out_of_bounds_decimal_value() -> VortexResult<()> { Ok(()) } -#[test] -fn decimal_i256_take_consumes_piecewise_indices() -> VortexResult<()> { +#[rstest] +#[case::u8(PrimitiveArray::from_iter([10u8, 20, 30, 40, 50]).into_array())] +#[case::f16(PrimitiveArray::from_iter([10.0, -20.0, 30.0, -40.0, 50.0].map(f16::from_f32)).into_array())] +#[case::f32(PrimitiveArray::from_iter([10f32, -20.0, 30.0, -40.0, 50.0]).into_array())] +#[case::f64(PrimitiveArray::from_iter([10f64, -20.0, 30.0, -40.0, 50.0]).into_array())] +#[case::decimal_i128(DecimalArray::new( + buffer![100i128, -200, 300, -400, 500], + DecimalDType::new(19, 2), + Validity::NonNullable, +).into_array())] +#[case::decimal_i256(DecimalArray::new( + Buffer::from_iter([100, -200, 300, -400, 500].map(i256::from_i128)), + DecimalDType::new(76, 2), + Validity::NonNullable, +).into_array())] +fn fixed_width_take_consumes_piecewise_indices( + #[case] values: ArrayRef, + #[values(false, true)] constant_length: bool, +) -> VortexResult<()> { let mut ctx = array_session().create_execution_ctx(); - let decimal_dtype = DecimalDType::new(76, 2); - let values = DecimalArray::new( - buffer![ - i256::from_i128(100), - i256::from_i128(200), - i256::from_i128(300), - i256::from_i128(400), - i256::from_i128(500), - ], - decimal_dtype, - Validity::NonNullable, - ); let starts = PrimitiveArray::from_iter([1u64, 3]).into_array(); - let lengths = PrimitiveArray::from_iter([2u64, 1]).into_array(); + let (lengths, output_len) = if constant_length { + (ConstantArray::new(2u64, 2).into_array(), 4) + } else { + (PrimitiveArray::from_iter([2u64, 1]).into_array(), 3) + }; let multipliers = ConstantArray::new(1u64, 2).into_array(); - let indices = PiecewiseSequenceArray::try_new(starts, lengths, multipliers, 3)?.into_array(); + let indices = + PiecewiseSequenceArray::try_new(starts, lengths, multipliers, output_len)?.into_array(); let taken = values.take(indices)?; - let expected = DecimalArray::new( - buffer![ - i256::from_i128(200), - i256::from_i128(300), - i256::from_i128(400), - ], - decimal_dtype, - Validity::NonNullable, - ); - assert_arrays_eq!(taken, expected, &mut ctx); + assert_arrays_eq!(taken, values.slice(1..1 + output_len)?, &mut ctx); Ok(()) } diff --git a/vortex-array/src/arrays/primitive/compute/fixed_width.rs b/vortex-array/src/arrays/primitive/compute/fixed_width.rs index 6f3202ecc3c..d829f6e07e6 100644 --- a/vortex-array/src/arrays/primitive/compute/fixed_width.rs +++ b/vortex-array/src/arrays/primitive/compute/fixed_width.rs @@ -1,6 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_buffer::Buffer; use vortex_buffer::ByteBuffer; use vortex_error::VortexResult; @@ -15,8 +16,10 @@ impl FixedWidthArray for Primitive { array.ptype().byte_width() } - fn values(array: ArrayView<'_, Self>) -> ByteBuffer { - array.buffer_handle().to_host_sync() + fn values(array: ArrayView<'_, Self>) -> Buffer { + let values = array.buffer_handle().to_host_sync(); + let alignment = values.alignment(); + Buffer::from_byte_buffer_aligned(values, alignment) } fn with_values(