From eb365f554e8ae16e97814c30482ae0b775cb7c07 Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Mon, 20 Jul 2026 13:15:26 +0000 Subject: [PATCH 1/8] feat: add mandatory demand mask to scalar_fn execution Introduce ExecutionArgs::demand(), a required Mask identifying the rows whose output values the caller will observe. Undemanded rows are don't-care: any well-typed value is acceptable, and fallible scalar fns must not raise a domain error attributable solely to undemanded rows. Demand never compacts: results keep the full row count. VecExecutionArgs::new now takes the demand mask explicitly (asserting its length matches the row count) and VecExecutionArgs::all constructs the all-demanded case. All construction sites pass AllTrue except the dynamic comparison fn, which forwards its incoming demand to its Binary delegate. TypedScalarFnInstance::execute asserts mask length at the execution choke point. No producer constructs a non-trivial mask yet; demand pushdown is deferred until all fallible fns honor the contract. Signed-off-by: Joe Isaacs Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01GhAvQmWLipNHBKjrZhxnKv --- .../src/arrays/scalar_fn/vtable/mod.rs | 2 +- .../src/arrays/scalar_fn/vtable/operations.rs | 2 +- .../src/arrays/scalar_fn/vtable/validity.rs | 2 +- vortex-array/src/scalar_fn/fns/dynamic.rs | 3 +- vortex-array/src/scalar_fn/typed.rs | 8 + vortex-array/src/scalar_fn/vtable.rs | 182 +++++++++++++++++- vortex-row/src/encoder.rs | 2 +- 7 files changed, 193 insertions(+), 8 deletions(-) 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/dynamic.rs b/vortex-array/src/scalar_fn/fns/dynamic.rs index 74058aa19f2..e652a773633 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); diff --git a/vortex-array/src/scalar_fn/typed.rs b/vortex-array/src/scalar_fn/typed.rs index 1b14a9d613d..3c2d8f9224c 100644 --- a/vortex-array/src/scalar_fn/typed.rs +++ b/vortex-array/src/scalar_fn/typed.rs @@ -127,6 +127,14 @@ 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(), + ); #[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..3881c99c451 100644 --- a/vortex-array/src/scalar_fn/vtable.rs +++ b/vortex-array/src/scalar_fn/vtable.rs @@ -14,6 +14,7 @@ use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_err; +use vortex_mask::Mask; use vortex_session::VortexSession; use crate::ArrayRef; @@ -122,6 +123,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 +335,49 @@ 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 mask is always [`ExecutionArgs::row_count`] long. 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) -> &Mask; } /// A concrete [`ExecutionArgs`] backed by a `Vec`. pub struct VecExecutionArgs { inputs: Vec, row_count: usize, + demand: Mask, } 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`. + pub fn new(inputs: Vec, row_count: usize, demand: Mask) -> Self { + assert_eq!( + demand.len(), + row_count, + "Demand mask length must match the execution row count" + ); + 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, Mask::new_true(row_count)) } } @@ -359,6 +399,10 @@ impl ExecutionArgs for VecExecutionArgs { fn row_count(&self) -> usize { self.row_count } + + fn demand(&self) -> &Mask { + &self.demand + } } #[derive(Clone, Debug, PartialEq, Eq, Hash)] @@ -398,3 +442,135 @@ impl ScalarFnVTableExt for V {} /// A reference to the name of a child expression. pub type ChildName = ArcRef; + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use parking_lot::Mutex; + use vortex_buffer::buffer; + use vortex_error::VortexResult; + use vortex_error::vortex_err; + use vortex_mask::Mask; + use vortex_session::registry::CachedId; + + use super::*; + use crate::ExecutionCtx; + use crate::IntoArray; + use crate::VortexSessionExecute; + use crate::array_session; + use crate::scalar_fn::ScalarFnVTableExt; + + /// A scalar fn that records the demand mask it was executed with. + #[derive(Clone)] + struct DemandProbe { + observed: Arc>>, + } + + impl ScalarFnVTable for DemandProbe { + type Options = EmptyOptions; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("test.demand_probe"); + *ID + } + + fn arity(&self, _options: &Self::Options) -> Arity { + Arity::Exact(1) + } + + fn child_name(&self, _options: &Self::Options, _child_idx: usize) -> ChildName { + ChildName::from("input") + } + + fn return_dtype(&self, _options: &Self::Options, args: &[DType]) -> VortexResult { + Ok(args[0].clone()) + } + + fn execute( + &self, + _options: &Self::Options, + args: &dyn ExecutionArgs, + _ctx: &mut ExecutionCtx, + ) -> VortexResult { + *self.observed.lock() = Some(args.demand().clone()); + args.get(0) + } + } + + #[test] + fn test_all_demands_every_row() { + let args = VecExecutionArgs::all(vec![buffer![1i32, 2, 3].into_array()], 3); + assert_eq!(args.demand().len(), 3); + assert!(args.demand().all_true()); + } + + #[test] + #[should_panic(expected = "Demand mask length must match")] + fn test_demand_length_mismatch_panics() { + drop(VecExecutionArgs::new( + vec![buffer![1i32, 2, 3].into_array()], + 3, + Mask::new_true(2), + )); + } + + #[test] + fn test_demand_reaches_scalar_fn() -> VortexResult<()> { + let observed = Arc::new(Mutex::new(None)); + let probe = DemandProbe { + observed: Arc::clone(&observed), + }; + let scalar_fn = probe.bind(EmptyOptions); + + let demand = Mask::from_iter([true, false, true]); + let args = VecExecutionArgs::new(vec![buffer![1i32, 2, 3].into_array()], 3, demand.clone()); + scalar_fn.execute(&args, &mut array_session().create_execution_ctx())?; + + let seen = observed + .lock() + .clone() + .ok_or_else(|| vortex_err!("probe was not executed"))?; + assert_eq!(seen.to_bit_buffer(), demand.to_bit_buffer()); + Ok(()) + } + + /// An [`ExecutionArgs`] whose demand mask length disagrees with its row count. + struct MismatchedDemandArgs { + input: ArrayRef, + demand: Mask, + } + + impl ExecutionArgs for MismatchedDemandArgs { + fn get(&self, _index: usize) -> VortexResult { + Ok(self.input.clone()) + } + + fn num_inputs(&self) -> usize { + 1 + } + + fn row_count(&self) -> usize { + self.input.len() + } + + fn demand(&self) -> &Mask { + &self.demand + } + } + + #[test] + #[should_panic(expected = "does not match row count")] + fn test_execute_rejects_mismatched_demand() { + let probe = DemandProbe { + observed: Arc::new(Mutex::new(None)), + }; + let scalar_fn = probe.bind(EmptyOptions); + + let args = MismatchedDemandArgs { + input: buffer![1i32, 2, 3].into_array(), + demand: Mask::new_true(2), + }; + drop(scalar_fn.execute(&args, &mut array_session().create_execution_ctx())); + } +} 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))) } } From ced7a88c440b7b92a441e4590234516d878ea41a Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Mon, 20 Jul 2026 13:19:38 +0000 Subject: [PATCH 2/8] review: use Mask equality in test, tighten demand doc, pin dynamic delegate forwarding Signed-off-by: Joe Isaacs Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01GhAvQmWLipNHBKjrZhxnKv --- vortex-array/src/scalar_fn/fns/dynamic.rs | 20 ++++++++++++++++++++ vortex-array/src/scalar_fn/vtable.rs | 5 +++-- 2 files changed, 23 insertions(+), 2 deletions(-) diff --git a/vortex-array/src/scalar_fn/fns/dynamic.rs b/vortex-array/src/scalar_fn/fns/dynamic.rs index e652a773633..b7685098df9 100644 --- a/vortex-array/src/scalar_fn/fns/dynamic.rs +++ b/vortex-array/src/scalar_fn/fns/dynamic.rs @@ -275,6 +275,7 @@ mod tests { use vortex_buffer::buffer; use vortex_error::VortexResult; + use vortex_mask::Mask; use super::*; use crate::IntoArray; @@ -356,6 +357,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 = Mask::from_iter([true, false, true]); + 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/vtable.rs b/vortex-array/src/scalar_fn/vtable.rs index 3881c99c451..bbbcf9f4d5f 100644 --- a/vortex-array/src/scalar_fn/vtable.rs +++ b/vortex-array/src/scalar_fn/vtable.rs @@ -338,7 +338,8 @@ pub trait ExecutionArgs { /// Returns the demand mask: the rows whose output values the caller will observe. /// - /// The mask is always [`ExecutionArgs::row_count`] long. Undemanded rows are + /// Implementations must return a mask of exactly [`ExecutionArgs::row_count`] length; + /// typed dispatch asserts this. 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 @@ -531,7 +532,7 @@ mod tests { .lock() .clone() .ok_or_else(|| vortex_err!("probe was not executed"))?; - assert_eq!(seen.to_bit_buffer(), demand.to_bit_buffer()); + assert_eq!(seen, demand); Ok(()) } From a22957292b237640a3786b53c2e4c8dc02829635 Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Tue, 21 Jul 2026 09:44:06 +0000 Subject: [PATCH 3/8] feat: honor demand mask in checked numeric operations Intersect the demand mask with the valid-lanes mask in the Binary checked arithmetic family (add/sub/mul/div) so domain errors at undemanded lanes are suppressed and one-pass kernels skip undemanded work. The all-true demand path is unchanged. Signed-off-by: Joe Isaacs Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01GhAvQmWLipNHBKjrZhxnKv --- vortex-array/src/scalar_fn/fns/binary/mod.rs | 8 +- .../src/scalar_fn/fns/binary/numeric.rs | 140 +++++++++++++++--- 2 files changed, 126 insertions(+), 22 deletions(-) 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..4b5ffc34407 100644 --- a/vortex-array/src/scalar_fn/fns/binary/numeric.rs +++ b/vortex-array/src/scalar_fn/fns/binary/numeric.rs @@ -105,10 +105,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: &Mask, ctx: &mut ExecutionCtx, ) -> VortexResult { let ptype = PType::try_from(lhs.dtype())?; @@ -122,10 +126,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 +137,7 @@ pub(crate) fn execute_numeric( fn execute_checked_typed( lhs: &ArrayRef, rhs: &ArrayRef, + demand: &Mask, ctx: &mut ExecutionCtx, ) -> VortexResult where @@ -157,28 +162,36 @@ where let validity = lhs.validity().and(rhs.validity())?; let valid_rows = validity.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 +302,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 +320,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 +338,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()), @@ -925,6 +938,7 @@ fn check_numeric_errors(failed: bool, error: &'static str) -> VortexResult<()> { mod test { use vortex_buffer::buffer; use vortex_error::VortexResult; + use vortex_mask::Mask; use crate::ArrayRef; use crate::IntoArray; @@ -936,9 +950,25 @@ mod test { 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: Mask, + ) -> 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 +1145,80 @@ 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, Mask::from_iter([false, true]))?; + + 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, Mask::from_iter([true, false])); + + 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, Mask::from_iter([false, true]))?; + + 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, Mask::new_false(2))?; + + 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, Mask::from_iter([true, false])); + + 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, Mask::from_iter([false, true]))?; + + assert_eq!(result.execute_scalar(1, &mut ctx)?, Scalar::from(5i32)); + 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, Mask::new_true(2))?; + + 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(); From 0c56959ed5ec9020e723a152b2205a6b68bf5237 Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Tue, 21 Jul 2026 09:47:14 +0000 Subject: [PATCH 4/8] review: cover demanded-error one-pass and constant-array shapes, reuse mask buffer Signed-off-by: Joe Isaacs Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01GhAvQmWLipNHBKjrZhxnKv --- .../src/scalar_fn/fns/binary/numeric.rs | 22 ++++++++++++++++++- 1 file changed, 21 insertions(+), 1 deletion(-) diff --git a/vortex-array/src/scalar_fn/fns/binary/numeric.rs b/vortex-array/src/scalar_fn/fns/binary/numeric.rs index 4b5ffc34407..eec03b24f22 100644 --- a/vortex-array/src/scalar_fn/fns/binary/numeric.rs +++ b/vortex-array/src/scalar_fn/fns/binary/numeric.rs @@ -167,7 +167,7 @@ where let care_lanes = if demand.all_true() { valid_rows } else { - &valid_rows & demand + valid_rows & demand }; let checked = match (&lhs, &rhs) { @@ -1207,6 +1207,26 @@ mod test { 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, Mask::from_iter([true, false])); + + 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, Mask::from_iter([false, true]))?; + + 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(); From 716ac07bdf7c686af702781d942560821cae5c33 Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Tue, 21 Jul 2026 10:32:04 +0000 Subject: [PATCH 5/8] refactor: collapse per-shape checked kernels into one care-lane dispatch The operand split (PrimitiveOperand) and the validity-and-demand merge (care_lanes) both happen before the kernels run, so the three shape dispatchers and twelve shape-specific kernel wrappers reduce to a single checked_op generic over a lane accessor. Slices and constants are bound outside the closures so codegen matches the old direct-slice kernels (verified with the binary_ops divan bench). Signed-off-by: Joe Isaacs Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01GhAvQmWLipNHBKjrZhxnKv --- .../src/scalar_fn/fns/binary/numeric.rs | 223 ++++-------------- 1 file changed, 44 insertions(+), 179 deletions(-) diff --git a/vortex-array/src/scalar_fn/fns/binary/numeric.rs b/vortex-array/src/scalar_fn/fns/binary/numeric.rs index eec03b24f22..31d6f654a97 100644 --- a/vortex-array/src/scalar_fn/fns/binary/numeric.rs +++ b/vortex-array/src/scalar_fn/fns/binary/numeric.rs @@ -174,15 +174,24 @@ where ( PrimitiveOperand::Array { values: lhs, .. }, PrimitiveOperand::Array { values: rhs, .. }, - ) => checked_array_array::(lhs, rhs, &care_lanes), + ) => { + let (lhs, rhs) = (lhs.as_slice(), rhs.as_slice()); + checked_op::(len, &care_lanes, |idx| (lhs[idx], rhs[idx])) + } ( PrimitiveOperand::Array { values: lhs, .. }, PrimitiveOperand::Constant { value: rhs, .. }, - ) => checked_array_constant::(lhs, *rhs, &care_lanes), + ) => { + let (lhs, rhs) = (lhs.as_slice(), *rhs); + checked_op::(len, &care_lanes, |idx| (lhs[idx], rhs)) + } ( PrimitiveOperand::Constant { value: lhs, .. }, PrimitiveOperand::Array { values: rhs, .. }, - ) => checked_constant_array::(*lhs, rhs, &care_lanes), + ) => { + let (lhs, rhs) = (*lhs, rhs.as_slice()); + checked_op::(len, &care_lanes, |idx| (lhs, rhs[idx])) + } ( PrimitiveOperand::Constant { value: lhs, .. }, PrimitiveOperand::Constant { value: rhs, .. }, @@ -302,132 +311,48 @@ impl 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 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()), - AllOr::Some(valid_bits) if Op::CHECKED_VALUE_LOOP => { - checked_array_array_valid_lanes_one_pass::(lhs, rhs, valid_bits) - } - AllOr::Some(valid_bits) => checked_array_array_valid_lanes::(lhs, rhs, valid_bits), - } -} - -fn checked_array_constant(lhs: &[T], rhs: T, care_lanes: &Mask) -> CheckedValues -where - T: NativePType, - Op: CheckedPrimitiveOp, -{ - 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()), - AllOr::Some(valid_bits) if Op::CHECKED_VALUE_LOOP => { - checked_array_constant_valid_lanes_one_pass::(lhs, rhs, valid_bits) - } - AllOr::Some(valid_bits) => { - checked_array_constant_valid_lanes::(lhs, rhs, valid_bits) - } - } -} - -fn checked_constant_array(lhs: T, rhs: &[T], care_lanes: &Mask) -> CheckedValues +/// Dispatch a checked op over the care lanes, reading operands via `operands_at`. +/// +/// The one-pass strategy short-circuits on the first observable error; the default +/// split strategy computes every lane branch-free and re-scans only care lanes for +/// errors, so undemanded and null lanes never fail the operation. +fn checked_op(len: usize, care_lanes: &Mask, operands_at: F) -> CheckedValues where T: NativePType, Op: CheckedPrimitiveOp, + F: Fn(usize) -> (T, T) + Copy, { 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()), - AllOr::Some(valid_bits) if Op::CHECKED_VALUE_LOOP => { - checked_constant_array_valid_lanes_one_pass::(lhs, rhs, valid_bits) + AllOr::All if Op::CHECKED_VALUE_LOOP => checked_all_lanes(len, |idx| { + let (lhs, rhs) = operands_at(idx); + Op::checked(lhs, rhs) + }), + AllOr::All => collect_all_lanes(len, |idx| { + let (lhs, rhs) = operands_at(idx); + Op::apply(lhs, rhs) + }), + AllOr::None => CheckedValues::zeroed(len), + AllOr::Some(care_bits) if Op::CHECKED_VALUE_LOOP => { + checked_valid_lanes(len, care_bits, |idx| { + let (lhs, rhs) = operands_at(idx); + Op::checked(lhs, rhs) + }) } - AllOr::Some(valid_bits) => { - checked_constant_array_valid_lanes::(lhs, rhs, valid_bits) + AllOr::Some(care_bits) => { + let mut checked = collect_all_lanes(len, |idx| { + let (lhs, rhs) = operands_at(idx); + Op::apply(lhs, rhs) + }); + checked.failed = checked.failed + && any_valid_error(len, care_bits, |idx| { + let (lhs, rhs) = operands_at(idx); + Op::apply(lhs, rhs).1 + }); + checked } } } -fn checked_array_array_all_lanes(lhs: &[T], rhs: &[T]) -> CheckedValues -where - T: NativePType, - Op: CheckedPrimitiveOp, -{ - collect_all_lanes(lhs.len(), |idx| Op::apply(lhs[idx], rhs[idx])) -} - -fn checked_array_array_valid_lanes( - lhs: &[T], - rhs: &[T], - valid_bits: &BitBuffer, -) -> CheckedValues -where - T: NativePType, - Op: CheckedPrimitiveOp, -{ - let mut checked = collect_all_lanes(lhs.len(), |idx| Op::apply(lhs[idx], rhs[idx])); - - checked.failed = checked.failed - && any_valid_error(lhs.len(), valid_bits, |idx| Op::apply(lhs[idx], rhs[idx]).1); - checked -} - -fn checked_array_constant_all_lanes(lhs: &[T], rhs: T) -> CheckedValues -where - T: NativePType, - Op: CheckedPrimitiveOp, -{ - collect_all_lanes(lhs.len(), |idx| Op::apply(lhs[idx], rhs)) -} - -fn checked_array_constant_valid_lanes( - lhs: &[T], - rhs: T, - valid_bits: &BitBuffer, -) -> CheckedValues -where - T: NativePType, - Op: CheckedPrimitiveOp, -{ - let mut checked = collect_all_lanes(lhs.len(), |idx| Op::apply(lhs[idx], rhs)); - - checked.failed = - checked.failed && any_valid_error(lhs.len(), valid_bits, |idx| Op::apply(lhs[idx], rhs).1); - checked -} - -fn checked_constant_array_all_lanes(lhs: T, rhs: &[T]) -> CheckedValues -where - T: NativePType, - Op: CheckedPrimitiveOp, -{ - collect_all_lanes(rhs.len(), |idx| Op::apply(lhs, rhs[idx])) -} - -fn checked_constant_array_valid_lanes( - lhs: T, - rhs: &[T], - valid_bits: &BitBuffer, -) -> CheckedValues -where - T: NativePType, - Op: CheckedPrimitiveOp, -{ - let mut checked = collect_all_lanes(rhs.len(), |idx| Op::apply(lhs, rhs[idx])); - - checked.failed = - checked.failed && any_valid_error(rhs.len(), valid_bits, |idx| Op::apply(lhs, rhs[idx]).1); - checked -} - fn collect_all_lanes(len: usize, mut value_and_error_at: F) -> CheckedValues where T: NativePType, @@ -450,66 +375,6 @@ where } } -fn checked_array_array_one_pass(lhs: &[T], rhs: &[T]) -> CheckedValues -where - T: NativePType, - Op: CheckedPrimitiveOp, -{ - checked_all_lanes(lhs.len(), |idx| Op::checked(lhs[idx], rhs[idx])) -} - -fn checked_array_array_valid_lanes_one_pass( - lhs: &[T], - rhs: &[T], - valid_bits: &BitBuffer, -) -> CheckedValues -where - T: NativePType, - Op: CheckedPrimitiveOp, -{ - checked_valid_lanes(lhs.len(), valid_bits, |idx| Op::checked(lhs[idx], rhs[idx])) -} - -fn checked_array_constant_one_pass(lhs: &[T], rhs: T) -> CheckedValues -where - T: NativePType, - Op: CheckedPrimitiveOp, -{ - checked_all_lanes(lhs.len(), |idx| Op::checked(lhs[idx], rhs)) -} - -fn checked_array_constant_valid_lanes_one_pass( - lhs: &[T], - rhs: T, - valid_bits: &BitBuffer, -) -> CheckedValues -where - T: NativePType, - Op: CheckedPrimitiveOp, -{ - checked_valid_lanes(lhs.len(), valid_bits, |idx| Op::checked(lhs[idx], rhs)) -} - -fn checked_constant_array_one_pass(lhs: T, rhs: &[T]) -> CheckedValues -where - T: NativePType, - Op: CheckedPrimitiveOp, -{ - checked_all_lanes(rhs.len(), |idx| Op::checked(lhs, rhs[idx])) -} - -fn checked_constant_array_valid_lanes_one_pass( - lhs: T, - rhs: &[T], - valid_bits: &BitBuffer, -) -> CheckedValues -where - T: NativePType, - Op: CheckedPrimitiveOp, -{ - checked_valid_lanes(rhs.len(), valid_bits, |idx| Op::checked(lhs, rhs[idx])) -} - // Checked one-pass ops delay early exit until the end of a small block. This // keeps the loop generic while avoiding a branch-driven exit decision on every // lane; it is deliberately independent of mask density or input length. From a97717cd9ba5198ee0214087f6a9274b335435e5 Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Tue, 21 Jul 2026 10:35:21 +0000 Subject: [PATCH 6/8] Revert "refactor: collapse per-shape checked kernels into one care-lane dispatch" This reverts commit 32c5287897299552896518135705421b9de85df9, restoring the per-shape checked kernels. Signed-off-by: Joe Isaacs Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01GhAvQmWLipNHBKjrZhxnKv --- .../src/scalar_fn/fns/binary/numeric.rs | 223 ++++++++++++++---- 1 file changed, 179 insertions(+), 44 deletions(-) diff --git a/vortex-array/src/scalar_fn/fns/binary/numeric.rs b/vortex-array/src/scalar_fn/fns/binary/numeric.rs index 31d6f654a97..eec03b24f22 100644 --- a/vortex-array/src/scalar_fn/fns/binary/numeric.rs +++ b/vortex-array/src/scalar_fn/fns/binary/numeric.rs @@ -174,24 +174,15 @@ where ( PrimitiveOperand::Array { values: lhs, .. }, PrimitiveOperand::Array { values: rhs, .. }, - ) => { - let (lhs, rhs) = (lhs.as_slice(), rhs.as_slice()); - checked_op::(len, &care_lanes, |idx| (lhs[idx], rhs[idx])) - } + ) => checked_array_array::(lhs, rhs, &care_lanes), ( PrimitiveOperand::Array { values: lhs, .. }, PrimitiveOperand::Constant { value: rhs, .. }, - ) => { - let (lhs, rhs) = (lhs.as_slice(), *rhs); - checked_op::(len, &care_lanes, |idx| (lhs[idx], rhs)) - } + ) => checked_array_constant::(lhs, *rhs, &care_lanes), ( PrimitiveOperand::Constant { value: lhs, .. }, PrimitiveOperand::Array { values: rhs, .. }, - ) => { - let (lhs, rhs) = (*lhs, rhs.as_slice()); - checked_op::(len, &care_lanes, |idx| (lhs, rhs[idx])) - } + ) => checked_constant_array::(*lhs, rhs, &care_lanes), ( PrimitiveOperand::Constant { value: lhs, .. }, PrimitiveOperand::Constant { value: rhs, .. }, @@ -311,48 +302,132 @@ impl CheckedValues { } } -/// Dispatch a checked op over the care lanes, reading operands via `operands_at`. -/// -/// The one-pass strategy short-circuits on the first observable error; the default -/// split strategy computes every lane branch-free and re-scans only care lanes for -/// errors, so undemanded and null lanes never fail the operation. -fn checked_op(len: usize, care_lanes: &Mask, operands_at: F) -> CheckedValues +fn checked_array_array(lhs: &[T], rhs: &[T], care_lanes: &Mask) -> CheckedValues where T: NativePType, Op: CheckedPrimitiveOp, - F: Fn(usize) -> (T, T) + Copy, { + debug_assert_eq!(lhs.len(), rhs.len()); + match care_lanes.bit_buffer() { - AllOr::All if Op::CHECKED_VALUE_LOOP => checked_all_lanes(len, |idx| { - let (lhs, rhs) = operands_at(idx); - Op::checked(lhs, rhs) - }), - AllOr::All => collect_all_lanes(len, |idx| { - let (lhs, rhs) = operands_at(idx); - Op::apply(lhs, rhs) - }), - AllOr::None => CheckedValues::zeroed(len), - AllOr::Some(care_bits) if Op::CHECKED_VALUE_LOOP => { - checked_valid_lanes(len, care_bits, |idx| { - let (lhs, rhs) = operands_at(idx); - Op::checked(lhs, rhs) - }) + 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()), + AllOr::Some(valid_bits) if Op::CHECKED_VALUE_LOOP => { + checked_array_array_valid_lanes_one_pass::(lhs, rhs, valid_bits) + } + AllOr::Some(valid_bits) => checked_array_array_valid_lanes::(lhs, rhs, valid_bits), + } +} + +fn checked_array_constant(lhs: &[T], rhs: T, care_lanes: &Mask) -> CheckedValues +where + T: NativePType, + Op: CheckedPrimitiveOp, +{ + 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()), + AllOr::Some(valid_bits) if Op::CHECKED_VALUE_LOOP => { + checked_array_constant_valid_lanes_one_pass::(lhs, rhs, valid_bits) + } + AllOr::Some(valid_bits) => { + checked_array_constant_valid_lanes::(lhs, rhs, valid_bits) + } + } +} + +fn checked_constant_array(lhs: T, rhs: &[T], care_lanes: &Mask) -> CheckedValues +where + T: NativePType, + Op: CheckedPrimitiveOp, +{ + 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()), + AllOr::Some(valid_bits) if Op::CHECKED_VALUE_LOOP => { + checked_constant_array_valid_lanes_one_pass::(lhs, rhs, valid_bits) } - AllOr::Some(care_bits) => { - let mut checked = collect_all_lanes(len, |idx| { - let (lhs, rhs) = operands_at(idx); - Op::apply(lhs, rhs) - }); - checked.failed = checked.failed - && any_valid_error(len, care_bits, |idx| { - let (lhs, rhs) = operands_at(idx); - Op::apply(lhs, rhs).1 - }); - checked + AllOr::Some(valid_bits) => { + checked_constant_array_valid_lanes::(lhs, rhs, valid_bits) } } } +fn checked_array_array_all_lanes(lhs: &[T], rhs: &[T]) -> CheckedValues +where + T: NativePType, + Op: CheckedPrimitiveOp, +{ + collect_all_lanes(lhs.len(), |idx| Op::apply(lhs[idx], rhs[idx])) +} + +fn checked_array_array_valid_lanes( + lhs: &[T], + rhs: &[T], + valid_bits: &BitBuffer, +) -> CheckedValues +where + T: NativePType, + Op: CheckedPrimitiveOp, +{ + let mut checked = collect_all_lanes(lhs.len(), |idx| Op::apply(lhs[idx], rhs[idx])); + + checked.failed = checked.failed + && any_valid_error(lhs.len(), valid_bits, |idx| Op::apply(lhs[idx], rhs[idx]).1); + checked +} + +fn checked_array_constant_all_lanes(lhs: &[T], rhs: T) -> CheckedValues +where + T: NativePType, + Op: CheckedPrimitiveOp, +{ + collect_all_lanes(lhs.len(), |idx| Op::apply(lhs[idx], rhs)) +} + +fn checked_array_constant_valid_lanes( + lhs: &[T], + rhs: T, + valid_bits: &BitBuffer, +) -> CheckedValues +where + T: NativePType, + Op: CheckedPrimitiveOp, +{ + let mut checked = collect_all_lanes(lhs.len(), |idx| Op::apply(lhs[idx], rhs)); + + checked.failed = + checked.failed && any_valid_error(lhs.len(), valid_bits, |idx| Op::apply(lhs[idx], rhs).1); + checked +} + +fn checked_constant_array_all_lanes(lhs: T, rhs: &[T]) -> CheckedValues +where + T: NativePType, + Op: CheckedPrimitiveOp, +{ + collect_all_lanes(rhs.len(), |idx| Op::apply(lhs, rhs[idx])) +} + +fn checked_constant_array_valid_lanes( + lhs: T, + rhs: &[T], + valid_bits: &BitBuffer, +) -> CheckedValues +where + T: NativePType, + Op: CheckedPrimitiveOp, +{ + let mut checked = collect_all_lanes(rhs.len(), |idx| Op::apply(lhs, rhs[idx])); + + checked.failed = + checked.failed && any_valid_error(rhs.len(), valid_bits, |idx| Op::apply(lhs, rhs[idx]).1); + checked +} + fn collect_all_lanes(len: usize, mut value_and_error_at: F) -> CheckedValues where T: NativePType, @@ -375,6 +450,66 @@ where } } +fn checked_array_array_one_pass(lhs: &[T], rhs: &[T]) -> CheckedValues +where + T: NativePType, + Op: CheckedPrimitiveOp, +{ + checked_all_lanes(lhs.len(), |idx| Op::checked(lhs[idx], rhs[idx])) +} + +fn checked_array_array_valid_lanes_one_pass( + lhs: &[T], + rhs: &[T], + valid_bits: &BitBuffer, +) -> CheckedValues +where + T: NativePType, + Op: CheckedPrimitiveOp, +{ + checked_valid_lanes(lhs.len(), valid_bits, |idx| Op::checked(lhs[idx], rhs[idx])) +} + +fn checked_array_constant_one_pass(lhs: &[T], rhs: T) -> CheckedValues +where + T: NativePType, + Op: CheckedPrimitiveOp, +{ + checked_all_lanes(lhs.len(), |idx| Op::checked(lhs[idx], rhs)) +} + +fn checked_array_constant_valid_lanes_one_pass( + lhs: &[T], + rhs: T, + valid_bits: &BitBuffer, +) -> CheckedValues +where + T: NativePType, + Op: CheckedPrimitiveOp, +{ + checked_valid_lanes(lhs.len(), valid_bits, |idx| Op::checked(lhs[idx], rhs)) +} + +fn checked_constant_array_one_pass(lhs: T, rhs: &[T]) -> CheckedValues +where + T: NativePType, + Op: CheckedPrimitiveOp, +{ + checked_all_lanes(rhs.len(), |idx| Op::checked(lhs, rhs[idx])) +} + +fn checked_constant_array_valid_lanes_one_pass( + lhs: T, + rhs: &[T], + valid_bits: &BitBuffer, +) -> CheckedValues +where + T: NativePType, + Op: CheckedPrimitiveOp, +{ + checked_valid_lanes(rhs.len(), valid_bits, |idx| Op::checked(lhs, rhs[idx])) +} + // Checked one-pass ops delay early exit until the end of a small block. This // keeps the loop generic while avoiding a branch-driven exit decision on every // lane; it is deliberately independent of mask density or input length. From 4842cc145c08b79b070247cad170f4e778e68233 Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Tue, 21 Jul 2026 13:27:46 +0000 Subject: [PATCH 7/8] feat: represent the demand mask as a non-nullable boolean array ExecutionArgs::demand now returns an ArrayRef (Bool NonNullable) instead of a vortex-mask Mask, so demand can use any array encoding: constant, run-end, or a lazy expression. Typed dispatch asserts both the length and the dtype invariant. Consumers materialize once at the kernel boundary via child_to_validity + Validity::execute_mask, which keeps the constant all-true case O(1) (verified with the binary_ops bench). Signed-off-by: Joe Isaacs Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01GhAvQmWLipNHBKjrZhxnKv --- .../src/scalar_fn/fns/binary/numeric.rs | 77 ++++++++++++++--- vortex-array/src/scalar_fn/fns/dynamic.rs | 3 +- vortex-array/src/scalar_fn/typed.rs | 7 ++ vortex-array/src/scalar_fn/vtable.rs | 85 +++++++++++++------ 4 files changed, 132 insertions(+), 40 deletions(-) diff --git a/vortex-array/src/scalar_fn/fns/binary/numeric.rs b/vortex-array/src/scalar_fn/fns/binary/numeric.rs index eec03b24f22..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; @@ -112,7 +114,7 @@ pub(crate) fn execute_numeric( lhs: &ArrayRef, rhs: &ArrayRef, op: NumericOperator, - demand: &Mask, + demand: &ArrayRef, ctx: &mut ExecutionCtx, ) -> VortexResult { let ptype = PType::try_from(lhs.dtype())?; @@ -137,7 +139,7 @@ pub(crate) fn execute_numeric( fn execute_checked_typed( lhs: &ArrayRef, rhs: &ArrayRef, - demand: &Mask, + demand: &ArrayRef, ctx: &mut ExecutionCtx, ) -> VortexResult where @@ -162,12 +164,14 @@ 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 + valid_rows & &demand }; let checked = match (&lhs, &rhs) { @@ -938,13 +942,13 @@ fn check_numeric_errors(failed: bool, error: &'static str) -> VortexResult<()> { mod test { use vortex_buffer::buffer; use vortex_error::VortexResult; - use vortex_mask::Mask; use crate::ArrayRef; use crate::IntoArray; 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; @@ -960,7 +964,7 @@ mod test { lhs: ArrayRef, rhs: ArrayRef, op: Operator, - demand: Mask, + demand: ArrayRef, ) -> VortexResult { let len = lhs.len(); let args = VecExecutionArgs::new(vec![lhs, rhs], len, demand); @@ -1150,7 +1154,12 @@ mod test { 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, Mask::from_iter([false, true]))?; + 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)); @@ -1161,7 +1170,12 @@ mod 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, Mask::from_iter([true, false])); + let result = execute_with_demand( + lhs, + rhs, + Operator::Add, + BoolArray::from_iter([true, false]).into_array(), + ); assert!(result.is_err()); } @@ -1171,7 +1185,12 @@ mod test { 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, Mask::from_iter([false, true]))?; + 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(()) @@ -1181,7 +1200,12 @@ mod 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, Mask::new_false(2))?; + let result = execute_with_demand( + lhs, + rhs, + Operator::Add, + ConstantArray::new(false, 2).into_array(), + )?; assert_eq!(result.len(), 2); Ok(()) @@ -1191,7 +1215,12 @@ mod 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, Mask::from_iter([true, false])); + let result = execute_with_demand( + lhs, + rhs, + Operator::Add, + BoolArray::from_iter([true, false]).into_array(), + ); assert!(result.is_err()); } @@ -1201,7 +1230,12 @@ mod test { 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, Mask::from_iter([false, true]))?; + 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(()) @@ -1211,7 +1245,12 @@ mod 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, Mask::from_iter([true, false])); + let result = execute_with_demand( + lhs, + rhs, + Operator::Div, + BoolArray::from_iter([true, false]).into_array(), + ); assert!(result.is_err()); } @@ -1221,7 +1260,12 @@ mod test { 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, Mask::from_iter([false, true]))?; + 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(()) @@ -1233,7 +1277,12 @@ mod test { 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, Mask::new_true(2))?; + 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(()) diff --git a/vortex-array/src/scalar_fn/fns/dynamic.rs b/vortex-array/src/scalar_fn/fns/dynamic.rs index b7685098df9..346016b954d 100644 --- a/vortex-array/src/scalar_fn/fns/dynamic.rs +++ b/vortex-array/src/scalar_fn/fns/dynamic.rs @@ -275,7 +275,6 @@ mod tests { use vortex_buffer::buffer; use vortex_error::VortexResult; - use vortex_mask::Mask; use super::*; use crate::IntoArray; @@ -369,7 +368,7 @@ mod tests { default: true, }); - let demand = Mask::from_iter([true, false, 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); diff --git a/vortex-array/src/scalar_fn/typed.rs b/vortex-array/src/scalar_fn/typed.rs index 3c2d8f9224c..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; @@ -135,6 +136,12 @@ impl DynScalarFn for TypedScalarFnInstance { 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 bbbcf9f4d5f..228a06b2544 100644 --- a/vortex-array/src/scalar_fn/vtable.rs +++ b/vortex-array/src/scalar_fn/vtable.rs @@ -14,12 +14,14 @@ use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_err; -use vortex_mask::Mask; 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; @@ -338,23 +340,28 @@ pub trait ExecutionArgs { /// Returns the demand mask: the rows whose output values the caller will observe. /// - /// Implementations must return a mask of exactly [`ExecutionArgs::row_count`] length; - /// typed dispatch asserts this. 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. + /// 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) -> &Mask; + fn demand(&self) -> &ArrayRef; } /// A concrete [`ExecutionArgs`] backed by a `Vec`. pub struct VecExecutionArgs { inputs: Vec, row_count: usize, - demand: Mask, + demand: ArrayRef, } impl VecExecutionArgs { @@ -362,13 +369,19 @@ impl VecExecutionArgs { /// /// # Panics /// - /// Panics if `demand.len() != row_count`. - pub fn new(inputs: Vec, row_count: usize, demand: Mask) -> Self { + /// 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, @@ -378,7 +391,11 @@ impl VecExecutionArgs { /// Create a new `VecExecutionArgs` demanding every row. pub fn all(inputs: Vec, row_count: usize) -> Self { - Self::new(inputs, row_count, Mask::new_true(row_count)) + Self::new( + inputs, + row_count, + ConstantArray::new(true, row_count).into_array(), + ) } } @@ -401,7 +418,7 @@ impl ExecutionArgs for VecExecutionArgs { self.row_count } - fn demand(&self) -> &Mask { + fn demand(&self) -> &ArrayRef { &self.demand } } @@ -452,7 +469,6 @@ mod tests { use vortex_buffer::buffer; use vortex_error::VortexResult; use vortex_error::vortex_err; - use vortex_mask::Mask; use vortex_session::registry::CachedId; use super::*; @@ -460,12 +476,16 @@ mod tests { use crate::IntoArray; use crate::VortexSessionExecute; use crate::array_session; + use crate::arrays::BoolArray; + use crate::arrays::Constant; + use crate::assert_arrays_eq; + use crate::scalar::Scalar; use crate::scalar_fn::ScalarFnVTableExt; /// A scalar fn that records the demand mask it was executed with. #[derive(Clone)] struct DemandProbe { - observed: Arc>>, + observed: Arc>>, } impl ScalarFnVTable for DemandProbe { @@ -500,10 +520,16 @@ mod tests { } #[test] - fn test_all_demands_every_row() { + fn test_all_demands_every_row() -> VortexResult<()> { let args = VecExecutionArgs::all(vec![buffer![1i32, 2, 3].into_array()], 3); assert_eq!(args.demand().len(), 3); - assert!(args.demand().all_true()); + + let constant = args + .demand() + .as_opt::() + .ok_or_else(|| vortex_err!("all-demand should be a constant array"))?; + assert_eq!(constant.scalar(), &Scalar::from(true)); + Ok(()) } #[test] @@ -512,34 +538,45 @@ mod tests { drop(VecExecutionArgs::new( vec![buffer![1i32, 2, 3].into_array()], 3, - Mask::new_true(2), + BoolArray::from_iter([true, false]).into_array(), + )); + } + + #[test] + #[should_panic(expected = "must be a non-nullable boolean array")] + fn test_nullable_demand_panics() { + drop(VecExecutionArgs::new( + vec![buffer![1i32, 2, 3].into_array()], + 3, + BoolArray::from_iter([Some(true), Some(false), None]).into_array(), )); } #[test] fn test_demand_reaches_scalar_fn() -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); let observed = Arc::new(Mutex::new(None)); let probe = DemandProbe { observed: Arc::clone(&observed), }; let scalar_fn = probe.bind(EmptyOptions); - let demand = Mask::from_iter([true, false, true]); + let demand = BoolArray::from_iter([true, false, true]).into_array(); let args = VecExecutionArgs::new(vec![buffer![1i32, 2, 3].into_array()], 3, demand.clone()); - scalar_fn.execute(&args, &mut array_session().create_execution_ctx())?; + scalar_fn.execute(&args, &mut ctx)?; let seen = observed .lock() .clone() .ok_or_else(|| vortex_err!("probe was not executed"))?; - assert_eq!(seen, demand); + assert_arrays_eq!(seen, demand, &mut ctx); Ok(()) } /// An [`ExecutionArgs`] whose demand mask length disagrees with its row count. struct MismatchedDemandArgs { input: ArrayRef, - demand: Mask, + demand: ArrayRef, } impl ExecutionArgs for MismatchedDemandArgs { @@ -555,7 +592,7 @@ mod tests { self.input.len() } - fn demand(&self) -> &Mask { + fn demand(&self) -> &ArrayRef { &self.demand } } @@ -570,7 +607,7 @@ mod tests { let args = MismatchedDemandArgs { input: buffer![1i32, 2, 3].into_array(), - demand: Mask::new_true(2), + demand: BoolArray::from_iter([true, false]).into_array(), }; drop(scalar_fn.execute(&args, &mut array_session().create_execution_ctx())); } From fcdd188dd09429bcbffc41d0c6f84a6d3c4a8989 Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Tue, 21 Jul 2026 14:33:56 +0100 Subject: [PATCH 8/8] u Signed-off-by: Joe Isaacs --- vortex-array/src/scalar_fn/vtable.rs | 152 --------------------------- 1 file changed, 152 deletions(-) diff --git a/vortex-array/src/scalar_fn/vtable.rs b/vortex-array/src/scalar_fn/vtable.rs index 228a06b2544..26c023378d7 100644 --- a/vortex-array/src/scalar_fn/vtable.rs +++ b/vortex-array/src/scalar_fn/vtable.rs @@ -460,155 +460,3 @@ impl ScalarFnVTableExt for V {} /// A reference to the name of a child expression. pub type ChildName = ArcRef; - -#[cfg(test)] -mod tests { - use std::sync::Arc; - - use parking_lot::Mutex; - use vortex_buffer::buffer; - use vortex_error::VortexResult; - use vortex_error::vortex_err; - use vortex_session::registry::CachedId; - - use super::*; - use crate::ExecutionCtx; - use crate::IntoArray; - use crate::VortexSessionExecute; - use crate::array_session; - use crate::arrays::BoolArray; - use crate::arrays::Constant; - use crate::assert_arrays_eq; - use crate::scalar::Scalar; - use crate::scalar_fn::ScalarFnVTableExt; - - /// A scalar fn that records the demand mask it was executed with. - #[derive(Clone)] - struct DemandProbe { - observed: Arc>>, - } - - impl ScalarFnVTable for DemandProbe { - type Options = EmptyOptions; - - fn id(&self) -> ScalarFnId { - static ID: CachedId = CachedId::new("test.demand_probe"); - *ID - } - - fn arity(&self, _options: &Self::Options) -> Arity { - Arity::Exact(1) - } - - fn child_name(&self, _options: &Self::Options, _child_idx: usize) -> ChildName { - ChildName::from("input") - } - - fn return_dtype(&self, _options: &Self::Options, args: &[DType]) -> VortexResult { - Ok(args[0].clone()) - } - - fn execute( - &self, - _options: &Self::Options, - args: &dyn ExecutionArgs, - _ctx: &mut ExecutionCtx, - ) -> VortexResult { - *self.observed.lock() = Some(args.demand().clone()); - args.get(0) - } - } - - #[test] - fn test_all_demands_every_row() -> VortexResult<()> { - let args = VecExecutionArgs::all(vec![buffer![1i32, 2, 3].into_array()], 3); - assert_eq!(args.demand().len(), 3); - - let constant = args - .demand() - .as_opt::() - .ok_or_else(|| vortex_err!("all-demand should be a constant array"))?; - assert_eq!(constant.scalar(), &Scalar::from(true)); - Ok(()) - } - - #[test] - #[should_panic(expected = "Demand mask length must match")] - fn test_demand_length_mismatch_panics() { - drop(VecExecutionArgs::new( - vec![buffer![1i32, 2, 3].into_array()], - 3, - BoolArray::from_iter([true, false]).into_array(), - )); - } - - #[test] - #[should_panic(expected = "must be a non-nullable boolean array")] - fn test_nullable_demand_panics() { - drop(VecExecutionArgs::new( - vec![buffer![1i32, 2, 3].into_array()], - 3, - BoolArray::from_iter([Some(true), Some(false), None]).into_array(), - )); - } - - #[test] - fn test_demand_reaches_scalar_fn() -> VortexResult<()> { - let mut ctx = array_session().create_execution_ctx(); - let observed = Arc::new(Mutex::new(None)); - let probe = DemandProbe { - observed: Arc::clone(&observed), - }; - let scalar_fn = probe.bind(EmptyOptions); - - let demand = BoolArray::from_iter([true, false, true]).into_array(); - let args = VecExecutionArgs::new(vec![buffer![1i32, 2, 3].into_array()], 3, demand.clone()); - scalar_fn.execute(&args, &mut ctx)?; - - let seen = observed - .lock() - .clone() - .ok_or_else(|| vortex_err!("probe was not executed"))?; - assert_arrays_eq!(seen, demand, &mut ctx); - Ok(()) - } - - /// An [`ExecutionArgs`] whose demand mask length disagrees with its row count. - struct MismatchedDemandArgs { - input: ArrayRef, - demand: ArrayRef, - } - - impl ExecutionArgs for MismatchedDemandArgs { - fn get(&self, _index: usize) -> VortexResult { - Ok(self.input.clone()) - } - - fn num_inputs(&self) -> usize { - 1 - } - - fn row_count(&self) -> usize { - self.input.len() - } - - fn demand(&self) -> &ArrayRef { - &self.demand - } - } - - #[test] - #[should_panic(expected = "does not match row count")] - fn test_execute_rejects_mismatched_demand() { - let probe = DemandProbe { - observed: Arc::new(Mutex::new(None)), - }; - let scalar_fn = probe.bind(EmptyOptions); - - let args = MismatchedDemandArgs { - input: buffer![1i32, 2, 3].into_array(), - demand: BoolArray::from_iter([true, false]).into_array(), - }; - drop(scalar_fn.execute(&args, &mut array_session().create_execution_ctx())); - } -}