feat: add mandatory demand mask to scalar_fn execution

vortex-array/src/arrays/scalar_fn/vtable/mod.rs

@@ -156,7 +156,7 @@ impl VTable for ScalarFn {
 
     fn execute(array: Array<Self>, ctx: &mut ExecutionCtx) -> VortexResult<ExecutionResult> {
         ctx.log(format_args!("scalar_fn({}): executing", array.scalar_fn()));
-        let args = VecExecutionArgs::new(array.children(), array.len());
+        let args = VecExecutionArgs::all(array.children(), array.len());
         array
             .scalar_fn()
             .execute(&args, ctx)

vortex-array/src/arrays/scalar_fn/vtable/operations.rs

@@ -26,7 +26,7 @@ impl OperationsVTable<ScalarFn> for ScalarFn {
             .map(|child| Ok(ConstantArray::new(child.execute_scalar(index, ctx)?, 1).into_array()))
             .collect::<VortexResult<_>>()?;
 
-        let args = VecExecutionArgs::new(inputs, 1);
+        let args = VecExecutionArgs::all(inputs, 1);
         let result = array.scalar_fn().execute(&args, ctx)?;
 
         let scalar = match result.execute::<Columnar>(ctx)? {

vortex-array/src/arrays/scalar_fn/vtable/validity.rs

@@ -50,7 +50,7 @@ fn execute_expr(
         .map(|child| execute_expr(child, row_count, ctx))
         .collect::<VortexResult<_>>()?;
 
-    let args = VecExecutionArgs::new(inputs, row_count);
+    let args = VecExecutionArgs::all(inputs, row_count);
 
     Ok(expr.scalar_fn().execute(&args, ctx)?.into_array())
 }

vortex-array/src/scalar_fn/fns/binary/mod.rs

@@ -157,10 +157,10 @@ impl ScalarFnVTable for Binary {
             Operator::Gte => execute_compare(&lhs, &rhs, CompareOperator::Gte, ctx),
             Operator::And => execute_boolean(lhs, rhs, Operator::And, ctx),
             Operator::Or => execute_boolean(lhs, rhs, Operator::Or, ctx),
-            Operator::Add => execute_numeric(&lhs, &rhs, NumericOperator::Add, ctx),
-            Operator::Sub => execute_numeric(&lhs, &rhs, NumericOperator::Sub, ctx),
-            Operator::Mul => execute_numeric(&lhs, &rhs, NumericOperator::Mul, ctx),
-            Operator::Div => execute_numeric(&lhs, &rhs, NumericOperator::Div, ctx),
+            Operator::Add => execute_numeric(&lhs, &rhs, NumericOperator::Add, args.demand(), ctx),
+            Operator::Sub => execute_numeric(&lhs, &rhs, NumericOperator::Sub, args.demand(), ctx),
+            Operator::Mul => execute_numeric(&lhs, &rhs, NumericOperator::Mul, args.demand(), ctx),
+            Operator::Div => execute_numeric(&lhs, &rhs, NumericOperator::Div, args.demand(), ctx),
         }
     }
 

vortex-array/src/scalar_fn/fns/binary/numeric.rs

@@ -17,8 +17,10 @@ use crate::arrays::Constant;
 use crate::arrays::ConstantArray;
 use crate::arrays::PrimitiveArray;
 use crate::builtins::ArrayBuiltins;
+use crate::child_to_validity;
 use crate::dtype::DType;
 use crate::dtype::NativePType;
+use crate::dtype::Nullability;
 use crate::dtype::PType;
 use crate::dtype::half::f16;
 use crate::match_each_native_ptype;
@@ -105,10 +107,14 @@ impl<T: CheckedArithmetic> CheckedPrimitiveOp<T> for CheckedDiv {
 }
 
 /// Execute a numeric operation between two arrays.
+///
+/// Undemanded rows are don't-care per the [`crate::scalar_fn::ExecutionArgs::demand`]
+/// contract: their output values are unspecified and they never raise a domain error.
 pub(crate) fn execute_numeric(
     lhs: &ArrayRef,
     rhs: &ArrayRef,
     op: NumericOperator,
+    demand: &ArrayRef,
     ctx: &mut ExecutionCtx,
 ) -> VortexResult<ArrayRef> {
     let ptype = PType::try_from(lhs.dtype())?;
@@ -122,17 +128,18 @@ pub(crate) fn execute_numeric(
 
     match_each_native_ptype!(ptype, |T| {
         match op {
-            NumericOperator::Add => execute_checked_typed::<T, CheckedAdd>(lhs, rhs, ctx),
-            NumericOperator::Sub => execute_checked_typed::<T, CheckedSub>(lhs, rhs, ctx),
-            NumericOperator::Mul => execute_checked_typed::<T, CheckedMul>(lhs, rhs, ctx),
-            NumericOperator::Div => execute_checked_typed::<T, CheckedDiv>(lhs, rhs, ctx),
+            NumericOperator::Add => execute_checked_typed::<T, CheckedAdd>(lhs, rhs, demand, ctx),
+            NumericOperator::Sub => execute_checked_typed::<T, CheckedSub>(lhs, rhs, demand, ctx),
+            NumericOperator::Mul => execute_checked_typed::<T, CheckedMul>(lhs, rhs, demand, ctx),
+            NumericOperator::Div => execute_checked_typed::<T, CheckedDiv>(lhs, rhs, demand, ctx),
         }
     })
 }
 
 fn execute_checked_typed<T, Op>(
     lhs: &ArrayRef,
     rhs: &ArrayRef,
+    demand: &ArrayRef,
     ctx: &mut ExecutionCtx,
 ) -> VortexResult<ArrayRef>
 where
@@ -157,28 +164,38 @@ where
 
     let validity = lhs.validity().and(rhs.validity())?;
     let valid_rows = validity.execute_mask(len, ctx)?;
+    let demand =
+        child_to_validity(Some(demand), Nullability::NonNullable).execute_mask(len, ctx)?;
+    // Lanes that must produce correct values and whose errors are observable. Invalid
+    // lanes are null and undemanded lanes are don't-care, so both are excluded.
+    let care_lanes = if demand.all_true() {
+        valid_rows
+    } else {
+        valid_rows & &demand
+    };
 
     let checked = match (&lhs, &rhs) {
         (
             PrimitiveOperand::Array { values: lhs, .. },
             PrimitiveOperand::Array { values: rhs, .. },
-        ) => checked_array_array::<T, Op>(lhs, rhs, &valid_rows),
+        ) => checked_array_array::<T, Op>(lhs, rhs, &care_lanes),
         (
             PrimitiveOperand::Array { values: lhs, .. },
             PrimitiveOperand::Constant { value: rhs, .. },
-        ) => checked_array_constant::<T, Op>(lhs, *rhs, &valid_rows),
+        ) => checked_array_constant::<T, Op>(lhs, *rhs, &care_lanes),
         (
             PrimitiveOperand::Constant { value: lhs, .. },
             PrimitiveOperand::Array { values: rhs, .. },
-        ) => checked_constant_array::<T, Op>(*lhs, rhs, &valid_rows),
+        ) => checked_constant_array::<T, Op>(*lhs, rhs, &care_lanes),
         (
             PrimitiveOperand::Constant { value: lhs, .. },
             PrimitiveOperand::Constant { value: rhs, .. },
-        ) => {
-            let value = Op::checked(*lhs, *rhs)
-                .ok_or_else(|| vortex_err!(InvalidArgument: "{}", Op::ERROR))?;
-            return Ok(constant_result_array(value, len, &result_dtype));
-        }
+        ) => match Op::checked(*lhs, *rhs) {
+            Some(value) => return Ok(constant_result_array(value, len, &result_dtype)),
+            // The op fails on every lane, but no failing lane is observable.
+            None if care_lanes.all_false() => CheckedValues::zeroed(len),
+            None => vortex_bail!(InvalidArgument: "{}", Op::ERROR),
+        },
         (PrimitiveOperand::Null(_), _) | (_, PrimitiveOperand::Null(_)) => {
             CheckedValues::zeroed(len)
         }
@@ -289,14 +306,14 @@ impl<T: NativePType> CheckedValues<T> {
     }
 }
 
