feat(vortex-array): execute primitive interleave arrays

vortex-array/src/arrays/interleave/execute/bool.rs

@@ -7,10 +7,10 @@ use num_traits::AsPrimitive;
 use vortex_buffer::BitBuffer;
 use vortex_buffer::BitBufferMut;
 use vortex_error::VortexResult;
-use vortex_error::vortex_ensure;
 
 use super::super::Interleave;
 use super::super::InterleaveArrayExt;
+use super::validate_selectors;
 use crate::array::Array;
 use crate::arrays::Bool;
 use crate::arrays::BoolArray;
@@ -71,46 +71,18 @@ fn gather<A: AsPrimitive<usize>, R: AsPrimitive<usize>>(
     branches: &[A],
     rows: &[R],
 ) -> VortexResult<BitBufferMut> {
-    let len = validate_selectors(value_bits, branches, rows)?;
+    let len = validate_selectors(
+        value_bits.len(),
+        |branch| value_bits[branch].len(),
+        branches,
+        rows,
+    )?;
 
     // SAFETY: `validate_selectors` proved `branches.len() == rows.len() == len`, and for every
     // `i < len` that `branches[i] < value_bits.len()` and `rows[i] < value_bits[branches[i]].len()`.
     Ok(unsafe { gather_bits(len, value_bits, branches, rows) })
 }
 
-/// Validates the per-row selector bounds, returning the output length (`branches.len()`).
-///
-/// On success, `rows.len() == branches.len() == len` and, for every `i < len`,
-/// `branches[i] < value_bits.len()` and `rows[i] < value_bits[branches[i]].len()` — exactly the
-/// preconditions of [`gather_bits`]. Errors (rather than panics) on any out-of-bounds selector.
-fn validate_selectors<A: AsPrimitive<usize>, R: AsPrimitive<usize>>(
-    value_bits: &[BitBuffer],
-    branches: &[A],
-    rows: &[R],
-) -> VortexResult<usize> {
-    // The two selectors are validated to equal length at construction, which is the output length.
-    let len = branches.len();
-    vortex_ensure!(
-        rows.len() == len,
-        "interleave selectors differ in length: array_indices {len}, row_indices {}",
-        rows.len()
-    );
-
-    for i in 0..len {
-        let branch = branches[i].as_();
-        vortex_ensure!(
-            branch < value_bits.len(),
-            "interleave array index out of bounds"
-        );
-        vortex_ensure!(
-            rows[i].as_() < value_bits[branch].len(),
-            "interleave row index out of bounds"
-        );
-    }
-
-    Ok(len)
-}
-
 /// Gathers one bit per output from `bits[branches[i]]` at position `rows[i]`, packing 64 results per
 /// word with [`BitBufferMut::collect_bool`].
 ///

vortex-array/src/arrays/interleave/execute/mod.rs

@@ -5,14 +5,16 @@
 //!
 //! All values share a type (validated in [`Interleave::check`]), so the
 //! physical gather kernel is chosen from the first value. The selector types are an orthogonal
-//! concern handled within each kernel. Only boolean values are implemented today (see the [`bool`] module).
+//! concern handled within each kernel.
 //!
 //! [`Interleave::check`]: super::Interleave::check
-//! [`bool`]: module@crate::arrays::interleave::execute::bool
 
 mod bool;
+mod primitive;
 
+use num_traits::AsPrimitive;
 use vortex_error::VortexResult;
+use vortex_error::vortex_ensure;
 use vortex_error::vortex_panic;
 
 use super::Interleave;
@@ -28,12 +30,47 @@ pub(super) fn execute(
 ) -> VortexResult<ExecutionResult> {
     if array.value(0).dtype().is_boolean() {
         bool::execute(array, ctx)
+    } else if array.value(0).dtype().is_primitive() {
+        primitive::execute(array, ctx)
     } else {
         let value_dtype = array.value(0).dtype().clone();
         vortex_panic!(
-            "interleave execution is only implemented for boolean values; value dtype {} is not \
-             yet supported",
+            "interleave execution is not implemented for value dtype {}",
             value_dtype
         )
     }
 }
+
+/// Validate selector lengths and bounds, returning the common output length.
+///
+/// On success, `branches.len() == rows.len() == len`; for every `i < len`,
+/// `branches[i] < num_values` and `rows[i] < value_len(branches[i])`.
+fn validate_selectors<A, R, F>(
+    num_values: usize,
+    value_len: F,
+    branches: &[A],
+    rows: &[R],
+) -> VortexResult<usize>
+where
+    A: AsPrimitive<usize>,
+    R: AsPrimitive<usize>,
+    F: Fn(usize) -> usize,
+{
+    let len = branches.len();
+    vortex_ensure!(
+        rows.len() == len,
+        "interleave selectors differ in length: array_indices {len}, row_indices {}",
+        rows.len()
+    );
+
+    for i in 0..len {
+        let branch = branches[i].as_();
+        vortex_ensure!(branch < num_values, "interleave array index out of bounds");
+        vortex_ensure!(
+            rows[i].as_() < value_len(branch),
+            "interleave row index out of bounds"
+        );
+    }
+
+    Ok(len)
+}

