risingwave_frontend/optimizer/rule/
translate_apply_rule.rs

1// Copyright 2022 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 std::collections::HashMap;
16
17use risingwave_common::types::DataType;
18use risingwave_common::util::iter_util::ZipEqFast;
19use risingwave_pb::plan_common::JoinType;
20
21use super::prelude::{PlanRef, *};
22use super::{ApplyResult, FallibleRule};
23use crate::error::ErrorCode;
24use crate::expr::{ExprImpl, ExprType, FunctionCall, ImpureAnalyzer, InputRef};
25use crate::optimizer::plan_node::generic::{Agg, GenericPlanRef};
26use crate::optimizer::plan_node::{
27    LogicalApply, LogicalJoin, LogicalProject, LogicalScan, LogicalShare, PlanTreeNodeBinary,
28    PlanTreeNodeUnary, VisitExprsRecursive,
29};
30use crate::utils::{ColIndexMapping, Condition};
31
32/// General Unnesting based on the paper Unnesting Arbitrary Queries:
33/// Translate the apply into a canonical form.
34///
35/// Before:
36///
37/// ```text
38///     LogicalApply
39///    /            \
40///  LHS           RHS
41/// ```
42///
43/// After:
44///
45/// ```text
46///      LogicalJoin
47///    /            \
48///  LHS        LogicalApply
49///             /           \
50///          Domain         RHS
51/// ```
52pub struct TranslateApplyRule {
53    enable_share_plan: bool,
54}
55
56impl FallibleRule<Logical> for TranslateApplyRule {
57    fn apply(&self, plan: PlanRef) -> ApplyResult<PlanRef> {
58        let apply: &LogicalApply = plan.as_logical_apply()?;
59        if apply.translated() {
60            return ApplyResult::NotApplicable;
61        }
62        let mut left: PlanRef = apply.left();
63        let right: PlanRef = apply.right();
64        let apply_left_len = left.schema().len();
65        let correlated_indices = apply.correlated_indices();
66
67        let mut index_mapping =
68            ColIndexMapping::new(vec![None; apply_left_len], correlated_indices.len());
69        let mut data_types = HashMap::new();
70        let mut index = 0;
71
72        // First try to rewrite the left side of the apply.
73        // TODO: remove the rewrite and always use the general way to calculate the domain
74        //      after we support DAG.
75        let rewritten_left = self
76            .rewrite(
77                &left,
78                correlated_indices.clone(),
79                0,
80                &mut index_mapping,
81                &mut data_types,
82                &mut index,
83            )
84            .filter(|plan| {
85                // Rewriting removes Filter nodes, but a Scan can still contain a pushed-down
86                // impure predicate. Such a domain must share the outer evaluation as well.
87                let mut impurity = ImpureAnalyzer::default();
88                plan.visit_exprs_recursive(&mut impurity);
89                impurity.impure_expr_desc().is_none()
90            });
91        let domain: PlanRef = if let Some(rewritten_left) = rewritten_left {
92            // This `LogicalProject` is used to make sure that after `LogicalApply`'s left was
93            // rewritten, the new index of `correlated_index` is always at its position in
94            // `correlated_indices`.
95            let exprs = correlated_indices
96                .clone()
97                .into_iter()
98                .enumerate()
99                .map(|(i, correlated_index)| {
100                    let index = index_mapping.map(correlated_index);
101                    let data_type = rewritten_left.schema().fields()[index].data_type.clone();
102                    index_mapping.put(correlated_index, Some(i));
103                    InputRef::new(index, data_type).into()
104                })
105                .collect();
106            let project = LogicalProject::create(rewritten_left, exprs);
107            let distinct = Agg::new(vec![], (0..project.schema().len()).collect(), project);
108            distinct.into()
109        } else {
110            // The left side cannot be safely rewritten. Use the general way to calculate
111            // the domain: Distinct + Project + The Left of Apply.
112
113            if !self.enable_share_plan {
114                let mut impurity = ImpureAnalyzer::default();
115                left.visit_exprs_recursive(&mut impurity);
116                if let Some(expr) = impurity.impure_expr_desc() {
117                    // Without Share, the domain and the final join evaluate the outer plan
118                    // independently and may produce different correlated keys.
119                    return ApplyResult::Err(
120                        ErrorCode::NotSupported(
121                            format!(
122                                "correlated subquery would evaluate the impure outer expression ({expr}) more than once"
123                            ),
124                            "Store the outer query result in a table before running this correlated subquery."
125                                .into(),
126                        )
127                        .into(),
128                    );
129                }
130                if Self::has_row_limit(&left) {
131                    return ApplyResult::Err(
132                        ErrorCode::NotSupported(
133                            "correlated subquery would evaluate an outer LIMIT or TopN more than once"
134                                .into(),
135                            "Store the outer query result in a table before running this correlated subquery."
136                                .into(),
137                        )
138                        .into(),
139                    );
140                }
141            }
142
143            // Use Share
144            left = if self.enable_share_plan {
145                let logical_share = LogicalShare::new(left);
146                logical_share.into()
147            } else {
148                left
149            };
150
151            let exprs = correlated_indices
152                .clone()
153                .into_iter()
154                .map(|correlated_index| {
155                    let data_type = left.schema().fields()[correlated_index].data_type.clone();
156                    InputRef::new(correlated_index, data_type).into()
157                })
158                .collect();
159            let project = LogicalProject::create(left.clone(), exprs);
160            let distinct = Agg::new(vec![], (0..project.schema().len()).collect(), project);
161            distinct.into()
162        };
163
164        let eq_predicates = correlated_indices
165            .into_iter()
166            .enumerate()
167            .map(|(i, correlated_index)| {
168                let shifted_index = i + apply_left_len;
169                let data_type = domain.schema().fields()[i].data_type.clone();
170                let left = InputRef::new(correlated_index, data_type.clone());
171                let right = InputRef::new(shifted_index, data_type);
172                // use null-safe equal
173                FunctionCall::new_unchecked(
174                    ExprType::IsNotDistinctFrom,
175                    vec![left.into(), right.into()],
176                    DataType::Boolean,
177                )
178                .into()
179            })
180            .collect::<Vec<ExprImpl>>();
181
182        let new_apply = apply.clone_with_left_right(left, right);
183        let new_node = new_apply.translate_apply(domain, eq_predicates);
184        ApplyResult::Ok(new_node)
185    }
186}
187
188impl TranslateApplyRule {
189    pub fn create(enable_share_plan: bool) -> BoxedRule {
190        Box::new(TranslateApplyRule { enable_share_plan })
191    }
192
193    fn has_row_limit(plan: &PlanRef) -> bool {
194        // Even pure plans can select different rows: LIMIT has no ordering guarantee,
195        // and TopN may have ties. Do not duplicate them without proving determinism.
196        plan.as_logical_limit().is_some()
197            || plan.as_logical_top_n().is_some()
198            || plan.inputs().iter().any(Self::has_row_limit)
199    }
200
201    /// Rewrite `LogicalApply`'s left according to `correlated_indices`.
202    ///
203    /// The rewritten domain may be a superset of the correlated keys. In particular, filters
204    /// and row limits can be removed, so that sampling is only evaluated on the outer side.
205    /// Only traverse projections and row limits when Share is unavailable; otherwise, share
206    /// the outer result instead of potentially expanding the domain to the entire input.
207    fn rewrite(
208        &self,
209        plan: &PlanRef,
210        correlated_indices: Vec<usize>,
211        offset: usize,
212        index_mapping: &mut ColIndexMapping,
213        data_types: &mut HashMap<usize, DataType>,
214        index: &mut usize,
215    ) -> Option<PlanRef> {
216        if let Some(join) = plan.as_logical_join() {
217            self.rewrite_join(
218                join,
219                correlated_indices,
220                offset,
221                index_mapping,
222                data_types,
223                index,
224            )
225        } else if let Some(apply) = plan.as_logical_apply() {
226            self.rewrite_apply(
227                apply,
228                correlated_indices,
229                offset,
230                index_mapping,
231                data_types,
232                index,
233            )
234        } else if let Some(scan) = plan.as_logical_scan() {
235            Self::rewrite_scan(
236                scan,
237                correlated_indices,
238                offset,
239                index_mapping,
240                data_types,
241                index,
242            )
243        } else if let Some(filter) = plan.as_logical_filter() {
244            self.rewrite(
245                &filter.input(),
246                correlated_indices,
247                offset,
248                index_mapping,
249                data_types,
250                index,
251            )
252        } else if self.enable_share_plan {
253            None
254        } else if let Some(limit) = plan.as_logical_limit() {
255            self.rewrite(
256                &limit.input(),
257                correlated_indices,
258                offset,
259                index_mapping,
260                data_types,
261                index,
262            )
263        } else if let Some(top_n) = plan.as_logical_top_n() {
264            self.rewrite(
265                &top_n.input(),
266                correlated_indices,
267                offset,
268                index_mapping,
269                data_types,
270                index,
271            )
272        } else if let Some(project) = plan.as_logical_project() {
273            // Only the correlated columns matter. An unused random() ordering expression
274            // must not prevent us from deriving the domain from the underlying scan.
275            let input_indices = correlated_indices
276                .iter()
277                .map(|&i| project.exprs()[i].as_input_ref().map(|r| r.index()))
278                .collect::<Option<Vec<_>>>()?;
279            let mut input_mapping =
280                ColIndexMapping::empty(project.input().schema().len(), index_mapping.target_size());
281            let rewritten = self.rewrite(
282                &project.input(),
283                input_indices.clone(),
284                0,
285                &mut input_mapping,
286                data_types,
287                index,
288            )?;
289            for (output, input) in correlated_indices.into_iter().zip_eq_fast(input_indices) {
290                index_mapping.put(output + offset, Some(input_mapping.map(input)));
291            }
292            Some(rewritten)
293        } else {
294            // TODO: better to return an error.
295            None
296        }
297    }
298
299    fn rewrite_join(
300        &self,
301        join: &LogicalJoin,
302        required_col_idx: Vec<usize>,
303        mut offset: usize,
304        index_mapping: &mut ColIndexMapping,
305        data_types: &mut HashMap<usize, DataType>,
306        index: &mut usize,
307    ) -> Option<PlanRef> {
308        // TODO: Do we need to take the `on` into account?
309        let left_len = join.left().schema().len();
310        let (left_idxs, right_idxs): (Vec<_>, Vec<_>) = required_col_idx
311            .into_iter()
312            .partition(|idx| *idx < left_len);
313        let mut rewrite =
314            |plan: PlanRef, mut indices: Vec<usize>, is_right: bool| -> Option<PlanRef> {
315                if is_right {
316                    indices.iter_mut().for_each(|index| *index -= left_len);
317                    offset += left_len;
318                }
319                self.rewrite(&plan, indices, offset, index_mapping, data_types, index)
320            };
321        match (left_idxs.is_empty(), right_idxs.is_empty()) {
322            (true, false) => {
323                // Only accept join which doesn't generate null columns.
324                match join.join_type() {
325                    JoinType::Inner
326                    | JoinType::LeftSemi
327                    | JoinType::RightSemi
328                    | JoinType::LeftAnti
329                    | JoinType::RightAnti
330                    | JoinType::RightOuter
331                    | JoinType::AsofInner => rewrite(join.right(), right_idxs, true),
332                    JoinType::LeftOuter | JoinType::FullOuter | JoinType::AsofLeftOuter => None,
333                    JoinType::Unspecified => unreachable!(),
334                }
335            }
336            (false, true) => {
337                // Only accept join which doesn't generate null columns.
338                match join.join_type() {
339                    JoinType::Inner
340                    | JoinType::LeftSemi
341                    | JoinType::RightSemi
342                    | JoinType::LeftAnti
343                    | JoinType::RightAnti
344                    | JoinType::LeftOuter
345                    | JoinType::AsofInner
346                    | JoinType::AsofLeftOuter => rewrite(join.left(), left_idxs, false),
347                    JoinType::RightOuter | JoinType::FullOuter => None,
348                    JoinType::Unspecified => unreachable!(),
349                }
350            }
351            (false, false) => {
352                // Only accept join which doesn't generate null columns.
353                match join.join_type() {
354                    JoinType::Inner
355                    | JoinType::LeftSemi
356                    | JoinType::RightSemi
357                    | JoinType::LeftAnti
358                    | JoinType::RightAnti
359                    | JoinType::AsofInner => {
360                        let left = rewrite(join.left(), left_idxs, false)?;
361                        let right = rewrite(join.right(), right_idxs, true)?;
362                        let new_join =
363                            LogicalJoin::new(left, right, join.join_type(), Condition::true_cond());
364                        Some(new_join.into())
365                    }
366                    JoinType::LeftOuter
367                    | JoinType::RightOuter
368                    | JoinType::FullOuter
369                    | JoinType::AsofLeftOuter => None,
370                    JoinType::Unspecified => unreachable!(),
371                }
372            }
373            _ => None,
374        }
375    }
376
377    /// ```text
378    ///             LogicalApply
379    ///            /            \
380    ///     LogicalApply       RHS1
381    ///    /            \
382    ///  LHS           RHS2
383    /// ```
384    ///
385    /// A common structure of multi scalar subqueries is a chain of `LogicalApply`. To avoid exponential growth of the domain operator, we need to rewrite the apply and try to simplify the domain as much as possible.
386    /// We use a top-down apply order to rewrite the apply, so that we don't need to handle operator like project and aggregation generated by the domain calculation.
387    /// As a cost, we need to add a flag `translated` to the apply operator to remind `translate_apply_rule` that the apply has been translated.
388    fn rewrite_apply(
389        &self,
390        apply: &LogicalApply,
391        required_col_idx: Vec<usize>,
392        offset: usize,
393        index_mapping: &mut ColIndexMapping,
394        data_types: &mut HashMap<usize, DataType>,
395        index: &mut usize,
396    ) -> Option<PlanRef> {
397        // TODO: Do we need to take the `on` into account?
398        let left_len = apply.left().schema().len();
399        let (left_idxs, right_idxs): (Vec<_>, Vec<_>) = required_col_idx
400            .into_iter()
401            .partition(|idx| *idx < left_len);
402        if !left_idxs.is_empty() && right_idxs.is_empty() {
403            // Deal with multi scalar subqueries
404            match apply.join_type() {
405                JoinType::Inner
406                | JoinType::LeftSemi
407                | JoinType::LeftAnti
408                | JoinType::LeftOuter
409                | JoinType::AsofInner
410                | JoinType::AsofLeftOuter => {
411                    let plan = apply.left();
412                    self.rewrite(&plan, left_idxs, offset, index_mapping, data_types, index)
413                }
414                JoinType::RightOuter
415                | JoinType::RightAnti
416                | JoinType::RightSemi
417                | JoinType::FullOuter => None,
418                JoinType::Unspecified => unreachable!(),
419            }
420        } else {
421            None
422        }
423    }
424
425    fn rewrite_scan(
426        scan: &LogicalScan,
427        required_col_idx: Vec<usize>,
428        offset: usize,
429        index_mapping: &mut ColIndexMapping,
430        data_types: &mut HashMap<usize, DataType>,
431        index: &mut usize,
432    ) -> Option<PlanRef> {
433        for i in &required_col_idx {
434            let correlated_index = *i + offset;
435            index_mapping.put(correlated_index, Some(*index));
436            data_types.insert(
437                correlated_index,
438                scan.schema().fields()[*i].data_type.clone(),
439            );
440            *index += 1;
441        }
442
443        Some(scan.clone_with_output_indices(required_col_idx).into())
444    }
445}
risingwave_frontend/optimizer/rule/translate_apply_rule.rs

risingwave_frontend/optimizer/rule/
translate_apply_rule.rs