-fn checked_array_array<T, Op>(lhs: &[T], rhs: &[T], valid_rows: &Mask) -> CheckedValues<T>
+fn checked_array_array<T, Op>(lhs: &[T], rhs: &[T], care_lanes: &Mask) -> CheckedValues<T>
 where
     T: NativePType,
     Op: CheckedPrimitiveOp<T>,
 {
     debug_assert_eq!(lhs.len(), rhs.len());
 
-    match valid_rows.bit_buffer() {
+    match care_lanes.bit_buffer() {
         AllOr::All if Op::CHECKED_VALUE_LOOP => checked_array_array_one_pass::<T, Op>(lhs, rhs),
         AllOr::All => checked_array_array_all_lanes::<T, Op>(lhs, rhs),
         AllOr::None => CheckedValues::zeroed(lhs.len()),
@@ -307,12 +324,12 @@ where
     }
 }
 
-fn checked_array_constant<T, Op>(lhs: &[T], rhs: T, valid_rows: &Mask) -> CheckedValues<T>
+fn checked_array_constant<T, Op>(lhs: &[T], rhs: T, care_lanes: &Mask) -> CheckedValues<T>
 where
     T: NativePType,
     Op: CheckedPrimitiveOp<T>,
 {
-    match valid_rows.bit_buffer() {
+    match care_lanes.bit_buffer() {
         AllOr::All if Op::CHECKED_VALUE_LOOP => checked_array_constant_one_pass::<T, Op>(lhs, rhs),
         AllOr::All => checked_array_constant_all_lanes::<T, Op>(lhs, rhs),
         AllOr::None => CheckedValues::zeroed(lhs.len()),
@@ -325,12 +342,12 @@ where
     }
 }
 
-fn checked_constant_array<T, Op>(lhs: T, rhs: &[T], valid_rows: &Mask) -> CheckedValues<T>
+fn checked_constant_array<T, Op>(lhs: T, rhs: &[T], care_lanes: &Mask) -> CheckedValues<T>
 where
     T: NativePType,
     Op: CheckedPrimitiveOp<T>,
 {
-    match valid_rows.bit_buffer() {
+    match care_lanes.bit_buffer() {
         AllOr::All if Op::CHECKED_VALUE_LOOP => checked_constant_array_one_pass::<T, Op>(lhs, rhs),
         AllOr::All => checked_constant_array_all_lanes::<T, Op>(lhs, rhs),
         AllOr::None => CheckedValues::zeroed(rhs.len()),
@@ -931,14 +948,31 @@ mod test {
     use crate::RecursiveCanonical;
     use crate::VortexSessionExecute;
     use crate::array_session;
+    use crate::arrays::BoolArray;
     use crate::arrays::ConstantArray;
     use crate::arrays::PrimitiveArray;
     use crate::assert_arrays_eq;
     use crate::builtins::ArrayBuiltins;
     use crate::scalar::Scalar;
+    use crate::scalar_fn::ScalarFnVTableExt;
+    use crate::scalar_fn::VecExecutionArgs;
+    use crate::scalar_fn::fns::binary::Binary;
     use crate::scalar_fn::fns::operators::Operator;
     use crate::validity::Validity;
 
+    fn execute_with_demand(
+        lhs: ArrayRef,
+        rhs: ArrayRef,
+        op: Operator,
+        demand: ArrayRef,
+    ) -> VortexResult<ArrayRef> {
+        let len = lhs.len();
+        let args = VecExecutionArgs::new(vec![lhs, rhs], len, demand);
+        Binary
+            .bind(op)
+            .execute(&args, &mut array_session().create_execution_ctx())
+    }
+
     fn sub_scalar(array: &ArrayRef, scalar: impl Into<Scalar>) -> VortexResult<ArrayRef> {
         array
             .binary(
@@ -1115,6 +1149,145 @@ mod test {
         assert!(result.is_err());
     }
 
+    #[test]
+    fn test_add_overflow_on_undemanded_lane_is_ok() -> VortexResult<()> {
+        let mut ctx = array_session().create_execution_ctx();
+        let lhs = buffer![u8::MAX, 1].into_array();
+        let rhs = buffer![1u8, 1].into_array();
+        let result = execute_with_demand(
+            lhs,
+            rhs,
+            Operator::Add,
+            BoolArray::from_iter([false, true]).into_array(),
+        )?;
+
+        assert_eq!(result.len(), 2);
+        assert_eq!(result.execute_scalar(1, &mut ctx)?, Scalar::from(2u8));
+        Ok(())
+    }
+
+    #[test]
+    fn test_add_overflow_on_demanded_lane_errors() {
+        let lhs = buffer![u8::MAX, 1].into_array();
+        let rhs = buffer![1u8, 1].into_array();
+        let result = execute_with_demand(
+            lhs,
+            rhs,
+            Operator::Add,
+            BoolArray::from_iter([true, false]).into_array(),
+        );
+
+        assert!(result.is_err());
+    }
+
+    #[test]
+    fn test_add_scalar_overflow_on_undemanded_lane_is_ok() -> VortexResult<()> {
+        let mut ctx = array_session().create_execution_ctx();
+        let lhs = buffer![u8::MAX, 1].into_array();
+        let rhs = ConstantArray::new(1u8, 2).into_array();
+        let result = execute_with_demand(
+            lhs,
+            rhs,
+            Operator::Add,
+            BoolArray::from_iter([false, true]).into_array(),
+        )?;
+
+        assert_eq!(result.execute_scalar(1, &mut ctx)?, Scalar::from(2u8));
+        Ok(())
+    }
+
+    #[test]
+    fn test_add_constant_overflow_fully_undemanded_is_ok() -> VortexResult<()> {
+        let lhs = ConstantArray::new(u8::MAX, 2).into_array();
+        let rhs = ConstantArray::new(1u8, 2).into_array();
+        let result = execute_with_demand(
+            lhs,
+            rhs,
+            Operator::Add,
+            ConstantArray::new(false, 2).into_array(),
+        )?;
+
+        assert_eq!(result.len(), 2);
+        Ok(())
+    }
+
+    #[test]
+    fn test_add_constant_overflow_on_demanded_lane_errors() {
+        let lhs = ConstantArray::new(u8::MAX, 2).into_array();
+        let rhs = ConstantArray::new(1u8, 2).into_array();
+        let result = execute_with_demand(
+            lhs,
+            rhs,
+            Operator::Add,
+            BoolArray::from_iter([true, false]).into_array(),
+        );
+
+        assert!(result.is_err());
+    }
+
+    #[test]
+    fn test_div_by_zero_on_undemanded_lane_is_ok() -> VortexResult<()> {
+        let mut ctx = array_session().create_execution_ctx();
+        let lhs = buffer![10i32, 10].into_array();
+        let rhs = buffer![0i32, 2].into_array();
+        let result = execute_with_demand(
+            lhs,
+            rhs,
+            Operator::Div,
+            BoolArray::from_iter([false, true]).into_array(),
+        )?;
+
+        assert_eq!(result.execute_scalar(1, &mut ctx)?, Scalar::from(5i32));
+        Ok(())
+    }
+
+    #[test]
+    fn test_div_by_zero_on_demanded_lane_errors() {
+        let lhs = buffer![10i32, 10].into_array();
+        let rhs = buffer![0i32, 2].into_array();
+        let result = execute_with_demand(
+            lhs,
+            rhs,
+            Operator::Div,
+            BoolArray::from_iter([true, false]).into_array(),
+        );
+
+        assert!(result.is_err());
+    }
+
+    #[test]
+    fn test_scalar_add_array_overflow_on_undemanded_lane_is_ok() -> VortexResult<()> {
+        let mut ctx = array_session().create_execution_ctx();
+        let lhs = ConstantArray::new(1u8, 2).into_array();
+        let rhs = buffer![u8::MAX, 1].into_array();
+        let result = execute_with_demand(
+            lhs,
+            rhs,
+            Operator::Add,
+            BoolArray::from_iter([false, true]).into_array(),
+        )?;
+
+        assert_eq!(result.execute_scalar(1, &mut ctx)?, Scalar::from(2u8));
+        Ok(())
+    }
+
+    #[test]
+    fn test_add_overflow_on_null_demanded_lane_is_ok() -> VortexResult<()> {
+        let mut ctx = array_session().create_execution_ctx();
+        let lhs = PrimitiveArray::new(buffer![u8::MAX, 1], Validity::from_iter([false, true]))
+            .into_array();
+        let rhs = buffer![1u8, 1].into_array();
+        let result = execute_with_demand(
+            lhs,
+            rhs,
+            Operator::Add,
+            ConstantArray::new(true, 2).into_array(),
+        )?;
+
+        assert_eq!(result.execute_scalar(1, &mut ctx)?, Scalar::from(Some(2u8)));
+        Ok(())
+    }
+
     #[test]
     fn test_present_nullable_constant_preserves_nullable_output() {
         let mut ctx = array_session().create_execution_ctx();

vortex-array/src/scalar_fn/fns/dynamic.rs

@@ -104,7 +104,8 @@ impl ScalarFnVTable for DynamicComparison {
             let lhs = args.get(0)?;
             let rhs = ConstantArray::new(scalar, args.row_count()).into_array();
 
-            let delegate_args = VecExecutionArgs::new(vec![lhs, rhs], args.row_count());
+            let delegate_args =
+                VecExecutionArgs::new(vec![lhs, rhs], args.row_count(), args.demand().clone());
             return Binary
                 .bind(Operator::from(data.operator))
                 .execute(&delegate_args, ctx);
@@ -355,6 +356,25 @@ mod tests {
         Ok(())
     }
 
+    #[test]
+    fn execute_forwards_demand_to_delegate() -> VortexResult<()> {
+        // TODO: strengthen to observe the forwarded mask once Binary honors demand.
+        let scalar_fn = DynamicComparison.bind(DynamicComparisonExpr {
+            operator: CompareOperator::Lt,
+            rhs: Arc::new(Rhs {
+                value: Arc::new(|| Some(5i32.into())),
+                dtype: DType::Primitive(PType::I32, Nullability::NonNullable),
+            }),
+            default: true,
+        });
+
+        let demand = BoolArray::from_iter([true, false, true]).into_array();
+        let args = VecExecutionArgs::new(vec![buffer![1i32, 5, 10].into_array()], 3, demand);
+        let result = scalar_fn.execute(&args, &mut array_session().create_execution_ctx())?;
+        assert_eq!(result.len(), 3);
+        Ok(())
+    }
+
     #[test]
     fn execute_value_flips() -> VortexResult<()> {
         let mut ctx = array_session().create_execution_ctx();