Skip to main content

mz_sql/plan/
transform_ast.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//! Transformations of SQL ASTs.
11//!
12//! Most query optimizations are performed by the dataflow layer, but some
13//! are much easier to perform in SQL. Someday, we'll want our own SQL IR,
14//! but for now we just use the parser's AST directly.
15
16use itertools::Itertools;
17use mz_ore::id_gen::IdGen;
18use mz_ore::stack::{CheckedRecursion, RecursionGuard};
19use mz_repr::namespaces::{MZ_CATALOG_SCHEMA, MZ_UNSAFE_SCHEMA, PG_CATALOG_SCHEMA};
20use mz_sql_parser::ast::visit_mut::{self, VisitMut, VisitMutNode};
21use mz_sql_parser::ast::{
22    Expr, Function, FunctionArgs, HomogenizingFunction, Ident, IsExprConstruct, Op, OrderByExpr,
23    Query, Select, SelectItem, TableAlias, TableFactor, TableWithJoins, Value, WindowSpec,
24};
25use mz_sql_parser::ident;
26
27use crate::names::{Aug, PartialItemName, ResolvedDataType, ResolvedItemName};
28use crate::plan::{PlanError, StatementContext};
29use crate::{ORDINALITY_COL_NAME, normalize};
30
31pub(crate) fn transform<N>(scx: &StatementContext, node: &mut N) -> Result<(), PlanError>
32where
33    N: for<'a> VisitMutNode<'a, Aug>,
34{
35    let mut func_rewriter = FuncRewriter::new(scx);
36    node.visit_mut(&mut func_rewriter);
37    func_rewriter.status?;
38
39    let mut desugarer = Desugarer::new(scx);
40    node.visit_mut(&mut desugarer);
41    desugarer.status
42}
43
44// Transforms various functions to forms that are more easily handled by the
45// planner.
46//
47// Specifically:
48//
49//   * Rewrites the `mod` function to the `%` binary operator, so the modulus
50//     code only needs to handle the operator form.
51//
52//   * Rewrites the `nullif` function to a `CASE` statement, to reuse the code
53//     for planning equality of datums.
54//
55//   * Rewrites `avg(col)` to `sum(col) / count(col)`, so that we can pretend
56//     the `avg` aggregate function doesn't exist from here on out. This also
57//     has the nice side effect of reusing the division planning logic, which
58//     is not trivial for some types, like decimals.
59//
60//   * Rewrites the suite of standard deviation and variance functions in a
61//     manner similar to `avg`.
62//
63// TODO(sploiselle): rewrite these in terms of func::sql_op!
64struct FuncRewriter<'a> {
65    scx: &'a StatementContext<'a>,
66    status: Result<(), PlanError>,
67    rewriting_table_factor: bool,
68}
69
70impl<'a> FuncRewriter<'a> {
71    fn new(scx: &'a StatementContext<'a>) -> FuncRewriter<'a> {
72        FuncRewriter {
73            scx,
74            status: Ok(()),
75            rewriting_table_factor: false,
76        }
77    }
78
79    fn resolve_known_valid_data_type(&self, name: &PartialItemName) -> ResolvedDataType {
80        let item = self
81            .scx
82            .catalog
83            .resolve_type(name)
84            .expect("data type known to be valid");
85        let full_name = self.scx.catalog.resolve_full_name(item.name());
86        ResolvedDataType::Named {
87            id: item.id(),
88            qualifiers: item.name().qualifiers.clone(),
89            full_name,
90            modifiers: vec![],
91            print_id: true,
92        }
93    }
94
95    fn int32_data_type(&self) -> ResolvedDataType {
96        self.resolve_known_valid_data_type(&PartialItemName {
97            database: None,
98            schema: Some(PG_CATALOG_SCHEMA.into()),
99            item: "int4".into(),
100        })
101    }
102
103    // Divides `lhs` by `rhs` but replaces division-by-zero errors with NULL;
104    // note that this is semantically equivalent to `NULLIF(rhs, 0)`.
105    fn plan_divide(lhs: Expr<Aug>, rhs: Expr<Aug>) -> Expr<Aug> {
106        lhs.divide(Expr::Case {
107            operand: None,
108            conditions: vec![rhs.clone().equals(Expr::number("0"))],
109            results: vec![Expr::null()],
110            else_result: Some(Box::new(rhs)),
111        })
112    }
113
114    fn plan_agg(
115        &mut self,
116        name: ResolvedItemName,
117        expr: Expr<Aug>,
118        order_by: Vec<OrderByExpr<Aug>>,
119        filter: Option<Box<Expr<Aug>>>,
120        distinct: bool,
121        over: Option<WindowSpec<Aug>>,
122    ) -> Expr<Aug> {
123        if self.rewriting_table_factor && self.status.is_ok() {
124            self.status = Err(PlanError::Unstructured(
125                "aggregate functions are not supported in functions in FROM".to_string(),
126            ))
127        }
128        Expr::Function(Function {
129            name,
130            args: FunctionArgs::Args {
131                args: vec![expr],
132                order_by,
133            },
134            filter,
135            over,
136            distinct,
137        })
138    }
139
140    fn plan_avg(
141        &mut self,
142        expr: Expr<Aug>,
143        filter: Option<Box<Expr<Aug>>>,
144        distinct: bool,
145        over: Option<WindowSpec<Aug>>,
146    ) -> Expr<Aug> {
147        let sum = self
148            .plan_agg(
149                self.scx
150                    .dangerous_resolve_name(vec![PG_CATALOG_SCHEMA, "sum"]),
151                expr.clone(),
152                vec![],
153                filter.clone(),
154                distinct,
155                over.clone(),
156            )
157            .call_unary(
158                self.scx
159                    .dangerous_resolve_name(vec![MZ_UNSAFE_SCHEMA, "mz_avg_promotion"]),
160            );
161        let count = self.plan_agg(
162            self.scx
163                .dangerous_resolve_name(vec![PG_CATALOG_SCHEMA, "count"]),
164            expr,
165            vec![],
166            filter,
167            distinct,
168            over,
169        );
170        Self::plan_divide(sum, count)
171    }
172
173    /// Same as `plan_avg` but internally uses `mz_avg_promotion_internal_v1`.
174    fn plan_avg_internal_v1(
175        &mut self,
176        expr: Expr<Aug>,
177        filter: Option<Box<Expr<Aug>>>,
178        distinct: bool,
179        over: Option<WindowSpec<Aug>>,
180    ) -> Expr<Aug> {
181        let sum = self
182            .plan_agg(
183                self.scx
184                    .dangerous_resolve_name(vec![PG_CATALOG_SCHEMA, "sum"]),
185                expr.clone(),
186                vec![],
187                filter.clone(),
188                distinct,
189                over.clone(),
190            )
191            .call_unary(
192                self.scx
193                    .dangerous_resolve_name(vec![MZ_UNSAFE_SCHEMA, "mz_avg_promotion_internal_v1"]),
194            );
195        let count = self.plan_agg(
196            self.scx
197                .dangerous_resolve_name(vec![PG_CATALOG_SCHEMA, "count"]),
198            expr,
199            vec![],
200            filter,
201            distinct,
202            over,
203        );
204        Self::plan_divide(sum, count)
205    }
206
207    fn plan_variance(
208        &mut self,
209        expr: Expr<Aug>,
210        filter: Option<Box<Expr<Aug>>>,
211        distinct: bool,
212        sample: bool,
213        over: Option<WindowSpec<Aug>>,
214    ) -> Expr<Aug> {
215        // N.B. this variance calculation uses the "textbook" algorithm, which
216        // is known to accumulate problematic amounts of error. The numerically
217        // stable variants, the most well-known of which is Welford's, are
218        // however difficult to implement inside of Differential Dataflow, as
219        // they do not obviously support retractions efficiently (database-issues#436).
220        //
221        // The code below converts var_samp(x) into
222        //
223        //     (sum(x²) - sum(x)² / count(x)) / (count(x) - 1)
224        //
225        // and var_pop(x) into:
226        //
227        //     (sum(x²) - sum(x)² / count(x)) / count(x)
228        //
229        let expr = expr.call_unary(
230            self.scx
231                .dangerous_resolve_name(vec![MZ_UNSAFE_SCHEMA, "mz_avg_promotion"]),
232        );
233        let expr_squared = expr.clone().multiply(expr.clone());
234        let sum_squares = if distinct {
235            // With DISTINCT, all three component aggregates must deduplicate
236            // on the values of x, not on the values of their own inputs.
237            // sum(DISTINCT x) and count(DISTINCT x) do so naturally, but
238            // sum(DISTINCT x²) deduplicates on x², wrongly collapsing values
239            // that differ only in sign, e.g. -2 and 2. Squaring is injective
240            // on the non-negative values and, separately, on the negative
241            // values, so summing the two sign classes independently makes
242            // deduplication on x² agree with deduplication on x:
243            //
244            //     sum(DISTINCT CASE WHEN x >= 0 THEN x² END)
245            //       + sum(DISTINCT CASE WHEN x < 0 THEN x² END)
246            //
247            // Either sum is NULL when its sign class is empty, so the two are
248            // combined with COALESCE(..., 0). When there are no input rows at
249            // all this yields 0 instead of NULL, but the overall result is
250            // still NULL then because sum(DISTINCT x) below is NULL.
251            let case_squared = |condition| Expr::Case {
252                operand: None,
253                conditions: vec![condition],
254                results: vec![expr_squared.clone()],
255                else_result: None,
256            };
257            let sum_squares_nonneg = self.plan_agg(
258                self.scx
259                    .dangerous_resolve_name(vec![PG_CATALOG_SCHEMA, "sum"]),
260                case_squared(expr.clone().gt_eq(Expr::number("0"))),
261                vec![],
262                filter.clone(),
263                distinct,
264                over.clone(),
265            );
266            let sum_squares_neg = self.plan_agg(
267                self.scx
268                    .dangerous_resolve_name(vec![PG_CATALOG_SCHEMA, "sum"]),
269                case_squared(expr.clone().lt(Expr::number("0"))),
270                vec![],
271                filter.clone(),
272                distinct,
273                over.clone(),
274            );
275            let coalesce_zero = |sum| Expr::HomogenizingFunction {
276                function: HomogenizingFunction::Coalesce,
277                exprs: vec![sum, Expr::number("0")],
278            };
279            coalesce_zero(sum_squares_nonneg).binop(Op::bare("+"), coalesce_zero(sum_squares_neg))
280        } else {
281            self.plan_agg(
282                self.scx
283                    .dangerous_resolve_name(vec![PG_CATALOG_SCHEMA, "sum"]),
284                expr_squared,
285                vec![],
286                filter.clone(),
287                distinct,
288                over.clone(),
289            )
290        };
291        let sum = self.plan_agg(
292            self.scx
293                .dangerous_resolve_name(vec![PG_CATALOG_SCHEMA, "sum"]),
294            expr.clone(),
295            vec![],
296            filter.clone(),
297            distinct,
298            over.clone(),
299        );
300        let sum_squared = sum.clone().multiply(sum);
301        let count = self.plan_agg(
302            self.scx
303                .dangerous_resolve_name(vec![PG_CATALOG_SCHEMA, "count"]),
304            expr,
305            vec![],
306            filter,
307            distinct,
308            over,
309        );
310        let result = Self::plan_divide(
311            sum_squares.minus(Self::plan_divide(sum_squared, count.clone())),
312            if sample {
313                count.minus(Expr::number("1"))
314            } else {
315                count
316            },
317        );
318        // Result is _basically_ what we want, except
319        // that due to numerical inaccuracy, it might be a negative
320        // number very close to zero when it should mathematically be zero.
321        // This makes it so `stddev` fails as it tries to take the square root
322        // of a negative number.
323        // So, we need the following logic:
324        // If `result` is NULL, return NULL (no surprise here)
325        // Otherwise, if `result` is >0, return `result` (no surprise here either)
326        // Otherwise, return 0.
327        //
328        // Unfortunately, we can't use `GREATEST` directly for this,
329        // since `greatest(NULL, 0)` is 0, not NULL, so we need to
330        // create a `Case` expression that computes `result`
331        // twice. Hopefully the optimizer can deal with this!
332        let result_is_null = Expr::IsExpr {
333            expr: Box::new(result.clone()),
334            construct: IsExprConstruct::Null,
335            negated: false,
336        };
337        Expr::Case {
338            operand: None,
339            conditions: vec![result_is_null],
340            results: vec![Expr::Value(Value::Null)],
341            else_result: Some(Box::new(Expr::HomogenizingFunction {
342                function: HomogenizingFunction::Greatest,
343                exprs: vec![result, Expr::number("0")],
344            })),
345        }
346    }
347
348    fn plan_stddev(
349        &mut self,
350        expr: Expr<Aug>,
351        filter: Option<Box<Expr<Aug>>>,
352        distinct: bool,
353        sample: bool,
354        over: Option<WindowSpec<Aug>>,
355    ) -> Expr<Aug> {
356        self.plan_variance(expr, filter, distinct, sample, over)
357            .call_unary(
358                self.scx
359                    .dangerous_resolve_name(vec![PG_CATALOG_SCHEMA, "sqrt"]),
360            )
361    }
362
363    fn plan_bool_and(
364        &mut self,
365        expr: Expr<Aug>,
366        filter: Option<Box<Expr<Aug>>>,
367        distinct: bool,
368        over: Option<WindowSpec<Aug>>,
369    ) -> Expr<Aug> {
370        // The code below converts `bool_and(x)` into:
371        //
372        //     sum((NOT x)::int4) = 0
373        //
374        // It is tempting to use `count` instead, but count does not return NULL
375        // when all input values are NULL, as required.
376        //
377        // The `NOT x` expression has the side effect of implicitly casting `x`
378        // to `bool`. We intentionally do not write `NOT x::bool`, because that
379        // would perform an explicit cast, and to match PostgreSQL we must
380        // perform only an implicit cast.
381        let sum = self.plan_agg(
382            self.scx
383                .dangerous_resolve_name(vec![PG_CATALOG_SCHEMA, "sum"]),
384            expr.negate().cast(self.int32_data_type()),
385            vec![],
386            filter,
387            distinct,
388            over,
389        );
390        sum.equals(Expr::Value(Value::Number(0.to_string())))
391    }
392
393    fn plan_bool_or(
394        &mut self,
395        expr: Expr<Aug>,
396        filter: Option<Box<Expr<Aug>>>,
397        distinct: bool,
398        over: Option<WindowSpec<Aug>>,
399    ) -> Expr<Aug> {
400        // The code below converts `bool_or(x)`z into:
401        //
402        //     sum((x OR false)::int4) > 0
403        //
404        // It is tempting to use `count` instead, but count does not return NULL
405        // when all input values are NULL, as required.
406        //
407        // The `(x OR false)` expression implicitly casts `x` to `bool` without
408        // changing its logical value. It is tempting to use `x::bool` instead,
409        // but that performs an explicit cast, and to match PostgreSQL we must
410        // perform only an implicit cast.
411        let sum = self.plan_agg(
412            self.scx
413                .dangerous_resolve_name(vec![PG_CATALOG_SCHEMA, "sum"]),
414            expr.or(Expr::Value(Value::Boolean(false)))
415                .cast(self.int32_data_type()),
416            vec![],
417            filter,
418            distinct,
419            over,
420        );
421        sum.gt(Expr::Value(Value::Number(0.to_string())))
422    }
423
424    fn rewrite_function(&mut self, func: &Function<Aug>) -> Option<(Ident, Expr<Aug>)> {
425        if let Function {
426            name,
427            args: FunctionArgs::Args { args, order_by: _ },
428            filter,
429            distinct,
430            over,
431        } = func
432        {
433            let pg_catalog_id = self
434                .scx
435                .catalog
436                .resolve_schema(None, PG_CATALOG_SCHEMA)
437                .expect("pg_catalog schema exists")
438                .id();
439            let mz_catalog_id = self
440                .scx
441                .catalog
442                .resolve_schema(None, MZ_CATALOG_SCHEMA)
443                .expect("mz_catalog schema exists")
444                .id();
445            let name = match name {
446                ResolvedItemName::Item {
447                    qualifiers,
448                    full_name,
449                    ..
450                } => {
451                    if ![*pg_catalog_id, *mz_catalog_id].contains(&qualifiers.schema_spec) {
452                        return None;
453                    }
454                    full_name.item.clone()
455                }
456                _ => unreachable!(),
457            };
458
459            let filter = filter.clone();
460            let distinct = *distinct;
461            let over = over.clone();
462            let expr = if args.len() == 1 {
463                let arg = args[0].clone();
464                match name.as_str() {
465                    "avg_internal_v1" => self.plan_avg_internal_v1(arg, filter, distinct, over),
466                    "avg" => self.plan_avg(arg, filter, distinct, over),
467                    "variance" | "var_samp" => {
468                        self.plan_variance(arg, filter, distinct, true, over)
469                    }
470                    "var_pop" => self.plan_variance(arg, filter, distinct, false, over),
471                    "stddev" | "stddev_samp" => self.plan_stddev(arg, filter, distinct, true, over),
472                    "stddev_pop" => self.plan_stddev(arg, filter, distinct, false, over),
473                    "bool_and" => self.plan_bool_and(arg, filter, distinct, over),
474                    "bool_or" => self.plan_bool_or(arg, filter, distinct, over),
475                    _ => return None,
476                }
477            } else if args.len() == 2 {
478                let (lhs, rhs) = (args[0].clone(), args[1].clone());
479                match name.as_str() {
480                    "mod" => lhs.modulo(rhs),
481                    "pow" => Expr::call(
482                        self.scx
483                            .dangerous_resolve_name(vec![PG_CATALOG_SCHEMA, "power"]),
484                        vec![lhs, rhs],
485                    ),
486                    _ => return None,
487                }
488            } else {
489                return None;
490            };
491            Some((Ident::new_unchecked(name), expr))
492        } else {
493            None
494        }
495    }
496
497    fn rewrite_expr(&mut self, expr: &Expr<Aug>) -> Option<(Ident, Expr<Aug>)> {
498        match expr {
499            Expr::Function(function) => self.rewrite_function(function),
500            // Rewrites special keywords that SQL considers to be function calls
501            // to actual function calls. For example, `SELECT current_timestamp`
502            // is rewritten to `SELECT current_timestamp()`.
503            Expr::Identifier(ident) if ident.len() == 1 => {
504                let ident = normalize::ident(ident[0].clone());
505                let fn_ident = match ident.as_str() {
506                    "current_role" => Some("current_user"),
507                    "current_schema" | "current_timestamp" | "current_user" | "session_user"
508                    | "current_catalog" => Some(ident.as_str()),
509                    _ => None,
510                };
511                match fn_ident {
512                    None => None,
513                    Some(fn_ident) => {
514                        let expr = Expr::call_nullary(
515                            self.scx
516                                .dangerous_resolve_name(vec![PG_CATALOG_SCHEMA, fn_ident]),
517                        );
518                        Some((Ident::new_unchecked(ident), expr))
519                    }
520                }
521            }
522            _ => None,
523        }
524    }
525}
526
527impl<'ast> VisitMut<'ast, Aug> for FuncRewriter<'_> {
528    fn visit_select_item_mut(&mut self, item: &'ast mut SelectItem<Aug>) {
529        if let SelectItem::Expr { expr, alias: None } = item {
530            visit_mut::visit_expr_mut(self, expr);
531            if let Some((alias, expr)) = self.rewrite_expr(expr) {
532                *item = SelectItem::Expr {
533                    expr,
534                    alias: Some(alias),
535                };
536            }
537        } else {
538            visit_mut::visit_select_item_mut(self, item);
539        }
540    }
541
542    fn visit_table_with_joins_mut(&mut self, item: &'ast mut TableWithJoins<Aug>) {
543        visit_mut::visit_table_with_joins_mut(self, item);
544        match &mut item.relation {
545            TableFactor::Function {
546                function,
547                alias,
548                with_ordinality,
549            } => {
550                self.rewriting_table_factor = true;
551                // Functions that get rewritten must be rewritten as exprs
552                // because their catalog functions cannot be planned.
553                if let Some((ident, expr)) = self.rewrite_function(function) {
554                    let mut select = Select::default().project(SelectItem::Expr {
555                        expr,
556                        alias: Some(match &alias {
557                            Some(TableAlias { name, columns, .. }) => {
558                                columns.get(0).unwrap_or(name).clone()
559                            }
560                            None => ident,
561                        }),
562                    });
563
564                    if *with_ordinality {
565                        select = select.project(SelectItem::Expr {
566                            expr: Expr::Value(Value::Number("1".into())),
567                            alias: Some(ident!(ORDINALITY_COL_NAME)),
568                        });
569                    }
570
571                    item.relation = TableFactor::Derived {
572                        lateral: false,
573                        subquery: Box::new(Query {
574                            ctes: mz_sql_parser::ast::CteBlock::Simple(vec![]),
575                            body: mz_sql_parser::ast::SetExpr::Select(Box::new(select)),
576                            order_by: vec![],
577                            limit: None,
578                            offset: None,
579                        }),
580                        alias: alias.clone(),
581                    }
582                }
583                self.rewriting_table_factor = false;
584            }
585            _ => {}
586        }
587    }
588
589    fn visit_expr_mut(&mut self, expr: &'ast mut Expr<Aug>) {
590        visit_mut::visit_expr_mut(self, expr);
591        if let Some((_name, new_expr)) = self.rewrite_expr(expr) {
592            *expr = new_expr;
593        }
594    }
595}
596
597/// Removes syntax sugar to simplify the planner.
598///
599/// For example, `<expr> NOT IN (<subquery>)` is rewritten to `expr <> ALL
600/// (<subquery>)`.
601struct Desugarer<'a> {
602    scx: &'a StatementContext<'a>,
603    status: Result<(), PlanError>,
604    id_gen: IdGen,
605    recursion_guard: RecursionGuard,
606}
607
608impl<'a> CheckedRecursion for Desugarer<'a> {
609    fn recursion_guard(&self) -> &RecursionGuard {
610        &self.recursion_guard
611    }
612}
613
614impl<'a, 'ast> VisitMut<'ast, Aug> for Desugarer<'a> {
615    fn visit_expr_mut(&mut self, expr: &'ast mut Expr<Aug>) {
616        self.visit_internal(Self::visit_expr_mut_internal, expr);
617    }
618}
619
620impl<'a> Desugarer<'a> {
621    fn visit_internal<F, X>(&mut self, f: F, x: X)
622    where
623        F: Fn(&mut Self, X) -> Result<(), PlanError>,
624    {
625        if self.status.is_ok() {
626            // self.status could have changed from a deeper call, so don't blindly
627            // overwrite it with the result of this call.
628            let status = self.checked_recur_mut(|d| f(d, x));
629            if self.status.is_ok() {
630                self.status = status;
631            }
632        }
633    }
634
635    fn new(scx: &'a StatementContext) -> Desugarer<'a> {
636        Desugarer {
637            scx,
638            status: Ok(()),
639            id_gen: Default::default(),
640            recursion_guard: RecursionGuard::with_limit(1024), // chosen arbitrarily
641        }
642    }
643
644    fn visit_expr_mut_internal(&mut self, expr: &mut Expr<Aug>) -> Result<(), PlanError> {
645        // `($expr)` => `$expr`
646        while let Expr::Nested(e) = expr {
647            *expr = e.take();
648        }
649
650        // `$expr BETWEEN $low AND $high` => `$expr >= $low AND $expr <= $low`
651        // `$expr NOT BETWEEN $low AND $high` => `$expr < $low OR $expr > $low`
652        if let Expr::Between {
653            expr: e,
654            low,
655            high,
656            negated,
657        } = expr
658        {
659            if *negated {
660                *expr = Expr::lt(*e.clone(), low.take()).or(e.take().gt(high.take()));
661            } else {
662                *expr = e.clone().gt_eq(low.take()).and(e.take().lt_eq(high.take()));
663            }
664        }
665
666        // When `$expr` is a `ROW` constructor, we need to desugar as described
667        // below in order to enable the row comparision expansion at the end of
668        // this function. We don't do this desugaring unconditionally (i.e.,
669        // when `$expr` is not a `ROW` constructor) because the implementation
670        // in `plan_in_list` is more efficient when row comparison expansion is
671        // not required.
672        //
673        // `$expr IN ($list)` => `$expr = $list[0] OR $expr = $list[1] ... OR $expr = $list[n]`
674        // `$expr NOT IN ($list)` => `$expr <> $list[0] AND $expr <> $list[1] ... AND $expr <> $list[n]`
675        if let Expr::InList {
676            expr: e,
677            list,
678            negated,
679        } = expr
680        {
681            if let Expr::Row { .. } = &**e {
682                if *negated {
683                    *expr = list
684                        .drain(..)
685                        .map(|r| e.clone().not_equals(r))
686                        .reduce(|e1, e2| e1.and(e2))
687                        .expect("list known to contain at least one element");
688                } else {
689                    *expr = list
690                        .drain(..)
691                        .map(|r| e.clone().equals(r))
692                        .reduce(|e1, e2| e1.or(e2))
693                        .expect("list known to contain at least one element");
694                }
695            }
696        }
697
698        // `$expr IN ($subquery)` => `$expr = ANY ($subquery)`
699        // `$expr NOT IN ($subquery)` => `$expr <> ALL ($subquery)`
700        if let Expr::InSubquery {
701            expr: e,
702            subquery,
703            negated,
704        } = expr
705        {
706            if *negated {
707                *expr = Expr::AllSubquery {
708                    left: Box::new(e.take()),
709                    op: Op::bare("<>"),
710                    right: Box::new(subquery.take()),
711                };
712            } else {
713                *expr = Expr::AnySubquery {
714                    left: Box::new(e.take()),
715                    op: Op::bare("="),
716                    right: Box::new(subquery.take()),
717                };
718            }
719        }
720
721        // `$expr = ALL ($array_expr)`
722        // =>
723        // `$expr = ALL (SELECT elem FROM unnest($array_expr) _ (elem))`
724        //
725        // and analogously for other operators and ANY.
726        if let Expr::AnyExpr { left, op, right } | Expr::AllExpr { left, op, right } = expr {
727            let binding = ident!("elem");
728
729            let subquery = Query::select(
730                Select::default()
731                    .from(TableWithJoins {
732                        relation: TableFactor::Function {
733                            function: Function {
734                                name: self
735                                    .scx
736                                    .dangerous_resolve_name(vec![MZ_CATALOG_SCHEMA, "unnest"]),
737                                args: FunctionArgs::args(vec![right.take()]),
738                                filter: None,
739                                over: None,
740                                distinct: false,
741                            },
742                            alias: Some(TableAlias {
743                                name: ident!("_"),
744                                columns: vec![binding.clone()],
745                                strict: true,
746                            }),
747                            with_ordinality: false,
748                        },
749                        joins: vec![],
750                    })
751                    .project(SelectItem::Expr {
752                        expr: Expr::Identifier(vec![binding]),
753                        alias: None,
754                    }),
755            );
756
757            let left = Box::new(left.take());
758
759            let op = op.clone();
760
761            *expr = match expr {
762                Expr::AnyExpr { .. } => Expr::AnySubquery {
763                    left,
764                    op,
765                    right: Box::new(subquery),
766                },
767                Expr::AllExpr { .. } => Expr::AllSubquery {
768                    left,
769                    op,
770                    right: Box::new(subquery),
771                },
772                _ => unreachable!(),
773            };
774        }
775
776        // `$expr = ALL ($subquery)`
777        // =>
778        // `(SELECT mz_unsafe.mz_all($expr = $binding) FROM ($subquery) AS _ ($binding))
779        //
780        // and analogously for other operators and ANY.
781        if let Expr::AnySubquery { left, op, right } | Expr::AllSubquery { left, op, right } = expr
782        {
783            let left = match &mut **left {
784                Expr::Row { .. } => left.take(),
785                _ => Expr::Row {
786                    exprs: vec![left.take()],
787                },
788            };
789
790            let arity = match &left {
791                Expr::Row { exprs } => exprs.len(),
792                _ => unreachable!(),
793            };
794
795            let bindings: Vec<_> = (0..arity)
796                // Note: using unchecked is okay here because we know the value will be less than
797                // our maximum length.
798                .map(|col| {
799                    let unique_id = self.id_gen.allocate_id();
800                    Ident::new_unchecked(format!("right_col{col}_{unique_id}"))
801                })
802                .collect();
803
804            let subquery_unique_id = self.id_gen.allocate_id();
805            // Note: kay to use unchecked here because we know the value will be small enough.
806            let subquery_name = Ident::new_unchecked(format!("subquery{subquery_unique_id}"));
807            let select = Select::default()
808                .from(TableWithJoins::subquery(
809                    right.take(),
810                    TableAlias {
811                        name: subquery_name,
812                        columns: bindings.clone(),
813                        strict: true,
814                    },
815                ))
816                .project(SelectItem::Expr {
817                    expr: left
818                        .binop(
819                            op.clone(),
820                            Expr::Row {
821                                exprs: bindings
822                                    .into_iter()
823                                    .map(|b| Expr::Identifier(vec![b]))
824                                    .collect(),
825                            },
826                        )
827                        .call_unary(self.scx.dangerous_resolve_name(match expr {
828                            Expr::AnySubquery { .. } => vec![MZ_UNSAFE_SCHEMA, "mz_any"],
829                            Expr::AllSubquery { .. } => vec![MZ_UNSAFE_SCHEMA, "mz_all"],
830                            _ => unreachable!(),
831                        })),
832                    alias: None,
833                });
834
835            *expr = Expr::Subquery(Box::new(Query::select(select)));
836        }
837
838        // Expands row comparisons.
839        //
840        // ROW($l1, $l2, ..., $ln) = ROW($r1, $r2, ..., $rn)
841        // =>
842        // $l1 = $r1 AND $l2 = $r2 AND ... AND $ln = $rn
843        //
844        // ROW($l1, $l2, ..., $ln) < ROW($r1, $r2, ..., $rn)
845        // =>
846        // $l1 < $r1 OR ($l1 = $r1 AND ($l2 < $r2 OR ($l2 = $r2 AND ... ($ln < $rn))))
847        //
848        // ROW($l1, $l2, ..., $ln) <= ROW($r1, $r2, ..., $rn)
849        // =>
850        // $l1 < $r1 OR ($l1 = $r1 AND ($l2 < $r2 OR ($l2 = $r2 AND ... ($ln <= $rn))))
851        //
852        // and analogously for the inverse operations !=, >, and >=.
853        if let Expr::Op {
854            op,
855            expr1: left,
856            expr2: Some(right),
857        } = expr
858        {
859            if let (Expr::Row { exprs: left }, Expr::Row { exprs: right }) =
860                (&mut **left, &mut **right)
861            {
862                if matches!(normalize::op(op)?, "=" | "<>" | "<" | "<=" | ">" | ">=") {
863                    if left.len() != right.len() {
864                        sql_bail!("unequal number of entries in row expressions");
865                    }
866                    if left.is_empty() {
867                        assert!(right.is_empty());
868                        sql_bail!("cannot compare rows of zero length");
869                    }
870                }
871                match normalize::op(op)? {
872                    "=" | "<>" => {
873                        let mut pairs = left.iter_mut().zip_eq(right);
874                        let mut new = pairs
875                            .next()
876                            .map(|(l, r)| l.take().equals(r.take()))
877                            .expect("cannot compare rows of zero length");
878                        for (l, r) in pairs {
879                            new = l.take().equals(r.take()).and(new);
880                        }
881                        if normalize::op(op)? == "<>" {
882                            new = new.negate();
883                        }
884                        *expr = new;
885                    }
886                    "<" | "<=" | ">" | ">=" => {
887                        let strict_op = match normalize::op(op)? {
888                            "<" | "<=" => "<",
889                            ">" | ">=" => ">",
890                            _ => unreachable!(),
891                        };
892                        let (l, r) = (left.last_mut().unwrap(), right.last_mut().unwrap());
893                        let mut new = l.take().binop(op.clone(), r.take());
894                        for (l, r) in left
895                            .iter_mut()
896                            .rev()
897                            .zip_eq(right.into_iter().rev())
898                            .skip(1)
899                        {
900                            new = l
901                                .clone()
902                                .binop(Op::bare(strict_op), r.clone())
903                                .or(l.take().equals(r.take()).and(new));
904                        }
905                        *expr = new;
906                    }
907                    _ if left.len() == 1 && right.len() == 1 => {
908                        let left = left.remove(0);
909                        let right = right.remove(0);
910                        *expr = left.binop(op.clone(), right);
911                    }
912                    _ => (),
913                }
914            }
915        }
916
917        visit_mut::visit_expr_mut(self, expr);
918        Ok(())
919    }
920}