diff --git a/vortex-array/src/arrays/scalar_fn/vtable/mod.rs b/vortex-array/src/arrays/scalar_fn/vtable/mod.rs index d98e00975ea..d93ac66a402 100644 --- a/vortex-array/src/arrays/scalar_fn/vtable/mod.rs +++ b/vortex-array/src/arrays/scalar_fn/vtable/mod.rs @@ -156,7 +156,7 @@ impl VTable for ScalarFn { fn execute(array: Array, ctx: &mut ExecutionCtx) -> VortexResult { 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) diff --git a/vortex-array/src/arrays/scalar_fn/vtable/operations.rs b/vortex-array/src/arrays/scalar_fn/vtable/operations.rs index c19e3e2288c..de174fc2890 100644 --- a/vortex-array/src/arrays/scalar_fn/vtable/operations.rs +++ b/vortex-array/src/arrays/scalar_fn/vtable/operations.rs @@ -26,7 +26,7 @@ impl OperationsVTable for ScalarFn { .map(|child| Ok(ConstantArray::new(child.execute_scalar(index, ctx)?, 1).into_array())) .collect::>()?; - 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::(ctx)? { diff --git a/vortex-array/src/arrays/scalar_fn/vtable/validity.rs b/vortex-array/src/arrays/scalar_fn/vtable/validity.rs index 39ab8c2c466..6a03e025a3f 100644 --- a/vortex-array/src/arrays/scalar_fn/vtable/validity.rs +++ b/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::>()?; - let args = VecExecutionArgs::new(inputs, row_count); + let args = VecExecutionArgs::all(inputs, row_count); Ok(expr.scalar_fn().execute(&args, ctx)?.into_array()) } diff --git a/vortex-array/src/scalar_fn/fns/binary/mod.rs b/vortex-array/src/scalar_fn/fns/binary/mod.rs index 361c4fab240..7d8b88ac906 100644 --- a/vortex-array/src/scalar_fn/fns/binary/mod.rs +++ b/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), } } diff --git a/vortex-array/src/scalar_fn/fns/binary/numeric.rs b/vortex-array/src/scalar_fn/fns/binary/numeric.rs index ac76d2e4934..32ac14d74b4 100644 --- a/vortex-array/src/scalar_fn/fns/binary/numeric.rs +++ b/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 CheckedPrimitiveOp 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 { let ptype = PType::try_from(lhs.dtype())?; @@ -122,10 +128,10 @@ pub(crate) fn execute_numeric( match_each_native_ptype!(ptype, |T| { match op { - NumericOperator::Add => execute_checked_typed::(lhs, rhs, ctx), - NumericOperator::Sub => execute_checked_typed::(lhs, rhs, ctx), - NumericOperator::Mul => execute_checked_typed::(lhs, rhs, ctx), - NumericOperator::Div => execute_checked_typed::(lhs, rhs, ctx), + NumericOperator::Add => execute_checked_typed::(lhs, rhs, demand, ctx), + NumericOperator::Sub => execute_checked_typed::(lhs, rhs, demand, ctx), + NumericOperator::Mul => execute_checked_typed::(lhs, rhs, demand, ctx), + NumericOperator::Div => execute_checked_typed::(lhs, rhs, demand, ctx), } }) } @@ -133,6 +139,7 @@ pub(crate) fn execute_numeric( fn execute_checked_typed( lhs: &ArrayRef, rhs: &ArrayRef, + demand: &ArrayRef, ctx: &mut ExecutionCtx, ) -> VortexResult 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::(lhs, rhs, &valid_rows), + ) => checked_array_array::(lhs, rhs, &care_lanes), ( PrimitiveOperand::Array { values: lhs, .. }, PrimitiveOperand::Constant { value: rhs, .. }, - ) => checked_array_constant::(lhs, *rhs, &valid_rows), + ) => checked_array_constant::(lhs, *rhs, &care_lanes), ( PrimitiveOperand::Constant { value: lhs, .. }, PrimitiveOperand::Array { values: rhs, .. }, - ) => checked_constant_array::(*lhs, rhs, &valid_rows), + ) => checked_constant_array::(*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 CheckedValues { } } -fn checked_array_array(lhs: &[T], rhs: &[T], valid_rows: &Mask) -> CheckedValues +fn checked_array_array(lhs: &[T], rhs: &[T], care_lanes: &Mask) -> CheckedValues where T: NativePType, Op: CheckedPrimitiveOp, { 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::(lhs, rhs), AllOr::All => checked_array_array_all_lanes::(lhs, rhs), AllOr::None => CheckedValues::zeroed(lhs.len()), @@ -307,12 +324,12 @@ where } } -fn checked_array_constant(lhs: &[T], rhs: T, valid_rows: &Mask) -> CheckedValues +fn checked_array_constant(lhs: &[T], rhs: T, care_lanes: &Mask) -> CheckedValues where T: NativePType, Op: CheckedPrimitiveOp, { - match valid_rows.bit_buffer() { + match care_lanes.bit_buffer() { AllOr::All if Op::CHECKED_VALUE_LOOP => checked_array_constant_one_pass::(lhs, rhs), AllOr::All => checked_array_constant_all_lanes::(lhs, rhs), AllOr::None => CheckedValues::zeroed(lhs.len()), @@ -325,12 +342,12 @@ where } } -fn checked_constant_array(lhs: T, rhs: &[T], valid_rows: &Mask) -> CheckedValues +fn checked_constant_array(lhs: T, rhs: &[T], care_lanes: &Mask) -> CheckedValues where T: NativePType, Op: CheckedPrimitiveOp, { - match valid_rows.bit_buffer() { + match care_lanes.bit_buffer() { AllOr::All if Op::CHECKED_VALUE_LOOP => checked_constant_array_one_pass::(lhs, rhs), AllOr::All => checked_constant_array_all_lanes::(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 { + 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) -> VortexResult { 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(); diff --git a/vortex-array/src/scalar_fn/fns/dynamic.rs b/vortex-array/src/scalar_fn/fns/dynamic.rs index 74058aa19f2..346016b954d 100644 --- a/vortex-array/src/scalar_fn/fns/dynamic.rs +++ b/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(); diff --git a/vortex-array/src/scalar_fn/typed.rs b/vortex-array/src/scalar_fn/typed.rs index 1b14a9d613d..09886f202be 100644 --- a/vortex-array/src/scalar_fn/typed.rs +++ b/vortex-array/src/scalar_fn/typed.rs @@ -23,6 +23,7 @@ use vortex_error::VortexResult; use crate::ArrayRef; use crate::ExecutionCtx; use crate::dtype::DType; +use crate::dtype::Nullability; use crate::expr::Expression; use crate::scalar_fn::Arity; use crate::scalar_fn::ChildName; @@ -127,6 +128,20 @@ impl DynScalarFn for TypedScalarFnInstance { fn execute(&self, args: &dyn ExecutionArgs, ctx: &mut ExecutionCtx) -> VortexResult { let expected_row_count = args.row_count(); + assert_eq!( + args.demand().len(), + expected_row_count, + "Demand mask length {} does not match row count {} for {}", + args.demand().len(), + expected_row_count, + self.vtable.id(), + ); + assert_eq!( + args.demand().dtype(), + &DType::Bool(Nullability::NonNullable), + "Demand mask for {} must be a non-nullable boolean array", + self.vtable.id(), + ); #[cfg(debug_assertions)] let expected_dtype = { let args_dtypes: Vec = (0..args.num_inputs()) diff --git a/vortex-array/src/scalar_fn/vtable.rs b/vortex-array/src/scalar_fn/vtable.rs index c66afc34932..26c023378d7 100644 --- a/vortex-array/src/scalar_fn/vtable.rs +++ b/vortex-array/src/scalar_fn/vtable.rs @@ -18,7 +18,10 @@ use vortex_session::VortexSession; use crate::ArrayRef; use crate::ExecutionCtx; +use crate::IntoArray; +use crate::arrays::ConstantArray; use crate::dtype::DType; +use crate::dtype::Nullability; use crate::expr::Expression; use crate::expr::traversal::Node; use crate::scalar_fn::ScalarFnId; @@ -122,6 +125,14 @@ pub trait ScalarFnVTable: 'static + Sized + Clone + Send + Sync { /// Implementations are encouraged to check their inputs for constant arrays to perform /// more optimized execution. /// + /// # Demand + /// + /// [`ExecutionArgs::demand`] provides the rows whose output values the caller will + /// observe. Fallible implementations (see [`ScalarFnVTable::is_fallible`]) must not + /// raise a domain error attributable solely to undemanded rows. Infallible + /// implementations may ignore the demand mask, or use it to skip work. In all cases + /// the result must contain [`ExecutionArgs::row_count`] rows — demand never compacts. + /// /// If the input arguments cannot be directly used for execution (for example, an expression /// may require canonical input arrays), then the implementation should perform a single /// child execution and return a new [`crate::arrays::ScalarFnArray`] wrapping up the new child. @@ -326,18 +337,65 @@ pub trait ExecutionArgs { /// Returns the row count of the execution scope. fn row_count(&self) -> usize; + + /// Returns the demand mask: the rows whose output values the caller will observe. + /// + /// The demand mask is a non-nullable boolean array of exactly + /// [`ExecutionArgs::row_count`] length; typed dispatch asserts both. It may use any + /// encoding (constant, run-end, or a lazy expression); consumers typically + /// materialize it once, e.g. via [`child_to_validity`](crate::child_to_validity) + + /// [`Validity::execute_mask`](crate::validity::Validity::execute_mask). + /// + /// Undemanded rows are don't-care: the implementation may produce any well-typed + /// value (including null) at those positions, and fallible implementations must not + /// raise a domain error attributable solely to undemanded rows. Infallible + /// implementations may ignore this mask entirely; for them it is purely a + /// performance hint. + /// + /// Demand never compacts: the result of execution must still contain + /// [`ExecutionArgs::row_count`] rows. + fn demand(&self) -> &ArrayRef; } /// A concrete [`ExecutionArgs`] backed by a `Vec`. pub struct VecExecutionArgs { inputs: Vec, row_count: usize, + demand: ArrayRef, } impl VecExecutionArgs { - /// Create a new `VecExecutionArgs`. - pub fn new(inputs: Vec, row_count: usize) -> Self { - Self { inputs, row_count } + /// Create a new `VecExecutionArgs` with the given demand mask. + /// + /// # Panics + /// + /// Panics if `demand.len() != row_count` or the demand is not a non-nullable + /// boolean array. + pub fn new(inputs: Vec, row_count: usize, demand: ArrayRef) -> Self { + assert_eq!( + demand.len(), + row_count, + "Demand mask length must match the execution row count" + ); + assert_eq!( + demand.dtype(), + &DType::Bool(Nullability::NonNullable), + "Demand mask must be a non-nullable boolean array" + ); + Self { + inputs, + row_count, + demand, + } + } + + /// Create a new `VecExecutionArgs` demanding every row. + pub fn all(inputs: Vec, row_count: usize) -> Self { + Self::new( + inputs, + row_count, + ConstantArray::new(true, row_count).into_array(), + ) } } @@ -359,6 +417,10 @@ impl ExecutionArgs for VecExecutionArgs { fn row_count(&self) -> usize { self.row_count } + + fn demand(&self) -> &ArrayRef { + &self.demand + } } #[derive(Clone, Debug, PartialEq, Eq, Hash)] diff --git a/vortex-row/src/encoder.rs b/vortex-row/src/encoder.rs index c1d7a8f2a54..17142530bf0 100644 --- a/vortex-row/src/encoder.rs +++ b/vortex-row/src/encoder.rs @@ -119,7 +119,7 @@ impl RowEncoder { ); } } - Ok((options, VecExecutionArgs::new(cols.to_vec(), nrows))) + Ok((options, VecExecutionArgs::all(cols.to_vec(), nrows))) } }