risingwave_frontend/optimizer/rule/
grouping_sets_to_expand_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 itertools::Itertools;
16use risingwave_common::types::DataType;
17use risingwave_common::util::column_index_mapping::ColIndexMapping;
18use risingwave_expr::aggregate::{AggType, PbAggKind};
19
20use super::prelude::{PlanRef, *};
21use crate::expr::{Expr, ExprImpl, ExprType, FunctionCall, InputRef};
22use crate::optimizer::plan_node::generic::{Agg, GenericPlanNode, GenericPlanRef};
23use crate::optimizer::plan_node::*;
24pub struct GroupingSetsToExpandRule {}
25
26impl GroupingSetsToExpandRule {
27    pub fn create() -> BoxedRule {
28        Box::new(Self {})
29    }
30}
31
32impl Rule<Logical> for GroupingSetsToExpandRule {
33    fn apply(&self, plan: PlanRef) -> Option<PlanRef> {
34        let agg: &LogicalAgg = plan.as_logical_agg()?;
35        if agg.grouping_sets().is_empty() {
36            return None;
37        }
38        let (old_agg_calls, old_group_keys, grouping_sets, input, enable_two_phase) =
39            agg.clone().decompose();
40
41        let old_input_schema_len = input.schema().len();
42        let flag_col_idx = old_group_keys.len();
43
44        let column_subset = grouping_sets
45            .iter()
46            .map(|set| set.indices().collect_vec())
47            .collect_vec();
48
49        let expand = LogicalExpand::create(input, column_subset.clone());
50        let new_group_keys = {
51            let mut k = old_group_keys.clone();
52            // Add the expand flag.
53            k.extend(std::iter::once(expand.schema().len() - 1));
54            k
55        };
56
57        // Map from old input ref to expanded input (`LogicalExpand` prepends the same number of fields
58        // as expanded ones with NULLs before the real input fields).
59        let mut input_col_change =
60            ColIndexMapping::with_shift_offset(old_input_schema_len, old_input_schema_len as isize);
61
62        // Grouping agg calls need to be transformed into a project expression, and other agg calls
63        // need to shift their `input_ref`.
64        let mut project_agg_call_exprs = vec![];
65        let mut new_agg_calls = vec![];
66        for agg_call in old_agg_calls {
67            // Deal with grouping agg call for grouping sets.
68            if matches!(agg_call.agg_type, AggType::Builtin(PbAggKind::Grouping)) {
69                let mut grouping_values = vec![];
70                let args = agg_call
71                    .inputs
72                    .iter()
73                    .map(|input_ref| input_ref.index)
74                    .collect_vec();
75                for subset in &column_subset {
76                    let mut value = 0;
77                    for arg in &args {
78                        value <<= 1;
79                        if !subset.contains(arg) {
80                            value += 1;
81                        }
82                    }
83                    grouping_values.push(value);
84                }
85
86                let mut case_inputs = vec![];
87                for (i, grouping_value) in grouping_values.into_iter().enumerate() {
88                    let condition = ExprImpl::FunctionCall(
89                        FunctionCall::new_unchecked(
90                            ExprType::Equal,
91                            vec![
92                                ExprImpl::literal_bigint(i as i64),
93                                ExprImpl::InputRef(
94                                    InputRef::new(flag_col_idx, DataType::Int64).into(),
95                                ),
96                            ],
97                            DataType::Boolean,
98                        )
99                        .into(),
100                    );
101                    let value = ExprImpl::literal_int(grouping_value);
102                    case_inputs.push(condition);
103                    case_inputs.push(value);
104                }
105
106                let case_expr = ExprImpl::FunctionCall(
107                    FunctionCall::new_unchecked(ExprType::Case, case_inputs, DataType::Int32)
108                        .into(),
109                );
110                project_agg_call_exprs.push(case_expr);
111            } else {
112                let mut new_agg_call = agg_call;
113                // Shift agg_call to the original input columns
114                new_agg_call.inputs.iter_mut().for_each(|i| {
115                    let new_i = input_col_change.map(i.index());
116                    assert_eq!(expand.schema()[new_i].data_type(), i.return_type());
117                    *i = InputRef::new(new_i, i.return_type());
118                });
119                new_agg_call.order_by.iter_mut().for_each(|o| {
120                    o.column_index = input_col_change.map(o.column_index);
121                });
122                new_agg_call.filter = new_agg_call.filter.rewrite_expr(&mut input_col_change);
123                project_agg_call_exprs.push(ExprImpl::InputRef(
124                    InputRef::new(
125                        new_group_keys.len() + new_agg_calls.len(),
126                        new_agg_call.return_type.clone(),
127                    )
128                    .into(),
129                ));
130                new_agg_calls.push(new_agg_call);
131            }
132        }
133
134        let new_agg =
135            Agg::new(new_agg_calls, new_group_keys, expand).with_enable_two_phase(enable_two_phase);
136        let project_exprs = (0..old_group_keys.len())
137            .map(|i| ExprImpl::InputRef(InputRef::new(i, new_agg.schema()[i].data_type()).into()))
138            .chain(project_agg_call_exprs)
139            .collect();
140
141        let project = LogicalProject::new(new_agg.into(), project_exprs);
142
143        Some(project.into())
144    }
145}
risingwave_frontend/optimizer/rule/grouping_sets_to_expand_rule.rs

risingwave_frontend/optimizer/rule/
grouping_sets_to_expand_rule.rs
