Pack Boolean RowFn output directly

vortex-array/Cargo.toml

@@ -135,6 +135,10 @@ harness = false
 name = "binary_ops"
 harness = false
 
+[[bench]]
+name = "row_fn_output"
+harness = false
+
 [[bench]]
 name = "kleene_bool"
 harness = false

vortex-array/benches/row_fn_output.rs

@@ -0,0 +1,129 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+#![expect(
+    clippy::cast_possible_truncation,
+    reason = "benchmark fixtures reduce values below 1024 before narrowing"
+)]
+
+use std::marker::PhantomData;
+use std::sync::LazyLock;
+
+use divan::Bencher;
+use divan::counter::ItemsCount;
+use mimalloc::MiMalloc;
+use vortex_array::ArrayRef;
+use vortex_array::IntoArray;
+use vortex_array::VortexSessionExecute;
+use vortex_array::array_session;
+use vortex_array::arrays::ConstantArray;
+use vortex_array::arrays::PrimitiveArray;
+use vortex_array::dtype::DType;
+use vortex_array::dtype::NativePType;
+use vortex_array::scalar_fn::EmptyOptions;
+use vortex_array::scalar_fn::ScalarFnId;
+use vortex_array::scalar_fn::VecExecutionArgs;
+use vortex_array::scalar_fn::unstable::row::RowFn;
+use vortex_array::scalar_fn::unstable::row::RowVisitor;
+use vortex_array::scalar_fn::unstable::row::execute_rows;
+use vortex_error::VortexExpect;
+use vortex_error::VortexResult;
+use vortex_session::VortexSession;
+use vortex_session::registry::CachedId;
+
+#[global_allocator]
+static GLOBAL: MiMalloc = MiMalloc;
+
+fn main() {
+    LazyLock::force(&SESSION);
+    divan::main();
+}
+
+static SESSION: LazyLock<VortexSession> = LazyLock::new(array_session);
+
+const ROWS: usize = 1 << 14;
+
+const INPUT_SHAPES: &[InputShape] = &[
+    InputShape::PerRowPerRow,
+    InputShape::PerRowConstant,
+    InputShape::ConstantPerRow,
+];
+
+#[derive(Clone, Copy, Debug)]
+enum InputShape {
+    PerRowPerRow,
+    PerRowConstant,
+    ConstantPerRow,
+}
+
+trait BenchPrimitive: NativePType {
+    fn per_row(offset: usize) -> ArrayRef;
+
+    fn constant() -> ArrayRef;
+}
+
+impl BenchPrimitive for i32 {
+    fn per_row(offset: usize) -> ArrayRef {
+        PrimitiveArray::from_iter((0..ROWS).map(|index| ((index + offset) % 1024) as i32))
+            .into_array()
+    }
+
+    fn constant() -> ArrayRef {
+        ConstantArray::new(512_i32, ROWS).into_array()
+    }
+}
+
+impl BenchPrimitive for i64 {
+    fn per_row(offset: usize) -> ArrayRef {
+        PrimitiveArray::from_iter((0..ROWS).map(|index| ((index + offset) % 1024) as i64))
+            .into_array()
+    }
+
+    fn constant() -> ArrayRef {
+        ConstantArray::new(512_i64, ROWS).into_array()
+    }
+}
+
+#[derive(Clone)]
+struct InfallibleBool<T>(PhantomData<T>);
+
+impl<T: BenchPrimitive> RowFn for InfallibleBool<T> {
+    type Options = EmptyOptions;
+
+    const ARG_NAMES: &'static [&'static str] = &["lhs", "rhs"];
+    const INFALLIBLE: bool = true;
+
+    fn id(&self) -> ScalarFnId {
+        static ID: CachedId = CachedId::new("bench.row_fn_output.infallible_bool");
+        *ID
+    }
+
+    fn dispatch<V: RowVisitor>(
+        &self,
+        _options: &Self::Options,
+        _args: &[DType],
+        visitor: V,
+    ) -> VortexResult<V::VisitResult> {
+        visitor.visit::<(T, T), bool>(|(lhs, rhs)| lhs.is_lt(rhs))
+    }
+}
+
+#[vortex_bench_support::cpu_features]
+#[divan::bench(types = [i32, i64], args = INPUT_SHAPES)]
+fn infallible_bool<T: BenchPrimitive>(bencher: Bencher, &shape: &InputShape) {
+    let function = InfallibleBool::<T>(PhantomData);
+    let args = match shape {
+        InputShape::PerRowPerRow => vec![T::per_row(0), T::per_row(1)],
+        InputShape::PerRowConstant => vec![T::per_row(0), T::constant()],
+        InputShape::ConstantPerRow => vec![T::constant(), T::per_row(1)],
+    };
+    let args = VecExecutionArgs::new(args, ROWS);
+
+    bencher
+        .counter(ItemsCount::new(ROWS))
+        .with_inputs(|| (&args, SESSION.create_execution_ctx()))
+        .bench_refs(|(args, ctx)| {
+            execute_rows(&function, &EmptyOptions, *args, ctx)
+                .vortex_expect("row execution should succeed in benchmark")
+        });
+}

