query/optimizer/
expr_depth.rs

1// Copyright 2023 Greptime Team
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::sync::Arc;
16
17use datafusion_common::config::ConfigOptions;
18use datafusion_common::tree_node::{TreeNode, TreeNodeRecursion};
19use datafusion_common::{Result, plan_err};
20use datafusion_expr::{Expr, LogicalPlan, Partitioning};
21use datafusion_optimizer::analyzer::AnalyzerRule;
22use promql::extension_plan::EmptyMetric;
23
24use crate::dist_plan::{MergeScanLogicalPlan, MergeSortLogicalPlan};
25use crate::range_select::plan::RangeSelect;
26
27// A conservative optimizer input limit, not a guarantee for every recursive operation
28// performed while parsing, planning, or dropping an arbitrarily deep query.
29const MAX_EXPR_DEPTH: usize = 128;
30
31/// Checks the analyzer output, including the SessionState EXPLAIN path.
32#[derive(Debug)]
33pub(crate) struct ExprDepthRule;
34
35impl AnalyzerRule for ExprDepthRule {
36    fn name(&self) -> &str {
37        "expr_depth"
38    }
39
40    fn analyze(&self, plan: LogicalPlan, _: &ConfigOptions) -> Result<LogicalPlan> {
41        validate_plan(&plan)?;
42        Ok(plan)
43    }
44}
45
46/// Rejects deep expressions before recursive optimizer traversal, cloning, or display.
47/// Both plan and expression traversal are iterative. Subquery Arcs are retained so
48/// expressions supplied temporarily by extension nodes can also be inspected safely.
49pub(crate) fn validate_plan(plan: &LogicalPlan) -> Result<()> {
50    let mut subqueries = Vec::new();
51    validate_plan_exprs(plan, &mut subqueries)?;
52    while let Some(plan) = subqueries.pop() {
53        validate_plan_exprs(&plan, &mut subqueries)?;
54    }
55    Ok(())
56}
57
58fn validate_plan_exprs(plan: &LogicalPlan, subqueries: &mut Vec<Arc<LogicalPlan>>) -> Result<()> {
59    let mut plans = vec![plan];
60    while let Some(plan) = plans.pop() {
61        let mut check = |expr: &Expr| {
62            validate_expr(expr, subqueries)?;
63            Ok(TreeNodeRecursion::Continue)
64        };
65        match plan {
66            // apply_expressions clones range ordering expressions and extension
67            // expressions. Borrow the built-in nodes that can contain deep trees.
68            LogicalPlan::Repartition(repartition)
69                if let Partitioning::Range(range) = &repartition.partitioning_scheme =>
70            {
71                for sort in range.ordering() {
72                    check(&sort.expr)?;
73                }
74            }
75            LogicalPlan::Extension(extension) => {
76                let node = extension.node.as_any();
77                if let Some(range) = node.downcast_ref::<RangeSelect>() {
78                    for expr in &range.range_expr {
79                        check(&expr.expr)?;
80                    }
81                    check(&range.time_expr)?;
82                    for expr in &range.by {
83                        check(expr)?;
84                    }
85                } else if let Some(sort) = node.downcast_ref::<MergeSortLogicalPlan>() {
86                    for expr in &sort.expr {
87                        check(&expr.expr)?;
88                    }
89                } else if let Some(metric) = node.downcast_ref::<EmptyMetric>() {
90                    if let Some(expr) = metric.expr() {
91                        check(expr)?;
92                    }
93                } else if let Some(scan) = node.downcast_ref::<MergeScanLogicalPlan>() {
94                    // MergeScan deliberately hides this input from optimizer visitors.
95                    plans.push(scan.input());
96                } else {
97                    // The remaining built-in nodes only expose column expressions.
98                    // External extensions own the safety of their expressions() implementation.
99                    plan.apply_expressions(check)?;
100                }
101            }
102            LogicalPlan::Subquery(subquery) => {
103                for expr in &subquery.outer_ref_columns {
104                    check(expr)?;
105                }
106            }
107            _ => {
108                plan.apply_expressions(check)?;
109            }
110        }
111        plans.extend(plan.inputs());
112    }
113    Ok(())
114}
115
116fn validate_expr(expr: &Expr, subqueries: &mut Vec<Arc<LogicalPlan>>) -> Result<()> {
117    let mut expressions = vec![(expr, 1)];
118    while let Some((expr, depth)) = expressions.pop() {
119        if depth > MAX_EXPR_DEPTH {
120            return plan_err!(
121                "Expression depth exceeds the maximum supported depth of {MAX_EXPR_DEPTH}"
122            );
123        }
124        let subquery = match expr {
125            Expr::Exists(exists) => Some(&exists.subquery),
126            Expr::ScalarSubquery(subquery) => Some(subquery),
127            Expr::InSubquery(in_subquery) => Some(&in_subquery.subquery),
128            Expr::SetComparison(comparison) => Some(&comparison.subquery),
129            _ => None,
130        };
131        if let Some(subquery) = subquery {
132            subqueries.push(Arc::clone(&subquery.subquery));
133            expressions.extend(subquery.outer_ref_columns.iter().map(|expr| (expr, 1)));
134        }
135        expr.apply_children(|child| {
136            expressions.push((child, depth + 1));
137            Ok(TreeNodeRecursion::Continue)
138        })?;
139    }
140    Ok(())
141}
142
143#[cfg(test)]
144mod tests {
145    use arrow_schema::{DataType, Field, Schema};
146    use datafusion_common::DFSchema;
147    use datafusion_expr::expr::{Exists, InSubquery, SetComparison, SetQuantifier};
148    use datafusion_expr::{EmptyRelation, Filter, Operator, Subquery, lit};
149
150    use super::*;
151
152    fn filter(predicate: Expr) -> LogicalPlan {
153        LogicalPlan::Filter(
154            Filter::try_new(
155                predicate,
156                Arc::new(LogicalPlan::EmptyRelation(EmptyRelation {
157                    produce_one_row: true,
158                    schema: Arc::new(
159                        DFSchema::try_from(Schema::new(vec![Field::new(
160                            "value",
161                            DataType::Boolean,
162                            false,
163                        )]))
164                        .unwrap(),
165                    ),
166                })),
167            )
168            .unwrap(),
169        )
170    }
171
172    #[test]
173    fn test_depth_in_subqueries_and_hidden_input() {
174        std::thread::Builder::new()
175            .stack_size(2 * 1024 * 1024)
176            .spawn(|| {
177                for depth in [MAX_EXPR_DEPTH, MAX_EXPR_DEPTH + 1, 512] {
178                    let mut predicate = lit(true);
179                    for _ in 1..depth {
180                        predicate = predicate.or(lit(false));
181                    }
182                    // MergeScan hides its input, and apply_children on an Expr
183                    // does not visit the plan embedded in its subquery.
184                    let plan = MergeScanLogicalPlan::new(filter(predicate), false)
185                        .into_logical_plan();
186                    let subquery = Subquery {
187                        subquery: Arc::new(plan),
188                        outer_ref_columns: vec![],
189                        spans: Default::default(),
190                    };
191                    for expr in [
192                        Expr::Exists(Exists::new(subquery.clone(), false)),
193                        Expr::ScalarSubquery(subquery.clone()),
194                        Expr::InSubquery(InSubquery::new(
195                            Box::new(lit(true)),
196                            subquery.clone(),
197                            false,
198                        )),
199                        Expr::SetComparison(SetComparison::new(
200                            Box::new(lit(true)),
201                            subquery.clone(),
202                            Operator::Eq,
203                            SetQuantifier::Any,
204                        )),
205                    ] {
206                        let result = validate_plan(&filter(expr));
207                        if depth <= MAX_EXPR_DEPTH {
208                            result.unwrap();
209                        } else {
210                            assert!(matches!(
211                                result,
212                                Err(datafusion_common::DataFusionError::Plan(message))
213                                    if message == "Expression depth exceeds the maximum supported depth of 128"
214                            ));
215                        }
216                    }
217                }
218            })
219            .unwrap()
220            .join()
221            .unwrap();
222    }
223}
query/optimizer/expr_depth.rs

query/optimizer/
expr_depth.rs
