Do not hoist expressions from if bodies (MaterializeInc#3395)

Scalar CSE was hoisting common expressions in then and els expressions, without certainty that they would be evaluated. This led to queries like
```
SELECT
    CASE WHEN b = 0 THEN 0 ELSE 1/b END,
    CASE WHEN b != 0 THEN 1/b ELSE 0 END
FROM x
```
unconditionally evaluating 1/b, because it would be hoisted prematurely. We don't do that now. We could be more advanced and determine the expressions in the intersection of the then and els expressions, but that is future work for now.

Author: frankmcsherry

ci/slt/slt.sh

@@ -37,5 +37,6 @@ sqllogictest \
 sqllogictest \
     -v \
     test/sqllogictest/*.slt \
+    test/sqllogictest/transform/*.slt \
     test/sqllogictest/sqlite/test \
     | tee -a target/slt.log

ci/test/slt-fast.sh

@@ -25,6 +25,7 @@ export RUST_BACKTRACE=full
 
 tests=(
     test/sqllogictest/*.slt
+    test/sqllogictest/transform/*.slt
     # test/sqllogictest/sqlite/test/evidence/in1.test
     test/sqllogictest/sqlite/test/evidence/in2.test
     test/sqllogictest/sqlite/test/evidence/slt_lang_aggfunc.test

src/transform/src/cse/map.rs

@@ -62,39 +62,18 @@ impl Map {
 /// This method extracts all expression AST nodes as individual expressions,
 /// and builds others out of column references to them, avoiding any duplication.
 /// Once complete, expressions which are referred to only once are re-inlined.
+///
+/// Some care is taken to ensure that `if` expressions only memoize expresions
+/// they are certain to evaluate.
 fn memoize_and_reuse(
     exprs: &mut [ScalarExpr],
     input_arity: usize,
 ) -> (Vec<ScalarExpr>, Vec<usize>) {
     let mut projection = (0..input_arity).collect::<Vec<_>>();
     let mut scalars = Vec::new();
     for expr in exprs.iter_mut() {
-        expr.visit_mut(&mut |node| {
-            match node {
-                ScalarExpr::Column(index) => {
-                    // Column references need to be rewritten, but do not
-                    // need to be memoized.
-                    *index = projection[*index];
-                }
-                ScalarExpr::Literal(_, _) => {
-                    // Literals do not need to be memoized.
-                }
-                _ => {
-                    if let Some(position) = scalars.iter().position(|e| e == node) {
-                        // Any complex expression that already exists as a prior column can
-                        // be replaced by a reference to that column.
-                        *node = ScalarExpr::Column(input_arity + position);
-                    } else {
-                        // A complex expression that does not exist should be memoized, and
-                        // replaced by a reference to the column.
-                        scalars.push(std::mem::replace(
-                            node,
-                            ScalarExpr::Column(input_arity + scalars.len()),
-                        ));
-                    }
-                }
-            }
-        });
+        // Carefully memoize expressions that will certainly be evaluated.
+        memoize(expr, &mut scalars, &mut projection, input_arity);
 
         // At this point `expr` should be a column reference or a literal. It may not have
         // been added to `scalars` and we should do so if it is a literal to make sure we
@@ -115,6 +94,53 @@ fn memoize_and_reuse(
     (scalars, projection)
 }
 
+/// Visit and memoize expression nodes.
+///
+/// Importantly, we should not memoize expressions that may not be excluded by virtue of
+/// being guarded by `if` expressions.
+fn memoize(
+    expr: &mut ScalarExpr,
+    scalars: &mut Vec<ScalarExpr>,
+    projection: &mut Vec<usize>,
+    input_arity: usize,
+) {
+    match expr {
+        ScalarExpr::Column(index) => {
+            // Column references need to be rewritten, but do not need to be memoized.
+            *index = projection[*index];
+        }
+        ScalarExpr::Literal(_, _) => {
+            // Literals do not need to be memoized.
+        }
+        _ => {
+            // We should not eagerly memoize `if` branches that might not be taken.
+            // TODO: Memoize expressions in the intersection of `then` and `els`.
+            if let ScalarExpr::If {
+                cond,
+                then: _,
+                els: _,
+            } = expr
+            {
+                memoize(cond, scalars, projection, input_arity);
+            } else {
+                expr.visit1_mut(|e| memoize(e, scalars, projection, input_arity));
+            }
+            if let Some(position) = scalars.iter().position(|e| e == expr) {
+                // Any complex expression that already exists as a prior column can
+                // be replaced by a reference to that column.
+                *expr = ScalarExpr::Column(input_arity + position);
+            } else {
+                // A complex expression that does not exist should be memoized, and
+                // replaced by a reference to the column.
+                scalars.push(std::mem::replace(
+                    expr,
+                    ScalarExpr::Column(input_arity + scalars.len()),
+                ));
+            }
+        }
+    }
+}
+
 /// Replaces column references that occur once with the referenced expression.
 ///
 /// This method in-lines expressions that are only referenced once, and then

src/transform/src/lib.rs

@@ -252,9 +252,9 @@ impl Default for Optimizer {
                     Box::new(crate::map_lifting::LiteralLifting),
                 ],
             }),
-            Box::new(crate::fusion::project::Project),
             Box::new(crate::reduction_pushdown::ReductionPushdown),
             Box::new(crate::cse::map::Map),
+            Box::new(crate::fusion::project::Project),
             Box::new(crate::reduction::FoldConstants),
         ];
         Self { transforms }

test/sqllogictest/transform/scalar_cse.slt

@@ -60,3 +60,18 @@ FROM x
 | Project (#4, #5, #7, #8)
 
 EOF
+
+
+# Ensure we don't inline if-guarded expressions
+query T multiline
+EXPLAIN PLAN FOR SELECT
+    CASE WHEN b = 0 THEN 0 ELSE 1/b END,
+    CASE WHEN b != 0 THEN 1/b ELSE 0 END
+FROM x
+----
+%0 =
+| Get materialize.public.x (u1)
+| Map if (#1 = 0) then {0} else {(1 / #1)}, if (#1 != 0) then {(1 / #1)} else {0}
+| Project (#2, #3)
+
+EOF