vortex-array/src/scalar_fn/unstable/row/batch/tests.rs

@@ -6,6 +6,7 @@ use std::sync::atomic::AtomicUsize;
 use std::sync::atomic::Ordering;
 
 use rstest::rstest;
+use vortex_buffer::Buffer;
 use vortex_buffer::BufferMut;
 use vortex_error::VortexResult;
 use vortex_error::vortex_bail;
@@ -15,9 +16,11 @@ use vortex_session::registry::CachedId;
 
 use super::finalize_kernel_output;
 use crate::ArrayRef;
+use crate::ExecutionCtx;
 use crate::IntoArray;
 use crate::VortexSessionExecute;
 use crate::array_session;
+use crate::arrays::BoolArray;
 use crate::arrays::ConstantArray;
 use crate::arrays::ExtensionArray;
 use crate::arrays::FixedSizeListArray;
@@ -35,6 +38,7 @@ use crate::scalar_fn::ScalarFnId;
 use crate::scalar_fn::VecExecutionArgs;
 use crate::scalar_fn::unstable::row::FixedSizeListSink;
 use crate::scalar_fn::unstable::row::InitializedRow;
+use crate::scalar_fn::unstable::row::InputElement;
 use crate::scalar_fn::unstable::row::OutputElement;
 use crate::scalar_fn::unstable::row::OutputSink;
 use crate::scalar_fn::unstable::row::RowFn;
@@ -61,6 +65,51 @@ struct ValidOnlyIdentity;
 #[derive(Clone)]
 struct InvalidKernelOutput;
 
+#[derive(Clone)]
+struct PackedPositive;
+
+#[derive(Clone)]
+struct ValidOnlyPositive;
+
+#[derive(Clone)]
+struct NullaryTrue;
+
+struct ValidOnlyI64;
+
+// SAFETY: the view is a slice, and its reported length is the buffer length.
+unsafe impl InputElement for ValidOnlyI64 {
+    type Column = Buffer<i64>;
+    type View<'a> = &'a [i64];
+    type Elem<'a> = i64;
+
+    const DENSE_SAFE: bool = false;
+    const DECODE_INFALLIBLE: bool = true;
+
+    fn validate(dtype: &DType) -> VortexResult<()> {
+        <i64 as InputElement>::validate(dtype)
+    }
+
+    fn decode(array: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult<Self::Column> {
+        <i64 as InputElement>::decode(array, ctx)
+    }
+
+    fn can_decode_null_tolerant(_array: &ArrayRef) -> VortexResult<bool> {
+        Ok(true)
+    }
+
+    fn get(column: &Self::Column, index: usize) -> Self::Elem<'_> {
+        column[index]
+    }
+
+    fn view(column: &Self::Column) -> Self::View<'_> {
+        column.as_slice()
+    }
+
+    fn get_from_view<'a>(view: &Self::View<'a>, index: usize) -> Self::Elem<'a> {
+        view[index]
+    }
+}
+
 /// Produces a null row to exercise output validation at the row-function boundary.
 #[derive(Default)]
 struct NullProducingI64(i64);
@@ -274,6 +323,72 @@ fn test_fixed_size_list_sink_uses_runtime_width(
     Ok(())
 }
 
