feat[vortex-expr]: is_fallible analysis (#5546)

Signed-off-by: Joe Isaacs <[email protected]>

Author: joseph-isaacs

vortex-array/src/expr/analysis/fallible.rs

@@ -0,0 +1,76 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+use crate::expr::Expression;
+use crate::expr::analysis::BooleanLabels;
+use crate::expr::label_tree;
+
+pub fn label_is_fallible(expr: &Expression) -> BooleanLabels<'_> {
+    label_tree(expr, |expr| expr.is_fallible(), |acc, &child| acc | child)
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::expr::exprs::binary::checked_add;
+    use crate::expr::exprs::binary::eq;
+    use crate::expr::exprs::get_item::col;
+    use crate::expr::exprs::is_null::is_null;
+    use crate::expr::exprs::literal::lit;
+    use crate::expr::exprs::merge::DuplicateHandling;
+    use crate::expr::exprs::merge::merge_opts;
+    use crate::expr::exprs::not::not;
+
+    #[test]
+    fn not_is_not_fallible() {
+        let expr = not(col("x"));
+        let labels = label_is_fallible(&expr);
+        assert_eq!(labels.get(&expr), Some(&false));
+    }
+
+    #[test]
+    fn checked_add_defaults_to_fallible() {
+        let expr = checked_add(col("a"), col("b"));
+        let labels = label_is_fallible(&expr);
+        assert_eq!(labels.get(&expr), Some(&true));
+    }
+
+    #[test]
+    fn eq_not_fallible() {
+        let expr = eq(col("a"), lit(5));
+        let labels = label_is_fallible(&expr);
+        assert_eq!(labels.get(&expr), Some(&false));
+    }
+
+    #[test]
+    fn merge_with_error_handling_is_fallible() {
+        let expr = merge_opts([col("a"), col("b")], DuplicateHandling::Error);
+        let labels = label_is_fallible(&expr);
+        assert_eq!(labels.get(&expr), Some(&true));
+    }
+
+    #[test]
+    fn merge_with_rightmost_handling_is_not_fallible() {
+        let expr = merge_opts([col("a"), col("b")], DuplicateHandling::RightMost);
+        let labels = label_is_fallible(&expr);
+        assert_eq!(labels.get(&expr), Some(&false));
+    }
+
+    #[test]
+    fn nested_with_fallible_child() {
+        let child = checked_add(col("a"), col("b"));
+        let expr = not(child.clone());
+        let labels = label_is_fallible(&expr);
+        assert_eq!(labels.get(&child), Some(&true));
+        assert_eq!(labels.get(&expr), Some(&true));
+    }
+
+    #[test]
+    fn nested_without_fallible_child() {
+        let child = is_null(col("x"));
+        let expr = not(child.clone());
+        let labels = label_is_fallible(&expr);
+        assert_eq!(labels.get(&child), Some(&false));
+        assert_eq!(labels.get(&expr), Some(&false));
+    }
+}

vortex-array/src/expr/analysis/mod.rs

@@ -2,11 +2,14 @@
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
 pub mod annotation;
+mod fallible;
 pub mod immediate_access;
 mod labeling;
 mod null_sensitive;
 
 pub use annotation::*;
+pub use fallible::label_is_fallible;
 pub use immediate_access::*;
 pub use labeling::*;
-pub use null_sensitive::*;
+pub use null_sensitive::BooleanLabels;
+pub use null_sensitive::label_null_sensitive;

vortex-array/src/expr/analysis/null_sensitive.rs

@@ -6,13 +6,13 @@ use vortex_utils::aliases::hash_map::HashMap;
 use super::labeling::label_tree;
 use crate::expr::Expression;
 
-pub type NullSensitiveLabels<'a> = HashMap<&'a Expression, bool>;
+pub type BooleanLabels<'a> = HashMap<&'a Expression, bool>;
 
 /// Label each expression in the tree with whether it is null-sensitive.
 ///
 /// See [`crate::expr::VTable::is_null_sensitive`] for a definition of null sensitivity.
 /// This function operates on a tree of expressions, not just a single expression.
-pub fn label_null_sensitive(expr: &Expression) -> NullSensitiveLabels<'_> {
+pub fn label_null_sensitive(expr: &Expression) -> BooleanLabels<'_> {
     label_tree(
         expr,
         |expr| expr.is_null_sensitive(),

vortex-array/src/expr/exprs/binary.rs

@@ -256,6 +256,24 @@ impl VTable for Binary {
     fn is_null_sensitive(&self, _instance: &Self::Instance) -> bool {
         false
     }
+
+    fn is_fallible(&self, instance: &Self::Instance) -> bool {
+        // Opt-in not out for fallibility.
+        // Arithmetic operations could be better modelled here.
+        let infallible = matches!(
+            instance,
+            Operator::Eq
+                | Operator::NotEq
+                | Operator::Gt
+                | Operator::Gte
+                | Operator::Lt
+                | Operator::Lte
+                | Operator::And
+                | Operator::Or
+        );
+
+        !infallible
+    }
 }
 
 impl ExpressionView<'_, Binary> {