Consolidate fixed-width validation and width dispatch

Address review feedback on the shared fixed-width kernels:

- Document the contract of every FixedWidthArray hook.
- Validate the records buffer length once in a shared with_values helper
  instead of in each implementation, which are now pure constructors.
- Share the byte-width dispatch between take and filter through a single
  match_each_record_width macro.
- Fix the miri gate on the AVX2 tests, which was previously a no-op
  module-level ignore attribute.
- Test the arbitrary-width fallback paths of both take and filter, and
  restore 2-byte and float filter conformance coverage.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Signed-off-by: Connor Tsui <connor.tsui20@gmail.com>

Author: connortsui20

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

@@ -3,8 +3,6 @@
 
 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;
@@ -25,16 +23,9 @@ impl FixedWidthArray for Decimal {
     fn with_values(
         array: ArrayView<'_, Self>,
         values: ByteBuffer,
-        len: usize,
+        _len: usize,
         validity: Validity,
     ) -> VortexResult<DecimalArray> {
-        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(),

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

@@ -3,22 +3,60 @@
 
 use vortex_buffer::ByteBuffer;
 use vortex_error::VortexResult;
+use vortex_error::vortex_ensure_eq;
+use vortex_error::vortex_err;
 
 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.
+///
+/// The shared `take` and `filter` kernels view an implementing array as `array.len()` records of
+/// [`byte_width`] bytes each, which lets them move whole records without knowing the logical
+/// type.
+///
+/// [`byte_width`]: FixedWidthArray::byte_width
 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.
+    ///
+    /// The returned buffer must contain exactly `array.len() * byte_width` bytes.
     fn values(array: ArrayView<'_, Self>) -> ByteBuffer;
 
+    /// Rebuilds an array of this encoding from a records buffer, preserving the logical type of
+    /// `array`.
+    ///
+    /// Callers must provide a `values` buffer of exactly `len * byte_width` bytes and a
+    /// `validity` with logical length `len`. The shared kernels enforce the buffer length through
+    /// the module-level `with_values` helper rather than in each implementation.
     fn with_values(
         array: ArrayView<'_, Self>,
         values: ByteBuffer,
         len: usize,
         validity: Validity,
     ) -> VortexResult<Array<Self>>;
 }
+
+/// Rebuilds a fixed-width array from `len` records in `values`, validating the buffer length
+/// before dispatching to [`FixedWidthArray::with_values`].
+pub(crate) fn with_values<V: FixedWidthArray>(
+    array: ArrayView<'_, V>,
+    values: ByteBuffer,
+    len: usize,
+    validity: Validity,
+) -> VortexResult<Array<V>> {
+    let expected_len = len
+        .checked_mul(V::byte_width(array))
+        .ok_or_else(|| vortex_err!("Fixed-width values buffer length overflows usize"))?;
+    vortex_ensure_eq!(
+        values.len(),
+        expected_len,
+        "Fixed-width values buffer length {} does not match expected length {expected_len}",
+        values.len(),
+    );
+    V::with_values(array, values, len, validity)
+}

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

@@ -10,6 +10,8 @@ use vortex_error::VortexExpect;
 use vortex_mask::MaskValues;
 
 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_buffer_byte_compress;
@@ -28,48 +30,47 @@ pub(crate) fn filter<V: FixedWidthArray>(array: &Array<V>, mask: &Arc<MaskValues
             .vortex_expect("fixed-width validity should be derivable"),
         mask,
     );
-    V::with_values(array, values, mask.true_count(), validity)
+    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)
+    match_each_record_width!(
+        byte_width,
+        |W| {
+            let records = Buffer::<[u8; W]>::from_byte_buffer(values);
+            // Byte-compress processes eight records per mask byte and wins for narrow records,
+            // while wider records move enough bytes each for the plain filter to win.
+            let filtered = if W <= 4 {
+                filter_buffer_byte_compress(records, mask)
+            } else {
+                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)
+                }
+            }
         }
-    }
+    )
 }

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

@@ -2,8 +2,11 @@
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
 use rstest::rstest;
+use vortex_buffer::Buffer;
 use vortex_buffer::buffer;
+use vortex_mask::Mask;
 
+use super::filter_records;
 use crate::ArrayRef;
 use crate::IntoArray;
 use crate::VortexSessionExecute;
@@ -17,9 +20,30 @@ use crate::dtype::DecimalDType;
 use crate::dtype::i256;
 use crate::validity::Validity;
 
+#[test]
+fn filter_fallback_width_records() {
+    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);
+}
+
 #[rstest]
 #[case::primitive_i8(PrimitiveArray::from_iter([-2i8, -1, 0, 1, 2]).into_array())]
