mz_sql/plan/
explain.rs

1// Copyright Materialize, Inc. and contributors. All rights reserved.
2//
3// Use of this software is governed by the Business Source License
4// included in the LICENSE file.
5//
6// As of the Change Date specified in that file, in accordance with
7// the Business Source License, use of this software will be governed
8// by the Apache License, Version 2.0.
9
10//! `EXPLAIN` support for structures defined in this crate.
11
12use std::panic::AssertUnwindSafe;
13
14use mz_expr::explain::{ExplainContext, ExplainSinglePlan};
15use mz_expr::visit::{Visit, VisitChildren};
16use mz_expr::{Id, LocalId};
17use mz_ore::stack::RecursionLimitError;
18use mz_repr::SqlRelationType;
19use mz_repr::explain::{AnnotatedPlan, Explain, ExplainError, ScalarOps, UnsupportedFormat};
20
21use crate::plan::{HirRelationExpr, HirScalarExpr};
22
23mod text;
24
25impl<'a> Explain<'a> for HirRelationExpr {
26    type Context = ExplainContext<'a>;
27
28    type Text = ExplainSinglePlan<'a, HirRelationExpr>;
29
30    type Json = ExplainSinglePlan<'a, HirRelationExpr>;
31
32    type Dot = UnsupportedFormat;
33
34    fn explain_text(&'a mut self, context: &'a Self::Context) -> Result<Self::Text, ExplainError> {
35        self.as_explain_single_plan(context)
36    }
37
38    fn explain_json(&'a mut self, context: &'a Self::Context) -> Result<Self::Json, ExplainError> {
39        self.as_explain_single_plan(context)
40    }
41}
42
43impl<'a> HirRelationExpr {
44    fn as_explain_single_plan(
45        &'a mut self,
46        context: &'a ExplainContext<'a>,
47    ) -> Result<ExplainSinglePlan<'a, HirRelationExpr>, ExplainError> {
48        // unless raw plans are explicitly requested
49        // ensure that all nested subqueries are wrapped in Let blocks by calling
50        // `normalize_subqueries`
51        if !context.config.raw_plans {
52            mz_ore::panic::catch_unwind_str(AssertUnwindSafe(|| {
53                normalize_subqueries(self).map_err(|e| e.into())
54            }))
55            .unwrap_or_else(|panic| {
56                // A panic during optimization is always a bug; log an error so we learn about it.
57                // TODO(teskje): collect and log a backtrace from the panic site
58                tracing::error!("caught a panic during `normalize_subqueries`: {panic}");
59
60                let msg = format!("unexpected panic during `normalize_subqueries`: {panic}");
61                Err(ExplainError::UnknownError(msg))
62            })?
63        }
64
65        // TODO: use config values to infer requested
66        // plan annotations
67        let plan = AnnotatedPlan {
68            plan: self,
69            annotations: Default::default(),
70        };
71        Ok(ExplainSinglePlan { context, plan })
72    }
73}
74
75/// Normalize the way subqueries appear in [`HirScalarExpr::Exists`]
76/// or [`HirScalarExpr::Select`] variants.
77///
78/// After the transform is applied, subqueries are pulled as a value in
79/// a let binding enclosing the [`HirRelationExpr`] parent of the
80/// [`HirScalarExpr::Exists`] or [`HirScalarExpr::Select`] where the
81/// subquery appears, and the corresponding variant references the
82/// new binding with a [`HirRelationExpr::Get`].
83pub fn normalize_subqueries<'a>(expr: &'a mut HirRelationExpr) -> Result<(), RecursionLimitError> {
84    // A helper struct to represent accumulated `$local_id = $subquery`
85    // bindings that need to be installed in `let ... in $expr` nodes
86    // that wrap their parent $expr.
87    struct Binding {
88        local_id: LocalId,
89        subquery: HirRelationExpr,
90    }
91
92    // Context for the transformation
93    // - a stack of bindings
94    let mut bindings = Vec::<Binding>::new();
95    // - a generator of fresh local ids
96    let mut id_gen = id_gen(expr)?.peekable();
97
98    // Grow the `bindings` stack by collecting subqueries appearing in
99    // one of the HirScalarExpr children at the given HirRelationExpr.
100    // As part of this, the subquery is replaced by a `Get(id)` for a
101    // fresh local id.
102    let mut collect_subqueries = |expr: &mut HirRelationExpr, bindings: &mut Vec<Binding>| {
103        expr.try_visit_mut_children(|expr: &mut HirScalarExpr| {
104            use HirRelationExpr::Get;
105            use HirScalarExpr::{Exists, Select};
106            expr.visit_mut_post(&mut |expr: &mut HirScalarExpr| match expr {
107                Exists(expr, _) | Select(expr, _) => match expr.as_mut() {
108                    Get { .. } => (),
109                    expr => {
110                        // generate fresh local id
111                        let local_id = id_gen.next().unwrap();
112                        // generate a `Get(local_id)` to be used as a subquery replacement
113                        let mut subquery = Get {
114                            id: Id::Local(local_id.clone()),
115                            typ: SqlRelationType::empty(), // TODO (aalexandrov)
116                        };
117                        // swap the current subquery with the replacement
118                        std::mem::swap(expr, &mut subquery);
119                        // push a new $local_id = $subquery binding for a wrapping Let { ... }
120                        bindings.push(Binding { local_id, subquery });
121                    }
122                },
123                _ => (),
124            })
125        })
126    };
127
128    // Drain the `bindings` stack by wrapping the given `HirRelationExpr` with
129    // a sequence of `Let { ... }` nodes, one for each binding.
130    let insert_let_bindings = |expr: &mut HirRelationExpr, bindings: &mut Vec<Binding>| {
131        for binding in bindings.drain(..) {
132            let name = format!("subquery-{}", Into::<u64>::into(&binding.local_id));
133            let id = binding.local_id;
134            let value = Box::new(binding.subquery);
135            let body = Box::new(expr.take());
136            *expr = HirRelationExpr::Let {
137                name,
138                id,
139                value,
140                body,
141            }
142        }
143    };
144
145    expr.try_visit_mut_post(&mut |expr: &mut HirRelationExpr| {
146        // first grow bindings stack
147        collect_subqueries(expr, &mut bindings)?;
148        // then drain bindings stack
149        insert_let_bindings(expr, &mut bindings);
150        // done!
151        Ok(())
152    })
153}
154
155// Create an [`Iterator`] for [`LocalId`] values that are guaranteed to be
156// fresh within the scope of the given [`HirRelationExpr`].
157fn id_gen(
158    expr: &HirRelationExpr,
159) -> Result<impl Iterator<Item = LocalId> + use<>, RecursionLimitError> {
160    let mut max_id = 0_u64;
161
162    expr.visit_pre(&mut |expr| {
163        match expr {
164            HirRelationExpr::Let { id, .. } => max_id = std::cmp::max(max_id, id.into()),
165            _ => (),
166        };
167    })?;
168
169    Ok((max_id + 1..).map(LocalId::new))
170}
171
172impl ScalarOps for HirScalarExpr {
173    fn match_col_ref(&self) -> Option<usize> {
174        match self {
175            HirScalarExpr::Column(c, _name) if c.level == 0 => Some(c.column),
176            _ => None,
177        }
178    }
179
180    fn references(&self, column: usize) -> bool {
181        match self {
182            HirScalarExpr::Column(c, _name) => c.column == column && c.level == 0,
183            _ => false,
184        }
185    }
186}
mz_sql/plan/explain.rs

mz_sql/plan/
explain.rs
