Propagate aggregate equivalences

Author: frankmcsherry

src/transform/src/equivalence_propagation.rs

@@ -186,9 +186,23 @@ impl Analysis for EquivalencePropagation {
             }
             MirRelationExpr::Reduce {
                 group_key,
-                aggregates: _,
+                aggregates,
                 ..
             } => {
+                // All of the output columns are the result of applying expressions to input columns,
+                // and we require a game plan to emerge with any equivalences about the output columns.
+                //
+                // Our plan is to (pretend to) extend the input with the expressions in `group_key`,
+                // and the apply a projection that retains only these columns, preserving any equivalences
+                // among the grouping columns.
+                // Additionally, for each "equality preserving" aggregation (MIN, MAX, ANY, ALL) we will
+                // append their expression to those of the grouping keys, project to the keys and the
+                // aggregate expression, and permute it in to place (as it may not follow the grouping keys).
+                // We will finally stitch all of these classes together.
+                //
+                // Importantly, we handle the aggregate expressions in isolation, as equivalences that
+                // may hold (e.g. `col1 = col2`) may not hold after aggregation (`MIN(col1) v MAX(col2)`).
+
                 let input_arity = depends.results::<Arity>().unwrap()[results.len() - 1];
                 let mut equivalences = results.last().unwrap().clone();
                 // Introduce keys column equivalences as a map, then project to them as a projection.
@@ -197,13 +211,30 @@ impl Analysis for EquivalencePropagation {
                         .classes
                         .push(vec![MirScalarExpr::Column(input_arity + pos), expr.clone()]);
                 }
-                // TODO: MIN, MAX, ANY, ALL aggregates pass through all certain properties of their columns.
-                // They also pass through equivalences of them and other constant columns (e.g. key columns).
-                // However, it is not correct to simply project onto these columns, as relationships amongst
-                // aggregate columns may no longer be preserved. MAX(col) != MIN(col) even though col = col.
-                // TODO: COUNT ensures a non-null value.
-                equivalences.project(input_arity..(input_arity + group_key.len()));
-                equivalences
+
+                // Accumulate the equivalence classes for the output of `Reduce`.
+                let mut result = equivalences.clone();
+                result.project(input_arity..(input_arity + group_key.len()));
+
+                // Introduce aggregate expressions as columns and equate the columns and expressions.
+                for (pos, aggr) in aggregates.iter().enumerate() {
+                    if preserves_equivalences(&aggr.func) {
+                        let mut equivs = equivalences.clone();
+                        equivs.classes.push(vec![
+                            MirScalarExpr::Column(input_arity + group_key.len()),
+                            aggr.expr.clone(),
+                        ]);
+                        equivs.project(input_arity..(input_arity + group_key.len() + 1));
+                        let permutation = (0..group_key.len())
+                            .chain(Some(group_key.len() + pos))
+                            .collect::<Vec<_>>();
+                        equivs.permute(&permutation[..]);
+                        // Fold the classes in to
+                        result.classes.extend(equivs.classes.drain(..));
+                    }
+                }
+
+                result
             }
             MirRelationExpr::TopK { .. } => results.last().unwrap().clone(),
             MirRelationExpr::Negate { .. } => results.last().unwrap().clone(),
@@ -775,6 +806,8 @@ impl EquivalenceClasses {
     }
 
     /// Permutes each expression, looking up each column reference in `permutation` and replacing with what it finds.
+    ///
+    /// This is simpler than `project`, in that it only relabels column references and it does not la
     fn permute(&mut self, permutation: &[usize]) {
         for class in self.classes.iter_mut() {
             for expr in class.iter_mut() {
@@ -933,38 +966,42 @@ impl EquivalenceClasses {
     }
 }
 
-// fn preserves_equivalences(func: &AggregateFunc) -> bool {
-//     match func {
-//           AggregateFunc::MaxInt16
-//         | AggregateFunc::MaxInt32
-//         | AggregateFunc::MaxInt64
-//         | AggregateFunc::MaxUInt16
-//         | AggregateFunc::MaxUInt32
-//         | AggregateFunc::MaxUInt64
-//         | AggregateFunc::MaxMzTimestamp
-//         | AggregateFunc::MaxFloat32
-//         | AggregateFunc::MaxFloat64
-//         | AggregateFunc::MaxBool
-//         | AggregateFunc::MaxString
-//         | AggregateFunc::MaxDate
-//         | AggregateFunc::MaxTimestamp
-//         | AggregateFunc::MaxTimestampTz
-//         | AggregateFunc::MinInt16
-//         | AggregateFunc::MinInt32
-//         | AggregateFunc::MinInt64
-//         | AggregateFunc::MinUInt16
-//         | AggregateFunc::MinUInt32
-//         | AggregateFunc::MinUInt64
-//         | AggregateFunc::MinMzTimestamp
-//         | AggregateFunc::MinFloat32
-//         | AggregateFunc::MinFloat64
-//         | AggregateFunc::MinBool
-//         | AggregateFunc::MinString
-//         | AggregateFunc::MinDate
-//         | AggregateFunc::MinTimestamp
-//         | AggregateFunc::MinTimestampTz
-//         | AggregateFunc::Any
-//         | AggregateFunc::All => true,
-//         _ => false,
-//     }
-// }
+use mz_expr::AggregateFunc;
+
+/// True iff the aggregate function produces an instance of its input expression,
+/// and consequently preserves any equivalence involving
+fn preserves_equivalences(func: &AggregateFunc) -> bool {
+    match func {
+        AggregateFunc::MaxInt16
+        | AggregateFunc::MaxInt32
+        | AggregateFunc::MaxInt64
+        | AggregateFunc::MaxUInt16
+        | AggregateFunc::MaxUInt32
+        | AggregateFunc::MaxUInt64
+        | AggregateFunc::MaxMzTimestamp
+        | AggregateFunc::MaxFloat32
+        | AggregateFunc::MaxFloat64
+        | AggregateFunc::MaxBool
+        | AggregateFunc::MaxString
+        | AggregateFunc::MaxDate
+        | AggregateFunc::MaxTimestamp
+        | AggregateFunc::MaxTimestampTz
+        | AggregateFunc::MinInt16
+        | AggregateFunc::MinInt32
+        | AggregateFunc::MinInt64
+        | AggregateFunc::MinUInt16
+        | AggregateFunc::MinUInt32
+        | AggregateFunc::MinUInt64
+        | AggregateFunc::MinMzTimestamp
+        | AggregateFunc::MinFloat32
+        | AggregateFunc::MinFloat64
+        | AggregateFunc::MinBool
+        | AggregateFunc::MinString
+        | AggregateFunc::MinDate
+        | AggregateFunc::MinTimestamp
+        | AggregateFunc::MinTimestampTz
+        | AggregateFunc::Any
+        | AggregateFunc::All => true,
+        _ => false,
+    }
+}

test/testdrive/mz-arrangement-sharing.td

@@ -705,7 +705,7 @@ ReduceMinsMaxes
 "Arrange ReduceMinsMaxes"                     1
 "ArrangeAccumulable [val: empty]"             3
 "ArrangeBy[[Column(0), Column(2)]]"           1
-"ArrangeBy[[Column(0)]] [val: empty]"         2
+"ArrangeBy[[Column(0)]] [val: empty]"         1
 "ArrangeBy[[Column(0)]]-errors"              20
 "ArrangeBy[[Column(0)]]"                     39
 "ArrangeBy[[Column(1), Column(3)]]"           2