diff --git a/vortex-array/src/aggregate_fn/fns/is_constant/primitive.rs b/vortex-array/src/aggregate_fn/fns/is_constant/primitive.rs index 48c7423c0a7..c078f4dab09 100644 --- a/vortex-array/src/aggregate_fn/fns/is_constant/primitive.rs +++ b/vortex-array/src/aggregate_fn/fns/is_constant/primitive.rs @@ -1,10 +1,10 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_error::VortexExpect; + use crate::arrays::PrimitiveArray; -use crate::dtype::NativePType; -use crate::dtype::half::f16; -use crate::match_each_native_ptype; +use crate::match_each_ptype_width; cfg_if::cfg_if! { if #[cfg(target_feature = "avx2")] { @@ -14,9 +14,11 @@ cfg_if::cfg_if! { } } -/// Assumes any floating point has been cast into its bit representation for which != and !is_eq are the same -/// Assumes there's at least 1 value in the slice, which is an invariant of the entry level function. -pub fn compute_is_constant(values: &[T]) -> bool { +/// Checks whether every value equals the first value, comparing in groups of `WIDTH`. +/// +/// Callers that require bitwise equality must pass byte records or integer bit representations. +/// The slice must be nonempty and `WIDTH` must be nonzero. +pub fn compute_is_constant(values: &[T]) -> bool { let first_value = values[0]; let first_vec = &[first_value; WIDTH]; @@ -28,7 +30,7 @@ pub fn compute_is_constant(values: &[T]) -> } for value in remainder { - if !value.is_eq(first_value) { + if *value != first_value { return false; } } @@ -36,24 +38,43 @@ pub fn compute_is_constant(values: &[T]) -> true } -trait EqFloat { - type IntType; -} +pub(super) fn check_primitive_constant(array: &PrimitiveArray) -> bool { + let values = array + .buffer_handle() + .as_host_opt() + .vortex_expect("is_constant requires host values"); -impl EqFloat for f16 { - type IntType = u16; -} -impl EqFloat for f32 { - type IntType = u32; -} -impl EqFloat for f64 { - type IntType = u64; + match_each_ptype_width!( + array.ptype(), + |W| { + compute_is_constant::<_, { IS_CONST_LANE_WIDTH / W }>( + values.as_slice().as_chunks::().0, + ) + }, + _ => { unreachable!("primitive widths are specialized") }, + ) } -pub(super) fn check_primitive_constant(array: &PrimitiveArray) -> bool { - match_each_native_ptype!(array.ptype(), integral: |P| { - compute_is_constant::<_, {IS_CONST_LANE_WIDTH / size_of::

()}>(array.as_slice::

()) - }, floating: |P| { - compute_is_constant::<_, {IS_CONST_LANE_WIDTH / size_of::

()}>(unsafe { std::mem::transmute::<&[P], &[

::IntType]>(array.as_slice::

()) }) - }) +#[cfg(test)] +mod tests { + use rstest::rstest; + + use super::check_primitive_constant; + use crate::arrays::PrimitiveArray; + + #[rstest] + #[case::zero(&[0], true)] + #[case::signed_zero(&[0, 0x80000000], false)] + #[case::infinity(&[0x7f800000], true)] + #[case::same_nan(&[0x7fc00001], true)] + #[case::different_nan(&[0x7fc00001, 0x7fc00002], false)] + fn compares_stored_bits( + #[case] bits: &[u32], + #[case] expected: bool, + #[values(2, 16, 17, 65)] len: usize, + ) { + let values = bits.iter().copied().cycle().take(len).map(f32::from_bits); + let array = PrimitiveArray::from_iter(values); + assert_eq!(check_primitive_constant(&array), expected); + } } diff --git a/vortex-array/src/arrays/filter/execute/take/fixed_width.rs b/vortex-array/src/arrays/filter/execute/take/fixed_width.rs index 1f3fe8209db..81c876c848c 100644 --- a/vortex-array/src/arrays/filter/execute/take/fixed_width.rs +++ b/vortex-array/src/arrays/filter/execute/take/fixed_width.rs @@ -1,13 +1,13 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use vortex_buffer::Alignment; use vortex_buffer::BitBuffer; use vortex_buffer::Buffer; use vortex_buffer::BufferAllocatorRef; use vortex_buffer::BufferMut; use vortex_buffer::ByteBuffer; use vortex_error::VortexResult; -use vortex_error::vortex_panic; use vortex_mask::AllOr; use vortex_mask::Mask; @@ -26,9 +26,11 @@ use crate::arrays::decimal::DecimalArrayExt; use crate::arrays::filter::FilterArraySlotsExt; use crate::buffer::BufferHandle; use crate::dtype::IntegerPType; -use crate::dtype::i256; use crate::executor::ExecutionCtx; +use crate::match_each_decimal_type_width; +use crate::match_each_decimal_value_type; use crate::match_each_integer_ptype; +use crate::match_each_ptype_width; use crate::validity::Validity; #[inline] @@ -38,14 +40,22 @@ pub(in crate::arrays::filter) fn take_primitive( ctx: &mut ExecutionCtx, ) -> VortexResult { let child = array.child().clone().execute::(ctx)?; + let child_validity = child.validity()?; let ptype = child.ptype(); - let (taken, output_validity) = take_fixed_width_bytes( - child.buffer_handle().to_host_sync(), - ptype.byte_width(), - child.validity()?, - array.filter_mask(), - indices, - ctx, + let values = child.buffer_handle().to_host_sync(); + let (taken, output_validity) = match_each_ptype_width!( + ptype, + |W| { + take_fixed_width_records::( + values, + Alignment::new(W), + child_validity, + array.filter_mask(), + indices, + ctx, + ) + }, + _ => { unreachable!("primitive widths are specialized") }, )?; Ok(PrimitiveArray::from_byte_buffer(taken, ptype, output_validity).into_array()) @@ -58,18 +68,30 @@ pub(in crate::arrays::filter) fn take_decimal( ctx: &mut ExecutionCtx, ) -> VortexResult { let child = array.child().clone().execute::(ctx)?; + let child_validity = child.validity()?; let values_type = child.values_type(); - let (taken, output_validity) = take_fixed_width_bytes( - child.buffer_handle().to_host_sync(), - values_type.byte_width(), - child.validity()?, - array.filter_mask(), - indices, - ctx, + + // Decimal storage alignment can be smaller than its byte width. Retain the native alignment + // so contiguous record slices stay aligned without copying. + let alignment = match_each_decimal_value_type!(values_type, |T| { Alignment::of::() }); + let values = child.buffer_handle().to_host_sync(); + let (taken, output_validity) = match_each_decimal_type_width!( + values_type, + |W| { + take_fixed_width_records::( + values, + alignment, + child_validity, + array.filter_mask(), + indices, + ctx, + ) + }, + _ => { unreachable!("decimal widths are specialized") }, )?; - // SAFETY: Valid ranks copy existing decimal values, null ranks write default placeholders that - // are hidden by output validity, and the output validity was built for the take length. + // SAFETY: Valid ranks copy existing decimal values, null ranks write zero placeholders that are + // hidden by output validity, and the output validity was built for the take length. Ok(unsafe { DecimalArray::new_unchecked_handle( BufferHandle::new_host(taken), @@ -81,51 +103,27 @@ pub(in crate::arrays::filter) fn take_decimal( .into_array()) } -/// Takes values of `byte_width` bytes each. +/// Takes byte records with the array's required storage alignment. /// -/// The take copies values and does not inspect them, so all value types of one width share one -/// kernel. -fn take_fixed_width_bytes( - child: ByteBuffer, - byte_width: usize, - child_validity: Validity, - filter: &Mask, - indices: &PrimitiveArray, - ctx: &mut ExecutionCtx, -) -> VortexResult<(ByteBuffer, Validity)> { - match_each_integer_ptype!(indices.ptype(), |P| { - match byte_width { - 1 => take_fixed_width_as::(child, child_validity, filter, indices, ctx), - 2 => take_fixed_width_as::(child, child_validity, filter, indices, ctx), - 4 => take_fixed_width_as::(child, child_validity, filter, indices, ctx), - 8 => take_fixed_width_as::(child, child_validity, filter, indices, ctx), - 16 => take_fixed_width_as::(child, child_validity, filter, indices, ctx), - 32 => take_fixed_width_as::(child, child_validity, filter, indices, ctx), - _ => vortex_panic!("Unsupported fixed width of {byte_width} bytes"), - } - }) -} - -fn take_fixed_width_as( +/// `alignment` must divide `W` so slices between records remain aligned. +fn take_fixed_width_records( child: ByteBuffer, + alignment: Alignment, child_validity: Validity, filter: &Mask, indices: &PrimitiveArray, ctx: &mut ExecutionCtx, ) -> VortexResult<(ByteBuffer, Validity)> where - W: Copy + Default, - P: IntegerPType, + [u8; W]: Default, { - let (taken, output_validity) = take_fixed_width::( - Buffer::from_byte_buffer(child), - child_validity, - filter, - indices, - ctx, - )?; + let child = Buffer::<[u8; W]>::from_byte_buffer_aligned(child, alignment); + match_each_integer_ptype!(indices.ptype(), |P| { + let (taken, validity) = + take_fixed_width::<[u8; W], P>(child, child_validity, filter, indices, ctx)?; - Ok((taken.into_byte_buffer(), output_validity)) + Ok((taken.into_byte_buffer().aligned(alignment), validity)) + }) } fn take_fixed_width( diff --git a/vortex-array/src/arrays/filter/execute/take/tests.rs b/vortex-array/src/arrays/filter/execute/take/tests.rs index d869cde3f09..3a927aebeb1 100644 --- a/vortex-array/src/arrays/filter/execute/take/tests.rs +++ b/vortex-array/src/arrays/filter/execute/take/tests.rs @@ -4,7 +4,11 @@ use std::sync::LazyLock; use rstest::rstest; +use vortex_buffer::Alignment; +use vortex_buffer::Buffer; +use vortex_buffer::BufferMut; use vortex_buffer::buffer; +use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_mask::Mask; use vortex_session::VortexSession; @@ -23,14 +27,19 @@ use crate::arrays::Primitive; use crate::arrays::PrimitiveArray; use crate::arrays::StructArray; use crate::arrays::VarBinViewArray; +use crate::arrays::decimal::DecimalArrayExt; use crate::arrays::dict::TakeExecuteAdaptor; use crate::assert_arrays_eq; +use crate::buffer::BufferHandle; +use crate::dtype::BigCast; use crate::dtype::DecimalDType; +use crate::dtype::DecimalType; use crate::dtype::FieldNames; use crate::dtype::half::f16; use crate::dtype::i256; use crate::executor::ExecutionCtx; use crate::kernel::ExecuteParentKernel; +use crate::match_each_decimal_value_type; use crate::validity::Validity; static SESSION: LazyLock = LazyLock::new(VortexSession::empty); @@ -498,6 +507,66 @@ fn test_take_execute_kernel_handles_decimal_filter_child() -> VortexResult<()> { ) } +#[rstest] +fn decimal_contiguous_take_preserves_storage_and_shares_buffer( + #[values( + DecimalType::I8, + DecimalType::I16, + DecimalType::I32, + DecimalType::I64, + DecimalType::I128, + DecimalType::I256 + )] + values_type: DecimalType, +) -> VortexResult<()> { + let dtype = DecimalDType::new(2, 0); + let (child, expected) = match_each_decimal_value_type!(values_type, |T| { + let values = Buffer::from_iter([1i8, 2, 3, 4, 5].map(|value| { + ::from(value).vortex_expect("small values fit every decimal type") + })); + let alignment = Alignment::of::(); + let mut bytes = + BufferMut::::with_capacity(values.as_bytes().len() + alignment.as_usize()); + + // Offset by the native alignment to exercise storage that may not align to its byte width. + bytes.extend_from_slice(&vec![0; alignment.as_usize()]); + bytes.extend_from_slice(values.as_bytes()); + let bytes = bytes + .freeze() + .slice_with_alignment(alignment.as_usize().., alignment); + let child = DecimalArray::try_new_handle( + BufferHandle::new_host(bytes), + values_type, + dtype, + Validity::NonNullable, + )?; + let expected = DecimalArray::from_iter( + [3i8, 4].map(|value| { + ::from(value).vortex_expect("small values fit every decimal type") + }), + dtype, + ); + + (child, expected) + }); + let original_ptr = child.buffer_handle().as_host_opt().unwrap().as_ptr(); + let filter = + FilterArray::new(child.into_array(), Mask::from_slices(5, vec![(2, 5)])).into_array(); + let parent = DictArray::try_new(buffer![0u64, 1].into_array(), filter.clone())?.into_array(); + let mut ctx = SESSION.create_execution_ctx(); + let taken = execute_parent(&filter, &parent, 1, &mut ctx)? + .vortex_expect("decimal filter should execute its take parent") + .execute::(&mut ctx)?; + + assert_eq!(taken.values_type(), values_type); + assert_eq!( + taken.buffer_handle().as_host_opt().unwrap().as_ptr(), + original_ptr.wrapping_add(2 * values_type.byte_width()), + ); + assert_arrays_eq!(taken, expected, &mut ctx); + Ok(()) +} + #[test] fn test_take_execute_kernel_handles_fixed_size_list_filter_child() -> VortexResult<()> { assert_take_execute_maps_child_dtype( diff --git a/vortex-array/src/arrays/fixed_width/filter.rs b/vortex-array/src/arrays/fixed_width/filter.rs index 36af5cd531a..1c2b16b77d1 100644 --- a/vortex-array/src/arrays/fixed_width/filter.rs +++ b/vortex-array/src/arrays/fixed_width/filter.rs @@ -10,11 +10,11 @@ use vortex_mask::MaskValues; use vortex_mask::MaskValuesRef; use super::FixedWidthArray; -use super::match_each_record_width; use super::with_values; use crate::array::Array; use crate::arrays::filter::filter_buffer; use crate::arrays::filter::filter_validity; +use crate::match_each_power_of_two_width; #[cfg(test)] #[expect(clippy::cast_possible_truncation)] @@ -50,7 +50,7 @@ fn filter_records( ) -> ByteBuffer { let alignment = values.alignment(); - match_each_record_width!( + match_each_power_of_two_width!( byte_width, |W| { let records = Buffer::<[u8; W]>::from_byte_buffer(values); diff --git a/vortex-array/src/arrays/fixed_width/mod.rs b/vortex-array/src/arrays/fixed_width/mod.rs index ec19f05413c..c0f39120fcc 100644 --- a/vortex-array/src/arrays/fixed_width/mod.rs +++ b/vortex-array/src/arrays/fixed_width/mod.rs @@ -10,38 +10,3 @@ pub(crate) mod vtable; pub(crate) use self::array::FixedWidthArray; pub(crate) use self::array::with_values; - -/// Dispatches a runtime byte width to a compile-time `const $W: usize` for every record width -/// with a dedicated fixed-width kernel, falling back to `$fallback` for any other width. -macro_rules! match_each_record_width { - ($byte_width:expr, | $W:ident | $body:block,_ => $fallback:block) => { - match $byte_width { - 1 => { - const $W: usize = 1; - $body - } - 2 => { - const $W: usize = 2; - $body - } - 4 => { - const $W: usize = 4; - $body - } - 8 => { - const $W: usize = 8; - $body - } - 16 => { - const $W: usize = 16; - $body - } - 32 => { - const $W: usize = 32; - $body - } - _ => $fallback, - } - }; -} -pub(crate) use match_each_record_width; diff --git a/vortex-array/src/arrays/fixed_width/take/records.rs b/vortex-array/src/arrays/fixed_width/take/records.rs index 692707db9da..bc3f93746ab 100644 --- a/vortex-array/src/arrays/fixed_width/take/records.rs +++ b/vortex-array/src/arrays/fixed_width/take/records.rs @@ -9,8 +9,8 @@ use vortex_error::VortexResult; use vortex_error::vortex_err; use super::take_values; -use crate::arrays::fixed_width::match_each_record_width; use crate::dtype::UnsignedPType; +use crate::match_each_power_of_two_width; pub(super) fn take_byte_records( values: &ByteBuffer, @@ -21,7 +21,7 @@ pub(super) fn take_byte_records( ) -> VortexResult { let alignment = values.alignment(); - match_each_record_width!( + match_each_power_of_two_width!( byte_width, |W| { let records = Buffer::<[u8; W]>::from_byte_buffer(values.clone()); diff --git a/vortex-array/src/arrays/interleave/execute/primitive.rs b/vortex-array/src/arrays/interleave/execute/primitive.rs index c405df102df..3f75f8d4b60 100644 --- a/vortex-array/src/arrays/interleave/execute/primitive.rs +++ b/vortex-array/src/arrays/interleave/execute/primitive.rs @@ -4,6 +4,7 @@ //! Execution for primitive [`Interleave`] values. use num_traits::AsPrimitive; +use vortex_buffer::Alignment; use vortex_buffer::Buffer; use vortex_buffer::BufferAllocatorRef; use vortex_buffer::BufferMut; @@ -21,10 +22,10 @@ use crate::arrays::Constant; use crate::arrays::Primitive; use crate::arrays::PrimitiveArray; use crate::arrays::primitive::PrimitiveArrayExt; -use crate::dtype::NativePType; +use crate::dtype::ToBytes; use crate::executor::ExecutionCtx; use crate::executor::ExecutionResult; -use crate::match_each_native_ptype; +use crate::match_each_ptype_width; use crate::match_each_unsigned_integer_ptype; use crate::require_child; @@ -41,10 +42,17 @@ pub(super) fn execute( } let validity = array.as_ref().validity()?; - let output = match_each_native_ptype!(array.dtype().as_ptype(), |T| { - let values = gather_values::(&array, ctx.allocator())?; - VortexResult::Ok(PrimitiveArray::new(values, validity)) - })?; + let ptype = array.dtype().as_ptype(); + let values = match_each_ptype_width!( + ptype, + |W| { + gather_values::(&array, ctx.allocator())? + .into_byte_buffer() + .aligned(Alignment::new(W)) + }, + _ => { unreachable!("primitive widths are specialized") }, + ); + let output = PrimitiveArray::from_byte_buffer(values, ptype, validity); Ok(ExecutionResult::done(output)) } @@ -57,10 +65,10 @@ struct PrimitiveSource { row_mask: usize, } -fn gather_values( +fn gather_values( array: &Array, allocator: &BufferAllocatorRef, -) -> VortexResult> { +) -> VortexResult> { let values = (0..array.num_values()) .map(|i| { let value = array.value(i); @@ -70,8 +78,13 @@ fn gather_values( let payload = constant .scalar() .as_primitive() - .typed_value::() - .unwrap_or_default(); + .pvalue() + .map(|value| { + let mut bytes = [0u8; W]; + bytes.copy_from_slice(value.to_le_bytes()); + bytes + }) + .unwrap_or([0; W]); PrimitiveSource { data: Buffer::full_in(payload, 1, allocator.clone()), len, @@ -79,7 +92,9 @@ fn gather_values( } } else { PrimitiveSource { - data: value.as_::().to_buffer::(), + data: Buffer::<[u8; W]>::from_byte_buffer( + value.as_::().buffer_handle().to_host_sync(), + ), len, row_mask: usize::MAX, } @@ -108,7 +123,7 @@ fn gather( allocator: &BufferAllocatorRef, ) -> VortexResult> where - T: NativePType, + T: Copy, A: AsPrimitive, R: AsPrimitive, { diff --git a/vortex-array/src/arrays/primitive/array/patch.rs b/vortex-array/src/arrays/primitive/array/patch.rs index 91c8222c79e..a54e92ee84a 100644 --- a/vortex-array/src/arrays/primitive/array/patch.rs +++ b/vortex-array/src/arrays/primitive/array/patch.rs @@ -3,6 +3,8 @@ use std::ops::Range; +use vortex_buffer::Buffer; +use vortex_error::VortexExpect; use vortex_error::VortexResult; use crate::ExecutionCtx; @@ -12,7 +14,7 @@ use crate::dtype::IntegerPType; use crate::dtype::NativePType; use crate::dtype::UnsignedPType; use crate::match_each_integer_ptype; -use crate::match_each_native_ptype; +use crate::match_each_ptype_width; use crate::patches::PATCH_CHUNK_SIZE; use crate::patches::Patches; use crate::validity::Validity; @@ -22,6 +24,8 @@ impl PrimitiveArray { let patch_indices = patches.indices().clone().execute::(ctx)?; let patch_values = patches.values().clone().execute::(ctx)?; + assert_eq!(self.ptype(), patch_values.ptype()); + let patch_validity = patch_values.validity()?; let patched_validity = self.validity()?.patch( self.len(), @@ -30,19 +34,24 @@ impl PrimitiveArray { &patch_validity, ctx, )?; + Ok(match_each_integer_ptype!(patch_indices.ptype(), |I| { - match_each_native_ptype!(self.ptype(), |T| { - self.patch_typed::( - patch_indices, - patches.offset(), - patch_values, - patched_validity, - ) - }) + match_each_ptype_width!( + self.ptype(), + |W| { + self.patch_typed::( + patch_indices, + patches.offset(), + patch_values, + patched_validity, + ) + }, + _ => { unreachable!("primitive widths are specialized") }, + ) })) } - fn patch_typed( + fn patch_typed( self, patch_indices: PrimitiveArray, patch_indices_offset: usize, @@ -50,17 +59,32 @@ impl PrimitiveArray { patched_validity: Validity, ) -> Self where - T: NativePType, I: IntegerPType, { - let mut own_values = self.into_buffer_mut::(); + let ptype = self.ptype(); + let values = self.into_data_parts().buffer.into_host_sync(); + let alignment = values.alignment(); + let mut own_values = + Buffer::<[u8; W]>::from_byte_buffer_aligned(values, alignment).into_mut(); let patch_indices = patch_indices.as_slice::(); - let patch_values = patch_values.as_slice::(); + let patch_values = patch_values + .buffer_handle() + .as_host_opt() + .vortex_expect("patch values require host storage") + .as_slice() + .as_chunks::() + .0; + for (idx, value) in itertools::zip_eq(patch_indices, patch_values) { own_values[idx.as_() - patch_indices_offset] = *value; } - Self::new(own_values, patched_validity) + + Self::from_byte_buffer( + own_values.freeze().into_byte_buffer(), + ptype, + patched_validity, + ) } } diff --git a/vortex-array/src/arrays/primitive/compute/fill_null.rs b/vortex-array/src/arrays/primitive/compute/fill_null.rs index ea657448327..97cd53ed8b5 100644 --- a/vortex-array/src/arrays/primitive/compute/fill_null.rs +++ b/vortex-array/src/arrays/primitive/compute/fill_null.rs @@ -3,6 +3,7 @@ use std::ops::Not; +use vortex_buffer::Buffer; use vortex_error::VortexExpect; use vortex_error::VortexResult; @@ -13,7 +14,8 @@ use crate::array::ArrayView; use crate::arrays::BoolArray; use crate::arrays::Primitive; use crate::arrays::PrimitiveArray; -use crate::match_each_native_ptype; +use crate::dtype::ToBytes; +use crate::match_each_ptype_width; use crate::scalar::Scalar; use crate::scalar_fn::fns::fill_null::FillNullKernel; use crate::validity::Validity; @@ -29,17 +31,35 @@ impl FillNullKernel for Primitive { Ok(Some(match array.validity()? { Validity::Array(is_valid) => { let is_invalid = is_valid.execute::(ctx)?.into_bit_buffer().not(); - match_each_native_ptype!(array.ptype(), |T| { - let mut buffer = array.to_buffer::().into_mut(); - let fill_value = fill_value - .as_primitive() - .typed_value::() - .vortex_expect("top-level fill_null ensure non-null fill value"); - for invalid_index in is_invalid.set_indices() { - buffer[invalid_index] = fill_value; - } - PrimitiveArray::new(buffer.freeze(), result_validity).into_array() - }) + let fill_value = fill_value.as_primitive(); + assert_eq!(array.ptype(), fill_value.ptype()); + let fill_value = fill_value + .pvalue() + .vortex_expect("fill_null requires a non-null fill value"); + let values = array.buffer_handle().to_host_sync(); + let alignment = values.alignment(); + + match_each_ptype_width!( + array.ptype(), + |W| { + let mut buffer = + Buffer::<[u8; W]>::from_byte_buffer_aligned(values, alignment).into_mut(); + let mut fill_bytes = [0u8; W]; + fill_bytes.copy_from_slice(fill_value.to_le_bytes()); + + for invalid_index in is_invalid.set_indices() { + buffer[invalid_index] = fill_bytes; + } + + PrimitiveArray::from_byte_buffer( + buffer.freeze().into_byte_buffer(), + array.ptype(), + result_validity, + ) + .into_array() + }, + _ => { unreachable!("primitive widths are specialized") }, + ) } _ => unreachable!("checked in entry point"), })) @@ -47,7 +67,7 @@ impl FillNullKernel for Primitive { } #[cfg(test)] -mod test { +mod tests { use std::sync::LazyLock; use vortex_buffer::buffer; diff --git a/vortex-array/src/arrays/primitive/compute/zip.rs b/vortex-array/src/arrays/primitive/compute/zip.rs index a653f01099e..6d0456e299a 100644 --- a/vortex-array/src/arrays/primitive/compute/zip.rs +++ b/vortex-array/src/arrays/primitive/compute/zip.rs @@ -3,6 +3,7 @@ use std::mem::MaybeUninit; +use vortex_buffer::Alignment; use vortex_buffer::BufferMut; use vortex_error::VortexExpect; use vortex_error::VortexResult; @@ -16,8 +17,7 @@ use crate::array::ArrayView; use crate::arrays::Primitive; use crate::arrays::PrimitiveArray; use crate::builtins::ArrayBuiltins; -use crate::dtype::NativePType; -use crate::match_each_native_ptype; +use crate::match_each_ptype_width; use crate::scalar_fn::fns::zip::ZipKernel; use crate::scalar_fn::fns::zip::zip_validity; @@ -55,27 +55,43 @@ impl ZipKernel for Primitive { } let validity = zip_validity(if_true.validity()?, if_false.validity()?, &mask)?; - - // TODO(perf): inspect the mask's true_count (and validity) to special-case heavily-skewed + let true_values = if_true + .buffer_handle() + .as_host_opt() + .vortex_expect("zip requires host values"); + let false_values = if_false + .buffer_handle() + .as_host_opt() + .vortex_expect("zip requires host values"); + + // TODO(connor): inspect the mask's true_count (and validity) to special-case heavily-skewed // masks. When one side dominates (true_count near 0 or near len), it is cheaper to bulk - // copy — or mutate in place, if the dominant side is uniquely owned — that side's values + // copy (or mutate in place, if the dominant side is uniquely owned) that side's values // and validity, then conditionally pull in only the minority rows from the other side, // rather than blending every row. - let array = match_each_native_ptype!(if_true.ptype(), |T| { - let values = - select_values::(if_true.as_slice::(), if_false.as_slice::(), &mask); - PrimitiveArray::new(values.freeze(), validity).into_array() - }); + let values = match_each_ptype_width!( + if_true.ptype(), + |W| { + select_values( + true_values.as_slice().as_chunks::().0, + false_values.as_slice().as_chunks::().0, + &mask, + ) + .freeze() + .into_byte_buffer() + .aligned(Alignment::new(W)) + }, + _ => { unreachable!("primitive widths are specialized") }, + ); + let array = + PrimitiveArray::from_byte_buffer(values, if_true.ptype(), validity).into_array(); + Ok(Some(array)) } } /// Branchlessly blend `if_true` and `if_false` per row into a fresh value buffer. -fn select_values( - true_values: &[T], - false_values: &[T], - mask: &Mask, -) -> BufferMut { +fn select_values(true_values: &[T], false_values: &[T], mask: &Mask) -> BufferMut { let len = true_values.len(); let mut out = BufferMut::::with_capacity(len); { @@ -122,7 +138,7 @@ fn select_values( /// bit keeps `false_values[j]`. Slices are trimmed to the output length up front so the compiler can /// elide bounds checks and lower the body to a vector blend / conditional move. #[inline] -fn select_block( +fn select_block( word: u64, true_values: &[T], false_values: &[T], diff --git a/vortex-array/src/dtype/decimal/macros.rs b/vortex-array/src/dtype/decimal/macros.rs index 39629ac285e..27682caa7cb 100644 --- a/vortex-array/src/dtype/decimal/macros.rs +++ b/vortex-array/src/dtype/decimal/macros.rs @@ -1,6 +1,71 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +/// Dispatches a decimal storage type to a compile-time byte width for structural kernels. +/// +/// Accepts an owned or borrowed [`DecimalType`](crate::dtype::DecimalType), evaluated once, and +/// specializes 1, 2, 4, 8, 16, and 32 bytes. Kernels can operate on `[u8; W]` records while the caller +/// preserves decimal storage type, precision, and scale. +/// +/// The optional `widths = [...]` list and required fallback follow +/// [`match_each_power_of_two_width!`](crate::match_each_power_of_two_width). +/// +/// # Examples +/// +/// ``` +/// use vortex_array::dtype::DecimalType; +/// use vortex_array::match_each_decimal_type_width; +/// +/// let width = match_each_decimal_type_width!( +/// DecimalType::I128, +/// |W| { W }, +/// _ => { unreachable!() }, +/// ); +/// assert_eq!(width, 16); +/// ``` +/// +/// Logical precision and scale do not determine an existing buffer's physical width: +/// +/// ```compile_fail +/// use vortex_array::dtype::DecimalDType; +/// use vortex_array::match_each_decimal_type_width; +/// +/// match_each_decimal_type_width!(DecimalDType::new(10, 2), |W| { W }, _ => { unreachable!() }); +/// ``` +/// +/// Primitive types use [`match_each_ptype_width!`](crate::match_each_ptype_width): +/// +/// ```compile_fail +/// use vortex_array::dtype::PType; +/// use vortex_array::match_each_decimal_type_width; +/// +/// match_each_decimal_type_width!(PType::I64, |W| { W }, _ => { unreachable!() }); +/// ``` +#[macro_export] +macro_rules! match_each_decimal_type_width { + ( + $decimal_type:expr, + widths = [$($width:literal),+ $(,)?], + |$W:ident| $body:block, + $other:pat => $fallback:block $(,)? + ) => { + $crate::match_each_power_of_two_width!( + $crate::dtype::DecimalType::byte_width(&$decimal_type), + widths = [$($width),+], + |$W| $body, + $other => $fallback, + ) + }; + ($decimal_type:expr, |$W:ident| $body:block, $other:pat => $fallback:block $(,)?) => { + $crate::match_each_decimal_type_width!( + $decimal_type, + widths = [1, 2, 4, 8, 16, 32], + |$W| $body, + $other => $fallback, + ) + }; +} + /// Matches over each decimal value variant, binding the inner value to a variable. /// /// # Example diff --git a/vortex-array/src/dtype/ptype.rs b/vortex-array/src/dtype/ptype.rs index 03a0a9d1da8..d386b604261 100644 --- a/vortex-array/src/dtype/ptype.rs +++ b/vortex-array/src/dtype/ptype.rs @@ -352,6 +352,59 @@ native_float_ptype!(f16, F16); native_float_ptype!(f32, F32); native_float_ptype!(f64, F64); +/// Dispatches a [`PType`] to a compile-time byte width for structural kernels. +/// +/// Accepts an owned or borrowed [`PType`], evaluated once, and specializes 1, 2, 4, and 8 bytes. +/// Kernels can operate on `[u8; W]` records as bit patterns. Callers preserve the original dtype +/// when rebuilding arrays, including signedness and floating-point interpretation. +/// +/// The optional `widths = [...]` list and required fallback follow +/// [`match_each_power_of_two_width!`](crate::match_each_power_of_two_width). +/// +/// # Examples +/// +/// ``` +/// use vortex_array::dtype::PType; +/// use vortex_array::match_each_ptype_width; +/// +/// let width = match_each_ptype_width!(PType::F32, |W| { W }, _ => { unreachable!() }); +/// assert_eq!(width, 4); +/// ``` +/// +/// Decimal storage types use +/// [`match_each_decimal_type_width!`](crate::match_each_decimal_type_width): +/// +/// ```compile_fail +/// use vortex_array::dtype::DecimalType; +/// use vortex_array::match_each_ptype_width; +/// +/// match_each_ptype_width!(DecimalType::I32, |W| { W }, _ => { unreachable!() }); +/// ``` +#[macro_export] +macro_rules! match_each_ptype_width { + ( + $ptype:expr, + widths = [$($width:literal),+ $(,)?], + |$W:ident| $body:block, + $other:pat => $fallback:block $(,)? + ) => { + $crate::match_each_power_of_two_width!( + $crate::dtype::PType::byte_width(&$ptype), + widths = [$($width),+], + |$W| $body, + $other => $fallback, + ) + }; + ($ptype:expr, |$W:ident| $body:block, $other:pat => $fallback:block $(,)?) => { + $crate::match_each_ptype_width!( + $ptype, + widths = [1, 2, 4, 8], + |$W| $body, + $other => $fallback, + ) + }; +} + /// Macro to match over each PType, binding the corresponding native type (from `NativePType`) #[macro_export] macro_rules! match_each_native_ptype { diff --git a/vortex-array/src/lib.rs b/vortex-array/src/lib.rs index 53ccf53bb52..fb009dce9cb 100644 --- a/vortex-array/src/lib.rs +++ b/vortex-array/src/lib.rs @@ -88,6 +88,7 @@ pub use mask_future::*; pub use metadata::*; pub use smallvec; pub use vortex_array_macros::array_slots; +pub use vortex_buffer::match_each_power_of_two_width; use vortex_session::SessionExt; use vortex_session::VortexSession; use vortex_session::registry::Interner; diff --git a/vortex-array/tests/width_dispatch.rs b/vortex-array/tests/width_dispatch.rs new file mode 100644 index 00000000000..f00b5ab4be9 --- /dev/null +++ b/vortex-array/tests/width_dispatch.rs @@ -0,0 +1,226 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Public width dispatch and bit-preserving structural operations. + +#[cfg(test)] +mod tests { + use std::sync::LazyLock; + + use rstest::rstest; + use vortex_array::ExecutionCtx; + use vortex_array::IntoArray; + use vortex_array::VortexSessionExecute; + use vortex_array::arrays::ConstantArray; + use vortex_array::arrays::InterleaveArray; + use vortex_array::arrays::Primitive; + use vortex_array::arrays::PrimitiveArray; + use vortex_array::arrays::primitive::PrimitiveArrayExt; + use vortex_array::assert_arrays_eq; + use vortex_array::dtype::DecimalType; + use vortex_array::dtype::PType; + use vortex_array::match_each_decimal_type_width; + use vortex_array::match_each_power_of_two_width; + use vortex_array::match_each_ptype_width; + use vortex_array::patches::Patches; + use vortex_array::scalar_fn::fns::fill_null::FillNullKernel; + use vortex_array::scalar_fn::fns::zip::ZipKernel; + use vortex_array::validity::Validity; + use vortex_buffer::Alignment; + use vortex_buffer::Buffer; + use vortex_error::VortexExpect; + use vortex_error::VortexResult; + use vortex_mask::Mask; + use vortex_session::VortexSession; + + static SESSION: LazyLock = LazyLock::new(VortexSession::empty); + + /// Builds over-aligned storage from the low bytes of each supplied bit pattern. + fn primitive_from_bits(ptype: PType, bits: &[u64], validity: Validity) -> PrimitiveArray { + let width = ptype.byte_width(); + let values = bits + .iter() + .flat_map(|bits| bits.to_le_bytes().into_iter().take(width)); + let buffer = Buffer::from_iter(values).aligned(Alignment::new(64)); + + PrimitiveArray::from_byte_buffer(buffer, ptype, validity) + } + + #[track_caller] + fn assert_records(actual: &PrimitiveArray, ptype: PType, bits: &[u64], ctx: &mut ExecutionCtx) { + let expected = primitive_from_bits(ptype, bits, Validity::NonNullable); + assert_arrays_eq!((*actual).clone(), expected, ctx); + assert_eq!( + actual.buffer_handle().as_host_opt().unwrap().as_slice(), + expected.buffer_handle().as_host_opt().unwrap().as_slice(), + ); + } + + #[rstest] + fn accepts_owned_and_borrowed_ptypes( + #[values( + PType::U8, PType::U16, PType::U32, PType::U64, PType::I8, PType::I16, PType::I32, + PType::I64, PType::F16, PType::F32, PType::F64 + )] + ptype: PType, + ) { + let owned = match_each_ptype_width!(ptype, |W| { W }, _ => { unreachable!() }); + let borrowed = match_each_ptype_width!(&ptype, |W| { W }, _ => { unreachable!() }); + assert_eq!(owned, ptype.byte_width()); + assert_eq!(borrowed, owned); + } + + #[rstest] + fn accepts_owned_and_borrowed_decimal_types( + #[values( + DecimalType::I8, + DecimalType::I16, + DecimalType::I32, + DecimalType::I64, + DecimalType::I128, + DecimalType::I256 + )] + values_type: DecimalType, + ) { + let owned = match_each_decimal_type_width!(values_type, |W| { W }, _ => { unreachable!() }); + let borrowed = + match_each_decimal_type_width!(&values_type, |W| { W }, _ => { unreachable!() }); + assert_eq!(owned, values_type.byte_width()); + assert_eq!(borrowed, owned); + } + + #[test] + fn forwards_specializations_and_evaluates_physical_types_once() { + let mut ptype_calls = 0; + let primitive = match_each_ptype_width!( + { + ptype_calls += 1; + PType::F32 + }, + widths = [1, 8], + |W| { Ok(W) }, + other => { Err(other) }, + ); + let mut decimal_calls = 0; + let decimal = match_each_decimal_type_width!( + { + decimal_calls += 1; + DecimalType::I128 + }, + widths = [1, 8], + |W| { Ok(W) }, + other => { Err(other) }, + ); + assert_eq!(primitive, Err(4)); + assert_eq!(decimal, Err(16)); + assert_eq!((ptype_calls, decimal_calls), (1, 1)); + + let width = match_each_power_of_two_width!(16, |W| { W }, _ => { unreachable!() }); + assert_eq!(width, 16); + } + + #[rstest] + fn records_survive_structural_kernels( + #[values( + PType::U8, PType::U16, PType::U32, PType::U64, PType::I8, PType::I16, PType::I32, + PType::I64, PType::F16, PType::F32, PType::F64 + )] + ptype: PType, + #[values(false, true)] shared: bool, + ) -> VortexResult<()> { + let bits = match ptype { + PType::F16 => [0, 0x8000, 0x7c00, 0xfc00, 0x7e01, 0x7e02], + PType::F32 => [ + 0, 0x80000000, 0x7f800000, 0xff800000, 0x7fc00001, 0x7fc00002, + ], + PType::F64 => [ + 0, + 0x8000000000000000, + 0x7ff0000000000000, + 0xfff0000000000000, + 0x7ff8000000000001, + 0x7ff8000000000002, + ], + _ => [0, u64::MAX, 1, 2, 0x8080808080808080, 0x7f7f7f7f7f7f7f7f], + }; + let mut ctx = SESSION.create_execution_ctx(); + let input = primitive_from_bits(ptype, &bits, Validity::NonNullable); + let input_ptr = input.buffer_handle().as_host_opt().unwrap().as_ptr(); + let retained = shared.then(|| input.clone()); + let patches = Patches::new( + bits.len(), + 7, + PrimitiveArray::from_iter([8u32, 11]).into_array(), + primitive_from_bits(ptype, &[bits[5], bits[0]], Validity::NonNullable).into_array(), + None, + &mut ctx, + )?; + let patched = input.patch(&patches, &mut ctx)?; + assert_records( + &patched, + ptype, + &[bits[0], bits[5], bits[2], bits[3], bits[0], bits[5]], + &mut ctx, + ); + assert!( + patched + .buffer_handle() + .as_host_opt() + .unwrap() + .is_aligned(Alignment::new(64)) + ); + let output_ptr = patched.buffer_handle().as_host_opt().unwrap().as_ptr(); + if let Some(retained) = retained { + assert_ne!(output_ptr, input_ptr); + assert_records(&retained, ptype, &bits, &mut ctx); + } else { + assert_eq!(output_ptr, input_ptr); + } + + let if_false = primitive_from_bits(ptype, &bits, Validity::NonNullable).into_array(); + let mask = Mask::from_iter([true, false, true, false, true, false]).into_array(); + let zipped = ::zip(patched.as_view(), &if_false, &mask, &mut ctx)? + .vortex_expect("mixed primitive zip uses the primitive kernel") + .execute::(&mut ctx)?; + assert_records( + &zipped, + ptype, + &[bits[0], bits[1], bits[2], bits[3], bits[0], bits[5]], + &mut ctx, + ); + + let nullable = PrimitiveArray::from_buffer_handle( + zipped.buffer_handle().clone(), + ptype, + Validity::from_iter([true, false, true, false, true, true]), + ); + let fill = if_false.execute_scalar(5, &mut ctx)?; + let filled = ::fill_null(nullable.as_view(), &fill, &mut ctx)? + .vortex_expect("mixed primitive validity uses the fill-null kernel") + .execute::(&mut ctx)?; + assert_records( + &filled, + ptype, + &[bits[0], bits[5], bits[2], bits[5], bits[0], bits[5]], + &mut ctx, + ); + + let interleaved = InterleaveArray::try_new( + vec![ + ConstantArray::new(fill, 3).into_array(), + filled.into_array(), + ], + PrimitiveArray::from_iter([0u8, 1, 0, 1]).into_array(), + PrimitiveArray::from_iter([2u32, 2, 0, 0]).into_array(), + )? + .into_array() + .execute::(&mut ctx)?; + assert_records( + &interleaved, + ptype, + &[bits[5], bits[2], bits[5], bits[0]], + &mut ctx, + ); + Ok(()) + } +} diff --git a/vortex-buffer/src/macros.rs b/vortex-buffer/src/macros.rs index eedbcd7f82d..079329cfb00 100644 --- a/vortex-buffer/src/macros.rs +++ b/vortex-buffer/src/macros.rs @@ -1,6 +1,77 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +/// Dispatches a byte width to a compile-time `const W: usize`. +/// +/// The default specialized widths are 1, 2, 4, 8, 16, and 32 bytes. Pass `widths = [...]` to +/// specialize a different set of nonzero powers of two. The input is a `usize` expression evaluated +/// once. Any other width, including zero, reaches the required fallback without rounding. +/// +/// The body and fallback are ordinary match arms: they can move values, return, and propagate +/// errors from the enclosing function. Both must produce the same result type. +/// +/// # Examples +/// +/// ``` +/// use vortex_buffer::match_each_power_of_two_width; +/// +/// let width = match_each_power_of_two_width!( +/// 8, +/// |W| { Ok::<_, usize>(size_of::<[u8; W]>()) }, +/// other => { Err(other) }, +/// ); +/// assert_eq!(width, Ok(8)); +/// +/// let width = match_each_power_of_two_width!( +/// 64, +/// widths = [16, 64], +/// |W| { W }, +/// other => { other }, +/// ); +/// assert_eq!(width, 64); +/// ``` +/// +/// Specializations must be nonzero powers of two: +/// +/// ```compile_fail +/// use vortex_buffer::match_each_power_of_two_width; +/// +/// match_each_power_of_two_width!(3, widths = [3], |W| { W }, other => { other }); +/// ``` +#[macro_export] +macro_rules! match_each_power_of_two_width { + ( + $byte_width:expr, + widths = [$($width:literal),+ $(,)?], + |$W:ident| $body:block, + $other:pat => $fallback:block $(,)? + ) => {{ + let byte_width: usize = $byte_width; + match byte_width { + $( + $width => { + const $W: usize = $width; + const { assert!($W.is_power_of_two(), "width must be a nonzero power of two") }; + $body + }, + )+ + $other => $fallback, + } + }}; + ( + $byte_width:expr, + |$W:ident| $body:block, + $other:pat => $fallback:block $(,)? + ) => { + $crate::match_each_power_of_two_width!( + $byte_width, + widths = [1, 2, 4, 8, 16, 32], + |$W| $body, + $other => $fallback, + ) + }; +} + /// A macro for constructing buffers akin to `vec![..]`. #[macro_export] macro_rules! buffer { diff --git a/vortex-buffer/tests/width_dispatch.rs b/vortex-buffer/tests/width_dispatch.rs new file mode 100644 index 00000000000..74196409b64 --- /dev/null +++ b/vortex-buffer/tests/width_dispatch.rs @@ -0,0 +1,95 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Public byte-width dispatch, including fallback and control-flow behavior. + +#[cfg(test)] +mod tests { + use rstest::rstest; + use vortex_buffer::match_each_power_of_two_width; + use vortex_error::VortexResult; + use vortex_error::vortex_bail; + use vortex_error::vortex_err; + + fn dispatch_with_error(width: usize, value: VortexResult) -> VortexResult { + let result = match_each_power_of_two_width!( + width, + |W| { value? + W }, + other => { vortex_bail!("unsupported record width {other}") }, + ); + Ok(result) + } + + #[rstest] + #[case(1)] + #[case(2)] + #[case(4)] + #[case(8)] + #[case(16)] + #[case(32)] + fn specializes_record_width(#[case] width: usize) { + let records = match_each_power_of_two_width!( + width, + |W| { vec![[7u8; W]; 2].into_iter().flatten().collect::>() }, + _ => { panic!("expected a specialized width") }, + ); + assert_eq!(records, vec![7; width * 2]); + } + + #[rstest] + #[case(0)] + #[case(3)] + #[case(12)] + #[case(64)] + fn preserves_unsupported_width(#[case] width: usize) { + let result = match_each_power_of_two_width!( + width, + |W| { Ok(W) }, + other => { Err(other) }, + ); + assert_eq!(result, Err(width)); + } + + #[rstest] + #[case(2, Err(2))] + #[case(8, Ok(8))] + #[case(64, Ok(64))] + fn restricts_specializations(#[case] width: usize, #[case] expected: Result) { + let result = match_each_power_of_two_width!( + width, + widths = [8, 64], + |W| { Ok(W) }, + other => { Err(other) }, + ); + assert_eq!(result, expected); + } + + #[test] + fn evaluates_input_and_selected_body_once() { + let mut input_calls = 0; + let mut body_calls = 0; + let value = String::from("selected"); + let result = match_each_power_of_two_width!( + { + input_calls += 1; + 4 + }, + |W| { + body_calls += 1; + (W, value) + }, + _ => { panic!("expected a specialized width") }, + ); + + assert_eq!((input_calls, body_calls), (1, 1)); + assert_eq!(result, (4, String::from("selected"))); + } + + #[test] + fn propagates_errors_to_enclosing_function() -> VortexResult<()> { + assert_eq!(dispatch_with_error(8, Ok(3))?, 11); + assert!(dispatch_with_error(8, Err(vortex_err!("payload error"))).is_err()); + assert!(dispatch_with_error(3, Ok(3)).is_err()); + Ok(()) + } +}