Rewrite FlatMapElimination to simplify and generalize it

From Petros

Author: ggevay

src/expr/src/relation/func.rs

@@ -3792,19 +3792,16 @@ pub enum TableFunc {
     WithOrdinality(WithOrdinality),
 }
 
-/// Private enum variant of `TableFunc`. Don't construct this directly, but use
-/// `TableFunc::with_ordinality` instead.
-///
 /// Evaluates the inner table function, expands its results into unary (repeating each row as
 /// many times as the diff indicates), and appends an integer corresponding to the ordinal
 /// position (starting from 1). For example, it numbers the elements of a list when calling
 /// `unnest_list`.
 ///
-/// TODO(ggevay): This struct (and its field) is pub only temporarily, until we make
-/// `FlatMapElimination` not dive into it.
+/// Private enum variant of `TableFunc`. Don't construct this directly, but use
+/// `TableFunc::with_ordinality` instead.
 #[derive(Clone, Debug, Eq, PartialEq, Ord, PartialOrd, Serialize, Deserialize, Hash, MzReflect)]
-pub struct WithOrdinality {
-    pub inner: Box<TableFunc>,
+struct WithOrdinality {
+    inner: Box<TableFunc>,
 }
 
 impl TableFunc {

src/transform/src/canonicalization/flat_map_elimination.rs

@@ -7,16 +7,19 @@
 // the Business Source License, use of this software will be governed
 // by the Apache License, Version 2.0.
 
-//! Turns `FlatMap` into `Map` if only one row is produced by flatmap.
-//!
+//! 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.
 
+use itertools::Itertools;
 use mz_expr::visit::Visit;
 use mz_expr::{MirRelationExpr, MirScalarExpr, TableFunc};
-use mz_repr::{Datum, Diff, ScalarType};
+use mz_repr::{Diff, Row, RowArena};
 
 use crate::TransformCtx;
 
-/// Turns `FlatMap` into `Map` if only one row is produced by flatmap.
+/// See comment at the top of the file.
 #[derive(Debug)]
 pub struct FlatMapElimination;
 
@@ -42,70 +45,49 @@ impl crate::Transform for FlatMapElimination {
 }
 
 impl FlatMapElimination {
-    /// Turns `FlatMap` into `Map` if only one row is produced by flatmap.
+    /// See comment at the top of the file.
     pub fn action(relation: &mut MirRelationExpr) {
         if let MirRelationExpr::FlatMap { func, exprs, input } = relation {
-            let (func, with_ordinality) = if let TableFunc::WithOrdinality(with_ordinality) = func {
-                // get to the actual function, but remember that we have a WITH ORDINALITY clause.
-                (&*with_ordinality.inner, true)
-            } else {
-                (&*func, false)
-            };
-
-            if let TableFunc::GuardSubquerySize { .. } = func {
-                // (`with_ordinality` doesn't matter because this function never emits rows)
-                if let Some(1) = exprs[0].as_literal_int64() {
-                    relation.take_safely(None);
-                }
-            } else if let TableFunc::Wrap { width, .. } = func {
+            // Treat Wrap specially.
+            if let TableFunc::Wrap { width, .. } = func {
                 if *width >= exprs.len() {
                     *relation = input.take_dangerous().map(std::mem::take(exprs));
-                    if with_ordinality {
-                        *relation = relation.take_dangerous().map_one(MirScalarExpr::literal(
-                            Ok(Datum::Int64(1)),
-                            ScalarType::Int64,
-                        ));
-                    }
+                    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,
                 }
-            } else if is_supported_unnest(func) {
-                let func = func.clone();
-                let exprs = exprs.clone();
-                use mz_expr::MirScalarExpr;
-                use mz_repr::RowArena;
-                if let MirScalarExpr::Literal(Ok(row), ..) = &exprs[0] {
-                    let temp_storage = RowArena::default();
-                    if let Ok(mut iter) = func.eval(&[row.iter().next().unwrap()], &temp_storage) {
-                        match (iter.next(), iter.next()) {
-                            (None, _) => {
-                                // If there are no elements in the literal argument, no output.
-                                relation.take_safely(None);
-                            }
-                            (Some((row, Diff::ONE)), None) => {
-                                assert_eq!(func.output_type().column_types.len(), 1);
-                                *relation =
-                                    input.take_dangerous().map(vec![MirScalarExpr::Literal(
-                                        Ok(row),
-                                        func.output_type().column_types[0].clone(),
-                                    )]);
-                                if with_ordinality {
-                                    *relation =
-                                        relation.take_dangerous().map_one(MirScalarExpr::literal(
-                                            Ok(Datum::Int64(1)),
-                                            ScalarType::Int64,
-                                        ));
-                                }
-                            }
-                            _ => {}
-                        }
-                    };
+            }
+            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);
+                }
+                // The table function evaluated to a collection with more than 1 row; nothing to do.
+                _ => {}
             }
         }
     }
 }
-
-/// Returns `true` for `unnest_~` variants supported by [`FlatMapElimination`].
-fn is_supported_unnest(func: &TableFunc) -> bool {
-    use TableFunc::*;
-    matches!(func, UnnestArray { .. } | UnnestList { .. })
-}

test/sqllogictest/table_func.slt

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