+#[case::primitive_u16(PrimitiveArray::from_iter([1u16, 2, 3, 4, 5]).into_array())]
 #[case::primitive_i32(PrimitiveArray::from_iter([1i32, 2, 3, 4, 5]).into_array())]
+#[case::primitive_f32(PrimitiveArray::from_iter([0.1f32, 0.2, 0.3, 0.4, 0.5]).into_array())]
 #[case::primitive_nullable(PrimitiveArray::from_option_iter(
     [Some(1i64), None, Some(3), Some(4), None],
 ).into_array())]

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

@@ -9,3 +9,39 @@ pub(crate) mod take;
 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/avx2/tests.rs

@@ -1,8 +1,7 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
-#![cfg_attr(miri, ignore)]
-#![cfg(target_arch = "x86_64")]
+#![cfg(all(target_arch = "x86_64", not(miri)))]
 
 use super::take_avx2;
 

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

@@ -22,6 +22,7 @@ use self::scalar::take_values_scalar;
 use self::slices::take_slices;
 use self::slices::take_slices_constant_length;
 use super::FixedWidthArray;
+use super::with_values;
 use crate::ArrayRef;
 use crate::Columnar;
 use crate::ExecutionCtx;
@@ -144,7 +145,7 @@ pub(crate) fn take<V: FixedWidthArray>(
         )
     })?;
     Ok(Some(
-        V::with_values(array, values, indices.len(), validity)?.into_array(),
+        with_values(array, values, indices.len(), validity)?.into_array(),
     ))
 }
 
@@ -193,6 +194,6 @@ fn take_contiguous_ranges<V: FixedWidthArray>(
     }?;
     let validity = array.validity()?.take(indices_ref)?;
     Ok(Some(
-        V::with_values(array, taken, output_len, validity)?.into_array(),
+        with_values(array, taken, output_len, validity)?.into_array(),
     ))
 }

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

@@ -8,6 +8,7 @@ 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;
 
 pub(super) fn take_byte_records<I: UnsignedPType>(
@@ -18,39 +19,31 @@ pub(super) fn take_byte_records<I: UnsignedPType>(
 ) -> VortexResult<ByteBuffer> {
     let alignment = values.alignment();
 
-    macro_rules! take_typed_records {
-        ($width:expr) => {{
-            let records = Buffer::<[u8; $width]>::from_byte_buffer(values.clone());
+    match_each_record_width!(
+        byte_width,
+        |W| {
+            let records = Buffer::<[u8; W]>::from_byte_buffer(values.clone());
             debug_assert_eq!(records.len(), record_count);
             Ok(take_values(records.as_slice(), indices)
                 .into_byte_buffer()
                 .aligned(alignment))
-        }};
-    }
-
-    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::<u8>::with_capacity(output_len);
-    for index in indices {
-        let index = index.as_();
-        assert!(
-            index < record_count,
-            "take index {index} out of bounds for length {record_count}"
-        );
-        let start = index * byte_width;
-        result.extend_from_slice(&values[start..start + byte_width]);
-    }
-    Ok(result.freeze().into_byte_buffer().aligned(alignment))
+        },
+        _ => {
+            let output_len = indices
+                .len()
+                .checked_mul(byte_width)
+                .ok_or_else(|| vortex_err!("Fixed-width take output length overflows usize"))?;
+            let mut result = BufferMut::<u8>::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))
+        }
+    )
 }

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

@@ -45,6 +45,8 @@ fn take_eight_byte_values() {
 #[case(8)]
 #[case(16)]
 #[case(32)]
+#[case::fallback(3)]
+#[case::fallback_wide(12)]
 fn take_runtime_width_records(#[case] byte_width: usize) -> VortexResult<()> {
     let values = Buffer::from_iter((0u8..).take(3 * byte_width));
     let expected = values[2 * byte_width..3 * byte_width]
@@ -57,6 +59,13 @@ fn take_runtime_width_records(#[case] byte_width: usize) -> VortexResult<()> {
     Ok(())
 }
 
+#[test]
+#[should_panic(expected = "take index 3 out of bounds for length 3")]
+fn fallback_take_rejects_out_of_bounds_index() {
+    let values = Buffer::from_iter((0u8..).take(9)).into_byte_buffer();
+    drop(take_byte_records(&values, 3, 3, &[3u32]));
+}
+
 #[test]
 fn take_variable_length_slices() -> VortexResult<()> {
     let values = buffer![10u8, 11, 12, 13, 14].into_byte_buffer();