fix(array): preserve result nullability when simplifying and/or/zip

vortex-array/src/expr/optimize.rs

@@ -132,20 +132,24 @@ impl BoundExpression {
 
 #[cfg(test)]
 mod tests {
+    use rstest::rstest;
     use vortex_error::VortexResult;
     use vortex_error::vortex_err;
 
     use crate::dtype::DType;
     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;
     use crate::expr::lit;
     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(())
+    }
 }

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,

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)