risingwave_frontend/optimizer/rule/
agg_group_by_simplify_rule.rs

1// Copyright 2023 RisingWave Labs
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use risingwave_expr::aggregate::PbAggKind;
16
17use super::prelude::{PlanRef, *};
18use crate::expr::InputRef;
19use crate::optimizer::plan_node::generic::{Agg, GenericPlanRef};
20use crate::optimizer::plan_node::*;
21use crate::utils::{Condition, IndexSet};
22
23/// Use functional dependencies to simplify aggregation's group by
24/// Before:
25/// group by = [a, b, c], where b -> [a, c]
26/// After
27/// group by b, `first_value`(a), `first_value`(c),
28pub struct AggGroupBySimplifyRule {}
29impl Rule<Logical> for AggGroupBySimplifyRule {
30    fn apply(&self, plan: PlanRef) -> Option<PlanRef> {
31        let agg: &LogicalAgg = plan.as_logical_agg()?;
32        let (agg_calls, group_key, grouping_sets, agg_input, _two_phase) = agg.clone().decompose();
33        if !grouping_sets.is_empty() {
34            return None;
35        }
36        let functional_dependency = agg_input.functional_dependency();
37        let group_key = group_key.to_vec();
38        if group_key.is_empty() || !functional_dependency.is_key(&group_key) {
39            return None;
40        }
41        let mut minimized_group_key = functional_dependency.minimize_key(&group_key);
42        // The minimized key is empty when the input has at most one row. Keep one group key
43        // column, otherwise the grouped agg would become a simple agg, which outputs one row
44        // instead of zero rows on empty input.
45        if minimized_group_key.is_empty() {
46            minimized_group_key.push(group_key[0]);
47        }
48        if minimized_group_key.len() < group_key.len() {
49            let new_group_key = IndexSet::from(minimized_group_key);
50            let new_group_key_len = new_group_key.len();
51            let mut new_agg_calls = vec![];
52            for &i in &group_key {
53                if !new_group_key.contains(i) {
54                    let data_type = agg_input.schema().fields[i].data_type();
55                    new_agg_calls.push(PlanAggCall {
56                        agg_type: PbAggKind::InternalLastSeenValue.into(),
57                        return_type: data_type.clone(),
58                        inputs: vec![InputRef::new(i, data_type)],
59                        distinct: false,
60                        order_by: vec![],
61                        filter: Condition::true_cond(),
62                        direct_args: vec![],
63                    });
64                }
65            }
66            new_agg_calls.extend(agg_calls);
67
68            // Use project to align schema type
69            let mut out_fields = vec![];
70            let mut remained_group_key_offset = 0;
71            let mut removed_group_key_offset = new_group_key_len;
72            for &i in &group_key {
73                if new_group_key.contains(i) {
74                    out_fields.push(remained_group_key_offset);
75                    remained_group_key_offset += 1;
76                } else {
77                    out_fields.push(removed_group_key_offset);
78                    removed_group_key_offset += 1;
79                }
80            }
81            for i in group_key.len()..agg.base.schema().len() {
82                out_fields.push(i);
83            }
84            let new_agg = Agg::new(new_agg_calls, new_group_key, agg.input());
85
86            Some(LogicalProject::with_out_col_idx(new_agg.into(), out_fields.into_iter()).into())
87        } else {
88            None
89        }
90    }
91}
92
93impl AggGroupBySimplifyRule {
94    pub fn create() -> BoxedRule {
95        Box::new(AggGroupBySimplifyRule {})
96    }
97}
risingwave_frontend/optimizer/rule/agg_group_by_simplify_rule.rs

risingwave_frontend/optimizer/rule/
agg_group_by_simplify_rule.rs
