Merge pull request #30884 from ggevay/typ-demand-literal-constr

Eliminate more `typ` and `arity` calls

Author: ggevay

src/compute-types/src/plan/lowering.rs

@@ -548,8 +548,6 @@ impl Context {
                 equivalences,
                 implementation,
             } => {
-                let input_mapper = JoinInputMapper::new(inputs);
-
                 // Plan each of the join inputs independently.
                 // The `plans` get surfaced upwards, and the `input_keys` should
                 // be used as part of join planning / to validate the existing
@@ -564,6 +562,9 @@ impl Context {
                     input_keys.push(keys);
                 }
 
+                let input_mapper =
+                    JoinInputMapper::new_from_input_arities(input_arities.iter().copied());
+
                 // Extract temporal predicates as joins cannot currently absorb them.
                 let (plan, missing) = match implementation {
                     IndexedFilter(_coll_id, _idx_id, key, _val) => {

src/expr/src/linear.rs

@@ -329,6 +329,8 @@ impl MapFilterProject {
                 let (mfp, expr) = Self::extract_from_expression(input);
                 (mfp.project(outputs.iter().cloned()), expr)
             }
+            // TODO: The recursion is quadratic in the number of Map/Filter/Project operators due to
+            // this call to `arity()`.
             x => (Self::new(x.arity()), x),
         }
     }

src/expr/src/relation.rs

@@ -593,7 +593,7 @@ impl MirRelationExpr {
     /// Reports the column types of the relation given the column types of the
     /// input relations.
     ///
-    /// This method delegates to `try_col_with_input_cols`, panicing if an `Err`
+    /// This method delegates to `try_col_with_input_cols`, panicking if an `Err`
     /// variant is returned.
     pub fn col_with_input_cols<'a, I>(&self, input_types: I) -> Vec<ColumnType>
     where
@@ -1792,22 +1792,21 @@ impl MirRelationExpr {
             .unzip();
         assert_eq!(keys_and_values.arity() - self.arity(), data.len());
         self.let_in(id_gen, |_id_gen, get_keys| {
+            let get_keys_arity = get_keys.arity();
             Ok(MirRelationExpr::join(
                 vec![
                     // all the missing keys (with count 1)
                     keys_and_values
-                        .distinct_by((0..get_keys.arity()).collect())
+                        .distinct_by((0..get_keys_arity).collect())
                         .negate()
                         .union(get_keys.clone().distinct()),
                     // join with keys to get the correct counts
                     get_keys.clone(),
                 ],
-                (0..get_keys.arity())
-                    .map(|i| vec![(0, i), (1, i)])
-                    .collect(),
+                (0..get_keys_arity).map(|i| vec![(0, i), (1, i)]).collect(),
             )
             // get rid of the extra copies of columns from keys
-            .project((0..get_keys.arity()).collect())
+            .project((0..get_keys_arity).collect())
             // This join is logically equivalent to
             // `.map(<default_expr>)`, but using a join allows for
             // potential predicate pushdown and elision in the

src/sql/src/plan/lowering.rs

@@ -1399,7 +1399,8 @@ impl HirScalarExpr {
 
                 // Record input arity here so that any group_keys that need to mutate get_inner
                 // don't add those columns to the aggregate input.
-                let input_arity = get_inner.typ().arity();
+                let input_type = get_inner.typ();
+                let input_arity = input_type.arity();
                 // The reduction that computes the window function must be keyed on the columns
                 // from the outer context, plus the expressions in the partition key. The current
                 // subquery will be 'executed' for every distinct row from the outer context so
@@ -1428,8 +1429,6 @@ impl HirScalarExpr {
                 }
 
                 get_inner.let_in(id_gen, |id_gen, mut get_inner| {
-                    let input_type = get_inner.typ();
-
                     // Original columns of the relation
                     let fields: Box<_> = input_type
                         .column_types
@@ -1575,12 +1574,14 @@ impl HirScalarExpr {
                 let inner_arity = get_inner.arity();
                 let mut total_arity = inner_arity;
                 let mut join_inputs = vec![get_inner];
+                let mut join_input_arities = vec![inner_arity];
                 for (expr, subquery) in subqueries.into_iter() {
                     // Avoid lowering duplicated subqueries
                     if !subquery_map.contains_key(&expr) {
                         let subquery_arity = subquery.arity();
                         assert_eq!(subquery_arity, inner_arity + 1);
                         join_inputs.push(subquery);
+                        join_input_arities.push(subquery_arity);
                         total_arity += subquery_arity;
 
                         // Column with the value of the subquery
@@ -1590,7 +1591,8 @@ impl HirScalarExpr {
                 // Each subquery projects all the columns of the outer context (distinct_inner)
                 // plus 1 column, containing the result of the subquery. Those columns must be
                 // joined with the outer/main relation (get_inner).
-                let input_mapper = mz_expr::JoinInputMapper::new(&join_inputs);
+                let input_mapper =
+                    mz_expr::JoinInputMapper::new_from_input_arities(join_input_arities);
                 let equivalences = (0..inner_arity)
                     .map(|col| {
                         join_inputs

src/transform/src/demand.rs

@@ -176,7 +176,7 @@ impl Demand {
                 ),
                 MirRelationExpr::Map { input, scalars } => {
                     let relation_type = relation_type.as_ref().unwrap();
-                    let arity = input.arity();
+                    let arity = relation_type.arity() - scalars.len();
                     // contains columns whose supports have yet to be explored
                     let mut new_columns = columns.clone();
                     new_columns.retain(|c| *c >= arity);
@@ -219,7 +219,8 @@ impl Demand {
                     for expr in exprs {
                         expr.support_into(&mut columns);
                     }
-                    columns.retain(|c| *c < input.arity());
+                    let arity = input.arity();
+                    columns.retain(|c| *c < arity);
                     self.action(input, columns, gets)
                 }
                 MirRelationExpr::Filter { input, predicates } => {

src/transform/src/literal_constraints.rs

@@ -29,7 +29,7 @@ use mz_ore::collections::CollectionExt;
 use mz_ore::iter::IteratorExt;
 use mz_ore::stack::RecursionLimitError;
 use mz_ore::vec::swap_remove_multiple;
-use mz_repr::{GlobalId, Row};
+use mz_repr::{GlobalId, RelationType, Row};
 
 use crate::canonicalize_mfp::CanonicalizeMfp;
 use crate::notice::IndexTooWideForLiteralConstraints;
@@ -70,7 +70,9 @@ impl LiteralConstraints {
         relation.try_visit_mut_children(|e| self.action(e, transform_ctx))?;
 
         if let MirRelationExpr::Get {
-            id: Id::Global(id), ..
+            id: Id::Global(id),
+            ref typ,
+            ..
         } = *relation
         {
             let orig_mfp = mfp.clone();
@@ -87,10 +89,11 @@ impl LiteralConstraints {
                 mfp: &mut MapFilterProject,
                 orig_mfp: &MapFilterProject,
                 relation: &MirRelationExpr,
+                relation_type: RelationType,
             ) {
                 // undo list_of_predicates_to_and_of_predicates, distribute_and_over_or, unary_and
                 // (It undoes the latter 2 through `MirScalarExp::reduce`.)
-                LiteralConstraints::canonicalize_predicates(mfp, relation);
+                LiteralConstraints::canonicalize_predicates(mfp, relation, relation_type);
                 // undo inline_literal_constraints
                 mfp.optimize();
                 // We can usually undo, but sometimes not (see comment on `distribute_and_over_or`),
@@ -109,14 +112,16 @@ impl LiteralConstraints {
             // todo: We might want to also call `canonicalize_equivalences`,
             // see near the end of literal_constraints.slt.
 
+            let inp_typ = typ.clone();
+
             let key_val = Self::detect_literal_constraints(&mfp, id, transform_ctx);
 
             match key_val {
                 None => {
                     // We didn't find a usable index, so no chance to remove literal constraints.
                     // But, we might have removed contradicting OR args.
                     if removed_contradicting_or_args {
-                        undo_preparation(&mut mfp, &orig_mfp, relation);
+                        undo_preparation(&mut mfp, &orig_mfp, relation, inp_typ);
                     } else {
                         // We didn't remove anything, so let's go with the original MFP.
                         mfp = orig_mfp;
@@ -130,7 +135,7 @@ impl LiteralConstraints {
                     {
                         // We were able to remove the literal constraints or contradicting OR args,
                         // so we would like to use this new MFP, so we try undoing the preparation.
-                        undo_preparation(&mut mfp, &orig_mfp, relation);
+                        undo_preparation(&mut mfp, &orig_mfp, relation, inp_typ.clone());
                     } else {
                         // We were not able to remove the literal constraint, so `mfp` is
                         // equivalent to `orig_mfp`, but `orig_mfp` is often simpler (or the same).
@@ -140,7 +145,6 @@ impl LiteralConstraints {
                     // We transform the Get into a semi-join with a constant collection.
 
                     let inp_id = id.clone();
-                    let inp_typ = relation.typ();
                     let filter_list = MirRelationExpr::Constant {
                         rows: Ok(possible_vals.iter().map(|val| (val.clone(), 1)).collect()),
                         typ: mz_repr::RelationType {
@@ -164,10 +168,7 @@ impl LiteralConstraints {
                     if possible_vals.is_empty() {
                         // Even better than what we were hoping for: Found contradicting
                         // literal constraints, so the whole relation is empty.
-                        *relation = MirRelationExpr::Constant {
-                            rows: Ok(Vec::new()),
-                            typ: relation.typ(),
-                        };
+                        relation.take_safely(Some(inp_typ));
                     } else {
                         // The common case: We need to build the join which is the main point of
                         // this transform.
@@ -610,9 +611,16 @@ impl LiteralConstraints {
     }
 
     /// Call [mz_expr::canonicalize::canonicalize_predicates] on each of the predicates in the MFP.
-    fn canonicalize_predicates(mfp: &mut MapFilterProject, relation: &MirRelationExpr) {
+    fn canonicalize_predicates(
+        mfp: &mut MapFilterProject,
+        relation: &MirRelationExpr,
+        relation_type: RelationType,
+    ) {
         let (map, mut predicates, project) = mfp.as_map_filter_project();
-        let typ_after_map = relation.clone().map(map.clone()).typ();
+        let typ_after_map = relation
+            .clone()
+            .map(map.clone())
+            .typ_with_input_types(&[relation_type]);
         canonicalize_predicates(&mut predicates, &typ_after_map.column_types);
         // Rebuild the MFP with the new predicates.
         *mfp = MapFilterProject::new(mfp.input_arity)