vortex-array/src/arrays/interleave/execute/primitive.rs

@@ -0,0 +1,163 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+//! Execution for primitive [`Interleave`] values.
+
+use num_traits::AsPrimitive;
+use vortex_buffer::Buffer;
+use vortex_buffer::BufferMut;
+use vortex_error::VortexResult;
+
+use super::super::Interleave;
+use super::super::InterleaveArrayExt;
+use super::validate_selectors;
+use crate::array::Array;
+use crate::array::ArrayView;
+use crate::arrays::Constant;
+use crate::arrays::Primitive;
+use crate::arrays::PrimitiveArray;
+use crate::arrays::primitive::PrimitiveArrayExt;
+use crate::dtype::NativePType;
+use crate::executor::ExecutionCtx;
+use crate::executor::ExecutionResult;
+use crate::match_each_native_ptype;
+use crate::match_each_unsigned_integer_ptype;
+use crate::require_child;
+
+pub(super) fn execute(
+    mut array: Array<Interleave>,
+    _ctx: &mut ExecutionCtx,
+) -> VortexResult<ExecutionResult> {
+    let num_values = array.num_values();
+    array = require_child!(array, array.array_indices(), 0 => Primitive);
+    array = require_child!(array, array.row_indices(), 1 => Primitive);
+    for i in 0..num_values {
+        if array.value(i).as_opt::<Constant>().is_none() {
+            array = require_child!(array, array.value(i), i + 2 => Primitive);
+        }
+    }
+
+    let validity = array.as_ref().validity()?;
+    let output = match_each_native_ptype!(array.value(0).dtype().as_ptype(), |T| {
+        let values = gather_values::<T>(&array)?;
+        VortexResult::Ok(PrimitiveArray::new(values, validity))
+    })?;
+
+    Ok(ExecutionResult::done(output))
+}
+
+/// Physical primitive values; nullness remains in the source array's validity.
+enum PrimitiveValues<T> {
+    Buffer(Buffer<T>),
+    Constant { value: T, len: usize },
+}
+
+impl<T: Copy> PrimitiveValues<T> {
+    fn len(&self) -> usize {
+        match self {
+            Self::Buffer(values) => values.len(),
+            Self::Constant { len, .. } => *len,
+        }
+    }
+
+    /// Returns the physical value at `index` without bounds checking.
+    ///
+    /// # Safety
+    ///
+    /// `index` must be less than [`Self::len`].
+    unsafe fn value_unchecked(&self, index: usize) -> T {
+        match self {
+            // SAFETY: the caller guarantees that `index` is in bounds.
+            Self::Buffer(values) => *unsafe { values.get_unchecked(index) },
+            Self::Constant { value, .. } => *value,
+        }
+    }
+}
+
+fn gather_values<T: NativePType>(array: &Array<Interleave>) -> VortexResult<Buffer<T>> {
+    let values = (0..array.num_values())
+        .map(|i| {
+            let value = array.value(i);
+            if let Some(constant) = value.as_opt::<Constant>() {
+                PrimitiveValues::Constant {
+                    value: constant
+                        .scalar()
+                        .as_primitive()
+                        .typed_value::<T>()
+                        // Validity carries nullness; a null constant's payload is never observed.
+                        .unwrap_or_default(),
+                    len: value.len(),
+                }
+            } else {
+                PrimitiveValues::Buffer(value.as_::<Primitive>().to_buffer::<T>())
+            }
+        })
+        .collect::<Vec<_>>();
+    let branches = array.array_indices().as_::<Primitive>();
+    let rows = array.row_indices().as_::<Primitive>();
+
+    match_each_unsigned_integer_ptype!(branches.ptype(), |A| {
+        gather_rows::<T, A>(&values, branches.as_slice::<A>(), rows)
+    })
+}
+
+fn gather_rows<T, A>(
+    values: &[PrimitiveValues<T>],
+    branches: &[A],
+    rows: ArrayView<'_, Primitive>,
+) -> VortexResult<Buffer<T>>
+where
+    T: NativePType,
+    A: AsPrimitive<usize>,
+{
+    match_each_unsigned_integer_ptype!(rows.ptype(), |R| {
+        gather(values, branches, rows.as_slice::<R>())
+    })
+}
+
+fn gather<T, A, R>(
+    values: &[PrimitiveValues<T>],
+    branches: &[A],
+    rows: &[R],
+) -> VortexResult<Buffer<T>>
+where
+    T: NativePType,
+    A: AsPrimitive<usize>,
+    R: AsPrimitive<usize>,
+{
+    let len = validate_selectors(values.len(), |branch| values[branch].len(), branches, rows)?;
+
+    // SAFETY: `validate_selectors` proved `branches.len() == rows.len() == len`, and for every
+    // `i < len` that `branches[i] < values.len()` and `rows[i] < values[branches[i]].len()`.
+    Ok(unsafe { gather_unchecked(len, values, branches, rows) })
+}
+
+/// Gathers one primitive value per output from `values[branches[i]]` at position `rows[i]`.
+///
+/// # Safety
+///
+/// `branches` and `rows` must both contain at least `len` elements. For every `i < len`,
+/// `branches[i] < values.len()` and `rows[i] < values[branches[i]].len()`.
+unsafe fn gather_unchecked<T, A, R>(
+    len: usize,
+    values: &[PrimitiveValues<T>],
+    branches: &[A],
+    rows: &[R],
+) -> Buffer<T>
+where
+    T: NativePType,
+    A: AsPrimitive<usize>,
+    R: AsPrimitive<usize>,
+{
+    let mut output = BufferMut::with_capacity(len);
+    for i in 0..len {
+        // SAFETY: the caller guarantees `i` is in bounds for both selectors, and that the selected
+        // branch and row are in bounds for `values` and the selected physical value buffer.
+        output.push(unsafe {
+            values
+                .get_unchecked(branches.get_unchecked(i).as_())
+                .value_unchecked(rows.get_unchecked(i).as_())
+        });
+    }
+    output.freeze()
+}

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

