sql: Convert transforms to use repr types

src/expr/src/relation.rs

@@ -380,6 +380,66 @@ impl MirRelationExpr {
         type_stack.pop().unwrap()
     }
 
+    /// Reports the representation type of the relation.
+    ///
+    /// This method determines the type through recursive traversal of the
+    /// relation expression, drawing from the types of base collections.
+    /// As such, this is not an especially cheap method, and should be used
+    /// judiciously.
+    ///
+    /// The relation type is computed incrementally with a recursive post-order
+    /// traversal, that accumulates the input types for the relations yet to be
+    /// visited in `type_stack`.
+    pub fn repr_typ(&self) -> ReprRelationType {
+        let mut type_stack = Vec::new();
+        #[allow(deprecated)]
+        self.visit_pre_post_nolimit(
+            &mut |e: &MirRelationExpr| -> Option<Vec<&MirRelationExpr>> {
+                match &e {
+                    MirRelationExpr::Let { body, .. } => {
+                        // Do not traverse the value sub-graph, since it's not relevant for
+                        // determining the relation type of Let operators.
+                        Some(vec![&*body])
+                    }
+                    MirRelationExpr::LetRec { body, .. } => {
+                        // Do not traverse the value sub-graph, since it's not relevant for
+                        // determining the relation type of Let operators.
+                        Some(vec![&*body])
+                    }
+                    _ => None,
+                }
+            },
+            &mut |e: &MirRelationExpr| {
+                match e {
+                    MirRelationExpr::Let { .. } => {
+                        let body_typ = type_stack.pop().unwrap();
+                        // Insert a dummy relation type for the value, since `typ_with_input_types`
+                        // won't look at it, but expects the relation type of the body to be second.
+                        type_stack.push(ReprRelationType::empty());
+                        type_stack.push(body_typ);
+                    }
+                    MirRelationExpr::LetRec { values, .. } => {
+                        let body_typ = type_stack.pop().unwrap();
+                        // Insert dummy relation types for the values, since `typ_with_input_types`
+                        // won't look at them, but expects the relation type of the body to be last.
+                        type_stack.extend(
+                            std::iter::repeat(ReprRelationType::empty()).take(values.len()),
+                        );
+                        type_stack.push(body_typ);
+                    }
+                    _ => {}
+                }
+                let num_inputs = e.num_inputs();
+                let relation_type =
+                    e.repr_typ_with_input_types(&type_stack[type_stack.len() - num_inputs..]);
+                type_stack.truncate(type_stack.len() - num_inputs);
+                type_stack.push(relation_type);
+            },
+        );
+        assert_eq!(type_stack.len(), 1);
+        type_stack.pop().unwrap()
+    }
+
     /// Reports the schema of the relation given the schema of the input relations.
     ///
     /// `input_types` is required to contain the schemas for the input relations of
@@ -404,6 +464,27 @@ impl MirRelationExpr {
             .with_keys(unique_keys)
     }
 
+    /// Reports the representation type of the relation given the repr types of the input relations.
+    ///
+    /// `input_types` is required to contain the schemas for the input relations of
+    /// the current relation in the same order as they are visited by `try_visit_children`
+    /// method, even though not all may be used for computing the schema of the
+    /// current relation. For example, `Let` expects two input types, one for the
+    /// value relation and one for the body, in that order, but only the one for the
+    /// body is used to determine the type of the `Let` relation.
+    ///
+    /// It is meant to be used during post-order traversals to compute relation
+    /// schemas incrementally.
+    pub fn repr_typ_with_input_types(&self, input_types: &[ReprRelationType]) -> ReprRelationType {
+        let column_types =
+            self.repr_col_with_input_repr_cols(input_types.iter().map(|i| &i.column_types));
+        let unique_keys = self.keys_with_input_keys(
+            input_types.iter().map(|i| i.arity()),
+            input_types.iter().map(|i| &i.keys),
+        );
+        ReprRelationType::new(column_types).with_keys(unique_keys)
+    }
+
     /// Reports the column types of the relation given the column types of the
     /// input relations.
     ///

src/expr/src/relation/canonicalize.rs

@@ -16,7 +16,7 @@ use std::rc::Rc;
 
 use itertools::Itertools;
 use mz_ore::soft_assert_or_log;
-use mz_repr::{SqlColumnType, SqlScalarType};
+use mz_repr::{ReprColumnType, ReprScalarType};
 
 use crate::visit::Visit;
 use crate::{MirScalarExpr, UnaryFunc, VariadicFunc, func};
@@ -40,7 +40,7 @@ pub fn canonicalize_equivalences<'a, I>(
     equivalences: &mut Vec<Vec<MirScalarExpr>>,
     input_column_types: I,
 ) where
-    I: Iterator<Item = &'a Vec<SqlColumnType>>,
+    I: Iterator<Item = &'a Vec<ReprColumnType>>,
 {
     let column_types = input_column_types
         .flat_map(|f| f.clone())
@@ -216,12 +216,12 @@ where
 /// null rejecting predicate for the same sub-expression.
 pub fn canonicalize_predicates(
     predicates: &mut Vec<MirScalarExpr>,
-    column_types: &[SqlColumnType],
+    column_types: &[ReprColumnType],
 ) {
     soft_assert_or_log!(
         predicates
             .iter()
-            .all(|p| p.typ(column_types).scalar_type == SqlScalarType::Bool),
+            .all(|p| p.repr_typ(column_types).scalar_type == ReprScalarType::Bool),
         "cannot canonicalize predicates that are not of type bool"
     );
 
@@ -369,7 +369,7 @@ pub fn canonicalize_predicates(
         (p.is_literal_false() || p.is_literal_null()) &&
         // This extra check is only needed if we determine that the soft-assert
         // at the top of this function would ever fail for a good reason.
-        p.typ(column_types).scalar_type == SqlScalarType::Bool
+        p.repr_typ(column_types).scalar_type == ReprScalarType::Bool
     }) {
         // all rows get filtered away if any predicate is null or false.
         *predicates = vec![MirScalarExpr::literal_false()]
@@ -390,7 +390,7 @@ fn replace_subexpr_and_reduce(
     predicate: &mut MirScalarExpr,
     replace_if_equal_to: &MirScalarExpr,
     replace_with: &MirScalarExpr,
-    column_types: &[SqlColumnType],
+    column_types: &[ReprColumnType],
 ) -> bool {
     let mut changed = false;
     #[allow(deprecated)]