From f3a3ece1afcb6c8def6b2581a34f3c1cf3c99b64 Mon Sep 17 00:00:00 2001 From: jackylee-ch Date: Sun, 4 Oct 2026 11:03:26 +0800 Subject: [PATCH 1/2] fix(array): preserve result nullability when simplifying and/or/zip Literal simplification could turn a well-typed nullable expression into a non-nullable one: the and/or annihilator and identity rules returned a fresh non-nullable literal or a bare child, and the zip literal-mask rule returned the selected branch, none of them retaining the nullability binding had inferred. Coerce each simplified result back to the bound expression's dtype, so the optimized plan keeps the same result schema. The regression cases come from the issue. Fixes #10009 Signed-off-by: jackylee-ch --- vortex-array/src/expr/optimize.rs | 33 ++++++++++++++++++++ vortex-array/src/scalar_fn/fns/binary/mod.rs | 14 +++++++-- vortex-array/src/scalar_fn/fns/zip/mod.rs | 16 +++++++--- 3 files changed, 57 insertions(+), 6 deletions(-) diff --git a/vortex-array/src/expr/optimize.rs b/vortex-array/src/expr/optimize.rs index 956fb8c0f10..89ecf9aade9 100644 --- a/vortex-array/src/expr/optimize.rs +++ b/vortex-array/src/expr/optimize.rs @@ -132,6 +132,7 @@ impl BoundExpression { #[cfg(test)] mod tests { + use rstest::rstest; use vortex_error::VortexResult; use vortex_error::vortex_err; @@ -139,6 +140,8 @@ mod tests { use crate::dtype::Nullability; use crate::dtype::PType; use crate::dtype::StructFields; + use crate::expr::Expression; + use crate::expr::and; use crate::expr::cast; use crate::expr::eq; use crate::expr::get_item; @@ -146,6 +149,7 @@ mod tests { use crate::expr::lt_eq; use crate::expr::or; use crate::expr::root; + use crate::expr::zip_expr; use crate::scalar::Scalar; use crate::scalar_fn::fns::literal::Literal; @@ -198,4 +202,33 @@ mod tests { assert_eq!(rhs, &Scalar::primitive(3.0f64, Nullability::NonNullable)); Ok(()) } + + #[rstest] + #[case::and_annihilator(and(root(), lit(false)), DType::Bool(Nullability::Nullable))] + #[case::or_annihilator(or(root(), lit(true)), DType::Bool(Nullability::Nullable))] + #[case::and_nullable_identity( + and(root(), lit(Scalar::from(Some(true)))), + DType::Bool(Nullability::NonNullable) + )] + #[case::or_nullable_identity( + or(root(), lit(Scalar::from(Some(false)))), + DType::Bool(Nullability::NonNullable) + )] + #[case::zip_true( + zip_expr(lit(true), lit(1i32), root()), + DType::Primitive(PType::I32, Nullability::Nullable) + )] + #[case::zip_false( + zip_expr(lit(false), root(), lit(1i32)), + DType::Primitive(PType::I32, Nullability::Nullable) + )] + fn literal_simplification_preserves_result_dtype( + #[case] expr: Expression, + #[case] scope: DType, + ) -> VortexResult<()> { + let original = expr.bind(&scope)?; + let optimized = expr.bind(&scope)?.optimize_recursive()?; + assert_eq!(optimized.dtype(), original.dtype()); + Ok(()) + } } diff --git a/vortex-array/src/scalar_fn/fns/binary/mod.rs b/vortex-array/src/scalar_fn/fns/binary/mod.rs index b99604544e8..9a9e13afc4f 100644 --- a/vortex-array/src/scalar_fn/fns/binary/mod.rs +++ b/vortex-array/src/scalar_fn/fns/binary/mod.rs @@ -241,7 +241,7 @@ impl ScalarFnVTable for Binary { // Other null cases either fall out of the identity/annihilator rules // above (`null AND true`, `null OR false`) or cannot be simplified under // Kleene semantics (`null AND x`, `null OR x` for non-literal `x`). - Ok(match operator { + let simplified = match operator { Operator::And => match (bool_literal(lhs), bool_literal(rhs)) { (Some(Some(false)), _) | (_, Some(Some(false))) => Some(bound::lit(false)), (Some(Some(true)), _) => Some(rhs.clone()), @@ -257,7 +257,17 @@ impl ScalarFnVTable for Binary { _ => None, }, _ => None, - }) + }; + + // The annihilator literal and the identity child can be less nullable than + // `lhs OP rhs`, so keep the result dtype that binding inferred. + Ok(simplified.map(|simplified| { + if simplified.dtype() == expr.dtype() { + simplified + } else { + bound::cast(simplified, expr.dtype().clone()) + } + })) } fn reduce(&self, operator: &Operator, node: &T) -> VortexResult> { diff --git a/vortex-array/src/scalar_fn/fns/zip/mod.rs b/vortex-array/src/scalar_fn/fns/zip/mod.rs index 9731f6be980..54fb04cfd79 100644 --- a/vortex-array/src/scalar_fn/fns/zip/mod.rs +++ b/vortex-array/src/scalar_fn/fns/zip/mod.rs @@ -28,6 +28,7 @@ use crate::builtins::ArrayBuiltins; use crate::dtype::DType; use crate::dtype::StructFields; use crate::expr::BoundExpression; +use crate::expr::bound; use crate::expr::display::ExprDisplay; use crate::scalar_fn::Arity; use crate::scalar_fn::ChildName; @@ -165,11 +166,18 @@ impl ScalarFnVTable for Zip { }; if let Some(mask_val) = mask_lit.as_bool().value() { - if mask_val { - return Ok(Some(expr.child(0).clone())); + let selected = if mask_val { + expr.child(0).clone() } else { - return Ok(Some(expr.child(1).clone())); - } + expr.child(1).clone() + }; + // The selected branch can be less nullable than the zip result, which unions + // both branches' nullability, so keep the dtype that binding inferred. + return Ok(Some(if selected.dtype() == expr.dtype() { + selected + } else { + bound::cast(selected, expr.dtype().clone()) + })); } Ok(None) From b614620d71998a59c4cbcea99df45e7086eac5b9 Mon Sep 17 00:00:00 2001 From: jackylee-ch Date: Sun, 4 Oct 2026 23:50:59 +0800 Subject: [PATCH 2/2] refactor(array): type the and/or simplify arms directly Build the annihilator literal and widen the identity child to the inferred result nullability inside each arm, instead of casting the whole result afterwards, as suggested in review. Signed-off-by: jackylee-ch --- vortex-array/src/scalar_fn/fns/binary/mod.rs | 43 +++++++++++--------- 1 file changed, 23 insertions(+), 20 deletions(-) diff --git a/vortex-array/src/scalar_fn/fns/binary/mod.rs b/vortex-array/src/scalar_fn/fns/binary/mod.rs index 9a9e13afc4f..9e0d062eb72 100644 --- a/vortex-array/src/scalar_fn/fns/binary/mod.rs +++ b/vortex-array/src/scalar_fn/fns/binary/mod.rs @@ -241,33 +241,36 @@ impl ScalarFnVTable for Binary { // Other null cases either fall out of the identity/annihilator rules // above (`null AND true`, `null OR false`) or cannot be simplified under // Kleene semantics (`null AND x`, `null OR x` for non-literal `x`). - let simplified = match operator { + // Simplification must keep the result nullability binding inferred: the annihilator + // literal and the identity child can both be less nullable than `lhs OP rhs`. + let result_nullable = lhs.dtype().is_nullable() || rhs.dtype().is_nullable(); + let nullability = Nullability::from(result_nullable); + let bool_lit = |value: bool| bound::lit(Scalar::bool(value, nullability)); + let widen = |child: &BoundExpression| { + if child.dtype().is_nullable() == result_nullable { + child.clone() + } else { + bound::cast(child.clone(), DType::Bool(nullability)) + } + }; + + Ok(match operator { Operator::And => match (bool_literal(lhs), bool_literal(rhs)) { - (Some(Some(false)), _) | (_, Some(Some(false))) => Some(bound::lit(false)), - (Some(Some(true)), _) => Some(rhs.clone()), - (_, Some(Some(true))) => Some(lhs.clone()), - (Some(None), Some(None)) => Some(lhs.clone()), + (Some(Some(false)), _) | (_, Some(Some(false))) => Some(bool_lit(false)), + (Some(Some(true)), _) => Some(widen(rhs)), + (_, Some(Some(true))) => Some(widen(lhs)), + (Some(None), Some(None)) => Some(widen(lhs)), _ => None, }, Operator::Or => match (bool_literal(lhs), bool_literal(rhs)) { - (Some(Some(true)), _) | (_, Some(Some(true))) => Some(bound::lit(true)), - (Some(Some(false)), _) => Some(rhs.clone()), - (_, Some(Some(false))) => Some(lhs.clone()), - (Some(None), Some(None)) => Some(lhs.clone()), + (Some(Some(true)), _) | (_, Some(Some(true))) => Some(bool_lit(true)), + (Some(Some(false)), _) => Some(widen(rhs)), + (_, Some(Some(false))) => Some(widen(lhs)), + (Some(None), Some(None)) => Some(widen(lhs)), _ => None, }, _ => None, - }; - - // The annihilator literal and the identity child can be less nullable than - // `lhs OP rhs`, so keep the result dtype that binding inferred. - Ok(simplified.map(|simplified| { - if simplified.dtype() == expr.dtype() { - simplified - } else { - bound::cast(simplified, expr.dtype().clone()) - } - })) + }) } fn reduce(&self, operator: &Operator, node: &T) -> VortexResult> {