feat[vortex-expr]: add a is_null_sensitive expression analysis (#5528)

Signed-off-by: Joe Isaacs <joe.isaacs@live.co.uk>

Author: joseph-isaacs

vortex-array/src/arrays/struct_/vtable/reduce.rs

@@ -16,12 +16,12 @@ use crate::IntoArray;
 use crate::arrays::ExprArray;
 use crate::arrays::StructArray;
 use crate::expr::Expression;
+use crate::expr::analysis::immediate_access::annotate_scope_access;
 use crate::expr::col;
 use crate::expr::root;
 use crate::expr::session::ExprSession;
 use crate::expr::transform::ExprOptimizer;
 use crate::expr::transform::PartitionedExpr;
-use crate::expr::transform::immediate_access::annotate_scope_access;
 use crate::expr::transform::partition;
 use crate::expr::transform::replace;
 use crate::expr::transform::replace_root_fields;

vortex-array/src/expr/analysis.rs

@@ -7,12 +7,12 @@ use vortex_error::VortexExpect;
 use vortex_utils::aliases::hash_set::HashSet;
 
 use crate::expr::Expression;
+use crate::expr::analysis::AnnotationFn;
+use crate::expr::analysis::Annotations;
+use crate::expr::descendent_annotations;
 use crate::expr::exprs::get_item::GetItem;
 use crate::expr::exprs::root::Root;
 use crate::expr::exprs::select::Select;
-use crate::expr::transform::annotations::AnnotationFn;
-use crate::expr::transform::annotations::Annotations;
-use crate::expr::transform::annotations::descendent_annotations;
 
 pub type FieldAccesses<'a> = Annotations<'a, FieldName>;
 

vortex-array/src/expr/transform/annotations.rs renamed to vortex-array/src/expr/analysis/annotation.rs

@@ -0,0 +1,120 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+use vortex_error::VortexExpect;
+use vortex_error::VortexResult;
+use vortex_utils::aliases::hash_map::HashMap;
+
+use crate::expr::Expression;
+use crate::expr::traversal::NodeExt;
+use crate::expr::traversal::NodeVisitor;
+use crate::expr::traversal::TraversalOrder;
+
+/// Label each node in an expression tree using a bottom-up traversal.
+///
+/// This function separates tree labeling into two distinct steps:
+/// 1. **Label Self**: Compute a label for each node based only on the node itself
+/// 2. **Merge Child**: Fold/accumulate labels from children into the node's self-label
+///
+/// The labeling process:
+/// - First, `self_label` is called on the node to produce its self-label
+/// - Then, for each child, `merge_child` is called with `(self_label, child_label)`
+///   to fold the child label into the self_label
+/// - This produces the final label for the node
+///
+/// # Parameters
+///
+/// - `expr`: The root expression to label
+/// - `self_label`: Function that computes a label for a single node
+/// - `merge_child`: Mutable function that folds child labels into an accumulator.
+///   Takes `(self_label, child_label)` and returns the updated accumulator.
+///   Called once per child, with the initial accumulator being the node's self-label.
+///
+pub fn label_tree<L: Clone>(
+    expr: &Expression,
+    self_label: impl Fn(&Expression) -> L,
+    mut merge_child: impl FnMut(L, &L) -> L,
+) -> HashMap<&Expression, L> {
+    let mut visitor = LabelingVisitor {
+        labels: Default::default(),
+        self_label,
+        merge_child: &mut merge_child,
+    };
+    expr.accept(&mut visitor)
+        .vortex_expect("LabelingVisitor is infallible");
+    visitor.labels
+}
+
+struct LabelingVisitor<'a, 'b, L, F, G>
+where
+    F: Fn(&Expression) -> L,
+    G: FnMut(L, &L) -> L,
+{
+    labels: HashMap<&'a Expression, L>,
+    self_label: F,
+    merge_child: &'b mut G,
+}
+
+impl<'a, 'b, L: Clone, F, G> NodeVisitor<'a> for LabelingVisitor<'a, 'b, L, F, G>
+where
+    F: Fn(&Expression) -> L,
+    G: FnMut(L, &L) -> L,
+{
+    type NodeTy = Expression;
+
+    fn visit_up(&mut self, node: &'a Expression) -> VortexResult<TraversalOrder> {
+        let self_label = (self.self_label)(node);
+
+        let final_label = node.children().iter().fold(self_label, |acc, child| {
+            let child_label = self
+                .labels
+                .get(child)
+                .vortex_expect("child must have label");
+            (self.merge_child)(acc, child_label)
+        });
+
+        self.labels.insert(node, final_label);
+
+        Ok(TraversalOrder::Continue)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::expr::exprs::binary::eq;
+    use crate::expr::exprs::get_item::col;
+    use crate::expr::exprs::literal::lit;
+
+    #[test]
+    fn test_tree_depth() {
+        // Expression: $.col1 = 5
+        // Tree: eq(get_item(root(), "col1"), lit(5))
+        // Depth: root = 1, get_item = 2, lit = 1, eq = 3
+        let expr = eq(col("col1"), lit(5));
+        let depths = label_tree(
+            &expr,
+            |_node| 1, // Each node has depth 1 by itself
+            |self_depth, child_depth| self_depth.max(*child_depth + 1),
+        );
+
+        // The root (eq) should have depth 3
+        assert_eq!(depths.get(&expr), Some(&3));
+    }
+
+    #[test]
+    fn test_node_count() {
+        // Count total nodes in subtree (including self)
+        // Tree: eq(get_item(root(), "col1"), lit(5))
+        // Nodes: eq, get_item, root, lit = 4
+        let expr = eq(col("col1"), lit(5));
+        let counts = label_tree(
+            &expr,
+            |_node| 1, // Each node counts as 1
+            |self_count, child_count| self_count + *child_count,
+        );
+
+        // Root should have count of 4 (eq, get_item, root, lit)
+        assert_eq!(counts.get(&expr), Some(&4));
+    }
+}

vortex-array/src/expr/transform/immediate_access.rs renamed to vortex-array/src/expr/analysis/immediate_access.rs

@@ -0,0 +1,12 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+pub mod annotation;
+pub mod immediate_access;
+mod labeling;
+mod null_sensitive;
+
+pub use annotation::*;
+pub use immediate_access::*;
+pub use labeling::*;
+pub use null_sensitive::*;

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

@@ -0,0 +1,62 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+use vortex_utils::aliases::hash_map::HashMap;
+
+use super::labeling::label_tree;
+use crate::expr::Expression;
+
+pub type NullSensitiveLabels<'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<'_> {
+    label_tree(
+        expr,
+        |expr| expr.is_null_sensitive(),
+        |acc, &child| acc | child,
+    )
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    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;
+
+    #[test]
+    fn test_null_sensitive_with_is_null() {
+        let expr = is_null(col("col1"));
+        let labels = label_null_sensitive(&expr);
+
+        // The root expression should be null-sensitive
+        assert_eq!(labels.get(&expr), Some(&true));
+    }
+
+    #[test]
+    fn test_null_sensitive_without_is_null() {
+        let expr = eq(col("col1"), lit(5));
+        let labels = label_null_sensitive(&expr);
+
+        // Since the default is conservative (true), all expressions are sensitive
+        assert_eq!(labels.get(&expr), Some(&true));
+    }
+
+    #[test]
+    fn test_null_sensitive_nested() {
+        let left = eq(col("col1"), lit(5));
+        let right = is_null(col("col2"));
+        let expr = eq(left.clone(), right.clone());
+
+        let labels = label_null_sensitive(&expr);
+
+        // With conservative defaults, all are sensitive
+        assert_eq!(labels.get(&left), Some(&true));
+        assert_eq!(labels.get(&right), Some(&true));
+        assert_eq!(labels.get(&expr), Some(&true));
+    }
+}

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

@@ -232,6 +232,12 @@ impl Expression {
         self.stat_expression(Stat::Max, catalog)
     }
 
+    /// Returns whether this expression itself is null-sensitive.
+    /// See [`VTable::is_null_sensitive`].
+    pub fn is_null_sensitive(&self) -> bool {
+        self.vtable.as_dyn().is_null_sensitive(self.data.as_ref())
+    }
+
     /// Format the expression as a compact string.
     ///
     /// Since this is a recursive formatter, it is exposed on the public Expression type.

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

@@ -151,6 +151,10 @@ impl VTable for Between {
     ) -> Option<Expression> {
         expr.to_binary_expr().stat_falsification(catalog)
     }
+
+    fn is_null_sensitive(&self, _instance: &Self::Instance) -> bool {
+        false
+    }
 }
 
 impl ExpressionView<'_, Between> {

vortex-array/src/expr/expression.rs

@@ -252,6 +252,10 @@ impl VTable for Binary {
             Operator::Add | Operator::Sub | Operator::Mul | Operator::Div => None,
         }
     }
+
+    fn is_null_sensitive(&self, _instance: &Self::Instance) -> bool {
+        false
+    }
 }
 
 impl ExpressionView<'_, Binary> {