Fold and/or constants on Arrays during .optimize() (#9997)

We have two ways of optimizing and/or: we have an optimizer of
expressions,
including simplify_untyped, and we have a boolean kernel which folds
constants
with regards to Kleene logic. However, sometimes we can fold constants
without
executing the array. This change implements constant reduction for
and/or with
an array reduce rule which applies when you call array's optimize():

and(x, const), and(const, x) folds to x or ConstantArray(constant), same
for
or().
and(const, const), or(const, const) folds to a ConstantArray as well.

Kleene NULLs are not handled. This work is delegated to boolean kernel
to keep
existing code clean.

As a result, .optimize() on array, and, consequently, .reduce() can
change
reduced array's dtype to a constant and dtype's nullability from
nullable to
non-nullable.

This also brings in an invariant arguments for and/or, if constant, must
be
booleans. I've fixed the test where and(..., lit(true)) was optimized,
which
isn't a valid expression anyway.

---------

Signed-off-by: Mikhail Kot <mikhail@spiraldb.com>

Author: myrrc

vortex-array/src/array/mod.rs

@@ -427,11 +427,16 @@ impl<V: VTable> DynArrayData for ArrayData<V> {
             this.encoding_id(),
             reduced.encoding_id()
         );
+        // Rules may reduce to a non-nullable constant
         vortex_ensure!(
-            reduced.dtype() == this.dtype(),
-            "Reduced array dtype mismatch from {} to {}",
+            reduced.dtype() == this.dtype()
+                || (reduced.dtype().eq_ignore_nullability(this.dtype())
+                    && !reduced.dtype().is_nullable()),
+            "Reduced array dtype mismatch from {} ({}) to {} ({})",
             this.encoding_id(),
-            reduced.encoding_id()
+            this.dtype(),
+            reduced.encoding_id(),
+            reduced.dtype()
         );
         Ok(Some(reduced))
     }

vortex-array/src/arrays/dict/compute/rules.rs

