risingwave_frontend/optimizer/plan_visitor/
apply_visitor.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 super::{DefaultBehavior, LogicalPlanVisitor, Merge};
16use crate::error::{ErrorCode, RwError};
17use crate::expr::{CorrelatedInputRef, ExprVisitor, FunctionCallWithLambda, Subquery};
18use crate::optimizer::plan_node::{
19    LogicalApply, LogicalPlanRef as PlanRef, PlanTreeNodeBinary, VisitExprsRecursive,
20};
21use crate::optimizer::plan_visitor::PlanVisitor;
22
23pub struct HasMaxOneRowApply();
24
25impl LogicalPlanVisitor for HasMaxOneRowApply {
26    type Result = bool;
27
28    type DefaultBehavior = impl DefaultBehavior<Self::Result>;
29
30    fn default_behavior() -> Self::DefaultBehavior {
31        Merge(|a, b| a | b)
32    }
33
34    fn visit_logical_apply(&mut self, plan: &LogicalApply) -> bool {
35        plan.max_one_row() | self.visit(plan.left()) | self.visit(plan.right())
36    }
37}
38
39#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Clone, Copy, Default)]
40enum CheckResult {
41    #[default]
42    Ok,
43    UnplannedSubquery,
44    UnresolvedCorrelatedInputRef,
45    CannotBeUnnested,
46    MoreThanOneRow,
47}
48
49impl From<CheckResult> for Result<(), RwError> {
50    fn from(val: CheckResult) -> Self {
51        let msg = match val {
52            CheckResult::Ok => return Ok(()),
53            CheckResult::UnplannedSubquery => "Subquery is not supported in this context.",
54            CheckResult::UnresolvedCorrelatedInputRef => {
55                "Subquery cannot be unnested: a correlated column reference is not resolved."
56            }
57            CheckResult::CannotBeUnnested => "Subquery cannot be unnested.",
58            CheckResult::MoreThanOneRow => "Scalar subquery might produce more than one row.",
59        };
60
61        Err(ErrorCode::InternalError(msg.to_owned()).into())
62    }
63}
64
65#[derive(Default)]
66pub struct CheckApplyElimination {
67    result: CheckResult,
68}
69
70impl LogicalPlanVisitor for CheckApplyElimination {
71    type Result = ();
72
73    type DefaultBehavior = impl DefaultBehavior<Self::Result>;
74
75    fn default_behavior() -> Self::DefaultBehavior {
76        Merge(std::cmp::max)
77    }
78
79    fn visit_logical_apply(&mut self, plan: &LogicalApply) {
80        // If there's a runtime max-one-row check on the right side, it's likely to be the
81        // reason for the failed unnesting. Report users with a more precise error message.
82        if plan.right().as_logical_max_one_row().is_some() {
83            self.result = CheckResult::MoreThanOneRow;
84        } else {
85            self.result = CheckResult::CannotBeUnnested;
86        }
87    }
88}
89
90/// Finds `Subquery` and `CorrelatedInputRef` expressions, which cannot be serialized. The planner
91/// must turn every subquery into an `Apply`, and unnesting must resolve every correlated reference.
92#[derive(Default)]
93struct SubqueryExprFinder {
94    has_subquery: bool,
95    has_correlated_input_ref: bool,
96}
97
98impl SubqueryExprFinder {
99    fn find(plan: &PlanRef) -> Self {
100        let mut finder = Self::default();
101        plan.visit_exprs_recursive(&mut finder);
102        finder
103    }
104}
105
106impl ExprVisitor for SubqueryExprFinder {
107    fn visit_subquery(&mut self, _: &Subquery) {
108        self.has_subquery = true;
109    }
110
111    fn visit_correlated_input_ref(&mut self, _: &CorrelatedInputRef) {
112        self.has_correlated_input_ref = true;
113    }
114
115    // The default implementation skips the lambda body, which is also serialized.
116    fn visit_function_call_with_lambda(&mut self, func_call: &FunctionCallWithLambda) {
117        func_call
118            .inputs_with_lambda_arg()
119            .for_each(|expr| self.visit_expr(expr));
120    }
121}
122
123#[easy_ext::ext(PlanCheckApplyEliminationExt)]
124impl PlanRef {
125    /// Checks that the planner has turned every subquery expression into an `Apply`. A subquery
126    /// left in an expression, e.g., in a lambda body, can be neither hashed nor unnested.
127    pub fn check_subquery_planned(&self) -> Result<(), RwError> {
128        if SubqueryExprFinder::find(self).has_subquery {
129            CheckResult::UnplannedSubquery.into()
130        } else {
131            Ok(())
132        }
133    }
134
135    /// Checks if all `LogicalApply` nodes in the plan have been eliminated, that is,
136    /// subqueries are successfully unnested, and that no correlated column reference is left.
137    pub fn check_apply_elimination(&self) -> Result<(), RwError> {
138        let mut visitor = CheckApplyElimination::default();
139        visitor.visit(self.clone());
140        let mut result = visitor.result;
141        if SubqueryExprFinder::find(self).has_correlated_input_ref {
142            result = result.max(CheckResult::UnresolvedCorrelatedInputRef);
143        }
144        result.into()
145    }
146}
risingwave_frontend/optimizer/plan_visitor/apply_visitor.rs

risingwave_frontend/optimizer/plan_visitor/
apply_visitor.rs
