Skip to main content

vortex_array/expr/analysis/
strict.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use super::BooleanLabels;
5use super::labeling::label_tree;
6use crate::expr::Expression;
7
8/// Label each expression with whether its entire subtree is strict.
9///
10/// A subtree is strict only when the node's scalar function and every child subtree are strict.
11/// See [`crate::scalar_fn::ScalarFnVTable::is_strict`] for the scalar-function contract.
12pub fn label_strict(expr: &Expression) -> BooleanLabels<'_> {
13    label_tree(
14        expr,
15        |expr| match expr {
16            Expression::Scalar { scalar_fn, .. } => scalar_fn.signature().is_strict(),
17            // Vacuously strict.
18            Expression::Root => true,
19        },
20        |acc, &child| acc & child,
21    )
22}
23
24#[cfg(test)]
25mod tests {
26    use super::*;
27    use crate::expr::col;
28    use crate::expr::eq;
29    use crate::expr::is_null;
30    use crate::expr::lit;
31
32    #[test]
33    fn test_non_strict_with_is_null() {
34        let expr = is_null(col("col1"));
35        let labels = label_strict(&expr);
36
37        assert_eq!(labels.get(&expr), Some(&false));
38    }
39
40    #[test]
41    fn test_strict_expression() {
42        let expr = eq(lit(4), lit(5));
43        let labels = label_strict(&expr);
44
45        assert_eq!(labels.get(&expr), Some(&true));
46    }
47
48    #[test]
49    fn test_non_strict_child_makes_parent_subtree_non_strict() {
50        let left = eq(lit(4), lit(5));
51        let right = is_null(col("col2"));
52        let expr = eq(left.clone(), right.clone());
53
54        let labels = label_strict(&expr);
55
56        assert_eq!(labels.get(&left), Some(&true));
57        assert_eq!(labels.get(&right), Some(&false));
58        assert_eq!(labels.get(&expr), Some(&false));
59    }
60}