feat(vortex-geo): add make-line and length scalar functions

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,43 @@ 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 output length.
+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,85 @@
+// 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::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(
+    array: Array<Interleave>,
+    _ctx: &mut ExecutionCtx,
+) -> VortexResult<ExecutionResult> {
+    let num_values = array.num_values();
+    let mut array = array;
+    array = require_child!(array, array.array_indices(), 0 => Primitive);
+    array = require_child!(array, array.row_indices(), 1 => Primitive);
+    for i in 0..num_values {
+        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).as_::<Primitive>().ptype(), |T| {
+        let values = gather_values::<T>(&array)?;
+        VortexResult::Ok(PrimitiveArray::new(values, validity))
+    })?;
+
+    Ok(ExecutionResult::done(output))
+}
+
+fn gather_values<T: NativePType>(array: &Array<Interleave>) -> VortexResult<Buffer<T>> {
+    let buffers = (0..array.num_values())
+        .map(|i| array.value(i).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>(&buffers, branches.as_slice::<A>(), rows)
+    })
+}
+
+fn gather_rows<T, A>(
+    values: &[Buffer<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: &[Buffer<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)?;
+    let mut output = BufferMut::with_capacity(len);
+    for i in 0..len {
+        output.push(values[branches[i].as_()][rows[i].as_()]);
+    }
+    Ok(output.freeze())
+}

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

@@ -719,17 +719,18 @@ 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();
-        interleaved.execute::<Canonical>(&mut ctx).ok();
+        assert_arrays_eq!(interleaved, expected, &mut ctx);
+        Ok(())
     }
 }

vortex-geo/Cargo.toml

@@ -58,5 +58,13 @@ harness = false
 name = "distance"
 harness = false
 
+[[bench]]
+name = "make_line"
+harness = false
+
+[[bench]]
+name = "length"
+harness = false
+
 [lints]
 workspace = true

vortex-geo/benches/length.rs

@@ -0,0 +1,100 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+//! Microbenchmarks for native `ST_Length` over LineStrings.
+//!
+//! The two-vertex case tracks ordinary route segments, while the longer-line case captures the
+//! per-vertex traversal cost. The nullable case measures strict null propagation separately.
+//!
+//! Run with `cargo bench -p vortex-geo --bench length`.
+
+#![expect(clippy::unwrap_used)]
+
+use std::sync::LazyLock;
+
+use divan::Bencher;
+use divan::counter::ItemsCount;
+use mimalloc::MiMalloc;
+use vortex_array::ArrayRef;
+use vortex_array::Canonical;
+use vortex_array::ExecutionCtx;
+use vortex_array::IntoArray;
+use vortex_array::VortexSessionExecute;
+use vortex_array::arrays::MaskedArray;
+use vortex_array::validity::Validity;
+use vortex_geo::scalar_fn::length::GeoLength;
+use vortex_geo::test_harness::geo_session;
+use vortex_geo::test_harness::linestring_column;
+use vortex_session::VortexSession;
+
+// Scalar function execution allocates its output inside the timed region, so use the vendored
+// allocator instead of measuring glibc differences between CodSpeed runner images.
+#[global_allocator]
+static GLOBAL: MiMalloc = MiMalloc;
+
+static SESSION: LazyLock<VortexSession> = LazyLock::new(geo_session);
+
+const ROWS: usize = 512;
+
+fn main() {
+    divan::main();
+}
+
+/// A deterministic vertex ordinate.
+fn ordinate(i: usize) -> f64 {
+    (i.wrapping_mul(2_654_435_761) % 10_000) as f64 / 100.0
+}
+
+fn linestrings(vertices: usize) -> ArrayRef {
+    linestring_column(
+        (0..ROWS)
+            .map(|row| {
+                (0..vertices)
+                    .map(|vertex| (ordinate(row + vertex), ordinate(row + vertex + 1)))
+                    .collect()
+            })
+            .collect(),
+    )
+    .unwrap()
+}
+
+fn lengths(lines: &ArrayRef, ctx: &mut ExecutionCtx) -> ArrayRef {
+    GeoLength::try_new_array(lines.clone())
+        .unwrap()
+        .into_array()
+        .execute::<Canonical>(ctx)
+        .unwrap()
+        .into_array()
+}
+
+#[divan::bench]
+fn two_vertex_lines(bencher: Bencher) {
+    let lines = linestrings(2);
+    let mut ctx = SESSION.create_execution_ctx();
+    bencher
+        .counter(ItemsCount::new(ROWS))
+        .bench_local(|| lengths(&lines, &mut ctx));
+}
+
+#[divan::bench]
+fn sixteen_vertex_lines(bencher: Bencher) {
+    let lines = linestrings(16);
+    let mut ctx = SESSION.create_execution_ctx();
+    bencher
+        .counter(ItemsCount::new(ROWS))
+        .bench_local(|| lengths(&lines, &mut ctx));
+}
+
+#[divan::bench]
+fn nullable_two_vertex_lines(bencher: Bencher) {
+    let lines = MaskedArray::try_new(
+        linestrings(2),
+        Validity::from_iter((0..ROWS).map(|i| !i.is_multiple_of(8))),
+    )
+    .unwrap()
+    .into_array();
+    let mut ctx = SESSION.create_execution_ctx();
+    bencher
+        .counter(ItemsCount::new(ROWS))
+        .bench_local(|| lengths(&lines, &mut ctx));
+}