Perform fixed width array operations in type width units instead of byte units

vortex-array/src/arrays/decimal/compute/fixed_width.rs

@@ -1,6 +1,7 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
+use vortex_buffer::Buffer;
 use vortex_buffer::ByteBuffer;
 use vortex_error::VortexResult;
 
@@ -16,8 +17,10 @@ impl FixedWidthArray for Decimal {
         array.values_type().byte_width()
     }
 
-    fn values(array: ArrayView<'_, Self>) -> ByteBuffer {
-        array.buffer_handle().to_host_sync()
+    fn values<T: Copy>(array: ArrayView<'_, Self>) -> Buffer<T> {
+        let values = array.buffer_handle().to_host_sync();
+        let alignment = values.alignment();
+        Buffer::from_byte_buffer_aligned(values, alignment)
     }
 
     fn with_values(

vortex-array/src/arrays/fixed_width/array.rs

@@ -1,6 +1,7 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
+use vortex_buffer::Buffer;
 use vortex_buffer::ByteBuffer;
 use vortex_error::VortexResult;
 use vortex_error::vortex_ensure_eq;
@@ -22,10 +23,11 @@ pub(crate) trait FixedWidthArray: VTable {
     /// Returns the number of bytes each record occupies.
     fn byte_width(array: ArrayView<'_, Self>) -> usize;
 
-    /// Returns the records of `array` as a single host-resident byte buffer.
+    /// Returns the storage of `array` as a host-resident buffer of `T`, preserving its alignment.
     ///
     /// The returned buffer must contain exactly `array.len() * byte_width` bytes.
-    fn values(array: ArrayView<'_, Self>) -> ByteBuffer;
+    /// Its byte length and alignment must be compatible with `T`.
+    fn values<T: Copy>(array: ArrayView<'_, Self>) -> Buffer<T>;
 
     /// Rebuilds an array of this encoding from a records buffer, preserving the logical type of
     /// `array`.

vortex-array/src/arrays/fixed_width/filter.rs

@@ -1,27 +1,35 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
-use vortex_buffer::Buffer;
-use vortex_buffer::BufferMut;
 use vortex_buffer::ByteBuffer;
 use vortex_error::VortexExpect;
+use vortex_error::vortex_panic;
 use vortex_mask::MaskValues;
 use vortex_mask::MaskValuesRef;
 
 use super::FixedWidthArray;
-use super::match_each_record_width;
 use super::with_values;
 use crate::array::Array;
+use crate::array::ArrayView;
 use crate::arrays::filter::filter_buffer;
 use crate::arrays::filter::filter_validity;
+use crate::dtype::i256;
 
 #[cfg(test)]
 #[expect(clippy::cast_possible_truncation)]
 mod tests;
 
 pub(crate) fn filter<V: FixedWidthArray>(array: &Array<V>, mask: &MaskValuesRef) -> Array<V> {
     let array = array.as_view();
-    let values = filter_records(V::values(array), V::byte_width(array), mask.as_ref());
+    let values = match V::byte_width(array) {
+        1 => filter_records::<V, u8>(array, mask.as_ref()),
+        2 => filter_records::<V, u16>(array, mask.as_ref()),
+        4 => filter_records::<V, u32>(array, mask.as_ref()),
+        8 => filter_records::<V, u64>(array, mask.as_ref()),
+        16 => filter_records::<V, u128>(array, mask.as_ref()),
+        32 => filter_records::<V, i256>(array, mask.as_ref()),
+        byte_width => vortex_panic!("Unsupported fixed-width byte width: {byte_width}"),
+    };
     let validity = filter_validity(
         array
             .validity()
@@ -32,39 +40,15 @@ pub(crate) fn filter<V: FixedWidthArray>(array: &Array<V>, mask: &MaskValuesRef)
         .vortex_expect("filtering fixed-width values preserves array invariants")
 }
 
-fn filter_records(values: ByteBuffer, byte_width: usize, mask: &MaskValues) -> ByteBuffer {
+fn filter_records<V: FixedWidthArray, T: Copy>(
+    array: ArrayView<'_, V>,
+    mask: &MaskValues,
+) -> ByteBuffer {
+    let values = V::values::<T>(array);
     let alignment = values.alignment();
-
-    match_each_record_width!(
-        byte_width,
-        |W| {
-            let records = Buffer::<[u8; W]>::from_byte_buffer(values);
-            // `filter_buffer` picks between in-place compaction, cached indices/slices,
-            // byte-compress, and bitmap iteration based on record width and mask density.
-            let filtered = filter_buffer(records, mask);
-            filtered.into_byte_buffer().aligned(alignment)
-        },
-        _ => {
-            match values.try_into_mut() {
-                Ok(mut values) => {
-                    let mut destination = 0;
-                    mask.bit_buffer().for_each_set_index(|index| {
-                        let source = index * byte_width;
-                        values.copy_within(source..source + byte_width, destination);
-                        destination += byte_width;
-                    });
-                    values.truncate(destination);
-                    values.freeze().into_byte_buffer().aligned(alignment)
-                }
-                Err(values) => {
-                    let mut filtered = BufferMut::with_capacity(mask.true_count() * byte_width);
-                    mask.bit_buffer().for_each_set_index(|index| {
-                        let start = index * byte_width;
-                        filtered.extend_from_slice(&values[start..start + byte_width]);
-                    });
-                    filtered.freeze().into_byte_buffer().aligned(alignment)
-                }
-            }
-        }
-    )
+    // `filter_buffer` picks between in-place compaction, cached indices/slices,
+    // byte-compress, and bitmap iteration based on record width and mask density.
+    filter_buffer(values, mask)
+        .into_byte_buffer()
+        .aligned(alignment)
 }

vortex-array/src/arrays/fixed_width/filter/tests.rs

@@ -2,41 +2,61 @@
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
 use rstest::rstest;
+use vortex_buffer::Alignment;
 use vortex_buffer::Buffer;
 use vortex_buffer::buffer;
+use vortex_error::VortexResult;
 use vortex_mask::Mask;
 
-use super::filter_records;
 use crate::ArrayRef;
 use crate::IntoArray;
 use crate::VortexSessionExecute;
+use crate::array::Array;
 use crate::array_session;
 use crate::arrays::DecimalArray;
 use crate::arrays::PrimitiveArray;
+use crate::arrays::fixed_width::FixedWidthArray;
+use crate::arrays::fixed_width::with_values;
 use crate::compute::conformance::filter::LARGE_SIZE;
 use crate::compute::conformance::filter::MEDIUM_SIZE;
 use crate::compute::conformance::filter::test_filter_conformance;
 use crate::dtype::DecimalDType;
 use crate::dtype::i256;
 use crate::validity::Validity;
 
-#[test]
-fn filter_fallback_width_records() {
+#[rstest]
+#[case::u8(PrimitiveArray::from_iter([10u8, 20, 30, 40]))]
+#[case::u16(PrimitiveArray::from_iter([10u16, 20, 30, 40]))]
+#[case::u32(PrimitiveArray::from_iter([10u32, 20, 30, 40]))]
+#[case::u64(PrimitiveArray::from_iter([10u64, 20, 30, 40]))]
+#[case::u128(DecimalArray::new(
+    buffer![10i128, -20, -30, -40],
+    DecimalDType::new(19, 0),
+    Validity::NonNullable,
+))]
+#[case::i256(DecimalArray::new(
+    Buffer::from_iter([10, -20, -30, -40].map(|value| i256::from_parts(1, value))),
+    DecimalDType::new(76, 0),
+    Validity::NonNullable,
+))]
+fn filter_typed_records<V: FixedWidthArray>(#[case] array: Array<V>) -> VortexResult<()> {
     let Mask::Values(mask) = Mask::from_iter([true, false, true, false]) else {
         panic!("a mixed mask must have mask values");
     };
-    let expected = [0u8, 1, 2, 6, 7, 8];
-
-    // A uniquely owned buffer takes the in-place `copy_within` path.
-    let owned = Buffer::from_iter(0u8..12);
-    let filtered = filter_records(owned, 3, &mask);
-    assert_eq!(filtered.as_slice(), &expected);
-
-    // Retaining a second reference forces the copying path instead.
-    let shared = Buffer::from_iter(0u8..12);
-    let _retained = shared.clone();
-    let filtered = filter_records(shared, 3, &mask);
-    assert_eq!(filtered.as_slice(), &expected);
+    let byte_width = V::byte_width(array.as_view());
+    let values = V::values::<u8>(array.as_view()).aligned(Alignment::new(64));
+    let expected = values[..byte_width]
+        .iter()
+        .chain(&values[2 * byte_width..3 * byte_width])
+        .copied()
+        .collect::<Vec<_>>();
+    let array = with_values(array.as_view(), values, array.len(), Validity::NonNullable)?;
+    let alignment = V::values::<u8>(array.as_view()).alignment();
+    let filtered = super::filter(&array, &mask).into_array();
+    let buffers = filtered.buffers();
+    assert_eq!(buffers[0].as_slice(), expected);
+    assert_eq!(buffers[0].alignment(), alignment);
+    Ok(())
 }
 
 #[rstest]

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;

vortex-array/src/arrays/fixed_width/take/mod.rs

@@ -15,9 +15,10 @@ use std::sync::LazyLock;
 use vortex_buffer::Buffer;
 use vortex_error::VortexResult;
 use vortex_error::vortex_bail;
+use vortex_error::vortex_ensure;
 use vortex_mask::Mask;
 
-use self::records::take_byte_records;
+use self::records::take_records;
 use self::scalar::take_values_scalar;
 use self::slices::take_slices;
 use self::slices::take_slices_constant_length;
@@ -38,6 +39,7 @@ use crate::builtins::ArrayBuiltins;
 use crate::dtype::DType;
 use crate::dtype::UnsignedPType;
 use crate::dtype::half::f16;
+use crate::dtype::i256;
 use crate::match_each_unsigned_integer_ptype;
 use crate::scalar::Scalar;
 
@@ -71,7 +73,9 @@ macro_rules! impl_fixed_width_take_value {
     };
 }
 
-impl_fixed_width_take_value!(u8, u16, u32, u64, i8, i16, i32, i64, f16, f32, f64,);
+impl_fixed_width_take_value!(
+    u8, u16, u32, u64, u128, i8, i16, i32, i64, i256, f16, f32, f64,
+);
 
 // SAFETY: Byte arrays have no padding and every byte is initialized.
 unsafe impl<const N: usize> FixedWidthTakeValue for [u8; N] {}
@@ -96,10 +100,19 @@ pub(crate) fn take<V: FixedWidthArray>(
     indices: &ArrayRef,
     ctx: &mut ExecutionCtx,
 ) -> VortexResult<Option<ArrayRef>> {
-    if let Some(piecewise_indices) = indices.as_opt::<PiecewiseSequence>()
-        && let Some(taken) = take_contiguous_ranges(array, piecewise_indices, indices, ctx)?
-    {
-        return Ok(Some(taken));
+    if let Some(piecewise_indices) = indices.as_opt::<PiecewiseSequence>() {
+        let taken = match V::byte_width(array) {
+            1 => take_contiguous_ranges::<V, u8>(array, piecewise_indices, indices, ctx)?,
+            2 => take_contiguous_ranges::<V, u16>(array, piecewise_indices, indices, ctx)?,
+            4 => take_contiguous_ranges::<V, u32>(array, piecewise_indices, indices, ctx)?,
+            8 => take_contiguous_ranges::<V, u64>(array, piecewise_indices, indices, ctx)?,
+            16 => take_contiguous_ranges::<V, u128>(array, piecewise_indices, indices, ctx)?,
+            32 => take_contiguous_ranges::<V, i256>(array, piecewise_indices, indices, ctx)?,
+            _ => None,
+        };
+        if taken.is_some() {
+            return Ok(taken);
+        }
     }
 
     let DType::Primitive(ptype, nullability) = indices.dtype() else {
@@ -135,21 +148,23 @@ pub(crate) fn take<V: FixedWidthArray>(
         .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::<I>(),
-        )
-    })?;
+    let values = match V::byte_width(array) {
+        1 => take_records::<V, u8>(array, &indices)?,
+        2 => take_records::<V, u16>(array, &indices)?,
+        4 => take_records::<V, u32>(array, &indices)?,
+        8 => take_records::<V, u64>(array, &indices)?,
+        16 => take_records::<V, u128>(array, &indices)?,
+        32 => take_records::<V, i256>(array, &indices)?,
+        _ => return Ok(None),
+    };
     Ok(Some(
         with_values(array, values, indices.len(), validity)?.into_array(),
     ))
 }
 
-fn take_contiguous_ranges<V: FixedWidthArray>(
+// Avoid duplicating the starts/lengths dispatch in every record-width arm of `take`.
+#[inline(never)]
+fn take_contiguous_ranges<V: FixedWidthArray, T: Copy>(
     array: ArrayView<'_, V>,
     indices: ArrayView<'_, PiecewiseSequence>,
     indices_ref: &ArrayRef,
@@ -159,21 +174,17 @@ fn take_contiguous_ranges<V: FixedWidthArray>(
         return Ok(None);
     };
 
-    let values = V::values(array);
-    let byte_width = V::byte_width(array);
+    let values = V::values::<T>(array);
+    vortex_ensure!(
+        values.len() == array.len(),
+        "Fixed-width values buffer length does not match record count"
+    );
     let output_len = indices_ref.len();
     let taken = match lengths {
         Columnar::Constant(lengths) => {
             let length = constant_unsigned_usize(&lengths);
             match_each_unsigned_integer_ptype!(starts.ptype(), |S| {
-                take_slices_constant_length(
-                    &values,
-                    byte_width,
-                    array.len(),
-                    starts.as_slice::<S>(),
-                    length,
-                    output_len,
-                )
+                take_slices_constant_length(&values, starts.as_slice::<S>(), length, output_len)
             })
         }
         Columnar::Canonical(lengths) => {
@@ -182,8 +193,6 @@ fn take_contiguous_ranges<V: FixedWidthArray>(
                 match_each_unsigned_integer_ptype!(lengths.ptype(), |L| {
                     take_slices(
                         &values,
-                        byte_width,
-                        array.len(),
                         starts.as_slice::<S>(),
                         lengths.as_slice::<L>(),
                         output_len,
@@ -194,6 +203,6 @@ fn take_contiguous_ranges<V: FixedWidthArray>(
     }?;
     let validity = array.validity()?.take(indices_ref)?;
     Ok(Some(
-        with_values(array, taken, output_len, validity)?.into_array(),
+        with_values(array, taken.into_byte_buffer(), output_len, validity)?.into_array(),
     ))
 }