Add comments and do a bit of refactoring in FlatMapElimination

Author: ggevay

src/expr/src/scalar.rs

@@ -648,6 +648,8 @@ impl MirScalarExpr {
         mem::replace(self, MirScalarExpr::literal_null(ScalarType::String))
     }
 
+    /// If the expression is a literal, this returns the literal's Datum or the literal's EvalError.
+    /// Otherwise, it returns None.
     pub fn as_literal(&self) -> Option<Result<Datum<'_>, &EvalError>> {
         if let MirScalarExpr::Literal(lit, _column_type) = self {
             Some(lit.as_ref().map(|row| row.unpack_first()))
@@ -656,6 +658,12 @@ impl MirScalarExpr {
         }
     }
 
+    /// Flattens the two failure modes of `as_literal` into one layer of Option: returns the
+    /// literal's Datum only if the expression is a literal, and it's not a literal error.
+    pub fn as_literal_non_error(&self) -> Option<Datum<'_>> {
+        self.as_literal().map(|eval_err| eval_err.ok()).flatten()
+    }
+
     pub fn as_literal_owned(&self) -> Option<Result<Row, EvalError>> {
         if let MirScalarExpr::Literal(lit, _column_type) = self {
             Some(lit.clone())

src/transform/src/canonicalization/flat_map_elimination.rs

@@ -7,10 +7,13 @@
 // the Business Source License, use of this software will be governed
 // by the Apache License, Version 2.0.
 
-//! For a `FlatMap` whose args are all constants, turns it into `Map` if only 1 row is produced by
-//! the table function, or turns it into an empty constant if 0 rows are produced by the table
-//! function. Additionally, a `Wrap` whose width is larger than its number of arguments can be
-//! removed.
+//! For a `FlatMap` where the table function's arguments are all constants, turns it into `Map` if
+//! only 1 row is produced by the table function, or turns it into an empty constant collection if 0
+//! rows are produced by the table function.
+//!
+//! It does an additional optimization on the `Wrap` table function: when `Wrap`'s width is larger
+//! than its number of arguments, it removes the `FlatMap Wrap ...`, because such `Wrap`s would have
+//! no effect.
 
 use itertools::Itertools;
 use mz_expr::visit::Visit;
@@ -19,7 +22,7 @@ use mz_repr::{Diff, Row, RowArena};
 
 use crate::TransformCtx;
 
-/// See comment at the top of the file.
+/// Attempts to eliminate FlatMaps that are sure to have 0 or 1 results on each input row.
 #[derive(Debug)]
 pub struct FlatMapElimination;
 
@@ -45,48 +48,55 @@ impl crate::Transform for FlatMapElimination {
 }
 
 impl FlatMapElimination {
-    /// See comment at the top of the file.
+    /// Apply `FlatMapElimination` to the root of the given `MirRelationExpr`.
     pub fn action(relation: &mut MirRelationExpr) {
+        // Treat Wrap specially: we can sometimes optimize it out even when it has non-literal
+        // arguments.
+        //
+        // (No need to look for WithOrdinality here, as that never occurs with Wrap: users can't
+        // call Wrap directly; we only create calls to Wrap ourselves, and we don't use
+        // WithOrdinality on it.)
         if let MirRelationExpr::FlatMap { func, exprs, input } = relation {
-            // Treat Wrap specially.
             if let TableFunc::Wrap { width, .. } = func {
                 if *width >= exprs.len() {
                     *relation = input.take_dangerous().map(std::mem::take(exprs));
-                    return;
                 }
             }
-            // For all other table functions, check for all arguments being literals.
-            let mut args = vec![];
-            for e in exprs {
-                match e.as_literal() {
-                    Some(Ok(datum)) => args.push(datum),
-                    // Give up if any arg is not a literal, or if it's a literal error.
-                    _ => return,
-                }
-            }
-            let temp_storage = RowArena::new();
-            let (first, second) = match func.eval(&args, &temp_storage) {
-                Ok(mut r) => (r.next(), r.next()),
-                // don't play with errors
-                Err(_) => return,
-            };
-            match (first, second) {
-                // The table function evaluated to an empty collection.
-                (None, None) => {
-                    relation.take_safely(None);
-                }
-                // The table function evaluated to a collection with exactly 1 row.
-                (Some((first_row, Diff::ONE)), None) => {
-                    let types = func.output_type().column_types;
-                    let map_exprs = first_row
-                        .into_iter()
-                        .zip_eq(types)
-                        .map(|(d, typ)| MirScalarExpr::Literal(Ok(Row::pack_slice(&[d])), typ))
-                        .collect();
-                    *relation = input.take_dangerous().map(map_exprs);
+        }
+        // For all other table functions (and Wraps that are not covered by the above), check
+        // whether all arguments are literals (with no errors), in which case we'll evaluate the
+        // table function and check how many output rows it has, and maybe turn the FlatMap into
+        // something simpler.
+        if let MirRelationExpr::FlatMap { func, exprs, input } = relation {
+            if let Some(args) = exprs
+                .iter()
+                .map(|e| e.as_literal_non_error())
+                .collect::<Option<Vec<_>>>()
+            {
+                let temp_storage = RowArena::new();
+                let (first, second) = match func.eval(&args, &temp_storage) {
+                    Ok(mut r) => (r.next(), r.next()),
+                    // don't play with errors
+                    Err(_) => return,
+                };
+                match (first, second) {
+                    // The table function evaluated to an empty collection.
+                    (None, _) => {
+                        relation.take_safely(None);
+                    }
+                    // The table function evaluated to a collection with exactly 1 row.
+                    (Some((first_row, Diff::ONE)), None) => {
+                        let types = func.output_type().column_types;
+                        let map_exprs = first_row
+                            .into_iter()
+                            .zip_eq(types)
+                            .map(|(d, typ)| MirScalarExpr::Literal(Ok(Row::pack_slice(&[d])), typ))
+                            .collect();
+                        *relation = input.take_dangerous().map(map_exprs);
+                    }
+                    // The table function evaluated to a collection with more than 1 row; nothing to do.
+                    _ => {}
                 }
-                // The table function evaluated to a collection with more than 1 row; nothing to do.
-                _ => {}
             }
         }
     }

src/transform/tests/test_transforms/flat_map_elimination.spec

@@ -152,4 +152,4 @@ FlatMap generate_series(5, 6, 1)
   Get t0
 ----
 FlatMap generate_series(5, 6, 1)
-  Get
+  Get t0

test/sqllogictest/table_func.slt

@@ -1715,7 +1715,7 @@ query T multiline
 EXPLAIN OPTIMIZED PLAN WITH (NO FAST PATH) FOR
 SELECT *
 FROM x, generate_series(1, x.a)
-WHERE x.a = 1
+WHERE x.a = 1;
 ----
 Explained Query:
   Project (#0, #1, #0)
@@ -1732,11 +1732,28 @@ query T multiline
 EXPLAIN OPTIMIZED PLAN WITH (NO FAST PATH) FOR
 SELECT *
 FROM x, generate_series(5, x.a)
-WHERE x.a = 1
+WHERE x.a = 1;
 ----
 Explained Query:
   Constant <empty>
 
 Target cluster: quickstart
 
 EOF
+
+query T multiline
+EXPLAIN OPTIMIZED PLAN WITH (NO FAST PATH) FOR
+SELECT *
+FROM x, LATERAL (VALUES ((1), (x.a)))
+WHERE x.a = 1;
+----
+Explained Query:
+  Project (#0, #1, #0, #0)
+    Filter (#0{a} = 1)
+      ReadStorage materialize.public.x
+
+Source materialize.public.x
+
+Target cluster: quickstart
+
+EOF