@@ -369,8 +369,8 @@ mod tests {
         let expr = Binary::try_new(dict, const_false, Operator::And)?.into_array();
 
         let mut ctx = array_session().create_execution_ctx();
-        // Kleene AND: null AND false == false, so all three rows must be valid `false`.
-        let expected = expr.clone().execute::<BoolArray>(&mut ctx)?.into_array();
+        // Kleene AND: null AND false == false
+        let expected = BoolArray::from_iter([false, false, false]).into_array();
         let optimized = expr.optimize()?;
         assert_arrays_eq!(optimized, expected, &mut ctx);
         Ok(())
@@ -386,8 +386,8 @@ mod tests {
         let expr = Binary::try_new(dict, const_true, Operator::Or)?.into_array();
 
         let mut ctx = array_session().create_execution_ctx();
-        // Kleene OR: null OR true == true, so all three rows must be valid `true`.
-        let expected = expr.clone().execute::<BoolArray>(&mut ctx)?.into_array();
+        // Kleene OR: null OR true == true
+        let expected = BoolArray::from_iter([true, true, true]).into_array();
         let optimized = expr.optimize()?;
         assert_arrays_eq!(optimized, expected, &mut ctx);
         Ok(())

vortex-array/src/arrays/scalar_fn/rules.rs

@@ -29,8 +29,10 @@ use crate::validity::Validity;
 
 pub(super) const RULES: ReduceRuleSet<ScalarFn> = ReduceRuleSet::new(&[
     &ScalarFnPackToStructRule,
-    &ScalarFnAbstractReduceRule,
     &IsNullReduceRule,
+    // Ordering is important. ScalarFn::reduce() must be called after all other
+    // optimizations
+    &ScalarFnAbstractReduceRule,
 ]);
 
 pub(super) const PARENT_RULES: ParentRuleSet<ScalarFn> = ParentRuleSet::new(&[

vortex-array/src/expr/transform/bound_partition.rs

@@ -312,7 +312,7 @@ mod tests {
     use crate::dtype::PType::I32;
     use crate::dtype::StructFields;
     use crate::expr::analysis::make_bound_free_field_annotator;
-    use crate::expr::and;
+    use crate::expr::checked_add;
     use crate::expr::col;
     use crate::expr::get_item;
     use crate::expr::lit;
@@ -409,7 +409,7 @@ mod tests {
 
     #[rstest]
     fn test_expr_top_level_ref_get_item_add(dtype: DType) {
-        let expr = and(get_item("y", get_item("a", root())), lit(1));
+        let expr = checked_add(get_item("y", get_item("a", root())), lit(1));
         let partitioned = partition_by_field(expr.bind(&dtype).unwrap(), &dtype).unwrap();
 
         // Whole expr is a single split
@@ -418,7 +418,7 @@ mod tests {
 
     #[rstest]
     fn test_expr_top_level_ref_get_item_add_cannot_split(dtype: DType) {
-        let expr = and(get_item("y", get_item("a", root())), get_item("b", root()));
+        let expr = checked_add(get_item("y", get_item("a", root())), get_item("b", root()));
         let partitioned = partition_by_field(expr.bind(&dtype).unwrap(), &dtype).unwrap();
 
         // One for id.a and id.b

vortex-array/src/expr/transform/partition.rs

@@ -311,7 +311,7 @@ mod tests {
     fn test_expr_top_level_ref_get_item_add(dtype: DType) {
         let fields = dtype.as_struct_fields_opt().unwrap();
 
-        let expr = and(get_item("y", get_item("a", root())), lit(1));
+        let expr = and(get_item("y", get_item("a", root())), lit(true));
         let partitioned = partition(expr, &dtype, make_free_field_annotator(fields)).unwrap();
 
         // Whole expr is a single split

vortex-array/src/scalar_fn/fns/binary/mod.rs

@@ -27,6 +27,7 @@ use crate::proto::expr as pb;
 use crate::scalar_fn::Arity;
 use crate::scalar_fn::ChildName;
 use crate::scalar_fn::ExecutionArgs;
+use crate::scalar_fn::ReduceNode;
 use crate::scalar_fn::ScalarFnId;
 use crate::scalar_fn::ScalarFnVTable;
 use crate::scalar_fn::ScalarFnVTableExt;
@@ -68,6 +69,22 @@ impl Binary {
     }
 }
 
+/// Kleene and/or lookup table where both arguments are constant and non-NULL.
+/// is_and, left, right
+const KLEENE_LUT: [[[bool; 2]; 2]; 2] = [
+    [[false, true], [true, true]],   // or
+    [[false, false], [false, true]], // and
+];
+
+/// Kleene and/or reduction where one argument is constant non-NULL and other
+/// is non-constant.
+fn kleene_one_const<T: ReduceNode>(node: T, constant: bool, is_and: bool) -> T {
+    match (is_and, constant) {
+        (true, true) | (false, false) => node,
+        (is_and, constant) => node.new_constant((constant && !is_and).into()),
+    }
+}
+
 impl ScalarFnVTable for Binary {
     type Options = Operator;
 
@@ -152,6 +169,12 @@ impl ScalarFnVTable for Binary {
             vortex_bail!("Cannot compare different DTypes {} and {}", lhs, rhs);
         }
 
+        if matches!(operator, Operator::And | Operator::Or)
+            && !(lhs.is_boolean() && rhs.is_boolean())
+        {
+            vortex_bail!("'{operator}' requires bool operands, got {lhs} and {rhs}",);
+        }
+
         Ok(DType::Bool((lhs.is_nullable() || rhs.is_nullable()).into()))
     }
 
@@ -254,18 +277,43 @@ impl ScalarFnVTable for Binary {
         operator: &Operator,
         expression: &Expression,
     ) -> VortexResult<Option<Expression>> {
+        if matches!(operator, Operator::And | Operator::Or) {
+            return Ok(None); // AND and OR are kleene logic
+        }
         let lhs = expression.child(0).validity()?;
         let rhs = expression.child(1).validity()?;
+        Ok(Some(and(lhs, rhs)))
+    }
 
-        Ok(match operator {
-            // AND and OR are kleene logic.
-            Operator::And => None,
-            Operator::Or => None,
-            _ => {
-                // All other binary operators are null if either side is null.
-                Some(and(lhs, rhs))
-            }
-        })
+    fn reduce<T: ReduceNode>(&self, operator: &Operator, node: &T) -> VortexResult<Option<T>> {
+        if !matches!(operator, Operator::And | Operator::Or) {
+            return Ok(None);
+        }
+        let left = node.child(0);
+        let right = node.child(1);
+
+        let left_const = left.as_constant();
+        let right_const = right.as_constant();
+
+        // We don't handle Kleene NULL reduction here (.value() returns None
+        // for NULL). This will be reduced in the boolean kernel during
+        // execution. Not handling the case keeps the code much simpler.
+        let left_const = left_const
+            .and_then(|s| s.value().cloned())
+            .map(|v| v.as_bool());
+        let right_const = right_const
+            .and_then(|s| s.value().cloned())
+            .map(|v| v.as_bool());
+        let is_and = *operator == Operator::And;
+
+        Ok(Some(match (left_const, right_const) {
+            (None, None) => return Ok(None),
+            (Some(left_const), Some(right_const)) => left.new_constant(
+                KLEENE_LUT[is_and as usize][left_const as usize][right_const as usize].into(),
+            ),
+            (Some(constant), None) => kleene_one_const(right, constant, is_and),
+            (None, Some(constant)) => kleene_one_const(left, constant, is_and),
+        }))
     }
 
     fn is_strict(&self, operator: &Operator) -> bool {
@@ -296,8 +344,13 @@ mod tests {
     use vortex_error::VortexResult;
 
     use super::*;
+    use crate::IntoArray;
     use crate::VortexSessionExecute;
     use crate::array_session;
+    use crate::arrays::Bool;
+    use crate::arrays::BoolArray;
+    use crate::arrays::ConstantArray;
+    use crate::arrays::PrimitiveArray;
     use crate::assert_arrays_eq;
     use crate::builtins::ArrayBuiltins;
     use crate::dtype::DType;
@@ -316,7 +369,9 @@ mod tests {
     use crate::expr::or;
     use crate::expr::or_collect;
     use crate::expr::test_harness;
+    use crate::optimizer::ArrayOptimizer;
     use crate::scalar::Scalar;
+
     #[test]
     fn and_collect_balanced() {
         let values = vec![lit(1), lit(2), lit(3), lit(4), lit(5)];
@@ -472,40 +527,22 @@ mod tests {
 
         // Create a struct array with one element for testing.
         let lhs_struct = StructArray::from_fields(&[
-            (
-                "a",
-                crate::arrays::PrimitiveArray::from_iter([1i32]).into_array(),
-            ),
-            (
-                "b",
-                crate::arrays::PrimitiveArray::from_iter([3i32]).into_array(),
-            ),
+            ("a", PrimitiveArray::from_iter([1i32]).into_array()),
+            ("b", PrimitiveArray::from_iter([3i32]).into_array()),
         ])
         .unwrap()
         .into_array();
 
         let rhs_struct_equal = StructArray::from_fields(&[
-            (
-                "a",
-                crate::arrays::PrimitiveArray::from_iter([1i32]).into_array(),
-            ),
-            (
-                "b",
-                crate::arrays::PrimitiveArray::from_iter([3i32]).into_array(),
-            ),
+            ("a", PrimitiveArray::from_iter([1i32]).into_array()),
+            ("b", PrimitiveArray::from_iter([3i32]).into_array()),
         ])
         .unwrap()
         .into_array();
 
         let rhs_struct_different = StructArray::from_fields(&[
-            (
-                "a",
-                crate::arrays::PrimitiveArray::from_iter([1i32]).into_array(),
-            ),
-            (
-                "b",
-                crate::arrays::PrimitiveArray::from_iter([4i32]).into_array(),
-            ),
+            ("a", PrimitiveArray::from_iter([1i32]).into_array()),
+            ("b", PrimitiveArray::from_iter([4i32]).into_array()),
         ])
         .unwrap()
         .into_array();
@@ -535,8 +572,6 @@ mod tests {
     #[test]
     fn test_or_kleene_validity() {
         let mut ctx = array_session().create_execution_ctx();
-        use crate::IntoArray;
-        use crate::arrays::BoolArray;
         use crate::arrays::StructArray;
         use crate::expr::col;
 
@@ -630,14 +665,79 @@ mod tests {
     fn test_scalar_subtract_float_underflow_is_ok() {
         use vortex_buffer::buffer;
 
-        use crate::IntoArray;
-        use crate::arrays::ConstantArray;
-
         let values = buffer![f32::MIN, 2.0, 3.0].into_array();
         let rhs1 = ConstantArray::new(Scalar::from(1.0f32), 3).into_array();
         let _results = values.binary(rhs1, Operator::Sub).unwrap();
         let values = buffer![f32::MIN, 2.0, 3.0].into_array();
         let rhs2 = ConstantArray::new(Scalar::from(f32::MAX), 3).into_array();
         let _results = values.binary(rhs2, Operator::Sub).unwrap();
     }
+
+    #[test]
+    fn test_and_reduce() -> VortexResult<()> {
+        let mut ctx = array_session().create_execution_ctx();
+        let left = BoolArray::from_iter([true, true, false]).into_array();
+        let right = ConstantArray::new(true, 3).into_array();
+
+        let array = Binary::try_new(left.clone(), right.clone(), Operator::And)?
+            .into_array()
+            .optimize()?; // calls reduce()
+        assert!(array.is::<Bool>()); // and(left, const) -> left
+        assert_arrays_eq!(array, left, &mut ctx);
+
+        let array = Binary::try_new(right.clone(), left.clone(), Operator::And)?
+            .into_array()
+            .optimize()?;
+        assert!(array.is::<Bool>()); // and(const, left) -> left
+        assert_arrays_eq!(array, left, &mut ctx);
+
+        let array = Binary::try_new(left.clone(), right.clone(), Operator::Or)?
+            .into_array()
+            .optimize()?;
+        assert_eq!(array.as_constant(), Some(true.into())); // or(left, const) -> true
+
+        let array = Binary::try_new(right.clone(), left, Operator::Or)?
+            .into_array()
+            .optimize()?;
+        assert_eq!(array.as_constant(), Some(true.into())); // or(const, left) -> true
+
+        let array = Binary::try_new(right.clone(), right.clone(), Operator::Or)?
+            .into_array()
+            .optimize()?;
+        assert_eq!(array.as_constant(), Some(true.into())); // or(const, const) -> const
+
+        let left_false = ConstantArray::new(false, 3).into_array();
+
+        let array = Binary::try_new(right.clone(), left_false.clone(), Operator::Or)?
+            .into_array()
+            .optimize()?;
+        assert_eq!(array.as_constant(), Some(true.into()));
+        let array = Binary::try_new(left_false, right, Operator::And)?
+            .into_array()
+            .optimize()?;
+        assert_eq!(array.as_constant(), Some(false.into()));
+
+        Ok(())
+    }
+
+    #[test]
+    fn test_and_reduce_nullable() -> VortexResult<()> {
+        let mut ctx = array_session().create_execution_ctx();
+        let nullable_true =
+            ConstantArray::new(Scalar::bool(true, Nullability::Nullable), 3).into_array();
+        let right = BoolArray::from_iter([true, false, true]).into_array();
+
+        let array = Binary::try_new(nullable_true, right.clone(), Operator::And)?
+            .into_array()
+            .optimize()?;
+        assert_arrays_eq!(array, right, &mut ctx);
+        Ok(())
+    }
+
+    #[test]
+    fn test_and_reject_non_bool() {
+        let lhs = ConstantArray::new(7i32, 3).into_array();
+        let rhs = PrimitiveArray::from_iter([1i32, 2, 3]).into_array();
+        assert!(Binary::try_new(lhs, rhs, Operator::And).is_err());
+    }
 }