Skip to content

Commit e6eeaae

Browse files
feat(vortex-array): union_child_validities expr helper (#8829)
## What Adds `vortex_array::expr::union_child_validities(expr)`: the conjunction of an expression's child validities (result is null iff any operand is null). This is the common `ScalarFnVTable::validity()` for null-propagating kernels (comparisons, arithmetic, most geo/tensor ops), letting the planner derive the output null mask without executing the kernel. Migrates `vortex-tensor`'s four scalar functions (`InnerProduct`, `CosineSimilarity`, `L2Denorm`, `L2Norm`) onto it, replacing their inline `and(child validities)`. Signed-off-by: Nemo Yu <zyu379@wisc.edu>
1 parent ef056d5 commit e6eeaae

5 files changed

Lines changed: 25 additions & 16 deletions

File tree

vortex-array/src/expr/exprs.rs

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@ use std::sync::Arc;
77
use std::sync::LazyLock;
88

99
use vortex_error::VortexExpect;
10+
use vortex_error::VortexResult;
1011
use vortex_error::vortex_panic;
1112
use vortex_utils::iter::ReduceBalancedIterExt;
1213

@@ -434,6 +435,22 @@ where
434435
iter.into_iter().reduce_balanced(and)
435436
}
436437

438+
/// The conjunction of an expression's child validities — i.e. the validity of a scalar function
439+
/// whose result is null exactly when any operand is null.
440+
///
441+
/// This is the `ScalarFnVTable::validity` for kernels that propagate nulls and never produce a
442+
/// null from non-null inputs (comparisons, arithmetic, most geo and tensor ops). Returning it lets
443+
/// the planner derive the output's null mask without executing the kernel. Yields `None` when the
444+
/// expression has no children.
445+
pub fn union_child_validities(expression: &Expression) -> VortexResult<Option<Expression>> {
446+
let child_validities = expression
447+
.children()
448+
.iter()
449+
.map(Expression::validity)
450+
.collect::<VortexResult<Vec<_>>>()?;
451+
Ok(and_collect(child_validities))
452+
}
453+
437454
/// Create a new [`Binary`] using the [`Add`](Operator::Add) operator.
438455
///
439456
/// ## Example usage

vortex-tensor/src/scalar_fns/cosine_similarity.rs

Lines changed: 2 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,7 @@ use vortex_array::arrays::scalar_fn::plugin::ScalarFnArrayVTable;
1515
use vortex_array::dtype::DType;
1616
use vortex_array::dtype::Nullability;
1717
use vortex_array::expr::Expression;
18-
use vortex_array::expr::and;
18+
use vortex_array::expr::union_child_validities;
1919
use vortex_array::match_each_float_ptype;
2020
use vortex_array::scalar_fn::Arity;
2121
use vortex_array::scalar_fn::ChildName;
@@ -180,10 +180,7 @@ impl ScalarFnVTable for CosineSimilarity {
180180
expression: &Expression,
181181
) -> VortexResult<Option<Expression>> {
182182
// The result is null if either input tensor is null.
183-
let lhs_validity = expression.child(0).validity()?;
184-
let rhs_validity = expression.child(1).validity()?;
185-
186-
Ok(Some(and(lhs_validity, rhs_validity)))
183+
union_child_validities(expression)
187184
}
188185

189186
fn is_null_sensitive(&self, _options: &Self::Options) -> bool {

vortex-tensor/src/scalar_fns/inner_product.rs

Lines changed: 2 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,7 @@ use vortex_array::dtype::DType;
1818
use vortex_array::dtype::NativePType;
1919
use vortex_array::dtype::Nullability;
2020
use vortex_array::expr::Expression;
21-
use vortex_array::expr::and;
21+
use vortex_array::expr::union_child_validities;
2222
use vortex_array::match_each_float_ptype;
2323
use vortex_array::scalar_fn::Arity;
2424
use vortex_array::scalar_fn::ChildName;
@@ -164,10 +164,7 @@ impl ScalarFnVTable for InnerProduct {
164164
expression: &Expression,
165165
) -> VortexResult<Option<Expression>> {
166166
// The result is null if either input tensor is null.
167-
let lhs_validity = expression.child(0).validity()?;
168-
let rhs_validity = expression.child(1).validity()?;
169-
170-
Ok(Some(and(lhs_validity, rhs_validity)))
167+
union_child_validities(expression)
171168
}
172169

173170
fn is_null_sensitive(&self, _options: &Self::Options) -> bool {

vortex-tensor/src/scalar_fns/l2_denorm.rs

Lines changed: 2 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -31,7 +31,7 @@ use vortex_array::dtype::NativePType;
3131
use vortex_array::dtype::Nullability;
3232
use vortex_array::dtype::proto::dtype as pb;
3333
use vortex_array::expr::Expression;
34-
use vortex_array::expr::and;
34+
use vortex_array::expr::union_child_validities;
3535
use vortex_array::match_each_float_ptype;
3636
use vortex_array::scalar::Scalar;
3737
use vortex_array::scalar::ScalarValue;
@@ -256,10 +256,7 @@ impl ScalarFnVTable for L2Denorm {
256256
_options: &Self::Options,
257257
expression: &Expression,
258258
) -> VortexResult<Option<Expression>> {
259-
let normalized_validity = expression.child(0).validity()?;
260-
let norms_validity = expression.child(1).validity()?;
261-
262-
Ok(Some(and(normalized_validity, norms_validity)))
259+
union_child_validities(expression)
263260
}
264261

265262
fn is_null_sensitive(&self, _options: &Self::Options) -> bool {

vortex-tensor/src/scalar_fns/l2_norm.rs

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,7 @@ use vortex_array::dtype::NativePType;
2525
use vortex_array::dtype::Nullability;
2626
use vortex_array::dtype::proto::dtype as pb;
2727
use vortex_array::expr::Expression;
28+
use vortex_array::expr::union_child_validities;
2829
use vortex_array::match_each_float_ptype;
2930
use vortex_array::scalar::Scalar;
3031
use vortex_array::scalar_fn::Arity;
@@ -184,7 +185,7 @@ impl ScalarFnVTable for L2Norm {
184185
expression: &Expression,
185186
) -> VortexResult<Option<Expression>> {
186187
// The result is null if the input tensor is null.
187-
Ok(Some(expression.child(0).validity()?))
188+
union_child_validities(expression)
188189
}
189190

190191
fn is_null_sensitive(&self, _options: &Self::Options) -> bool {

0 commit comments

Comments
 (0)