+impl RowFn for PackedPositive {
+    type Options = EmptyOptions;
+
+    const ARG_NAMES: &'static [&'static str] = &["value"];
+    const INFALLIBLE: bool = true;
+
+    fn id(&self) -> ScalarFnId {
+        static ID: CachedId = CachedId::new("test.packed_positive");
+        *ID
+    }
+
+    fn dispatch<V: RowVisitor>(
+        &self,
+        _options: &Self::Options,
+        _args: &[DType],
+        visitor: V,
+    ) -> VortexResult<V::VisitResult> {
+        visitor.visit::<(i64,), bool>(|(value,)| value > 0)
+    }
+}
+
+impl RowFn for ValidOnlyPositive {
+    type Options = EmptyOptions;
+
+    const ARG_NAMES: &'static [&'static str] = &["value"];
+    const INFALLIBLE: bool = true;
+
+    fn id(&self) -> ScalarFnId {
+        static ID: CachedId = CachedId::new("test.valid_only_positive");
+        *ID
+    }
+
+    fn dispatch<V: RowVisitor>(
+        &self,
+        _options: &Self::Options,
+        _args: &[DType],
+        visitor: V,
+    ) -> VortexResult<V::VisitResult> {
+        visitor.visit::<(ValidOnlyI64,), bool>(|(value,)| {
+            assert_ne!(value, i64::MIN, "an invalid row was evaluated");
+            value > 0
+        })
+    }
+}
+
+impl RowFn for NullaryTrue {
+    type Options = EmptyOptions;
+
+    const ARG_NAMES: &'static [&'static str] = &[];
+    const INFALLIBLE: bool = true;
+
+    fn id(&self) -> ScalarFnId {
+        static ID: CachedId = CachedId::new("test.nullary_true");
+        *ID
+    }
+
+    fn dispatch<V: RowVisitor>(
+        &self,
+        _options: &Self::Options,
+        _args: &[DType],
+        visitor: V,
+    ) -> VortexResult<V::VisitResult> {
+        visitor.visit::<(), bool>(|()| true)
+    }
+}
+
 #[test]
 fn test_finalize_kernel_output_rejects_nested_dtype_mismatch() -> VortexResult<()> {
     static ID: CachedId = CachedId::new("test.finalize_kernel_output");
@@ -320,6 +435,73 @@ fn test_kernel_output_rejects_nulls_at_function_boundary() -> VortexResult<()> {
     Ok(())
 }
 
+#[test]
+fn test_bool_output_builds_packed_values() -> VortexResult<()> {
+    let input = PrimitiveArray::new(
+        vec![1_i64, -1, 2, 0, 3],
+        Validity::from_iter([true, true, false, true, true]),
+    )
+    .into_array();
+    let args = VecExecutionArgs::new(vec![input], 5);
+    let mut ctx = array_session().create_execution_ctx();
+
+    let actual = execute_rows(&PackedPositive, &EmptyOptions, &args, &mut ctx)?;
+    let expected =
+        BoolArray::from_iter([Some(true), Some(false), None, Some(false), Some(true)]).into_array();
+
+    assert_arrays_eq!(&actual, &expected, &mut ctx);
+    Ok(())
+}
+
+#[rstest]
+#[case::empty(0)]
+#[case::one(1)]
+#[case::below_word(63)]
+#[case::exact_word(64)]
+#[case::above_word(65)]
+#[case::multiword_remainder(130)]
+fn test_bool_output_word_boundaries(#[case] len: usize) -> VortexResult<()> {
+    let values: Vec<_> = (0..len).map(|index| index as i64 - 32).collect();
+    let input = PrimitiveArray::from_iter(values.iter().copied()).into_array();
+    let args = VecExecutionArgs::new(vec![input], len);
+    let mut ctx = array_session().create_execution_ctx();
+
+    let actual = execute_rows(&PackedPositive, &EmptyOptions, &args, &mut ctx)?;
+    let expected = BoolArray::from_iter(values.iter().map(|value| *value > 0)).into_array();
+
+    assert_arrays_eq!(&actual, &expected, &mut ctx);
+    Ok(())
+}
+
+#[test]
+fn test_bool_output_handles_nullary_rows() -> VortexResult<()> {
+    let args = VecExecutionArgs::new(vec![], 65);
+    let mut ctx = array_session().create_execution_ctx();
+
+    let actual = execute_rows(&NullaryTrue, &EmptyOptions, &args, &mut ctx)?;
+    let expected = BoolArray::from_iter(std::iter::repeat_n(true, 65)).into_array();
+
+    assert_arrays_eq!(&actual, &expected, &mut ctx);
+    Ok(())
+}
+
+#[test]
+fn test_valid_only_bool_output_skips_invalid_rows() -> VortexResult<()> {
+    let input = PrimitiveArray::new(
+        vec![1_i64, i64::MIN, -1],
+        Validity::from_iter([true, false, true]),
+    )
+    .into_array();
+    let args = VecExecutionArgs::new(vec![input], 3);
+    let mut ctx = array_session().create_execution_ctx();
+
+    let actual = execute_rows(&ValidOnlyPositive, &EmptyOptions, &args, &mut ctx)?;
+    let expected = BoolArray::from_iter([Some(true), None, Some(false)]).into_array();
+
+    assert_arrays_eq!(&actual, &expected, &mut ctx);
+    Ok(())
+}
+
 #[test]
 fn test_deferred_owned_execution_retries_null_row_failure() -> VortexResult<()> {
     let function = DeferredAdd::default();