Pack deferred Boolean RowFn output directly

vortex-array/benches/row_fn_output.rs

@@ -28,6 +28,7 @@ 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_error::vortex_ensure;
 use vortex_session::VortexSession;
 use vortex_session::registry::CachedId;
 
@@ -87,6 +88,12 @@ impl BenchPrimitive for i64 {
 #[derive(Clone)]
 struct InfallibleBool<T>(PhantomData<T>);
 
+#[derive(Clone)]
+struct DeferredBool<T>(PhantomData<T>);
+
+#[derive(Clone)]
+struct DeferredI64;
+
 impl<T: BenchPrimitive> RowFn for InfallibleBool<T> {
     type Options = EmptyOptions;
 
@@ -108,22 +115,106 @@ impl<T: BenchPrimitive> RowFn for InfallibleBool<T> {
     }
 }
 
+impl<T: BenchPrimitive> RowFn for DeferredBool<T> {
+    type Options = EmptyOptions;
+
+    const ARG_NAMES: &'static [&'static str] = &["lhs", "rhs"];
+    const INFALLIBLE: bool = false;
+
+    fn id(&self) -> ScalarFnId {
+        static ID: CachedId = CachedId::new("bench.row_fn_output.deferred_bool");
+        *ID
+    }
+
+    fn dispatch<V: RowVisitor>(
+        &self,
+        _options: &Self::Options,
+        _args: &[DType],
+        visitor: V,
+    ) -> VortexResult<V::VisitResult> {
+        visitor.visit_deferred::<(T, T), bool, bool>(
+            |(lhs, rhs)| (lhs.is_lt(rhs), lhs.is_lt(T::default())),
+            |negative| {
+                vortex_ensure!(
+                    !negative,
+                    "deferred-bool benchmark inputs must be nonnegative"
+                );
+                Ok(())
+            },
+        )
+    }
+}
+
+impl RowFn for DeferredI64 {
+    type Options = EmptyOptions;
+
+    const ARG_NAMES: &'static [&'static str] = &["lhs", "rhs"];
+    const INFALLIBLE: bool = false;
+
+    fn id(&self) -> ScalarFnId {
+        static ID: CachedId = CachedId::new("bench.row_fn_output.deferred_i64");
+        *ID
+    }
+
+    fn dispatch<V: RowVisitor>(
+        &self,
+        _options: &Self::Options,
+        _args: &[DType],
+        visitor: V,
+    ) -> VortexResult<V::VisitResult> {
+        visitor.visit_deferred::<(i64, i64), i64, bool>(
+            |(lhs, rhs)| lhs.overflowing_add(rhs),
+            |overflowed| {
+                vortex_ensure!(
+                    !overflowed,
+                    "deferred-i64 benchmark inputs must not overflow"
+                );
+                Ok(())
+            },
+        )
+    }
+}
+
 #[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);
+    bench_row_fn(bencher, &function, make_args::<T>(shape));
+}
+
+#[vortex_bench_support::cpu_features]
+#[divan::bench(types = [i32, i64], args = INPUT_SHAPES)]
+fn deferred_bool<T: BenchPrimitive>(bencher: Bencher, &shape: &InputShape) {
+    let function = DeferredBool::<T>(PhantomData);
+    bench_row_fn(bencher, &function, make_args::<T>(shape));
+}
+
+#[vortex_bench_support::cpu_features]
+#[divan::bench(args = INPUT_SHAPES)]
+fn deferred_i64(bencher: Bencher, &shape: &InputShape) {
+    bench_row_fn(bencher, &DeferredI64, make_args::<i64>(shape));
+}
+
+fn make_args<T: BenchPrimitive>(shape: InputShape) -> VecExecutionArgs {
     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);
 
