mz_sql/plan/
explain.rs

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

mz_sql/plan/
explain.rs
