From 152c18fd27e2267affc7b72a2be202c61c60c77b Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Thu, 23 Jul 2026 13:55:24 -0400 Subject: [PATCH 1/2] Share fixed-width VTable plumbing Signed-off-by: Connor Tsui --- vortex-array/src/arrays/decimal/vtable/mod.rs | 36 +++----------- vortex-array/src/arrays/fixed_width/mod.rs | 6 +++ vortex-array/src/arrays/fixed_width/vtable.rs | 47 +++++++++++++++++++ vortex-array/src/arrays/mod.rs | 2 + .../src/arrays/primitive/vtable/mod.rs | 45 ++++-------------- 5 files changed, 70 insertions(+), 66 deletions(-) create mode 100644 vortex-array/src/arrays/fixed_width/mod.rs create mode 100644 vortex-array/src/arrays/fixed_width/vtable.rs diff --git a/vortex-array/src/arrays/decimal/vtable/mod.rs b/vortex-array/src/arrays/decimal/vtable/mod.rs index ec157b28c66..aa51a0d0cb2 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; @@ -85,17 +84,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 +96,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::single_buffer(buffers)?; Ok( ArrayParts::new(self.clone(), array.dtype().clone(), array.len(), data) .with_slots(array.slots().iter().cloned().collect()), @@ -169,19 +157,9 @@ impl VTable for Decimal { _session: &VortexSession, ) -> VortexResult> { let metadata = DecimalMetadata::decode(metadata)?; - if buffers.len() != 1 { - vortex_bail!("Expected 1 buffer, got {}", buffers.len()); - } - 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 values = fixed_width::single_buffer(buffers)?; + + 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) 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..61ee10a29bc --- /dev/null +++ b/vortex-array/src/arrays/fixed_width/mod.rs @@ -0,0 +1,6 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Shared structural operations for fixed-width canonical arrays. + +pub(crate) mod vtable; 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..79d6ed20072 --- /dev/null +++ b/vortex-array/src/arrays/fixed_width/vtable.rs @@ -0,0 +1,47 @@ +// 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::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 single_buffer(buffers: &[BufferHandle]) -> VortexResult { + vortex_ensure!( + buffers.len() == 1, + "Expected 1 buffer, got {}", + buffers.len() + ); + Ok(buffers[0].clone()) +} + +pub(crate) fn deserialize_validity( + nullability: Nullability, + len: usize, + children: &dyn ArrayChildren, +) -> VortexResult { + match children.len() { + 0 => Ok(Validity::from(nullability)), + 1 => Ok(Validity::Array(children.get(0, &Validity::DTYPE, len)?)), + child_count => vortex_bail!("Expected 0 or 1 child, got {child_count}"), + } +} diff --git a/vortex-array/src/arrays/mod.rs b/vortex-array/src/arrays/mod.rs index 7c31a1b48b1..045f7c505b6 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/vtable/mod.rs b/vortex-array/src/arrays/primitive/vtable/mod.rs index 50189bf0037..e0e095c58e0 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; @@ -74,17 +73,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 +85,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::single_buffer(buffers)?; Ok( ArrayParts::new(self.clone(), array.dtype().clone(), array.len(), data) .with_slots(array.slots().iter().cloned().collect()), @@ -157,19 +145,9 @@ impl VTable for Primitive { metadata.len() ); } - if buffers.len() != 1 { - vortex_bail!("Expected 1 buffer, got {}", buffers.len()); - } - 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 buffer = fixed_width::single_buffer(buffers)?; + + let validity = fixed_width::deserialize_validity(dtype.nullability(), len, children)?; let ptype = PType::try_from(dtype)?; @@ -188,14 +166,7 @@ impl VTable for Primitive { ); } - vortex_ensure!( - buffer.is_aligned_to(Alignment::new(ptype.byte_width())), - "PrimitiveArray::build: Buffer (align={}) must be aligned to {}", - buffer.alignment(), - ptype.byte_width() - ); - - // SAFETY: checked ahead of time + // SAFETY: the buffer length and alignment are checked above. let slots = PrimitiveData::make_slots(&validity, len); let data = unsafe { PrimitiveData::new_unchecked_from_handle(buffer, ptype, validity) }; Ok(ArrayParts::new(self.clone(), dtype.clone(), len, data).with_slots(slots)) From 2d5d094ea290772cfb13425f739f98f9f6098eff Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Thu, 23 Jul 2026 13:55:49 -0400 Subject: [PATCH 2/2] Share fixed-width compute kernels Signed-off-by: Connor Tsui --- .../src/arrays/decimal/compute/fixed_width.rs | 45 ++ .../src/arrays/decimal/compute/mod.rs | 2 +- .../src/arrays/decimal/compute/take.rs | 357 ------------ .../src/arrays/filter/execute/buffer.rs | 2 +- .../arrays/filter/execute/byte_compress.rs | 2 +- .../src/arrays/filter/execute/decimal.rs | 70 --- vortex-array/src/arrays/filter/execute/mod.rs | 11 +- .../src/arrays/filter/execute/primitive.rs | 116 ---- vortex-array/src/arrays/filter/mod.rs | 3 + vortex-array/src/arrays/fixed_width/array.rs | 24 + vortex-array/src/arrays/fixed_width/filter.rs | 75 +++ .../src/arrays/fixed_width/filter/tests.rs | 56 ++ vortex-array/src/arrays/fixed_width/mod.rs | 5 + .../arrays/fixed_width/take/avx2/gather.rs | 215 +++++++ .../src/arrays/fixed_width/take/avx2/mod.rs | 141 +++++ .../src/arrays/fixed_width/take/avx2/tests.rs | 114 ++++ .../src/arrays/fixed_width/take/mod.rs | 198 +++++++ .../src/arrays/fixed_width/take/records.rs | 56 ++ .../src/arrays/fixed_width/take/scalar.rs | 27 + .../src/arrays/fixed_width/take/slices.rs | 110 ++++ .../src/arrays/fixed_width/take/tests.rs | 196 +++++++ .../arrays/primitive/compute/fixed_width.rs | 43 ++ .../src/arrays/primitive/compute/mod.rs | 2 +- .../src/arrays/primitive/compute/take/avx2.rs | 536 ------------------ .../src/arrays/primitive/compute/take/mod.rs | 453 --------------- 25 files changed, 1317 insertions(+), 1542 deletions(-) create mode 100644 vortex-array/src/arrays/decimal/compute/fixed_width.rs delete mode 100644 vortex-array/src/arrays/decimal/compute/take.rs delete mode 100644 vortex-array/src/arrays/filter/execute/decimal.rs delete mode 100644 vortex-array/src/arrays/filter/execute/primitive.rs create mode 100644 vortex-array/src/arrays/fixed_width/array.rs create mode 100644 vortex-array/src/arrays/fixed_width/filter.rs create mode 100644 vortex-array/src/arrays/fixed_width/filter/tests.rs create mode 100644 vortex-array/src/arrays/fixed_width/take/avx2/gather.rs create mode 100644 vortex-array/src/arrays/fixed_width/take/avx2/mod.rs create mode 100644 vortex-array/src/arrays/fixed_width/take/avx2/tests.rs create mode 100644 vortex-array/src/arrays/fixed_width/take/mod.rs create mode 100644 vortex-array/src/arrays/fixed_width/take/records.rs create mode 100644 vortex-array/src/arrays/fixed_width/take/scalar.rs create mode 100644 vortex-array/src/arrays/fixed_width/take/slices.rs create mode 100644 vortex-array/src/arrays/fixed_width/take/tests.rs create mode 100644 vortex-array/src/arrays/primitive/compute/fixed_width.rs delete mode 100644 vortex-array/src/arrays/primitive/compute/take/avx2.rs delete mode 100644 vortex-array/src/arrays/primitive/compute/take/mod.rs diff --git a/vortex-array/src/arrays/decimal/compute/fixed_width.rs b/vortex-array/src/arrays/decimal/compute/fixed_width.rs new file mode 100644 index 00000000000..85ed44e482d --- /dev/null +++ b/vortex-array/src/arrays/decimal/compute/fixed_width.rs @@ -0,0 +1,45 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_buffer::ByteBuffer; +use vortex_error::VortexResult; +use vortex_error::vortex_ensure; +use vortex_error::vortex_err; + +use crate::array::ArrayView; +use crate::arrays::Decimal; +use crate::arrays::DecimalArray; +use crate::arrays::fixed_width::FixedWidthArray; +use crate::buffer::BufferHandle; +use crate::validity::Validity; + +impl FixedWidthArray for Decimal { + fn byte_width(array: ArrayView<'_, Self>) -> usize { + array.values_type().byte_width() + } + + fn values(array: ArrayView<'_, Self>) -> ByteBuffer { + array.buffer_handle().to_host_sync() + } + + fn with_values( + array: ArrayView<'_, Self>, + values: ByteBuffer, + len: usize, + validity: Validity, + ) -> VortexResult { + let expected_len = len + .checked_mul(array.values_type().byte_width()) + .ok_or_else(|| vortex_err!("Decimal values buffer length overflows usize"))?; + vortex_ensure!( + values.len() == expected_len, + "Decimal values buffer length does not match output length" + ); + DecimalArray::try_new_handle( + BufferHandle::new_host(values), + array.values_type(), + array.decimal_dtype(), + validity, + ) + } +} diff --git a/vortex-array/src/arrays/decimal/compute/mod.rs b/vortex-array/src/arrays/decimal/compute/mod.rs index 0c2f32b4c32..a6e1cfb2848 100644 --- a/vortex-array/src/arrays/decimal/compute/mod.rs +++ b/vortex-array/src/arrays/decimal/compute/mod.rs @@ -4,9 +4,9 @@ mod between; mod cast; mod fill_null; +mod fixed_width; mod mask; pub mod rules; -mod take; #[cfg(test)] mod tests { diff --git a/vortex-array/src/arrays/decimal/compute/take.rs b/vortex-array/src/arrays/decimal/compute/take.rs deleted file mode 100644 index d7bde288261..00000000000 --- a/vortex-array/src/arrays/decimal/compute/take.rs +++ /dev/null @@ -1,357 +0,0 @@ -// 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::Decimal; -use crate::arrays::DecimalArray; -use crate::arrays::PiecewiseSequence; -use crate::arrays::PrimitiveArray; -use crate::arrays::dict::TakeExecute; -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::validity::Validity; - -impl TakeExecute for Decimal { - fn take( - array: ArrayView<'_, Decimal>, - 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)); - } - - let indices = indices.clone().execute::(ctx)?; - let validity = array.validity()?.take(&indices.clone().into_array())?; - - // 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()); - // 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) } - }) - }); - - Ok(Some(decimal.into_array())) - } -} - -fn take_contiguous_ranges( - array: ArrayView<'_, Decimal>, - indices: ArrayView<'_, PiecewiseSequence>, - indices_ref: &ArrayRef, - ctx: &mut ExecutionCtx, -) -> VortexResult> { - let Some((starts, lengths)) = maybe_contiguous_slices(indices, ctx)? else { - return Ok(None); - }; - let validity = array.validity()?.take(indices_ref)?; - let output_len = indices_ref.len(); - let taken = match lengths { - Columnar::Constant(lengths) => { - let length = constant_unsigned_usize(&lengths); - take_slices_constant_length(array, &starts, length, validity, output_len)? - } - Columnar::Canonical(lengths) => { - let lengths = lengths.into_primitive(); - take_slices(array, &starts, &lengths, validity, output_len)? - } - }; - Ok(Some(taken)) -} - -fn take_slices_constant_length( - array: ArrayView<'_, Decimal>, - starts: &PrimitiveArray, - length: usize, - validity: Validity, - output_len: usize, -) -> VortexResult { - match_each_decimal_value_type!(array.values_type(), |D| { - take_slices_constant_length_typed::(array, starts, length, validity, output_len) - }) -} - -fn take_slices_constant_length_typed( - array: ArrayView<'_, Decimal>, - starts: &PrimitiveArray, - length: usize, - validity: Validity, - output_len: usize, -) -> VortexResult -where - D: NativeDecimalType, -{ - match_each_unsigned_integer_ptype!(starts.ptype(), |S| { - let values = take_slices_constant_length_to_buffer::( - starts.as_slice::(), - length, - array.buffer::().as_slice(), - output_len, - )?; - - // SAFETY: contiguous gather preserves the decimal dtype and value representation. - Ok( - unsafe { DecimalArray::new_unchecked(values, array.decimal_dtype(), validity) } - .into_array(), - ) - }) -} - -fn take_slices( - array: ArrayView<'_, Decimal>, - starts: &PrimitiveArray, - lengths: &PrimitiveArray, - validity: Validity, - output_len: usize, -) -> VortexResult { - match_each_decimal_value_type!(array.values_type(), |D| { - take_slices_typed::(array, starts, lengths, validity, output_len) - }) -} - -fn take_slices_typed( - array: ArrayView<'_, Decimal>, - starts: &PrimitiveArray, - lengths: &PrimitiveArray, - validity: Validity, - output_len: usize, -) -> VortexResult -where - D: NativeDecimalType, -{ - match_each_unsigned_integer_ptype!(starts.ptype(), |S| { - take_slices_start_typed::(array, starts, lengths, validity, output_len) - }) -} - -fn take_slices_start_typed( - array: ArrayView<'_, Decimal>, - starts: &PrimitiveArray, - lengths: &PrimitiveArray, - validity: Validity, - output_len: usize, -) -> VortexResult -where - D: NativeDecimalType, - S: UnsignedPType, -{ - match_each_unsigned_integer_ptype!(lengths.ptype(), |L| { - let values = take_slices_to_buffer::( - starts.as_slice::(), - lengths.as_slice::(), - array.buffer::().as_slice(), - output_len, - )?; - - // SAFETY: contiguous gather preserves the decimal dtype and value representation. - Ok( - unsafe { DecimalArray::new_unchecked(values, array.decimal_dtype(), validity) } - .into_array(), - ) - }) -} - -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; - use vortex_buffer::Buffer; - use vortex_buffer::buffer; - - use crate::IntoArray; - use crate::VortexSessionExecute; - use crate::array_session; - use crate::arrays::DecimalArray; - use crate::arrays::PrimitiveArray; - use crate::assert_arrays_eq; - use crate::compute::conformance::take::test_take_conformance; - use crate::dtype::DecimalDType; - use crate::validity::Validity; - - #[test] - fn test_take() { - 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 = buffer![0, 2, 3].into_array(); - let taken = array.take(indices).unwrap(); - - let expected = DecimalArray::from_iter([10i128, 12, 13], ddtype); - assert_arrays_eq!(expected, taken, &mut ctx); - } - - #[test] - fn test_take_null_indices() { - let mut ctx = array_session().create_execution_ctx(); - let ddtype = DecimalDType::new(19, 1); - let array = DecimalArray::new( - buffer![i128::MAX, 11i128, 12i128, 13i128], - ddtype, - Validity::NonNullable, - ); - - let indices = PrimitiveArray::from_option_iter([None, Some(2), Some(3)]).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], - DecimalDType::new(19, 2), - Validity::NonNullable, - ))] - #[case(DecimalArray::new( - buffer![10i64, 20i64, 30i64, 40i64, 50i64], - DecimalDType::new(10, 1), - Validity::NonNullable, - ))] - #[case(DecimalArray::new( - buffer![1i32, 2i32, 3i32, 4i32, 5i32], - DecimalDType::new(5, 0), - Validity::NonNullable, - ))] - #[case(DecimalArray::new( - buffer![1000i128, 2000i128, 3000i128, 4000i128, 5000i128], - DecimalDType::new(19, 3), - Validity::from_iter([true, false, true, true, false]), - ))] - #[case(DecimalArray::new( - buffer![42i128], - DecimalDType::new(19, 0), - Validity::NonNullable, - ))] - #[case({ - let values: Vec = (0..100).map(|i| i * 1000).collect(); - DecimalArray::new( - Buffer::from_iter(values), - DecimalDType::new(19, 4), - Validity::NonNullable, - ) - })] - fn test_take_decimal_conformance(#[case] array: DecimalArray) { - test_take_conformance( - &array.into_array(), - &mut array_session().create_execution_ctx(), - ); - } -} diff --git a/vortex-array/src/arrays/filter/execute/buffer.rs b/vortex-array/src/arrays/filter/execute/buffer.rs index 73ee3026654..0e1779a4d3f 100644 --- a/vortex-array/src/arrays/filter/execute/buffer.rs +++ b/vortex-array/src/arrays/filter/execute/buffer.rs @@ -15,7 +15,7 @@ use crate::arrays::filter::execute::slice; /// /// This will attempt to filter in-place (via [`Buffer::try_into_mut`]) when the buffer has /// exclusive ownership, avoiding an extra allocation. -pub(super) fn filter_buffer(buffer: Buffer, mask: &MaskValues) -> Buffer { +pub(crate) fn filter_buffer(buffer: Buffer, mask: &MaskValues) -> Buffer { match buffer.try_into_mut() { Ok(mut buffer_mut) => { let new_len = slice::filter_slice_mut_by_mask_values(buffer_mut.as_mut_slice(), mask); diff --git a/vortex-array/src/arrays/filter/execute/byte_compress.rs b/vortex-array/src/arrays/filter/execute/byte_compress.rs index 6c58ce4f33b..e75ae1a35ae 100644 --- a/vortex-array/src/arrays/filter/execute/byte_compress.rs +++ b/vortex-array/src/arrays/filter/execute/byte_compress.rs @@ -45,7 +45,7 @@ static BYTE_COMPRESS_LUT: &[([u8; 8], u8); 256] = &{ /// /// Processes the mask one byte at a time (8 source elements per byte), /// using a precomputed permutation to compact selected elements. -pub(super) fn filter_buffer(buffer: Buffer, mask: &MaskValues) -> Buffer { +pub(crate) fn filter_buffer(buffer: Buffer, mask: &MaskValues) -> Buffer { debug_assert_eq!(buffer.len(), mask.len()); let src = buffer.as_slice(); diff --git a/vortex-array/src/arrays/filter/execute/decimal.rs b/vortex-array/src/arrays/filter/execute/decimal.rs deleted file mode 100644 index 5ffb4a963b2..00000000000 --- a/vortex-array/src/arrays/filter/execute/decimal.rs +++ /dev/null @@ -1,70 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -use std::sync::Arc; - -use vortex_error::VortexExpect; -use vortex_mask::MaskValues; - -use crate::arrays::DecimalArray; -use crate::arrays::filter::execute::buffer; -use crate::arrays::filter::execute::filter_validity; -use crate::match_each_decimal_value_type; - -pub fn filter_decimal(array: &DecimalArray, mask: &Arc) -> DecimalArray { - let filtered_validity = filter_validity( - array - .validity() - .vortex_expect("decimal validity should be derivable"), - mask, - ); - - match_each_decimal_value_type!(array.values_type(), |T| { - let filtered_buffer = buffer::filter_buffer(array.buffer::(), mask.as_ref()); - - // SAFETY: We filter both the validity and the buffer with the same mask, so they must have - // the same length, and since the buffer came from an existing and valid `DecimalArray` the - // values must all be be representable by the decimal type. - unsafe { - DecimalArray::new_unchecked(filtered_buffer, array.decimal_dtype(), filtered_validity) - } - }) -} - -#[cfg(test)] -mod test { - use crate::IntoArray; - use crate::VortexSessionExecute; - use crate::array_session; - use crate::arrays::filter::execute::decimal::DecimalArray; - use crate::compute::conformance::filter::test_filter_conformance; - use crate::dtype::DecimalDType; - - #[test] - fn test_filter_decimal128_conformance() { - let decimal_dtype = DecimalDType::new(38, 2); - let values = vec![ - Some(12345i128), - Some(67890), - Some(-12345), - Some(0), - Some(99999), - ]; - let array = DecimalArray::from_option_iter(values, decimal_dtype); - test_filter_conformance( - &array.into_array(), - &mut array_session().create_execution_ctx(), - ); - } - - #[test] - fn test_filter_decimal128_with_nulls_conformance() { - let decimal_dtype = DecimalDType::new(38, 4); - let values = vec![Some(12345i128), None, Some(-12345), Some(0), None]; - let array = DecimalArray::from_option_iter(values, decimal_dtype); - test_filter_conformance( - &array.into_array(), - &mut array_session().create_execution_ctx(), - ); - } -} diff --git a/vortex-array/src/arrays/filter/execute/mod.rs b/vortex-array/src/arrays/filter/execute/mod.rs index 048b117933e..174500a31d3 100644 --- a/vortex-array/src/arrays/filter/execute/mod.rs +++ b/vortex-array/src/arrays/filter/execute/mod.rs @@ -24,25 +24,24 @@ use crate::arrays::NullArray; use crate::arrays::VariantArray; use crate::arrays::extension::ExtensionArrayExt; use crate::arrays::filter::FilterArrayExt; +use crate::arrays::fixed_width; use crate::arrays::variant::VariantArrayExt; use crate::scalar::Scalar; use crate::validity::Validity; mod bitbuffer; mod bool; -mod buffer; +pub(crate) mod buffer; pub(crate) mod byte_compress; -mod decimal; mod fixed_size_list; mod listview; -mod primitive; mod slice; mod struct_; pub mod take; mod varbinview; /// A helper function that lazily filters a [`Validity`] with selection mask values. -fn filter_validity(validity: Validity, mask: &Arc) -> Validity { +pub(crate) fn filter_validity(validity: Validity, mask: &Arc) -> Validity { validity .filter(&Mask::Values(Arc::clone(mask))) .vortex_expect("Somehow unable to wrap filter around a validity array") @@ -87,8 +86,8 @@ pub(super) fn execute_filter(canonical: Canonical, mask: &Arc) -> Ca match canonical { Canonical::Null(_) => Canonical::Null(NullArray::new(mask.true_count())), Canonical::Bool(a) => Canonical::Bool(bool::filter_bool(&a, mask)), - Canonical::Primitive(a) => Canonical::Primitive(primitive::filter_primitive(&a, mask)), - Canonical::Decimal(a) => Canonical::Decimal(decimal::filter_decimal(&a, mask)), + Canonical::Primitive(a) => Canonical::Primitive(fixed_width::filter::filter(&a, mask)), + Canonical::Decimal(a) => Canonical::Decimal(fixed_width::filter::filter(&a, mask)), Canonical::VarBinView(a) => Canonical::VarBinView(varbinview::filter_varbinview(&a, mask)), Canonical::List(a) => Canonical::List(listview::filter_listview(&a, mask)), Canonical::FixedSizeList(a) => { diff --git a/vortex-array/src/arrays/filter/execute/primitive.rs b/vortex-array/src/arrays/filter/execute/primitive.rs deleted file mode 100644 index b7abd8efa4b..00000000000 --- a/vortex-array/src/arrays/filter/execute/primitive.rs +++ /dev/null @@ -1,116 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -use std::sync::Arc; - -use vortex_error::VortexExpect; -use vortex_mask::MaskValues; - -use crate::arrays::PrimitiveArray; -use crate::arrays::filter::execute::buffer; -use crate::arrays::filter::execute::byte_compress; -use crate::arrays::filter::execute::filter_validity; -use crate::dtype::NativePType; -use crate::dtype::PType; -use crate::match_each_native_ptype; -use crate::validity::Validity; - -pub fn filter_primitive(array: &PrimitiveArray, mask: &Arc) -> PrimitiveArray { - let validity = array - .validity() - .vortex_expect("primitive validity should be derivable"); - let filtered_validity = filter_validity(validity, mask); - - match array.ptype() { - // Byte-compress avoids materializing indices/slices and processes 8 elements per mask byte. - PType::U8 => filter_byte_compress::(array, filtered_validity, mask), - PType::I8 => filter_byte_compress::(array, filtered_validity, mask), - PType::U16 => filter_byte_compress::(array, filtered_validity, mask), - PType::I16 => filter_byte_compress::(array, filtered_validity, mask), - PType::U32 => filter_byte_compress::(array, filtered_validity, mask), - PType::I32 => filter_byte_compress::(array, filtered_validity, mask), - _ => match_each_native_ptype!(array.ptype(), |T| { - let filtered_buffer = buffer::filter_buffer(array.to_buffer::(), mask.as_ref()); - - // SAFETY: We filter both the validity and the buffer with the same mask, so they must - // have the same length. - unsafe { PrimitiveArray::new_unchecked(filtered_buffer, filtered_validity) } - }), - } -} - -fn filter_byte_compress( - array: &PrimitiveArray, - filtered_validity: Validity, - mask: &Arc, -) -> PrimitiveArray { - let filtered_buffer = byte_compress::filter_buffer(array.to_buffer::(), mask.as_ref()); - - // SAFETY: We filter both the validity and the buffer with the same mask, so they must have the - // same length. - unsafe { PrimitiveArray::new_unchecked(filtered_buffer, filtered_validity) } -} - -#[cfg(test)] -#[expect(clippy::cast_possible_truncation)] -mod test { - use itertools::Itertools; - use rstest::rstest; - use vortex_mask::Mask; - - use crate::IntoArray; - use crate::VortexSessionExecute; - use crate::array_session; - use crate::arrays::PrimitiveArray; - use crate::compute::conformance::filter::LARGE_SIZE; - use crate::compute::conformance::filter::MEDIUM_SIZE; - use crate::compute::conformance::filter::test_filter_conformance; - - #[test] - fn filter_run_variant_mixed_test() { - let mut ctx = array_session().create_execution_ctx(); - let mask = [true, true, false, true, true, true, false, true]; - let arr = PrimitiveArray::from_iter([1u32, 24, 54, 2, 3, 2, 3, 2]); - - let filtered = arr - .filter(Mask::from_iter(mask)) - .unwrap() - .execute::(&mut ctx) - .unwrap(); - assert_eq!( - filtered.len(), - mask.iter().filter(|x| **x).collect_vec().len() - ); - - let rust_arr = arr.as_slice::(); - assert_eq!( - filtered.as_slice::().to_vec(), - mask.iter() - .enumerate() - .filter(|(_idx, b)| **b) - .map(|m| rust_arr[m.0]) - .collect_vec() - ) - } - - #[rstest] - #[case(PrimitiveArray::from_iter([-2i8, -1, 0, 1, 2]))] - #[case(PrimitiveArray::from_iter([1u16, 2, 3, 4, 5]))] - #[case(PrimitiveArray::from_iter([-2i16, -1, 0, 1, 2]))] - #[case(PrimitiveArray::from_iter([1i32, 2, 3, 4, 5]))] - #[case(PrimitiveArray::from_option_iter([Some(1i32), None, Some(3), Some(4), None]))] - #[case(PrimitiveArray::from_iter([42u64]))] - #[case(PrimitiveArray::from_iter(0..MEDIUM_SIZE as i32))] - #[case(PrimitiveArray::from_option_iter( - (0..MEDIUM_SIZE).map(|i| if i % 3 == 0 { None } else { Some(i as i64) })) - )] - #[case(PrimitiveArray::from_iter(0..LARGE_SIZE as u32))] - #[case(PrimitiveArray::from_iter([0.1f32, 0.2, 0.3, 0.4, 0.5]))] - #[case(PrimitiveArray::from_option_iter([Some(1.1f64), None, Some(2.2), Some(3.3), None]))] - fn test_filter_primitive_conformance(#[case] array: PrimitiveArray) { - test_filter_conformance( - &array.into_array(), - &mut array_session().create_execution_ctx(), - ); - } -} diff --git a/vortex-array/src/arrays/filter/mod.rs b/vortex-array/src/arrays/filter/mod.rs index 9510ca92f5b..c4e9a0ea97a 100644 --- a/vortex-array/src/arrays/filter/mod.rs +++ b/vortex-array/src/arrays/filter/mod.rs @@ -8,6 +8,9 @@ pub use array::FilterDataParts; pub use vtable::FilterArray; mod execute; +pub(crate) use execute::buffer::filter_buffer; +pub(crate) use execute::byte_compress::filter_buffer as filter_buffer_byte_compress; +pub(crate) use execute::filter_validity; mod kernel; pub use kernel::FilterExecuteAdaptor; diff --git a/vortex-array/src/arrays/fixed_width/array.rs b/vortex-array/src/arrays/fixed_width/array.rs new file mode 100644 index 00000000000..980d66c3527 --- /dev/null +++ b/vortex-array/src/arrays/fixed_width/array.rs @@ -0,0 +1,24 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_buffer::ByteBuffer; +use vortex_error::VortexResult; + +use crate::array::Array; +use crate::array::ArrayView; +use crate::array::VTable; +use crate::validity::Validity; + +/// Type-specific access needed by shared fixed-width structural compute. +pub(crate) trait FixedWidthArray: VTable { + fn byte_width(array: ArrayView<'_, Self>) -> usize; + + fn values(array: ArrayView<'_, Self>) -> ByteBuffer; + + fn with_values( + array: ArrayView<'_, Self>, + values: ByteBuffer, + len: usize, + validity: Validity, + ) -> VortexResult>; +} diff --git a/vortex-array/src/arrays/fixed_width/filter.rs b/vortex-array/src/arrays/fixed_width/filter.rs new file mode 100644 index 00000000000..90edf155cc4 --- /dev/null +++ b/vortex-array/src/arrays/fixed_width/filter.rs @@ -0,0 +1,75 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::sync::Arc; + +use vortex_buffer::Buffer; +use vortex_buffer::BufferMut; +use vortex_buffer::ByteBuffer; +use vortex_error::VortexExpect; +use vortex_mask::MaskValues; + +use super::FixedWidthArray; +use crate::array::Array; +use crate::arrays::filter::filter_buffer; +use crate::arrays::filter::filter_buffer_byte_compress; +use crate::arrays::filter::filter_validity; + +#[cfg(test)] +#[expect(clippy::cast_possible_truncation)] +mod tests; + +pub(crate) fn filter(array: &Array, mask: &Arc) -> Array { + let array = array.as_view(); + let values = filter_records(V::values(array), V::byte_width(array), mask); + let validity = filter_validity( + array + .validity() + .vortex_expect("fixed-width validity should be derivable"), + mask, + ); + V::with_values(array, values, mask.true_count(), validity) + .vortex_expect("filtering fixed-width values preserves array invariants") +} + +fn filter_records(values: ByteBuffer, byte_width: usize, mask: &MaskValues) -> ByteBuffer { + let alignment = values.alignment(); + + macro_rules! filter_typed_records { + ($filter:path, $width:expr) => {{ + let records = Buffer::<[u8; $width]>::from_byte_buffer(values); + return $filter(records, mask).into_byte_buffer().aligned(alignment); + }}; + } + + match byte_width { + 1 => filter_typed_records!(filter_buffer_byte_compress, 1), + 2 => filter_typed_records!(filter_buffer_byte_compress, 2), + 4 => filter_typed_records!(filter_buffer_byte_compress, 4), + 8 => filter_typed_records!(filter_buffer, 8), + 16 => filter_typed_records!(filter_buffer, 16), + 32 => filter_typed_records!(filter_buffer, 32), + _ => {} + } + + 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) + } + } +} diff --git a/vortex-array/src/arrays/fixed_width/filter/tests.rs b/vortex-array/src/arrays/fixed_width/filter/tests.rs new file mode 100644 index 00000000000..fe8dc64f628 --- /dev/null +++ b/vortex-array/src/arrays/fixed_width/filter/tests.rs @@ -0,0 +1,56 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use rstest::rstest; +use vortex_buffer::buffer; + +use crate::ArrayRef; +use crate::IntoArray; +use crate::VortexSessionExecute; +use crate::array_session; +use crate::arrays::DecimalArray; +use crate::arrays::PrimitiveArray; +use crate::compute::conformance::filter::LARGE_SIZE; +use crate::compute::conformance::filter::MEDIUM_SIZE; +use crate::compute::conformance::filter::test_filter_conformance; +use crate::dtype::DecimalDType; +use crate::dtype::i256; +use crate::validity::Validity; + +#[rstest] +#[case::primitive_i8(PrimitiveArray::from_iter([-2i8, -1, 0, 1, 2]).into_array())] +#[case::primitive_i32(PrimitiveArray::from_iter([1i32, 2, 3, 4, 5]).into_array())] +#[case::primitive_nullable(PrimitiveArray::from_option_iter( + [Some(1i64), None, Some(3), Some(4), None], +).into_array())] +#[case::primitive_large(PrimitiveArray::from_iter(0..LARGE_SIZE as u32).into_array())] +#[case::primitive_medium(PrimitiveArray::from_iter(0..MEDIUM_SIZE as i64).into_array())] +#[case::decimal_i32(DecimalArray::new( + buffer![123i32, 456, -123, 0, 999], + DecimalDType::new(8, 2), + Validity::NonNullable, +).into_array())] +#[case::decimal_i64(DecimalArray::new( + buffer![12345i64, 67890, -12345, 0, 99999], + DecimalDType::new(18, 2), + Validity::NonNullable, +).into_array())] +#[case::decimal_i128(DecimalArray::new( + buffer![12345i128, 67890, -12345, 0, 99999], + DecimalDType::new(38, 4), + Validity::from_iter([true, false, true, true, false]), +).into_array())] +#[case::decimal_i256(DecimalArray::new( + buffer![ + i256::from_i128(12345), + i256::from_i128(67890), + i256::from_i128(-12345), + i256::ZERO, + i256::from_i128(99999), + ], + DecimalDType::new(76, 4), + Validity::NonNullable, +).into_array())] +fn fixed_width_filter_conformance(#[case] array: ArrayRef) { + test_filter_conformance(&array, &mut array_session().create_execution_ctx()); +} diff --git a/vortex-array/src/arrays/fixed_width/mod.rs b/vortex-array/src/arrays/fixed_width/mod.rs index 61ee10a29bc..bfb588a16b5 100644 --- a/vortex-array/src/arrays/fixed_width/mod.rs +++ b/vortex-array/src/arrays/fixed_width/mod.rs @@ -3,4 +3,9 @@ //! Shared structural operations for fixed-width canonical arrays. +mod array; +pub(crate) mod filter; +pub(crate) mod take; pub(crate) mod vtable; + +pub(crate) use self::array::FixedWidthArray; diff --git a/vortex-array/src/arrays/fixed_width/take/avx2/gather.rs b/vortex-array/src/arrays/fixed_width/take/avx2/gather.rs new file mode 100644 index 00000000000..3b051dc3039 --- /dev/null +++ b/vortex-array/src/arrays/fixed_width/take/avx2/gather.rs @@ -0,0 +1,215 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::arch::x86_64::_mm_loadu_si128; +use std::arch::x86_64::_mm_setzero_si128; +use std::arch::x86_64::_mm_shuffle_epi32; +use std::arch::x86_64::_mm_storeu_si128; +use std::arch::x86_64::_mm_unpacklo_epi64; +use std::arch::x86_64::_mm256_andnot_si256; +use std::arch::x86_64::_mm256_cmpgt_epi32; +use std::arch::x86_64::_mm256_cmpgt_epi64; +use std::arch::x86_64::_mm256_cvtepu8_epi32; +use std::arch::x86_64::_mm256_cvtepu8_epi64; +use std::arch::x86_64::_mm256_cvtepu16_epi32; +use std::arch::x86_64::_mm256_cvtepu16_epi64; +use std::arch::x86_64::_mm256_cvtepu32_epi64; +use std::arch::x86_64::_mm256_extracti128_si256; +use std::arch::x86_64::_mm256_loadu_si256; +use std::arch::x86_64::_mm256_mask_i32gather_epi32; +use std::arch::x86_64::_mm256_mask_i64gather_epi32; +use std::arch::x86_64::_mm256_mask_i64gather_epi64; +use std::arch::x86_64::_mm256_movemask_epi8; +use std::arch::x86_64::_mm256_set1_epi32; +use std::arch::x86_64::_mm256_set1_epi64x; +use std::arch::x86_64::_mm256_setzero_si256; +use std::arch::x86_64::_mm256_storeu_si256; +use std::convert::identity; + +/// Moves one SIMD block of fixed-width values into an output buffer. +pub(super) trait GatherFn { + /// The number of data elements written on each iteration. + const WIDTH: usize; + /// The number of indices read on each iteration. + const STRIDE: usize = Self::WIDTH; + + /// Gather values from `src` into `dst`. + /// + /// # Safety + /// + /// `indices` must be readable for `STRIDE` elements. `src` must contain `max_idx + 1` + /// elements, and `dst` must be writable for `WIDTH` elements. + /// + /// Returns whether every gathered index was at most `max_idx`. + unsafe fn gather( + indices: *const Index, + max_idx: Index, + src: *const Lane, + dst: *mut Lane, + ) -> bool; +} + +/// AVX2 gather implementations for 32- and 64-bit value lanes. +pub(super) enum Avx2Gather {} + +macro_rules! impl_gather { + ($idx:ty, $({$value:ty => load: $load:ident, extend: $extend:ident, splat: $splat:ident, zero_vec: $zero_vec:ident, mask_indices: $mask_indices:ident, mask_cvt: |$mask_var:ident| $mask_cvt:block, gather: $masked_gather:ident, store: $store:ident, WIDTH = $WIDTH:literal, STRIDE = $STRIDE:literal }),+) => { + $( + impl_gather!(single; $idx, $value, load: $load, extend: $extend, splat: $splat, zero_vec: $zero_vec, mask_indices: $mask_indices, mask_cvt: |$mask_var| $mask_cvt, gather: $masked_gather, store: $store, WIDTH = $WIDTH, STRIDE = $STRIDE); + )* + }; + (single; $idx:ty, $value:ty, load: $load:ident, extend: $extend:ident, splat: $splat:ident, zero_vec: $zero_vec:ident, mask_indices: $mask_indices:ident, mask_cvt: |$mask_var:ident| $mask_cvt:block, gather: $masked_gather:ident, store: $store:ident, WIDTH = $WIDTH:literal, STRIDE = $STRIDE:literal) => { + impl GatherFn<$idx, $value> for Avx2Gather { + const WIDTH: usize = $WIDTH; + const STRIDE: usize = $STRIDE; + + #[allow(unused_unsafe, clippy::cast_possible_truncation)] + #[inline(always)] + unsafe fn gather( + indices: *const $idx, + max_idx: $idx, + src: *const $value, + dst: *mut $value, + ) -> bool { + const { + assert!($WIDTH <= $STRIDE, "dst cannot advance by more than the stride"); + } + + const SCALE: i32 = std::mem::size_of::<$value>() as i32; + + let indices_vec = unsafe { $load(indices.cast()) }; + let indices_vec = unsafe { $extend(indices_vec) }; + let max_idx_vec = unsafe { $splat(max_idx as _) }; + let valid_mask = unsafe { + _mm256_andnot_si256( + $mask_indices(indices_vec, max_idx_vec), + $splat(-1), + ) + }; + let all_valid = unsafe { _mm256_movemask_epi8(valid_mask) == -1 }; + let valid_mask = { + let $mask_var = valid_mask; + $mask_cvt + }; + let values_vec = unsafe { + $masked_gather::( + $zero_vec(), + src.cast(), + indices_vec, + valid_mask, + ) + }; + unsafe { $store(dst.cast(), values_vec) }; + all_valid + } + } + }; +} + +impl_gather!(u8, + { u32 => + load: _mm_loadu_si128, + extend: _mm256_cvtepu8_epi32, + splat: _mm256_set1_epi32, + zero_vec: _mm256_setzero_si256, + mask_indices: _mm256_cmpgt_epi32, + mask_cvt: |x| { x }, + gather: _mm256_mask_i32gather_epi32, + store: _mm256_storeu_si256, + WIDTH = 8, STRIDE = 16 + }, + { u64 => + load: _mm_loadu_si128, + extend: _mm256_cvtepu8_epi64, + splat: _mm256_set1_epi64x, + zero_vec: _mm256_setzero_si256, + mask_indices: _mm256_cmpgt_epi64, + mask_cvt: |x| { x }, + gather: _mm256_mask_i64gather_epi64, + store: _mm256_storeu_si256, + WIDTH = 4, STRIDE = 16 + } +); + +impl_gather!(u16, + { u32 => + load: _mm_loadu_si128, + extend: _mm256_cvtepu16_epi32, + splat: _mm256_set1_epi32, + zero_vec: _mm256_setzero_si256, + mask_indices: _mm256_cmpgt_epi32, + mask_cvt: |x| { x }, + gather: _mm256_mask_i32gather_epi32, + store: _mm256_storeu_si256, + WIDTH = 8, STRIDE = 8 + }, + { u64 => + load: _mm_loadu_si128, + extend: _mm256_cvtepu16_epi64, + splat: _mm256_set1_epi64x, + zero_vec: _mm256_setzero_si256, + mask_indices: _mm256_cmpgt_epi64, + mask_cvt: |x| { x }, + gather: _mm256_mask_i64gather_epi64, + store: _mm256_storeu_si256, + WIDTH = 4, STRIDE = 8 + } +); + +impl_gather!(u32, + { u32 => + load: _mm256_loadu_si256, + extend: identity, + splat: _mm256_set1_epi32, + zero_vec: _mm256_setzero_si256, + mask_indices: _mm256_cmpgt_epi32, + mask_cvt: |x| { x }, + gather: _mm256_mask_i32gather_epi32, + store: _mm256_storeu_si256, + WIDTH = 8, STRIDE = 8 + }, + { u64 => + load: _mm_loadu_si128, + extend: _mm256_cvtepu32_epi64, + splat: _mm256_set1_epi64x, + zero_vec: _mm256_setzero_si256, + mask_indices: _mm256_cmpgt_epi64, + mask_cvt: |x| { x }, + gather: _mm256_mask_i64gather_epi64, + store: _mm256_storeu_si256, + WIDTH = 4, STRIDE = 4 + } +); + +impl_gather!(u64, + { u32 => + load: _mm256_loadu_si256, + extend: identity, + splat: _mm256_set1_epi64x, + zero_vec: _mm_setzero_si128, + mask_indices: _mm256_cmpgt_epi64, + mask_cvt: |mask| { + unsafe { + let lo_bits = _mm256_extracti128_si256::<0>(mask); + let hi_bits = _mm256_extracti128_si256::<1>(mask); + let lo_packed = _mm_shuffle_epi32::<0b01_01_01_01>(lo_bits); + let hi_packed = _mm_shuffle_epi32::<0b01_01_01_01>(hi_bits); + _mm_unpacklo_epi64(lo_packed, hi_packed) + } + }, + gather: _mm256_mask_i64gather_epi32, + store: _mm_storeu_si128, + WIDTH = 4, STRIDE = 4 + }, + { u64 => + load: _mm256_loadu_si256, + extend: identity, + splat: _mm256_set1_epi64x, + zero_vec: _mm256_setzero_si256, + mask_indices: _mm256_cmpgt_epi64, + mask_cvt: |x| { x }, + gather: _mm256_mask_i64gather_epi64, + store: _mm256_storeu_si256, + WIDTH = 4, STRIDE = 4 + } +); diff --git a/vortex-array/src/arrays/fixed_width/take/avx2/mod.rs b/vortex-array/src/arrays/fixed_width/take/avx2/mod.rs new file mode 100644 index 00000000000..5335337d4d6 --- /dev/null +++ b/vortex-array/src/arrays/fixed_width/take/avx2/mod.rs @@ -0,0 +1,141 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! An AVX2 implementation of take operation using gather instructions. +//! +//! Only enabled for x86_64 hosts and it is gated at runtime behind feature detection to +//! ensure AVX2 instructions are available. + +#![cfg(any(target_arch = "x86_64", target_arch = "x86"))] + +mod gather; +#[cfg(test)] +mod tests; + +use std::arch::x86_64::__m256i; + +use vortex_buffer::Alignment; +use vortex_buffer::Buffer; +use vortex_buffer::BufferMut; + +use self::gather::Avx2Gather; +use self::gather::GatherFn; +use super::FixedWidthTakeValue; +use super::take_values_scalar; +use crate::dtype::UnsignedPType; +use crate::match_each_unsigned_integer_ptype; + +/// Takes the specified indices into a new [`Buffer`] using AVX2 SIMD. +/// +/// An AVX2 gather only moves raw bytes, so signedness and float-ness are irrelevant — only the +/// byte width of `V` matters. Any 4-byte value rides the gather through the `u32` lane and any +/// 8-byte value through the `u64` lane, regardless of its actual type. Values 1 or 2 bytes wide +/// (AVX2 has no sub-32-bit gather) and wider than 8 bytes fall back to the scalar kernel. +/// +/// [`FixedWidthTakeValue`] guarantees that the complete representation is initialized before it is +/// read through an integer lane. +/// +/// # Safety +/// +/// The caller must ensure the `avx2` feature is enabled. +#[target_feature(enable = "avx2")] +pub(super) unsafe fn take_avx2( + buffer: &[V], + indices: &[I], +) -> Buffer { + if buffer.is_empty() { + assert!( + indices.is_empty(), + "cannot take a non-empty set of indices from an empty buffer" + ); + return Buffer::empty(); + } + + // Dispatch on the gather lane width. The index type must still be concretized to select the + // right `GatherFn` impl, so re-dispatch it with `match_each_unsigned_integer_ptype!`. + macro_rules! dispatch { + ($lane:ty) => {{ + match_each_unsigned_integer_ptype!(I::PTYPE, |Idx| { + // SAFETY: `Idx` has the same `PTYPE` as `I`, so this is a no-op reinterpret of the + // index slice into the concrete type the gather impl is keyed on. + let indices = unsafe { std::mem::transmute::<&[I], &[Idx]>(indices) }; + exec_take::(buffer, indices) + }) + }}; + } + + match size_of::() { + 4 => dispatch!(u32), + 8 => dispatch!(u64), + // 1/2-byte and >8-byte values have no AVX2 gather lane, so fall back to scalar. + _ => take_values_scalar(buffer, indices), + } +} + +/// AVX2 core inner loop for a given index type `Idx`, output element type `Out`, and gather +/// `Lane` type. +/// +/// `Out` is the element type written to the output buffer; `Lane` (`u32` or `u64`) is the +/// integer type the gather intrinsics operate on. The caller must pair them so that +/// `size_of::() == size_of::()` (the only caller, [`take_avx2`], picks `Lane` from +/// `size_of::()`). Each valid lane copies the initialized representation of an existing +/// `Out`; invalid lanes are masked and cause a panic before the output buffer is initialized. +/// Gather instructions tolerate the source's potentially weaker alignment. +#[inline(always)] +fn exec_take(values: &[Out], indices: &[Idx]) -> Buffer +where + Out: FixedWidthTakeValue, + Idx: UnsignedPType, + Gather: GatherFn, +{ + assert_eq!( + size_of::(), + size_of::(), + "gather lane and output element must have the same size" + ); + + let indices_len = indices.len(); + let max_index = Idx::from(values.len() - 1).unwrap_or_else(|| Idx::max_value()); + let mut buffer = + BufferMut::::with_capacity_aligned(indices_len, Alignment::of::<__m256i>()); + let buf_uninit = buffer.spare_capacity_mut(); + + let mut offset = 0; + // Loop terminates STRIDE elements before end of the indices array because the `GatherFn` + // might read up to STRIDE src elements at a time, even though it only advances WIDTH elements + // in the dst. + while offset + Gather::STRIDE < indices_len { + // SAFETY: `gather` preconditions satisfied: + // 1. `(indices + offset)..(indices + offset + STRIDE)` is in-bounds for indices + // allocation. + // 2. `buffer` has same len as indices so `buffer + offset + WIDTH` is always valid. + // 3. `size_of::() == size_of::()` (asserted above), so the `Lane`-typed + // pointers address the same bytes as the `Out`-typed `values`/`buffer` allocations. + let all_valid = unsafe { + Gather::gather( + indices.as_ptr().add(offset), + max_index, + values.as_ptr().cast::(), + buf_uninit.as_mut_ptr().add(offset).cast::(), + ) + }; + assert!(all_valid, "take index out of bounds"); + offset += Gather::WIDTH; + } + + // Remainder. + while offset < indices_len { + buf_uninit[offset].write(values[indices[offset].as_()]); + offset += 1; + } + + assert_eq!(offset, indices_len); + + // SAFETY: All elements have been initialized. + unsafe { buffer.set_len(indices_len) }; + + // Do not expose the temporary SIMD over-alignment as part of the returned buffer. + buffer = buffer.aligned(Alignment::of::()); + + buffer.freeze() +} diff --git a/vortex-array/src/arrays/fixed_width/take/avx2/tests.rs b/vortex-array/src/arrays/fixed_width/take/avx2/tests.rs new file mode 100644 index 00000000000..cf8dc197bb3 --- /dev/null +++ b/vortex-array/src/arrays/fixed_width/take/avx2/tests.rs @@ -0,0 +1,114 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +#![cfg_attr(miri, ignore)] +#![cfg(target_arch = "x86_64")] + +use super::take_avx2; + +macro_rules! test_cases { + (index_type => $IDX:ty, value_types => $($VAL:ty),+) => { + paste::paste! { + $( + #[test] + #[allow(clippy::cast_possible_truncation)] + fn []() { + let values: Vec<$VAL> = (1..=127).map(|x| x as $VAL).collect(); + let indices: Vec<$IDX> = (0..127).collect(); + + let result = unsafe { take_avx2(&values, &indices) }; + assert_eq!(&values, result.as_slice()); + } + + #[test] + #[should_panic(expected = "cannot take a non-empty set of indices")] + #[allow(clippy::cast_possible_truncation)] + fn []() { + let values: Vec<$VAL> = vec![]; + let indices: Vec<$IDX> = (0..127).collect(); + drop(unsafe { take_avx2(&values, &indices) }); + } + + #[test] + #[should_panic(expected = "take index out of bounds")] + #[allow(clippy::cast_possible_truncation)] + fn []() { + let values: Vec<$VAL> = (1..=127).map(|x| x as $VAL).collect(); + let indices: Vec<$IDX> = (127..=254).collect(); + drop(unsafe { take_avx2(&values, &indices) }); + } + )+ + } + }; +} + +test_cases!( + index_type => u8, + value_types => u32, i32, u64, i64, f32, f64 +); +test_cases!( + index_type => u16, + value_types => u32, i32, u64, i64, f32, f64 +); +test_cases!( + index_type => u32, + value_types => u32, i32, u64, i64, f32, f64 +); +test_cases!( + index_type => u64, + value_types => u32, i32, u64, i64, f32, f64 +); + +#[test] +fn last_valid_u8_index() { + let values: Vec = (0..=255).collect(); + let indices: Vec = vec![255; 20]; + + let result = unsafe { take_avx2(&values, &indices) }; + assert_eq!(&[255; 20], result.as_slice()); +} + +#[test] +fn last_valid_u16_index() { + let values: Vec = (0..=65535).collect(); + let indices: Vec = vec![65535; 20]; + + let result = unsafe { take_avx2(&values, &indices) }; + assert_eq!(&[65535; 20], result.as_slice()); +} + +#[test] +#[should_panic(expected = "take index out of bounds")] +fn invalid_index_only_in_simd_block() { + let values = vec![10u32, 20, 30]; + let indices = vec![3u32, 0, 1, 2, 0, 1, 2, 0, 1]; + + drop(unsafe { take_avx2(&values, &indices) }); +} + +#[test] +fn simd_array_u8x4() { + let values: Vec<[u8; 4]> = (1u32..=200).map(u32::to_le_bytes).collect(); + let indices: Vec = (0..200).collect(); + + let result = unsafe { take_avx2(&values, &indices) }; + assert_eq!(values.as_slice(), result.as_slice()); +} + +#[test] +fn scalar_fallback_u16() { + let values: Vec = (1..=300).collect(); + let indices: Vec = (0..300).collect(); + + let result = unsafe { take_avx2(&values, &indices) }; + assert_eq!(values.as_slice(), result.as_slice()); +} + +#[test] +fn scalar_fallback_array_u8x16() { + let values: Vec<[u8; 16]> = (0u128..200).map(u128::to_le_bytes).collect(); + let indices: Vec = (0..200).collect(); + + let result = unsafe { take_avx2(&values, &indices) }; + assert_eq!(values.as_slice(), result.as_slice()); +} diff --git a/vortex-array/src/arrays/fixed_width/take/mod.rs b/vortex-array/src/arrays/fixed_width/take/mod.rs new file mode 100644 index 00000000000..6f79775ad4a --- /dev/null +++ b/vortex-array/src/arrays/fixed_width/take/mod.rs @@ -0,0 +1,198 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +#[cfg(any(target_arch = "x86_64", target_arch = "x86"))] +mod avx2; +mod records; +mod scalar; +mod slices; +#[cfg(test)] +mod tests; + +#[cfg(any(target_arch = "x86_64", target_arch = "x86"))] +use std::sync::LazyLock; + +use vortex_buffer::Buffer; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_mask::Mask; + +use self::records::take_byte_records; +use self::scalar::take_values_scalar; +use self::slices::take_slices; +use self::slices::take_slices_constant_length; +use super::FixedWidthArray; +use crate::ArrayRef; +use crate::Columnar; +use crate::ExecutionCtx; +use crate::IntoArray; +use crate::array::ArrayView; +use crate::arrays::ConstantArray; +use crate::arrays::PiecewiseSequence; +use crate::arrays::PrimitiveArray; +use crate::arrays::dict::TakeExecute; +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::UnsignedPType; +use crate::dtype::half::f16; +use crate::match_each_unsigned_integer_ptype; +use crate::scalar::Scalar; + +#[cfg(any(target_arch = "x86_64", target_arch = "x86"))] +static HAS_AVX2: LazyLock = LazyLock::new(|| is_x86_feature_detected!("avx2")); + +impl TakeExecute for V { + fn take( + array: ArrayView<'_, Self>, + indices: &ArrayRef, + ctx: &mut ExecutionCtx, + ) -> VortexResult> { + take(array, indices, ctx) + } +} + +/// A fixed-width value whose initialized bytes may be moved through integer SIMD lanes. +/// +/// # Safety +/// +/// Implementors must have no uninitialized bytes. The shared AVX2 gather reads the complete +/// representation through a same-width integer lane before writing those bytes back unchanged. +pub(crate) unsafe trait FixedWidthTakeValue: Copy {} + +macro_rules! impl_fixed_width_take_value { + ($($ty:ty),+ $(,)?) => { + $( + // SAFETY: These scalar representations contain no padding or uninitialized bytes. + unsafe impl FixedWidthTakeValue for $ty {} + )+ + }; +} + +impl_fixed_width_take_value!(u8, u16, u32, u64, i8, i16, i32, i64, f16, f32, f64,); + +// SAFETY: Byte arrays have no padding and every byte is initialized. +unsafe impl FixedWidthTakeValue for [u8; N] {} + +pub(crate) fn take_values( + values: &[T], + indices: &[I], +) -> Buffer { + #[cfg(any(target_arch = "x86_64", target_arch = "x86"))] + if *HAS_AVX2 { + // SAFETY: AVX2 was detected above and `FixedWidthTakeValue` guarantees an initialized byte + // representation. The AVX2 dispatcher retains Primitive's existing scalar fallbacks and + // out-of-bounds behavior for every value width. + return unsafe { avx2::take_avx2(values, indices) }; + } + + take_values_scalar(values, indices) +} + +pub(crate) fn take( + array: ArrayView<'_, V>, + 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)); + } + + let DType::Primitive(ptype, nullability) = indices.dtype() else { + vortex_bail!("Invalid indices dtype: {}", indices.dtype()) + }; + if !ptype.is_int() { + vortex_bail!("Invalid indices dtype: {}", indices.dtype()) + } + + let indices_validity = indices.validity()?; + 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 indices = if ptype.is_unsigned_int() { + indices_nulls_zeroed.execute::(ctx)? + } else { + indices_nulls_zeroed + .cast(DType::Primitive(ptype.to_unsigned(), *nullability))? + .execute::(ctx)? + }; + let validity = array + .validity()? + .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::(), + ) + })?; + Ok(Some( + V::with_values(array, values, indices.len(), validity)?.into_array(), + )) +} + +fn take_contiguous_ranges( + array: ArrayView<'_, V>, + indices: ArrayView<'_, PiecewiseSequence>, + indices_ref: &ArrayRef, + ctx: &mut ExecutionCtx, +) -> VortexResult> { + let Some((starts, lengths)) = maybe_contiguous_slices(indices, ctx)? else { + return Ok(None); + }; + + let values = V::values(array); + let byte_width = V::byte_width(array); + 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, + ) + }) + } + Columnar::Canonical(lengths) => { + let lengths = lengths.into_primitive(); + match_each_unsigned_integer_ptype!(starts.ptype(), |S| { + match_each_unsigned_integer_ptype!(lengths.ptype(), |L| { + take_slices( + &values, + byte_width, + array.len(), + starts.as_slice::(), + lengths.as_slice::(), + output_len, + ) + }) + }) + } + }?; + let validity = array.validity()?.take(indices_ref)?; + Ok(Some( + V::with_values(array, taken, 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 new file mode 100644 index 00000000000..f6f91f25529 --- /dev/null +++ b/vortex-array/src/arrays/fixed_width/take/records.rs @@ -0,0 +1,56 @@ +// 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 super::take_values; +use crate::dtype::UnsignedPType; + +pub(super) fn take_byte_records( + values: &ByteBuffer, + byte_width: usize, + record_count: usize, + indices: &[I], +) -> VortexResult { + let alignment = values.alignment(); + + macro_rules! take_typed_records { + ($width:expr) => {{ + let records = Buffer::<[u8; $width]>::from_byte_buffer(values.clone()); + debug_assert_eq!(records.len(), record_count); + Ok(take_values(records.as_slice(), indices) + .into_byte_buffer() + .aligned(alignment)) + }}; + } + + match byte_width { + 1 => return take_typed_records!(1), + 2 => return take_typed_records!(2), + 4 => return take_typed_records!(4), + 8 => return take_typed_records!(8), + 16 => return take_typed_records!(16), + 32 => return take_typed_records!(32), + _ => {} + } + + 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)) +} diff --git a/vortex-array/src/arrays/fixed_width/take/scalar.rs b/vortex-array/src/arrays/fixed_width/take/scalar.rs new file mode 100644 index 00000000000..7a1d1865be3 --- /dev/null +++ b/vortex-array/src/arrays/fixed_width/take/scalar.rs @@ -0,0 +1,27 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_buffer::Buffer; +use vortex_buffer::BufferMut; + +use crate::dtype::IntegerPType; + +#[inline(always)] +pub(crate) fn take_values_scalar( + values: &[T], + indices: &[I], +) -> Buffer { + // 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 result_ptr = result.spare_capacity_mut().as_mut_ptr().cast::(); + + for (output_index, index) in indices.iter().enumerate() { + // SAFETY: `indices.len()` elements were reserved and each output position is written once. + unsafe { result_ptr.add(output_index).write(values[index.as_()]) }; + } + + // SAFETY: the loop initialized every element in the reserved output range. + unsafe { result.set_len(indices.len()) }; + result.freeze() +} diff --git a/vortex-array/src/arrays/fixed_width/take/slices.rs b/vortex-array/src/arrays/fixed_width/take/slices.rs new file mode 100644 index 00000000000..a4d9ceda84a --- /dev/null +++ b/vortex-array/src/arrays/fixed_width/take/slices.rs @@ -0,0 +1,110 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::ptr; + +use itertools::Itertools as _; +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, + starts: &[S], + lengths: &[L], + output_len: usize, +) -> 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) +} + +pub(super) fn take_slices_constant_length( + values: &ByteBuffer, + byte_width: usize, + record_count: usize, + starts: &[S], + length: usize, + output_len: usize, +) -> VortexResult { + 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}" + ); + 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, + 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) + .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 cursor = 0usize; + + for (start, length) in slices { + let end = start + .checked_add(length) + .ok_or_else(|| vortex_err!("PiecewiseSequenceArray slice end overflows usize"))?; + vortex_ensure!( + 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]; + // SAFETY: `source` and the checked spare-capacity range have equal lengths and do not + // overlap. + unsafe { + ptr::copy_nonoverlapping( + source.as_ptr(), + spare[cursor..][..source.len()].as_mut_ptr().cast::(), + source.len(), + ); + } + cursor += source.len(); + } + + // 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() + ); + Ok(result.freeze().into_byte_buffer()) +} diff --git a/vortex-array/src/arrays/fixed_width/take/tests.rs b/vortex-array/src/arrays/fixed_width/take/tests.rs new file mode 100644 index 00000000000..98ac12d2a64 --- /dev/null +++ b/vortex-array/src/arrays/fixed_width/take/tests.rs @@ -0,0 +1,196 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use rstest::rstest; +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_session; +use crate::arrays::BoolArray; +use crate::arrays::ConstantArray; +use crate::arrays::DecimalArray; +use crate::arrays::PiecewiseSequenceArray; +use crate::arrays::PrimitiveArray; +use crate::assert_arrays_eq; +use crate::compute::conformance::take::test_take_conformance; +use crate::dtype::DecimalDType; +use crate::dtype::i256; +use crate::validity::Validity; + +#[test] +fn take_four_byte_records() { + let values = [[1u8, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12]]; + let taken = take_values(&values, &[2u32, 0]); + assert_eq!(taken.as_slice(), &[[9, 10, 11, 12], [1, 2, 3, 4]]); +} + +#[test] +fn take_eight_byte_values() { + let taken = take_values(&[10i64, 20, 30], &[1u16, 2, 0]); + assert_eq!(taken.as_slice(), &[20, 30, 10]); +} + +#[rstest] +#[case(1)] +#[case(2)] +#[case(4)] +#[case(8)] +#[case(16)] +#[case(32)] +fn take_runtime_width_records(#[case] byte_width: usize) -> VortexResult<()> { + let values = Buffer::from_iter((0u8..).take(3 * byte_width)); + 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); + Ok(()) +} + +#[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]); + 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(()) +} + +#[test] +fn null_index_skips_out_of_bounds_primitive_value() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let values = PrimitiveArray::from_iter([10i32, 20, 30]); + let indices = PrimitiveArray::new( + buffer![1u64, 3], + Validity::Array(BoolArray::from_iter([true, false]).into_array()), + ); + + let taken = values.take(indices.into_array())?; + + assert_arrays_eq!( + taken, + PrimitiveArray::from_option_iter([Some(20i32), None]).into_array(), + &mut ctx + ); + Ok(()) +} + +#[test] +fn null_index_skips_out_of_bounds_decimal_value() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let decimal_dtype = DecimalDType::new(19, 1); + let values = DecimalArray::new( + buffer![10i128, 20, 30], + decimal_dtype, + Validity::NonNullable, + ); + let indices = PrimitiveArray::new( + buffer![1u64, 3], + Validity::Array(BoolArray::from_iter([true, false]).into_array()), + ); + + let taken = values.take(indices.into_array())?; + + assert_arrays_eq!( + taken, + DecimalArray::from_option_iter([Some(20i128), None], decimal_dtype).into_array(), + &mut ctx + ); + Ok(()) +} + +#[test] +fn decimal_i256_take_consumes_piecewise_indices() -> 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 multipliers = ConstantArray::new(1u64, 2).into_array(); + let indices = PiecewiseSequenceArray::try_new(starts, lengths, multipliers, 3)?.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); + Ok(()) +} + +#[rstest] +#[case::primitive(PrimitiveArray::new( + buffer![0i32, 1, 2, 3, 4], + Validity::NonNullable, +).into_array())] +#[case::primitive_nullable(PrimitiveArray::from_option_iter( + [Some(1i64), None, Some(3), Some(4), None], +).into_array())] +#[case::decimal_i32(DecimalArray::new( + buffer![1i32, 2, 3, 4, 5], + DecimalDType::new(5, 0), + Validity::NonNullable, +).into_array())] +#[case::decimal_i64(DecimalArray::new( + buffer![10i64, 20, 30, 40, 50], + DecimalDType::new(10, 1), + Validity::NonNullable, +).into_array())] +#[case::decimal_i128(DecimalArray::new( + buffer![100i128, 200, 300, 400, 500], + DecimalDType::new(19, 2), + Validity::from_iter([true, false, true, true, false]), +).into_array())] +#[case::decimal_i256(DecimalArray::new( + buffer![ + i256::from_i128(100), + i256::from_i128(200), + i256::from_i128(300), + i256::from_i128(400), + i256::from_i128(500), + ], + DecimalDType::new(76, 2), + Validity::NonNullable, +).into_array())] +fn fixed_width_take_conformance(#[case] array: ArrayRef) { + test_take_conformance(&array, &mut array_session().create_execution_ctx()); +} diff --git a/vortex-array/src/arrays/primitive/compute/fixed_width.rs b/vortex-array/src/arrays/primitive/compute/fixed_width.rs new file mode 100644 index 00000000000..d112031a82b --- /dev/null +++ b/vortex-array/src/arrays/primitive/compute/fixed_width.rs @@ -0,0 +1,43 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_buffer::ByteBuffer; +use vortex_error::VortexResult; +use vortex_error::vortex_ensure; +use vortex_error::vortex_err; + +use crate::array::ArrayView; +use crate::arrays::Primitive; +use crate::arrays::PrimitiveArray; +use crate::arrays::fixed_width::FixedWidthArray; +use crate::validity::Validity; + +impl FixedWidthArray for Primitive { + fn byte_width(array: ArrayView<'_, Self>) -> usize { + array.ptype().byte_width() + } + + fn values(array: ArrayView<'_, Self>) -> ByteBuffer { + array.buffer_handle().to_host_sync() + } + + fn with_values( + array: ArrayView<'_, Self>, + values: ByteBuffer, + len: usize, + validity: Validity, + ) -> VortexResult { + let expected_len = len + .checked_mul(array.ptype().byte_width()) + .ok_or_else(|| vortex_err!("Primitive values buffer length overflows usize"))?; + vortex_ensure!( + values.len() == expected_len, + "Primitive values buffer length does not match output length" + ); + Ok(PrimitiveArray::from_byte_buffer( + values, + array.ptype(), + validity, + )) + } +} diff --git a/vortex-array/src/arrays/primitive/compute/mod.rs b/vortex-array/src/arrays/primitive/compute/mod.rs index 1ca4b17d3d5..382b42ee6e2 100644 --- a/vortex-array/src/arrays/primitive/compute/mod.rs +++ b/vortex-array/src/arrays/primitive/compute/mod.rs @@ -4,10 +4,10 @@ mod between; mod cast; mod fill_null; +mod fixed_width; mod mask; pub(crate) mod rules; mod slice; -mod take; mod zip; #[cfg(test)] diff --git a/vortex-array/src/arrays/primitive/compute/take/avx2.rs b/vortex-array/src/arrays/primitive/compute/take/avx2.rs deleted file mode 100644 index fb64e6748b0..00000000000 --- a/vortex-array/src/arrays/primitive/compute/take/avx2.rs +++ /dev/null @@ -1,536 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -//! An AVX2 implementation of take operation using gather instructions. -//! -//! Only enabled for x86_64 hosts and it is gated at runtime behind feature detection to -//! ensure AVX2 instructions are available. - -#![cfg(any(target_arch = "x86_64", target_arch = "x86"))] - -use std::arch::x86_64::__m256i; -use std::arch::x86_64::_mm_loadu_si128; -use std::arch::x86_64::_mm_setzero_si128; -use std::arch::x86_64::_mm_shuffle_epi32; -use std::arch::x86_64::_mm_storeu_si128; -use std::arch::x86_64::_mm_unpacklo_epi64; -use std::arch::x86_64::_mm256_andnot_si256; -use std::arch::x86_64::_mm256_cmpgt_epi32; -use std::arch::x86_64::_mm256_cmpgt_epi64; -use std::arch::x86_64::_mm256_cvtepu8_epi32; -use std::arch::x86_64::_mm256_cvtepu8_epi64; -use std::arch::x86_64::_mm256_cvtepu16_epi32; -use std::arch::x86_64::_mm256_cvtepu16_epi64; -use std::arch::x86_64::_mm256_cvtepu32_epi64; -use std::arch::x86_64::_mm256_extracti128_si256; -use std::arch::x86_64::_mm256_loadu_si256; -use std::arch::x86_64::_mm256_mask_i32gather_epi32; -use std::arch::x86_64::_mm256_mask_i64gather_epi32; -use std::arch::x86_64::_mm256_mask_i64gather_epi64; -use std::arch::x86_64::_mm256_set1_epi32; -use std::arch::x86_64::_mm256_set1_epi64x; -use std::arch::x86_64::_mm256_setzero_si256; -use std::arch::x86_64::_mm256_storeu_si256; -use std::convert::identity; - -use vortex_buffer::Alignment; -use vortex_buffer::Buffer; -use vortex_buffer::BufferMut; -use vortex_error::VortexResult; - -use crate::ArrayRef; -use crate::IntoArray; -use crate::array::ArrayView; -use crate::arrays::PrimitiveArray; -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; -use crate::match_each_native_ptype; -use crate::match_each_unsigned_integer_ptype; -use crate::validity::Validity; - -#[allow(unused)] -pub(super) struct TakeKernelAVX2; - -impl TakeImpl for TakeKernelAVX2 { - #[inline(always)] - fn take( - &self, - values: ArrayView<'_, Primitive>, - indices: ArrayView<'_, Primitive>, - validity: Validity, - ) -> VortexResult { - assert!(indices.ptype().is_unsigned_int()); - - Ok(match_each_unsigned_integer_ptype!(indices.ptype(), |I| { - match_each_native_ptype!(values.ptype(), |V| { - // SAFETY: This kernel is only selected when avx2 cpu-feature is detected. - unsafe { - take_primitive_avx2(values.as_slice::(), indices.as_slice::(), validity) - } - }) - }) - .into_array()) - } -} - -/// # Safety -/// -/// The caller must ensure that if the validity has a length, it is the same length as the indices, -/// and that the `avx2` feature is enabled. -#[target_feature(enable = "avx2")] -#[allow(unused)] -unsafe fn take_primitive_avx2( - values: &[V], - indices: &[I], - validity: Validity, -) -> PrimitiveArray -where - V: NativePType, - I: UnsignedPType, -{ - // SAFETY: The caller guarantees that the `avx2` feature is enabled. - let buffer = unsafe { take_avx2(values, indices) }; - - debug_assert!( - validity - .maybe_len() - .is_none_or(|validity_len| validity_len == buffer.len()) - ); - - // SAFETY: The caller ensures that the validity and indices have the same length, so the taken - // buffer and the validity must have the same length. - unsafe { PrimitiveArray::new_unchecked(buffer, validity) } -} - -// --------------------------------------------------------------------------- -// AVX2 SIMD take algorithm -// --------------------------------------------------------------------------- - -/// Takes the specified indices into a new [`Buffer`] using AVX2 SIMD. -/// -/// An AVX2 gather only moves raw bytes, so signedness and float-ness are irrelevant — only the -/// byte width of `V` matters. Any 4-byte value rides the gather through the `u32` lane and any -/// 8-byte value through the `u64` lane, regardless of its actual type. Values 1 or 2 bytes wide -/// (AVX2 has no sub-32-bit gather) and wider than 8 bytes (`i128`, decimals) fall back to the -/// scalar kernel. -/// -/// This treats `V` as plain-old-data: reinterpreting the gathered bytes as `V` is only sound -/// because every bit pattern is a valid `V`. All primitive and decimal-backing types satisfy -/// this, as does any `Copy` POD type the caller supplies. -/// -/// # Panics -/// -/// This function panics if any of the provided `indices` are out of bounds for `values`. -/// -/// # Safety -/// -/// The caller must ensure the `avx2` feature is enabled. -#[target_feature(enable = "avx2")] -#[doc(hidden)] -unsafe fn take_avx2(buffer: &[V], indices: &[I]) -> Buffer { - if buffer.is_empty() { - return Buffer::zeroed(indices.len()); - } - - // Dispatch on the gather lane width. The index type must still be concretized to select the - // right `GatherFn` impl, so re-dispatch it with `match_each_unsigned_integer_ptype!`. - macro_rules! dispatch { - ($lane:ty) => {{ - match_each_unsigned_integer_ptype!(I::PTYPE, |Idx| { - // SAFETY: `Idx` has the same `PTYPE` as `I`, so this is a no-op reinterpret of the - // index slice into the concrete type the gather impl is keyed on. - let indices = unsafe { std::mem::transmute::<&[I], &[Idx]>(indices) }; - exec_take::(buffer, indices) - }) - }}; - } - - match size_of::() { - 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), - } -} - -/// The main gather function that is used by the inner loop kernel for AVX2 gather. -trait GatherFn { - /// The number of data elements that are written to the `dst` on each loop iteration. - const WIDTH: usize; - /// The number of indices read from `indices` on each loop iteration. Depending on the - /// available instructions and bit-width we may stride by a larger amount than we actually - /// end up reading from `src` (governed by the `WIDTH` parameter). - const STRIDE: usize = Self::WIDTH; - - /// Gather values from `src` into the `dst` using the `indices`, optionally using SIMD - /// instructions. - /// - /// # Safety - /// - /// This function can read up to `STRIDE` elements through `indices`, and read/write up to - /// `WIDTH` elements through `src` and `dst` respectively. - unsafe fn gather(indices: *const Idx, max_idx: Idx, src: *const Values, dst: *mut Values); -} - -/// AVX2 version of [`GatherFn`] defined for 32- and 64-bit value types. -enum AVX2Gather {} - -macro_rules! impl_gather { - ($idx:ty, $({$value:ty => load: $load:ident, extend: $extend:ident, splat: $splat:ident, zero_vec: $zero_vec:ident, mask_indices: $mask_indices:ident, mask_cvt: |$mask_var:ident| $mask_cvt:block, gather: $masked_gather:ident, store: $store:ident, WIDTH = $WIDTH:literal, STRIDE = $STRIDE:literal }),+) => { - $( - impl_gather!(single; $idx, $value, load: $load, extend: $extend, splat: $splat, zero_vec: $zero_vec, mask_indices: $mask_indices, mask_cvt: |$mask_var| $mask_cvt, gather: $masked_gather, store: $store, WIDTH = $WIDTH, STRIDE = $STRIDE); - )* - }; - (single; $idx:ty, $value:ty, load: $load:ident, extend: $extend:ident, splat: $splat:ident, zero_vec: $zero_vec:ident, mask_indices: $mask_indices:ident, mask_cvt: |$mask_var:ident| $mask_cvt:block, gather: $masked_gather:ident, store: $store:ident, WIDTH = $WIDTH:literal, STRIDE = $STRIDE:literal) => { - impl GatherFn<$idx, $value> for AVX2Gather { - const WIDTH: usize = $WIDTH; - const STRIDE: usize = $STRIDE; - - #[allow(unused_unsafe, clippy::cast_possible_truncation)] - #[inline(always)] - unsafe fn gather(indices: *const $idx, max_idx: $idx, src: *const $value, dst: *mut $value) { - const { - assert!($WIDTH <= $STRIDE, "dst cannot advance by more than the stride"); - } - - const SCALE: i32 = std::mem::size_of::<$value>() as i32; - - let indices_vec = unsafe { $load(indices.cast()) }; - // Extend indices to fill vector register. - let indices_vec = unsafe { $extend(indices_vec) }; - - // Create a vec of the max idx. - let max_idx_vec = unsafe { $splat(max_idx as _) }; - // Create a mask for valid indices (where the max_idx > provided index). - let invalid_mask = unsafe { _mm256_andnot_si256($mask_indices(indices_vec, max_idx_vec), $splat(-1)) }; - let invalid_mask = { - let $mask_var = invalid_mask; - $mask_cvt - }; - let zero_vec = unsafe { $zero_vec() }; - - // Gather the values into new vector register, for masked positions - // it substitutes zero instead of accessing the src. - let values_vec = unsafe { $masked_gather::(zero_vec, src.cast(), indices_vec, invalid_mask) }; - - // Write the vec out to dst. - unsafe { $store(dst.cast(), values_vec) }; - } - } - }; -} - -// kernels for u8 indices -impl_gather!(u8, - // 32-bit values, loaded 8 at a time - { u32 => - load: _mm_loadu_si128, - extend: _mm256_cvtepu8_epi32, - splat: _mm256_set1_epi32, - zero_vec: _mm256_setzero_si256, - mask_indices: _mm256_cmpgt_epi32, - mask_cvt: |x| { x }, - gather: _mm256_mask_i32gather_epi32, - store: _mm256_storeu_si256, - WIDTH = 8, STRIDE = 16 - }, - - // 64-bit values, loaded 4 at a time - { u64 => - load: _mm_loadu_si128, - extend: _mm256_cvtepu8_epi64, - splat: _mm256_set1_epi64x, - zero_vec: _mm256_setzero_si256, - mask_indices: _mm256_cmpgt_epi64, - mask_cvt: |x| { x }, - gather: _mm256_mask_i64gather_epi64, - store: _mm256_storeu_si256, - WIDTH = 4, STRIDE = 16 - } -); - -// kernels for u16 indices -impl_gather!(u16, - // 32-bit values. 8x indices loaded at a time and 8x values written at a time. - { u32 => - load: _mm_loadu_si128, - extend: _mm256_cvtepu16_epi32, - splat: _mm256_set1_epi32, - zero_vec: _mm256_setzero_si256, - mask_indices: _mm256_cmpgt_epi32, - mask_cvt: |x| { x }, - gather: _mm256_mask_i32gather_epi32, - store: _mm256_storeu_si256, - WIDTH = 8, STRIDE = 8 - }, - - // 64-bit values. 8x indices loaded at a time and 4x values loaded at a time. - { u64 => - load: _mm_loadu_si128, - extend: _mm256_cvtepu16_epi64, - splat: _mm256_set1_epi64x, - zero_vec: _mm256_setzero_si256, - mask_indices: _mm256_cmpgt_epi64, - mask_cvt: |x| { x }, - gather: _mm256_mask_i64gather_epi64, - store: _mm256_storeu_si256, - WIDTH = 4, STRIDE = 8 - } -); - -// kernels for u32 indices -impl_gather!(u32, - // 32-bit values. 8x indices loaded at a time and 8x values written. - { u32 => - load: _mm256_loadu_si256, - extend: identity, - splat: _mm256_set1_epi32, - zero_vec: _mm256_setzero_si256, - mask_indices: _mm256_cmpgt_epi32, - mask_cvt: |x| { x }, - gather: _mm256_mask_i32gather_epi32, - store: _mm256_storeu_si256, - WIDTH = 8, STRIDE = 8 - }, - - // 64-bit values. - { u64 => - load: _mm_loadu_si128, - extend: _mm256_cvtepu32_epi64, - splat: _mm256_set1_epi64x, - zero_vec: _mm256_setzero_si256, - mask_indices: _mm256_cmpgt_epi64, - mask_cvt: |x| { x }, - gather: _mm256_mask_i64gather_epi64, - store: _mm256_storeu_si256, - WIDTH = 4, STRIDE = 4 - } -); - -// kernels for u64 indices -impl_gather!(u64, - { u32 => - load: _mm256_loadu_si256, - extend: identity, - splat: _mm256_set1_epi64x, - zero_vec: _mm_setzero_si128, - mask_indices: _mm256_cmpgt_epi64, - mask_cvt: |m| { - unsafe { - let lo_bits = _mm256_extracti128_si256::<0>(m); // lower half - let hi_bits = _mm256_extracti128_si256::<1>(m); // upper half - let lo_packed = _mm_shuffle_epi32::<0b01_01_01_01>(lo_bits); - let hi_packed = _mm_shuffle_epi32::<0b01_01_01_01>(hi_bits); - _mm_unpacklo_epi64(lo_packed, hi_packed) - } - }, - gather: _mm256_mask_i64gather_epi32, - store: _mm_storeu_si128, - WIDTH = 4, STRIDE = 4 - }, - - // 64-bit values. - { u64 => - load: _mm256_loadu_si256, - extend: identity, - splat: _mm256_set1_epi64x, - zero_vec: _mm256_setzero_si256, - mask_indices: _mm256_cmpgt_epi64, - mask_cvt: |x| { x }, - gather: _mm256_mask_i64gather_epi64, - store: _mm256_storeu_si256, - WIDTH = 4, STRIDE = 4 - } -); - -/// AVX2 core inner loop for a given index type `Idx`, output element type `Out`, and gather -/// `Lane` type. -/// -/// `Out` is the element type written to the output buffer; `Lane` (`u32` or `u64`) is the -/// integer type the gather intrinsics operate on. The caller must pair them so that -/// `size_of::() == size_of::()` (the only caller, [`take_avx2`], picks `Lane` from -/// `size_of::()`). The gather moves `size_of::()` raw bytes per element, which only -/// yields a valid `Out` because `Out` is plain-old-data (every bit pattern is a valid `Out`). -/// Pointers into the `Out`-typed slices are cast to `*const Lane`/`*mut Lane`; gather tolerates -/// the (possibly weaker) `Out` alignment. -#[inline(always)] -fn exec_take(values: &[Out], indices: &[Idx]) -> Buffer -where - Out: Copy, - Idx: UnsignedPType, - Gather: GatherFn, -{ - debug_assert_eq!( - size_of::(), - size_of::(), - "gather lane and output element must have the same size" - ); - - let indices_len = indices.len(); - let max_index = Idx::from(values.len()).unwrap_or_else(|| Idx::max_value()); - let mut buffer = - BufferMut::::with_capacity_aligned(indices_len, Alignment::of::<__m256i>()); - let buf_uninit = buffer.spare_capacity_mut(); - - let mut offset = 0; - // Loop terminates STRIDE elements before end of the indices array because the `GatherFn` - // might read up to STRIDE src elements at a time, even though it only advances WIDTH elements - // in the dst. - while offset + Gather::STRIDE < indices_len { - // SAFETY: `gather` preconditions satisfied: - // 1. `(indices + offset)..(indices + offset + STRIDE)` is in-bounds for indices - // allocation. - // 2. `buffer` has same len as indices so `buffer + offset + WIDTH` is always valid. - // 3. `size_of::() == size_of::()` (asserted above), so the `Lane`-typed - // pointers address the same bytes as the `Out`-typed `values`/`buffer` allocations. - unsafe { - Gather::gather( - indices.as_ptr().add(offset), - max_index, - values.as_ptr().cast::(), - buf_uninit.as_mut_ptr().add(offset).cast::(), - ) - }; - offset += Gather::WIDTH; - } - - // Remainder. - while offset < indices_len { - buf_uninit[offset].write(values[indices[offset].as_()]); - offset += 1; - } - - assert_eq!(offset, indices_len); - - // SAFETY: All elements have been initialized. - unsafe { buffer.set_len(indices_len) }; - - // Reset the buffer alignment to the output type. - // NOTE: if we don't do this, we pass back a Buffer which is over-aligned to the SIMD - // register width. The caller expects that this memory should be aligned to the value type - // so that we can slice it at value boundaries. - buffer = buffer.aligned(Alignment::of::()); - - buffer.freeze() -} - -#[cfg(test)] -#[cfg_attr(miri, ignore)] -#[cfg(target_arch = "x86_64")] -mod avx2_tests { - use super::*; - - macro_rules! test_cases { - (index_type => $IDX:ty, value_types => $($VAL:ty),+) => { - paste::paste! { - $( - // Test "happy path" take, valid indices on valid array. - #[test] - #[allow(clippy::cast_possible_truncation)] - fn []() { - let values: Vec<$VAL> = (1..=127).map(|x| x as $VAL).collect(); - let indices: Vec<$IDX> = (0..127).collect(); - - let result = unsafe { take_avx2(&values, &indices) }; - assert_eq!(&values, result.as_slice()); - } - - // Test take on empty array. - #[test] - #[should_panic] - #[allow(clippy::cast_possible_truncation)] - fn []() { - let values: Vec<$VAL> = vec![]; - let indices: Vec<$IDX> = (0..127).collect(); - let result = unsafe { take_avx2(&values, &indices) }; - assert!(result.is_empty()); - } - - // Test all invalid take indices mapping to zeros. - #[test] - #[should_panic] - #[allow(clippy::cast_possible_truncation)] - fn []() { - let values: Vec<$VAL> = (1..=127).map(|x| x as $VAL).collect(); - // All out of bounds indices. - let indices: Vec<$IDX> = (127..=254).collect(); - - let result = unsafe { take_avx2(&values, &indices) }; - assert_eq!(&[0 as $VAL; 127], result.as_slice()); - } - )+ - } - }; - } - - test_cases!( - index_type => u8, - value_types => u32, i32, u64, i64, f32, f64 - ); - test_cases!( - index_type => u16, - value_types => u32, i32, u64, i64, f32, f64 - ); - test_cases!( - index_type => u32, - value_types => u32, i32, u64, i64, f32, f64 - ); - test_cases!( - index_type => u64, - value_types => u32, i32, u64, i64, f32, f64 - ); - - #[test] - fn test_avx2_take_last_valid_index_u8() { - let values: Vec = (0..(255 + 1)).collect(); - let indices: Vec = vec![255; 20]; - - let result = unsafe { take_avx2(&values, &indices) }; - assert_eq!(&vec![255; indices.len()], result.as_slice()); - } - - #[test] - fn test_avx2_take_last_valid_index_u16() { - let values: Vec = (0..(65535 + 1)).collect(); - let indices: Vec = vec![65535; 20]; - - let result = unsafe { take_avx2(&values, &indices) }; - assert_eq!(&vec![65535; indices.len()], result.as_slice()); - } - - /// A `[u8; 4]` is a 4-byte `Copy` POD that is not a `NativePType`. This proves the kernel - /// gathers an arbitrary 4-byte value type through the `u32` SIMD lane. - #[test] - fn test_avx2_take_simd_array_u8x4() { - let values: Vec<[u8; 4]> = (1u32..=200).map(u32::to_le_bytes).collect(); - let indices: Vec = (0..200).collect(); - - let result = unsafe { take_avx2(&values, &indices) }; - assert_eq!(values.as_slice(), result.as_slice()); - } - - /// 2-byte values have no AVX2 gather, so they take the scalar fallback path and must still be - /// correct. - #[test] - fn test_avx2_take_scalar_fallback_u16() { - let values: Vec = (1..=300).collect(); - let indices: Vec = (0..300).collect(); - - let result = unsafe { take_avx2(&values, &indices) }; - assert_eq!(values.as_slice(), result.as_slice()); - } - - /// Values wider than 8 bytes (e.g. `i128`/decimal backing) exceed the gather lane and fall - /// back to the scalar kernel. - #[test] - fn test_avx2_take_scalar_fallback_array_u8x16() { - let values: Vec<[u8; 16]> = (0u128..200).map(u128::to_le_bytes).collect(); - let indices: Vec = (0..200).collect(); - - let result = unsafe { take_avx2(&values, &indices) }; - assert_eq!(values.as_slice(), result.as_slice()); - } -} diff --git a/vortex-array/src/arrays/primitive/compute/take/mod.rs b/vortex-array/src/arrays/primitive/compute/take/mod.rs deleted file mode 100644 index 5d663f55967..00000000000 --- a/vortex-array/src/arrays/primitive/compute/take/mod.rs +++ /dev/null @@ -1,453 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -#[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; -use crate::IntoArray; -use crate::array::ArrayView; -use crate::arrays::ConstantArray; -use crate::arrays::PiecewiseSequence; -use crate::arrays::Primitive; -use crate::arrays::PrimitiveArray; -use crate::arrays::dict::TakeExecute; -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; -use crate::match_each_integer_ptype; -use crate::match_each_native_ptype; -use crate::match_each_unsigned_integer_ptype; -use crate::scalar::Scalar; -use crate::validity::Validity; - -// Kernel selection happens on the first call to `take` and uses a combination of compile-time -// and runtime feature detection to infer the best kernel for the platform. -static PRIMITIVE_TAKE_KERNEL: LazyLock<&'static dyn TakeImpl> = LazyLock::new(|| { - #[cfg(any(target_arch = "x86_64", target_arch = "x86"))] - { - if is_x86_feature_detected!("avx2") { - &avx2::TakeKernelAVX2 - } else { - &TakeKernelScalar - } - } - - #[cfg(not(any(target_arch = "x86_64", target_arch = "x86")))] - { - &TakeKernelScalar - } -}); - -trait TakeImpl: Send + Sync { - fn take( - &self, - array: ArrayView<'_, Primitive>, - indices: ArrayView<'_, Primitive>, - validity: Validity, - ) -> VortexResult; -} - -struct TakeKernelScalar; - -impl TakeImpl for TakeKernelScalar { - fn take( - &self, - array: ArrayView<'_, Primitive>, - indices: ArrayView<'_, Primitive>, - validity: Validity, - ) -> 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::()); - Ok(PrimitiveArray::new(values, validity).into_array()) - }) - }) - } -} - -impl TakeExecute for Primitive { - fn take( - array: ArrayView<'_, Primitive>, - 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)); - } - - 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 validity = array - .validity()? - .take(&unsigned_indices.clone().into_array())? - .and(indices_validity)?; - // Delegate to the best kernel based on the target CPU - { - let unsigned_indices = unsigned_indices.as_view(); - PRIMITIVE_TAKE_KERNEL - .take(array, unsigned_indices, validity) - .map(Some) - } - } -} - -fn take_contiguous_ranges( - array: ArrayView<'_, Primitive>, - indices: ArrayView<'_, PiecewiseSequence>, - indices_ref: &ArrayRef, - ctx: &mut ExecutionCtx, -) -> VortexResult> { - let Some((starts, lengths)) = maybe_contiguous_slices(indices, ctx)? else { - return Ok(None); - }; - let validity = array.validity()?.take(indices_ref)?; - let output_len = indices_ref.len(); - let taken = match lengths { - Columnar::Constant(lengths) => { - let length = constant_unsigned_usize(&lengths); - take_slices_constant_length(array, &starts, length, validity, output_len)? - } - Columnar::Canonical(lengths) => { - let lengths = lengths.into_primitive(); - take_slices(array, &starts, &lengths, validity, output_len)? - } - }; - 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, - lengths: &PrimitiveArray, - validity: Validity, - output_len: usize, -) -> VortexResult { - match_each_native_ptype!(array.ptype(), |T| { - take_slices_typed::(array, starts, lengths, validity, output_len) - }) -} - -fn take_slices_typed( - array: ArrayView<'_, Primitive>, - starts: &PrimitiveArray, - lengths: &PrimitiveArray, - validity: Validity, - output_len: usize, -) -> VortexResult -where - T: NativePType, -{ - match_each_unsigned_integer_ptype!(starts.ptype(), |S| { - take_slices_start_typed::(array, starts, lengths, validity, output_len) - }) -} - -fn take_slices_start_typed( - array: ArrayView<'_, Primitive>, - starts: &PrimitiveArray, - lengths: &PrimitiveArray, - validity: Validity, - output_len: usize, -) -> VortexResult -where - T: NativePType, - S: UnsignedPType, -{ - match_each_unsigned_integer_ptype!(lengths.ptype(), |L| { - let values = take_slices_to_buffer::( - array.as_slice::(), - starts.as_slice::(), - lengths.as_slice::(), - output_len, - )?; - Ok(PrimitiveArray::new(values, validity).into_array()) - }) -} - -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, - length: usize, - validity: Validity, - output_len: usize, -) -> VortexResult { - match_each_native_ptype!(array.ptype(), |T| { - take_slices_constant_length_typed::(array, starts, length, validity, output_len) - }) -} - -fn take_slices_constant_length_typed( - array: ArrayView<'_, Primitive>, - starts: &PrimitiveArray, - length: usize, - validity: Validity, - output_len: usize, -) -> VortexResult -where - T: NativePType, -{ - match_each_unsigned_integer_ptype!(starts.ptype(), |S| { - let values = take_slices_constant_length_to_buffer::( - array.as_slice::(), - starts.as_slice::(), - length, - output_len, - )?; - Ok(PrimitiveArray::new(values, validity).into_array()) - }) -} - -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 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::compute::conformance::take::test_take_conformance; - use crate::scalar::Scalar; - use crate::validity::Validity; - - #[test] - fn test_take() { - let a = vec![1i32, 2, 3, 4, 5]; - let result = take_primitive_scalar(&a, &[0, 0, 4, 2]); - assert_eq!(result.as_slice(), &[1i32, 1, 5, 3]); - } - - #[test] - fn test_take_with_null_indices() { - let mut ctx = array_session().create_execution_ctx(); - let values = PrimitiveArray::new( - buffer![1i32, 2, 3, 4, 5], - Validity::Array(BoolArray::from_iter([true, true, false, false, true]).into_array()), - ); - let indices = PrimitiveArray::new( - buffer![0, 3, 4], - Validity::Array(BoolArray::from_iter([true, true, false]).into_array()), - ); - let actual = values.take(indices.into_array()).unwrap(); - assert_eq!( - actual.execute_scalar(0, &mut ctx).vortex_expect("no fail"), - Scalar::from(Some(1)) - ); - // position 3 is null - assert_eq!( - actual.execute_scalar(1, &mut ctx).vortex_expect("no fail"), - Scalar::null_native::() - ); - // the third index is null - assert_eq!( - actual.execute_scalar(2, &mut ctx).vortex_expect("no fail"), - Scalar::null_native::() - ); - } - - #[rstest] - #[case(PrimitiveArray::new(buffer![42i32], Validity::NonNullable))] - #[case(PrimitiveArray::new(buffer![0, 1], Validity::NonNullable))] - #[case(PrimitiveArray::new(buffer![0, 1, 2, 3, 4], Validity::NonNullable))] - #[case(PrimitiveArray::new(buffer![0, 1, 2, 3, 4, 5, 6, 7], Validity::NonNullable))] - #[case(PrimitiveArray::new(buffer![0, 1, 2, 3, 4], Validity::AllValid))] - #[case(PrimitiveArray::new( - buffer![0, 1, 2, 3, 4, 5], - Validity::Array(BoolArray::from_iter([true, false, true, false, true, true]).into_array()), - ))] - #[case(PrimitiveArray::from_option_iter([Some(1), None, Some(3), Some(4), None]))] - fn test_take_primitive_conformance(#[case] array: PrimitiveArray) { - test_take_conformance( - &array.into_array(), - &mut array_session().create_execution_ctx(), - ); - } -} - -#[cfg(test)] -mod tests { - use vortex_buffer::buffer; - - use crate::IntoArray; - use crate::VortexSessionExecute; - use crate::array_session; - use crate::arrays::BoolArray; - use crate::arrays::PrimitiveArray; - use crate::assert_arrays_eq; - use crate::validity::Validity; - - #[test] - fn take_null_index_skips_out_of_bounds_value() { - let mut ctx = array_session().create_execution_ctx(); - let values = PrimitiveArray::from_iter([10i32, 20, 30]); - let indices = PrimitiveArray::new( - buffer![1u64, 3], - Validity::Array(BoolArray::from_iter([true, false]).into_array()), - ); - - let taken = values.take(indices.into_array()).unwrap(); - - assert_arrays_eq!( - taken, - PrimitiveArray::from_option_iter([Some(20i32), None]).into_array(), - &mut ctx - ); - } -}