From f81df622521815085997a31bb8f56847558031b4 Mon Sep 17 00:00:00 2001 From: Nemo Yu Date: Fri, 17 Jul 2026 13:35:08 -0400 Subject: [PATCH] feat(vortex-array): union_child_validities expr helper MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add expr::union_child_validities — the conjunction of an expression's child validities (result null iff any operand is null), the common ScalarFnVTable validity() for null-propagating kernels. Migrate vortex-tensor's four scalar functions (InnerProduct, CosineSimilarity, L2Denorm, L2Norm) onto it, replacing their inline and(child validities). Signed-off-by: Nemo Yu --- vortex-array/src/expr/exprs.rs | 17 +++++++++++++++++ .../src/scalar_fns/cosine_similarity.rs | 7 ++----- vortex-tensor/src/scalar_fns/inner_product.rs | 7 ++----- vortex-tensor/src/scalar_fns/l2_denorm.rs | 7 ++----- vortex-tensor/src/scalar_fns/l2_norm.rs | 3 ++- 5 files changed, 25 insertions(+), 16 deletions(-) diff --git a/vortex-array/src/expr/exprs.rs b/vortex-array/src/expr/exprs.rs index da70427a0f7..9843277ede7 100644 --- a/vortex-array/src/expr/exprs.rs +++ b/vortex-array/src/expr/exprs.rs @@ -7,6 +7,7 @@ use std::sync::Arc; use std::sync::LazyLock; use vortex_error::VortexExpect; +use vortex_error::VortexResult; use vortex_error::vortex_panic; use vortex_utils::iter::ReduceBalancedIterExt; @@ -434,6 +435,22 @@ where iter.into_iter().reduce_balanced(and) } +/// The conjunction of an expression's child validities — i.e. the validity of a scalar function +/// whose result is null exactly when any operand is null. +/// +/// This is the `ScalarFnVTable::validity` for kernels that propagate nulls and never produce a +/// null from non-null inputs (comparisons, arithmetic, most geo and tensor ops). Returning it lets +/// the planner derive the output's null mask without executing the kernel. Yields `None` when the +/// expression has no children. +pub fn union_child_validities(expression: &Expression) -> VortexResult> { + let child_validities = expression + .children() + .iter() + .map(Expression::validity) + .collect::>>()?; + Ok(and_collect(child_validities)) +} + /// Create a new [`Binary`] using the [`Add`](Operator::Add) operator. /// /// ## Example usage diff --git a/vortex-tensor/src/scalar_fns/cosine_similarity.rs b/vortex-tensor/src/scalar_fns/cosine_similarity.rs index 82328a9b2d9..ffe198eb0fa 100644 --- a/vortex-tensor/src/scalar_fns/cosine_similarity.rs +++ b/vortex-tensor/src/scalar_fns/cosine_similarity.rs @@ -15,7 +15,7 @@ use vortex_array::arrays::scalar_fn::plugin::ScalarFnArrayVTable; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::expr::Expression; -use vortex_array::expr::and; +use vortex_array::expr::union_child_validities; use vortex_array::match_each_float_ptype; use vortex_array::scalar_fn::Arity; use vortex_array::scalar_fn::ChildName; @@ -180,10 +180,7 @@ impl ScalarFnVTable for CosineSimilarity { expression: &Expression, ) -> VortexResult> { // The result is null if either input tensor is null. - let lhs_validity = expression.child(0).validity()?; - let rhs_validity = expression.child(1).validity()?; - - Ok(Some(and(lhs_validity, rhs_validity))) + union_child_validities(expression) } fn is_null_sensitive(&self, _options: &Self::Options) -> bool { diff --git a/vortex-tensor/src/scalar_fns/inner_product.rs b/vortex-tensor/src/scalar_fns/inner_product.rs index cfea956f314..cb17d179161 100644 --- a/vortex-tensor/src/scalar_fns/inner_product.rs +++ b/vortex-tensor/src/scalar_fns/inner_product.rs @@ -18,7 +18,7 @@ use vortex_array::dtype::DType; use vortex_array::dtype::NativePType; use vortex_array::dtype::Nullability; use vortex_array::expr::Expression; -use vortex_array::expr::and; +use vortex_array::expr::union_child_validities; use vortex_array::match_each_float_ptype; use vortex_array::scalar_fn::Arity; use vortex_array::scalar_fn::ChildName; @@ -164,10 +164,7 @@ impl ScalarFnVTable for InnerProduct { expression: &Expression, ) -> VortexResult> { // The result is null if either input tensor is null. - let lhs_validity = expression.child(0).validity()?; - let rhs_validity = expression.child(1).validity()?; - - Ok(Some(and(lhs_validity, rhs_validity))) + union_child_validities(expression) } fn is_null_sensitive(&self, _options: &Self::Options) -> bool { diff --git a/vortex-tensor/src/scalar_fns/l2_denorm.rs b/vortex-tensor/src/scalar_fns/l2_denorm.rs index 4f4159b4a21..a4ac9e6c91e 100644 --- a/vortex-tensor/src/scalar_fns/l2_denorm.rs +++ b/vortex-tensor/src/scalar_fns/l2_denorm.rs @@ -31,7 +31,7 @@ use vortex_array::dtype::NativePType; use vortex_array::dtype::Nullability; use vortex_array::dtype::proto::dtype as pb; use vortex_array::expr::Expression; -use vortex_array::expr::and; +use vortex_array::expr::union_child_validities; use vortex_array::match_each_float_ptype; use vortex_array::scalar::Scalar; use vortex_array::scalar::ScalarValue; @@ -256,10 +256,7 @@ impl ScalarFnVTable for L2Denorm { _options: &Self::Options, expression: &Expression, ) -> VortexResult> { - let normalized_validity = expression.child(0).validity()?; - let norms_validity = expression.child(1).validity()?; - - Ok(Some(and(normalized_validity, norms_validity))) + union_child_validities(expression) } fn is_null_sensitive(&self, _options: &Self::Options) -> bool { diff --git a/vortex-tensor/src/scalar_fns/l2_norm.rs b/vortex-tensor/src/scalar_fns/l2_norm.rs index aa85ab4c78e..0eeb6777ab3 100644 --- a/vortex-tensor/src/scalar_fns/l2_norm.rs +++ b/vortex-tensor/src/scalar_fns/l2_norm.rs @@ -25,6 +25,7 @@ use vortex_array::dtype::NativePType; use vortex_array::dtype::Nullability; use vortex_array::dtype::proto::dtype as pb; use vortex_array::expr::Expression; +use vortex_array::expr::union_child_validities; use vortex_array::match_each_float_ptype; use vortex_array::scalar::Scalar; use vortex_array::scalar_fn::Arity; @@ -184,7 +185,7 @@ impl ScalarFnVTable for L2Norm { expression: &Expression, ) -> VortexResult> { // The result is null if the input tensor is null. - Ok(Some(expression.child(0).validity()?)) + union_child_validities(expression) } fn is_null_sensitive(&self, _options: &Self::Options) -> bool {