+    VecExecutionArgs::new(args, ROWS)
+}
+
+fn bench_row_fn<F: RowFn<Options = EmptyOptions>>(
+    bencher: Bencher,
+    function: &F,
+    args: VecExecutionArgs,
+) {
     bencher
         .counter(ItemsCount::new(ROWS))
         .with_inputs(|| (&args, SESSION.create_execution_ctx()))
         .bench_refs(|(args, ctx)| {
-            execute_rows(&function, &EmptyOptions, *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

@@ -68,6 +68,12 @@ struct InvalidKernelOutput;
 #[derive(Clone)]
 struct PackedPositive;
 
+#[derive(Clone)]
+struct PackedGreaterThan;
+
+#[derive(Clone)]
+struct DeferredGreaterThan;
+
 #[derive(Clone)]
 struct ValidOnlyPositive;
 
@@ -344,6 +350,56 @@ impl RowFn for PackedPositive {
     }
 }
 
+impl RowFn for PackedGreaterThan {
+    type Options = EmptyOptions;
+
+    const ARG_NAMES: &'static [&'static str] = &["lhs", "rhs"];
+    const INFALLIBLE: bool = true;
+
+    fn id(&self) -> ScalarFnId {
+        static ID: CachedId = CachedId::new("test.packed_greater_than");
+        *ID
+    }
+
+    fn dispatch<V: RowVisitor>(
+        &self,
+        _options: &Self::Options,
+        _args: &[DType],
+        visitor: V,
+    ) -> VortexResult<V::VisitResult> {
+        visitor.visit::<(i64, i64), bool>(|(lhs, rhs)| lhs > rhs)
+    }
+}
+
+impl RowFn for DeferredGreaterThan {
+    type Options = EmptyOptions;
+
+    const ARG_NAMES: &'static [&'static str] = &["lhs", "rhs"];
+    const INFALLIBLE: bool = false;
+
+    fn id(&self) -> ScalarFnId {
+        static ID: CachedId = CachedId::new("test.deferred_greater_than");
+        *ID
+    }
+
+    fn dispatch<V: RowVisitor>(
+        &self,
+        _options: &Self::Options,
+        _args: &[DType],
+        visitor: V,
+    ) -> VortexResult<V::VisitResult> {
+        visitor.visit_deferred::<(i64, i64), bool, bool>(
+            |(lhs, rhs)| (lhs > rhs, lhs == i64::MIN),
+            |failed| {
+                if failed {
+                    vortex_bail!(InvalidArgument: "deferred comparison failed");
+                }
+
+                Ok(())
+            },
+        )
+    }
+}
 impl RowFn for ValidOnlyPositive {
     type Options = EmptyOptions;
 
@@ -473,6 +529,39 @@ fn test_bool_output_word_boundaries(#[case] len: usize) -> VortexResult<()> {
     Ok(())
 }
 
+#[test]
+fn test_bool_output_handles_partial_constants() -> VortexResult<()> {
+    let values: Vec<_> = (0_i64..65).collect();
+    let varying = PrimitiveArray::from_iter(values.iter().copied()).into_array();
+    let constant = ConstantArray::new(32_i64, values.len()).into_array();
+    let mut ctx = array_session().create_execution_ctx();
+
+    let lhs_varying = VecExecutionArgs::new(vec![varying.clone(), constant.clone()], values.len());
+    let actual = execute_rows(&PackedGreaterThan, &EmptyOptions, &lhs_varying, &mut ctx)?;
+    let expected = BoolArray::from_iter(values.iter().map(|value| *value > 32)).into_array();
+    assert_arrays_eq!(&actual, &expected, &mut ctx);
+
+    let rhs_varying = VecExecutionArgs::new(vec![constant, varying], values.len());
+    let actual = execute_rows(&PackedGreaterThan, &EmptyOptions, &rhs_varying, &mut ctx)?;
+    let expected = BoolArray::from_iter(values.iter().map(|value| 32 > *value)).into_array();
+    assert_arrays_eq!(&actual, &expected, &mut ctx);
+
+    Ok(())
+}
+
+#[test]
+fn test_bool_output_handles_all_constant_input() -> VortexResult<()> {
+    let input = ConstantArray::new(7_i64, 65).into_array();
+    let args = VecExecutionArgs::new(vec![input], 65);
+    let mut ctx = array_session().create_execution_ctx();
+
+    let actual = execute_rows(&PackedPositive, &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_bool_output_handles_nullary_rows() -> VortexResult<()> {
     let args = VecExecutionArgs::new(vec![], 65);
@@ -502,6 +591,74 @@ fn test_valid_only_bool_output_skips_invalid_rows() -> VortexResult<()> {
     Ok(())
 }
 
+#[test]
+fn test_deferred_bool_output_builds_packed_values() -> VortexResult<()> {
+    let values: Vec<_> = (0_i64..65).collect();
+    let lhs = PrimitiveArray::from_iter(values.iter().copied()).into_array();
+    let rhs = PrimitiveArray::from_iter(std::iter::repeat_n(32_i64, values.len())).into_array();
+    let args = VecExecutionArgs::new(vec![lhs, rhs], values.len());
+    let mut ctx = array_session().create_execution_ctx();
+
+    let actual = execute_rows(&DeferredGreaterThan, &EmptyOptions, &args, &mut ctx)?;
+    let expected = BoolArray::from_iter(values.iter().map(|value| *value > 32)).into_array();
+
+    assert_arrays_eq!(&actual, &expected, &mut ctx);
+    Ok(())
+}
+
+#[test]
+fn test_deferred_bool_output_handles_partial_constants() -> VortexResult<()> {
+    let values: Vec<_> = (0_i64..65).collect();
+    let varying = PrimitiveArray::from_iter(values.iter().copied()).into_array();
+    let constant = ConstantArray::new(32_i64, values.len()).into_array();
+    let mut ctx = array_session().create_execution_ctx();
+
+    let lhs_varying = VecExecutionArgs::new(vec![varying.clone(), constant.clone()], values.len());
+    let actual = execute_rows(&DeferredGreaterThan, &EmptyOptions, &lhs_varying, &mut ctx)?;
+    let expected = BoolArray::from_iter(values.iter().map(|value| *value > 32)).into_array();
+    assert_arrays_eq!(&actual, &expected, &mut ctx);
+
+    let rhs_varying = VecExecutionArgs::new(vec![constant, varying], values.len());
+    let actual = execute_rows(&DeferredGreaterThan, &EmptyOptions, &rhs_varying, &mut ctx)?;
+    let expected = BoolArray::from_iter(values.iter().map(|value| 32 > *value)).into_array();
+    assert_arrays_eq!(&actual, &expected, &mut ctx);
+
+    Ok(())
+}
+
+#[test]
+fn test_deferred_bool_output_reports_valid_row_failure() -> VortexResult<()> {
+    let lhs = PrimitiveArray::from_iter([1_i64, i64::MIN, -1]).into_array();
+    let rhs = ConstantArray::new(0_i64, 3).into_array();
+    let args = VecExecutionArgs::new(vec![lhs, rhs], 3);
+    let mut ctx = array_session().create_execution_ctx();
+
+    let error = match execute_rows(&DeferredGreaterThan, &EmptyOptions, &args, &mut ctx) {
+        Err(error) => error.to_string(),
+        Ok(_) => vortex_bail!("a valid-row deferred failure was not reported"),
+    };
+
+    assert!(
+        error.contains("deferred comparison failed"),
+        "fallible execution must report its deferred error, got {error}",
+    );
+    Ok(())
+}
+
+#[test]
+fn test_deferred_owned_execution_handles_constant_lhs() -> VortexResult<()> {
+    let lhs = ConstantArray::new(10_i64, 3).into_array();
+    let rhs = PrimitiveArray::from_iter([1_i64, 2, 3]).into_array();
+    let args = VecExecutionArgs::new(vec![lhs, rhs], 3);
+    let mut ctx = array_session().create_execution_ctx();
+
+    let actual = execute_rows(&DeferredAdd::default(), &EmptyOptions, &args, &mut ctx)?;
+    let expected = PrimitiveArray::from_iter([11_i64, 12, 13]).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();