Skip to main content

mz_sql/plan/
lowering.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//! Lowering is the process of transforming a `HirRelationExpr`
11//! into a `MirRelationExpr`.
12//!
13//! The most crucial part of lowering is decorrelation; i.e.: rewriting a
14//! `HirScalarExpr` that may contain subqueries (e.g. `SELECT` or `EXISTS`)
15//! with instances of `MirScalarExpr` that contain none of these.
16//!
17//! Informally, a subquery should be viewed as a query that is executed in
18//! the context of some outer relation, for each row of that relation. The
19//! subqueries often contain references to the columns of the outer
20//! relation.
21//!
22//! The transformation we perform maintains an `outer` relation and then
23//! traverses the relation expression that may contain references to those
24//! outer columns. As subqueries are discovered, the current relation
25//! expression is recast as the outer expression until such a point as the
26//! scalar expression's evaluation can be determined and appended to each
27//! row of the previously outer relation.
28//!
29//! It is important that the outer columns (the initial columns) act as keys
30//! for all nested computation. When counts or other aggregations are
31//! performed, they should include not only the indicated keys but also all
32//! of the outer columns.
33//!
34//! The decorrelation transformation is initialized with an empty outer
35//! relation, but it seems entirely appropriate to decorrelate queries that
36//! contain "holes" from prepared statements, as if the query was a subquery
37//! against a relation containing the assignments of values to those holes.
38
39use std::collections::{BTreeMap, BTreeSet};
40use std::iter::repeat;
41
42use itertools::Itertools;
43use mz_expr::func::variadic;
44use mz_expr::visit::Visit;
45use mz_expr::{AccessStrategy, AggregateFunc, Columns, MirRelationExpr, MirScalarExpr, func};
46use mz_ore::collections::CollectionExt;
47use mz_ore::stack::maybe_grow;
48use mz_repr::*;
49
50use crate::optimizer_metrics::OptimizerMetrics;
51use crate::plan::hir::{
52    AggregateExpr, ColumnOrder, ColumnRef, HirRelationExpr, HirScalarExpr, JoinKind, WindowExprType,
53};
54use crate::plan::{PlanError, transform_hir};
55use crate::session::vars::SystemVars;
56
57mod variadic_left;
58
59/// Maps a leveled column reference to a specific column.
60///
61/// Leveled column references are nested, so that larger levels are
62/// found early in a record and level zero is found at the end.
63///
64/// The column map only stores references for levels greater than zero,
65/// and column references at level zero simply start at the first column
66/// after all prior references.
67#[derive(Debug, Clone)]
68struct ColumnMap {
69    inner: BTreeMap<ColumnRef, usize>,
70}
71
72impl ColumnMap {
73    fn empty() -> ColumnMap {
74        Self::new(BTreeMap::new())
75    }
76
77    fn new(inner: BTreeMap<ColumnRef, usize>) -> ColumnMap {
78        ColumnMap { inner }
79    }
80
81    fn get(&self, col_ref: &ColumnRef) -> usize {
82        if col_ref.level == 0 {
83            self.inner.len() + col_ref.column
84        } else {
85            self.inner[col_ref]
86        }
87    }
88
89    fn len(&self) -> usize {
90        self.inner.len()
91    }
92
93    /// Updates references in the `ColumnMap` for use in a nested scope. The
94    /// provided `arity` must specify the arity of the current scope.
95    fn enter_scope(&self, arity: usize) -> ColumnMap {
96        // From the perspective of the nested scope, all existing column
97        // references will be one level greater.
98        let existing = self
99            .inner
100            .clone()
101            .into_iter()
102            .update(|(col, _i)| col.level += 1);
103
104        // All columns in the current scope become explicit entries in the
105        // immediate parent scope.
106        let new = (0..arity).map(|i| {
107            (
108                ColumnRef {
109                    level: 1,
110                    column: i,
111                },
112                self.len() + i,
113            )
114        });
115
116        ColumnMap::new(existing.chain(new).collect())
117    }
118}
119
120/// Map with the CTEs currently in scope.
121type CteMap = BTreeMap<mz_expr::LocalId, CteDesc>;
122
123/// Information about needed when finding a reference to a CTE in scope.
124#[derive(Clone)]
125struct CteDesc {
126    /// The new ID assigned to the lowered version of the CTE, which may not match
127    /// the ID of the input CTE.
128    new_id: mz_expr::LocalId,
129    /// The relation type of the CTE including the columns from the outer
130    /// context at the beginning.
131    relation_type: ReprRelationType,
132    /// The outer relation the CTE was applied to.
133    outer_relation: MirRelationExpr,
134}
135
136#[derive(Debug, Clone, Copy)]
137pub struct Config {
138    /// Enable outer join lowering implemented in database-issues#6747.
139    pub enable_new_outer_join_lowering: bool,
140    /// Enable outer join lowering implemented in database-issues#7561.
141    pub enable_variadic_left_join_lowering: bool,
142    pub enable_cast_elimination: bool,
143    pub enable_simplify_quantified_comparisons: bool,
144    /// See the feature flag of the same name.
145    pub enable_fixed_correlated_cte_lowering: bool,
146}
147
148impl Default for Config {
149    fn default() -> Self {
150        Self {
151            enable_new_outer_join_lowering: false,
152            enable_variadic_left_join_lowering: false,
153            enable_cast_elimination: false,
154            enable_simplify_quantified_comparisons: false,
155            enable_fixed_correlated_cte_lowering: false,
156        }
157    }
158}
159
160impl From<&SystemVars> for Config {
161    fn from(vars: &SystemVars) -> Self {
162        Self {
163            enable_new_outer_join_lowering: vars.enable_new_outer_join_lowering(),
164            enable_variadic_left_join_lowering: vars.enable_variadic_left_join_lowering(),
165            enable_cast_elimination: vars.enable_cast_elimination(),
166            enable_simplify_quantified_comparisons: vars.enable_simplify_quantified_comparisons(),
167            enable_fixed_correlated_cte_lowering: vars.enable_fixed_correlated_cte_lowering(),
168        }
169    }
170}
171
172/// Context passed to the lowering. This is wired to most parts of the lowering.
173pub(crate) struct Context<'a> {
174    /// Feature flags affecting the behavior of lowering.
175    pub config: &'a Config,
176    /// Optional, because some callers don't have an `OptimizerMetrics` handy. When it's None, we
177    /// simply don't write metrics.
178    pub metrics: Option<&'a OptimizerMetrics>,
179}
180
181impl HirRelationExpr {
182    /// Rewrite `self` into a `MirRelationExpr`.
183    /// This requires rewriting all correlated subqueries (nested `HirRelationExpr`s) into flat queries
184    #[mz_ore::instrument(target = "optimizer", level = "trace", name = "hir_to_mir")]
185    pub fn lower<C: Into<Config>>(
186        self,
187        config: C,
188        metrics: Option<&OptimizerMetrics>,
189    ) -> Result<MirRelationExpr, PlanError> {
190        let context = Context {
191            config: &config.into(),
192            metrics,
193        };
194        let result = match self {
195            // We directly rewrite a Constant into the corresponding `MirRelationExpr::Constant`
196            // to ensure that the downstream optimizer can easily bypass most
197            // irrelevant optimizations (e.g. reduce folding) for this expression
198            // without having to re-learn the fact that it is just a constant,
199            // as it would if the constant were wrapped in a Let-Get pair.
200            HirRelationExpr::Constant { rows, typ } => {
201                let rows: Vec<_> = rows.into_iter().map(|row| (row, Diff::ONE)).collect();
202                MirRelationExpr::Constant {
203                    rows: Ok(rows),
204                    typ: ReprRelationType::from(&typ),
205                }
206            }
207            mut other => {
208                let mut id_gen = mz_ore::id_gen::IdGen::default();
209                transform_hir::split_subquery_predicates(&mut other)?;
210                transform_hir::try_simplify_quantified_comparisons(
211                    &mut other,
212                    context.config.enable_simplify_quantified_comparisons,
213                )?;
214                transform_hir::fuse_window_functions(&mut other, &context)?;
215                MirRelationExpr::constant(vec![vec![]], ReprRelationType::new(vec![])).let_in(
216                    &mut id_gen,
217                    |id_gen, get_outer| {
218                        other.applied_to(
219                            id_gen,
220                            get_outer,
221                            &ColumnMap::empty(),
222                            &mut CteMap::new(),
223                            &context,
224                        )
225                    },
226                )?
227            }
228        };
229
230        mz_repr::explain::trace_plan(&result);
231
232        Ok(result)
233    }
234
235    /// Return a `MirRelationExpr` which evaluates `self` once for each row of `get_outer`.
236    ///
237    /// For uncorrelated `self`, this should be the cross-product between `get_outer` and `self`.
238    /// When `self` references columns of `get_outer`, much more work needs to occur.
239    ///
240    /// The `col_map` argument contains mappings to some of the columns of `get_outer`, though
241    /// perhaps not all of them. It should be used as the basis of resolving column references,
242    /// but care must be taken when adding new columns that `get_outer.arity()` is where they
243    /// will start, rather than any function of `col_map`.
244    ///
245    /// The `get_outer` expression should be a `Get` with no duplicate rows, describing the distinct
246    /// assignment of values to outer rows.
247    fn applied_to(
248        self,
249        id_gen: &mut mz_ore::id_gen::IdGen,
250        get_outer: MirRelationExpr,
251        col_map: &ColumnMap,
252        cte_map: &mut CteMap,
253        context: &Context,
254    ) -> Result<MirRelationExpr, PlanError> {
255        maybe_grow(|| {
256            use MirRelationExpr as SR;
257
258            use HirRelationExpr::*;
259
260            if let MirRelationExpr::Get { .. } = &get_outer {
261            } else {
262                panic!(
263                    "get_outer: expected a MirRelationExpr::Get, found\n{}",
264                    get_outer.pretty(),
265                );
266            }
267            assert_eq!(col_map.len(), get_outer.arity());
268            Ok(match self {
269                Constant { rows, typ } => {
270                    // Constant expressions are not correlated with `get_outer`, and should be cross-products.
271                    get_outer.product(SR::Constant {
272                        rows: Ok(rows.into_iter().map(|row| (row, Diff::ONE)).collect()),
273                        typ: ReprRelationType::from(&typ),
274                    })
275                }
276                Get { id, typ } => match id {
277                    mz_expr::Id::Local(local_id) => {
278                        let cte_desc = cte_map.get(&local_id).unwrap();
279                        let get_cte = SR::Get {
280                            id: mz_expr::Id::Local(cte_desc.new_id.clone()),
281                            typ: cte_desc.relation_type.clone(),
282                            access_strategy: AccessStrategy::UnknownOrLocal,
283                        };
284                        if get_outer == cte_desc.outer_relation {
285                            // If the CTE was applied to the same exact relation, we can safely
286                            // return a `Get` relation.
287                            get_cte
288                        } else {
289                            // Otherwise, the new outer relation may contain more columns from some
290                            // intermediate scope placed between the definition of the CTE and this
291                            // reference of the CTE and/or more operations applied on top of the
292                            // outer relation.
293                            //
294                            // An example of the latter is the following query:
295                            //
296                            // SELECT *
297                            // FROM x,
298                            //      LATERAL(WITH a(m) as (SELECT max(y.a) FROM y WHERE y.a < x.a)
299                            //              SELECT (SELECT m FROM a) FROM y) b;
300                            //
301                            // When the CTE is lowered, the outer relation is `Get x`. But then,
302                            // the reference of the CTE is applied to `Distinct(Join(Get x, Get y), x.*)`
303                            // which has the same cardinality as `Get x`.
304                            //
305                            // In any case, `get_outer` is guaranteed to contain the columns of the
306                            // outer relation the CTE was applied to at its prefix. Since, we must
307                            // return a relation containing `get_outer`'s column at the beginning,
308                            // we must build a join between `get_outer` and `get_cte` on their common
309                            // columns.
310                            let oa = get_outer.arity();
311                            let cte_outer_columns = cte_desc.relation_type.arity() - typ.arity();
312                            let equivalences = (0..cte_outer_columns)
313                                .map(|pos| {
314                                    vec![
315                                        MirScalarExpr::column(pos),
316                                        MirScalarExpr::column(pos + oa),
317                                    ]
318                                })
319                                .collect();
320
321                            // Project out the second copy of the common between `get_outer` and
322                            // `cte_desc.outer_relation`.
323                            let projection = (0..oa)
324                                .chain(oa + cte_outer_columns..oa + cte_outer_columns + typ.arity())
325                                .collect_vec();
326                            SR::join_scalars(vec![get_outer, get_cte], equivalences)
327                                .project(projection)
328                        }
329                    }
330                    mz_expr::Id::Global(_) => {
331                        // Get statements are only to external sources, and are not correlated with `get_outer`.
332                        get_outer.product(SR::Get {
333                            id,
334                            typ: ReprRelationType::from(&typ),
335                            access_strategy: AccessStrategy::UnknownOrLocal,
336                        })
337                    }
338                },
339                Let {
340                    name: _,
341                    id,
342                    value,
343                    body,
344                } => {
345                    let value =
346                        value.applied_to(id_gen, get_outer.clone(), col_map, cte_map, context)?;
347                    value.let_in(id_gen, |id_gen, get_value| {
348                        let (new_id, typ) = if let MirRelationExpr::Get {
349                            id: mz_expr::Id::Local(id),
350                            typ,
351                            ..
352                        } = get_value
353                        {
354                            (id, typ)
355                        } else {
356                            panic!(
357                                "get_value: expected a MirRelationExpr::Get with local Id, found\n{}",
358                                get_value.pretty(),
359                            );
360                        };
361                        // Add the information about the CTE to the map and remove it when
362                        // it goes out of scope.
363                        let old_value = cte_map.insert(
364                            id.clone(),
365                            CteDesc {
366                                new_id,
367                                relation_type: typ,
368                                outer_relation: get_outer.clone(),
369                            },
370                        );
371                        let body = body.applied_to(id_gen, get_outer, col_map, cte_map, context);
372                        if let Some(old_value) = old_value {
373                            cte_map.insert(id, old_value);
374                        } else {
375                            cte_map.remove(&id);
376                        }
377                        body
378                    })?
379                }
380                LetRec {
381                    limit,
382                    bindings,
383                    body,
384                } => {
385                    let num_bindings = bindings.len();
386
387                    // We use the outer type with the HIR types to form MIR CTE types.
388                    let outer_column_types = get_outer.typ().column_types;
389
390                    // Rename and introduce all bindings.
391                    let mut shadowed_bindings = Vec::with_capacity(num_bindings);
392                    let mut mir_ids = Vec::with_capacity(num_bindings);
393                    for (_name, id, _value, typ) in bindings.iter() {
394                        let mir_id = mz_expr::LocalId::new(id_gen.allocate_id());
395                        mir_ids.push(mir_id);
396                        let shadowed = cte_map.insert(
397                            id.clone(),
398                            CteDesc {
399                                new_id: mir_id,
400                                relation_type: ReprRelationType::new(
401                                    outer_column_types
402                                        .iter()
403                                        .cloned()
404                                        .chain(typ.column_types.iter().map(ReprColumnType::from))
405                                        .collect::<Vec<_>>(),
406                                ),
407                                outer_relation: get_outer.clone(),
408                            },
409                        );
410                        shadowed_bindings.push((*id, shadowed));
411                    }
412
413                    let mut mir_values = Vec::with_capacity(num_bindings);
414                    for (_name, _id, value, _typ) in bindings.into_iter() {
415                        mir_values.push(value.applied_to(
416                            id_gen,
417                            get_outer.clone(),
418                            col_map,
419                            cte_map,
420                            context,
421                        )?);
422                    }
423
424                    let mir_body = body.applied_to(id_gen, get_outer, col_map, cte_map, context)?;
425
426                    // Remove our bindings and reinstate any shadowed bindings.
427                    for (id, shadowed) in shadowed_bindings {
428                        if let Some(shadowed) = shadowed {
429                            cte_map.insert(id, shadowed);
430                        } else {
431                            cte_map.remove(&id);
432                        }
433                    }
434
435                    MirRelationExpr::LetRec {
436                        ids: mir_ids,
437                        values: mir_values,
438                        // Copy the limit to each binding.
439                        limits: repeat(limit).take(num_bindings).collect(),
440                        body: Box::new(mir_body),
441                    }
442                }
443                Project { input, outputs } => {
444                    // Projections should be applied to the decorrelated `inner`, and to its columns,
445                    // which means rebasing `outputs` to start `get_outer.arity()` columns later.
446                    let input =
447                        input.applied_to(id_gen, get_outer.clone(), col_map, cte_map, context)?;
448                    let outputs = (0..get_outer.arity())
449                        .chain(outputs.into_iter().map(|i| get_outer.arity() + i))
450                        .collect::<Vec<_>>();
451                    input.project(outputs)
452                }
453                Map { input, mut scalars } => {
454                    // Scalar expressions may contain correlated subqueries. We must be cautious!
455
456                    // We lower scalars in chunks, and must keep track of the
457                    // arity of the HIR fragments lowered so far.
458                    let mut lowered_arity = input.arity();
459
460                    let mut input =
461                        input.applied_to(id_gen, get_outer, col_map, cte_map, context)?;
462
463                    // Lower subqueries in maximally sized batches, such as no subquery in the current
464                    // batch depends on columns from the same batch.
465                    // Note that subqueries in this projection may reference columns added by this
466                    // Map operator, so we need to ensure these columns exist before lowering the
467                    // subquery.
468                    while !scalars.is_empty() {
469                        let end_idx = scalars
470                            .iter_mut()
471                            .position(|s| {
472                                let mut requires_nonexistent_column = false;
473                                #[allow(deprecated)]
474                                s.visit_columns(0, &mut |depth, col| {
475                                    if col.level == depth {
476                                        requires_nonexistent_column |= col.column >= lowered_arity
477                                    }
478                                });
479                                requires_nonexistent_column
480                            })
481                            .unwrap_or(scalars.len());
482                        assert!(
483                            end_idx > 0,
484                            "a Map expression references itself or a later column; lowered_arity: {}, expressions: {:?}",
485                            lowered_arity,
486                            scalars
487                        );
488
489                        lowered_arity = lowered_arity + end_idx;
490                        let scalars = scalars.drain(0..end_idx).collect_vec();
491
492                        let old_arity = input.arity();
493                        let (with_subqueries, subquery_map) = HirScalarExpr::lower_subqueries(
494                            &scalars, id_gen, col_map, cte_map, input, context,
495                        )?;
496                        input = with_subqueries;
497
498                        // We will proceed sequentially through the scalar expressions, for each transforming
499                        // the decorrelated `input` into a relation with potentially more columns capable of
500                        // addressing the needs of the scalar expression.
501                        // Having done so, we add the scalar value of interest and trim off any other newly
502                        // added columns.
503                        //
504                        // The sequential traversal is present as expressions are allowed to depend on the
505                        // values of prior expressions.
506                        let mut scalar_columns = Vec::new();
507                        for scalar in scalars {
508                            let scalar = scalar.applied_to(
509                                id_gen,
510                                col_map,
511                                cte_map,
512                                &mut input,
513                                &Some(&subquery_map),
514                                context,
515                            )?;
516                            input = input.map_one(scalar);
517                            scalar_columns.push(input.arity() - 1);
518                        }
519
520                        // Discard any new columns added by the lowering of the scalar expressions
521                        input = input.project((0..old_arity).chain(scalar_columns).collect());
522                    }
523
524                    input
525                }
526                CallTable { func, exprs } => {
527                    // FlatMap expressions may contain correlated subqueries. Unlike Map they are not
528                    // allowed to refer to the results of previous expressions, and we have a simpler
529                    // implementation that appends all relevant columns first, then applies the flatmap
530                    // operator to the result, then strips off any columns introduce by subqueries.
531
532                    let mut input = get_outer;
533                    let old_arity = input.arity();
534
535                    let exprs = exprs
536                        .into_iter()
537                        .map(|e| e.applied_to(id_gen, col_map, cte_map, &mut input, &None, context))
538                        .collect::<Result<Vec<_>, _>>()?;
539
540                    let new_arity = input.arity();
541                    let output_arity = func.output_arity();
542                    input = input.flat_map(func, exprs);
543                    if old_arity != new_arity {
544                        // this means we added some columns to handle subqueries, and now we need to get rid of them
545                        input = input.project(
546                            (0..old_arity)
547                                .chain(new_arity..new_arity + output_arity)
548                                .collect(),
549                        );
550                    }
551                    input
552                }
553                Filter { input, predicates } => {
554                    // Filter expressions may contain correlated subqueries.
555                    // We extend `get_outer` with sufficient values to determine the value of the predicate,
556                    // then filter the results, then strip off any columns that were added for this purpose.
557                    let mut input =
558                        input.applied_to(id_gen, get_outer, col_map, cte_map, context)?;
559                    for predicate in predicates {
560                        let old_arity = input.arity();
561                        let predicate = predicate
562                            .applied_to(id_gen, col_map, cte_map, &mut input, &None, context)?;
563                        let new_arity = input.arity();
564                        input = input.filter(vec![predicate]);
565                        if old_arity != new_arity {
566                            // this means we added some columns to handle subqueries, and now we need to get rid of them
567                            input = input.project((0..old_arity).collect());
568                        }
569                    }
570                    input
571                }
572                Join {
573                    left,
574                    right,
575                    on,
576                    kind,
577                } if right.is_correlated() => {
578                    // A correlated join is a join in which the right expression has
579                    // access to the columns in the left expression. It turns out
580                    // this is *exactly* our branch operator, plus some additional
581                    // null handling in the case of left joins. (Right and full
582                    // lateral joins are not permitted.)
583                    //
584                    // As with normal joins, the `on` predicate may be correlated,
585                    // and we treat it as a filter that follows the branch.
586
587                    assert!(kind.can_be_correlated());
588
589                    let left = left.applied_to(id_gen, get_outer, col_map, cte_map, context)?;
590                    left.let_in(id_gen, |id_gen, get_left| {
591                        let apply_requires_distinct_outer = false;
592                        let mut join = branch(
593                            id_gen,
594                            get_left.clone(),
595                            col_map,
596                            cte_map,
597                            *right,
598                            apply_requires_distinct_outer,
599                            context,
600                            |id_gen, right, get_left, col_map, cte_map, context| {
601                                right.applied_to(id_gen, get_left, col_map, cte_map, context)
602                            },
603                        )?;
604
605                        // Plan the `on` predicate.
606                        let old_arity = join.arity();
607                        let on =
608                            on.applied_to(id_gen, col_map, cte_map, &mut join, &None, context)?;
609                        join = join.filter(vec![on]);
610                        let new_arity = join.arity();
611                        if old_arity != new_arity {
612                            // This means we added some columns to handle
613                            // subqueries, and now we need to get rid of them.
614                            join = join.project((0..old_arity).collect());
615                        }
616
617                        // If a left join, reintroduce any rows from the left that
618                        // are missing, with nulls filled in for the right columns.
619                        if let JoinKind::LeftOuter { .. } = kind {
620                            let default = join
621                                .typ()
622                                .column_types
623                                .into_iter()
624                                .skip(get_left.arity())
625                                .map(|typ| (Datum::Null, typ.scalar_type))
626                                .collect();
627                            get_left.lookup(id_gen, join, default)
628                        } else {
629                            Ok::<_, PlanError>(join)
630                        }
631                    })?
632                }
633                Join {
634                    left,
635                    right,
636                    on,
637                    kind,
638                } => {
639                    if context.config.enable_variadic_left_join_lowering {
640                        // Attempt to extract a stack of left joins.
641                        if let JoinKind::LeftOuter = kind {
642                            let mut rights = vec![(&*right, &on)];
643                            let mut left_test = &left;
644                            while let Join {
645                                left,
646                                right,
647                                on,
648                                kind: JoinKind::LeftOuter,
649                            } = &**left_test
650                            {
651                                rights.push((&**right, on));
652                                left_test = left;
653                            }
654                            if rights.len() > 1 {
655                                // Defensively clone `cte_map` as it may be mutated.
656                                let cte_map_clone = cte_map.clone();
657                                if let Ok(Some(magic)) = variadic_left::attempt_left_join_magic(
658                                    left_test,
659                                    rights,
660                                    id_gen,
661                                    get_outer.clone(),
662                                    col_map,
663                                    cte_map,
664                                    context,
665                                ) {
666                                    return Ok(magic);
667                                } else {
668                                    cte_map.clone_from(&cte_map_clone);
669                                }
670                            }
671                        }
672                    }
673
674                    // Both join expressions should be decorrelated, and then joined by their
675                    // leading columns to form only those pairs corresponding to the same row
676                    // of `get_outer`.
677                    //
678                    // The `on` predicate may contain correlated subqueries, and we treat it
679                    // as though it was a filter, with the caveat that we also translate outer
680                    // joins in this step. The post-filtration results need to be considered
681                    // against the records present in the left and right (decorrelated) inputs,
682                    // depending on the type of join.
683                    let oa = get_outer.arity();
684                    let left =
685                        left.applied_to(id_gen, get_outer.clone(), col_map, cte_map, context)?;
686                    let lt = left.typ().column_types.into_iter().skip(oa).collect_vec();
687                    let la = lt.len();
688                    left.let_in(id_gen, |id_gen, get_left| {
689                        let right_col_map = col_map.enter_scope(0);
690                        let right = right.applied_to(
691                            id_gen,
692                            get_outer.clone(),
693                            &right_col_map,
694                            cte_map,
695                            context,
696                        )?;
697                        let rt = right.typ().column_types.into_iter().skip(oa).collect_vec();
698                        let ra = rt.len();
699                        right.let_in(id_gen, |id_gen, get_right| {
700                            let mut product = SR::join(
701                                vec![get_left.clone(), get_right.clone()],
702                                (0..oa).map(|i| vec![(0, i), (1, i)]).collect(),
703                            )
704                            // Project away the repeated copy of get_outer's columns.
705                            .project(
706                                (0..(oa + la))
707                                    .chain((oa + la + oa)..(oa + la + oa + ra))
708                                    .collect(),
709                            );
710
711                            // Decorrelate and lower the `on` clause.
712                            let on = on.applied_to(
713                                id_gen,
714                                col_map,
715                                cte_map,
716                                &mut product,
717                                &None,
718                                context,
719                            )?;
720                            // Collect the types of all subqueries appearing in
721                            // the `on` clause. The subquery results were
722                            // appended to `product` in the `on.applied_to(...)`
723                            // call above.
724                            let on_subquery_types = product
725                                .typ()
726                                .column_types
727                                .drain(oa + la + ra..)
728                                .collect_vec();
729                            // Remember if `on` had any subqueries.
730                            let on_has_subqueries = !on_subquery_types.is_empty();
731
732                            // Attempt an efficient equijoin implementation, in which outer joins are
733                            // more efficiently rendered than in general. This can return `None` if
734                            // such a plan is not possible, for example if `on` does not describe an
735                            // equijoin between columns of `left` and `right`.
736                            if kind != JoinKind::Inner {
737                                if let Some(joined) = attempt_outer_equijoin(
738                                    get_left.clone(),
739                                    get_right.clone(),
740                                    on.clone(),
741                                    on_subquery_types,
742                                    kind.clone(),
743                                    oa,
744                                    id_gen,
745                                    context,
746                                )? {
747                                    if let Some(metrics) = context.metrics {
748                                        metrics.inc_outer_join_lowering("equi");
749                                    }
750                                    return Ok(joined);
751                                }
752                            }
753
754                            // Otherwise, perform a more general join.
755                            if let Some(metrics) = context.metrics {
756                                metrics.inc_outer_join_lowering("general");
757                            }
758                            let mut join = product.filter(vec![on]);
759                            if on_has_subqueries {
760                                // This means that `on.applied_to(...)` appended
761                                // some columns to handle subqueries, and now we
762                                // need to get rid of them.
763                                join = join.project((0..oa + la + ra).collect());
764                            }
765                            join.let_in(id_gen, |id_gen, get_join| {
766                                let mut result = get_join.clone();
767                                if let JoinKind::LeftOuter { .. } | JoinKind::FullOuter { .. } =
768                                    kind
769                                {
770                                    let left_outer = get_left.clone().anti_lookup::<PlanError>(
771                                        id_gen,
772                                        get_join.clone(),
773                                        rt.into_iter()
774                                            .map(|typ| (Datum::Null, typ.scalar_type))
775                                            .collect(),
776                                    )?;
777                                    result = result.union(left_outer);
778                                }
779                                if let JoinKind::RightOuter | JoinKind::FullOuter = kind {
780                                    let right_outer = get_right
781                                        .clone()
782                                        .anti_lookup::<PlanError>(
783                                            id_gen,
784                                            get_join
785                                                // need to swap left and right to make the anti_lookup work
786                                                .project(
787                                                    (0..oa)
788                                                        .chain((oa + la)..(oa + la + ra))
789                                                        .chain((oa)..(oa + la))
790                                                        .collect(),
791                                                ),
792                                            lt.into_iter()
793                                                .map(|typ| (Datum::Null, typ.scalar_type))
794                                                .collect(),
795                                        )?
796                                        // swap left and right back again
797                                        .project(
798                                            (0..oa)
799                                                .chain((oa + ra)..(oa + ra + la))
800                                                .chain((oa)..(oa + ra))
801                                                .collect(),
802                                        );
803                                    result = result.union(right_outer);
804                                }
805                                Ok::<MirRelationExpr, PlanError>(result)
806                            })
807                        })
808                    })?
809                }
810                Union { base, inputs } => {
811                    // Union is uncomplicated.
812                    SR::Union {
813                        base: Box::new(base.applied_to(
814                            id_gen,
815                            get_outer.clone(),
816                            col_map,
817                            cte_map,
818                            context,
819                        )?),
820                        inputs: inputs
821                            .into_iter()
822                            .map(|input| {
823                                input.applied_to(
824                                    id_gen,
825                                    get_outer.clone(),
826                                    col_map,
827                                    cte_map,
828                                    context,
829                                )
830                            })
831                            .collect::<Result<Vec<_>, _>>()?,
832                    }
833                }
834                Reduce {
835                    input,
836                    group_key,
837                    aggregates,
838                    expected_group_size,
839                } => {
840                    // Reduce may contain expressions with correlated subqueries.
841                    // In addition, here an empty reduction key signifies that we need to supply default values
842                    // in the case that there are no results (as in a SQL aggregation without an explicit GROUP BY).
843                    let mut input =
844                        input.applied_to(id_gen, get_outer.clone(), col_map, cte_map, context)?;
845                    let applied_group_key = (0..get_outer.arity())
846                        .chain(group_key.iter().map(|i| get_outer.arity() + i))
847                        .collect();
848                    let applied_aggregates = aggregates
849                        .into_iter()
850                        .map(|aggregate| {
851                            aggregate.applied_to(id_gen, col_map, cte_map, &mut input, context)
852                        })
853                        .collect::<Result<Vec<_>, _>>()?;
854                    let input_type = input.typ();
855                    let default = applied_aggregates
856                        .iter()
857                        .map(|agg| {
858                            (
859                                agg.func.default(),
860                                agg.typ(&input_type.column_types).scalar_type,
861                            )
862                        })
863                        .collect();
864                    // NOTE we don't need to remove any extra columns from aggregate.applied_to above because the reduce will do that anyway
865                    let mut reduced =
866                        input.reduce(applied_group_key, applied_aggregates, expected_group_size);
867
868                    // Introduce default values in the case the group key is empty.
869                    if group_key.is_empty() {
870                        reduced = get_outer.lookup::<PlanError>(id_gen, reduced, default)?;
871                    }
872                    reduced
873                }
874                Distinct { input } => {
875                    // Distinct is uncomplicated.
876                    input
877                        .applied_to(id_gen, get_outer, col_map, cte_map, context)?
878                        .distinct()
879                }
880                TopK {
881                    input,
882                    group_key,
883                    order_key,
884                    limit,
885                    offset,
886                    expected_group_size,
887                } => {
888                    // TopK is uncomplicated, except that we must group by the columns of `get_outer` as well.
889                    let mut input =
890                        input.applied_to(id_gen, get_outer.clone(), col_map, cte_map, context)?;
891                    let mut applied_group_key: Vec<_> = (0..get_outer.arity())
892                        .chain(group_key.iter().map(|i| get_outer.arity() + i))
893                        .collect();
894                    let applied_order_key = order_key
895                        .iter()
896                        .map(|column_order| ColumnOrder {
897                            column: column_order.column + get_outer.arity(),
898                            desc: column_order.desc,
899                            nulls_last: column_order.nulls_last,
900                        })
901                        .collect();
902
903                    let old_arity = input.arity();
904
905                    // Lower `limit`, which may introduce new columns if is a correlated subquery.
906                    let mut limit_mir = None;
907                    if let Some(limit) = limit {
908                        limit_mir = Some(
909                            limit
910                                .applied_to(id_gen, col_map, cte_map, &mut input, &None, context)?,
911                        );
912                    }
913
914                    let new_arity = input.arity();
915                    // Extend the key to contain any new columns.
916                    applied_group_key.extend(old_arity..new_arity);
917
918                    let offset = offset
919                        .try_into_literal_int64()
920                        .expect("Should be a Literal by this time")
921                        .try_into()
922                        .expect("Should have checked non-negativity of OFFSET clause already");
923                    let mut result = input.top_k(
924                        applied_group_key,
925                        applied_order_key,
926                        limit_mir,
927                        offset,
928                        expected_group_size,
929                    );
930
931                    // If new columns were added for `limit` we must remove them.
932                    if old_arity != new_arity {
933                        result = result.project((0..old_arity).collect());
934                    }
935
936                    result
937                }
938                Negate { input } => {
939                    // Negate is uncomplicated.
940                    input
941                        .applied_to(id_gen, get_outer, col_map, cte_map, context)?
942                        .negate()
943                }
944                Threshold { input } => {
945                    // Threshold is uncomplicated.
946                    input
947                        .applied_to(id_gen, get_outer, col_map, cte_map, context)?
948                        .threshold()
949                }
950            })
951        })
952    }
953}
954
955impl HirScalarExpr {
956    /// Rewrite `self` into a `mz_expr::ScalarExpr` which can be applied to the modified `inner`.
957    ///
958    /// This method is responsible for decorrelating subqueries in `self` by introducing further columns
959    /// to `inner`, and rewriting `self` to refer to its physical columns (specified by `usize` positions).
960    /// The most complicated logic is for the scalar expressions that involve subqueries, each of which are
961    /// documented in more detail closer to their logic.
962    ///
963    /// This process presumes that `inner` is the result of decorrelation, meaning its first several columns
964    /// may be inherited from outer relations. The `col_map` column map should provide specific offsets where
965    /// each of these references can be found.
966    fn applied_to(
967        self,
968        id_gen: &mut mz_ore::id_gen::IdGen,
969        col_map: &ColumnMap,
970        cte_map: &mut CteMap,
971        inner: &mut MirRelationExpr,
972        subquery_map: &Option<&BTreeMap<HirScalarExpr, usize>>,
973        context: &Context,
974    ) -> Result<MirScalarExpr, PlanError> {
975        maybe_grow(|| {
976            use MirScalarExpr as SS;
977
978            use HirScalarExpr::*;
979
980            if let Some(subquery_map) = subquery_map {
981                if let Some(col) = subquery_map.get(&self) {
982                    return Ok(SS::column(*col));
983                }
984            }
985
986            Ok::<MirScalarExpr, PlanError>(match self {
987                Column(col_ref, name) => SS::Column(col_map.get(&col_ref), name),
988                Literal(row, typ, _name) => SS::Literal(Ok(row), ReprColumnType::from(&typ)),
989                Parameter(_, _name) => {
990                    panic!("cannot decorrelate expression with unbound parameters")
991                }
992                CallUnmaterializable(func, _name) => SS::CallUnmaterializable(func),
993                CallUnary {
994                    func,
995                    expr,
996                    name: _,
997                } => {
998                    let inner =
999                        expr.applied_to(id_gen, col_map, cte_map, inner, subquery_map, context)?;
1000                    if context.config.enable_cast_elimination && func.is_eliminable_cast() {
1001                        inner
1002                    } else {
1003                        SS::CallUnary {
1004                            func,
1005                            expr: Box::new(inner),
1006                        }
1007                    }
1008                }
1009                CallBinary {
1010                    func,
1011                    expr1,
1012                    expr2,
1013                    name: _,
1014                } => SS::CallBinary {
1015                    func,
1016                    expr1: Box::new(expr1.applied_to(
1017                        id_gen,
1018                        col_map,
1019                        cte_map,
1020                        inner,
1021                        subquery_map,
1022                        context,
1023                    )?),
1024                    expr2: Box::new(expr2.applied_to(
1025                        id_gen,
1026                        col_map,
1027                        cte_map,
1028                        inner,
1029                        subquery_map,
1030                        context,
1031                    )?),
1032                },
1033                CallVariadic {
1034                    func,
1035                    exprs,
1036                    name: _,
1037                } => SS::call_variadic(
1038                    func,
1039                    exprs
1040                        .into_iter()
1041                        .map(|expr| {
1042                            expr.applied_to(id_gen, col_map, cte_map, inner, subquery_map, context)
1043                        })
1044                        .collect::<Result<_, _>>()?,
1045                ),
1046                If {
1047                    cond,
1048                    then,
1049                    els,
1050                    name,
1051                } => {
1052                    // The `If` case is complicated by the fact that we do not want to
1053                    // apply the `then` or `else` logic to tuples that respectively do
1054                    // not or do pass the `cond` test. Our strategy is to independently
1055                    // decorrelate the `then` and `else` logic, and apply each to tuples
1056                    // that respectively pass and do not pass the `cond` logic (which is
1057                    // executed, and so decorrelated, for all tuples).
1058                    //
1059                    // Informally, we turn the `if` statement into:
1060                    //
1061                    //   let then_case = inner.filter(cond).map(then);
1062                    //   let else_case = inner.filter(!cond).map(else);
1063                    //   return then_case.concat(else_case);
1064                    //
1065                    // We only require this if either expression would result in any
1066                    // computation beyond the expr itself, which we will interpret as
1067                    // "introduces additional columns". In the absence of correlation,
1068                    // we should just retain a `ScalarExpr::If` expression; the inverse
1069                    // transformation as above is complicated to recover after the fact,
1070                    // and we would benefit from not introducing the complexity.
1071
1072                    let inner_arity = inner.arity();
1073                    let cond_expr =
1074                        cond.applied_to(id_gen, col_map, cte_map, inner, subquery_map, context)?;
1075
1076                    // Defensive copies, in case we mangle these in decorrelation.
1077                    let inner_clone = inner.clone();
1078                    let then_clone = then.clone();
1079                    let else_clone = els.clone();
1080
1081                    let cond_arity = inner.arity();
1082                    let then_expr =
1083                        then.applied_to(id_gen, col_map, cte_map, inner, subquery_map, context)?;
1084                    let else_expr =
1085                        els.applied_to(id_gen, col_map, cte_map, inner, subquery_map, context)?;
1086
1087                    if cond_arity == inner.arity() {
1088                        // If no additional columns were added, we simply return the
1089                        // `If` variant with the updated expressions.
1090                        SS::If {
1091                            cond: Box::new(cond_expr),
1092                            then: Box::new(then_expr),
1093                            els: Box::new(else_expr),
1094                        }
1095                    } else {
1096                        // If columns were added, we need a more careful approach, as
1097                        // described above. First, we need to de-correlate each of
1098                        // the two expressions independently, and apply their cases
1099                        // as `MirRelationExpr::Map` operations.
1100
1101                        *inner = inner_clone.let_in(id_gen, |id_gen, get_inner| {
1102                            // Restrict to records satisfying `cond_expr` and apply `then` as a map.
1103                            let mut then_inner = get_inner.clone().filter(vec![cond_expr.clone()]);
1104                            let then_expr = then_clone.applied_to(
1105                                id_gen,
1106                                col_map,
1107                                cte_map,
1108                                &mut then_inner,
1109                                subquery_map,
1110                                context,
1111                            )?;
1112                            let then_arity = then_inner.arity();
1113                            then_inner = then_inner
1114                                .map_one(then_expr)
1115                                .project((0..inner_arity).chain(Some(then_arity)).collect());
1116
1117                            // Restrict to records not satisfying `cond_expr` and apply `els` as a map.
1118                            let mut else_inner = get_inner.filter(vec![SS::call_variadic(
1119                                variadic::Or,
1120                                vec![
1121                                    cond_expr.clone().call_binary(SS::literal_false(), func::Eq),
1122                                    cond_expr.clone().call_is_null(),
1123                                ],
1124                            )]);
1125                            let else_expr = else_clone.applied_to(
1126                                id_gen,
1127                                col_map,
1128                                cte_map,
1129                                &mut else_inner,
1130                                subquery_map,
1131                                context,
1132                            )?;
1133                            let else_arity = else_inner.arity();
1134                            else_inner = else_inner
1135                                .map_one(else_expr)
1136                                .project((0..inner_arity).chain(Some(else_arity)).collect());
1137
1138                            // concatenate the two results.
1139                            Ok::<MirRelationExpr, PlanError>(then_inner.union(else_inner))
1140                        })?;
1141
1142                        SS::Column(inner_arity, name)
1143                    }
1144                }
1145
1146                // Subqueries!
1147                // These are surprisingly subtle. Things to be careful of:
1148
1149                // Anything in the subquery that cares about row counts (Reduce/Distinct/Negate/Threshold) must not:
1150                // * change the row counts of the outer query
1151                // * accidentally compute its own value using the row counts of the outer query
1152                // Use `branch` to calculate the subquery once for each __distinct__ key in the outer
1153                // query and then join the answers back on to the original rows of the outer query.
1154
1155                // When the subquery would return 0 rows for some row in the outer query, `subquery.applied_to(get_inner)` will not have any corresponding row.
1156                // Use `lookup` if you need to add default values for cases when the subquery returns 0 rows.
1157                Exists(expr, name) => {
1158                    let apply_requires_distinct_outer = true;
1159                    *inner = apply_existential_subquery(
1160                        id_gen,
1161                        inner.take_dangerous(),
1162                        col_map,
1163                        cte_map,
1164                        *expr,
1165                        apply_requires_distinct_outer,
1166                        context,
1167                    )?;
1168                    SS::Column(inner.arity() - 1, name)
1169                }
1170
1171                Select(expr, name) => {
1172                    let apply_requires_distinct_outer = true;
1173                    *inner = apply_scalar_subquery(
1174                        id_gen,
1175                        inner.take_dangerous(),
1176                        col_map,
1177                        cte_map,
1178                        *expr,
1179                        apply_requires_distinct_outer,
1180                        context,
1181                    )?;
1182                    SS::Column(inner.arity() - 1, name)
1183                }
1184                Windowing(expr, _name) => {
1185                    let partition_by = expr.partition_by;
1186                    let order_by = expr.order_by;
1187
1188                    // argument lowering for scalar window functions
1189                    // (We need to specify the & _ in the arguments because of this problem:
1190                    // https://users.rust-lang.org/t/the-implementation-of-fnonce-is-not-general-enough/72141/3 )
1191                    let scalar_lower_args =
1192                        |_id_gen: &mut _,
1193                         _col_map: &_,
1194                         _cte_map: &mut _,
1195                         _get_inner: &mut _,
1196                         _subquery_map: &Option<&_>,
1197                         order_by_mir: Vec<MirScalarExpr>,
1198                         original_row_record,
1199                         original_row_record_type: SqlScalarType| {
1200                            let agg_input = MirScalarExpr::call_variadic(
1201                                variadic::ListCreate {
1202                                    elem_type: original_row_record_type.clone(),
1203                                },
1204                                vec![original_row_record],
1205                            );
1206                            let mut agg_input = vec![agg_input];
1207                            agg_input.extend(order_by_mir.clone());
1208                            let agg_input = MirScalarExpr::call_variadic(
1209                                variadic::RecordCreate {
1210                                    field_names: (0..agg_input.len())
1211                                        .map(|_| ColumnName::from(UNKNOWN_COLUMN_NAME))
1212                                        .collect_vec(),
1213                                },
1214                                agg_input,
1215                            );
1216                            let list_type = SqlScalarType::List {
1217                                element_type: Box::new(original_row_record_type),
1218                                custom_id: None,
1219                            };
1220                            let agg_input_type = SqlScalarType::Record {
1221                                fields: std::iter::once(&list_type)
1222                                    .map(|t| {
1223                                        (
1224                                            ColumnName::from(UNKNOWN_COLUMN_NAME),
1225                                            t.clone().nullable(false),
1226                                        )
1227                                    })
1228                                    .collect(),
1229                                custom_id: None,
1230                            }
1231                            .nullable(false);
1232
1233                            Ok((agg_input, agg_input_type))
1234                        };
1235
1236                    // argument lowering for value window functions and aggregate window functions
1237                    let value_or_aggr_lower_args = |hir_encoded_args: Box<HirScalarExpr>| {
1238                        |id_gen: &mut _,
1239                         col_map: &_,
1240                         cte_map: &mut _,
1241                         get_inner: &mut _,
1242                         subquery_map: &Option<&_>,
1243                         order_by_mir: Vec<MirScalarExpr>,
1244                         original_row_record,
1245                         original_row_record_type| {
1246                            // Creates [((OriginalRow, EncodedArgs), OrderByExprs...)]
1247
1248                            // Compute the encoded args for all rows
1249                            let mir_encoded_args = hir_encoded_args.applied_to(
1250                                id_gen,
1251                                col_map,
1252                                cte_map,
1253                                get_inner,
1254                                subquery_map,
1255                                context,
1256                            )?;
1257                            let mir_encoded_args_type = mir_encoded_args
1258                                .sql_typ(&get_inner.sql_typ().column_types)
1259                                .scalar_type;
1260
1261                            // Build a new record that has two fields:
1262                            // 1. the original row in a record
1263                            // 2. the encoded args (which can be either a single value, or a record
1264                            //    if the window function has multiple arguments, such as `lag`)
1265                            let fn_input_record_fields: Box<[_]> =
1266                                [original_row_record_type, mir_encoded_args_type]
1267                                    .iter()
1268                                    .map(|t| {
1269                                        (
1270                                            ColumnName::from(UNKNOWN_COLUMN_NAME),
1271                                            t.clone().nullable(false),
1272                                        )
1273                                    })
1274                                    .collect();
1275                            let fn_input_record = MirScalarExpr::call_variadic(
1276                                variadic::RecordCreate {
1277                                    field_names: fn_input_record_fields
1278                                        .iter()
1279                                        .map(|(n, _)| n.clone())
1280                                        .collect_vec(),
1281                                },
1282                                vec![original_row_record, mir_encoded_args],
1283                            );
1284                            let fn_input_record_type = SqlScalarType::Record {
1285                                fields: fn_input_record_fields,
1286                                custom_id: None,
1287                            }
1288                            .nullable(false);
1289
1290                            // Build a new record with the record above + the ORDER BY exprs
1291                            // This follows the standard encoding of ORDER BY exprs used by aggregate functions
1292                            let mut agg_input = vec![fn_input_record];
1293                            agg_input.extend(order_by_mir.clone());
1294                            let agg_input = MirScalarExpr::call_variadic(
1295                                variadic::RecordCreate {
1296                                    field_names: (0..agg_input.len())
1297                                        .map(|_| ColumnName::from(UNKNOWN_COLUMN_NAME))
1298                                        .collect_vec(),
1299                                },
1300                                agg_input,
1301                            );
1302
1303                            let agg_input_type = SqlScalarType::Record {
1304                                fields: [(
1305                                    ColumnName::from(UNKNOWN_COLUMN_NAME),
1306                                    fn_input_record_type.nullable(false),
1307                                )]
1308                                .into(),
1309                                custom_id: None,
1310                            }
1311                            .nullable(false);
1312
1313                            Ok((agg_input, agg_input_type))
1314                        }
1315                    };
1316
1317                    match expr.func {
1318                        WindowExprType::Scalar(scalar_window_expr) => {
1319                            let mir_aggr_func = scalar_window_expr.into_expr();
1320                            Self::window_func_applied_to(
1321                                id_gen,
1322                                col_map,
1323                                cte_map,
1324                                inner,
1325                                subquery_map,
1326                                partition_by,
1327                                order_by,
1328                                mir_aggr_func,
1329                                scalar_lower_args,
1330                                context,
1331                            )?
1332                        }
1333                        WindowExprType::Value(value_window_expr) => {
1334                            let (hir_encoded_args, mir_aggr_func) = value_window_expr.into_expr();
1335
1336                            Self::window_func_applied_to(
1337                                id_gen,
1338                                col_map,
1339                                cte_map,
1340                                inner,
1341                                subquery_map,
1342                                partition_by,
1343                                order_by,
1344                                mir_aggr_func,
1345                                value_or_aggr_lower_args(hir_encoded_args),
1346                                context,
1347                            )?
1348                        }
1349                        WindowExprType::Aggregate(aggr_window_expr) => {
1350                            let (hir_encoded_args, mir_aggr_func) = aggr_window_expr.into_expr();
1351
1352                            Self::window_func_applied_to(
1353                                id_gen,
1354                                col_map,
1355                                cte_map,
1356                                inner,
1357                                subquery_map,
1358                                partition_by,
1359                                order_by,
1360                                mir_aggr_func,
1361                                value_or_aggr_lower_args(hir_encoded_args),
1362                                context,
1363                            )?
1364                        }
1365                    }
1366                }
1367            })
1368        })
1369    }
1370
1371    fn window_func_applied_to<F>(
1372        id_gen: &mut mz_ore::id_gen::IdGen,
1373        col_map: &ColumnMap,
1374        cte_map: &mut CteMap,
1375        inner: &mut MirRelationExpr,
1376        subquery_map: &Option<&BTreeMap<HirScalarExpr, usize>>,
1377        partition_by: Vec<HirScalarExpr>,
1378        order_by: Vec<HirScalarExpr>,
1379        mir_aggr_func: AggregateFunc,
1380        lower_args: F,
1381        context: &Context,
1382    ) -> Result<MirScalarExpr, PlanError>
1383    where
1384        F: FnOnce(
1385            &mut mz_ore::id_gen::IdGen,
1386            &ColumnMap,
1387            &mut CteMap,
1388            &mut MirRelationExpr,
1389            &Option<&BTreeMap<HirScalarExpr, usize>>,
1390            Vec<MirScalarExpr>,
1391            MirScalarExpr,
1392            SqlScalarType,
1393        ) -> Result<(MirScalarExpr, SqlColumnType), PlanError>,
1394    {
1395        // Example MIRs for a window function (specifically, a window aggregation):
1396        //
1397        // CREATE TABLE t7(x INT, y INT);
1398        //
1399        // explain decorrelated plan for select sum(x*y) over (partition by x+y order by x-y, x/y) from t7;
1400        //
1401        // Decorrelated Plan
1402        // Project (#3)
1403        //   Map (#2)
1404        //     Project (#3..=#5)
1405        //       Map (record_get[0](record_get[1](#2)), record_get[1](record_get[1](#2)), record_get[0](#2))
1406        //         FlatMap unnest_list(#1)
1407        //           Reduce group_by=[#2] aggregates=[window_agg[sum order_by=[#0 asc nulls_last, #1 asc nulls_last]](row(row(row(#0, #1), (#0 * #1)), (#0 - #1), (#0 / #1)))]
1408        //             Map ((#0 + #1))
1409        //               CrossJoin
1410        //                 Constant
1411        //                   - ()
1412        //                 Get materialize.public.t7
1413        //
1414        // The same query after optimizations:
1415        //
1416        // explain select sum(x*y) over (partition by x+y order by x-y, x/y) from t7;
1417        //
1418        // Optimized Plan
1419        // Explained Query:
1420        //   Project (#2)
1421        //     Map (record_get[0](#1))
1422        //       FlatMap unnest_list(#0)
1423        //         Project (#1)
1424        //           Reduce group_by=[(#0 + #1)] aggregates=[window_agg[sum order_by=[#0 asc nulls_last, #1 asc nulls_last]](row(row(row(#0, #1), (#0 * #1)), (#0 - #1), (#0 / #1)))]
1425        //             ReadStorage materialize.public.t7
1426        //
1427        // The `row(row(row(...), ...), ...)` stuff means the following:
1428        // `row(row(row(<original row>), <arguments to window function>), <order by values>...)`
1429        //   - The <arguments to window function> can be either a single value or itself a
1430        //     `row` if there are multiple arguments.
1431        //   - The <order by values> are _not_ wrapped in a `row`, even if there are more than one
1432        //     ORDER BY columns.
1433        //   - The <original row> currently always captures the entire original row. This should
1434        //     improve when we make `ProjectionPushdown` smarter, see
1435        //     https://github.com/MaterializeInc/database-issues/issues/5090
1436        //
1437        // TODO:
1438        // We should probably introduce some dedicated Datum constructor functions instead of `row`
1439        // to make MIR plans and MIR construction/manipulation code more readable. Additionally, we
1440        // might even introduce dedicated Datum enum variants, so that the rendering code also
1441        // becomes more readable (and possibly slightly more performant).
1442
1443        *inner = inner
1444            .take_dangerous()
1445            .let_in(id_gen, |id_gen, mut get_inner| {
1446                let order_by_mir = order_by
1447                    .into_iter()
1448                    .map(|o| {
1449                        o.applied_to(
1450                            id_gen,
1451                            col_map,
1452                            cte_map,
1453                            &mut get_inner,
1454                            subquery_map,
1455                            context,
1456                        )
1457                    })
1458                    .collect::<Result<Vec<_>, _>>()?;
1459
1460                // Record input arity here so that any group_keys that need to mutate get_inner
1461                // don't add those columns to the aggregate input.
1462                let input_type = get_inner.sql_typ();
1463                let input_arity = input_type.arity();
1464                // The reduction that computes the window function must be keyed on the columns
1465                // from the outer context, plus the expressions in the partition key. The current
1466                // subquery will be 'executed' for every distinct row from the outer context so
1467                // by putting the outer columns in the grouping key we isolate each re-execution.
1468                let mut group_key = col_map
1469                    .inner
1470                    .iter()
1471                    .map(|(_, outer_col)| *outer_col)
1472                    .sorted()
1473                    .collect_vec();
1474                for p in partition_by {
1475                    let key = p.applied_to(
1476                        id_gen,
1477                        col_map,
1478                        cte_map,
1479                        &mut get_inner,
1480                        subquery_map,
1481                        context,
1482                    )?;
1483                    if let MirScalarExpr::Column(c, _name) = key {
1484                        group_key.push(c);
1485                    } else {
1486                        get_inner = get_inner.map_one(key);
1487                        group_key.push(get_inner.arity() - 1);
1488                    }
1489                }
1490
1491                get_inner.let_in(id_gen, |id_gen, mut get_inner| {
1492                    // Original columns of the relation
1493                    let fields: Box<_> = input_type
1494                        .column_types
1495                        .iter()
1496                        .take(input_arity)
1497                        .map(|t| (ColumnName::from(UNKNOWN_COLUMN_NAME), t.clone()))
1498                        .collect();
1499
1500                    // Original row made into a record
1501                    let original_row_record = MirScalarExpr::call_variadic(
1502                        variadic::RecordCreate {
1503                            field_names: fields.iter().map(|(name, _)| name.clone()).collect_vec(),
1504                        },
1505                        (0..input_arity).map(MirScalarExpr::column).collect_vec(),
1506                    );
1507                    let original_row_record_type = SqlScalarType::Record {
1508                        fields,
1509                        custom_id: None,
1510                    };
1511
1512                    let (agg_input, agg_input_type) = lower_args(
1513                        id_gen,
1514                        col_map,
1515                        cte_map,
1516                        &mut get_inner,
1517                        subquery_map,
1518                        order_by_mir,
1519                        original_row_record,
1520                        original_row_record_type,
1521                    )?;
1522
1523                    let aggregate = mz_expr::AggregateExpr {
1524                        func: mir_aggr_func,
1525                        expr: agg_input,
1526                        distinct: false,
1527                    };
1528
1529                    // Actually call reduce with the window function
1530                    // The output of the aggregation function should be a list of tuples that has
1531                    // the result in the first position, and the original row in the second position
1532                    let mut reduce = get_inner
1533                        .reduce(group_key.clone(), vec![aggregate.clone()], None)
1534                        .flat_map(
1535                            mz_expr::TableFunc::UnnestList {
1536                                el_typ: aggregate
1537                                    .func
1538                                    .output_sql_type(agg_input_type)
1539                                    .scalar_type
1540                                    .unwrap_list_element_type()
1541                                    .clone(),
1542                            },
1543                            vec![MirScalarExpr::column(group_key.len())],
1544                        );
1545                    let record_col = reduce.arity() - 1;
1546
1547                    // Unpack the record output by the window function
1548                    for c in 0..input_arity {
1549                        reduce = reduce.take_dangerous().map_one(MirScalarExpr::CallUnary {
1550                            func: mz_expr::UnaryFunc::RecordGet(mz_expr::func::RecordGet(c)),
1551                            expr: Box::new(MirScalarExpr::CallUnary {
1552                                func: mz_expr::UnaryFunc::RecordGet(mz_expr::func::RecordGet(1)),
1553                                expr: Box::new(MirScalarExpr::column(record_col)),
1554                            }),
1555                        });
1556                    }
1557
1558                    // Append the column with the result of the window function.
1559                    reduce = reduce.take_dangerous().map_one(MirScalarExpr::CallUnary {
1560                        func: mz_expr::UnaryFunc::RecordGet(mz_expr::func::RecordGet(0)),
1561                        expr: Box::new(MirScalarExpr::column(record_col)),
1562                    });
1563
1564                    let agg_col = record_col + 1 + input_arity;
1565                    Ok::<_, PlanError>(reduce.project((record_col + 1..agg_col + 1).collect_vec()))
1566                })
1567            })?;
1568        Ok(MirScalarExpr::column(inner.arity() - 1))
1569    }
1570
1571    /// Applies the subqueries in the given list of scalar expressions to every distinct
1572    /// value of the given relation and returns a join of the given relation with all
1573    /// the subqueries found, and the mapping of scalar expressions with columns projected
1574    /// by the returned join that will hold their results.
1575    fn lower_subqueries(
1576        exprs: &[Self],
1577        id_gen: &mut mz_ore::id_gen::IdGen,
1578        col_map: &ColumnMap,
1579        cte_map: &mut CteMap,
1580        inner: MirRelationExpr,
1581        context: &Context,
1582    ) -> Result<(MirRelationExpr, BTreeMap<HirScalarExpr, usize>), PlanError> {
1583        let mut subquery_map = BTreeMap::new();
1584        let output = inner.let_in(id_gen, |id_gen, get_inner| {
1585            let mut subqueries = Vec::new();
1586            let distinct_inner = get_inner.clone().distinct();
1587            for expr in exprs.iter() {
1588                expr.visit_pre_post(
1589                    &mut |e| match e {
1590                        // For simplicity, subqueries within a conditional statement will be
1591                        // lowered when lowering the conditional expression.
1592                        HirScalarExpr::If { .. } => Some(vec![]),
1593                        _ => None,
1594                    },
1595                    &mut |e| match e {
1596                        HirScalarExpr::Select(expr, _name) => {
1597                            let apply_requires_distinct_outer = false;
1598                            let subquery = apply_scalar_subquery(
1599                                id_gen,
1600                                distinct_inner.clone(),
1601                                col_map,
1602                                cte_map,
1603                                (**expr).clone(),
1604                                apply_requires_distinct_outer,
1605                                context,
1606                            )
1607                            .unwrap();
1608
1609                            subqueries.push((e.clone(), subquery));
1610                        }
1611                        HirScalarExpr::Exists(expr, _name) => {
1612                            let apply_requires_distinct_outer = false;
1613                            let subquery = apply_existential_subquery(
1614                                id_gen,
1615                                distinct_inner.clone(),
1616                                col_map,
1617                                cte_map,
1618                                (**expr).clone(),
1619                                apply_requires_distinct_outer,
1620                                context,
1621                            )
1622                            .unwrap();
1623                            subqueries.push((e.clone(), subquery));
1624                        }
1625                        _ => {}
1626                    },
1627                );
1628            }
1629
1630            if subqueries.is_empty() {
1631                Ok::<MirRelationExpr, PlanError>(get_inner)
1632            } else {
1633                let inner_arity = get_inner.arity();
1634                let mut total_arity = inner_arity;
1635                let mut join_inputs = vec![get_inner];
1636                let mut join_input_arities = vec![inner_arity];
1637                for (expr, subquery) in subqueries.into_iter() {
1638                    // Avoid lowering duplicated subqueries
1639                    if !subquery_map.contains_key(&expr) {
1640                        let subquery_arity = subquery.arity();
1641                        assert_eq!(subquery_arity, inner_arity + 1);
1642                        join_inputs.push(subquery);
1643                        join_input_arities.push(subquery_arity);
1644                        total_arity += subquery_arity;
1645
1646                        // Column with the value of the subquery
1647                        subquery_map.insert(expr, total_arity - 1);
1648                    }
1649                }
1650                // Each subquery projects all the columns of the outer context (distinct_inner)
1651                // plus 1 column, containing the result of the subquery. Those columns must be
1652                // joined with the outer/main relation (get_inner).
1653                let input_mapper =
1654                    mz_expr::JoinInputMapper::new_from_input_arities(join_input_arities);
1655                let equivalences = (0..inner_arity)
1656                    .map(|col| {
1657                        join_inputs
1658                            .iter()
1659                            .enumerate()
1660                            .map(|(input, _)| {
1661                                MirScalarExpr::column(input_mapper.map_column_to_global(col, input))
1662                            })
1663                            .collect_vec()
1664                    })
1665                    .collect_vec();
1666                Ok(MirRelationExpr::join_scalars(join_inputs, equivalences))
1667            }
1668        })?;
1669        Ok((output, subquery_map))
1670    }
1671
1672    /// Rewrites `self` into a `mz_expr::ScalarExpr`.
1673    ///
1674    /// Returns an _internal_ error if the expression contains
1675    /// - a subquery
1676    /// - a column reference to an outer level
1677    /// - a parameter
1678    /// - a window function call
1679    ///
1680    /// Should succeed if [`HirScalarExpr::is_constant`] would return true on `self`.
1681    ///
1682    /// Set `enable_cast_elimination` to remove casts that are noops in MIR.
1683    pub fn lower_uncorrelated<C: Into<Config>>(
1684        self,
1685        config: C,
1686    ) -> Result<MirScalarExpr, PlanError> {
1687        let config = config.into();
1688
1689        use MirScalarExpr as SS;
1690
1691        use HirScalarExpr::*;
1692
1693        Ok(match self {
1694            Column(ColumnRef { level: 0, column }, name) => SS::Column(column, name),
1695            Literal(datum, typ, _name) => SS::Literal(Ok(datum), ReprColumnType::from(&typ)),
1696            CallUnmaterializable(func, _name) => SS::CallUnmaterializable(func),
1697            CallUnary {
1698                func,
1699                expr,
1700                name: _,
1701            } => {
1702                let inner = expr.lower_uncorrelated(config)?;
1703
1704                if config.enable_cast_elimination && func.is_eliminable_cast() {
1705                    inner
1706                } else {
1707                    SS::CallUnary {
1708                        func,
1709                        expr: Box::new(inner),
1710                    }
1711                }
1712            }
1713            CallBinary {
1714                func,
1715                expr1,
1716                expr2,
1717                name: _,
1718            } => SS::CallBinary {
1719                func,
1720                expr1: Box::new(expr1.lower_uncorrelated(config)?),
1721                expr2: Box::new(expr2.lower_uncorrelated(config)?),
1722            },
1723            CallVariadic {
1724                func,
1725                exprs,
1726                name: _,
1727            } => SS::call_variadic(
1728                func,
1729                exprs
1730                    .into_iter()
1731                    .map(|expr| expr.lower_uncorrelated(config))
1732                    .collect::<Result<_, _>>()?,
1733            ),
1734            If {
1735                cond,
1736                then,
1737                els,
1738                name: _,
1739            } => SS::If {
1740                cond: Box::new(cond.lower_uncorrelated(config)?),
1741                then: Box::new(then.lower_uncorrelated(config)?),
1742                els: Box::new(els.lower_uncorrelated(config)?),
1743            },
1744            Select { .. } | Exists { .. } | Parameter(..) | Column(..) | Windowing(..) => {
1745                sql_bail!(
1746                    "Internal error: unexpected HirScalarExpr in lower_uncorrelated: {:?}",
1747                    self
1748                );
1749            }
1750        })
1751    }
1752}
1753
1754/// Prepare to apply `inner` to `outer`. Note that `inner` is a correlated (SQL)
1755/// expression, while `outer` is a non-correlated (dataflow) expression. `inner`
1756/// will, in effect, be executed once for every distinct row in `outer`, and the
1757/// results will be joined with `outer`. Note that columns in `outer` that are
1758/// not depended upon by `inner` are thrown away before the distinct, so that we
1759/// don't perform needless computation of `inner`.
1760///
1761/// `branch` will inspect the contents of `inner` to determine whether `inner`
1762/// is not multiplicity sensitive (roughly, contains only maps, filters,
1763/// projections, and calls to table functions). If it is not multiplicity
1764/// sensitive, `branch` will *not* distinctify outer. If this is problematic,
1765/// e.g. because the `apply` callback itself introduces multiplicity-sensitive
1766/// operations that were not present in `inner`, then set
1767/// `apply_requires_distinct_outer` to ensure that `branch` chooses the plan
1768/// that distinctifies `outer`.
1769///
1770/// The caller must supply the `apply` function that applies the rewritten
1771/// `inner` to `outer`.
1772fn branch<F>(
1773    id_gen: &mut mz_ore::id_gen::IdGen,
1774    outer: MirRelationExpr,
1775    col_map: &ColumnMap,
1776    cte_map: &mut CteMap,
1777    inner: HirRelationExpr,
1778    apply_requires_distinct_outer: bool,
1779    context: &Context,
1780    apply: F,
1781) -> Result<MirRelationExpr, PlanError>
1782where
1783    F: FnOnce(
1784        &mut mz_ore::id_gen::IdGen,
1785        HirRelationExpr,
1786        MirRelationExpr,
1787        &ColumnMap,
1788        &mut CteMap,
1789        &Context,
1790    ) -> Result<MirRelationExpr, PlanError>,
1791{
1792    // TODO: It would be nice to have a version of this code w/o optimizations,
1793    // at the least for purposes of understanding. It was difficult for one reader
1794    // to understand the required properties of `outer` and `col_map`.
1795
1796    // If the inner expression is sufficiently simple, it is safe to apply it
1797    // *directly* to outer, rather than applying it to the distinctified key
1798    // (see below).
1799    //
1800    // As an example, consider the following two queries:
1801    //
1802    //     CREATE TABLE t (a int, b int);
1803    //     SELECT a, series FROM t, generate_series(1, t.b) series;
1804    //
1805    // The "simple" path for the `SELECT` yields
1806    //
1807    //     %0 =
1808    //     | Get t
1809    //     | FlatMap generate_series(1, #1)
1810    //
1811    // while the non-simple path yields:
1812    //
1813    //    %0 =
1814    //    | Get t
1815    //
1816    //    %1 =
1817    //    | Get t
1818    //    | Distinct group=(#1)
1819    //    | FlatMap generate_series(1, #0)
1820    //
1821    //    %2 =
1822    //    | LeftJoin %1 %2 (= #1 #2)
1823    //
1824    // There is a tradeoff here: the simple plan is stateless, but the non-
1825    // simple plan may do (much) less computation if there are only a few
1826    // distinct values of `t.b`.
1827    //
1828    // We apply a very simple heuristic here and take the simple path if `inner`
1829    // contains only maps, filters, projections, and calls to table functions.
1830    // The intuition is that straightforward usage of table functions should
1831    // take the simple path, while everything else should not. (In theory we
1832    // think this transformation is valid as long as `inner` does not contain a
1833    // Reduce, Distinct, or TopK node, but it is not always an optimization in
1834    // the general case.)
1835    //
1836    // TODO(benesch): this should all be handled by a proper optimizer, but
1837    // detecting the moment of decorrelation in the optimizer right now is too
1838    // hard.
1839    let mut is_simple = true;
1840    #[allow(deprecated)]
1841    inner.visit(0, &mut |expr, _| match expr {
1842        HirRelationExpr::Constant { .. }
1843        | HirRelationExpr::Project { .. }
1844        | HirRelationExpr::Map { .. }
1845        | HirRelationExpr::Filter { .. }
1846        | HirRelationExpr::CallTable { .. } => (),
1847        _ => is_simple = false,
1848    });
1849    if is_simple && !apply_requires_distinct_outer {
1850        let new_col_map = col_map.enter_scope(outer.arity() - col_map.len());
1851        return outer.let_in(id_gen, |id_gen, get_outer| {
1852            apply(id_gen, inner, get_outer, &new_col_map, cte_map, context)
1853        });
1854    }
1855
1856    // The key consists of the columns from the outer expression upon which the
1857    // inner relation depends. We discover these dependencies by walking the
1858    // inner relation expression and looking for column references whose level
1859    // escapes inner.
1860    //
1861    // At the end of this process, `key` contains the decorrelated position of
1862    // each outer column, according to the passed-in `col_map`, and
1863    // `new_col_map` maps each outer column to its new ordinal position in key.
1864    let mut outer_cols = BTreeSet::new();
1865    #[allow(deprecated)]
1866    inner.visit_columns(0, &mut |depth, col| {
1867        // Test if the column reference escapes the subquery.
1868        if col.level > depth {
1869            outer_cols.insert(ColumnRef {
1870                level: col.level - depth,
1871                column: col.column,
1872            });
1873        }
1874    });
1875    // Collect all the outer columns referenced by any CTE referenced by
1876    // the inner relation. The `Get` arm of `applied_to` reconciles a CTE
1877    // reference with the relation the CTE was applied to by equating the
1878    // first `cte_outer_arity` columns of both sides, so it relies on the
1879    // outer columns of every referenced CTE sitting at the beginning of the
1880    // branch key, in order. We track the maximum such arity to arrange for
1881    // that below.
1882    let mut max_cte_outer_arity = 0;
1883    {
1884        let mut visit = |e: &HirRelationExpr| match e {
1885            HirRelationExpr::Get {
1886                id: mz_expr::Id::Local(id),
1887                ..
1888            } => {
1889                if let Some(cte_desc) = cte_map.get(id) {
1890                    let cte_outer_arity = cte_desc.outer_relation.arity();
1891                    max_cte_outer_arity = max_cte_outer_arity.max(cte_outer_arity);
1892                    outer_cols.extend(
1893                        col_map
1894                            .inner
1895                            .iter()
1896                            .filter(|(_, position)| **position < cte_outer_arity)
1897                            .map(|(c, _)| {
1898                                // `col_map` maps column references to column positions in
1899                                // `outer`'s projection.
1900                                // `outer_cols` is meant to contain the external column
1901                                // references in `inner`.
1902                                // Since `inner` defines a new scope, any column reference
1903                                // in `col_map` is one level deeper when seen from within
1904                                // `inner`, hence the +1.
1905                                ColumnRef {
1906                                    level: c.level + 1,
1907                                    column: c.column,
1908                                }
1909                            }),
1910                    );
1911                }
1912            }
1913            HirRelationExpr::Let { id, .. } => {
1914                // Note: if ID uniqueness is not guaranteed, we can't use `visit` since
1915                // we would need to remove the old CTE with the same ID temporarily while
1916                // traversing the definition of the new CTE under the same ID.
1917                assert!(!cte_map.contains_key(id));
1918            }
1919            _ => {}
1920        };
1921        if context.config.enable_fixed_correlated_cte_lowering {
1922            // `visit_post` descends into scalar subqueries, so this finds CTE
1923            // references at any depth inside `inner`. Full depth matters: a
1924            // nested `branch()` inside `inner` relies on the key computed
1925            // here already carrying its referenced CTEs' outer columns as a
1926            // prefix.
1927            inner.visit_post(&mut visit);
1928        } else {
1929            // The deprecated `visit` does not descend into scalar subqueries,
1930            // so CTE references inside them are missed and the branch key can
1931            // lack outer columns that the CTE reference's reconciliation join
1932            // needs, producing wrong correlations (SQL-349).
1933            #[allow(deprecated)]
1934            inner.visit(0, &mut |e, _| visit(e));
1935        }
1936    }
1937    let mut new_col_map = BTreeMap::new();
1938    let mut key = vec![];
1939    if context.config.enable_fixed_correlated_cte_lowering {
1940        // The `Get` arm reconciles a CTE reference by equating key slots
1941        // `0..cte_outer_arity` with the CTE's outer columns, positionally, so
1942        // it needs every referenced CTE's outer columns at the head of the
1943        // key, in their original `outer` order. Those columns are exactly the
1944        // ones at `outer` positions `0..max_cte_outer_arity` (see the discovery
1945        // above), so we split the discovered columns into that prefix and the
1946        // rest, and emit the prefix first, sorted by position. That makes
1947        // `key[i] == i` over the prefix, which is what the `Get` arm assumes.
1948        //
1949        // A single prefix covers all referenced CTEs because their outer
1950        // relations are nested prefixes of one another (nested scopes only ever
1951        // append columns). The `rest`'s order is free: the final join in
1952        // `branch` compensates for any key order, and only the `Get` arm cares
1953        // about absolute slot positions (we keep `outer_cols`'s order).
1954        let mut prefix = Vec::new();
1955        let mut rest = Vec::new();
1956        for col in outer_cols {
1957            // `position` is `col`'s index in `outer`'s projection.
1958            let position = col_map.get(&ColumnRef {
1959                // `outer_cols` holds references as seen from within `inner`,
1960                // one level deeper than `outer`'s scope (the discovery walk and
1961                // `visit_columns` both record them that way), so undo that with
1962                // `-1` before looking `col` up in `outer`'s `col_map`.
1963                level: col.level - 1,
1964                column: col.column,
1965            });
1966            if position < max_cte_outer_arity {
1967                prefix.push((position, col));
1968            } else {
1969                rest.push((position, col));
1970            }
1971        }
1972        prefix.sort_unstable_by_key(|(position, _)| *position);
1973        // The discovery inserts every position `0..cte_outer_arity` of each
1974        // referenced CTE, and `col_map` is a bijection onto `0..col_map.len()`,
1975        // so the prefix positions are exactly `0..max_cte_outer_arity`, each
1976        // once, reading `0, 1, ..., max-1` after the sort. If they don't, the
1977        // `Get` arm's reconciliation would join on the wrong columns and
1978        // silently produce wrong results, so fail the query instead. There is
1979        // no safe fallback: the unpartitioned key order is the bug this is
1980        // fixing.
1981        if !prefix
1982            .iter()
1983            .map(|(position, _)| *position)
1984            .eq(0..max_cte_outer_arity)
1985        {
1986            return Err(PlanError::Internal(format!(
1987                "CTE outer columns are not a contiguous prefix of the branch key: \
1988                 prefix positions {:?}, max_cte_outer_arity {}",
1989                prefix
1990                    .iter()
1991                    .map(|(position, _)| *position)
1992                    .collect::<Vec<_>>(),
1993                max_cte_outer_arity,
1994            )));
1995        }
1996        // Give each chosen outer column its key slot: `new_col_map` records the
1997        // slot, `key` records which `outer` position feeds it. Over the prefix
1998        // the slot (`key.len()`) equals `position`, so `key[i] == i`; over the
1999        // rest they may differ, which is fine.
2000        for (position, col) in prefix.into_iter().chain(rest) {
2001            new_col_map.insert(col, key.len());
2002            key.push(position);
2003        }
2004    } else {
2005        // Note: this order can break the `Get` arm's prefix assumption when a
2006        // level-1 column sorts before a referenced CTE's outer column
2007        // (SQL-349).
2008        for col in outer_cols {
2009            new_col_map.insert(col, key.len());
2010            key.push(col_map.get(&ColumnRef {
2011                // Note: `outer_cols` contains the external column references within `inner`.
2012                // We must compensate for `inner`'s scope when translating column references
2013                // as seen within `inner` to column references as seen from `outer`'s context,
2014                // hence the -1.
2015                level: col.level - 1,
2016                column: col.column,
2017            }));
2018        }
2019    }
2020    let new_col_map = ColumnMap::new(new_col_map);
2021    outer.let_in(id_gen, |id_gen, get_outer| {
2022        let keyed_outer = if key.is_empty() {
2023            // Don't depend on outer at all if the branch is not correlated,
2024            // which yields vastly better query plans. Note that this is a bit
2025            // weird in that the branch will be computed even if outer has no
2026            // rows, whereas if it had been correlated it would not (and *could*
2027            // not) have been computed if outer had no rows, but the callers of
2028            // this function don't mind these somewhat-weird semantics.
2029            MirRelationExpr::constant(vec![vec![]], ReprRelationType::new(vec![]))
2030        } else {
2031            get_outer.clone().distinct_by(key.clone())
2032        };
2033        keyed_outer.let_in(id_gen, |id_gen, get_keyed_outer| {
2034            let oa = get_outer.arity();
2035            let branch = apply(
2036                id_gen,
2037                inner,
2038                get_keyed_outer,
2039                &new_col_map,
2040                cte_map,
2041                context,
2042            )?;
2043            let ba = branch.arity();
2044            let joined = MirRelationExpr::join(
2045                vec![get_outer.clone(), branch],
2046                key.iter()
2047                    .enumerate()
2048                    .map(|(i, &k)| vec![(0, k), (1, i)])
2049                    .collect(),
2050            )
2051            // throw away the right-hand copy of the key we just joined on
2052            .project((0..oa).chain((oa + key.len())..(oa + ba)).collect());
2053            Ok(joined)
2054        })
2055    })
2056}
2057
2058fn apply_scalar_subquery(
2059    id_gen: &mut mz_ore::id_gen::IdGen,
2060    outer: MirRelationExpr,
2061    col_map: &ColumnMap,
2062    cte_map: &mut CteMap,
2063    scalar_subquery: HirRelationExpr,
2064    apply_requires_distinct_outer: bool,
2065    context: &Context,
2066) -> Result<MirRelationExpr, PlanError> {
2067    branch(
2068        id_gen,
2069        outer,
2070        col_map,
2071        cte_map,
2072        scalar_subquery,
2073        apply_requires_distinct_outer,
2074        context,
2075        |id_gen, expr, get_inner, col_map, cte_map, context| {
2076            // compute for every row in get_inner
2077            let select = expr.applied_to(id_gen, get_inner.clone(), col_map, cte_map, context)?;
2078            let col_type = select.sql_typ().column_types.into_last();
2079
2080            let inner_arity = get_inner.arity();
2081            // We must determine a count for each `get_inner` prefix,
2082            // and report an error if that count exceeds one.
2083            let guarded = select.let_in(id_gen, |_id_gen, get_select| {
2084                // Count for each `get_inner` prefix.
2085                let counts = get_select.clone().reduce(
2086                    (0..inner_arity).collect::<Vec<_>>(),
2087                    vec![mz_expr::AggregateExpr {
2088                        func: mz_expr::AggregateFunc::Count,
2089                        expr: MirScalarExpr::literal_true(),
2090                        distinct: false,
2091                    }],
2092                    None,
2093                );
2094
2095                // Errors should result from counts > 1.
2096                let errors = counts
2097                    .flat_map(
2098                        mz_expr::TableFunc::GuardSubquerySize {
2099                            column_type: col_type.clone().scalar_type,
2100                        },
2101                        vec![MirScalarExpr::column(inner_arity)],
2102                    )
2103                    .project(
2104                        (0..inner_arity)
2105                            .chain(Some(inner_arity + 1))
2106                            .collect::<Vec<_>>(),
2107                    );
2108                // Return `get_select` and any errors added in.
2109                Ok::<_, PlanError>(get_select.union(errors))
2110            })?;
2111            // append Null to anything that didn't return any rows
2112            let default = vec![(Datum::Null, ReprScalarType::from(&col_type.scalar_type))];
2113            get_inner.lookup(id_gen, guarded, default)
2114        },
2115    )
2116}
2117
2118fn apply_existential_subquery(
2119    id_gen: &mut mz_ore::id_gen::IdGen,
2120    outer: MirRelationExpr,
2121    col_map: &ColumnMap,
2122    cte_map: &mut CteMap,
2123    subquery_expr: HirRelationExpr,
2124    apply_requires_distinct_outer: bool,
2125    context: &Context,
2126) -> Result<MirRelationExpr, PlanError> {
2127    branch(
2128        id_gen,
2129        outer,
2130        col_map,
2131        cte_map,
2132        subquery_expr,
2133        apply_requires_distinct_outer,
2134        context,
2135        |id_gen, expr, get_inner, col_map, cte_map, context| {
2136            let exists = expr
2137                // compute for every row in get_inner
2138                .applied_to(id_gen, get_inner.clone(), col_map, cte_map, context)?
2139                // throw away actual values and just remember whether or not there were __any__ rows
2140                .distinct_by((0..get_inner.arity()).collect())
2141                // Append true to anything that returned any rows.
2142                .map(vec![MirScalarExpr::literal_true()]);
2143
2144            // append False to anything that didn't return any rows
2145            get_inner.lookup(id_gen, exists, vec![(Datum::False, ReprScalarType::Bool)])
2146        },
2147    )
2148}
2149
2150impl AggregateExpr {
2151    fn applied_to(
2152        self,
2153        id_gen: &mut mz_ore::id_gen::IdGen,
2154        col_map: &ColumnMap,
2155        cte_map: &mut CteMap,
2156        inner: &mut MirRelationExpr,
2157        context: &Context,
2158    ) -> Result<mz_expr::AggregateExpr, PlanError> {
2159        let AggregateExpr {
2160            func,
2161            expr,
2162            distinct,
2163        } = self;
2164
2165        Ok(mz_expr::AggregateExpr {
2166            func: func.into_expr(),
2167            expr: expr.applied_to(id_gen, col_map, cte_map, inner, &None, context)?,
2168            distinct,
2169        })
2170    }
2171}
2172
2173/// Attempts an efficient outer join, if `on` has equijoin structure.
2174///
2175/// Both `left` and `right` are decorrelated inputs.
2176///
2177/// The first `oa` columns correspond to an outer context: we should do the
2178/// outer join independently for each prefix. In the case that `on` contains
2179/// just some equality tests between columns of `left` and `right` and some
2180/// local predicates, we can employ a relatively simple plan.
2181///
2182/// The last `on_subquery_types.len()` columns correspond to results from
2183/// subqueries defined in the `on` clause - we treat those as theta-join
2184/// conditions that prohibit the use of the simple plan attempted here.
2185fn attempt_outer_equijoin(
2186    left: MirRelationExpr,
2187    right: MirRelationExpr,
2188    on: MirScalarExpr,
2189    on_subquery_types: Vec<ReprColumnType>,
2190    kind: JoinKind,
2191    oa: usize,
2192    id_gen: &mut mz_ore::id_gen::IdGen,
2193    context: &Context,
2194) -> Result<Option<MirRelationExpr>, PlanError> {
2195    // TODO(database-issues#6827): In theory, we can be smarter and also handle `on`
2196    // predicates that reference subqueries as long as these subqueries don't
2197    // reference `left` and `right` at the same time.
2198    //
2199    // TODO(database-issues#6828): This code can be improved as follows:
2200    //
2201    // 1. Move the `canonicalize_predicates(...)` call to `applied_to`.
2202    // 2. Use the canonicalized `on` predicate in the non-equijoin based
2203    //    lowering strategy.
2204    // 3. Move the `OnPredicates::new(...)` call to `applied_to`.
2205    // 4. Pass the classified `OnPredicates` as a parameter.
2206    // 5. Guard calls of this function with `on_predicates.is_equijoin()`.
2207    //
2208    // Steps (1 + 2) require further investigation because we might change the
2209    // error semantics in case the `on` predicate contains a literal error..
2210
2211    let l_type = left.typ();
2212    let r_type = right.typ();
2213    let la = l_type.column_types.len() - oa;
2214    let ra = r_type.column_types.len() - oa;
2215    let sa = on_subquery_types.len();
2216
2217    // The output type contains [outer, left, right, sa] attributes.
2218    let mut output_type = Vec::with_capacity(oa + la + ra + sa);
2219    output_type.extend(l_type.column_types);
2220    output_type.extend(r_type.column_types.into_iter().skip(oa));
2221    output_type.extend(on_subquery_types);
2222
2223    // Generally healthy to do, but specifically `USING` conditions sometimes
2224    // put an `AND true` at the end of the `ON` condition.
2225    //
2226    // TODO(aalexandrov): maybe we should already be doing this in `applied_to`.
2227    // However, in that case it's not clear that we won't see regressions if
2228    // `on` simplifies to a literal error.
2229    let mut on = vec![on];
2230    mz_expr::canonicalize::canonicalize_predicates(&mut on, &output_type);
2231
2232    // Form the left and right types without the outer attributes.
2233    output_type.drain(0..oa);
2234    let lt = output_type.drain(0..la).collect_vec();
2235    let rt = output_type.drain(0..ra).collect_vec();
2236    assert!(output_type.len() == sa);
2237
2238    let on_predicates = OnPredicates::new(oa, la, ra, sa, on.clone(), context);
2239    if !on_predicates.is_equijoin(context) {
2240        return Ok(None);
2241    }
2242
2243    // If we've gotten this far, we can do the clever thing.
2244    // We'll want to use left and right multiple times
2245    let result = left.let_in(id_gen, |id_gen, get_left| {
2246        right.let_in(id_gen, |id_gen, get_right| {
2247            // TODO: we know that we can re-use the arrangements of left and right
2248            // needed for the inner join with each of the conditional outer joins.
2249            // It is not clear whether we should hint that, or just let the planner
2250            // and optimizer run and see what happens.
2251
2252            // We'll want the inner join (minus repeated columns)
2253            let join = MirRelationExpr::join(
2254                vec![get_left.clone(), get_right.clone()],
2255                (0..oa).map(|i| vec![(0, i), (1, i)]).collect(),
2256            )
2257            // remove those columns from `right` repeating the first `oa` columns.
2258            .project(
2259                (0..(oa + la))
2260                    .chain((oa + la + oa)..(oa + la + oa + ra))
2261                    .collect(),
2262            )
2263            // apply the filter constraints here, to ensure nulls are not matched.
2264            .filter(on);
2265
2266            // We'll want to re-use the results of the join multiple times.
2267            join.let_in(id_gen, |id_gen, get_join| {
2268                let mut result = get_join.clone();
2269
2270                // A collection of keys present in both left and right collections.
2271                let join_keys = on_predicates.join_keys();
2272                let both_keys_arity = join_keys.len();
2273                let both_keys = get_join.restrict(join_keys).distinct();
2274
2275                // The plan is now to determine the left and right rows matched in the
2276                // inner join, subtract them from left and right respectively, pad what
2277                // remains with nulls, and fold them in to `result`.
2278
2279                both_keys.let_in(id_gen, |_id_gen, get_both| {
2280                    if let JoinKind::LeftOuter { .. } | JoinKind::FullOuter = kind {
2281                        // Rows in `left` matched in the inner equijoin. This is
2282                        // a semi-join between `left` and `both_keys`.
2283                        let left_present = MirRelationExpr::join_scalars(
2284                            vec![
2285                                get_left
2286                                    .clone()
2287                                    // Push local predicates.
2288                                    .filter(on_predicates.lhs()),
2289                                get_both.clone(),
2290                            ],
2291                            itertools::zip_eq(
2292                                on_predicates.eq_lhs(),
2293                                (0..both_keys_arity).map(|k| MirScalarExpr::column(oa + la + k)),
2294                            )
2295                            .map(|(l_key, b_key)| [l_key, b_key].to_vec())
2296                            .collect(),
2297                        )
2298                        .project((0..(oa + la)).collect());
2299
2300                        // Determine the types of nulls to use as filler.
2301                        let right_fill = rt
2302                            .into_iter()
2303                            .map(|typ| MirScalarExpr::literal_null(typ.scalar_type))
2304                            .collect();
2305                        // Add to `result` absent elements, filled with typed nulls.
2306                        result = left_present
2307                            .negate()
2308                            .union(get_left.clone())
2309                            .map(right_fill)
2310                            .union(result);
2311                    }
2312
2313                    if let JoinKind::RightOuter | JoinKind::FullOuter = kind {
2314                        // Rows in `right` matched in the inner equijoin. This
2315                        // is a semi-join between `right` and `both_keys`.
2316                        let right_present = MirRelationExpr::join_scalars(
2317                            vec![
2318                                get_right
2319                                    .clone()
2320                                    // Push local predicates.
2321                                    .filter(on_predicates.rhs()),
2322                                get_both,
2323                            ],
2324                            itertools::zip_eq(
2325                                on_predicates.eq_rhs(),
2326                                (0..both_keys_arity).map(|k| MirScalarExpr::column(oa + ra + k)),
2327                            )
2328                            .map(|(r_key, b_key)| [r_key, b_key].to_vec())
2329                            .collect(),
2330                        )
2331                        .project((0..(oa + ra)).collect());
2332
2333                        // Determine the types of nulls to use as filler.
2334                        let left_fill = lt
2335                            .into_iter()
2336                            .map(|typ| MirScalarExpr::literal_null(typ.scalar_type))
2337                            .collect();
2338
2339                        // Add to `result` absent elements, prepended with typed nulls.
2340                        result = right_present
2341                            .negate()
2342                            .union(get_right.clone())
2343                            .map(left_fill)
2344                            // Permute left fill before right values.
2345                            .project(
2346                                itertools::chain!(
2347                                    0..oa,                 // Preserve `outer`.
2348                                    oa + ra..oa + la + ra, // Increment the next `la` cols by `ra`.
2349                                    oa..oa + ra            // Decrement the next `ra` cols by `la`.
2350                                )
2351                                .collect(),
2352                            )
2353                            .union(result)
2354                    }
2355
2356                    Ok::<_, PlanError>(result)
2357                })
2358            })
2359        })
2360    })?;
2361    Ok(Some(result))
2362}
2363
2364/// A struct that represents the predicates in the `on` clause in a form
2365/// suitable for efficient planning outer joins with equijoin predicates.
2366struct OnPredicates {
2367    /// A store for classified `ON` predicates.
2368    ///
2369    /// Predicates that reference a single side are adjusted to assume an
2370    /// `outer × <side>` schema.
2371    predicates: Vec<OnPredicate>,
2372    /// Number of outer context columns.
2373    oa: usize,
2374}
2375
2376impl OnPredicates {
2377    const I_OUT: usize = 0; // outer context input position
2378    const I_LHS: usize = 1; // lhs input position
2379    const I_RHS: usize = 2; // rhs input position
2380    const I_SUB: usize = 3; // on subqueries input position
2381
2382    /// Classify the predicates in the `on` clause of an outer join.
2383    ///
2384    /// The other parameters are arities of the input parts:
2385    ///
2386    /// - `oa` is the arity of the `outer` context.
2387    /// - `la` is the arity of the `left` input.
2388    /// - `ra` is the arity of the `right` input.
2389    /// - `sa` is the arity of the `on` subqueries.
2390    ///
2391    /// The constructor assumes that:
2392    ///
2393    /// 1. The `on` parameter will be applied on a result that has the following
2394    ///    schema `outer × left × right × on_subqueries`.
2395    /// 2. The `on` parameter is already adjusted to assume that schema.
2396    /// 3. The `on` parameter is obtained by canonicalizing the original `on:
2397    ///    MirScalarExpr` with `canonicalize_predicates`.
2398    fn new(
2399        oa: usize,
2400        la: usize,
2401        ra: usize,
2402        sa: usize,
2403        on: Vec<MirScalarExpr>,
2404        _context: &Context,
2405    ) -> Self {
2406        use mz_expr::BinaryFunc::Eq;
2407
2408        // Re-bind those locally for more compact pattern matching.
2409        const I_LHS: usize = OnPredicates::I_LHS;
2410        const I_RHS: usize = OnPredicates::I_RHS;
2411
2412        // Self parameters.
2413        let mut predicates = Vec::with_capacity(on.len());
2414
2415        // Helpers for populating `predicates`.
2416        let inner_join_mapper = mz_expr::JoinInputMapper::new_from_input_arities([oa, la, ra, sa]);
2417        let rhs_permutation = itertools::chain!(0..oa + la, oa..oa + ra).collect::<Vec<_>>();
2418        let lookup_inputs = |expr: &MirScalarExpr| -> Vec<usize> {
2419            inner_join_mapper
2420                .lookup_inputs(expr)
2421                .filter(|&i| i != Self::I_OUT)
2422                .collect()
2423        };
2424        let has_subquery_refs = |expr: &MirScalarExpr| -> bool {
2425            inner_join_mapper
2426                .lookup_inputs(expr)
2427                .any(|i| i == Self::I_SUB)
2428        };
2429
2430        // Iterate over `on` elements and populate `predicates`.
2431        for mut predicate in on {
2432            if predicate.might_error() {
2433                tracing::debug!(case = "thetajoin (error)", "OnPredicates::new");
2434                // Treat predicates that can produce a literal error as Theta.
2435                predicates.push(OnPredicate::Theta(predicate));
2436            } else if has_subquery_refs(&predicate) {
2437                tracing::debug!(case = "thetajoin (subquery)", "OnPredicates::new");
2438                // Treat predicates referencing an `on` subquery as Theta.
2439                predicates.push(OnPredicate::Theta(predicate));
2440            } else if let MirScalarExpr::CallBinary {
2441                func: Eq(_),
2442                expr1,
2443                expr2,
2444            } = &mut predicate
2445            {
2446                // Obtain the non-outer inputs referenced by each side.
2447                let inputs1 = lookup_inputs(expr1);
2448                let inputs2 = lookup_inputs(expr2);
2449
2450                match (&inputs1[..], &inputs2[..]) {
2451                    // Neither side references an input. This could be a
2452                    // constant expression or an expression that depends only on
2453                    // the outer context.
2454                    ([], []) => {
2455                        predicates.push(OnPredicate::Const(predicate));
2456                    }
2457                    // Both sides reference different inputs.
2458                    ([I_LHS], [I_RHS]) => {
2459                        let lhs = expr1.take();
2460                        let mut rhs = expr2.take();
2461                        rhs.permute(&rhs_permutation);
2462                        predicates.push(OnPredicate::Eq(lhs.clone(), rhs.clone()));
2463                        predicates.push(OnPredicate::LhsConsequence(lhs.call_is_null().not()));
2464                        predicates.push(OnPredicate::RhsConsequence(rhs.call_is_null().not()));
2465                    }
2466                    // Both sides reference different inputs (swapped).
2467                    ([I_RHS], [I_LHS]) => {
2468                        let lhs = expr2.take();
2469                        let mut rhs = expr1.take();
2470                        rhs.permute(&rhs_permutation);
2471                        predicates.push(OnPredicate::Eq(lhs.clone(), rhs.clone()));
2472                        predicates.push(OnPredicate::LhsConsequence(lhs.call_is_null().not()));
2473                        predicates.push(OnPredicate::RhsConsequence(rhs.call_is_null().not()));
2474                    }
2475                    // Both sides reference the left input or no input.
2476                    ([I_LHS], [I_LHS]) | ([I_LHS], []) | ([], [I_LHS]) => {
2477                        predicates.push(OnPredicate::Lhs(predicate));
2478                    }
2479                    // Both sides reference the right input or no input.
2480                    ([I_RHS], [I_RHS]) | ([I_RHS], []) | ([], [I_RHS]) => {
2481                        predicate.permute(&rhs_permutation);
2482                        predicates.push(OnPredicate::Rhs(predicate));
2483                    }
2484                    // At least one side references more than one input.
2485                    _ => {
2486                        tracing::debug!(case = "thetajoin (eq)", "OnPredicates::new");
2487                        predicates.push(OnPredicate::Theta(predicate));
2488                    }
2489                }
2490            } else {
2491                // Obtain the non-outer inputs referenced by this predicate.
2492                let inputs = lookup_inputs(&predicate);
2493
2494                match &inputs[..] {
2495                    // The predicate references no inputs. This could be a
2496                    // constant expression or an expression that depends only on
2497                    // the outer context.
2498                    [] => {
2499                        predicates.push(OnPredicate::Const(predicate));
2500                    }
2501                    // The predicate references only the left input.
2502                    [I_LHS] => {
2503                        predicates.push(OnPredicate::Lhs(predicate));
2504                    }
2505                    // The predicate references only the right input.
2506                    [I_RHS] => {
2507                        predicate.permute(&rhs_permutation);
2508                        predicates.push(OnPredicate::Rhs(predicate));
2509                    }
2510                    // The predicate references both inputs.
2511                    _ => {
2512                        tracing::debug!(case = "thetajoin (non-eq)", "OnPredicates::new");
2513                        predicates.push(OnPredicate::Theta(predicate));
2514                    }
2515                }
2516            }
2517        }
2518
2519        Self { predicates, oa }
2520    }
2521
2522    /// Check if the predicates can be lowered with an equijoin-based strategy.
2523    fn is_equijoin(&self, context: &Context) -> bool {
2524        // Count each `OnPredicate` variant in `self.predicates`.
2525        let (const_cnt, lhs_cnt, rhs_cnt, eq_cnt, eq_cols, theta_cnt) =
2526            self.predicates.iter().fold(
2527                (0, 0, 0, 0, 0, 0),
2528                |(const_cnt, lhs_cnt, rhs_cnt, eq_cnt, eq_cols, theta_cnt), p| {
2529                    (
2530                        const_cnt + usize::from(matches!(p, OnPredicate::Const(..))),
2531                        lhs_cnt + usize::from(matches!(p, OnPredicate::Lhs(..))),
2532                        rhs_cnt + usize::from(matches!(p, OnPredicate::Rhs(..))),
2533                        eq_cnt + usize::from(matches!(p, OnPredicate::Eq(..))),
2534                        eq_cols
2535                            + usize::from(matches!(
2536                                p,
2537                                OnPredicate::Eq(lhs, rhs) if lhs.is_column() && rhs.is_column()
2538                            )),
2539                        theta_cnt + usize::from(matches!(p, OnPredicate::Theta(..))),
2540                    )
2541                },
2542            );
2543
2544        let is_equijion = if context.config.enable_new_outer_join_lowering {
2545            // New classifier.
2546            eq_cnt > 0 && theta_cnt == 0
2547        } else {
2548            // Old classifier.
2549            eq_cnt > 0 && eq_cnt == eq_cols && theta_cnt + const_cnt + lhs_cnt + rhs_cnt == 0
2550        };
2551
2552        // Log an entry only if this is an equijoin according to the new classifier.
2553        if eq_cnt > 0 && theta_cnt == 0 {
2554            tracing::debug!(
2555                const_cnt,
2556                lhs_cnt,
2557                rhs_cnt,
2558                eq_cnt,
2559                eq_cols,
2560                theta_cnt,
2561                "OnPredicates::is_equijoin"
2562            );
2563        }
2564
2565        is_equijion
2566    }
2567
2568    /// Return an [`MirRelationExpr`] list that represents the keys for the
2569    /// equijoin. The list will contain the outer columns as a prefix.
2570    fn join_keys(&self) -> JoinKeys {
2571        // We could return either the `lhs` or the `rhs` of the keys used to
2572        // form the inner join as they are equated by the join condition.
2573        let join_keys = self.eq_lhs().collect::<Vec<_>>();
2574
2575        if join_keys.iter().all(|k| k.is_column()) {
2576            tracing::debug!(case = "outputs", "OnPredicates::join_keys");
2577            JoinKeys::Outputs(join_keys.iter().flat_map(|k| k.as_column()).collect())
2578        } else {
2579            tracing::debug!(case = "scalars", "OnPredicates::join_keys");
2580            JoinKeys::Scalars(join_keys)
2581        }
2582    }
2583
2584    /// Return an iterator over the left-hand sides of all [`OnPredicate::Eq`]
2585    /// conditions in the predicates list.
2586    ///
2587    /// The iterator will start with column references to the outer columns as a
2588    /// prefix.
2589    fn eq_lhs(&self) -> impl Iterator<Item = MirScalarExpr> + '_ {
2590        itertools::chain(
2591            (0..self.oa).map(MirScalarExpr::column),
2592            self.predicates.iter().filter_map(|e| match e {
2593                OnPredicate::Eq(lhs, _) => Some(lhs.clone()),
2594                _ => None,
2595            }),
2596        )
2597    }
2598
2599    /// Return an iterator over the right-hand sides of all [`OnPredicate::Eq`]
2600    /// conditions in the predicates list.
2601    ///
2602    /// The iterator will start with column references to the outer columns as a
2603    /// prefix.
2604    fn eq_rhs(&self) -> impl Iterator<Item = MirScalarExpr> + '_ {
2605        itertools::chain(
2606            (0..self.oa).map(MirScalarExpr::column),
2607            self.predicates.iter().filter_map(|e| match e {
2608                OnPredicate::Eq(_, rhs) => Some(rhs.clone()),
2609                _ => None,
2610            }),
2611        )
2612    }
2613
2614    /// Return an iterator over the [`OnPredicate::Lhs`], [`OnPredicate::LhsConsequence`] and
2615    /// [`OnPredicate::Const`] conditions in the predicates list.
2616    fn lhs(&self) -> impl Iterator<Item = MirScalarExpr> + '_ {
2617        self.predicates.iter().filter_map(|p| match p {
2618            // We treat Const predicates local to both inputs.
2619            OnPredicate::Const(p) => Some(p.clone()),
2620            OnPredicate::Lhs(p) => Some(p.clone()),
2621            OnPredicate::LhsConsequence(p) => Some(p.clone()),
2622            _ => None,
2623        })
2624    }
2625
2626    /// Return an iterator over the [`OnPredicate::Rhs`], [`OnPredicate::RhsConsequence`] and
2627    /// [`OnPredicate::Const`] conditions in the predicates list.
2628    fn rhs(&self) -> impl Iterator<Item = MirScalarExpr> + '_ {
2629        self.predicates.iter().filter_map(|p| match p {
2630            // We treat Const predicates local to both inputs.
2631            OnPredicate::Const(p) => Some(p.clone()),
2632            OnPredicate::Rhs(p) => Some(p.clone()),
2633            OnPredicate::RhsConsequence(p) => Some(p.clone()),
2634            _ => None,
2635        })
2636    }
2637}
2638
2639enum OnPredicate {
2640    // A predicate that is either constant or references only outer columns.
2641    Const(MirScalarExpr),
2642    // A local predicate on the left-hand side of the join, i.e., it references only the left input
2643    // and possibly outer columns.
2644    //
2645    // This is one of the original predicates from the ON clause.
2646    //
2647    // One _must_ apply this predicate.
2648    Lhs(MirScalarExpr),
2649    // A local predicate on the left-hand side of the join, i.e., it references only the left input
2650    // and possibly outer columns.
2651    //
2652    // This is not one of the original predicates from the ON clause, but is just a consequence
2653    // of an original predicate in the ON clause, where the original predicate references both
2654    // inputs, but the consequence references only the left input.
2655    //
2656    // For example, the original predicate `input1.x = input2.a` has the consequence
2657    // `input1.x IS NOT NULL`. Applying such a consequence before the input is fed into the join
2658    // prevents null skew, and also makes more CSE opportunities available when the left input's key
2659    // doesn't have a NOT NULL constraint, saving us an arrangement.
2660    //
2661    // Applying the predicate is optional, because the original predicate will be applied anyway.
2662    LhsConsequence(MirScalarExpr),
2663    // A local predicate on the right-hand side of the join.
2664    //
2665    // This is one of the original predicates from the ON clause.
2666    //
2667    // One _must_ apply this predicate.
2668    Rhs(MirScalarExpr),
2669    // A consequence of an original ON predicate, see above.
2670    RhsConsequence(MirScalarExpr),
2671    // An equality predicate between the two sides.
2672    Eq(MirScalarExpr, MirScalarExpr),
2673    // a non-equality predicate between the two sides.
2674    #[allow(dead_code)]
2675    Theta(MirScalarExpr),
2676}
2677
2678/// A set of join keys referencing an input.
2679///
2680/// This is used in the [`MirRelationExpr::Join`] lowering code in order to
2681/// avoid changes (and thereby possible regressions) in plans that have equijoin
2682/// predicates consisting only of column refs.
2683///
2684/// If we were running `CanonicalizeMfp` as part of `NormalizeOps` we might be
2685/// able to get rid of this code, but as it stands `Map` simplification seems
2686/// more cumbersome than `Project` simplification, so do this just to be sure.
2687enum JoinKeys {
2688    // A predicate that is either constant or references only outer columns.
2689    Outputs(Vec<usize>),
2690    // A local predicate on the left-hand side of the join.
2691    Scalars(Vec<MirScalarExpr>),
2692}
2693
2694impl JoinKeys {
2695    fn len(&self) -> usize {
2696        match self {
2697            JoinKeys::Outputs(outputs) => outputs.len(),
2698            JoinKeys::Scalars(scalars) => scalars.len(),
2699        }
2700    }
2701}
2702
2703/// Extension methods for [`MirRelationExpr`] required in the HIR ⇒ MIR lowering
2704/// code.
2705trait LoweringExt {
2706    /// See [`MirRelationExpr::restrict`].
2707    fn restrict(self, join_keys: JoinKeys) -> Self;
2708}
2709
2710impl LoweringExt for MirRelationExpr {
2711    /// Restrict the set of columns of an input to the sequence of [`JoinKeys`].
2712    fn restrict(self, join_keys: JoinKeys) -> Self {
2713        let num_keys = join_keys.len();
2714        match join_keys {
2715            JoinKeys::Outputs(outputs) => self.project(outputs),
2716            JoinKeys::Scalars(scalars) => {
2717                let input_arity = self.arity();
2718                let outputs = (input_arity..input_arity + num_keys).collect();
2719                self.map(scalars).project(outputs)
2720            }
2721        }
2722    }
2723}