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..9e0d062eb72 100644 --- a/vortex-array/src/scalar_fn/fns/binary/mod.rs +++ b/vortex-array/src/scalar_fn/fns/binary/mod.rs @@ -241,19 +241,32 @@ 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`). + // 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, 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)