@@ -463,6 +463,7 @@ mod tests {
     use crate::arrays::BoolArray;
     use crate::arrays::PrimitiveArray;
     use crate::assert_arrays_eq;
+    use crate::dtype::PType;
 
     /// Reference (oracle) implementation of the interleave spec, used only to validate the optimized
     /// [execute](super::execute) path. It is intentionally simple and slow: it pulls each output
@@ -719,17 +720,55 @@ mod tests {
     }
 
     #[test]
-    #[should_panic(expected = "only implemented for boolean values")]
-    fn non_boolean_value_execution_panics() {
-        // Execution dispatches on the value type: primitive values have no kernel yet.
-        let v0 = PrimitiveArray::from_iter([1u32]).into_array();
-        let v1 = PrimitiveArray::from_iter([2u32]).into_array();
-        let array_indices = PrimitiveArray::from_iter([0u32, 1]).into_array();
-        let row_indices = PrimitiveArray::from_iter([0u32, 0]).into_array();
-        let interleaved = InterleaveArray::try_new(vec![v0, v1], array_indices, row_indices)
-            .vortex_expect("primitive values should construct")
-            .into_array();
+    fn executes_primitive_values() -> VortexResult<()> {
+        let v0 = PrimitiveArray::from_iter([1.0f64, 2.0]).into_array();
+        let v1 = PrimitiveArray::from_option_iter([Some(10.0f64), None]).into_array();
+        let array_indices = PrimitiveArray::from_iter([0u8, 1, 0, 1]).into_array();
+        let row_indices = PrimitiveArray::from_iter([0u32, 0, 1, 1]).into_array();
+        let interleaved =
+            InterleaveArray::try_new(vec![v0, v1], array_indices, row_indices)?.into_array();
+        let expected =
+            PrimitiveArray::from_option_iter([Some(1.0f64), Some(10.0), Some(2.0), None])
+                .into_array();
+        let mut ctx = array_session().create_execution_ctx();
+        assert_arrays_eq!(interleaved, expected, &mut ctx);
+        Ok(())
+    }
+
+    #[test]
+    fn executes_primitive_constant_values() -> VortexResult<()> {
+        let constant = ConstantArray::new(1.0f64, 2).into_array();
+        let column = PrimitiveArray::from_iter([10.0f64, 20.0]).into_array();
+        let array_indices = PrimitiveArray::from_iter([0u8, 1, 0, 1]).into_array();
+        let row_indices = PrimitiveArray::from_iter([0u32, 0, 1, 1]).into_array();
+        let interleaved =
+            InterleaveArray::try_new(vec![constant, column], array_indices, row_indices)?
+                .into_array();
+        let expected = PrimitiveArray::from_iter([1.0f64, 10.0, 1.0, 20.0]).into_array();
         let mut ctx = array_session().create_execution_ctx();
-        interleaved.execute::<Canonical>(&mut ctx).ok();
+
+        assert_arrays_eq!(interleaved, expected, &mut ctx);
+        Ok(())
+    }
+
+    #[test]
+    fn executes_null_primitive_constant_values() -> VortexResult<()> {
+        let constant = ConstantArray::new(
+            Scalar::null(DType::Primitive(PType::F64, Nullability::Nullable)),
+            2,
+        )
+        .into_array();
+        let column = PrimitiveArray::from_iter([10.0f64, 20.0]).into_array();
+        let array_indices = PrimitiveArray::from_iter([0u8, 1, 0, 1]).into_array();
+        let row_indices = PrimitiveArray::from_iter([0u32, 0, 1, 1]).into_array();
+        let interleaved =
+            InterleaveArray::try_new(vec![constant, column], array_indices, row_indices)?
+                .into_array();
+        let expected =
+            PrimitiveArray::from_option_iter([None, Some(10.0f64), None, Some(20.0)]).into_array();
+        let mut ctx = array_session().create_execution_ctx();
+
+        assert_arrays_eq!(interleaved, expected, &mut ctx);
+        Ok(())
     }
 }