mz_transform/
reduce_elision.rs

Help
1// Copyright Materialize, Inc. and contributors. All rights reserved.
2//
3// Use of this software is governed by the Business Source License
4// included in the LICENSE file.
5//
6// As of the Change Date specified in that file, in accordance with
7// the Business Source License, use of this software will be governed
8// by the Apache License, Version 2.0.
9
10//! Removes `Reduce` when the input has as unique keys the keys of the reduce.
11//!
12//! When a reduce has grouping keys that are contained within a
13//! set of columns that form unique keys for the input, the reduce
14//! can be simplified to a map operation.
15
16use itertools::Itertools;
17use mz_expr::MirRelationExpr;
18
19use crate::TransformCtx;
20use crate::analysis::{DerivedBuilder, DerivedView};
21use crate::analysis::{RelationType, UniqueKeys};
22
23/// Removes `Reduce` when the input has as unique keys the keys of the reduce.
24#[derive(Debug)]
25pub struct ReduceElision;
26
27impl crate::Transform for ReduceElision {
28    fn name(&self) -> &'static str {
29        "ReduceElision"
30    }
31
32    #[mz_ore::instrument(
33        target = "optimizer",
34        level = "debug",
35        fields(path.segment = "reduce_elision")
36    )]
37    fn actually_perform_transform(
38        &self,
39        relation: &mut MirRelationExpr,
40        ctx: &mut TransformCtx,
41    ) -> Result<(), crate::TransformError> {
42        // Assemble type information once for the whole expression.
43        let mut builder = DerivedBuilder::new(ctx.features);
44        builder.require(RelationType);
45        builder.require(UniqueKeys);
46        let derived = builder.visit(relation);
47        let derived_view = derived.as_view();
48
49        self.action(relation, derived_view);
50
51        mz_repr::explain::trace_plan(&*relation);
52        Ok(())
53    }
54}
55
56impl ReduceElision {
57    /// Removes `Reduce` when the input has as unique keys the keys of the reduce.
58    pub fn action(&self, relation: &mut MirRelationExpr, derived: DerivedView) {
59        let mut todo = vec![(relation, derived)];
60        while let Some((expr, view)) = todo.pop() {
61            let mut replaced = false;
62            if let MirRelationExpr::Reduce {
63                input,
64                group_key,
65                aggregates,
66                monotonic: _,
67                expected_group_size: _,
68            } = expr
69            {
70                let input_type = view
71                    .last_child()
72                    .value::<RelationType>()
73                    .expect("RelationType required")
74                    .as_ref()
75                    .expect("Expression not well-typed");
76                let input_keys = view
77                    .last_child()
78                    .value::<UniqueKeys>()
79                    .expect("UniqueKeys required");
80
81                if input_keys.iter().any(|keys| {
82                    keys.iter()
83                        .all(|k| group_key.contains(&mz_expr::MirScalarExpr::Column(*k)))
84                }) {
85                    let map_scalars = aggregates
86                        .iter()
87                        .map(|a| a.on_unique(input_type))
88                        .collect_vec();
89
90                    let mut result = input.take_dangerous();
91
92                    let input_arity = input_type.len();
93
94                    // Append the group keys, then any `map_scalars`, then project
95                    // to put them all in the right order.
96                    let mut new_scalars = group_key.clone();
97                    new_scalars.extend(map_scalars);
98                    result = result.map(new_scalars).project(
99                        (input_arity..(input_arity + (group_key.len() + aggregates.len())))
100                            .collect(),
101                    );
102
103                    *expr = result;
104                    replaced = true;
105
106                    // // NB: The following is borked because of smart builders.
107                    // if let MirRelationExpr::Project { input, .. } = expr {
108                    //     if let MirRelationExpr::Map { input, .. } = &mut **input {
109                    //         todo.push((&mut **input, view.last_child()))
110                    //     }
111                    // }
112                }
113            }
114
115            // This gets around an awkward borrow of both `expr` and `input` above.
116            if !replaced {
117                todo.extend(expr.children_mut().rev().zip(view.children_rev()));
118            }
119        }
120    }
121}
mz_transform/reduce_elision.rs

mz_transform/
reduce_elision.rs
