From cfbd5133d9c8668a1bdb5e36af38cc3a3bb10c01 Mon Sep 17 00:00:00 2001 From: Connor Tsui Date: Wed, 21 Oct 2026 11:05:05 -0400 Subject: [PATCH] Attach RowFn validity during Boolean output construction Let the packed Boolean collectors consume the batch validity when constructing their result. Keep custom output validation before masking and attach validity directly to other canonical outputs. Extracted from #9979. Signed-off-by: Connor Tsui --- .../unstable/row/batch/execute/constant.rs | 9 ++- .../unstable/row/batch/execute/dense.rs | 39 ++++++++++++- .../unstable/row/batch/execute/mod.rs | 6 +- .../unstable/row/batch/execute/output.rs | 2 +- .../unstable/row/batch/execute/valid_only.rs | 13 ++++- .../unstable/row/execute/packed_bool.rs | 11 ++-- .../scalar_fn/unstable/row/execute/tests.rs | 58 ++++++++++++++++--- .../scalar_fn/unstable/row/visitor/execute.rs | 14 ++++- .../src/scalar_fn/unstable/row/vtable.rs | 54 ++++++++++++++++- 9 files changed, 181 insertions(+), 25 deletions(-) diff --git a/vortex-array/src/scalar_fn/unstable/row/batch/execute/constant.rs b/vortex-array/src/scalar_fn/unstable/row/batch/execute/constant.rs index bcc61068562..bc8094a1b01 100644 --- a/vortex-array/src/scalar_fn/unstable/row/batch/execute/constant.rs +++ b/vortex-array/src/scalar_fn/unstable/row/batch/execute/constant.rs @@ -9,16 +9,21 @@ use crate::ArrayRef; use crate::ExecutionCtx; use crate::IntoArray; use crate::arrays::ConstantArray; +use crate::validity::Validity; impl RowFnExecutionArgs { /// Execute all-constant inputs by evaluating one row and broadcasting the validated result. pub(super) fn execute_all_constant( &self, - kernel: impl Fn(BorrowedRowFnArgs<'_>, &mut ExecutionCtx) -> VortexResult, + kernel: impl FnOnce( + BorrowedRowFnArgs<'_>, + &mut Option, + &mut ExecutionCtx, + ) -> VortexResult, ctx: &mut ExecutionCtx, ) -> VortexResult { let result = self.validate_kernel_output( - kernel(self.execution_args(&self.inputs, 1), ctx)?, + kernel(self.execution_args(&self.inputs, 1), &mut None, ctx)?, 1, ctx, )?; diff --git a/vortex-array/src/scalar_fn/unstable/row/batch/execute/dense.rs b/vortex-array/src/scalar_fn/unstable/row/batch/execute/dense.rs index cd5ab2aeea1..e539d268647 100644 --- a/vortex-array/src/scalar_fn/unstable/row/batch/execute/dense.rs +++ b/vortex-array/src/scalar_fn/unstable/row/batch/execute/dense.rs @@ -9,8 +9,12 @@ use vortex_mask::MaskValuesRef; use super::super::RowFnExecutionArgs; use super::super::args::BorrowedRowFnArgs; use crate::ArrayRef; +use crate::Canonical; use crate::ExecutionCtx; +use crate::IntoArray; +use crate::arrays::masked::mask_validity_canonical; use crate::builtins::ArrayBuiltins; +use crate::canonical::AnyCanonical; use crate::scalar_fn::unstable::row::execute::DenseAttempt; use crate::validity::Validity; @@ -18,10 +22,26 @@ impl RowFnExecutionArgs { /// Run every stored payload, then attach the input validity without materializing its mask. pub(super) fn execute_dense( &self, - kernel: impl Fn(BorrowedRowFnArgs<'_>, &mut ExecutionCtx) -> VortexResult, + kernel: impl FnOnce( + BorrowedRowFnArgs<'_>, + &mut Option, + &mut ExecutionCtx, + ) -> VortexResult, ctx: &mut ExecutionCtx, ) -> VortexResult { - let values = kernel(self.execution_args(&self.inputs, self.row_count), ctx)?; + // Offer a partial validity to the kernel. A collector that takes it builds the output with + // that validity attached, so the output needs neither a second construction nor an + // all-valid check. + let mut validity = + matches!(self.validity, Validity::Array(_)).then(|| self.validity.clone()); + let values = kernel( + self.execution_args(&self.inputs, self.row_count), + &mut validity, + ctx, + )?; + if validity.is_none() && matches!(self.validity, Validity::Array(_)) { + return self.finalize_output(values, self.row_count); + } self.finalize_dense_output(values, ctx) } @@ -107,7 +127,20 @@ impl RowFnExecutionArgs { Validity::NonNullable | Validity::AllValid => { self.finalize_output(values, self.row_count) } - Validity::Array(valid) => self.finalize_output(values.mask(valid)?, self.row_count), + // A canonical kernel output takes the validity directly. Masking it lazily would add + // a scalar-function node that the caller must execute again just to attach it. + Validity::Array(valid) => { + let values = match values.as_opt::() { + Some(canonical) => mask_validity_canonical( + Canonical::from(canonical), + Validity::Array(valid), + ctx, + )? + .into_array(), + None => values.mask(valid)?, + }; + self.finalize_output(values, self.row_count) + } Validity::AllInvalid => { unreachable!("all-invalid validity is handled before dense row execution") } diff --git a/vortex-array/src/scalar_fn/unstable/row/batch/execute/mod.rs b/vortex-array/src/scalar_fn/unstable/row/batch/execute/mod.rs index 7298abfc9bc..d1be51a240b 100644 --- a/vortex-array/src/scalar_fn/unstable/row/batch/execute/mod.rs +++ b/vortex-array/src/scalar_fn/unstable/row/batch/execute/mod.rs @@ -36,7 +36,11 @@ impl RowFnExecutionArgs { /// and dtype. pub(crate) fn execute( &self, - kernel: impl Fn(BorrowedRowFnArgs<'_>, &mut ExecutionCtx) -> VortexResult, + kernel: impl FnOnce( + BorrowedRowFnArgs<'_>, + &mut Option, + &mut ExecutionCtx, + ) -> VortexResult, execute_dense_attempt: impl FnOnce( BorrowedRowFnArgs<'_>, &mut ExecutionCtx, diff --git a/vortex-array/src/scalar_fn/unstable/row/batch/execute/output.rs b/vortex-array/src/scalar_fn/unstable/row/batch/execute/output.rs index 6afea5798dd..b77ad6f5285 100644 --- a/vortex-array/src/scalar_fn/unstable/row/batch/execute/output.rs +++ b/vortex-array/src/scalar_fn/unstable/row/batch/execute/output.rs @@ -68,7 +68,7 @@ pub(crate) fn finalize_kernel_output( ) -> VortexResult { validate_output(id, result_dtype, expected_len, &values)?; vortex_ensure!( - values.all_valid(ctx)?, + !values.dtype().is_nullable() || values.all_valid(ctx)?, "the {id} row kernel must produce only valid rows, got at least one null row", ); diff --git a/vortex-array/src/scalar_fn/unstable/row/batch/execute/valid_only.rs b/vortex-array/src/scalar_fn/unstable/row/batch/execute/valid_only.rs index 62c5afb7458..88be928a551 100644 --- a/vortex-array/src/scalar_fn/unstable/row/batch/execute/valid_only.rs +++ b/vortex-array/src/scalar_fn/unstable/row/batch/execute/valid_only.rs @@ -11,12 +11,17 @@ use crate::ArrayRef; use crate::ExecutionCtx; use crate::IntoArray; use crate::builtins::ArrayBuiltins; +use crate::validity::Validity; impl RowFnExecutionArgs { /// Resolve validity and try direct valid-row execution before filtered execution. pub(super) fn execute_valid_only( &self, - kernel: impl Fn(BorrowedRowFnArgs<'_>, &mut ExecutionCtx) -> VortexResult, + kernel: impl FnOnce( + BorrowedRowFnArgs<'_>, + &mut Option, + &mut ExecutionCtx, + ) -> VortexResult, try_valid_rows: impl FnOnce( BorrowedRowFnArgs<'_>, MaskValuesRef, @@ -34,7 +39,11 @@ impl RowFnExecutionArgs { let valid_rows = match validity { // An empty mask is both all-valid and all-null. Preserve the all-valid behavior. Mask::AllTrue(_) | Mask::AllFalse(0) => { - let values = kernel(self.execution_args(&self.inputs, self.row_count), ctx)?; + let values = kernel( + self.execution_args(&self.inputs, self.row_count), + &mut None, + ctx, + )?; let values = self.validate_kernel_output(values, self.row_count, ctx)?; return self.finalize_output(values, self.row_count); diff --git a/vortex-array/src/scalar_fn/unstable/row/execute/packed_bool.rs b/vortex-array/src/scalar_fn/unstable/row/execute/packed_bool.rs index 5fb3d0cadc1..f99bc1b0369 100644 --- a/vortex-array/src/scalar_fn/unstable/row/execute/packed_bool.rs +++ b/vortex-array/src/scalar_fn/unstable/row/execute/packed_bool.rs @@ -28,10 +28,11 @@ use crate::scalar_fn::unstable::row::IndexedElementTuple; use crate::scalar_fn::unstable::row::types::decoded_source; use crate::validity::Validity; -/// Decode every input column, then pack infallible Boolean outputs. +/// Decode every input column, then pack infallible Boolean outputs with `validity` attached. pub(crate) fn execute_owned_infallible_bool( args: &dyn ExecutionArgs, ctx: &mut ExecutionCtx, + validity: Validity, apply: impl Fn(Args::Elems<'_>) -> bool, ) -> VortexResult where @@ -55,13 +56,15 @@ where BitBuffer::collect_bool_in(row_count, collect, ctx.allocator().clone()) }; - Ok(BoolArray::new(values, Validity::NonNullable).into_array()) + Ok(BoolArray::try_new(values, validity)?.into_array()) } -/// Decode every input column, then pack Boolean outputs while combining failure evidence. +/// Decode every input column, then pack Boolean outputs with `validity` attached while combining +/// failure evidence. pub(crate) fn execute_owned_bool( args: &dyn ExecutionArgs, ctx: &mut ExecutionCtx, + validity: Validity, prepare: impl FnOnce(Args::ConstElems<'_>) -> Prepared, apply: impl Fn(&Prepared, Args::Elems<'_>) -> (bool, Fail), finish_failure: impl FnOnce(Fail) -> VortexResult<()>, @@ -104,7 +107,7 @@ where finish_failure(failure)?; - Ok(BoolArray::new(values, Validity::NonNullable).into_array()) + Ok(BoolArray::try_new(values, validity)?.into_array()) } /// Pack a dense Boolean attempt and report rejected failure evidence to the batch executor. diff --git a/vortex-array/src/scalar_fn/unstable/row/execute/tests.rs b/vortex-array/src/scalar_fn/unstable/row/execute/tests.rs index 40ab20e255e..3f9f460e61d 100644 --- a/vortex-array/src/scalar_fn/unstable/row/execute/tests.rs +++ b/vortex-array/src/scalar_fn/unstable/row/execute/tests.rs @@ -32,6 +32,7 @@ use crate::IntoArray; use crate::VortexSessionExecute; use crate::array_session; use crate::arrays::Bool; +use crate::arrays::BoolArray; use crate::arrays::ConstantArray; use crate::arrays::FixedSizeList; use crate::arrays::Primitive; @@ -190,15 +191,22 @@ fn packed_boolean_payload_uses_allocator( .with_allocator(allocator); let output = match (multiversioned, traversal) { - (false, Traversal::Infallible) => { - execute_owned_infallible_bool::<(i64,), false>(&args, &mut ctx, |(value,)| value > 2)? - } - (true, Traversal::Infallible) => { - execute_owned_infallible_bool::<(i64,), true>(&args, &mut ctx, |(value,)| value > 2)? - } + (false, Traversal::Infallible) => execute_owned_infallible_bool::<(i64,), false>( + &args, + &mut ctx, + Validity::NonNullable, + |(value,)| value > 2, + )?, + (true, Traversal::Infallible) => execute_owned_infallible_bool::<(i64,), true>( + &args, + &mut ctx, + Validity::NonNullable, + |(value,)| value > 2, + )?, (false, Traversal::Fallible) => execute_owned_bool::<(i64,), (), bool, false>( &args, &mut ctx, + Validity::NonNullable, |_| (), |_, (value,)| (value > 2, false), |_| Ok(()), @@ -206,6 +214,7 @@ fn packed_boolean_payload_uses_allocator( (true, Traversal::Fallible) => execute_owned_bool::<(i64,), (), bool, true>( &args, &mut ctx, + Validity::NonNullable, |_| (), |_, (value,)| (value > 2, false), |_| Ok(()), @@ -347,8 +356,12 @@ fn empty_outputs_and_zero_width_rows_do_not_allocate_payloads() -> VortexResult< let primitive = execute_owned_infallible::<(i64,), i64, ()>(&args, &mut ctx, |_| (), |_, (value,)| value)?; - let boolean = - execute_owned_infallible_bool::<(i64,), false>(&args, &mut ctx, |(value,)| value > 0)?; + let boolean = execute_owned_infallible_bool::<(i64,), false>( + &args, + &mut ctx, + Validity::NonNullable, + |(value,)| value > 0, + )?; let scalar = execute_sink::<(i64,), (), UninitElementSink, _>( &args, &(), @@ -495,3 +508,32 @@ fn zero_sized_output_uses_its_own_storage(#[case] traversal: Traversal) -> Vorte Ok(()) } + +#[rstest] +fn packed_boolean_attaches_validity( + #[values(false, true)] multiversioned: bool, +) -> VortexResult<()> { + let values = PrimitiveArray::from_iter([1_i64, 2, 3, 4]).into_array(); + let args = VecExecutionArgs::new(vec![values], 4); + let validity = Validity::from_iter([true, false, true, true]); + let mut ctx = array_session().create_execution_ctx(); + + let output = if multiversioned { + execute_owned_infallible_bool::<(i64,), true>(&args, &mut ctx, validity, |(value,)| { + value > 2 + })? + } else { + execute_owned_infallible_bool::<(i64,), false>(&args, &mut ctx, validity, |(value,)| { + value > 2 + })? + }; + + assert!(output.is::()); + let expected = BoolArray::new( + [false, false, true, true].into_iter().collect(), + Validity::from_iter([true, false, true, true]), + ) + .into_array(); + assert_arrays_eq!(&output, &expected, &mut ctx); + Ok(()) +} diff --git a/vortex-array/src/scalar_fn/unstable/row/visitor/execute.rs b/vortex-array/src/scalar_fn/unstable/row/visitor/execute.rs index b5539d5b29e..00fa4313df6 100644 --- a/vortex-array/src/scalar_fn/unstable/row/visitor/execute.rs +++ b/vortex-array/src/scalar_fn/unstable/row/visitor/execute.rs @@ -44,6 +44,7 @@ use crate::scalar_fn::unstable::row::execute::execute_owned_valid_rows; use crate::scalar_fn::unstable::row::execute::execute_sink; use crate::scalar_fn::unstable::row::execute::execute_sink_filtered; use crate::scalar_fn::unstable::row::execute::execute_sink_valid_rows; +use crate::validity::Validity; /// The runtime visit that decodes every column once and runs the selected row loop. pub(crate) struct ExecuteRows<'args, 'ctx, F: RowFn> { @@ -59,6 +60,9 @@ pub(crate) struct ExecuteRows<'args, 'ctx, F: RowFn> { /// The output dtype declared by [`RowVisitor::with_output_dtype`], if any. output_dtype: Option, + /// A batch validity that a collector able to attach it takes for its output. + validity: &'ctx mut Option, + /// The execution context used to decode the input columns. ctx: &'ctx mut ExecutionCtx, @@ -71,6 +75,7 @@ impl<'args, 'ctx, F: RowFn> ExecuteRows<'args, 'ctx, F> { args: &'args dyn ExecutionArgs, dtypes: &'args [DType], plan: &'args BatchPlan, + validity: &'ctx mut Option, ctx: &'ctx mut ExecutionCtx, ) -> Self { Self { @@ -78,6 +83,7 @@ impl<'args, 'ctx, F: RowFn> ExecuteRows<'args, 'ctx, F> { dtypes, plan, output_dtype: None, + validity, ctx, function: PhantomData, } @@ -129,7 +135,12 @@ impl RowVisitor for ExecuteRows<'_, '_, F> { )?; self.plan.ensure_reproduced_by(&visited)?; - execute_owned_infallible_bool::(self.args, self.ctx, apply) + execute_owned_infallible_bool::( + self.args, + self.ctx, + self.validity.take().unwrap_or(Validity::NonNullable), + apply, + ) } fn visit_prepared_into( @@ -205,6 +216,7 @@ impl RowVisitor for ExecuteRows<'_, '_, F> { execute_owned_bool::( self.args, self.ctx, + self.validity.take().unwrap_or(Validity::NonNullable), prepare, apply, finish_failure, diff --git a/vortex-array/src/scalar_fn/unstable/row/vtable.rs b/vortex-array/src/scalar_fn/unstable/row/vtable.rs index 41950737f49..ff6ea65b249 100644 --- a/vortex-array/src/scalar_fn/unstable/row/vtable.rs +++ b/vortex-array/src/scalar_fn/unstable/row/vtable.rs @@ -34,6 +34,7 @@ use crate::scalar_fn::ScalarFnId; use crate::scalar_fn::ScalarFnVTable; use crate::scalar_fn::union_child_validities; use crate::scalar_fn::unstable::row::execute::DenseAttempt; +use crate::validity::Validity; impl ScalarFnVTable for F { type Options = F::Options; @@ -125,7 +126,7 @@ pub fn execute_rows( let batch = prepare_batch(function, options, args)?; batch.execute( - |args, ctx| execute_row_kernel(function, options, args, ctx), + |args, validity, ctx| execute_row_kernel(function, options, args, validity, ctx), |args, ctx| execute_dense_attempt(function, options, args, ctx), |args, valid, ctx| try_execute_valid_rows(function, options, args, valid, ctx), |args, valid, ctx| execute_filtered_rows(function, options, args, valid, ctx), @@ -146,7 +147,8 @@ fn execute_nullary_rows( // A nullary function has no input validity to propagate, so its kernel output is the finished // column and this path labels it directly. - let values = plan.relabel_output(execute_row_kernel(function, options, args, ctx)?)?; + let values = + plan.relabel_output(execute_row_kernel(function, options, args, &mut None, ctx)?)?; finalize_kernel_output(RowFn::id(function), &result_dtype, row_count, values, ctx) } @@ -163,16 +165,21 @@ fn ensure_arity(function: &F, actual: usize) -> VortexResult<()> { Ok(()) } +/// Execute the row kernel over `args`. +/// +/// A collector that can attach the validity in `validity` takes it and builds its output with it, +/// leaving `None` behind. Otherwise the output is non-nullable and `validity` is left unchanged. fn execute_row_kernel( function: &F, options: &F::Options, args: BorrowedRowFnArgs<'_>, + validity: &mut Option, ctx: &mut ExecutionCtx, ) -> VortexResult { function.dispatch( options, args.dtypes(), - ExecuteRows::::new(&args, args.dtypes(), args.plan(), ctx), + ExecuteRows::::new(&args, args.dtypes(), args.plan(), validity, ctx), ) } @@ -244,6 +251,7 @@ mod tests { use crate::IntoArray; use crate::VortexSessionExecute; use crate::array_session; + use crate::arrays::BoolArray; use crate::arrays::PrimitiveArray; use crate::assert_arrays_eq; use crate::dtype::DType; @@ -429,6 +437,46 @@ mod tests { Ok(()) } + #[derive(Clone)] + struct GreaterThanTwo; + + impl RowFn for GreaterThanTwo { + type Options = EmptyOptions; + + const ARG_NAMES: &'static [&'static str] = &["value"]; + const INFALLIBLE: bool = true; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("test.greater_than_two"); + *ID + } + + fn dispatch( + &self, + _options: &Self::Options, + _args: &[DType], + visitor: V, + ) -> VortexResult { + visitor.visit_bool::<(i64,), false>(|(value,)| value > 2) + } + } + + #[test] + fn test_partially_valid_packed_bool_output_is_canonical() -> VortexResult<()> { + let validity = Validity::from_iter([true, false, true, true]); + let input = PrimitiveArray::new(vec![1_i64, 3, 3, 0], validity.clone()).into_array(); + let args = VecExecutionArgs::new(vec![input], 4); + let mut ctx = array_session().create_execution_ctx(); + + let actual = execute_rows(&GreaterThanTwo, &EmptyOptions, &args, &mut ctx)?; + let expected = + BoolArray::new([false, true, true, false].into_iter().collect(), validity).into_array(); + + assert!(actual.is_canonical(), "{actual}"); + assert_arrays_eq!(&actual, &expected, &mut ctx); + Ok(()) + } + #[track_caller] fn assert_arity_error(error: VortexError) { assert!(