Dispatch structural kernels by physical width

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<T: NativePType, const WIDTH: usize>(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<T: Copy + PartialEq, const WIDTH: usize>(values: &[T]) -> bool {
     let first_value = values[0];
     let first_vec = &[first_value; WIDTH];
 
@@ -28,32 +30,51 @@ pub fn compute_is_constant<T: NativePType, const WIDTH: usize>(values: &[T]) ->
     }
 
     for value in remainder {
-        if !value.is_eq(first_value) {
+        if *value != first_value {
             return false;
         }
     }
 
     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::<W>().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::<P>()}>(array.as_slice::<P>())
-    }, floating: |P| {
-        compute_is_constant::<_, {IS_CONST_LANE_WIDTH / size_of::<P>()}>(unsafe { std::mem::transmute::<&[P], &[<P as EqFloat>::IntType]>(array.as_slice::<P>()) })
-    })
+#[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);
+    }
 }

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<ArrayRef> {
     let child = array.child().clone().execute::<PrimitiveArray>(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::<W>(
+                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<ArrayRef> {
     let child = array.child().clone().execute::<DecimalArray>(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::<T>() });
+    let values = child.buffer_handle().to_host_sync();
+    let (taken, output_validity) = match_each_decimal_type_width!(
+        values_type,
+        |W| {
+            take_fixed_width_records::<W>(
+                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::<u8, P>(child, child_validity, filter, indices, ctx),
-            2 => take_fixed_width_as::<u16, P>(child, child_validity, filter, indices, ctx),
-            4 => take_fixed_width_as::<u32, P>(child, child_validity, filter, indices, ctx),
-            8 => take_fixed_width_as::<u64, P>(child, child_validity, filter, indices, ctx),
-            16 => take_fixed_width_as::<u128, P>(child, child_validity, filter, indices, ctx),
-            32 => take_fixed_width_as::<i256, P>(child, child_validity, filter, indices, ctx),
-            _ => vortex_panic!("Unsupported fixed width of {byte_width} bytes"),
-        }
-    })
-}
-
-fn take_fixed_width_as<W, P>(
+/// `alignment` must divide `W` so slices between records remain aligned.
+fn take_fixed_width_records<const W: usize>(
     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::<W, P>(
-        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<T, P>(

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<VortexSession> = 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| {
+            <T as BigCast>::from(value).vortex_expect("small values fit every decimal type")
+        }));
+        let alignment = Alignment::of::<T>();
+        let mut bytes =
+            BufferMut::<u8>::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| {
+                <T as BigCast>::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::<DecimalArray>(&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(

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);