Skip to main content

mz_sql/plan/
query.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//! SQL `Query`s are the declarative, computational part of SQL.
11//! This module turns `Query`s into `HirRelationExpr`s - a more explicit, algebraic way of
12//! describing computation.
13
14//! Functions named plan_* are typically responsible for handling a single node of the SQL ast.
15//! E.g. `plan_query` is responsible for handling `sqlparser::ast::Query`.
16//! plan_* functions which correspond to operations on relations typically return a `HirRelationExpr`.
17//! plan_* functions which correspond to operations on scalars typically return a `HirScalarExpr`
18//! and a `SqlScalarType`. (The latter is because it's not always possible to infer from a
19//! `HirScalarExpr` what the intended type is - notably in the case of decimals where the
20//! scale/precision are encoded only in the type).
21
22//! Aggregates are particularly twisty.
23//!
24//! In SQL, a GROUP BY turns any columns not in the group key into vectors of
25//! values. Then anywhere later in the scope, an aggregate function can be
26//! applied to that group. Inside the arguments of an aggregate function, other
27//! normal functions are applied element-wise over the vectors. Thus, `SELECT
28//! sum(foo.x + foo.y) FROM foo GROUP BY x` means adding the scalar `x` to the
29//! vector `y` and summing the results.
30//!
31//! In `HirRelationExpr`, aggregates can only be applied immediately at the time
32//! of grouping.
33//!
34//! To deal with this, whenever we see a SQL GROUP BY we look ahead for
35//! aggregates and precompute them in the `HirRelationExpr::Reduce`. When we
36//! reach the same aggregates during normal planning later on, we look them up
37//! in an `ExprContext` to find the precomputed versions.
38
39use std::borrow::Cow;
40use std::cell::RefCell;
41use std::collections::{BTreeMap, BTreeSet};
42use std::convert::{TryFrom, TryInto};
43use std::num::NonZeroU64;
44use std::rc::Rc;
45use std::sync::{Arc, LazyLock};
46use std::{iter, mem};
47
48use itertools::Itertools;
49use mz_expr::func::variadic::{
50    ArrayCreate, ArrayIndex, Coalesce, Greatest, Least, ListCreate, ListIndex, ListSliceLinear,
51    MapBuild, RecordCreate,
52};
53use mz_expr::virtual_syntax::AlgExcept;
54use mz_expr::{
55    Eval, Id, LetRecLimit, LocalId, MapFilterProject, MirScalarExpr, REPEAT_ROW_NAME,
56    RowSetFinishing, TableFunc, func as expr_func,
57};
58use mz_ore::collections::CollectionExt;
59use mz_ore::error::ErrorExt;
60use mz_ore::id_gen::IdGen;
61use mz_ore::option::FallibleMapExt;
62use mz_ore::stack::{CheckedRecursion, RecursionGuard};
63use mz_ore::str::StrExt;
64use mz_repr::adt::char::CharLength;
65use mz_repr::adt::numeric::{NUMERIC_DATUM_MAX_PRECISION, NumericMaxScale};
66use mz_repr::adt::timestamp::TimestampPrecision;
67use mz_repr::adt::varchar::VarCharMaxLength;
68use mz_repr::namespaces::MZ_CATALOG_SCHEMA;
69use mz_repr::{
70    CatalogItemId, ColumnIndex, ColumnName, Datum, RelationDesc, RelationVersionSelector,
71    ReprColumnType, Row, RowArena, SqlColumnType, SqlRelationType, SqlScalarType,
72    UNKNOWN_COLUMN_NAME, strconv,
73};
74use mz_sql_parser::ast::display::AstDisplay;
75use mz_sql_parser::ast::visit::Visit;
76use mz_sql_parser::ast::visit_mut::{self, VisitMut};
77use mz_sql_parser::ast::{
78    AsOf, Assignment, AstInfo, CastFailureMode, CreateWebhookSourceBody, CreateWebhookSourceCheck,
79    CreateWebhookSourceHeader, CreateWebhookSourceSecret, CteBlock, DeleteStatement, Distinct,
80    Expr, Function, FunctionArgs, HomogenizingFunction, Ident, InsertSource, IsExprConstruct, Join,
81    JoinConstraint, JoinOperator, Limit, MapEntry, MutRecBlock, MutRecBlockOption,
82    MutRecBlockOptionName, OrderByExpr, Query, Select, SelectItem, SelectOption, SelectOptionName,
83    SetExpr, SetOperator, ShowStatement, SubscriptPosition, TableAlias, TableFactor,
84    TableWithJoins, UnresolvedItemName, UpdateStatement, Value, Values, WindowFrame,
85    WindowFrameBound, WindowFrameUnits, WindowSpec, visit,
86};
87use mz_sql_parser::ident;
88
89use crate::catalog::{CatalogItemType, CatalogType, SessionCatalog};
90use crate::func::{self, Func, FuncSpec, TableFuncImpl};
91use crate::names::{
92    Aug, FullItemName, PartialItemName, ResolvedDataType, ResolvedItemName, SchemaSpecifier,
93};
94use crate::plan::PlanError::InvalidWmrRecursionLimit;
95use crate::plan::error::PlanError;
96use crate::plan::hir::{
97    AbstractColumnType, AbstractExpr, AggregateExpr, AggregateFunc, AggregateWindowExpr,
98    BinaryFunc, CoercibleScalarExpr, CoercibleScalarType, ColumnOrder, ColumnRef, Hir,
99    HirRelationExpr, HirScalarExpr, JoinKind, ScalarWindowExpr, ScalarWindowFunc, UnaryFunc,
100    ValueWindowExpr, ValueWindowFunc, VariadicFunc, WindowExpr, WindowExprType,
101};
102use crate::plan::plan_utils::{self, GroupSizeHints, JoinSide};
103use crate::plan::scope::{Scope, ScopeItem, ScopeUngroupedColumn};
104use crate::plan::statement::{StatementContext, StatementDesc, show};
105use crate::plan::typeconv::{self, CastContext, plan_hypothetical_cast};
106use crate::plan::{
107    Params, PlanContext, QueryWhen, ShowCreatePlan, WebhookValidation, WebhookValidationSecret,
108    literal, transform_ast,
109};
110use crate::session::vars::ENABLE_WITH_ORDINALITY_LEGACY_FALLBACK;
111use crate::session::vars::{self, FeatureFlag};
112use crate::{ORDINALITY_COL_NAME, normalize};
113
114#[derive(Debug)]
115pub struct PlannedRootQuery<E> {
116    pub expr: E,
117    pub desc: RelationDesc,
118    pub finishing: RowSetFinishing<HirScalarExpr, HirScalarExpr>,
119    pub scope: Scope,
120}
121
122/// Plans a top-level query, returning the `HirRelationExpr` describing the query
123/// plan, the `RelationDesc` describing the shape of the result set, a
124/// `RowSetFinishing` describing post-processing that must occur before results
125/// are sent to the client, and the types of the parameters in the query, if any
126/// were present.
127///
128/// Note that the returned `RelationDesc` describes the expression after
129/// applying the returned `RowSetFinishing`.
130#[mz_ore::instrument(target = "compiler", level = "trace", name = "ast_to_hir")]
131pub fn plan_root_query(
132    scx: &StatementContext,
133    mut query: Query<Aug>,
134    lifetime: QueryLifetime,
135) -> Result<PlannedRootQuery<HirRelationExpr>, PlanError> {
136    transform_ast::transform(scx, &mut query)?;
137    let mut qcx = QueryContext::root(scx, lifetime);
138    let PlannedQuery {
139        mut expr,
140        scope,
141        order_by,
142        limit,
143        offset,
144        project,
145        group_size_hints,
146    } = plan_query(&mut qcx, &query)?;
147
148    let mut finishing = RowSetFinishing {
149        limit,
150        offset,
151        project,
152        order_by,
153    };
154
155    // Attempt to push the finishing's ordering past its projection. This allows
156    // data to be projected down on the workers rather than the coordinator. It
157    // also improves the optimizer's demand analysis, as the optimizer can only
158    // reason about demand information in `expr` (i.e., it can't see
159    // `finishing.project`).
160    try_push_projection_order_by(&mut expr, &mut finishing.project, &mut finishing.order_by);
161
162    if lifetime.is_maintained() {
163        expr.finish_maintained(&mut finishing, group_size_hints);
164    }
165
166    let typ = qcx.relation_type(&expr);
167    let typ = SqlRelationType::new(
168        finishing
169            .project
170            .iter()
171            .map(|i| typ.column_types[*i].clone())
172            .collect(),
173    );
174    let desc = RelationDesc::new(typ, scope.column_names());
175
176    Ok(PlannedRootQuery {
177        expr,
178        desc,
179        finishing,
180        scope,
181    })
182}
183
184/// Attempts to push a projection through an order by.
185///
186/// The returned bool indicates whether the pushdown was successful or not.
187/// Successful pushdown requires that all the columns referenced in `order_by`
188/// are included in `project`.
189///
190/// When successful, `expr` is wrapped in a projection node, `order_by` is
191/// rewritten to account for the pushed-down projection, and `project` is
192/// replaced with the trivial projection. When unsuccessful, no changes are made
193/// to any of the inputs.
194fn try_push_projection_order_by(
195    expr: &mut HirRelationExpr,
196    project: &mut Vec<usize>,
197    order_by: &mut Vec<ColumnOrder>,
198) -> bool {
199    let mut unproject = vec![None; expr.arity()];
200    for (out_i, in_i) in project.iter().copied().enumerate() {
201        unproject[in_i] = Some(out_i);
202    }
203    if order_by
204        .iter()
205        .all(|ob| ob.column < unproject.len() && unproject[ob.column].is_some())
206    {
207        let trivial_project = (0..project.len()).collect();
208        *expr = expr.take().project(mem::replace(project, trivial_project));
209        for ob in order_by {
210            ob.column = unproject[ob.column].unwrap();
211        }
212        true
213    } else {
214        false
215    }
216}
217
218pub fn plan_insert_query(
219    scx: &StatementContext,
220    table_name: ResolvedItemName,
221    columns: Vec<Ident>,
222    source: InsertSource<Aug>,
223    returning: Vec<SelectItem<Aug>>,
224) -> Result<
225    (
226        CatalogItemId,
227        HirRelationExpr,
228        PlannedRootQuery<Vec<HirScalarExpr>>,
229    ),
230    PlanError,
231> {
232    let mut qcx = QueryContext::root(scx, QueryLifetime::OneShot);
233    let table = scx.get_item_by_resolved_name(&table_name)?;
234
235    // Validate the target of the insert.
236    if table.item_type() != CatalogItemType::Table {
237        sql_bail!(
238            "cannot insert into {} '{}'",
239            table.item_type(),
240            table_name.full_name_str()
241        );
242    }
243    let desc = table
244        .relation_desc()
245        .ok_or_else(|| sql_err!("item does not have a relation description"))?;
246    let mut defaults = table
247        .writable_table_details()
248        .ok_or_else(|| {
249            sql_err!(
250                "cannot insert into non-writeable table '{}'",
251                table_name.full_name_str()
252            )
253        })?
254        .to_vec();
255
256    for default in &mut defaults {
257        transform_ast::transform(scx, default)?;
258    }
259
260    if table.id().is_system() {
261        sql_bail!(
262            "cannot insert into system table '{}'",
263            table_name.full_name_str()
264        );
265    }
266
267    let columns: Vec<_> = columns.into_iter().map(normalize::column_name).collect();
268
269    // Validate target column order.
270    let mut source_types = Vec::with_capacity(columns.len());
271    let mut ordering = Vec::with_capacity(columns.len());
272
273    if columns.is_empty() {
274        // Columns in source query must be in order. Let's guess the full shape and truncate to the
275        // right size later after planning the source query
276        source_types.extend(desc.iter_types().map(|x| &x.scalar_type));
277        ordering.extend(0..desc.arity());
278    } else {
279        let column_by_name: BTreeMap<&ColumnName, (usize, &SqlColumnType)> = desc
280            .iter()
281            .enumerate()
282            .map(|(idx, (name, typ))| (name, (idx, typ)))
283            .collect();
284
285        for c in &columns {
286            if let Some((idx, typ)) = column_by_name.get(c) {
287                ordering.push(*idx);
288                source_types.push(&typ.scalar_type);
289            } else {
290                sql_bail!(
291                    "column {} of relation {} does not exist",
292                    c.quoted(),
293                    table_name.full_name_str().quoted()
294                );
295            }
296        }
297        if let Some(dup) = columns.iter().duplicates().next() {
298            sql_bail!("column {} specified more than once", dup.quoted());
299        }
300    };
301
302    // Plan the source.
303    let expr = match source {
304        InsertSource::Query(mut query) => {
305            transform_ast::transform(scx, &mut query)?;
306
307            match query {
308                // Special-case simple VALUES clauses as PostgreSQL does.
309                Query {
310                    body: SetExpr::Values(Values(values)),
311                    ctes,
312                    order_by,
313                    limit: None,
314                    offset: None,
315                } if ctes.is_empty() && order_by.is_empty() => {
316                    let names: Vec<_> = ordering.iter().map(|i| desc.get_name(*i)).collect();
317                    plan_values_insert(&qcx, &names, &source_types, &values)?
318                }
319                _ => {
320                    let (expr, _scope) = plan_nested_query(&mut qcx, &query)?;
321                    expr
322                }
323            }
324        }
325        InsertSource::DefaultValues => {
326            HirRelationExpr::constant(vec![vec![]], SqlRelationType::empty())
327        }
328    };
329
330    let expr_arity = expr.arity();
331
332    // Validate that the arity of the source query is at most the size of declared columns or the
333    // size of the table if none are declared
334    let max_columns = if columns.is_empty() {
335        desc.arity()
336    } else {
337        columns.len()
338    };
339    if expr_arity > max_columns {
340        sql_bail!("INSERT has more expressions than target columns");
341    }
342    // But it should never have less than the declared columns (or zero)
343    if expr_arity < columns.len() {
344        sql_bail!("INSERT has more target columns than expressions");
345    }
346
347    // Trim now that we know for sure the correct arity of the source query
348    source_types.truncate(expr_arity);
349    ordering.truncate(expr_arity);
350
351    // Ensure the types of the source query match the types of the target table,
352    // installing assignment casts where necessary and possible.
353    let expr = cast_relation(&qcx, CastContext::Assignment, expr, source_types).map_err(|e| {
354        sql_err!(
355            "column {} is of type {} but expression is of type {}",
356            desc.get_name(ordering[e.column]).quoted(),
357            qcx.humanize_sql_scalar_type(&e.target_type, false),
358            qcx.humanize_sql_scalar_type(&e.source_type, false),
359        )
360    })?;
361
362    // Fill in any omitted columns and rearrange into correct order
363    let mut map_exprs = vec![];
364    let mut project_key = Vec::with_capacity(desc.arity());
365
366    // Maps from table column index to position in the source query
367    let col_to_source: BTreeMap<_, _> = ordering.iter().enumerate().map(|(a, b)| (b, a)).collect();
368
369    let column_details = desc.iter_types().zip_eq(defaults).enumerate();
370    for (col_idx, (col_typ, default)) in column_details {
371        if let Some(src_idx) = col_to_source.get(&col_idx) {
372            project_key.push(*src_idx);
373        } else {
374            let hir = plan_default_expr(scx, &default, &col_typ.scalar_type)?;
375            project_key.push(expr_arity + map_exprs.len());
376            map_exprs.push(hir);
377        }
378    }
379
380    let returning = {
381        let (scope, typ) = if let ResolvedItemName::Item {
382            full_name,
383            version: _,
384            ..
385        } = table_name
386        {
387            let scope = Scope::from_source(Some(full_name.clone().into()), desc.iter_names());
388            let typ = desc.typ().clone();
389            (scope, typ)
390        } else {
391            (Scope::empty(), SqlRelationType::empty())
392        };
393        let ecx = &ExprContext {
394            qcx: &qcx,
395            name: "RETURNING clause",
396            scope: &scope,
397            relation_type: &typ,
398            allow_aggregates: false,
399            allow_subqueries: false,
400            allow_parameters: true,
401            allow_windows: false,
402        };
403        let table_func_names = BTreeMap::new();
404        let mut output_columns = vec![];
405        let mut new_exprs = vec![];
406        let mut new_type = SqlRelationType::empty();
407        for mut si in returning {
408            transform_ast::transform(scx, &mut si)?;
409            for (select_item, column_name) in expand_select_item(ecx, &si, &table_func_names)? {
410                let expr = match &select_item {
411                    ExpandedSelectItem::InputOrdinal(i) => HirScalarExpr::column(*i),
412                    ExpandedSelectItem::Expr(expr) => plan_expr(ecx, expr)?.type_as_any(ecx)?,
413                };
414                output_columns.push(column_name);
415                let typ = ecx.column_type(&expr);
416                new_type.column_types.push(typ);
417                new_exprs.push(expr);
418            }
419        }
420        let desc = RelationDesc::new(new_type, output_columns);
421        let desc_arity = desc.arity();
422        PlannedRootQuery {
423            expr: new_exprs,
424            desc,
425            finishing: HirRelationExpr::trivial_row_set_finishing_hir(desc_arity),
426            scope,
427        }
428    };
429
430    Ok((
431        table.id(),
432        expr.map(map_exprs).project(project_key),
433        returning,
434    ))
435}
436
437/// Determines the mapping between some external data and a Materialize relation.
438///
439/// Returns the following:
440/// * [`CatalogItemId`] for the destination table.
441/// * [`RelationDesc`] representing the shape of the __input__ data we are copying from.
442/// * The [`ColumnIndex`]es that the source data maps to. TODO(cf2): We don't need this mapping
443///   since we now return a [`MapFilterProject`].
444/// * [`MapFilterProject`] which will map and project the input data to match the shape of the
445///   destination table.
446///
447pub fn plan_copy_item(
448    scx: &StatementContext,
449    item_name: ResolvedItemName,
450    columns: Vec<Ident>,
451) -> Result<
452    (
453        CatalogItemId,
454        RelationDesc,
455        Vec<ColumnIndex>,
456        Option<MapFilterProject>,
457    ),
458    PlanError,
459> {
460    let item = scx.get_item_by_resolved_name(&item_name)?;
461    let fullname = scx.catalog.resolve_full_name(item.name());
462    let table_desc = match item.relation_desc() {
463        Some(desc) => desc.into_owned(),
464        None => {
465            return Err(PlanError::InvalidDependency {
466                name: fullname.to_string(),
467                item_type: item.item_type().to_string(),
468            });
469        }
470    };
471    let mut ordering = Vec::with_capacity(columns.len());
472
473    // TODO(cf2): The logic here to create the `source_desc` and the MFP are a bit duplicated and
474    // should be simplified. The reason they are currently separate code paths is so we can roll
475    // out `COPY ... FROM <url>` without touching the current `COPY ... FROM ... STDIN` behavior.
476
477    // If we're copying data into a table that users can write into (e.g. not a `CREATE TABLE ...
478    // FROM SOURCE ...`), then we generate an MFP.
479    //
480    // Note: This method is called for both `COPY INTO <table> FROM` and `COPY <expr> TO <external>`
481    // so it's not always guaranteed that our `item` is a table.
482    let mfp = if let Some(table_defaults) = item.writable_table_details() {
483        let mut table_defaults = table_defaults.to_vec();
484
485        for default in &mut table_defaults {
486            transform_ast::transform(scx, default)?;
487        }
488
489        // Fill in any omitted columns and rearrange into correct order
490        let source_column_names: Vec<_> = columns
491            .iter()
492            .cloned()
493            .map(normalize::column_name)
494            .collect();
495
496        let mut default_exprs = Vec::new();
497        let mut project_keys = Vec::with_capacity(table_desc.arity());
498
499        // For each column in the destination table, either project it from the source data, or provide
500        // an expression to fill in a default value.
501        let column_details = table_desc.iter().zip_eq(table_defaults);
502        for ((col_name, col_type), col_default) in column_details {
503            let maybe_src_idx = source_column_names.iter().position(|name| name == col_name);
504            if let Some(src_idx) = maybe_src_idx {
505                project_keys.push(src_idx);
506            } else {
507                // If one a column from the table does not exist in the source data, then a default
508                // value will get appended to the end of the input Row from the source data.
509                let hir = plan_default_expr(scx, &col_default, &col_type.scalar_type)?;
510                let mir = hir.lower_uncorrelated(scx.catalog.system_vars())?;
511                project_keys.push(source_column_names.len() + default_exprs.len());
512                default_exprs.push(mir);
513            }
514        }
515
516        let mfp = MapFilterProject::new(source_column_names.len())
517            .map(default_exprs)
518            .project(project_keys);
519        Some(mfp)
520    } else {
521        None
522    };
523
524    // Create a mapping from input data to the table we're copying into.
525    let source_desc = if columns.is_empty() {
526        let indexes = (0..table_desc.arity()).map(ColumnIndex::from_raw);
527        ordering.extend(indexes);
528
529        // The source data should be in the same order as the table.
530        table_desc
531    } else {
532        let columns: Vec<_> = columns.into_iter().map(normalize::column_name).collect();
533        let column_by_name: BTreeMap<&ColumnName, (ColumnIndex, &SqlColumnType)> = table_desc
534            .iter_all()
535            .map(|(idx, name, typ)| (name, (*idx, typ)))
536            .collect();
537
538        let mut names = Vec::with_capacity(columns.len());
539        let mut source_types = Vec::with_capacity(columns.len());
540
541        for c in &columns {
542            if let Some((idx, typ)) = column_by_name.get(c) {
543                ordering.push(*idx);
544                source_types.push((*typ).clone());
545                names.push(c.clone());
546            } else {
547                sql_bail!(
548                    "column {} of relation {} does not exist",
549                    c.quoted(),
550                    item_name.full_name_str().quoted()
551                );
552            }
553        }
554        if let Some(dup) = columns.iter().duplicates().next() {
555            sql_bail!("column {} specified more than once", dup.quoted());
556        }
557
558        // The source data is a different shape than the destination table.
559        RelationDesc::new(SqlRelationType::new(source_types), names)
560    };
561
562    Ok((item.id(), source_desc, ordering, mfp))
563}
564
565/// See the doc comment on [`plan_copy_item`] for the details of what this function returns.
566///
567/// TODO(cf3): Merge this method with [`plan_copy_item`].
568pub fn plan_copy_from(
569    scx: &StatementContext,
570    table_name: ResolvedItemName,
571    columns: Vec<Ident>,
572) -> Result<
573    (
574        CatalogItemId,
575        RelationDesc,
576        Vec<ColumnIndex>,
577        Option<MapFilterProject>,
578    ),
579    PlanError,
580> {
581    let table = scx.get_item_by_resolved_name(&table_name)?;
582
583    // Validate the target of the insert.
584    if table.item_type() != CatalogItemType::Table {
585        sql_bail!(
586            "cannot insert into {} '{}'",
587            table.item_type(),
588            table_name.full_name_str()
589        );
590    }
591
592    let _ = table.writable_table_details().ok_or_else(|| {
593        sql_err!(
594            "cannot insert into non-writeable table '{}'",
595            table_name.full_name_str()
596        )
597    })?;
598
599    if table.id().is_system() {
600        sql_bail!(
601            "cannot insert into system table '{}'",
602            table_name.full_name_str()
603        );
604    }
605    let (id, desc, ordering, mfp) = plan_copy_item(scx, table_name, columns)?;
606
607    Ok((id, desc, ordering, mfp))
608}
609
610/// Builds a plan that adds the default values for the missing columns and re-orders
611/// the datums in the given rows to match the order in the target table.
612pub fn plan_copy_from_rows(
613    pcx: &PlanContext,
614    catalog: &dyn SessionCatalog,
615    target_id: CatalogItemId,
616    target_name: String,
617    columns: Vec<ColumnIndex>,
618    rows: Vec<mz_repr::Row>,
619) -> Result<HirRelationExpr, PlanError> {
620    let scx = StatementContext::new(Some(pcx), catalog);
621
622    // Always copy at the latest version of the table.
623    let table = catalog
624        .try_get_item(&target_id)
625        .ok_or_else(|| PlanError::CopyFromTargetTableDropped { target_name })?
626        .at_version(RelationVersionSelector::Latest);
627
628    let mut defaults = table
629        .writable_table_details()
630        .ok_or_else(|| sql_err!("cannot copy into non-writeable table"))?
631        .to_vec();
632
633    for default in &mut defaults {
634        transform_ast::transform(&scx, default)?;
635    }
636
637    let desc = table
638        .relation_desc()
639        .ok_or_else(|| sql_err!("item does not have a relation description"))?;
640    let column_types = columns
641        .iter()
642        .map(|x| desc.get_type(x).clone())
643        .map(|mut x| {
644            // Null constraint is enforced later, when inserting the row in the table.
645            // Without this, an assert is hit during lowering.
646            x.nullable = true;
647            x
648        })
649        .collect();
650    let typ = SqlRelationType::new(column_types);
651    let expr = HirRelationExpr::Constant {
652        rows,
653        typ: typ.clone(),
654    };
655
656    // Exit early with just the raw constant if we know that all columns are present
657    // and in the correct order. This lets us bypass expensive downstream optimizations
658    // more easily, as at every stage we know this expression is nothing more than
659    // a constant (as opposed to e.g. a constant with with an identity map and identity
660    // projection).
661    let default: Vec<_> = (0..desc.arity()).map(ColumnIndex::from_raw).collect();
662    if columns == default {
663        return Ok(expr);
664    }
665
666    // Fill in any omitted columns and rearrange into correct order
667    let mut map_exprs = vec![];
668    let mut project_key = Vec::with_capacity(desc.arity());
669
670    // Maps from table column index to position in the source query
671    let col_to_source: BTreeMap<_, _> = columns.iter().enumerate().map(|(a, b)| (b, a)).collect();
672
673    let column_details = desc.iter_all().zip_eq(defaults);
674    for ((col_idx, _col_name, col_typ), default) in column_details {
675        if let Some(src_idx) = col_to_source.get(&col_idx) {
676            project_key.push(*src_idx);
677        } else {
678            let hir = plan_default_expr(&scx, &default, &col_typ.scalar_type)?;
679            project_key.push(typ.arity() + map_exprs.len());
680            map_exprs.push(hir);
681        }
682    }
683
684    Ok(expr.map(map_exprs).project(project_key))
685}
686
687/// Common information used for DELETE, UPDATE, and INSERT INTO ... SELECT plans.
688pub struct ReadThenWritePlan {
689    pub id: CatalogItemId,
690    /// Read portion of query.
691    ///
692    /// NOTE: Even if the WHERE filter is left off, we still need to perform a read to generate
693    /// retractions.
694    pub selection: HirRelationExpr,
695    /// Map from column index to SET expression. Empty for DELETE statements.
696    pub assignments: BTreeMap<usize, HirScalarExpr>,
697    pub finishing: RowSetFinishing,
698}
699
700pub fn plan_delete_query(
701    scx: &StatementContext,
702    mut delete_stmt: DeleteStatement<Aug>,
703) -> Result<ReadThenWritePlan, PlanError> {
704    transform_ast::transform(scx, &mut delete_stmt)?;
705
706    let qcx = QueryContext::root(scx, QueryLifetime::OneShot);
707    plan_mutation_query_inner(
708        qcx,
709        delete_stmt.table_name,
710        delete_stmt.alias,
711        delete_stmt.using,
712        vec![],
713        delete_stmt.selection,
714    )
715}
716
717pub fn plan_update_query(
718    scx: &StatementContext,
719    mut update_stmt: UpdateStatement<Aug>,
720) -> Result<ReadThenWritePlan, PlanError> {
721    transform_ast::transform(scx, &mut update_stmt)?;
722
723    let qcx = QueryContext::root(scx, QueryLifetime::OneShot);
724
725    plan_mutation_query_inner(
726        qcx,
727        update_stmt.table_name,
728        update_stmt.alias,
729        vec![],
730        update_stmt.assignments,
731        update_stmt.selection,
732    )
733}
734
735pub fn plan_mutation_query_inner(
736    qcx: QueryContext,
737    table_name: ResolvedItemName,
738    alias: Option<TableAlias>,
739    using: Vec<TableWithJoins<Aug>>,
740    assignments: Vec<Assignment<Aug>>,
741    selection: Option<Expr<Aug>>,
742) -> Result<ReadThenWritePlan, PlanError> {
743    // Get ID and version of the relation desc.
744    let (id, version) = match table_name {
745        ResolvedItemName::Item { id, version, .. } => (id, version),
746        _ => sql_bail!("cannot mutate non-user table"),
747    };
748
749    // Perform checks on item with given ID.
750    let item = qcx.scx.get_item(&id).at_version(version);
751    if item.item_type() != CatalogItemType::Table {
752        sql_bail!(
753            "cannot mutate {} '{}'",
754            item.item_type(),
755            table_name.full_name_str()
756        );
757    }
758    let _ = item.writable_table_details().ok_or_else(|| {
759        sql_err!(
760            "cannot mutate non-writeable table '{}'",
761            table_name.full_name_str()
762        )
763    })?;
764    if id.is_system() {
765        sql_bail!(
766            "cannot mutate system table '{}'",
767            table_name.full_name_str()
768        );
769    }
770
771    // Derive structs for operation from validated table
772    let (mut get, scope) = qcx.resolve_table_name(table_name)?;
773    let scope = plan_table_alias(scope, alias.as_ref())?;
774    let desc = item.relation_desc().expect("table has desc");
775    let relation_type = qcx.relation_type(&get);
776
777    if using.is_empty() {
778        if let Some(expr) = selection {
779            let ecx = &ExprContext {
780                qcx: &qcx,
781                name: "WHERE clause",
782                scope: &scope,
783                relation_type: &relation_type,
784                allow_aggregates: false,
785                allow_subqueries: true,
786                allow_parameters: true,
787                allow_windows: false,
788            };
789            let expr = plan_expr(ecx, &expr)?.type_as(ecx, &SqlScalarType::Bool)?;
790            get = get.filter(vec![expr]);
791        }
792    } else {
793        get = handle_mutation_using_clause(&qcx, selection, using, get, scope.clone())?;
794    }
795
796    let mut sets = BTreeMap::new();
797    for Assignment { id, value } in assignments {
798        // Get the index and type of the column.
799        let name = normalize::column_name(id);
800        match desc.get_by_name(&name) {
801            Some((idx, typ)) => {
802                let ecx = &ExprContext {
803                    qcx: &qcx,
804                    name: "SET clause",
805                    scope: &scope,
806                    relation_type: &relation_type,
807                    allow_aggregates: false,
808                    allow_subqueries: false,
809                    allow_parameters: true,
810                    allow_windows: false,
811                };
812                let expr = plan_expr(ecx, &value)?.cast_to(
813                    ecx,
814                    CastContext::Assignment,
815                    &typ.scalar_type,
816                )?;
817
818                if sets.insert(idx, expr).is_some() {
819                    sql_bail!("column {} set twice", name)
820                }
821            }
822            None => sql_bail!("unknown column {}", name),
823        };
824    }
825
826    let finishing = RowSetFinishing {
827        order_by: vec![],
828        limit: None,
829        offset: 0,
830        project: (0..desc.arity()).collect(),
831    };
832
833    Ok(ReadThenWritePlan {
834        id,
835        selection: get,
836        finishing,
837        assignments: sets,
838    })
839}
840
841// Adjust `get` to perform an existential subquery on `using` accounting for
842// `selection`.
843//
844// If `USING`, we essentially want to rewrite the query as a correlated
845// existential subquery, i.e.
846// ```
847// ...WHERE EXISTS (SELECT 1 FROM <using> WHERE <selection>)
848// ```
849// However, we can't do that directly because of esoteric rules w/r/t `lateral`
850// subqueries.
851// https://github.com/postgres/postgres/commit/158b7fa6a34006bdc70b515e14e120d3e896589b
852fn handle_mutation_using_clause(
853    qcx: &QueryContext,
854    selection: Option<Expr<Aug>>,
855    using: Vec<TableWithJoins<Aug>>,
856    get: HirRelationExpr,
857    outer_scope: Scope,
858) -> Result<HirRelationExpr, PlanError> {
859    // Plan `USING` as a cross-joined `FROM` without knowledge of the
860    // statement's `FROM` target. This prevents `lateral` subqueries from
861    // "seeing" the `FROM` target.
862    let (mut using_rel_expr, using_scope) =
863        using.into_iter().try_fold(plan_join_identity(), |l, twj| {
864            let (left, left_scope) = l;
865            plan_join(
866                qcx,
867                left,
868                left_scope,
869                &Join {
870                    relation: TableFactor::NestedJoin {
871                        join: Box::new(twj),
872                        alias: None,
873                    },
874                    join_operator: JoinOperator::CrossJoin,
875                },
876            )
877        })?;
878
879    if let Some(expr) = selection {
880        // Join `FROM` with `USING` tables, like `USING..., FROM`. This gives us
881        // PG-like semantics e.g. expressing ambiguous column references. We put
882        // `USING...` first for no real reason, but making a different decision
883        // would require adjusting the column references on this relation
884        // differently.
885        let on = HirScalarExpr::literal_true();
886        let joined = using_rel_expr
887            .clone()
888            .join(get.clone(), on, JoinKind::Inner);
889        let joined_scope = using_scope.product(outer_scope)?;
890        let joined_relation_type = qcx.relation_type(&joined);
891
892        let ecx = &ExprContext {
893            qcx,
894            name: "WHERE clause",
895            scope: &joined_scope,
896            relation_type: &joined_relation_type,
897            allow_aggregates: false,
898            allow_subqueries: true,
899            allow_parameters: true,
900            allow_windows: false,
901        };
902
903        // Plan the filter expression on `FROM, USING...`.
904        let mut expr = plan_expr(ecx, &expr)?.type_as(ecx, &SqlScalarType::Bool)?;
905
906        // Rewrite all column referring to the `FROM` section of `joined` (i.e.
907        // those to the right of `using_rel_expr`) to instead be correlated to
908        // the outer relation, i.e. `get`.
909        let using_rel_arity = qcx.relation_type(&using_rel_expr).arity();
910        // local import to not get confused with `mz_sql_parser::ast::visit::Visit`
911        use mz_expr::visit::Visit;
912        expr.visit_mut_post(&mut |e| {
913            if let HirScalarExpr::Column(c, _name) = e {
914                if c.column >= using_rel_arity {
915                    c.level += 1;
916                    c.column -= using_rel_arity;
917                };
918            }
919        });
920
921        // Filter `USING` tables like `<using_rel_expr> WHERE <expr>`. Note that
922        // this filters the `USING` tables, _not_ the joined `USING..., FROM`
923        // relation.
924        using_rel_expr = using_rel_expr.filter(vec![expr]);
925    } else {
926        // Check that scopes are at compatible (i.e. do not double-reference
927        // same table), despite lack of selection
928        let _joined_scope = using_scope.product(outer_scope)?;
929    }
930    // From pg: Since the result [of EXISTS (<subquery>)] depends only on
931    // whether any rows are returned, and not on the contents of those rows,
932    // the output list of the subquery is normally unimportant.
933    //
934    // This means we don't need to worry about projecting/mapping any
935    // additional expressions here.
936    //
937    // https://www.postgresql.org/docs/14/functions-subquery.html
938
939    // Filter `get` like `...WHERE EXISTS (<using_rel_expr>)`.
940    Ok(get.filter(vec![using_rel_expr.exists()]))
941}
942
943#[derive(Debug)]
944pub(crate) struct CastRelationError {
945    pub(crate) column: usize,
946    pub(crate) source_type: SqlScalarType,
947    pub(crate) target_type: SqlScalarType,
948}
949
950/// Cast a relation from one type to another using the specified type of cast.
951///
952/// The length of `target_types` must match the arity of `expr`.
953pub(crate) fn cast_relation<'a, I>(
954    qcx: &QueryContext,
955    ccx: CastContext,
956    expr: HirRelationExpr,
957    target_types: I,
958) -> Result<HirRelationExpr, CastRelationError>
959where
960    I: IntoIterator<Item = &'a SqlScalarType>,
961{
962    let ecx = &ExprContext {
963        qcx,
964        name: "values",
965        scope: &Scope::empty(),
966        relation_type: &qcx.relation_type(&expr),
967        allow_aggregates: false,
968        allow_subqueries: true,
969        allow_parameters: true,
970        allow_windows: false,
971    };
972    let mut map_exprs = vec![];
973    let mut project_key = vec![];
974    for (i, target_typ) in target_types.into_iter().enumerate() {
975        let expr = HirScalarExpr::column(i);
976        // We plan every cast and check the evaluated expressions rather than
977        // checking the types directly because of some complex casting rules
978        // between types not expressed in `SqlScalarType` equality.
979        match typeconv::plan_cast(ecx, ccx, expr.clone(), target_typ) {
980            Ok(cast_expr) => {
981                if expr == cast_expr {
982                    // Cast between types was unnecessary
983                    project_key.push(i);
984                } else {
985                    // Cast between types required
986                    project_key.push(ecx.relation_type.arity() + map_exprs.len());
987                    map_exprs.push(cast_expr);
988                }
989            }
990            Err(_) => {
991                return Err(CastRelationError {
992                    column: i,
993                    source_type: ecx.scalar_type(&expr),
994                    target_type: target_typ.clone(),
995                });
996            }
997        }
998    }
999    Ok(expr.map(map_exprs).project(project_key))
1000}
1001
1002/// Plans an expression in the AS OF position of a `SELECT` or `SUBSCRIBE`, or `CREATE MATERIALIZED
1003/// VIEW` statement.
1004pub fn plan_as_of(
1005    scx: &StatementContext,
1006    as_of: Option<AsOf<Aug>>,
1007) -> Result<QueryWhen, PlanError> {
1008    match as_of {
1009        None => Ok(QueryWhen::Immediately),
1010        Some(as_of) => match as_of {
1011            AsOf::At(expr) => Ok(QueryWhen::AtTimestamp(plan_as_of_or_up_to(scx, expr)?)),
1012            AsOf::AtLeast(expr) => Ok(QueryWhen::AtLeastTimestamp(plan_as_of_or_up_to(scx, expr)?)),
1013        },
1014    }
1015}
1016
1017/// Plans and evaluates a scalar expression in a OneShot context to a non-null MzTimestamp.
1018///
1019/// Produces [`PlanError::InvalidAsOfUpTo`] if the expression is
1020/// - not a constant,
1021/// - not castable to MzTimestamp,
1022/// - is null,
1023/// - contains an unmaterializable function,
1024/// - some other evaluation error occurs, e.g., a division by 0,
1025/// - contains aggregates, subqueries, parameters, or window function calls.
1026pub fn plan_as_of_or_up_to(
1027    scx: &StatementContext,
1028    mut expr: Expr<Aug>,
1029) -> Result<mz_repr::Timestamp, PlanError> {
1030    let scope = Scope::empty();
1031    let desc = RelationDesc::empty();
1032    // (Even for a SUBSCRIBE, we need QueryLifetime::OneShot, because the AS OF or UP TO is
1033    // evaluated only once.)
1034    let qcx = QueryContext::root(scx, QueryLifetime::OneShot);
1035    transform_ast::transform(scx, &mut expr)?;
1036    let ecx = &ExprContext {
1037        qcx: &qcx,
1038        name: "AS OF or UP TO",
1039        scope: &scope,
1040        relation_type: desc.typ(),
1041        allow_aggregates: false,
1042        allow_subqueries: false,
1043        allow_parameters: false,
1044        allow_windows: false,
1045    };
1046    let hir = plan_expr(ecx, &expr)?.cast_to(
1047        ecx,
1048        CastContext::Assignment,
1049        &SqlScalarType::MzTimestamp,
1050    )?;
1051    if hir.contains_unmaterializable() {
1052        bail_unsupported!("calling an unmaterializable function in AS OF or UP TO");
1053    }
1054    // At this point, we definitely have a constant expression:
1055    // - it can't contain any unmaterializable functions;
1056    // - it can't refer to any columns.
1057    // But the following can still fail due to a variety of reasons: most commonly, the cast can
1058    // fail, but also a null might appear, or some other evaluation error can happen, e.g., a
1059    // division by 0.
1060    let timestamp = hir
1061        .into_literal_mz_timestamp()
1062        .ok_or_else(|| PlanError::InvalidAsOfUpTo)?;
1063    Ok(timestamp)
1064}
1065
1066/// Plans an expression in the AS position of a `CREATE SECRET`.
1067pub fn plan_secret_as(
1068    scx: &StatementContext,
1069    mut expr: Expr<Aug>,
1070) -> Result<MirScalarExpr, PlanError> {
1071    let scope = Scope::empty();
1072    let desc = RelationDesc::empty();
1073    let qcx = QueryContext::root(scx, QueryLifetime::OneShot);
1074
1075    transform_ast::transform(scx, &mut expr)?;
1076
1077    let ecx = &ExprContext {
1078        qcx: &qcx,
1079        name: "AS",
1080        scope: &scope,
1081        relation_type: desc.typ(),
1082        allow_aggregates: false,
1083        allow_subqueries: false,
1084        allow_parameters: false,
1085        allow_windows: false,
1086    };
1087    let expr = plan_expr(ecx, &expr)?
1088        .type_as(ecx, &SqlScalarType::Bytes)?
1089        .lower_uncorrelated(scx.catalog.system_vars())?;
1090    Ok(expr)
1091}
1092
1093/// Plans an expression in the CHECK position of a `CREATE SOURCE ... FROM WEBHOOK`.
1094pub fn plan_webhook_validate_using(
1095    scx: &StatementContext,
1096    validate_using: CreateWebhookSourceCheck<Aug>,
1097) -> Result<WebhookValidation, PlanError> {
1098    let qcx = QueryContext::root(scx, QueryLifetime::Source);
1099
1100    let CreateWebhookSourceCheck {
1101        options,
1102        using: mut expr,
1103    } = validate_using;
1104
1105    let mut column_typs = vec![];
1106    let mut column_names = vec![];
1107
1108    let (bodies, headers, secrets) = options
1109        .map(|o| (o.bodies, o.headers, o.secrets))
1110        .unwrap_or_default();
1111
1112    // Append all of the bodies so they can be used in the expression.
1113    let mut body_tuples = vec![];
1114    for CreateWebhookSourceBody { alias, use_bytes } in bodies {
1115        let scalar_type = use_bytes
1116            .then_some(SqlScalarType::Bytes)
1117            .unwrap_or(SqlScalarType::String);
1118        let name = alias
1119            .map(|a| a.into_string())
1120            .unwrap_or_else(|| "body".to_string());
1121
1122        column_typs.push(SqlColumnType {
1123            scalar_type,
1124            nullable: false,
1125        });
1126        column_names.push(name);
1127
1128        // Store the column index so we can be sure to provide this body correctly.
1129        let column_idx = column_typs.len() - 1;
1130        // Double check we're consistent with column names.
1131        assert_eq!(
1132            column_idx,
1133            column_names.len() - 1,
1134            "body column names and types don't match"
1135        );
1136        body_tuples.push((column_idx, use_bytes));
1137    }
1138
1139    // Append all of the headers so they can be used in the expression.
1140    let mut header_tuples = vec![];
1141
1142    for CreateWebhookSourceHeader { alias, use_bytes } in headers {
1143        let value_type = use_bytes
1144            .then_some(SqlScalarType::Bytes)
1145            .unwrap_or(SqlScalarType::String);
1146        let name = alias
1147            .map(|a| a.into_string())
1148            .unwrap_or_else(|| "headers".to_string());
1149
1150        column_typs.push(SqlColumnType {
1151            scalar_type: SqlScalarType::Map {
1152                value_type: Box::new(value_type),
1153                custom_id: None,
1154            },
1155            nullable: false,
1156        });
1157        column_names.push(name);
1158
1159        // Store the column index so we can be sure to provide this body correctly.
1160        let column_idx = column_typs.len() - 1;
1161        // Double check we're consistent with column names.
1162        assert_eq!(
1163            column_idx,
1164            column_names.len() - 1,
1165            "header column names and types don't match"
1166        );
1167        header_tuples.push((column_idx, use_bytes));
1168    }
1169
1170    // Append all secrets so they can be used in the expression.
1171    let mut validation_secrets = vec![];
1172
1173    for CreateWebhookSourceSecret {
1174        secret,
1175        alias,
1176        use_bytes,
1177    } in secrets
1178    {
1179        // Either provide the secret to the validation expression as Bytes or a String.
1180        let scalar_type = use_bytes
1181            .then_some(SqlScalarType::Bytes)
1182            .unwrap_or(SqlScalarType::String);
1183
1184        column_typs.push(SqlColumnType {
1185            scalar_type,
1186            nullable: false,
1187        });
1188        let ResolvedItemName::Item {
1189            id,
1190            full_name: FullItemName { item, .. },
1191            ..
1192        } = secret
1193        else {
1194            return Err(PlanError::InvalidSecret(Box::new(secret)));
1195        };
1196
1197        // Plan the expression using the secret's alias, if one is provided.
1198        let name = if let Some(alias) = alias {
1199            alias.into_string()
1200        } else {
1201            item
1202        };
1203        column_names.push(name);
1204
1205        // Get the column index that corresponds for this secret, so we can make sure to provide the
1206        // secrets in the correct order during evaluation.
1207        let column_idx = column_typs.len() - 1;
1208        // Double check that our column names and types match.
1209        assert_eq!(
1210            column_idx,
1211            column_names.len() - 1,
1212            "column names and types don't match"
1213        );
1214
1215        validation_secrets.push(WebhookValidationSecret {
1216            id,
1217            column_idx,
1218            use_bytes,
1219        });
1220    }
1221
1222    let relation_typ = SqlRelationType::new(column_typs);
1223    let desc = RelationDesc::new(relation_typ, column_names.clone());
1224    let scope = Scope::from_source(None, column_names);
1225
1226    transform_ast::transform(scx, &mut expr)?;
1227
1228    let ecx = &ExprContext {
1229        qcx: &qcx,
1230        name: "CHECK",
1231        scope: &scope,
1232        relation_type: desc.typ(),
1233        allow_aggregates: false,
1234        allow_subqueries: false,
1235        allow_parameters: false,
1236        allow_windows: false,
1237    };
1238    let expr = plan_expr(ecx, &expr)?
1239        .type_as(ecx, &SqlScalarType::Bool)?
1240        .lower_uncorrelated(scx.catalog.system_vars())?;
1241    let validation = WebhookValidation {
1242        expression: expr,
1243        relation_desc: desc,
1244        bodies: body_tuples,
1245        headers: header_tuples,
1246        secrets: validation_secrets,
1247    };
1248    Ok(validation)
1249}
1250
1251pub fn plan_default_expr(
1252    scx: &StatementContext,
1253    expr: &Expr<Aug>,
1254    target_ty: &SqlScalarType,
1255) -> Result<HirScalarExpr, PlanError> {
1256    let qcx = QueryContext::root(scx, QueryLifetime::OneShot);
1257    let ecx = &ExprContext {
1258        qcx: &qcx,
1259        name: "DEFAULT expression",
1260        scope: &Scope::empty(),
1261        relation_type: &SqlRelationType::empty(),
1262        allow_aggregates: false,
1263        allow_subqueries: false,
1264        allow_parameters: false,
1265        allow_windows: false,
1266    };
1267    let hir = plan_expr(ecx, expr)?.cast_to(ecx, CastContext::Assignment, target_ty)?;
1268    Ok(hir)
1269}
1270
1271pub fn plan_params<'a>(
1272    scx: &'a StatementContext,
1273    params: Vec<Expr<Aug>>,
1274    desc: &StatementDesc,
1275) -> Result<Params, PlanError> {
1276    if params.len() != desc.param_types.len() {
1277        sql_bail!(
1278            "expected {} params, got {}",
1279            desc.param_types.len(),
1280            params.len()
1281        );
1282    }
1283
1284    let qcx = QueryContext::root(scx, QueryLifetime::OneShot);
1285
1286    let mut datums = Row::default();
1287    let mut packer = datums.packer();
1288    let mut actual_types = Vec::new();
1289    let temp_storage = &RowArena::new();
1290    for (i, (mut expr, expected_ty)) in params.into_iter().zip_eq(&desc.param_types).enumerate() {
1291        transform_ast::transform(scx, &mut expr)?;
1292
1293        let ecx = execute_expr_context(&qcx);
1294        let ex = plan_expr(&ecx, &expr)?.type_as_any(&ecx)?;
1295        let actual_ty = ecx.scalar_type(&ex);
1296        if plan_hypothetical_cast(&ecx, *EXECUTE_CAST_CONTEXT, &actual_ty, expected_ty).is_none() {
1297            return Err(PlanError::WrongParameterType(
1298                i + 1,
1299                ecx.humanize_sql_scalar_type(expected_ty, false),
1300                ecx.humanize_sql_scalar_type(&actual_ty, false),
1301            ));
1302        }
1303        let ex = ex.lower_uncorrelated(scx.catalog.system_vars())?;
1304        let evaled = ex.eval(&[], temp_storage)?;
1305        packer.push(evaled);
1306        actual_types.push(actual_ty);
1307    }
1308    Ok(Params {
1309        datums,
1310        execute_types: actual_types,
1311        expected_types: desc.param_types.clone(),
1312    })
1313}
1314
1315static EXECUTE_CONTEXT_SCOPE: LazyLock<Scope> = LazyLock::new(Scope::empty);
1316static EXECUTE_CONTEXT_REL_TYPE: LazyLock<SqlRelationType> = LazyLock::new(SqlRelationType::empty);
1317
1318/// Returns an `ExprContext` for the expressions in the parameters of an EXECUTE statement.
1319pub(crate) fn execute_expr_context<'a>(qcx: &'a QueryContext<'a>) -> ExprContext<'a> {
1320    ExprContext {
1321        qcx,
1322        name: "EXECUTE",
1323        scope: &EXECUTE_CONTEXT_SCOPE,
1324        relation_type: &EXECUTE_CONTEXT_REL_TYPE,
1325        allow_aggregates: false,
1326        allow_subqueries: false,
1327        allow_parameters: false,
1328        allow_windows: false,
1329    }
1330}
1331
1332/// The CastContext used when matching up the types of parameters passed to EXECUTE.
1333///
1334/// This is an assignment cast also in Postgres, see
1335/// <https://github.com/MaterializeInc/database-issues/issues/9266>
1336pub(crate) static EXECUTE_CAST_CONTEXT: LazyLock<CastContext> =
1337    LazyLock::new(|| CastContext::Assignment);
1338
1339pub fn plan_index_exprs<'a>(
1340    scx: &'a StatementContext,
1341    on_desc: &RelationDesc,
1342    exprs: Vec<Expr<Aug>>,
1343) -> Result<Vec<mz_expr::MirScalarExpr>, PlanError> {
1344    let scope = Scope::from_source(None, on_desc.iter_names());
1345    let qcx = QueryContext::root(scx, QueryLifetime::Index);
1346
1347    let ecx = &ExprContext {
1348        qcx: &qcx,
1349        name: "CREATE INDEX",
1350        scope: &scope,
1351        relation_type: on_desc.typ(),
1352        allow_aggregates: false,
1353        allow_subqueries: false,
1354        allow_parameters: false,
1355        allow_windows: false,
1356    };
1357    let repr_col_types: Vec<ReprColumnType> = on_desc
1358        .typ()
1359        .column_types
1360        .iter()
1361        .map(ReprColumnType::from)
1362        .collect();
1363    let mut out = vec![];
1364    for mut expr in exprs {
1365        transform_ast::transform(scx, &mut expr)?;
1366        let expr = plan_expr_or_col_index(ecx, &expr)?;
1367        let mut expr = expr.lower_uncorrelated(scx.catalog.system_vars())?;
1368        expr.reduce(&repr_col_types);
1369        out.push(expr);
1370    }
1371    Ok(out)
1372}
1373
1374fn plan_expr_or_col_index(ecx: &ExprContext, e: &Expr<Aug>) -> Result<HirScalarExpr, PlanError> {
1375    match check_col_index(ecx.name, e, ecx.relation_type.column_types.len())? {
1376        Some(column) => Ok(HirScalarExpr::column(column)),
1377        _ => plan_expr(ecx, e)?.type_as_any(ecx),
1378    }
1379}
1380
1381fn check_col_index(name: &str, e: &Expr<Aug>, max: usize) -> Result<Option<usize>, PlanError> {
1382    match e {
1383        Expr::Value(Value::Number(n)) => {
1384            let n = n.parse::<usize>().map_err(|e| {
1385                sql_err!("unable to parse column reference in {}: {}: {}", name, n, e)
1386            })?;
1387            if n < 1 || n > max {
1388                sql_bail!(
1389                    "column reference {} in {} is out of range (1 - {})",
1390                    n,
1391                    name,
1392                    max
1393                );
1394            }
1395            Ok(Some(n - 1))
1396        }
1397        _ => Ok(None),
1398    }
1399}
1400
1401struct PlannedQuery {
1402    expr: HirRelationExpr,
1403    scope: Scope,
1404    order_by: Vec<ColumnOrder>,
1405    limit: Option<HirScalarExpr>,
1406    /// `offset` is either
1407    /// - an Int64 literal
1408    /// - or contains parameters. (If it contains parameters, then after parameter substitution it
1409    ///   should also be `is_constant` and reduce to an Int64 literal, but we check this only
1410    ///   later.)
1411    offset: HirScalarExpr,
1412    project: Vec<usize>,
1413    group_size_hints: GroupSizeHints,
1414}
1415
1416fn plan_query(qcx: &mut QueryContext, q: &Query<Aug>) -> Result<PlannedQuery, PlanError> {
1417    qcx.checked_recur_mut(|qcx| plan_query_inner(qcx, q))
1418}
1419
1420fn plan_query_inner(qcx: &mut QueryContext, q: &Query<Aug>) -> Result<PlannedQuery, PlanError> {
1421    // Plan CTEs and introduce bindings to `qcx.ctes`. Returns shadowed bindings
1422    // for the identifiers, so that they can be re-installed before returning.
1423    let cte_bindings = plan_ctes(qcx, q)?;
1424
1425    let limit = match &q.limit {
1426        None => None,
1427        Some(Limit {
1428            quantity,
1429            with_ties: false,
1430        }) => {
1431            let ecx = &ExprContext {
1432                qcx,
1433                name: "LIMIT",
1434                scope: &Scope::empty(),
1435                relation_type: &SqlRelationType::empty(),
1436                allow_aggregates: false,
1437                allow_subqueries: true,
1438                allow_parameters: true,
1439                allow_windows: false,
1440            };
1441            let limit = plan_expr(ecx, quantity)?;
1442            let limit = limit.cast_to(ecx, CastContext::Explicit, &SqlScalarType::Int64)?;
1443
1444            let limit = if limit.is_constant() {
1445                let arena = RowArena::new();
1446                let limit = limit.lower_uncorrelated(qcx.scx.catalog.system_vars())?;
1447
1448                // TODO: Don't use ? on eval, but instead wrap the error and add the information
1449                // that the error happened in a LIMIT clause, so that we have better error msg for
1450                // something like `SELECT 5 LIMIT 'aaa'`.
1451                match limit.eval(&[], &arena)? {
1452                    d @ Datum::Int64(v) if v >= 0 => {
1453                        HirScalarExpr::literal(d, SqlScalarType::Int64)
1454                    }
1455                    d @ Datum::Null => HirScalarExpr::literal(d, SqlScalarType::Int64),
1456                    Datum::Int64(_) => sql_bail!("LIMIT must not be negative"),
1457                    _ => sql_bail!("constant LIMIT expression must reduce to an INT or NULL value"),
1458                }
1459            } else {
1460                // Gate non-constant LIMIT expressions behind a feature flag
1461                qcx.scx
1462                    .require_feature_flag(&vars::ENABLE_EXPRESSIONS_IN_LIMIT_SYNTAX)?;
1463                limit
1464            };
1465
1466            Some(limit)
1467        }
1468        Some(Limit {
1469            quantity: _,
1470            with_ties: true,
1471        }) => bail_unsupported!("FETCH ... WITH TIES"),
1472    };
1473
1474    let offset = match &q.offset {
1475        None => HirScalarExpr::literal(Datum::Int64(0), SqlScalarType::Int64),
1476        Some(offset) => {
1477            let ecx = &ExprContext {
1478                qcx,
1479                name: "OFFSET",
1480                scope: &Scope::empty(),
1481                relation_type: &SqlRelationType::empty(),
1482                allow_aggregates: false,
1483                allow_subqueries: false,
1484                allow_parameters: true,
1485                allow_windows: false,
1486            };
1487            let offset = plan_expr(ecx, offset)?;
1488            let offset = offset.cast_to(ecx, CastContext::Explicit, &SqlScalarType::Int64)?;
1489
1490            let offset = if offset.is_constant() {
1491                // Simplify it to a literal or error out. (E.g., the cast inserted above may fail.)
1492                let offset_value = offset_into_value(offset)?;
1493                HirScalarExpr::literal(Datum::Int64(offset_value), SqlScalarType::Int64)
1494            } else {
1495                // The only case when this is allowed to not be a constant is if it contains
1496                // parameters. (In which case, we'll later check that it's a constant after
1497                // parameter binding.)
1498                if !offset.contains_parameters() {
1499                    return Err(PlanError::InvalidOffset(format!(
1500                        "must be simplifiable to a constant, possibly after parameter binding, got {}",
1501                        offset
1502                    )));
1503                }
1504                offset
1505            };
1506            offset
1507        }
1508    };
1509
1510    let mut planned_query = match &q.body {
1511        SetExpr::Select(s) => {
1512            // Extract query options.
1513            let select_option_extracted = SelectOptionExtracted::try_from(s.options.clone())?;
1514            let group_size_hints = GroupSizeHints::try_from(select_option_extracted)?;
1515
1516            let plan = plan_select_from_where(qcx, *s.clone(), q.order_by.clone())?;
1517            PlannedQuery {
1518                expr: plan.expr,
1519                scope: plan.scope,
1520                order_by: plan.order_by,
1521                project: plan.project,
1522                limit,
1523                offset,
1524                group_size_hints,
1525            }
1526        }
1527        _ => {
1528            let (expr, scope) = plan_set_expr(qcx, &q.body)?;
1529            let ecx = &ExprContext {
1530                qcx,
1531                name: "ORDER BY clause of a set expression",
1532                scope: &scope,
1533                relation_type: &qcx.relation_type(&expr),
1534                allow_aggregates: false,
1535                allow_subqueries: true,
1536                allow_parameters: true,
1537                allow_windows: false,
1538            };
1539            let output_columns: Vec<_> = scope.column_names().enumerate().collect();
1540            let (order_by, map_exprs) = plan_order_by_exprs(ecx, &q.order_by, &output_columns)?;
1541            let project = (0..ecx.relation_type.arity()).collect();
1542            PlannedQuery {
1543                expr: expr.map(map_exprs),
1544                scope,
1545                order_by,
1546                limit,
1547                project,
1548                offset,
1549                group_size_hints: GroupSizeHints::default(),
1550            }
1551        }
1552    };
1553
1554    // Both introduce `Let` bindings atop `result` and re-install shadowed bindings.
1555    match &q.ctes {
1556        CteBlock::Simple(_) => {
1557            for (id, value, shadowed_val) in cte_bindings.into_iter().rev() {
1558                if let Some(cte) = qcx.ctes.remove(&id) {
1559                    planned_query.expr = HirRelationExpr::Let {
1560                        name: cte.name,
1561                        id: id.clone(),
1562                        value: Box::new(value),
1563                        body: Box::new(planned_query.expr),
1564                    };
1565                }
1566                if let Some(shadowed_val) = shadowed_val {
1567                    qcx.ctes.insert(id, shadowed_val);
1568                }
1569            }
1570        }
1571        CteBlock::MutuallyRecursive(MutRecBlock { options, ctes: _ }) => {
1572            let MutRecBlockOptionExtracted {
1573                recursion_limit,
1574                return_at_recursion_limit,
1575                error_at_recursion_limit,
1576                seen: _,
1577            } = MutRecBlockOptionExtracted::try_from(options.clone())?;
1578            let limit = match (
1579                recursion_limit,
1580                return_at_recursion_limit,
1581                error_at_recursion_limit,
1582            ) {
1583                (None, None, None) => None,
1584                (Some(max_iters), None, None) => {
1585                    Some((max_iters, LetRecLimit::RETURN_AT_LIMIT_DEFAULT))
1586                }
1587                (None, Some(max_iters), None) => Some((max_iters, true)),
1588                (None, None, Some(max_iters)) => Some((max_iters, false)),
1589                _ => {
1590                    return Err(InvalidWmrRecursionLimit(
1591                        "More than one recursion limit given. \
1592                         Please give at most one of RECURSION LIMIT, \
1593                         ERROR AT RECURSION LIMIT, \
1594                         RETURN AT RECURSION LIMIT."
1595                            .to_owned(),
1596                    ));
1597                }
1598            }
1599            .try_map(|(max_iters, return_at_limit)| {
1600                Ok::<LetRecLimit, PlanError>(LetRecLimit {
1601                    max_iters: NonZeroU64::new(*max_iters).ok_or(InvalidWmrRecursionLimit(
1602                        "Recursion limit has to be greater than 0.".to_owned(),
1603                    ))?,
1604                    return_at_limit: *return_at_limit,
1605                })
1606            })?;
1607
1608            let mut bindings = Vec::new();
1609            for (id, value, shadowed_val) in cte_bindings.into_iter() {
1610                if let Some(cte) = qcx.ctes.remove(&id) {
1611                    bindings.push((cte.name, id, value, cte.desc.into_typ()));
1612                }
1613                if let Some(shadowed_val) = shadowed_val {
1614                    qcx.ctes.insert(id, shadowed_val);
1615                }
1616            }
1617            if !bindings.is_empty() {
1618                planned_query.expr = HirRelationExpr::LetRec {
1619                    limit,
1620                    bindings,
1621                    body: Box::new(planned_query.expr),
1622                }
1623            }
1624        }
1625    }
1626
1627    Ok(planned_query)
1628}
1629
1630/// Converts an OFFSET expression into a value.
1631pub(crate) fn offset_into_value(offset: HirScalarExpr) -> Result<i64, PlanError> {
1632    let offset = offset
1633        .try_into_literal_int64()
1634        .map_err(|err| PlanError::InvalidOffset(err.to_string_with_causes()))?;
1635    if offset < 0 {
1636        return Err(negative_offset_error(offset));
1637    }
1638    Ok(offset)
1639}
1640
1641pub(crate) fn negative_offset_error(offset: i64) -> PlanError {
1642    PlanError::InvalidOffset(format!("must not be negative, got {}", offset))
1643}
1644
1645generate_extracted_config!(
1646    MutRecBlockOption,
1647    (RecursionLimit, u64),
1648    (ReturnAtRecursionLimit, u64),
1649    (ErrorAtRecursionLimit, u64)
1650);
1651
1652/// Creates plans for CTEs and introduces them to `qcx.ctes`.
1653///
1654/// Returns for each identifier a planned `HirRelationExpr` value, and an optional
1655/// shadowed value that can be reinstalled once the planning has completed.
1656pub fn plan_ctes(
1657    qcx: &mut QueryContext,
1658    q: &Query<Aug>,
1659) -> Result<Vec<(LocalId, HirRelationExpr, Option<CteDesc>)>, PlanError> {
1660    // Accumulate planned expressions and shadowed descriptions.
1661    let mut result = Vec::new();
1662    // Retain the old descriptions of CTE bindings so that we can restore them
1663    // after we're done planning this SELECT.
1664    let mut shadowed_descs = BTreeMap::new();
1665
1666    // A reused identifier indicates a reused name.
1667    if let Some(ident) = q.ctes.bound_identifiers().duplicates().next() {
1668        sql_bail!(
1669            "WITH query name {} specified more than once",
1670            normalize::ident_ref(ident).quoted()
1671        )
1672    }
1673
1674    match &q.ctes {
1675        CteBlock::Simple(ctes) => {
1676            // Plan all CTEs, introducing the types for non-recursive CTEs as we go.
1677            for cte in ctes.iter() {
1678                let cte_name = normalize::ident(cte.alias.name.clone());
1679                let (val, scope) = plan_nested_query(qcx, &cte.query)?;
1680                let typ = qcx.relation_type(&val);
1681                let mut desc = RelationDesc::new(typ, scope.column_names());
1682                plan_utils::maybe_rename_columns(
1683                    format!("CTE {}", cte.alias.name),
1684                    &mut desc,
1685                    &cte.alias.columns,
1686                )?;
1687                // Capture the prior value if it exists, so that it can be re-installed.
1688                let shadowed = qcx.ctes.insert(
1689                    cte.id,
1690                    CteDesc {
1691                        name: cte_name,
1692                        desc,
1693                    },
1694                );
1695
1696                result.push((cte.id, val, shadowed));
1697            }
1698        }
1699        CteBlock::MutuallyRecursive(MutRecBlock { options: _, ctes }) => {
1700            // Insert column types into `qcx.ctes` first for recursive bindings.
1701            for cte in ctes.iter() {
1702                let cte_name = normalize::ident(cte.name.clone());
1703                let mut desc_columns = Vec::with_capacity(cte.columns.capacity());
1704                for column in cte.columns.iter() {
1705                    desc_columns.push((
1706                        normalize::column_name(column.name.clone()),
1707                        SqlColumnType {
1708                            scalar_type: scalar_type_from_sql(qcx.scx, &column.data_type)?,
1709                            nullable: true,
1710                        },
1711                    ));
1712                }
1713                let desc = RelationDesc::from_names_and_types(desc_columns);
1714                let shadowed = qcx.ctes.insert(
1715                    cte.id,
1716                    CteDesc {
1717                        name: cte_name,
1718                        desc,
1719                    },
1720                );
1721                // Capture the prior value if it exists, so that it can be re-installed.
1722                if let Some(shadowed) = shadowed {
1723                    shadowed_descs.insert(cte.id, shadowed);
1724                }
1725            }
1726
1727            // Plan all CTEs and validate the proposed types.
1728            for cte in ctes.iter() {
1729                let (val, _scope) = plan_nested_query(qcx, &cte.query)?;
1730
1731                let proposed_typ = qcx.ctes[&cte.id].desc.typ();
1732
1733                if proposed_typ.column_types.iter().any(|c| !c.nullable) {
1734                    // Once WMR CTEs support NOT NULL constraints, check that
1735                    // nullability of derived column types are compatible.
1736                    sql_bail!(
1737                        "[internal error]: WMR CTEs do not support NOT NULL constraints on proposed column types"
1738                    );
1739                }
1740
1741                if !proposed_typ.keys.is_empty() {
1742                    // Once WMR CTEs support keys, check that keys exactly
1743                    // overlap.
1744                    sql_bail!("[internal error]: WMR CTEs do not support keys");
1745                }
1746
1747                // Validate that the derived and proposed types are the same.
1748                let derived_typ = qcx.relation_type(&val);
1749
1750                let type_err = |proposed_typ: &SqlRelationType, derived_typ: SqlRelationType| {
1751                    let cte_name = normalize::ident(cte.name.clone());
1752                    let proposed_typ = proposed_typ
1753                        .column_types
1754                        .iter()
1755                        .map(|ty| qcx.humanize_sql_scalar_type(&ty.scalar_type, false))
1756                        .collect::<Vec<_>>();
1757                    let inferred_typ = derived_typ
1758                        .column_types
1759                        .iter()
1760                        .map(|ty| qcx.humanize_sql_scalar_type(&ty.scalar_type, false))
1761                        .collect::<Vec<_>>();
1762                    Err(PlanError::RecursiveTypeMismatch(
1763                        cte_name,
1764                        proposed_typ,
1765                        inferred_typ,
1766                    ))
1767                };
1768
1769                if derived_typ.column_types.len() != proposed_typ.column_types.len() {
1770                    return type_err(proposed_typ, derived_typ);
1771                }
1772
1773                // Cast derived types to proposed types or error.
1774                let val = match cast_relation(
1775                    qcx,
1776                    // Choose `CastContext::Assignment`` because the user has
1777                    // been explicit about the types they expect. Choosing
1778                    // `CastContext::Implicit` is not "strong" enough to impose
1779                    // typmods from proposed types onto values.
1780                    CastContext::Assignment,
1781                    val,
1782                    proposed_typ.column_types.iter().map(|c| &c.scalar_type),
1783                ) {
1784                    Ok(val) => val,
1785                    Err(_) => return type_err(proposed_typ, derived_typ),
1786                };
1787
1788                result.push((cte.id, val, shadowed_descs.remove(&cte.id)));
1789            }
1790        }
1791    }
1792
1793    Ok(result)
1794}
1795
1796pub fn plan_nested_query(
1797    qcx: &mut QueryContext,
1798    q: &Query<Aug>,
1799) -> Result<(HirRelationExpr, Scope), PlanError> {
1800    let PlannedQuery {
1801        mut expr,
1802        scope,
1803        order_by,
1804        limit,
1805        offset,
1806        project,
1807        group_size_hints,
1808    } = qcx.checked_recur_mut(|qcx| plan_query(qcx, q))?;
1809    // A nested query is an unordered relation. Its `ORDER BY` is only observable
1810    // in combination with a row-limiting clause (`LIMIT`/`OFFSET`), which selects
1811    // which rows survive. Without such a clause the ordering has no defined
1812    // meaning, so it is dropped rather than materialized into a `TopK`.
1813    //
1814    // NOTE: This diverges from PostgreSQL, where an order-sensitive aggregate
1815    // (`array_agg`, `string_agg`, ...) in the outer query observes a sorted
1816    // subquery's output as an executor artifact. That behavior is not guaranteed
1817    // by the SQL standard and PostgreSQL itself documents it as fragile. Callers
1818    // that need a specific aggregation order must use the in-aggregate
1819    // `agg(value ORDER BY ...)` form instead.
1820    if limit.is_some()
1821        || !offset
1822            .clone()
1823            .try_into_literal_int64()
1824            .is_ok_and(|offset| offset == 0)
1825    {
1826        expr = HirRelationExpr::top_k(
1827            expr,
1828            vec![],
1829            order_by,
1830            limit,
1831            offset,
1832            group_size_hints.limit_input_group_size,
1833        );
1834    }
1835    Ok((expr.project(project), scope))
1836}
1837
1838fn plan_set_expr(
1839    qcx: &mut QueryContext,
1840    q: &SetExpr<Aug>,
1841) -> Result<(HirRelationExpr, Scope), PlanError> {
1842    match q {
1843        SetExpr::Select(select) => {
1844            let order_by_exprs = Vec::new();
1845            let plan = plan_select_from_where(qcx, *select.clone(), order_by_exprs)?;
1846            // We didn't provide any `order_by_exprs`, so `plan_select_from_where`
1847            // should not have planned any ordering.
1848            assert!(plan.order_by.is_empty());
1849            Ok((plan.expr.project(plan.project), plan.scope))
1850        }
1851        SetExpr::SetOperation {
1852            op,
1853            all,
1854            left,
1855            right,
1856        } => {
1857            // Plan the LHS and RHS.
1858            let (left_expr, left_scope) = qcx.checked_recur_mut(|qcx| plan_set_expr(qcx, left))?;
1859            let (right_expr, right_scope) =
1860                qcx.checked_recur_mut(|qcx| plan_set_expr(qcx, right))?;
1861
1862            // Validate that the LHS and RHS are the same width.
1863            let left_type = qcx.relation_type(&left_expr);
1864            let right_type = qcx.relation_type(&right_expr);
1865            if left_type.arity() != right_type.arity() {
1866                sql_bail!(
1867                    "each {} query must have the same number of columns: {} vs {}",
1868                    op,
1869                    left_type.arity(),
1870                    right_type.arity(),
1871                );
1872            }
1873
1874            // Match the types of the corresponding columns on the LHS and RHS
1875            // using the normal type coercion rules. This is equivalent to
1876            // `coerce_homogeneous_exprs`, but implemented in terms of
1877            // `HirRelationExpr` rather than `HirScalarExpr`.
1878            let left_ecx = &ExprContext {
1879                qcx,
1880                name: &op.to_string(),
1881                scope: &left_scope,
1882                relation_type: &left_type,
1883                allow_aggregates: false,
1884                allow_subqueries: false,
1885                allow_parameters: false,
1886                allow_windows: false,
1887            };
1888            let right_ecx = &ExprContext {
1889                qcx,
1890                name: &op.to_string(),
1891                scope: &right_scope,
1892                relation_type: &right_type,
1893                allow_aggregates: false,
1894                allow_subqueries: false,
1895                allow_parameters: false,
1896                allow_windows: false,
1897            };
1898            let mut left_casts = vec![];
1899            let mut right_casts = vec![];
1900            for (i, (left_type, right_type)) in left_type
1901                .column_types
1902                .iter()
1903                .zip_eq(right_type.column_types.iter())
1904                .enumerate()
1905            {
1906                let types = &[
1907                    CoercibleScalarType::Coerced(left_type.scalar_type.clone()),
1908                    CoercibleScalarType::Coerced(right_type.scalar_type.clone()),
1909                ];
1910                let target =
1911                    typeconv::guess_best_common_type(&left_ecx.with_name(&op.to_string()), types)?;
1912                match typeconv::plan_cast(
1913                    left_ecx,
1914                    CastContext::Implicit,
1915                    HirScalarExpr::column(i),
1916                    &target,
1917                ) {
1918                    Ok(expr) => left_casts.push(expr),
1919                    Err(_) => sql_bail!(
1920                        "{} types {} and {} cannot be matched",
1921                        op,
1922                        qcx.humanize_sql_scalar_type(&left_type.scalar_type, false),
1923                        qcx.humanize_sql_scalar_type(&target, false),
1924                    ),
1925                }
1926                match typeconv::plan_cast(
1927                    right_ecx,
1928                    CastContext::Implicit,
1929                    HirScalarExpr::column(i),
1930                    &target,
1931                ) {
1932                    Ok(expr) => right_casts.push(expr),
1933                    Err(_) => sql_bail!(
1934                        "{} types {} and {} cannot be matched",
1935                        op,
1936                        qcx.humanize_sql_scalar_type(&target, false),
1937                        qcx.humanize_sql_scalar_type(&right_type.scalar_type, false),
1938                    ),
1939                }
1940            }
1941            let lhs = if left_casts
1942                .iter()
1943                .enumerate()
1944                .any(|(i, e)| e != &HirScalarExpr::column(i))
1945            {
1946                let project_key: Vec<_> = (left_type.arity()..left_type.arity() * 2).collect();
1947                left_expr.map(left_casts).project(project_key)
1948            } else {
1949                left_expr
1950            };
1951            let rhs = if right_casts
1952                .iter()
1953                .enumerate()
1954                .any(|(i, e)| e != &HirScalarExpr::column(i))
1955            {
1956                let project_key: Vec<_> = (right_type.arity()..right_type.arity() * 2).collect();
1957                right_expr.map(right_casts).project(project_key)
1958            } else {
1959                right_expr
1960            };
1961
1962            let relation_expr = match op {
1963                SetOperator::Union => {
1964                    if *all {
1965                        lhs.union(rhs)
1966                    } else {
1967                        lhs.union(rhs).distinct()
1968                    }
1969                }
1970                SetOperator::Except => Hir::except(all, lhs, rhs),
1971                SetOperator::Intersect => {
1972                    // Planning below duplicates whichever input ends up on the left, so a
1973                    // left-deep chain of INTERSECTs doubles the plan at every level, i.e. would be
1974                    // exponential. INTERSECT is commutative, so put the cheaper input on the left.
1975                    // A subtree is then only duplicated when it is the smaller of the two, so plan
1976                    // size obeys T(a + b) <= 2*T(a) + T(b) for input sizes a <= b. The worst case
1977                    // is balanced trees, where this solves to O(n^log2(3)) instead of O(2^n).
1978                    let (lhs, rhs) = if lhs.relation_node_count() > rhs.relation_node_count() {
1979                        (rhs, lhs)
1980                    } else {
1981                        (lhs, rhs)
1982                    };
1983                    // TODO: Let's not duplicate the left-hand expression into TWO dataflows!
1984                    // The optimizer de-duplicates at some point, but it would be good to already
1985                    // not duplicate here.
1986                    //
1987                    // Also note that we do *not* need another threshold() at the end of the method chain
1988                    // because the right-hand side of the outer union only produces existing records,
1989                    // i.e., the record counts for differential data flow definitely remain non-negative.
1990                    let left_clone = lhs.clone();
1991                    if *all {
1992                        lhs.union(left_clone.union(rhs.negate()).threshold().negate())
1993                    } else {
1994                        lhs.union(left_clone.union(rhs.negate()).threshold().negate())
1995                            .distinct()
1996                    }
1997                }
1998            };
1999            let scope = Scope::from_source(
2000                None,
2001                // Column names are taken from the left, as in Postgres.
2002                left_scope.column_names(),
2003            );
2004
2005            Ok((relation_expr, scope))
2006        }
2007        SetExpr::Values(Values(values)) => plan_values(qcx, values),
2008        SetExpr::Table(name) => {
2009            let (expr, scope) = qcx.resolve_table_name(name.clone())?;
2010            Ok((expr, scope))
2011        }
2012        SetExpr::Query(query) => {
2013            let (expr, scope) = plan_nested_query(qcx, query)?;
2014            Ok((expr, scope))
2015        }
2016        SetExpr::Show(stmt) => {
2017            // The create SQL definition of involving this query, will have the explicit `SHOW`
2018            // command in it. Many `SHOW` commands will expand into a sub-query that involves the
2019            // current schema of the executing user. When Materialize restarts and tries to re-plan
2020            // these queries, it will only have access to the raw `SHOW` command and have no idea
2021            // what schema to use. As a result Materialize will fail to boot.
2022            //
2023            // Some `SHOW` commands are ok, like `SHOW CLUSTERS`, and there are probably other ways
2024            // around this issue. Such as expanding the `SHOW` command in the SQL definition.
2025            // However, banning show commands in views gives us more flexibility to change their
2026            // output.
2027            //
2028            // TODO(jkosh44) Add message to error that prints out an equivalent view definition
2029            // with all show commands expanded into their equivalent SELECT statements.
2030            if !qcx.lifetime.allow_show() {
2031                return Err(PlanError::ShowCommandInView);
2032            }
2033
2034            // Some SHOW statements are a SELECT query. Others produces Rows
2035            // directly. Convert both of these to the needed Hir and Scope.
2036            fn to_hirscope(
2037                plan: ShowCreatePlan,
2038                desc: StatementDesc,
2039            ) -> Result<(HirRelationExpr, Scope), PlanError> {
2040                let rows = vec![plan.row.iter().collect::<Vec<_>>()];
2041                let desc = desc.relation_desc.ok_or_else(|| {
2042                    internal_err!("statement description missing relation descriptor")
2043                })?;
2044                let scope = Scope::from_source(None, desc.iter_names());
2045                let expr = HirRelationExpr::constant(rows, desc.into_typ());
2046                Ok((expr, scope))
2047            }
2048
2049            match stmt.clone() {
2050                ShowStatement::ShowColumns(stmt) => {
2051                    show::show_columns(qcx.scx, stmt)?.plan_hir(qcx)
2052                }
2053                ShowStatement::ShowCreateConnection(stmt) => to_hirscope(
2054                    show::plan_show_create_connection(qcx.scx, stmt.clone())?,
2055                    show::describe_show_create_connection(qcx.scx, stmt)?,
2056                ),
2057                ShowStatement::ShowCreateCluster(stmt) => to_hirscope(
2058                    show::plan_show_create_cluster(qcx.scx, stmt.clone())?,
2059                    show::describe_show_create_cluster(qcx.scx, stmt)?,
2060                ),
2061                ShowStatement::ShowCreateIndex(stmt) => to_hirscope(
2062                    show::plan_show_create_index(qcx.scx, stmt.clone())?,
2063                    show::describe_show_create_index(qcx.scx, stmt)?,
2064                ),
2065                ShowStatement::ShowCreateSink(stmt) => to_hirscope(
2066                    show::plan_show_create_sink(qcx.scx, stmt.clone())?,
2067                    show::describe_show_create_sink(qcx.scx, stmt)?,
2068                ),
2069                ShowStatement::ShowCreateMetricSink(stmt) => to_hirscope(
2070                    show::plan_show_create_metric_sink(qcx.scx, stmt.clone())?,
2071                    show::describe_show_create_metric_sink(qcx.scx, stmt)?,
2072                ),
2073                ShowStatement::ShowCreateSource(stmt) => to_hirscope(
2074                    show::plan_show_create_source(qcx.scx, stmt.clone())?,
2075                    show::describe_show_create_source(qcx.scx, stmt)?,
2076                ),
2077                ShowStatement::ShowCreateTable(stmt) => to_hirscope(
2078                    show::plan_show_create_table(qcx.scx, stmt.clone())?,
2079                    show::describe_show_create_table(qcx.scx, stmt)?,
2080                ),
2081                ShowStatement::ShowCreateView(stmt) => to_hirscope(
2082                    show::plan_show_create_view(qcx.scx, stmt.clone())?,
2083                    show::describe_show_create_view(qcx.scx, stmt)?,
2084                ),
2085                ShowStatement::ShowCreateMaterializedView(stmt) => to_hirscope(
2086                    show::plan_show_create_materialized_view(qcx.scx, stmt.clone())?,
2087                    show::describe_show_create_materialized_view(qcx.scx, stmt)?,
2088                ),
2089                ShowStatement::ShowCreateType(stmt) => to_hirscope(
2090                    show::plan_show_create_type(qcx.scx, stmt.clone())?,
2091                    show::describe_show_create_type(qcx.scx, stmt)?,
2092                ),
2093                ShowStatement::ShowObjects(stmt) => {
2094                    show::show_objects(qcx.scx, stmt)?.plan_hir(qcx)
2095                }
2096                ShowStatement::ShowVariable(_) => bail_unsupported!("SHOW variable in subqueries"),
2097                ShowStatement::InspectShard(_) => sql_bail!("unsupported INSPECT statement"),
2098            }
2099        }
2100    }
2101}
2102
2103/// Plans a `VALUES` clause that appears in a `SELECT` statement.
2104fn plan_values(
2105    qcx: &QueryContext,
2106    values: &[Vec<Expr<Aug>>],
2107) -> Result<(HirRelationExpr, Scope), PlanError> {
2108    assert!(!values.is_empty());
2109
2110    let ecx = &ExprContext {
2111        qcx,
2112        name: "VALUES",
2113        scope: &Scope::empty(),
2114        relation_type: &SqlRelationType::empty(),
2115        allow_aggregates: false,
2116        allow_subqueries: true,
2117        allow_parameters: true,
2118        allow_windows: false,
2119    };
2120
2121    let ncols = values[0].len();
2122    let nrows = values.len();
2123
2124    // Arrange input expressions by columns, not rows, so that we can
2125    // call `coerce_homogeneous_exprs` on each column.
2126    let mut cols = vec![vec![]; ncols];
2127    for row in values {
2128        if row.len() != ncols {
2129            sql_bail!(
2130                "VALUES expression has varying number of columns: {} vs {}",
2131                row.len(),
2132                ncols
2133            );
2134        }
2135        for (i, v) in row.iter().enumerate() {
2136            cols[i].push(v);
2137        }
2138    }
2139
2140    // Plan each column.
2141    let mut col_iters = Vec::with_capacity(ncols);
2142    let mut col_types = Vec::with_capacity(ncols);
2143    for col in &cols {
2144        let col = coerce_homogeneous_exprs(ecx, plan_exprs(ecx, col)?, None)?;
2145        let mut col_type = ecx.column_type(&col[0]);
2146        for val in &col[1..] {
2147            col_type = col_type.sql_union(&ecx.column_type(val))?; // HIR deliberately not using `union`
2148        }
2149        col_types.push(col_type);
2150        col_iters.push(col.into_iter());
2151    }
2152
2153    // Build constant relation.
2154    let mut exprs = vec![];
2155    for _ in 0..nrows {
2156        for i in 0..ncols {
2157            exprs.push(col_iters[i].next().unwrap());
2158        }
2159    }
2160    let out = HirRelationExpr::CallTable {
2161        func: TableFunc::Wrap {
2162            width: ncols,
2163            types: col_types,
2164        },
2165        exprs,
2166    };
2167
2168    // Build column names.
2169    let mut scope = Scope::empty();
2170    for i in 0..ncols {
2171        let name = format!("column{}", i + 1);
2172        scope.items.push(ScopeItem::from_column_name(name));
2173    }
2174
2175    Ok((out, scope))
2176}
2177
2178/// Plans a `VALUES` clause that appears at the top level of an `INSERT`
2179/// statement.
2180///
2181/// This is special-cased in PostgreSQL and different enough from `plan_values`
2182/// that it is easier to use a separate function entirely. Unlike a normal
2183/// `VALUES` clause, each value is coerced to the type of the target table
2184/// via an assignment cast.
2185///
2186/// See: <https://github.com/postgres/postgres/blob/ad77039fa/src/backend/parser/analyze.c#L504-L518>
2187fn plan_values_insert(
2188    qcx: &QueryContext,
2189    target_names: &[&ColumnName],
2190    target_types: &[&SqlScalarType],
2191    values: &[Vec<Expr<Aug>>],
2192) -> Result<HirRelationExpr, PlanError> {
2193    assert!(!values.is_empty());
2194
2195    if !values.iter().map(|row| row.len()).all_equal() {
2196        sql_bail!("VALUES lists must all be the same length");
2197    }
2198
2199    let ecx = &ExprContext {
2200        qcx,
2201        name: "VALUES",
2202        scope: &Scope::empty(),
2203        relation_type: &SqlRelationType::empty(),
2204        allow_aggregates: false,
2205        allow_subqueries: true,
2206        allow_parameters: true,
2207        allow_windows: false,
2208    };
2209
2210    let mut exprs = vec![];
2211    let mut types = vec![];
2212    for row in values {
2213        if row.len() > target_names.len() {
2214            sql_bail!("INSERT has more expressions than target columns");
2215        }
2216        for (column, val) in row.into_iter().enumerate() {
2217            let target_type = &target_types[column];
2218            let val = plan_expr(ecx, val)?;
2219            let val = typeconv::plan_coerce(ecx, val, target_type)?;
2220            let source_type = &ecx.scalar_type(&val);
2221            let val = match typeconv::plan_cast(ecx, CastContext::Assignment, val, target_type) {
2222                Ok(val) => val,
2223                Err(_) => sql_bail!(
2224                    "column {} is of type {} but expression is of type {}",
2225                    target_names[column].quoted(),
2226                    qcx.humanize_sql_scalar_type(target_type, false),
2227                    qcx.humanize_sql_scalar_type(source_type, false),
2228                ),
2229            };
2230            if column >= types.len() {
2231                types.push(ecx.column_type(&val));
2232            } else {
2233                types[column] = types[column].sql_union(&ecx.column_type(&val))?; // HIR deliberately not using `union`
2234            }
2235            exprs.push(val);
2236        }
2237    }
2238
2239    Ok(HirRelationExpr::CallTable {
2240        func: TableFunc::Wrap {
2241            width: values[0].len(),
2242            types,
2243        },
2244        exprs,
2245    })
2246}
2247
2248fn plan_join_identity() -> (HirRelationExpr, Scope) {
2249    let typ = SqlRelationType::new(vec![]);
2250    let expr = HirRelationExpr::constant(vec![vec![]], typ);
2251    let scope = Scope::empty();
2252    (expr, scope)
2253}
2254
2255/// Describes how to execute a SELECT query.
2256///
2257/// `order_by` describes how to order the rows in `expr` *before* applying the
2258/// projection. The `scope` describes the columns in `expr` *after* the
2259/// projection has been applied.
2260#[derive(Debug)]
2261struct SelectPlan {
2262    expr: HirRelationExpr,
2263    scope: Scope,
2264    order_by: Vec<ColumnOrder>,
2265    project: Vec<usize>,
2266}
2267
2268generate_extracted_config!(
2269    SelectOption,
2270    (ExpectedGroupSize, u64),
2271    (AggregateInputGroupSize, u64),
2272    (DistinctOnInputGroupSize, u64),
2273    (LimitInputGroupSize, u64)
2274);
2275
2276/// Plans a SELECT query. The SELECT query may contain an intrusive ORDER BY clause.
2277///
2278/// Normally, the ORDER BY clause occurs after the columns specified in the
2279/// SELECT list have been projected. In a query like
2280///
2281///   CREATE TABLE (a int, b int)
2282///   (SELECT a FROM t) UNION (SELECT a FROM t) ORDER BY a
2283///
2284/// it is valid to refer to `a`, because it is explicitly selected, but it would
2285/// not be valid to refer to unselected column `b`.
2286///
2287/// But PostgreSQL extends the standard to permit queries like
2288///
2289///   SELECT a FROM t ORDER BY b
2290///
2291/// where expressions in the ORDER BY clause can refer to *both* input columns
2292/// and output columns.
2293fn plan_select_from_where(
2294    qcx: &QueryContext,
2295    mut s: Select<Aug>,
2296    mut order_by_exprs: Vec<OrderByExpr<Aug>>,
2297) -> Result<SelectPlan, PlanError> {
2298    // TODO: Both `s` and `order_by_exprs` are not references because the
2299    // AggregateTableFuncVisitor needs to be able to rewrite the expressions for
2300    // table function support (the UUID mapping). Attempt to change this so callers
2301    // don't need to clone the Select.
2302
2303    // Extract query options.
2304    let select_option_extracted = SelectOptionExtracted::try_from(s.options.clone())?;
2305    let group_size_hints = GroupSizeHints::try_from(select_option_extracted)?;
2306
2307    // Step 1. Handle FROM clause, including joins.
2308    let (mut relation_expr, mut from_scope) =
2309        s.from.iter().try_fold(plan_join_identity(), |l, twj| {
2310            let (left, left_scope) = l;
2311            plan_join(
2312                qcx,
2313                left,
2314                left_scope,
2315                &Join {
2316                    relation: TableFactor::NestedJoin {
2317                        join: Box::new(twj.clone()),
2318                        alias: None,
2319                    },
2320                    join_operator: JoinOperator::CrossJoin,
2321                },
2322            )
2323        })?;
2324
2325    // Step 2. Handle WHERE clause.
2326    if let Some(selection) = &s.selection {
2327        let ecx = &ExprContext {
2328            qcx,
2329            name: "WHERE clause",
2330            scope: &from_scope,
2331            relation_type: &qcx.relation_type(&relation_expr),
2332            allow_aggregates: false,
2333            allow_subqueries: true,
2334            allow_parameters: true,
2335            allow_windows: false,
2336        };
2337        let expr = plan_expr(ecx, selection)
2338            .map_err(|e| sql_err!("WHERE clause error: {}", e))?
2339            .type_as(ecx, &SqlScalarType::Bool)?;
2340        relation_expr = relation_expr.filter(vec![expr]);
2341    }
2342
2343    // Step 3. Gather aggregates and table functions.
2344    // (But skip window aggregates.)
2345    let (aggregates, table_funcs) = {
2346        let mut visitor = AggregateTableFuncVisitor::new(qcx.scx);
2347        visitor.visit_select_mut(&mut s);
2348        for o in order_by_exprs.iter_mut() {
2349            visitor.visit_order_by_expr_mut(o);
2350        }
2351        visitor.into_result()?
2352    };
2353    let mut table_func_names: BTreeMap<String, Ident> = BTreeMap::new();
2354    // Table functions in the SELECT list apply to the output of the reduce
2355    // (GROUP BY, aggregates, HAVING), but their columns must already be in
2356    // scope when the SELECT list is expanded and GROUP BY items are planned,
2357    // so the join is planned here regardless. Step 5 decides whether the
2358    // reduce consumes this join or the saved pre-join relation, and in the
2359    // latter case Step 8.5 plans the join again on top of the reduce.
2360    let pre_table_funcs_arity = from_scope.len();
2361    let mut pre_table_funcs_relation = None;
2362    let mut table_funcs_deferred = false;
2363    if !table_funcs.is_empty() {
2364        let (expr, scope) = plan_scalar_table_funcs(
2365            qcx,
2366            &table_funcs,
2367            &mut table_func_names,
2368            &relation_expr,
2369            &from_scope,
2370        )?;
2371        if !aggregates.is_empty() || !s.group_by.is_empty() || s.having.is_some() {
2372            pre_table_funcs_relation = Some(relation_expr.clone());
2373        }
2374        relation_expr = relation_expr.join(expr, HirScalarExpr::literal_true(), JoinKind::Inner);
2375        from_scope = from_scope.product(scope)?;
2376    }
2377
2378    // Step 4. Expand SELECT clause.
2379    let projection = {
2380        let ecx = &ExprContext {
2381            qcx,
2382            name: "SELECT clause",
2383            scope: &from_scope,
2384            relation_type: &qcx.relation_type(&relation_expr),
2385            allow_aggregates: true,
2386            allow_subqueries: true,
2387            allow_parameters: true,
2388            allow_windows: true,
2389        };
2390        let mut out = vec![];
2391        for si in &s.projection {
2392            if *si == SelectItem::Wildcard && s.from.is_empty() {
2393                sql_bail!("SELECT * with no tables specified is not valid");
2394            }
2395            out.extend(expand_select_item(ecx, si, &table_func_names)?);
2396        }
2397        out
2398    };
2399
2400    // Step 5. Handle GROUP BY clause.
2401    // This will also plan the aggregates gathered in Step 3.
2402    // See an overview of how aggregates are planned in the doc comment at the top of the file.
2403    let (mut group_scope, select_all_mapping) = {
2404        // Compute GROUP BY expressions.
2405        let ecx = &ExprContext {
2406            qcx,
2407            name: "GROUP BY clause",
2408            scope: &from_scope,
2409            relation_type: &qcx.relation_type(&relation_expr),
2410            allow_aggregates: false,
2411            allow_subqueries: true,
2412            allow_parameters: true,
2413            allow_windows: false,
2414        };
2415        let mut group_key = vec![];
2416        let mut group_exprs: BTreeMap<HirScalarExpr, ScopeItem> = BTreeMap::new();
2417        let mut group_hir_exprs = vec![];
2418        let mut group_scope = Scope::empty();
2419        let mut select_all_mapping = BTreeMap::new();
2420
2421        for group_expr in &s.group_by {
2422            let (group_expr, expr) = plan_group_by_expr(ecx, group_expr, &projection)?;
2423            let new_column = group_key.len();
2424
2425            // Multiple AST expressions can map to the same HIR expression, e.g.
2426            // `GROUP BY 1, 1` or `GROUP BY a, 1` where the positional reference
2427            // `1` resolves to the same column as `a`. When we already have a
2428            // ScopeItem for this HIR expression we must deduplicate: skip adding
2429            // a second group key for it. We must not gate this on `group_expr`
2430            // being `Some`, because a positional reference to an input column
2431            // (e.g. `GROUP BY 1`) has no AST expression to record yet still
2432            // needs to dedup against the existing key; gating on `Some` here is
2433            // what made `GROUP BY 1, 1` panic on the `group_hir_exprs.len() ==
2434            // group_exprs.len()` assertion below.
2435            if let Some(existing_scope_item) = group_exprs.get_mut(&expr) {
2436                // If this AST expression is a named expression (not a bare
2437                // positional reference), record it on the existing ScopeItem so
2438                // name resolution can find it.
2439                if let Some(group_expr) = group_expr {
2440                    existing_scope_item.exprs.insert(group_expr.clone());
2441                }
2442                continue;
2443            }
2444
2445            let mut scope_item = if let HirScalarExpr::Column(
2446                ColumnRef {
2447                    level: 0,
2448                    column: old_column,
2449                },
2450                _name,
2451            ) = &expr
2452            {
2453                // If we later have `SELECT foo.*` then we have to find all
2454                // the `foo` items in `from_scope` and figure out where they
2455                // ended up in `group_scope`. This is really hard to do
2456                // right using SQL name resolution, so instead we just track
2457                // the movement here.
2458                select_all_mapping.insert(*old_column, new_column);
2459                let scope_item = ecx.scope.items[*old_column].clone();
2460                scope_item
2461            } else {
2462                ScopeItem::empty()
2463            };
2464
2465            if let Some(group_expr) = group_expr.cloned() {
2466                scope_item.exprs.insert(group_expr);
2467            }
2468
2469            group_key.push(from_scope.len() + group_exprs.len());
2470            group_hir_exprs.push(expr.clone());
2471            group_exprs.insert(expr, scope_item);
2472        }
2473
2474        assert_eq!(group_hir_exprs.len(), group_exprs.len());
2475        for expr in &group_hir_exprs {
2476            if let Some(scope_item) = group_exprs.remove(expr) {
2477                group_scope.items.push(scope_item);
2478            }
2479        }
2480
2481        // Plan aggregates.
2482        let ecx = &ExprContext {
2483            qcx,
2484            name: "aggregate function",
2485            scope: &from_scope,
2486            relation_type: &qcx.relation_type(&relation_expr.clone().map(group_hir_exprs.clone())),
2487            allow_aggregates: false,
2488            allow_subqueries: true,
2489            allow_parameters: true,
2490            allow_windows: false,
2491        };
2492        let mut agg_exprs = vec![];
2493        for sql_function in aggregates {
2494            if sql_function.over.is_some() {
2495                unreachable!(
2496                    "Window aggregate; AggregateTableFuncVisitor explicitly filters these out"
2497                );
2498            }
2499            agg_exprs.push(plan_aggregate_common(ecx, &sql_function)?);
2500            group_scope
2501                .items
2502                .push(ScopeItem::from_expr(Expr::Function(sql_function.clone())));
2503        }
2504        if !agg_exprs.is_empty() || !group_key.is_empty() || s.having.is_some() {
2505            // Table functions join after the reduce only when no group key or
2506            // aggregate references their columns, e.g. GROUP BY on a SELECT
2507            // list alias of a table function.
2508            if let Some(pre_relation_expr) = pre_table_funcs_relation.take() {
2509                let mut references_table_funcs = false;
2510                let mut check = |column: usize| {
2511                    if column >= pre_table_funcs_arity {
2512                        references_table_funcs = true;
2513                    }
2514                };
2515                for expr in &group_hir_exprs {
2516                    expr.visit_columns_referring_to_root_level(&mut check);
2517                }
2518                for agg_expr in &agg_exprs {
2519                    agg_expr
2520                        .expr
2521                        .visit_columns_referring_to_root_level(&mut check);
2522                }
2523                if !references_table_funcs {
2524                    relation_expr = pre_relation_expr;
2525                    // The group keys point past the table functions' columns,
2526                    // which the saved relation does not have.
2527                    for (i, key) in group_key.iter_mut().enumerate() {
2528                        *key = pre_table_funcs_arity + i;
2529                    }
2530                    table_funcs_deferred = true;
2531                }
2532            }
2533
2534            // apply GROUP BY / aggregates
2535            relation_expr = relation_expr.map(group_hir_exprs).reduce(
2536                group_key,
2537                agg_exprs,
2538                group_size_hints.aggregate_input_group_size,
2539            );
2540
2541            // For every old column that wasn't a group key, add a scope item
2542            // that errors when referenced. We can't simply drop these items
2543            // from scope. These items need to *exist* because they might shadow
2544            // variables in outer scopes that would otherwise be valid to
2545            // reference, but accessing them needs to produce an error.
2546            // Deferred table functions' columns come back into scope in Step
2547            // 8.5, so they must not be recorded as ungrouped.
2548            let ungrouped_arity = if table_funcs_deferred {
2549                pre_table_funcs_arity
2550            } else {
2551                from_scope.len()
2552            };
2553            for i in 0..ungrouped_arity {
2554                if !select_all_mapping.contains_key(&i) {
2555                    let scope_item = &ecx.scope.items[i];
2556                    group_scope.ungrouped_columns.push(ScopeUngroupedColumn {
2557                        table_name: scope_item.table_name.clone(),
2558                        column_name: scope_item.column_name.clone(),
2559                        allow_unqualified_references: scope_item.allow_unqualified_references,
2560                    });
2561                }
2562            }
2563
2564            (group_scope, select_all_mapping)
2565        } else {
2566            // if no GROUP BY, aggregates or having then all columns remain in scope
2567            (
2568                from_scope.clone(),
2569                (0..from_scope.len()).map(|i| (i, i)).collect(),
2570            )
2571        }
2572    };
2573
2574    // Step 6. Handle HAVING clause.
2575    if let Some(ref having) = s.having {
2576        let ecx = &ExprContext {
2577            qcx,
2578            name: "HAVING clause",
2579            scope: &group_scope,
2580            relation_type: &qcx.relation_type(&relation_expr),
2581            allow_aggregates: true,
2582            allow_subqueries: true,
2583            allow_parameters: true,
2584            allow_windows: false,
2585        };
2586        let expr = plan_expr(ecx, having)?.type_as(ecx, &SqlScalarType::Bool)?;
2587        relation_expr = relation_expr.filter(vec![expr]);
2588    }
2589
2590    // Step 7. Gather window functions from SELECT, ORDER BY, and QUALIFY, and plan them.
2591    // (This includes window aggregations.)
2592    //
2593    // Note that window functions can be present only in SELECT, ORDER BY, or QUALIFY (including
2594    // DISTINCT ON), because they are executed after grouped aggregations and HAVING.
2595    //
2596    // Also note that window functions in the ORDER BY can't refer to columns introduced in the
2597    // SELECT. This is because when an output column appears in ORDER BY, it can only stand alone,
2598    // and can't be part of a bigger expression.
2599    // See https://www.postgresql.org/docs/current/queries-order.html:
2600    // "Note that an output column name has to stand alone, that is, it cannot be used in an
2601    // expression"
2602    let window_funcs = {
2603        let mut visitor = WindowFuncCollector::default();
2604        // The `visit_select` call visits both `SELECT` and `QUALIFY` (and many other things, but
2605        // window functions are excluded from other things by `allow_windows` being false when
2606        // planning those before this code).
2607        visitor.visit_select(&s);
2608        for o in order_by_exprs.iter() {
2609            visitor.visit_order_by_expr(o);
2610        }
2611        visitor.into_result()
2612    };
2613    for window_func in window_funcs {
2614        let ecx = &ExprContext {
2615            qcx,
2616            name: "window function",
2617            scope: &group_scope,
2618            relation_type: &qcx.relation_type(&relation_expr),
2619            allow_aggregates: true,
2620            allow_subqueries: true,
2621            allow_parameters: true,
2622            allow_windows: true,
2623        };
2624        relation_expr = relation_expr.map(vec![plan_expr(ecx, &window_func)?.type_as_any(ecx)?]);
2625        group_scope.items.push(ScopeItem::from_expr(window_func));
2626    }
2627    // From this point on, we shouldn't encounter _valid_ window function calls, because those have
2628    // been already planned now. However, we should still set `allow_windows: true` for the
2629    // remaining planning of `QUALIFY`, `SELECT`, and `ORDER BY`, in order to have a correct error
2630    // msg if an OVER clause is missing from a window function.
2631
2632    // Step 8. Handle QUALIFY clause. (very similar to HAVING)
2633    if let Some(ref qualify) = s.qualify {
2634        let ecx = &ExprContext {
2635            qcx,
2636            name: "QUALIFY clause",
2637            scope: &group_scope,
2638            relation_type: &qcx.relation_type(&relation_expr),
2639            allow_aggregates: true,
2640            allow_subqueries: true,
2641            allow_parameters: true,
2642            allow_windows: true,
2643        };
2644        let expr = plan_expr(ecx, qualify)?.type_as(ecx, &SqlScalarType::Bool)?;
2645        relation_expr = relation_expr.filter(vec![expr]);
2646    }
2647
2648    // Step 8.5. Join the table functions deferred in Step 5. Planning them
2649    // again rebinds their arguments' column references to the reduced
2650    // relation.
2651    if table_funcs_deferred {
2652        let (expr, scope) = plan_scalar_table_funcs(
2653            qcx,
2654            &table_funcs,
2655            &mut table_func_names,
2656            &relation_expr,
2657            &group_scope,
2658        )?;
2659        relation_expr = relation_expr.join(expr, HirScalarExpr::literal_true(), JoinKind::Inner);
2660        // `product` resets `ungrouped_columns`, but the ungrouped column
2661        // errors from the reduce must survive for the SELECT list.
2662        let ungrouped_columns = mem::take(&mut group_scope.ungrouped_columns);
2663        group_scope = group_scope.product(scope)?;
2664        group_scope.ungrouped_columns = ungrouped_columns;
2665    }
2666
2667    // Step 9. Handle SELECT clause.
2668    let output_columns = {
2669        let mut new_exprs = vec![];
2670        let mut new_type = qcx.relation_type(&relation_expr);
2671        let mut output_columns = vec![];
2672        for (select_item, column_name) in &projection {
2673            let ecx = &ExprContext {
2674                qcx,
2675                name: "SELECT clause",
2676                scope: &group_scope,
2677                relation_type: &new_type,
2678                allow_aggregates: true,
2679                allow_subqueries: true,
2680                allow_parameters: true,
2681                allow_windows: true,
2682            };
2683            let expr = match select_item {
2684                ExpandedSelectItem::InputOrdinal(i) => {
2685                    if let Some(column) = select_all_mapping.get(i).copied() {
2686                        HirScalarExpr::column(column)
2687                    } else {
2688                        return Err(PlanError::ungrouped_column(&from_scope.items[*i]));
2689                    }
2690                }
2691                ExpandedSelectItem::Expr(expr) => plan_expr(ecx, expr)?.type_as_any(ecx)?,
2692            };
2693            if let HirScalarExpr::Column(ColumnRef { level: 0, column }, _name) = expr {
2694                // Simple column reference; no need to map on a new expression.
2695                output_columns.push((column, column_name));
2696            } else {
2697                // Complicated expression that requires a map expression. We
2698                // update `group_scope` as we go so that future expressions that
2699                // are textually identical to this one can reuse it. This
2700                // duplicate detection is required for proper determination of
2701                // ambiguous column references with SQL92-style `ORDER BY`
2702                // items. See `plan_order_by_or_distinct_expr` for more.
2703                let typ = ecx.column_type(&expr);
2704                new_type.column_types.push(typ);
2705                new_exprs.push(expr);
2706                output_columns.push((group_scope.len(), column_name));
2707                group_scope
2708                    .items
2709                    .push(ScopeItem::from_expr(select_item.as_expr().cloned()));
2710            }
2711        }
2712        relation_expr = relation_expr.map(new_exprs);
2713        output_columns
2714    };
2715    let mut project_key: Vec<_> = output_columns.iter().map(|(i, _name)| *i).collect();
2716
2717    // Step 10. Handle intrusive ORDER BY and DISTINCT.
2718    let order_by = {
2719        let relation_type = qcx.relation_type(&relation_expr);
2720        let (mut order_by, mut map_exprs) = plan_order_by_exprs(
2721            &ExprContext {
2722                qcx,
2723                name: "ORDER BY clause",
2724                scope: &group_scope,
2725                relation_type: &relation_type,
2726                allow_aggregates: true,
2727                allow_subqueries: true,
2728                allow_parameters: true,
2729                allow_windows: true,
2730            },
2731            &order_by_exprs,
2732            &output_columns,
2733        )?;
2734
2735        match s.distinct {
2736            None => relation_expr = relation_expr.map(map_exprs),
2737            Some(Distinct::EntireRow) => {
2738                if relation_type.arity() == 0 {
2739                    sql_bail!("SELECT DISTINCT must have at least one column");
2740                }
2741                // `SELECT DISTINCT` only distincts on the columns in the SELECT
2742                // list, so we can't proceed if `ORDER BY` has introduced any
2743                // columns for arbitrary expressions. This matches PostgreSQL.
2744                if !try_push_projection_order_by(
2745                    &mut relation_expr,
2746                    &mut project_key,
2747                    &mut order_by,
2748                ) {
2749                    sql_bail!(
2750                        "for SELECT DISTINCT, ORDER BY expressions must appear in select list"
2751                    );
2752                }
2753                assert!(map_exprs.is_empty());
2754                relation_expr = relation_expr.distinct();
2755            }
2756            Some(Distinct::On(exprs)) => {
2757                // The table functions deferred in Step 5 join below this TopK,
2758                // so the distinct would collapse their expansion rather than
2759                // expand the rows the distinct picks. PostgreSQL instead
2760                // evaluates a SELECT list table function after the distinct
2761                // whenever the query has an ORDER BY and the function's output
2762                // is not itself a distinct or sort key. Reject these queries
2763                // rather than answer them differently.
2764                if table_funcs_deferred && !order_by_exprs.is_empty() {
2765                    bail_unsupported!(
2766                        "SELECT list table function with DISTINCT ON and ORDER BY over an aggregation"
2767                    );
2768                }
2769
2770                let ecx = &ExprContext {
2771                    qcx,
2772                    name: "DISTINCT ON clause",
2773                    scope: &group_scope,
2774                    relation_type: &qcx.relation_type(&relation_expr),
2775                    allow_aggregates: true,
2776                    allow_subqueries: true,
2777                    allow_parameters: true,
2778                    allow_windows: true,
2779                };
2780
2781                let mut distinct_exprs = vec![];
2782                for expr in &exprs {
2783                    let expr = plan_order_by_or_distinct_expr(ecx, expr, &output_columns)?;
2784                    distinct_exprs.push(expr);
2785                }
2786
2787                let mut distinct_key = vec![];
2788
2789                // If both `DISTINCT ON` and `ORDER BY` are specified, then the
2790                // `DISTINCT ON` expressions must match the initial `ORDER BY`
2791                // expressions, though the order of `DISTINCT ON` expressions
2792                // does not matter. This matches PostgreSQL and leaves the door
2793                // open to a future optimization where the `DISTINCT ON` and
2794                // `ORDER BY` operations happen in one pass.
2795                //
2796                // On the bright side, any columns that have already been
2797                // computed by `ORDER BY` can be reused in the distinct key.
2798                let arity = relation_type.arity();
2799                for ord in order_by.iter().take(distinct_exprs.len()) {
2800                    // The unusual construction of `expr` here is to ensure the
2801                    // temporary column expression lives long enough.
2802                    let mut expr = &HirScalarExpr::column(ord.column);
2803                    if ord.column >= arity {
2804                        expr = &map_exprs[ord.column - arity];
2805                    };
2806                    match distinct_exprs.iter().position(move |e| e == expr) {
2807                        None => sql_bail!(
2808                            "SELECT DISTINCT ON expressions must match initial ORDER BY expressions"
2809                        ),
2810                        Some(pos) => {
2811                            distinct_exprs.remove(pos);
2812                        }
2813                    }
2814                    distinct_key.push(ord.column);
2815                }
2816
2817                // Add any remaining `DISTINCT ON` expressions to the key.
2818                for expr in distinct_exprs {
2819                    // If the expression is a reference to an existing column,
2820                    // do not introduce a new column to support it.
2821                    let column = match expr {
2822                        HirScalarExpr::Column(ColumnRef { level: 0, column }, _name) => column,
2823                        _ => {
2824                            map_exprs.push(expr);
2825                            arity + map_exprs.len() - 1
2826                        }
2827                    };
2828                    distinct_key.push(column);
2829                }
2830
2831                // `DISTINCT ON` is semantically a TopK with limit 1. The
2832                // columns in `ORDER BY` that are not part of the distinct key,
2833                // if there are any, determine the ordering within each group,
2834                // per PostgreSQL semantics.
2835                let distinct_len = distinct_key.len();
2836                relation_expr = HirRelationExpr::top_k(
2837                    relation_expr.map(map_exprs),
2838                    distinct_key,
2839                    order_by.iter().skip(distinct_len).cloned().collect(),
2840                    Some(HirScalarExpr::literal(
2841                        Datum::Int64(1),
2842                        SqlScalarType::Int64,
2843                    )),
2844                    HirScalarExpr::literal(Datum::Int64(0), SqlScalarType::Int64),
2845                    group_size_hints.distinct_on_input_group_size,
2846                );
2847            }
2848        }
2849
2850        order_by
2851    };
2852
2853    // Construct a clean scope to expose outwards, where all of the state that
2854    // accumulated in the scope during planning of this SELECT is erased. The
2855    // clean scope has at most one name for each column, and the names are not
2856    // associated with any table.
2857    let scope = Scope::from_source(None, projection.into_iter().map(|(_expr, name)| name));
2858
2859    Ok(SelectPlan {
2860        expr: relation_expr,
2861        scope,
2862        order_by,
2863        project: project_key,
2864    })
2865}
2866
2867fn plan_scalar_table_funcs(
2868    qcx: &QueryContext,
2869    table_funcs: &BTreeMap<Function<Aug>, String>,
2870    table_func_names: &mut BTreeMap<String, Ident>,
2871    relation_expr: &HirRelationExpr,
2872    from_scope: &Scope,
2873) -> Result<(HirRelationExpr, Scope), PlanError> {
2874    let rows_from_qcx = qcx.derived_context(from_scope.clone(), qcx.relation_type(relation_expr));
2875
2876    for (table_func, id) in table_funcs.iter() {
2877        table_func_names.insert(
2878            id.clone(),
2879            // TODO(parkmycar): Re-visit after having `FullItemName` use `Ident`s.
2880            Ident::new_unchecked(table_func.name.full_item_name().item.clone()),
2881        );
2882    }
2883    // If there's only a single table function, we can skip generating
2884    // ordinality columns.
2885    if table_funcs.len() == 1 {
2886        let (table_func, id) = table_funcs.iter().next().unwrap();
2887        let (expr, mut scope) =
2888            plan_solitary_table_function(&rows_from_qcx, table_func, None, false)?;
2889
2890        // A single table-function might return several columns as a record
2891        let num_cols = scope.len();
2892        for i in 0..scope.len() {
2893            scope.items[i].table_name = Some(PartialItemName {
2894                database: None,
2895                schema: None,
2896                item: id.clone(),
2897            });
2898            scope.items[i].from_single_column_function = num_cols == 1;
2899            scope.items[i].allow_unqualified_references = false;
2900        }
2901        return Ok((expr, scope));
2902    }
2903    if table_funcs.keys().any(is_repeat_row) {
2904        // Note: Would be also caught by WITH ORDINALITY checking for `repeat_row`, but then the
2905        // error message would be misleading, because it would refer to WITH ORDINALITY.
2906        bail_unsupported!(format!(
2907            "{} in a SELECT clause with multiple table functions",
2908            REPEAT_ROW_NAME
2909        ));
2910    }
2911    // Otherwise, plan as usual, emulating the ROWS FROM behavior
2912    let (expr, mut scope, num_cols) =
2913        plan_rows_from_internal(&rows_from_qcx, table_funcs.keys(), None)?;
2914
2915    // Munge the scope so table names match with the generated ids.
2916    let mut i = 0;
2917    for (id, num_cols) in table_funcs.values().zip_eq(num_cols) {
2918        for _ in 0..num_cols {
2919            scope.items[i].table_name = Some(PartialItemName {
2920                database: None,
2921                schema: None,
2922                item: id.clone(),
2923            });
2924            scope.items[i].from_single_column_function = num_cols == 1;
2925            scope.items[i].allow_unqualified_references = false;
2926            i += 1;
2927        }
2928        // Ordinality column. This doubles as the
2929        // `is_exists_column_for_a_table_function_that_was_in_the_target_list` later on
2930        // because it only needs to be NULL or not.
2931        scope.items[i].table_name = Some(PartialItemName {
2932            database: None,
2933            schema: None,
2934            item: id.clone(),
2935        });
2936        scope.items[i].is_exists_column_for_a_table_function_that_was_in_the_target_list = true;
2937        scope.items[i].allow_unqualified_references = false;
2938        i += 1;
2939    }
2940    // Coalesced ordinality column.
2941    scope.items[i].allow_unqualified_references = false;
2942    Ok((expr, scope))
2943}
2944
2945/// Plans an expression in a `GROUP BY` clause.
2946///
2947/// For historical reasons, PostgreSQL allows `GROUP BY` expressions to refer to
2948/// names/expressions defined in the `SELECT` clause. These special cases are
2949/// handled by this function; see comments within the implementation for
2950/// details.
2951fn plan_group_by_expr<'a>(
2952    ecx: &ExprContext,
2953    group_expr: &'a Expr<Aug>,
2954    projection: &'a [(ExpandedSelectItem, ColumnName)],
2955) -> Result<(Option<&'a Expr<Aug>>, HirScalarExpr), PlanError> {
2956    let plan_projection = |column: usize| match &projection[column].0 {
2957        ExpandedSelectItem::InputOrdinal(column) => Ok((None, HirScalarExpr::column(*column))),
2958        ExpandedSelectItem::Expr(expr) => {
2959            Ok((Some(expr.as_ref()), plan_expr(ecx, expr)?.type_as_any(ecx)?))
2960        }
2961    };
2962
2963    // Check if the expression is a numeric literal, as in `GROUP BY 1`. This is
2964    // a special case that means to use the ith item in the SELECT clause.
2965    if let Some(column) = check_col_index(ecx.name, group_expr, projection.len())? {
2966        return plan_projection(column);
2967    }
2968
2969    // Check if the expression is a simple identifier, as in `GROUP BY foo`.
2970    // The `foo` can refer to *either* an input column or an output column. If
2971    // both exist, the input column is preferred.
2972    match group_expr {
2973        Expr::Identifier(names) => match plan_identifier(ecx, names) {
2974            Err(PlanError::UnknownColumn {
2975                table: None,
2976                column,
2977                similar,
2978            }) => {
2979                // The expression was a simple identifier that did not match an
2980                // input column. See if it matches an output column.
2981                let mut iter = projection.iter().map(|(_expr, name)| name);
2982                if let Some(i) = iter.position(|n| *n == column) {
2983                    if iter.any(|n| *n == column) {
2984                        Err(PlanError::AmbiguousColumn(column))
2985                    } else {
2986                        plan_projection(i)
2987                    }
2988                } else {
2989                    // The name didn't match an output column either. Return the
2990                    // "unknown column" error.
2991                    Err(PlanError::UnknownColumn {
2992                        table: None,
2993                        column,
2994                        similar,
2995                    })
2996                }
2997            }
2998            res => Ok((Some(group_expr), res?)),
2999        },
3000        _ => Ok((
3001            Some(group_expr),
3002            plan_expr(ecx, group_expr)?.type_as_any(ecx)?,
3003        )),
3004    }
3005}
3006
3007/// Plans a slice of `ORDER BY` expressions.
3008///
3009/// See `plan_order_by_or_distinct_expr` for details on the `output_columns`
3010/// parameter.
3011///
3012/// Returns the determined column orderings and a list of scalar expressions
3013/// that must be mapped onto the underlying relation expression.
3014pub(crate) fn plan_order_by_exprs(
3015    ecx: &ExprContext,
3016    order_by_exprs: &[OrderByExpr<Aug>],
3017    output_columns: &[(usize, &ColumnName)],
3018) -> Result<(Vec<ColumnOrder>, Vec<HirScalarExpr>), PlanError> {
3019    let mut order_by = vec![];
3020    let mut map_exprs = vec![];
3021    for obe in order_by_exprs {
3022        let expr = plan_order_by_or_distinct_expr(ecx, &obe.expr, output_columns)?;
3023        // If the expression is a reference to an existing column,
3024        // do not introduce a new column to support it.
3025        let column = match expr {
3026            HirScalarExpr::Column(ColumnRef { level: 0, column }, _name) => column,
3027            _ => {
3028                map_exprs.push(expr);
3029                ecx.relation_type.arity() + map_exprs.len() - 1
3030            }
3031        };
3032        order_by.push(resolve_desc_and_nulls_last(obe, column));
3033    }
3034    Ok((order_by, map_exprs))
3035}
3036
3037/// Plans an expression that appears in an `ORDER BY` or `DISTINCT ON` clause.
3038///
3039/// The `output_columns` parameter describes, in order, the physical index and
3040/// name of each expression in the `SELECT` list. For example, `[(3, "a")]`
3041/// corresponds to a `SELECT` list with a single entry named "a" that can be
3042/// found at index 3 in the underlying relation expression.
3043///
3044/// There are three cases to handle.
3045///
3046///    1. A simple numeric literal, as in `ORDER BY 1`. This is an ordinal
3047///       reference to the specified output column.
3048///    2. An unqualified identifier, as in `ORDER BY a`. This is a reference to
3049///       an output column, if it exists; otherwise it is a reference to an
3050///       input column.
3051///    3. An arbitrary expression, as in `ORDER BY -a`. Column references in
3052///       arbitrary expressions exclusively refer to input columns, never output
3053///       columns.
3054fn plan_order_by_or_distinct_expr(
3055    ecx: &ExprContext,
3056    expr: &Expr<Aug>,
3057    output_columns: &[(usize, &ColumnName)],
3058) -> Result<HirScalarExpr, PlanError> {
3059    if let Some(i) = check_col_index(ecx.name, expr, output_columns.len())? {
3060        return Ok(HirScalarExpr::column(output_columns[i].0));
3061    }
3062
3063    if let Expr::Identifier(names) = expr {
3064        if let [name] = &names[..] {
3065            let name = normalize::column_name(name.clone());
3066            let mut iter = output_columns.iter().filter(|(_, n)| **n == name);
3067            if let Some((i, _)) = iter.next() {
3068                match iter.next() {
3069                    // Per SQL92, names are not considered ambiguous if they
3070                    // refer to identical target list expressions, as in
3071                    // `SELECT a + 1 AS foo, a + 1 AS foo ... ORDER BY foo`.
3072                    Some((i2, _)) if i != i2 => return Err(PlanError::AmbiguousColumn(name)),
3073                    _ => return Ok(HirScalarExpr::column(*i)),
3074                }
3075            }
3076        }
3077    }
3078
3079    plan_expr(ecx, expr)?.type_as_any(ecx)
3080}
3081
3082fn plan_table_with_joins(
3083    qcx: &QueryContext,
3084    table_with_joins: &TableWithJoins<Aug>,
3085) -> Result<(HirRelationExpr, Scope), PlanError> {
3086    let (mut expr, mut scope) = plan_table_factor(qcx, &table_with_joins.relation)?;
3087    for join in &table_with_joins.joins {
3088        let (new_expr, new_scope) = plan_join(qcx, expr, scope, join)?;
3089        expr = new_expr;
3090        scope = new_scope;
3091    }
3092    Ok((expr, scope))
3093}
3094
3095fn plan_table_factor(
3096    qcx: &QueryContext,
3097    table_factor: &TableFactor<Aug>,
3098) -> Result<(HirRelationExpr, Scope), PlanError> {
3099    match table_factor {
3100        TableFactor::Table { name, alias } => {
3101            let (expr, scope) = qcx.resolve_table_name(name.clone())?;
3102            let scope = plan_table_alias(scope, alias.as_ref())?;
3103            Ok((expr, scope))
3104        }
3105
3106        TableFactor::Function {
3107            function,
3108            alias,
3109            with_ordinality,
3110        } => plan_solitary_table_function(qcx, function, alias.as_ref(), *with_ordinality),
3111
3112        TableFactor::RowsFrom {
3113            functions,
3114            alias,
3115            with_ordinality,
3116        } => plan_rows_from(qcx, functions, alias.as_ref(), *with_ordinality),
3117
3118        TableFactor::Derived {
3119            lateral,
3120            subquery,
3121            alias,
3122        } => {
3123            let mut qcx = (*qcx).clone();
3124            if !lateral {
3125                // Since this derived table was not marked as `LATERAL`,
3126                // make elements in outer scopes invisible until we reach the
3127                // next lateral barrier.
3128                for scope in &mut qcx.outer_scopes {
3129                    if scope.lateral_barrier {
3130                        break;
3131                    }
3132                    scope.items.clear();
3133                }
3134            }
3135            qcx.outer_scopes[0].lateral_barrier = true;
3136            let (expr, scope) = plan_nested_query(&mut qcx, subquery)?;
3137            let scope = plan_table_alias(scope, alias.as_ref())?;
3138            Ok((expr, scope))
3139        }
3140
3141        TableFactor::NestedJoin { join, alias } => {
3142            let (expr, scope) = plan_table_with_joins(qcx, join)?;
3143            let scope = plan_table_alias(scope, alias.as_ref())?;
3144            Ok((expr, scope))
3145        }
3146    }
3147}
3148
3149/// Plans a `ROWS FROM` expression.
3150///
3151/// `ROWS FROM` concatenates table functions into a single table, filling in
3152/// `NULL`s in places where one table function has fewer rows than another. We
3153/// can achieve this by augmenting each table function with a row number, doing
3154/// a `FULL JOIN` between each table function on the row number and eventually
3155/// projecting away the row number columns. Concretely, the following query
3156/// using `ROWS FROM`
3157///
3158/// ```sql
3159/// SELECT
3160///     *
3161/// FROM
3162///     ROWS FROM (
3163///         generate_series(1, 2),
3164///         information_schema._pg_expandarray(ARRAY[9]),
3165///         generate_series(3, 6)
3166///     );
3167/// ```
3168///
3169/// is equivalent to the following query that does not use `ROWS FROM`:
3170///
3171/// ```sql
3172/// SELECT
3173///     gs1.generate_series, expand.x, expand.n, gs2.generate_series
3174/// FROM
3175///     generate_series(1, 2) WITH ORDINALITY AS gs1
3176///     FULL JOIN information_schema._pg_expandarray(ARRAY[9]) WITH ORDINALITY AS expand
3177///         ON gs1.ordinality = expand.ordinality
3178///     FULL JOIN generate_series(3, 6) WITH ORDINALITY AS gs3
3179///         ON coalesce(gs1.ordinality, expand.ordinality) = gs3.ordinality;
3180/// ```
3181///
3182/// Note the call to `coalesce` in the last join condition, which ensures that
3183/// `gs3` will align with whichever of `gs1` or `expand` has more rows.
3184///
3185/// This function creates a HirRelationExpr that follows the structure of the
3186/// latter query.
3187///
3188/// `with_ordinality` can be used to have the output expression contain a
3189/// single coalesced ordinality column at the end of the entire expression.
3190fn plan_rows_from(
3191    qcx: &QueryContext,
3192    functions: &[Function<Aug>],
3193    alias: Option<&TableAlias>,
3194    with_ordinality: bool,
3195) -> Result<(HirRelationExpr, Scope), PlanError> {
3196    // The `repeat_row` function is not supported in ROWS FROM.
3197    if functions.iter().any(is_repeat_row) {
3198        // Note: Would be also caught by WITH ORDINALITY checking for `repeat_row`, but then the
3199        // error message would be misleading, because it would refer to WITH ORDINALITY instead of
3200        // ROWS FROM.
3201        bail_unsupported!(format!("{} in ROWS FROM", REPEAT_ROW_NAME));
3202    }
3203
3204    // If there's only a single table function, planning proceeds as if `ROWS
3205    // FROM` hadn't been written at all.
3206    if let [function] = functions {
3207        return plan_solitary_table_function(qcx, function, alias, with_ordinality);
3208    }
3209
3210    // Per PostgreSQL, all scope items take the name of the first function
3211    // (unless aliased).
3212    // See: https://github.com/postgres/postgres/blob/639a86e36/src/backend/parser/parse_relation.c#L1701-L1705
3213    let (expr, mut scope, num_cols) = plan_rows_from_internal(
3214        qcx,
3215        functions,
3216        Some(functions[0].name.full_item_name().clone()),
3217    )?;
3218
3219    // Columns tracks the set of columns we will keep in the projection.
3220    let mut columns = Vec::new();
3221    let mut offset = 0;
3222    // Retain table function's non-ordinality columns.
3223    for (idx, cols) in num_cols.into_iter().enumerate() {
3224        for i in 0..cols {
3225            columns.push(offset + i);
3226        }
3227        offset += cols + 1;
3228
3229        // Remove the ordinality column from the scope, accounting for previous scope
3230        // changes from this loop.
3231        scope.items.remove(offset - idx - 1);
3232    }
3233
3234    // If `WITH ORDINALITY` was specified, include the coalesced ordinality
3235    // column. Otherwise remove it from the scope.
3236    if with_ordinality {
3237        columns.push(offset);
3238    } else {
3239        scope.items.pop();
3240    }
3241
3242    let expr = expr.project(columns);
3243
3244    let scope = plan_table_alias(scope, alias)?;
3245    Ok((expr, scope))
3246}
3247
3248fn is_repeat_row(f: &Function<Aug>) -> bool {
3249    f.name.full_name_str().as_str() == format!("{}.{}", MZ_CATALOG_SCHEMA, REPEAT_ROW_NAME)
3250}
3251
3252/// Plans an expression coalescing multiple table functions. Each table
3253/// function is followed by its row ordinality. The entire expression is
3254/// followed by the coalesced row ordinality.
3255///
3256/// The returned Scope will set all item's table_name's to the `table_name`
3257/// parameter if it is `Some`. If `None`, they will be the name of each table
3258/// function.
3259///
3260/// The returned `Vec<usize>` is the number of (non-ordinality) columns from
3261/// each table function.
3262///
3263/// For example, with table functions tf1 returning 1 column (a) and tf2
3264/// returning 2 columns (b, c), this function will return an expr 6 columns:
3265///
3266/// - tf1.a
3267/// - tf1.ordinality
3268/// - tf2.b
3269/// - tf2.c
3270/// - tf2.ordinality
3271/// - coalesced_ordinality
3272///
3273/// And a `Vec<usize>` of `[1, 2]`.
3274fn plan_rows_from_internal<'a>(
3275    qcx: &QueryContext,
3276    functions: impl IntoIterator<Item = &'a Function<Aug>>,
3277    table_name: Option<FullItemName>,
3278) -> Result<(HirRelationExpr, Scope, Vec<usize>), PlanError> {
3279    let mut functions = functions.into_iter();
3280    let mut num_cols = Vec::new();
3281
3282    // Join together each of the table functions in turn. The last column is
3283    // always the column to join against and is maintained to be the coalescence
3284    // of the row number column for all prior functions.
3285    let (mut left_expr, mut left_scope) =
3286        plan_table_function_internal(qcx, functions.next().unwrap(), true, table_name.clone())?;
3287    num_cols.push(left_scope.len() - 1);
3288    // Create the coalesced ordinality column.
3289    left_expr = left_expr.map(vec![HirScalarExpr::column(left_scope.len() - 1)]);
3290    left_scope
3291        .items
3292        .push(ScopeItem::from_column_name(ORDINALITY_COL_NAME));
3293
3294    for function in functions {
3295        // The right hand side of a join must be planned in a new scope.
3296        let qcx = qcx.empty_derived_context();
3297        let (right_expr, mut right_scope) =
3298            plan_table_function_internal(&qcx, function, true, table_name.clone())?;
3299        num_cols.push(right_scope.len() - 1);
3300        let left_col = left_scope.len() - 1;
3301        let right_col = left_scope.len() + right_scope.len() - 1;
3302        let on = HirScalarExpr::call_binary(
3303            HirScalarExpr::column(left_col),
3304            HirScalarExpr::column(right_col),
3305            expr_func::Eq,
3306        );
3307        left_expr = left_expr
3308            .join(right_expr, on, JoinKind::FullOuter)
3309            .map(vec![HirScalarExpr::call_variadic(
3310                Coalesce,
3311                vec![
3312                    HirScalarExpr::column(left_col),
3313                    HirScalarExpr::column(right_col),
3314                ],
3315            )]);
3316
3317        // Project off the previous iteration's coalesced column, but keep both of this
3318        // iteration's ordinality columns.
3319        left_expr = left_expr.project(
3320            (0..left_col) // non-coalesced ordinality columns from left function
3321                .chain(left_col + 1..right_col + 2) // non-ordinality columns from right function
3322                .collect(),
3323        );
3324        // Move the coalesced ordinality column.
3325        right_scope.items.push(left_scope.items.pop().unwrap());
3326
3327        left_scope.items.extend(right_scope.items);
3328    }
3329
3330    Ok((left_expr, left_scope, num_cols))
3331}
3332
3333/// Plans a table function that appears alone, i.e., that is not part of a `ROWS
3334/// FROM` clause that contains other table functions. Special aliasing rules
3335/// apply.
3336fn plan_solitary_table_function(
3337    qcx: &QueryContext,
3338    function: &Function<Aug>,
3339    alias: Option<&TableAlias>,
3340    with_ordinality: bool,
3341) -> Result<(HirRelationExpr, Scope), PlanError> {
3342    let (expr, mut scope) = plan_table_function_internal(qcx, function, with_ordinality, None)?;
3343
3344    let single_column_function = scope.len() == 1 + if with_ordinality { 1 } else { 0 };
3345    if single_column_function {
3346        let item = &mut scope.items[0];
3347
3348        // Mark that the function only produced a single column. This impacts
3349        // whole-row references.
3350        item.from_single_column_function = true;
3351
3352        // Strange special case for solitary table functions that output one
3353        // column whose name matches the name of the table function. If a table
3354        // alias is provided, the column name is changed to the table alias's
3355        // name. Concretely, the following query returns a column named `x`
3356        // rather than a column named `generate_series`:
3357        //
3358        //     SELECT * FROM generate_series(1, 5) AS x
3359        //
3360        // Note that this case does not apply to e.g. `jsonb_array_elements`,
3361        // since its output column is explicitly named `value`, not
3362        // `jsonb_array_elements`.
3363        //
3364        // Note also that we may (correctly) change the column name again when
3365        // we plan the table alias below if the `alias.columns` is non-empty.
3366        if let Some(alias) = alias {
3367            if let ScopeItem {
3368                table_name: Some(table_name),
3369                column_name,
3370                ..
3371            } = item
3372            {
3373                if table_name.item.as_str() == column_name.as_str() {
3374                    *column_name = normalize::column_name(alias.name.clone());
3375                }
3376            }
3377        }
3378    }
3379
3380    let scope = plan_table_alias(scope, alias)?;
3381    Ok((expr, scope))
3382}
3383
3384/// Plans a table function.
3385///
3386/// You generally should call `plan_rows_from` or `plan_solitary_table_function`
3387/// instead to get the appropriate aliasing behavior.
3388fn plan_table_function_internal(
3389    qcx: &QueryContext,
3390    Function {
3391        name,
3392        args,
3393        filter,
3394        over,
3395        distinct,
3396    }: &Function<Aug>,
3397    with_ordinality: bool,
3398    table_name: Option<FullItemName>,
3399) -> Result<(HirRelationExpr, Scope), PlanError> {
3400    // The parser rejects FILTER, OVER, and DISTINCT in every table function
3401    // position (`FROM f(...)`, `ROWS FROM (...)`), and table functions in
3402    // scalar position are only lifted into a `FROM` clause when all three are
3403    // absent, so these are defensive.
3404    if filter.is_some() {
3405        sql_bail!("FILTER is not allowed for table functions in FROM");
3406    }
3407    if over.is_some() {
3408        sql_bail!("OVER is not allowed for table functions in FROM");
3409    }
3410    if *distinct {
3411        sql_bail!("DISTINCT is not allowed for table functions in FROM");
3412    }
3413
3414    let ecx = &ExprContext {
3415        qcx,
3416        name: "table function arguments",
3417        scope: &Scope::empty(),
3418        relation_type: &SqlRelationType::empty(),
3419        allow_aggregates: false,
3420        allow_subqueries: true,
3421        allow_parameters: true,
3422        allow_windows: false,
3423    };
3424
3425    let scalar_args = match args {
3426        FunctionArgs::Star => sql_bail!("{} does not accept * as an argument", name),
3427        FunctionArgs::Args { args, order_by } => {
3428            if !order_by.is_empty() {
3429                sql_bail!(
3430                    "ORDER BY specified, but {} is not an aggregate function",
3431                    name
3432                );
3433            }
3434            plan_exprs(ecx, args)?
3435        }
3436    };
3437
3438    let table_name = match table_name {
3439        Some(table_name) => table_name.item,
3440        None => name.full_item_name().item.clone(),
3441    };
3442
3443    let scope_name = Some(PartialItemName {
3444        database: None,
3445        schema: None,
3446        item: table_name,
3447    });
3448
3449    let (expr, mut scope) = match resolve_func(ecx, name, args)? {
3450        Func::Table(impls) => {
3451            let tf = func::select_impl(ecx, FuncSpec::Func(name), impls, scalar_args, vec![])?;
3452            let scope = Scope::from_source(scope_name.clone(), tf.column_names);
3453            let expr = match tf.imp {
3454                TableFuncImpl::CallTable { mut func, exprs } => {
3455                    if with_ordinality {
3456                        func = TableFunc::with_ordinality(func.clone()).ok_or(
3457                            PlanError::Unsupported {
3458                                feature: format!("WITH ORDINALITY on {}", func),
3459                                discussion_no: None,
3460                            },
3461                        )?;
3462                    }
3463                    HirRelationExpr::CallTable { func, exprs }
3464                }
3465                TableFuncImpl::Expr(expr) => {
3466                    if !with_ordinality {
3467                        expr
3468                    } else {
3469                        // The table function is defined by a SQL query (i.e., TableFuncImpl::Expr),
3470                        // so we can't use the new `WITH ORDINALITY` implementation. We can fall
3471                        // back to the legacy implementation or error out the query.
3472                        if qcx
3473                            .scx
3474                            .is_feature_flag_enabled(&ENABLE_WITH_ORDINALITY_LEGACY_FALLBACK)
3475                        {
3476                            // Note that this can give an incorrect ordering, and also has an extreme
3477                            // performance problem in some cases. See the doc comment of
3478                            // `TableFuncImpl`.
3479                            tracing::error!(
3480                                %name,
3481                                "Using the legacy WITH ORDINALITY / ROWS FROM implementation for a table function",
3482                            );
3483                            expr.map(vec![HirScalarExpr::windowing(WindowExpr {
3484                                func: WindowExprType::Scalar(ScalarWindowExpr {
3485                                    func: ScalarWindowFunc::RowNumber,
3486                                    order_by: vec![],
3487                                }),
3488                                partition_by: vec![],
3489                                order_by: vec![],
3490                            })])
3491                        } else {
3492                            bail_unsupported!(format!(
3493                                "WITH ORDINALITY or ROWS FROM with {}",
3494                                name
3495                            ));
3496                        }
3497                    }
3498                }
3499            };
3500            (expr, scope)
3501        }
3502        Func::Scalar(impls) => {
3503            let expr = func::select_impl(ecx, FuncSpec::Func(name), impls, scalar_args, vec![])?;
3504            let output = expr.typ(
3505                &qcx.outer_relation_types,
3506                &SqlRelationType::new(vec![]),
3507                &qcx.scx.param_types.borrow(),
3508            );
3509
3510            let relation = SqlRelationType::new(vec![output]);
3511
3512            let function_ident = Ident::new(name.full_item_name().item.clone())?;
3513            let column_name = normalize::column_name(function_ident);
3514            let name = column_name.to_string();
3515
3516            let scope = Scope::from_source(scope_name.clone(), vec![column_name]);
3517
3518            let mut func = TableFunc::TabletizedScalar { relation, name };
3519            if with_ordinality {
3520                func = TableFunc::with_ordinality(func.clone()).ok_or(PlanError::Unsupported {
3521                    feature: format!("WITH ORDINALITY on {}", func),
3522                    discussion_no: None,
3523                })?;
3524            }
3525            (
3526                HirRelationExpr::CallTable {
3527                    func,
3528                    exprs: vec![expr],
3529                },
3530                scope,
3531            )
3532        }
3533        o => sql_bail!(
3534            "{} functions are not supported in functions in FROM",
3535            o.class()
3536        ),
3537    };
3538
3539    if with_ordinality {
3540        scope
3541            .items
3542            .push(ScopeItem::from_name(scope_name, "ordinality"));
3543    }
3544
3545    Ok((expr, scope))
3546}
3547
3548fn plan_table_alias(mut scope: Scope, alias: Option<&TableAlias>) -> Result<Scope, PlanError> {
3549    if let Some(TableAlias {
3550        name,
3551        columns,
3552        strict,
3553    }) = alias
3554    {
3555        if (columns.len() > scope.items.len()) || (*strict && columns.len() != scope.items.len()) {
3556            sql_bail!(
3557                "{} has {} columns available but {} columns specified",
3558                name,
3559                scope.items.len(),
3560                columns.len()
3561            );
3562        }
3563
3564        let table_name = normalize::ident(name.to_owned());
3565        for (i, item) in scope.items.iter_mut().enumerate() {
3566            item.table_name = if item.allow_unqualified_references {
3567                Some(PartialItemName {
3568                    database: None,
3569                    schema: None,
3570                    item: table_name.clone(),
3571                })
3572            } else {
3573                // Columns that prohibit unqualified references are special
3574                // columns from the output of a NATURAL or USING join that can
3575                // only be referenced by their full, pre-join name. Applying an
3576                // alias to the output of that join renders those columns
3577                // inaccessible, which we accomplish here by setting the
3578                // table name to `None`.
3579                //
3580                // Concretely, consider:
3581                //
3582                //      CREATE TABLE t1 (a int);
3583                //      CREATE TABLE t2 (a int);
3584                //  (1) SELECT ... FROM (t1 NATURAL JOIN t2);
3585                //  (2) SELECT ... FROM (t1 NATURAL JOIN t2) AS t;
3586                //
3587                // In (1), the join has no alias. The underlying columns from
3588                // either side of the join can be referenced as `t1.a` and
3589                // `t2.a`, respectively, and the unqualified name `a` refers to
3590                // a column whose value is `coalesce(t1.a, t2.a)`.
3591                //
3592                // In (2), the join is aliased as `t`. The columns from either
3593                // side of the join (`t1.a` and `t2.a`) are inaccessible, and
3594                // the coalesced column can be named as either `a` or `t.a`.
3595                //
3596                // We previously had a bug [0] that mishandled this subtle
3597                // logic.
3598                //
3599                // NOTE(benesch): We could in theory choose to project away
3600                // those inaccessible columns and drop them from the scope
3601                // entirely, but that would require that this function also
3602                // take and return the `HirRelationExpr` that is being aliased,
3603                // which is a rather large refactor.
3604                //
3605                // [0]: https://github.com/MaterializeInc/database-issues/issues/4887
3606                None
3607            };
3608            item.column_name = columns
3609                .get(i)
3610                .map(|a| normalize::column_name(a.clone()))
3611                .unwrap_or_else(|| item.column_name.clone());
3612        }
3613    }
3614    Ok(scope)
3615}
3616
3617// `table_func_names` is a mapping from a UUID to the original function
3618// name. The UUIDs are identifiers that have been rewritten from some table
3619// function expression, and this mapping restores the original names.
3620fn invent_column_name(
3621    ecx: &ExprContext,
3622    expr: &Expr<Aug>,
3623    table_func_names: &BTreeMap<String, Ident>,
3624) -> Result<Option<ColumnName>, PlanError> {
3625    // We follow PostgreSQL exactly here, which has some complicated rules
3626    // around "high" and "low" quality names. Low quality names override other
3627    // low quality names but not high quality names.
3628    //
3629    // See: https://github.com/postgres/postgres/blob/1f655fdc3/src/backend/parser/parse_target.c#L1716-L1728
3630
3631    #[derive(Debug)]
3632    enum NameQuality {
3633        Low,
3634        High,
3635    }
3636
3637    fn invent(
3638        ecx: &ExprContext,
3639        expr: &Expr<Aug>,
3640        table_func_names: &BTreeMap<String, Ident>,
3641    ) -> Result<Option<(ColumnName, NameQuality)>, PlanError> {
3642        Ok(match expr {
3643            Expr::Identifier(names) => {
3644                if let [name] = names.as_slice() {
3645                    if let Some(table_func_name) = table_func_names.get(name.as_str()) {
3646                        return Ok(Some((
3647                            normalize::column_name(table_func_name.clone()),
3648                            NameQuality::High,
3649                        )));
3650                    }
3651                }
3652                names
3653                    .last()
3654                    .map(|n| (normalize::column_name(n.clone()), NameQuality::High))
3655            }
3656            Expr::Value(v) => match v {
3657                // Per PostgreSQL, `bool` and `interval` literals take on the name
3658                // of their type, but not other literal types.
3659                Value::Boolean(_) => Some(("bool".into(), NameQuality::High)),
3660                Value::Interval(_) => Some(("interval".into(), NameQuality::High)),
3661                _ => None,
3662            },
3663            Expr::Function(func) => {
3664                let (schema, item) = match &func.name {
3665                    ResolvedItemName::Item {
3666                        qualifiers,
3667                        full_name,
3668                        ..
3669                    } => (&qualifiers.schema_spec, full_name.item.clone()),
3670                    // Name resolution should have rejected anything other than
3671                    // `Item` for a function call.
3672                    _ => {
3673                        bail_internal!("function name did not resolve to an item: {:?}", func.name)
3674                    }
3675                };
3676
3677                if schema == &SchemaSpecifier::from(ecx.qcx.scx.catalog.get_mz_internal_schema_id())
3678                    || schema
3679                        == &SchemaSpecifier::from(ecx.qcx.scx.catalog.get_mz_unsafe_schema_id())
3680                {
3681                    None
3682                } else {
3683                    Some((item.into(), NameQuality::High))
3684                }
3685            }
3686            Expr::HomogenizingFunction { function, .. } => Some((
3687                function.to_string().to_lowercase().into(),
3688                NameQuality::High,
3689            )),
3690            Expr::NullIf { .. } => Some(("nullif".into(), NameQuality::High)),
3691            Expr::Array { .. } => Some(("array".into(), NameQuality::High)),
3692            Expr::List { .. } => Some(("list".into(), NameQuality::High)),
3693            Expr::Map { .. } | Expr::MapSubquery(_) => Some(("map".into(), NameQuality::High)),
3694            Expr::Cast {
3695                expr, data_type, ..
3696            } => match invent(ecx, expr, table_func_names)? {
3697                Some((name, NameQuality::High)) => Some((name, NameQuality::High)),
3698                _ => Some((data_type.unqualified_item_name().into(), NameQuality::Low)),
3699            },
3700            Expr::Case { else_result, .. } => {
3701                let inner = match else_result.as_ref() {
3702                    Some(else_result) => invent(ecx, else_result, table_func_names)?,
3703                    None => None,
3704                };
3705                match inner {
3706                    Some((name, NameQuality::High)) => Some((name, NameQuality::High)),
3707                    _ => Some(("case".into(), NameQuality::Low)),
3708                }
3709            }
3710            Expr::FieldAccess { field, .. } => {
3711                Some((normalize::column_name(field.clone()), NameQuality::High))
3712            }
3713            Expr::Exists { .. } => Some(("exists".into(), NameQuality::High)),
3714            Expr::Subscript { expr, .. } => invent(ecx, expr, table_func_names)?,
3715            Expr::Subquery(query) | Expr::ListSubquery(query) | Expr::ArraySubquery(query) => {
3716                // A bit silly to have to plan the query here just to get its column
3717                // name, since we throw away the planned expression, but fixing this
3718                // requires a separate semantic analysis phase.
3719                //
3720                // We deliberately swallow planning errors here: if the subquery
3721                // doesn't plan, we just don't invent a name for it; the real
3722                // planning attempt elsewhere will surface the error.
3723                let Ok((_expr, scope)) = plan_nested_query(&mut ecx.derived_query_context(), query)
3724                else {
3725                    return Ok(None);
3726                };
3727                scope
3728                    .items
3729                    .first()
3730                    .map(|name| (name.column_name.clone(), NameQuality::High))
3731            }
3732            Expr::Row { .. } => Some(("row".into(), NameQuality::High)),
3733            _ => None,
3734        })
3735    }
3736
3737    Ok(invent(ecx, expr, table_func_names)?.map(|(name, _quality)| name))
3738}
3739
3740#[derive(Debug)]
3741enum ExpandedSelectItem<'a> {
3742    InputOrdinal(usize),
3743    Expr(Cow<'a, Expr<Aug>>),
3744}
3745
3746impl ExpandedSelectItem<'_> {
3747    fn as_expr(&self) -> Option<&Expr<Aug>> {
3748        match self {
3749            ExpandedSelectItem::InputOrdinal(_) => None,
3750            ExpandedSelectItem::Expr(expr) => Some(expr),
3751        }
3752    }
3753}
3754
3755fn expand_select_item<'a>(
3756    ecx: &ExprContext,
3757    s: &'a SelectItem<Aug>,
3758    table_func_names: &BTreeMap<String, Ident>,
3759) -> Result<Vec<(ExpandedSelectItem<'a>, ColumnName)>, PlanError> {
3760    match s {
3761        SelectItem::Expr {
3762            expr: Expr::QualifiedWildcard(table_name),
3763            alias: _,
3764        } => {
3765            *ecx.qcx.scx.ambiguous_columns.borrow_mut() = true;
3766            let table_name =
3767                normalize::unresolved_item_name(UnresolvedItemName(table_name.clone()))?;
3768            let out: Vec<_> = ecx
3769                .scope
3770                .items
3771                .iter()
3772                .enumerate()
3773                .filter(|(_i, item)| item.is_from_table(&table_name))
3774                .map(|(i, item)| {
3775                    let name = item.column_name.clone();
3776                    (ExpandedSelectItem::InputOrdinal(i), name)
3777                })
3778                .collect();
3779            if out.is_empty() {
3780                sql_bail!("no table named '{}' in scope", table_name);
3781            }
3782            Ok(out)
3783        }
3784        SelectItem::Expr {
3785            expr: Expr::WildcardAccess(sql_expr),
3786            alias: _,
3787        } => {
3788            *ecx.qcx.scx.ambiguous_columns.borrow_mut() = true;
3789            // A bit silly to have to plan the expression here just to get its
3790            // type, since we throw away the planned expression, but fixing this
3791            // requires a separate semantic analysis phase. Luckily this is an
3792            // uncommon operation and the PostgreSQL docs have a warning that
3793            // this operation is slow in Postgres too.
3794            let expr = plan_expr(ecx, sql_expr)?.type_as_any(ecx)?;
3795            let fields = match ecx.scalar_type(&expr) {
3796                SqlScalarType::Record { fields, .. } => fields,
3797                ty => sql_bail!(
3798                    "type {} is not composite",
3799                    ecx.humanize_sql_scalar_type(&ty, false)
3800                ),
3801            };
3802            let mut skip_cols: BTreeSet<ColumnName> = BTreeSet::new();
3803            if let Expr::Identifier(ident) = sql_expr.as_ref() {
3804                if let [name] = ident.as_slice() {
3805                    if let Ok(items) = ecx.scope.items_from_table(
3806                        &[],
3807                        &PartialItemName {
3808                            database: None,
3809                            schema: None,
3810                            item: name.as_str().to_string(),
3811                        },
3812                    ) {
3813                        for (_, item) in items {
3814                            if item
3815                                .is_exists_column_for_a_table_function_that_was_in_the_target_list
3816                            {
3817                                skip_cols.insert(item.column_name.clone());
3818                            }
3819                        }
3820                    }
3821                }
3822            }
3823            let items = fields
3824                .iter()
3825                .filter_map(|(name, _ty)| {
3826                    if skip_cols.contains(name) {
3827                        None
3828                    } else {
3829                        let item = ExpandedSelectItem::Expr(Cow::Owned(Expr::FieldAccess {
3830                            expr: sql_expr.clone(),
3831                            field: normalize::column_name_ident(name),
3832                        }));
3833                        Some((item, name.clone()))
3834                    }
3835                })
3836                .collect();
3837            Ok(items)
3838        }
3839        SelectItem::Wildcard => {
3840            *ecx.qcx.scx.ambiguous_columns.borrow_mut() = true;
3841            let items: Vec<_> = ecx
3842                .scope
3843                .items
3844                .iter()
3845                .enumerate()
3846                .filter(|(_i, item)| item.allow_unqualified_references)
3847                .map(|(i, item)| {
3848                    let name = item.column_name.clone();
3849                    (ExpandedSelectItem::InputOrdinal(i), name)
3850                })
3851                .collect();
3852
3853            Ok(items)
3854        }
3855        SelectItem::Expr { expr, alias } => {
3856            let name = match alias.clone().map(normalize::column_name) {
3857                Some(name) => name,
3858                None => invent_column_name(ecx, expr, table_func_names)?
3859                    .unwrap_or_else(|| UNKNOWN_COLUMN_NAME.into()),
3860            };
3861            Ok(vec![(ExpandedSelectItem::Expr(Cow::Borrowed(expr)), name)])
3862        }
3863    }
3864}
3865
3866fn plan_join(
3867    left_qcx: &QueryContext,
3868    left: HirRelationExpr,
3869    left_scope: Scope,
3870    join: &Join<Aug>,
3871) -> Result<(HirRelationExpr, Scope), PlanError> {
3872    const ON_TRUE: JoinConstraint<Aug> = JoinConstraint::On(Expr::Value(Value::Boolean(true)));
3873    let (kind, constraint) = match &join.join_operator {
3874        JoinOperator::CrossJoin => (JoinKind::Inner, &ON_TRUE),
3875        JoinOperator::Inner(constraint) => (JoinKind::Inner, constraint),
3876        JoinOperator::LeftOuter(constraint) => (JoinKind::LeftOuter, constraint),
3877        JoinOperator::RightOuter(constraint) => (JoinKind::RightOuter, constraint),
3878        JoinOperator::FullOuter(constraint) => (JoinKind::FullOuter, constraint),
3879    };
3880
3881    let mut right_qcx = left_qcx.derived_context(left_scope.clone(), left_qcx.relation_type(&left));
3882    if !kind.can_be_correlated() {
3883        for item in &mut right_qcx.outer_scopes[0].items {
3884            // Per PostgreSQL (and apparently SQL:2008), we can't simply remove
3885            // these items from scope. These items need to *exist* because they
3886            // might shadow variables in outer scopes that would otherwise be
3887            // valid to reference, but accessing them needs to produce an error.
3888            item.error_if_referenced =
3889                Some(|table, column| PlanError::WrongJoinTypeForLateralColumn {
3890                    table: table.cloned(),
3891                    column: column.clone(),
3892                });
3893        }
3894    }
3895    let (right, right_scope) = plan_table_factor(&right_qcx, &join.relation)?;
3896
3897    let (expr, scope) = match constraint {
3898        JoinConstraint::On(expr) => {
3899            let product_scope = left_scope.product(right_scope)?;
3900            let ecx = &ExprContext {
3901                qcx: left_qcx,
3902                name: "ON clause",
3903                scope: &product_scope,
3904                relation_type: &SqlRelationType::new(
3905                    left_qcx
3906                        .relation_type(&left)
3907                        .column_types
3908                        .into_iter()
3909                        .chain(right_qcx.relation_type(&right).column_types)
3910                        .collect(),
3911                ),
3912                allow_aggregates: false,
3913                allow_subqueries: true,
3914                allow_parameters: true,
3915                allow_windows: false,
3916            };
3917            let on = plan_expr(ecx, expr)?.type_as(ecx, &SqlScalarType::Bool)?;
3918            let joined = left.join(right, on, kind);
3919            (joined, product_scope)
3920        }
3921        JoinConstraint::Using { columns, alias } => {
3922            let column_names = columns
3923                .iter()
3924                .map(|ident| normalize::column_name(ident.clone()))
3925                .collect::<Vec<_>>();
3926
3927            plan_using_constraint(
3928                &column_names,
3929                left_qcx,
3930                left,
3931                left_scope,
3932                &right_qcx,
3933                right,
3934                right_scope,
3935                kind,
3936                alias.as_ref(),
3937            )?
3938        }
3939        JoinConstraint::Natural => {
3940            // We shouldn't need to set ambiguous_columns on both the right and left qcx since they
3941            // have the same scx. However, it doesn't hurt to be safe.
3942            *left_qcx.scx.ambiguous_columns.borrow_mut() = true;
3943            *right_qcx.scx.ambiguous_columns.borrow_mut() = true;
3944            let left_column_names = left_scope.column_names();
3945            let right_column_names: BTreeSet<_> = right_scope.column_names().collect();
3946            let column_names: Vec<_> = left_column_names
3947                .filter(|col| right_column_names.contains(col))
3948                .cloned()
3949                .collect();
3950            plan_using_constraint(
3951                &column_names,
3952                left_qcx,
3953                left,
3954                left_scope,
3955                &right_qcx,
3956                right,
3957                right_scope,
3958                kind,
3959                None,
3960            )?
3961        }
3962    };
3963    Ok((expr, scope))
3964}
3965
3966// See page 440 of ANSI SQL 2016 spec for details on scoping of using/natural joins
3967#[allow(clippy::too_many_arguments)]
3968fn plan_using_constraint(
3969    column_names: &[ColumnName],
3970    left_qcx: &QueryContext,
3971    left: HirRelationExpr,
3972    left_scope: Scope,
3973    right_qcx: &QueryContext,
3974    right: HirRelationExpr,
3975    right_scope: Scope,
3976    kind: JoinKind,
3977    alias: Option<&Ident>,
3978) -> Result<(HirRelationExpr, Scope), PlanError> {
3979    let mut both_scope = left_scope.clone().product(right_scope.clone())?;
3980
3981    // Cargo culting PG here; no discernable reason this must fail, but PG does
3982    // so we do, as well.
3983    let mut unique_column_names = BTreeSet::new();
3984    for c in column_names {
3985        if !unique_column_names.insert(c) {
3986            return Err(PlanError::Unsupported {
3987                feature: format!(
3988                    "column name {} appears more than once in USING clause",
3989                    c.quoted()
3990                ),
3991                discussion_no: None,
3992            });
3993        }
3994    }
3995
3996    let alias_item_name = alias.map(|alias| PartialItemName {
3997        database: None,
3998        schema: None,
3999        item: alias.clone().to_string(),
4000    });
4001
4002    if let Some(alias_item_name) = &alias_item_name {
4003        for partial_item_name in both_scope.table_names() {
4004            if partial_item_name.matches(alias_item_name) {
4005                sql_bail!(
4006                    "table name \"{}\" specified more than once",
4007                    alias_item_name
4008                )
4009            }
4010        }
4011    }
4012
4013    let ecx = &ExprContext {
4014        qcx: right_qcx,
4015        name: "USING clause",
4016        scope: &both_scope,
4017        relation_type: &SqlRelationType::new(
4018            left_qcx
4019                .relation_type(&left)
4020                .column_types
4021                .into_iter()
4022                .chain(right_qcx.relation_type(&right).column_types)
4023                .collect(),
4024        ),
4025        allow_aggregates: false,
4026        allow_subqueries: false,
4027        allow_parameters: false,
4028        allow_windows: false,
4029    };
4030
4031    let mut join_exprs = vec![];
4032    let mut map_exprs = vec![];
4033    let mut new_items = vec![];
4034    let mut join_cols = vec![];
4035    let mut hidden_cols = vec![];
4036
4037    for column_name in column_names {
4038        // the two sides will have different names (e.g., `t1.a` and `t2.a`)
4039        let (lhs, lhs_name) = left_scope.resolve_using_column(
4040            column_name,
4041            JoinSide::Left,
4042            &mut left_qcx.name_manager.borrow_mut(),
4043        )?;
4044        let (mut rhs, rhs_name) = right_scope.resolve_using_column(
4045            column_name,
4046            JoinSide::Right,
4047            &mut right_qcx.name_manager.borrow_mut(),
4048        )?;
4049
4050        // Adjust the RHS reference to its post-join location.
4051        rhs.column += left_scope.len();
4052
4053        // Join keys must be resolved to same type.
4054        let mut exprs = coerce_homogeneous_exprs(
4055            &ecx.with_name(&format!(
4056                "NATURAL/USING join column {}",
4057                column_name.quoted()
4058            )),
4059            vec![
4060                CoercibleScalarExpr::Coerced(HirScalarExpr::named_column(
4061                    lhs,
4062                    Arc::clone(&lhs_name),
4063                )),
4064                CoercibleScalarExpr::Coerced(HirScalarExpr::named_column(
4065                    rhs,
4066                    Arc::clone(&rhs_name),
4067                )),
4068            ],
4069            None,
4070        )?;
4071        let (expr1, expr2) = (exprs.remove(0), exprs.remove(0));
4072
4073        match kind {
4074            JoinKind::LeftOuter { .. } | JoinKind::Inner { .. } => {
4075                join_cols.push(lhs.column);
4076                hidden_cols.push(rhs.column);
4077            }
4078            JoinKind::RightOuter => {
4079                join_cols.push(rhs.column);
4080                hidden_cols.push(lhs.column);
4081            }
4082            JoinKind::FullOuter => {
4083                // Create a new column that will be the coalesced value of left
4084                // and right.
4085                join_cols.push(both_scope.items.len() + map_exprs.len());
4086                hidden_cols.push(lhs.column);
4087                hidden_cols.push(rhs.column);
4088                map_exprs.push(HirScalarExpr::call_variadic(
4089                    Coalesce,
4090                    vec![expr1.clone(), expr2.clone()],
4091                ));
4092                new_items.push(ScopeItem::from_column_name(column_name));
4093            }
4094        }
4095
4096        // If a `join_using_alias` is present, add a new scope item that accepts
4097        // only table-qualified references for each specified join column.
4098        // Unlike regular table aliases, a `join_using_alias` should not hide the
4099        // names of the joined relations.
4100        if alias_item_name.is_some() {
4101            let new_item_col = both_scope.items.len() + new_items.len();
4102            join_cols.push(new_item_col);
4103            hidden_cols.push(new_item_col);
4104
4105            new_items.push(ScopeItem::from_name(
4106                alias_item_name.clone(),
4107                column_name.clone().to_string(),
4108            ));
4109
4110            // The aliased column `alias.col` must take the same value as the
4111            // unqualified join output column `col`. For INNER and LEFT joins
4112            // that's the LHS value, for RIGHT joins it's the RHS value, and
4113            // for FULL OUTER joins it's COALESCE(lhs, rhs). Using `lhs`
4114            // unconditionally produces wrong results for RIGHT/FULL joins on
4115            // rows where the LHS side is NULL.
4116            let alias_expr = match kind {
4117                JoinKind::LeftOuter { .. } | JoinKind::Inner { .. } => {
4118                    HirScalarExpr::named_column(lhs, Arc::clone(&lhs_name))
4119                }
4120                JoinKind::RightOuter => HirScalarExpr::named_column(rhs, Arc::clone(&rhs_name)),
4121                JoinKind::FullOuter => {
4122                    HirScalarExpr::call_variadic(Coalesce, vec![expr1.clone(), expr2.clone()])
4123                }
4124            };
4125            map_exprs.push(alias_expr);
4126        }
4127
4128        join_exprs.push(expr1.call_binary(expr2, expr_func::Eq));
4129    }
4130    both_scope.items.extend(new_items);
4131
4132    // The columns from the secondary side of the join remain accessible by
4133    // their table-qualified name, but not by their column name alone. They are
4134    // also excluded from `SELECT *`.
4135    for c in hidden_cols {
4136        both_scope.items[c].allow_unqualified_references = false;
4137    }
4138
4139    // Reproject all returned elements to the front of the list.
4140    let project_key = join_cols
4141        .into_iter()
4142        .chain(0..both_scope.items.len())
4143        .unique()
4144        .collect::<Vec<_>>();
4145
4146    both_scope = both_scope.project(&project_key);
4147
4148    let on = HirScalarExpr::variadic_and(join_exprs);
4149
4150    let both = left
4151        .join(right, on, kind)
4152        .map(map_exprs)
4153        .project(project_key);
4154    Ok((both, both_scope))
4155}
4156
4157pub fn plan_expr<'a>(
4158    ecx: &'a ExprContext,
4159    e: &Expr<Aug>,
4160) -> Result<CoercibleScalarExpr, PlanError> {
4161    ecx.checked_recur(|ecx| plan_expr_inner(ecx, e))
4162}
4163
4164fn plan_expr_inner<'a>(
4165    ecx: &'a ExprContext,
4166    e: &Expr<Aug>,
4167) -> Result<CoercibleScalarExpr, PlanError> {
4168    if let Some((i, item)) = ecx.scope.resolve_expr(e) {
4169        // We've already calculated this expression.
4170        return Ok(HirScalarExpr::named_column(
4171            i,
4172            ecx.qcx.name_manager.borrow_mut().intern_scope_item(item),
4173        )
4174        .into());
4175    }
4176
4177    match e {
4178        // Names.
4179        Expr::Identifier(names) | Expr::QualifiedWildcard(names) => {
4180            Ok(plan_identifier(ecx, names)?.into())
4181        }
4182
4183        // Literals.
4184        Expr::Value(val) => plan_literal(val),
4185        Expr::Parameter(n) => plan_parameter(ecx, *n),
4186        Expr::Array(exprs) => plan_array(ecx, exprs, None),
4187        Expr::List(exprs) => plan_list(ecx, exprs, None),
4188        Expr::Map(exprs) => plan_map(ecx, exprs, None),
4189        Expr::Row { exprs } => plan_row(ecx, exprs),
4190
4191        // Generalized functions, operators, and casts.
4192        Expr::Op { op, expr1, expr2 } => {
4193            Ok(plan_op(ecx, normalize::op(op)?, expr1, expr2.as_deref())?.into())
4194        }
4195        Expr::Cast {
4196            expr,
4197            data_type,
4198            failure_mode,
4199        } => plan_cast(ecx, expr, data_type, *failure_mode),
4200        Expr::Function(func) => Ok(plan_function(ecx, func)?.into()),
4201
4202        // Special functions and operators.
4203        Expr::Not { expr } => plan_not(ecx, expr),
4204        Expr::And { left, right } => plan_and(ecx, left, right),
4205        Expr::Or { left, right } => plan_or(ecx, left, right),
4206        Expr::IsExpr {
4207            expr,
4208            construct,
4209            negated,
4210        } => Ok(plan_is_expr(ecx, expr, construct, *negated)?.into()),
4211        Expr::Case {
4212            operand,
4213            conditions,
4214            results,
4215            else_result,
4216        } => Ok(plan_case(ecx, operand, conditions, results, else_result)?.into()),
4217        Expr::HomogenizingFunction { function, exprs } => {
4218            plan_homogenizing_function(ecx, function, exprs)
4219        }
4220        Expr::NullIf { l_expr, r_expr } => Ok(plan_case(
4221            ecx,
4222            &None,
4223            &[l_expr.clone().equals(*r_expr.clone())],
4224            &[Expr::null()],
4225            &Some(Box::new(*l_expr.clone())),
4226        )?
4227        .into()),
4228        Expr::FieldAccess { expr, field } => plan_field_access(ecx, expr, field),
4229        Expr::WildcardAccess(expr) => plan_expr(ecx, expr),
4230        Expr::Subscript { expr, positions } => plan_subscript(ecx, expr, positions),
4231        Expr::Like {
4232            expr,
4233            pattern,
4234            escape,
4235            case_insensitive,
4236            negated,
4237        } => Ok(plan_like(
4238            ecx,
4239            expr,
4240            pattern,
4241            escape.as_deref(),
4242            *case_insensitive,
4243            *negated,
4244        )?
4245        .into()),
4246
4247        Expr::InList {
4248            expr,
4249            list,
4250            negated,
4251        } => plan_in_list(ecx, expr, list, negated),
4252
4253        // Subqueries.
4254        Expr::Exists(query) => plan_exists(ecx, query),
4255        Expr::Subquery(query) => plan_subquery(ecx, query),
4256        Expr::ListSubquery(query) => plan_list_subquery(ecx, query),
4257        Expr::MapSubquery(query) => plan_map_subquery(ecx, query),
4258        Expr::ArraySubquery(query) => plan_array_subquery(ecx, query),
4259        Expr::Collate { expr, collation } => plan_collate(ecx, expr, collation),
4260        Expr::Nested(_) => bail_internal!("Expr::Nested should have been desugared"),
4261        Expr::InSubquery { .. } => {
4262            bail_internal!("Expr::InSubquery should have been desugared")
4263        }
4264        Expr::AnyExpr { .. } => {
4265            bail_internal!("Expr::AnyExpr should have been desugared")
4266        }
4267        Expr::AllExpr { .. } => {
4268            bail_internal!("Expr::AllExpr should have been desugared")
4269        }
4270        Expr::AnySubquery { .. } => {
4271            bail_internal!("Expr::AnySubquery should have been desugared")
4272        }
4273        Expr::AllSubquery { .. } => {
4274            bail_internal!("Expr::AllSubquery should have been desugared")
4275        }
4276        Expr::Between { .. } => {
4277            bail_internal!("Expr::Between should have been desugared")
4278        }
4279    }
4280}
4281
4282fn plan_parameter(ecx: &ExprContext, n: usize) -> Result<CoercibleScalarExpr, PlanError> {
4283    if !ecx.allow_parameters {
4284        // It might be clearer to return an error like "cannot use parameter
4285        // here", but this is how PostgreSQL does it, and so for now we follow
4286        // PostgreSQL.
4287        return Err(PlanError::UnknownParameter(n));
4288    }
4289    if n == 0 || n > 65536 {
4290        return Err(PlanError::UnknownParameter(n));
4291    }
4292    if ecx.param_types().borrow().contains_key(&n) {
4293        Ok(HirScalarExpr::parameter(n).into())
4294    } else {
4295        Ok(CoercibleScalarExpr::Parameter(n))
4296    }
4297}
4298
4299fn plan_row(ecx: &ExprContext, exprs: &[Expr<Aug>]) -> Result<CoercibleScalarExpr, PlanError> {
4300    let mut out = vec![];
4301    for e in exprs {
4302        out.push(plan_expr(ecx, e)?);
4303    }
4304    Ok(CoercibleScalarExpr::LiteralRecord(out))
4305}
4306
4307/// Plans `CAST(expr AS data_type)` or, under
4308/// [`CastFailureMode::NullFallback`], `TRY_CAST(expr AS data_type)`.
4309fn plan_cast(
4310    ecx: &ExprContext,
4311    expr: &Expr<Aug>,
4312    data_type: &ResolvedDataType,
4313    failure_mode: CastFailureMode,
4314) -> Result<CoercibleScalarExpr, PlanError> {
4315    let name = match failure_mode {
4316        CastFailureMode::Error => "CAST",
4317        CastFailureMode::NullFallback => {
4318            ecx.qcx.scx.require_feature_flag(&vars::ENABLE_TRY_CAST)?;
4319            "TRY_CAST"
4320        }
4321    };
4322    let to_scalar_type = scalar_type_from_sql(ecx.qcx.scx, data_type)?;
4323    let keeps_hint = |empty: bool| failure_mode == CastFailureMode::Error || empty;
4324    let expr = match expr {
4325        // Special case a direct cast of an ARRAY, LIST, or MAP expression so
4326        // we can pass in the target type as a type hint. This is
4327        // a limited form of the coercion that we do for string literals
4328        // via CoercibleScalarExpr. We used to let CoercibleScalarExpr
4329        // handle ARRAY/LIST/MAP coercion too, but doing so causes
4330        // PostgreSQL compatibility trouble.
4331        //
4332        // See: https://github.com/postgres/postgres/blob/31f403e95/src/backend/parser/parse_expr.c#L2762-L2768
4333        //
4334        // The hint makes the element casts part of the literal, where a
4335        // failure mode cannot reach them. Under `NullFallback` a non-empty
4336        // ARRAY or LIST literal is planned on its own instead, so the cast
4337        // that gets wrapped is the whole-value one, and a bad element yields
4338        // NULL for the whole value just as it does for a non-literal source.
4339        // An empty literal has no element casts, so it keeps the hint, which
4340        // is the only thing that can type it. A non-empty MAP literal has no
4341        // whole-value cast to fall back on, since there is no cast between
4342        // map types, so under `NullFallback` it is rejected rather than
4343        // converted with the hint's element-level strictness.
4344        Expr::Array(exprs) if keeps_hint(exprs.is_empty()) => {
4345            plan_array(ecx, exprs, Some(&to_scalar_type))?
4346        }
4347        Expr::List(exprs) if keeps_hint(exprs.is_empty()) => {
4348            plan_list(ecx, exprs, Some(&to_scalar_type))?
4349        }
4350        Expr::Map(exprs) if keeps_hint(exprs.is_empty()) => {
4351            plan_map(ecx, exprs, Some(&to_scalar_type))?
4352        }
4353        Expr::Map(_) => {
4354            bail_unsupported!("non-empty MAP constructor as the argument of TRY_CAST")
4355        }
4356        expr => plan_expr(ecx, expr)?,
4357    };
4358    let ecx = &ecx.with_name(name);
4359    // A string literal reaches its cast through coercion, so the failure mode
4360    // has to apply there too, or `TRY_CAST('abc' AS int4)` would still error.
4361    let expr = typeconv::plan_coerce_with_failure_mode(ecx, expr, &to_scalar_type, failure_mode)?;
4362    let expr = typeconv::plan_cast_with_failure_mode(
4363        ecx,
4364        CastContext::Explicit,
4365        failure_mode,
4366        expr,
4367        &to_scalar_type,
4368    )?;
4369    Ok(expr.into())
4370}
4371
4372fn plan_not(ecx: &ExprContext, expr: &Expr<Aug>) -> Result<CoercibleScalarExpr, PlanError> {
4373    let ecx = ecx.with_name("NOT argument");
4374    Ok(plan_expr(&ecx, expr)?
4375        .type_as(&ecx, &SqlScalarType::Bool)?
4376        .call_unary(UnaryFunc::Not(expr_func::Not))
4377        .into())
4378}
4379
4380fn plan_and(
4381    ecx: &ExprContext,
4382    left: &Expr<Aug>,
4383    right: &Expr<Aug>,
4384) -> Result<CoercibleScalarExpr, PlanError> {
4385    let ecx = ecx.with_name("AND argument");
4386    Ok(HirScalarExpr::variadic_and(vec![
4387        plan_expr(&ecx, left)?.type_as(&ecx, &SqlScalarType::Bool)?,
4388        plan_expr(&ecx, right)?.type_as(&ecx, &SqlScalarType::Bool)?,
4389    ])
4390    .into())
4391}
4392
4393fn plan_or(
4394    ecx: &ExprContext,
4395    left: &Expr<Aug>,
4396    right: &Expr<Aug>,
4397) -> Result<CoercibleScalarExpr, PlanError> {
4398    let ecx = ecx.with_name("OR argument");
4399    Ok(HirScalarExpr::variadic_or(vec![
4400        plan_expr(&ecx, left)?.type_as(&ecx, &SqlScalarType::Bool)?,
4401        plan_expr(&ecx, right)?.type_as(&ecx, &SqlScalarType::Bool)?,
4402    ])
4403    .into())
4404}
4405
4406fn plan_in_list(
4407    ecx: &ExprContext,
4408    lhs: &Expr<Aug>,
4409    list: &Vec<Expr<Aug>>,
4410    negated: &bool,
4411) -> Result<CoercibleScalarExpr, PlanError> {
4412    let ecx = ecx.with_name("IN list");
4413    let or = HirScalarExpr::variadic_or(
4414        list.into_iter()
4415            .map(|e| {
4416                let eq = lhs.clone().equals(e.clone());
4417                plan_expr(&ecx, &eq)?.type_as(&ecx, &SqlScalarType::Bool)
4418            })
4419            .collect::<Result<Vec<HirScalarExpr>, PlanError>>()?,
4420    );
4421    Ok(if *negated {
4422        or.call_unary(UnaryFunc::Not(expr_func::Not))
4423    } else {
4424        or
4425    }
4426    .into())
4427}
4428
4429fn plan_homogenizing_function(
4430    ecx: &ExprContext,
4431    function: &HomogenizingFunction,
4432    exprs: &[Expr<Aug>],
4433) -> Result<CoercibleScalarExpr, PlanError> {
4434    assert!(!exprs.is_empty()); // `COALESCE()` is a syntax error
4435    let expr = HirScalarExpr::call_variadic(
4436        match function {
4437            HomogenizingFunction::Coalesce => VariadicFunc::from(Coalesce),
4438            HomogenizingFunction::Greatest => VariadicFunc::from(Greatest),
4439            HomogenizingFunction::Least => VariadicFunc::from(Least),
4440        },
4441        coerce_homogeneous_exprs(
4442            &ecx.with_name(&function.to_string().to_lowercase()),
4443            plan_exprs(ecx, exprs)?,
4444            None,
4445        )?,
4446    );
4447    Ok(expr.into())
4448}
4449
4450fn plan_field_access(
4451    ecx: &ExprContext,
4452    expr: &Expr<Aug>,
4453    field: &Ident,
4454) -> Result<CoercibleScalarExpr, PlanError> {
4455    let field = normalize::column_name(field.clone());
4456    let expr = plan_expr(ecx, expr)?.type_as_any(ecx)?;
4457    let ty = ecx.scalar_type(&expr);
4458    let i = match &ty {
4459        SqlScalarType::Record { fields, .. } => {
4460            fields.iter().position(|(name, _ty)| *name == field)
4461        }
4462        ty => sql_bail!(
4463            "column notation applied to type {}, which is not a composite type",
4464            ecx.humanize_sql_scalar_type(ty, false)
4465        ),
4466    };
4467    match i {
4468        None => sql_bail!(
4469            "field {} not found in data type {}",
4470            field,
4471            ecx.humanize_sql_scalar_type(&ty, false)
4472        ),
4473        Some(i) => Ok(expr
4474            .call_unary(UnaryFunc::RecordGet(expr_func::RecordGet(i)))
4475            .into()),
4476    }
4477}
4478
4479fn plan_subscript(
4480    ecx: &ExprContext,
4481    expr: &Expr<Aug>,
4482    positions: &[SubscriptPosition<Aug>],
4483) -> Result<CoercibleScalarExpr, PlanError> {
4484    assert!(
4485        !positions.is_empty(),
4486        "subscript expression must contain at least one position"
4487    );
4488
4489    let ecx = &ecx.with_name("subscripting");
4490    let expr = plan_expr(ecx, expr)?.type_as_any(ecx)?;
4491    let ty = ecx.scalar_type(&expr);
4492    match &ty {
4493        SqlScalarType::Array(..) | SqlScalarType::Int2Vector => plan_subscript_array(
4494            ecx,
4495            expr,
4496            positions,
4497            // Int2Vector uses 0-based indexing, while arrays use 1-based indexing, so we need to
4498            // adjust all Int2Vector subscript operations by 1 (both w/r/t input and the values we
4499            // track in its backing data).
4500            if ty == SqlScalarType::Int2Vector {
4501                1
4502            } else {
4503                0
4504            },
4505        ),
4506        SqlScalarType::Jsonb => plan_subscript_jsonb(ecx, expr, positions),
4507        SqlScalarType::List { element_type, .. } => {
4508            // `elem_type_name` is used only in error msgs, so we set `postgres_compat` to false.
4509            let elem_type_name = ecx.humanize_sql_scalar_type(element_type, false);
4510            let n_layers = ty.unwrap_list_n_layers();
4511            plan_subscript_list(ecx, expr, positions, n_layers, &elem_type_name)
4512        }
4513        ty => sql_bail!(
4514            "cannot subscript type {}",
4515            ecx.humanize_sql_scalar_type(ty, false)
4516        ),
4517    }
4518}
4519
4520// All subscript positions are of the form [<expr>(:<expr>?)?]; extract all
4521// expressions from those that look like indexes (i.e. `[<expr>]`) or error if
4522// any were slices (i.e. included colon).
4523fn extract_scalar_subscript_from_positions<'a>(
4524    positions: &'a [SubscriptPosition<Aug>],
4525    expr_type_name: &str,
4526) -> Result<Vec<&'a Expr<Aug>>, PlanError> {
4527    let mut scalar_subscripts = Vec::with_capacity(positions.len());
4528    for p in positions {
4529        if p.explicit_slice {
4530            sql_bail!("{} subscript does not support slices", expr_type_name);
4531        }
4532        assert!(
4533            p.end.is_none(),
4534            "index-appearing subscripts cannot have end value"
4535        );
4536        scalar_subscripts.push(p.start.as_ref().expect("has start if not slice"));
4537    }
4538    Ok(scalar_subscripts)
4539}
4540
4541fn plan_subscript_array(
4542    ecx: &ExprContext,
4543    expr: HirScalarExpr,
4544    positions: &[SubscriptPosition<Aug>],
4545    offset: i64,
4546) -> Result<CoercibleScalarExpr, PlanError> {
4547    let mut exprs = Vec::with_capacity(positions.len() + 1);
4548    exprs.push(expr);
4549
4550    // Subscripting arrays doesn't yet support slicing, so we always want to
4551    // extract scalars or error.
4552    let indexes = extract_scalar_subscript_from_positions(positions, "array")?;
4553
4554    for i in indexes {
4555        exprs.push(plan_expr(ecx, i)?.cast_to(
4556            ecx,
4557            CastContext::Explicit,
4558            &SqlScalarType::Int64,
4559        )?);
4560    }
4561
4562    Ok(HirScalarExpr::call_variadic(ArrayIndex { offset }, exprs).into())
4563}
4564
4565fn plan_subscript_list(
4566    ecx: &ExprContext,
4567    mut expr: HirScalarExpr,
4568    positions: &[SubscriptPosition<Aug>],
4569    mut remaining_layers: usize,
4570    elem_type_name: &str,
4571) -> Result<CoercibleScalarExpr, PlanError> {
4572    let mut i = 0;
4573
4574    while i < positions.len() {
4575        // Take all contiguous index operations, i.e. find next slice operation.
4576        let j = positions[i..]
4577            .iter()
4578            .position(|p| p.explicit_slice)
4579            .unwrap_or(positions.len() - i);
4580        if j != 0 {
4581            let indexes = extract_scalar_subscript_from_positions(&positions[i..i + j], "")?;
4582            let (n, e) = plan_index_list(
4583                ecx,
4584                expr,
4585                indexes.as_slice(),
4586                remaining_layers,
4587                elem_type_name,
4588            )?;
4589            remaining_layers = n;
4590            expr = e;
4591            i += j;
4592        }
4593
4594        // Take all contiguous slice operations, i.e. find next index operation.
4595        let j = positions[i..]
4596            .iter()
4597            .position(|p| !p.explicit_slice)
4598            .unwrap_or(positions.len() - i);
4599        if j != 0 {
4600            expr = plan_slice_list(
4601                ecx,
4602                expr,
4603                &positions[i..i + j],
4604                remaining_layers,
4605                elem_type_name,
4606            )?;
4607            i += j;
4608        }
4609    }
4610
4611    Ok(expr.into())
4612}
4613
4614fn plan_index_list(
4615    ecx: &ExprContext,
4616    expr: HirScalarExpr,
4617    indexes: &[&Expr<Aug>],
4618    n_layers: usize,
4619    elem_type_name: &str,
4620) -> Result<(usize, HirScalarExpr), PlanError> {
4621    let depth = indexes.len();
4622
4623    if depth > n_layers {
4624        if n_layers == 0 {
4625            sql_bail!("cannot subscript type {}", elem_type_name)
4626        } else {
4627            sql_bail!(
4628                "cannot index into {} layers; list only has {} layer{}",
4629                depth,
4630                n_layers,
4631                if n_layers == 1 { "" } else { "s" }
4632            )
4633        }
4634    }
4635
4636    let mut exprs = Vec::with_capacity(depth + 1);
4637    exprs.push(expr);
4638
4639    for i in indexes {
4640        exprs.push(plan_expr(ecx, i)?.cast_to(
4641            ecx,
4642            CastContext::Explicit,
4643            &SqlScalarType::Int64,
4644        )?);
4645    }
4646
4647    Ok((
4648        n_layers - depth,
4649        HirScalarExpr::call_variadic(ListIndex, exprs),
4650    ))
4651}
4652
4653fn plan_slice_list(
4654    ecx: &ExprContext,
4655    expr: HirScalarExpr,
4656    slices: &[SubscriptPosition<Aug>],
4657    n_layers: usize,
4658    elem_type_name: &str,
4659) -> Result<HirScalarExpr, PlanError> {
4660    if n_layers == 0 {
4661        sql_bail!("cannot subscript type {}", elem_type_name)
4662    }
4663
4664    // first arg will be list
4665    let mut exprs = Vec::with_capacity(slices.len() + 1);
4666    exprs.push(expr);
4667    // extract (start, end) parts from collected slices
4668    let extract_position_or_default = |position, default| -> Result<HirScalarExpr, PlanError> {
4669        Ok(match position {
4670            Some(p) => {
4671                plan_expr(ecx, p)?.cast_to(ecx, CastContext::Explicit, &SqlScalarType::Int64)?
4672            }
4673            None => HirScalarExpr::literal(Datum::Int64(default), SqlScalarType::Int64),
4674        })
4675    };
4676    for p in slices {
4677        let start = extract_position_or_default(p.start.as_ref(), 1)?;
4678        let end = extract_position_or_default(p.end.as_ref(), i64::MAX - 1)?;
4679        exprs.push(start);
4680        exprs.push(end);
4681    }
4682
4683    Ok(HirScalarExpr::call_variadic(ListSliceLinear, exprs))
4684}
4685
4686fn plan_like(
4687    ecx: &ExprContext,
4688    expr: &Expr<Aug>,
4689    pattern: &Expr<Aug>,
4690    escape: Option<&Expr<Aug>>,
4691    case_insensitive: bool,
4692    not: bool,
4693) -> Result<HirScalarExpr, PlanError> {
4694    use CastContext::Implicit;
4695    let ecx = ecx.with_name("LIKE argument");
4696    let expr = plan_expr(&ecx, expr)?;
4697    let haystack = match ecx.scalar_type(&expr) {
4698        CoercibleScalarType::Coerced(ref ty @ SqlScalarType::Char { length }) => expr
4699            .type_as(&ecx, ty)?
4700            .call_unary(UnaryFunc::PadChar(expr_func::PadChar { length })),
4701        _ => expr.cast_to(&ecx, Implicit, &SqlScalarType::String)?,
4702    };
4703    let mut pattern = plan_expr(&ecx, pattern)?.cast_to(&ecx, Implicit, &SqlScalarType::String)?;
4704    if let Some(escape) = escape {
4705        pattern = pattern.call_binary(
4706            plan_expr(&ecx, escape)?.cast_to(&ecx, Implicit, &SqlScalarType::String)?,
4707            expr_func::LikeEscape,
4708        );
4709    }
4710    let func: BinaryFunc = if case_insensitive {
4711        expr_func::IsLikeMatchCaseInsensitive.into()
4712    } else {
4713        expr_func::IsLikeMatchCaseSensitive.into()
4714    };
4715    let like = haystack.call_binary(pattern, func);
4716    if not {
4717        Ok(like.call_unary(UnaryFunc::Not(expr_func::Not)))
4718    } else {
4719        Ok(like)
4720    }
4721}
4722
4723fn plan_subscript_jsonb(
4724    ecx: &ExprContext,
4725    expr: HirScalarExpr,
4726    positions: &[SubscriptPosition<Aug>],
4727) -> Result<CoercibleScalarExpr, PlanError> {
4728    use CastContext::Implicit;
4729    use SqlScalarType::{Int64, String};
4730
4731    // JSONB doesn't support the slicing syntax, so simply error if you
4732    // encounter any explicit slices.
4733    let subscripts = extract_scalar_subscript_from_positions(positions, "jsonb")?;
4734
4735    let mut exprs = Vec::with_capacity(subscripts.len());
4736    for s in subscripts {
4737        let subscript = plan_expr(ecx, s)?;
4738        let subscript = if let Ok(subscript) = subscript.clone().cast_to(ecx, Implicit, &String) {
4739            subscript
4740        } else if let Ok(subscript) = subscript.cast_to(ecx, Implicit, &Int64) {
4741            // Integers are converted to a string here and then re-parsed as an
4742            // integer by `JsonbGetPath`. Weird, but this is how PostgreSQL says to
4743            // do it.
4744            typeconv::to_string(ecx, subscript)?
4745        } else {
4746            sql_bail!("jsonb subscript type must be coercible to integer or text");
4747        };
4748        exprs.push(subscript);
4749    }
4750
4751    // Subscripting works like `expr #> ARRAY[subscript]` rather than
4752    // `expr->subscript` as you might expect.
4753    let expr = expr.call_binary(
4754        HirScalarExpr::call_variadic(
4755            ArrayCreate {
4756                elem_type: SqlScalarType::String,
4757            },
4758            exprs,
4759        ),
4760        expr_func::JsonbGetPath,
4761    );
4762    Ok(expr.into())
4763}
4764
4765fn plan_exists(ecx: &ExprContext, query: &Query<Aug>) -> Result<CoercibleScalarExpr, PlanError> {
4766    if !ecx.allow_subqueries {
4767        sql_bail!("{} does not allow subqueries", ecx.name)
4768    }
4769    let mut qcx = ecx.derived_query_context();
4770    let (expr, _scope) = plan_nested_query(&mut qcx, query)?;
4771    Ok(expr.exists().into())
4772}
4773
4774fn plan_subquery(ecx: &ExprContext, query: &Query<Aug>) -> Result<CoercibleScalarExpr, PlanError> {
4775    if !ecx.allow_subqueries {
4776        sql_bail!("{} does not allow subqueries", ecx.name)
4777    }
4778    let mut qcx = ecx.derived_query_context();
4779    let (expr, _scope) = plan_nested_query(&mut qcx, query)?;
4780    let column_types = qcx.relation_type(&expr).column_types;
4781    if column_types.len() != 1 {
4782        sql_bail!(
4783            "Expected subselect to return 1 column, got {} columns",
4784            column_types.len()
4785        );
4786    }
4787    Ok(expr.select().into())
4788}
4789
4790fn plan_list_subquery(
4791    ecx: &ExprContext,
4792    query: &Query<Aug>,
4793) -> Result<CoercibleScalarExpr, PlanError> {
4794    plan_vector_like_subquery(
4795        ecx,
4796        query,
4797        |_| false,
4798        |elem_type| ListCreate { elem_type }.into(),
4799        |order_by| AggregateFunc::ListConcat { order_by },
4800        expr_func::ListListConcat.into(),
4801        |elem_type| {
4802            HirScalarExpr::literal(
4803                Datum::empty_list(),
4804                SqlScalarType::List {
4805                    element_type: Box::new(elem_type),
4806                    custom_id: None,
4807                },
4808            )
4809        },
4810        "list",
4811    )
4812}
4813
4814fn plan_array_subquery(
4815    ecx: &ExprContext,
4816    query: &Query<Aug>,
4817) -> Result<CoercibleScalarExpr, PlanError> {
4818    plan_vector_like_subquery(
4819        ecx,
4820        query,
4821        |elem_type| {
4822            matches!(
4823                elem_type,
4824                SqlScalarType::Char { .. }
4825                    | SqlScalarType::Array { .. }
4826                    | SqlScalarType::List { .. }
4827                    | SqlScalarType::Map { .. }
4828            )
4829        },
4830        |elem_type| ArrayCreate { elem_type }.into(),
4831        |order_by| AggregateFunc::ArrayConcat { order_by },
4832        expr_func::ArrayArrayConcat.into(),
4833        |elem_type| {
4834            HirScalarExpr::literal(
4835                Datum::empty_array(),
4836                SqlScalarType::Array(Box::new(elem_type)),
4837            )
4838        },
4839        "[]",
4840    )
4841}
4842
4843/// Generic function used to plan both array subqueries and list subqueries
4844fn plan_vector_like_subquery<F1, F2, F3, F4>(
4845    ecx: &ExprContext,
4846    query: &Query<Aug>,
4847    is_unsupported_type: F1,
4848    vector_create: F2,
4849    aggregate_concat: F3,
4850    binary_concat: BinaryFunc,
4851    empty_literal: F4,
4852    vector_type_string: &str,
4853) -> Result<CoercibleScalarExpr, PlanError>
4854where
4855    F1: Fn(&SqlScalarType) -> bool,
4856    F2: Fn(SqlScalarType) -> VariadicFunc,
4857    F3: Fn(Vec<ColumnOrder>) -> AggregateFunc,
4858    F4: Fn(SqlScalarType) -> HirScalarExpr,
4859{
4860    if !ecx.allow_subqueries {
4861        sql_bail!("{} does not allow subqueries", ecx.name)
4862    }
4863
4864    let mut qcx = ecx.derived_query_context();
4865    let mut planned_query = plan_query(&mut qcx, query)?;
4866    if planned_query.limit.is_some()
4867        || !planned_query
4868            .offset
4869            .clone()
4870            .try_into_literal_int64()
4871            .is_ok_and(|offset| offset == 0)
4872    {
4873        planned_query.expr = HirRelationExpr::top_k(
4874            planned_query.expr,
4875            vec![],
4876            planned_query.order_by.clone(),
4877            planned_query.limit,
4878            planned_query.offset,
4879            planned_query.group_size_hints.limit_input_group_size,
4880        );
4881    }
4882
4883    if planned_query.project.len() != 1 {
4884        sql_bail!(
4885            "Expected subselect to return 1 column, got {} columns",
4886            planned_query.project.len()
4887        );
4888    }
4889
4890    let project_column = *planned_query.project.get(0).unwrap();
4891    let elem_type = qcx
4892        .relation_type(&planned_query.expr)
4893        .column_types
4894        .get(project_column)
4895        .cloned()
4896        .unwrap()
4897        .scalar_type();
4898
4899    if is_unsupported_type(&elem_type) {
4900        bail_unsupported!(format!(
4901            "cannot build array from subquery because return type {}{}",
4902            ecx.humanize_sql_scalar_type(&elem_type, false),
4903            vector_type_string
4904        ));
4905    }
4906
4907    // `ColumnRef`s in `aggregation_exprs` refers to the columns produced by planning the
4908    // subquery above.
4909    let aggregation_exprs: Vec<_> = iter::once(HirScalarExpr::call_variadic(
4910        vector_create(elem_type.clone()),
4911        vec![HirScalarExpr::column(project_column)],
4912    ))
4913    .chain(
4914        planned_query
4915            .order_by
4916            .iter()
4917            .map(|co| HirScalarExpr::column(co.column)),
4918    )
4919    .collect();
4920
4921    // However, column references for `aggregation_projection` and `aggregation_order_by`
4922    // are with reference to the `exprs` of the aggregation expression.  Here that is
4923    // `aggregation_exprs`.
4924    let aggregation_projection = vec![0];
4925    let aggregation_order_by = planned_query
4926        .order_by
4927        .into_iter()
4928        .enumerate()
4929        .map(|(i, order)| ColumnOrder { column: i, ..order })
4930        .collect();
4931
4932    let reduced_expr = planned_query
4933        .expr
4934        .reduce(
4935            vec![],
4936            vec![AggregateExpr {
4937                func: aggregate_concat(aggregation_order_by),
4938                expr: Box::new(HirScalarExpr::call_variadic(
4939                    RecordCreate {
4940                        field_names: iter::repeat(ColumnName::from(""))
4941                            .take(aggregation_exprs.len())
4942                            .collect(),
4943                    },
4944                    aggregation_exprs,
4945                )),
4946                distinct: false,
4947            }],
4948            None,
4949        )
4950        .project(aggregation_projection);
4951
4952    // If `expr` has no rows, return an empty array/list rather than NULL.
4953    Ok(reduced_expr
4954        .select()
4955        .call_binary(empty_literal(elem_type), binary_concat)
4956        .into())
4957}
4958
4959fn plan_map_subquery(
4960    ecx: &ExprContext,
4961    query: &Query<Aug>,
4962) -> Result<CoercibleScalarExpr, PlanError> {
4963    if !ecx.allow_subqueries {
4964        sql_bail!("{} does not allow subqueries", ecx.name)
4965    }
4966
4967    let mut qcx = ecx.derived_query_context();
4968    let mut query = plan_query(&mut qcx, query)?;
4969    if query.limit.is_some()
4970        || !query
4971            .offset
4972            .clone()
4973            .try_into_literal_int64()
4974            .is_ok_and(|offset| offset == 0)
4975    {
4976        query.expr = HirRelationExpr::top_k(
4977            query.expr,
4978            vec![],
4979            query.order_by.clone(),
4980            query.limit,
4981            query.offset,
4982            query.group_size_hints.limit_input_group_size,
4983        );
4984    }
4985    if query.project.len() != 2 {
4986        sql_bail!(
4987            "expected map subquery to return 2 columns, got {} columns",
4988            query.project.len()
4989        );
4990    }
4991
4992    let query_types = qcx.relation_type(&query.expr).column_types;
4993    let key_column = query.project[0];
4994    let key_type = query_types[key_column].clone().scalar_type();
4995    let value_column = query.project[1];
4996    let value_type = query_types[value_column].clone().scalar_type();
4997
4998    if key_type != SqlScalarType::String {
4999        sql_bail!("cannot build map from subquery because first column is not of type text");
5000    }
5001
5002    let aggregation_exprs: Vec<_> = iter::once(HirScalarExpr::call_variadic(
5003        RecordCreate {
5004            field_names: vec![ColumnName::from("key"), ColumnName::from("value")],
5005        },
5006        vec![
5007            HirScalarExpr::column(key_column),
5008            HirScalarExpr::column(value_column),
5009        ],
5010    ))
5011    .chain(
5012        query
5013            .order_by
5014            .iter()
5015            .map(|co| HirScalarExpr::column(co.column)),
5016    )
5017    .collect();
5018
5019    let expr = query
5020        .expr
5021        .reduce(
5022            vec![],
5023            vec![AggregateExpr {
5024                func: AggregateFunc::MapAgg {
5025                    order_by: query
5026                        .order_by
5027                        .into_iter()
5028                        .enumerate()
5029                        .map(|(i, order)| ColumnOrder { column: i, ..order })
5030                        .collect(),
5031                    value_type: value_type.clone(),
5032                },
5033                expr: Box::new(HirScalarExpr::call_variadic(
5034                    RecordCreate {
5035                        field_names: iter::repeat(ColumnName::from(""))
5036                            .take(aggregation_exprs.len())
5037                            .collect(),
5038                    },
5039                    aggregation_exprs,
5040                )),
5041                distinct: false,
5042            }],
5043            None,
5044        )
5045        .project(vec![0]);
5046
5047    // If `expr` has no rows, return an empty map rather than NULL.
5048    let expr = HirScalarExpr::call_variadic(
5049        Coalesce,
5050        vec![
5051            expr.select(),
5052            HirScalarExpr::literal(
5053                Datum::empty_map(),
5054                SqlScalarType::Map {
5055                    value_type: Box::new(value_type),
5056                    custom_id: None,
5057                },
5058            ),
5059        ],
5060    );
5061
5062    Ok(expr.into())
5063}
5064
5065fn plan_collate(
5066    ecx: &ExprContext,
5067    expr: &Expr<Aug>,
5068    collation: &UnresolvedItemName,
5069) -> Result<CoercibleScalarExpr, PlanError> {
5070    if collation.0.len() == 2
5071        && collation.0[0] == ident!(mz_repr::namespaces::PG_CATALOG_SCHEMA)
5072        && collation.0[1] == ident!("default")
5073    {
5074        plan_expr(ecx, expr)
5075    } else {
5076        bail_unsupported!("COLLATE");
5077    }
5078}
5079
5080/// Plans a slice of expressions.
5081///
5082/// This function is a simple convenience function for mapping [`plan_expr`]
5083/// over a slice of expressions. The planned expressions are returned in the
5084/// same order as the input. If any of the expressions fail to plan, returns an
5085/// error instead.
5086fn plan_exprs<E>(ecx: &ExprContext, exprs: &[E]) -> Result<Vec<CoercibleScalarExpr>, PlanError>
5087where
5088    E: std::borrow::Borrow<Expr<Aug>>,
5089{
5090    let mut out = vec![];
5091    for expr in exprs {
5092        out.push(plan_expr(ecx, expr.borrow())?);
5093    }
5094    Ok(out)
5095}
5096
5097/// Plans an `ARRAY` expression.
5098fn plan_array(
5099    ecx: &ExprContext,
5100    exprs: &[Expr<Aug>],
5101    type_hint: Option<&SqlScalarType>,
5102) -> Result<CoercibleScalarExpr, PlanError> {
5103    // Plan each element expression.
5104    let mut out = vec![];
5105    for expr in exprs {
5106        out.push(match expr {
5107            // Special case nested ARRAY expressions so we can plumb
5108            // the type hint through.
5109            Expr::Array(exprs) => plan_array(ecx, exprs, type_hint.clone())?,
5110            _ => plan_expr(ecx, expr)?,
5111        });
5112    }
5113
5114    // Attempt to make use of the type hint.
5115    let type_hint = match type_hint {
5116        // The user has provided an explicit cast to an array type. We know the
5117        // element type to coerce to. Need to be careful, though: if there's
5118        // evidence that any of the array elements are themselves arrays, we
5119        // want to coerce to the array type, not the element type.
5120        Some(SqlScalarType::Array(elem_type)) => {
5121            let multidimensional = out.iter().any(|e| {
5122                matches!(
5123                    ecx.scalar_type(e),
5124                    CoercibleScalarType::Coerced(SqlScalarType::Array(_))
5125                )
5126            });
5127            if multidimensional {
5128                type_hint
5129            } else {
5130                Some(&**elem_type)
5131            }
5132        }
5133        // The user provided an explicit cast to a non-array type. We'll have to
5134        // guess what the correct type for the array. Our caller will then
5135        // handle converting that array type to the desired non-array type.
5136        Some(_) => None,
5137        // No type hint. We'll have to guess the correct type for the array.
5138        None => None,
5139    };
5140
5141    // Coerce all elements to the same type.
5142    let (elem_type, exprs) = if exprs.is_empty() {
5143        if let Some(elem_type) = type_hint {
5144            (elem_type.clone(), vec![])
5145        } else {
5146            sql_bail!("cannot determine type of empty array");
5147        }
5148    } else {
5149        let out = coerce_homogeneous_exprs(&ecx.with_name("ARRAY"), out, type_hint)?;
5150        (ecx.scalar_type(&out[0]), out)
5151    };
5152
5153    // Arrays of `char` type are disallowed due to a known limitation:
5154    // https://github.com/MaterializeInc/database-issues/issues/2360.
5155    //
5156    // Arrays of `list` and `map` types are disallowed due to mind-bending
5157    // semantics.
5158    if matches!(
5159        elem_type,
5160        SqlScalarType::Char { .. } | SqlScalarType::List { .. } | SqlScalarType::Map { .. }
5161    ) {
5162        bail_unsupported!(format!(
5163            "{}[]",
5164            ecx.humanize_sql_scalar_type(&elem_type, false)
5165        ));
5166    }
5167
5168    Ok(HirScalarExpr::call_variadic(ArrayCreate { elem_type }, exprs).into())
5169}
5170
5171fn plan_list(
5172    ecx: &ExprContext,
5173    exprs: &[Expr<Aug>],
5174    type_hint: Option<&SqlScalarType>,
5175) -> Result<CoercibleScalarExpr, PlanError> {
5176    let (elem_type, exprs) = if exprs.is_empty() {
5177        if let Some(SqlScalarType::List { element_type, .. }) = type_hint {
5178            (element_type.without_modifiers(), vec![])
5179        } else {
5180            sql_bail!("cannot determine type of empty list");
5181        }
5182    } else {
5183        let type_hint = match type_hint {
5184            Some(SqlScalarType::List { element_type, .. }) => Some(&**element_type),
5185            _ => None,
5186        };
5187
5188        let mut out = vec![];
5189        for expr in exprs {
5190            out.push(match expr {
5191                // Special case nested LIST expressions so we can plumb
5192                // the type hint through.
5193                Expr::List(exprs) => plan_list(ecx, exprs, type_hint)?,
5194                _ => plan_expr(ecx, expr)?,
5195            });
5196        }
5197        let out = coerce_homogeneous_exprs(&ecx.with_name("LIST"), out, type_hint)?;
5198        (ecx.scalar_type(&out[0]).without_modifiers(), out)
5199    };
5200
5201    if matches!(elem_type, SqlScalarType::Char { .. }) {
5202        bail_unsupported!("char list");
5203    }
5204
5205    Ok(HirScalarExpr::call_variadic(ListCreate { elem_type }, exprs).into())
5206}
5207
5208fn plan_map(
5209    ecx: &ExprContext,
5210    entries: &[MapEntry<Aug>],
5211    type_hint: Option<&SqlScalarType>,
5212) -> Result<CoercibleScalarExpr, PlanError> {
5213    let (value_type, exprs) = if entries.is_empty() {
5214        if let Some(SqlScalarType::Map { value_type, .. }) = type_hint {
5215            (value_type.without_modifiers(), vec![])
5216        } else {
5217            sql_bail!("cannot determine type of empty map");
5218        }
5219    } else {
5220        let type_hint = match type_hint {
5221            Some(SqlScalarType::Map { value_type, .. }) => Some(&**value_type),
5222            _ => None,
5223        };
5224
5225        let mut keys = vec![];
5226        let mut values = vec![];
5227        for MapEntry { key, value } in entries {
5228            let key = plan_expr(ecx, key)?.type_as(ecx, &SqlScalarType::String)?;
5229            let value = match value {
5230                // Special case nested MAP expressions so we can plumb
5231                // the type hint through.
5232                Expr::Map(entries) => plan_map(ecx, entries, type_hint)?,
5233                _ => plan_expr(ecx, value)?,
5234            };
5235            keys.push(key);
5236            values.push(value);
5237        }
5238        let values = coerce_homogeneous_exprs(&ecx.with_name("MAP"), values, type_hint)?;
5239        let value_type = ecx.scalar_type(&values[0]).without_modifiers();
5240        let out = itertools::interleave(keys, values).collect();
5241        (value_type, out)
5242    };
5243
5244    if matches!(value_type, SqlScalarType::Char { .. }) {
5245        bail_unsupported!("char map");
5246    }
5247
5248    let expr = HirScalarExpr::call_variadic(MapBuild { value_type }, exprs);
5249    Ok(expr.into())
5250}
5251
5252/// Coerces a list of expressions such that all input expressions will be cast
5253/// to the same type. If successful, returns a new list of expressions in the
5254/// same order as the input, where each expression has the appropriate casts to
5255/// make them all of a uniform type.
5256///
5257/// If `force_type` is `Some`, the expressions are forced to the specified type
5258/// via an explicit cast. Otherwise the best common type is guessed via
5259/// [`typeconv::guess_best_common_type`] and conversions are attempted via
5260/// implicit casts
5261///
5262/// Note that this is our implementation of Postgres' type conversion for
5263/// ["`UNION`, `CASE`, and Related Constructs"][union-type-conv], though it
5264/// isn't yet used in all of those cases.
5265///
5266/// [union-type-conv]:
5267/// https://www.postgresql.org/docs/12/typeconv-union-case.html
5268pub fn coerce_homogeneous_exprs(
5269    ecx: &ExprContext,
5270    exprs: Vec<CoercibleScalarExpr>,
5271    force_type: Option<&SqlScalarType>,
5272) -> Result<Vec<HirScalarExpr>, PlanError> {
5273    assert!(!exprs.is_empty());
5274
5275    let target_holder;
5276    let target = match force_type {
5277        Some(t) => t,
5278        None => {
5279            let types: Vec<_> = exprs.iter().map(|e| ecx.scalar_type(e)).collect();
5280            target_holder = typeconv::guess_best_common_type(ecx, &types)?;
5281            &target_holder
5282        }
5283    };
5284
5285    // Try to cast all expressions to `target`.
5286    let mut out = Vec::new();
5287    for expr in exprs {
5288        let arg = typeconv::plan_coerce(ecx, expr, target)?;
5289        let ccx = match force_type {
5290            None => CastContext::Implicit,
5291            Some(_) => CastContext::Explicit,
5292        };
5293        match typeconv::plan_cast(ecx, ccx, arg.clone(), target) {
5294            Ok(expr) => out.push(expr),
5295            Err(_) => sql_bail!(
5296                "{} could not convert type {} to {}",
5297                ecx.name,
5298                ecx.humanize_sql_scalar_type(&ecx.scalar_type(&arg), false),
5299                ecx.humanize_sql_scalar_type(target, false),
5300            ),
5301        }
5302    }
5303    Ok(out)
5304}
5305
5306/// Creates a `ColumnOrder` from an `OrderByExpr` and column index.
5307/// Column index is specified by the caller, but `desc` and `nulls_last` is figured out here.
5308pub(crate) fn resolve_desc_and_nulls_last<T: AstInfo>(
5309    obe: &OrderByExpr<T>,
5310    column: usize,
5311) -> ColumnOrder {
5312    let desc = !obe.asc.unwrap_or(true);
5313    ColumnOrder {
5314        column,
5315        desc,
5316        // https://www.postgresql.org/docs/14/queries-order.html
5317        //   "NULLS FIRST is the default for DESC order, and NULLS LAST otherwise"
5318        nulls_last: obe.nulls_last.unwrap_or(!desc),
5319    }
5320}
5321
5322/// Plans the ORDER BY clause of a window function.
5323///
5324/// Unfortunately, we have to create two HIR structs from an AST OrderByExpr:
5325/// A ColumnOrder has asc/desc and nulls first/last, but can't represent an HirScalarExpr, just
5326/// a column reference by index. Therefore, we return both HirScalarExprs and ColumnOrders.
5327/// Note that the column references in the ColumnOrders point NOT to input columns, but into the
5328/// `Vec<HirScalarExpr>` that we return.
5329fn plan_function_order_by(
5330    ecx: &ExprContext,
5331    order_by: &[OrderByExpr<Aug>],
5332) -> Result<(Vec<HirScalarExpr>, Vec<ColumnOrder>), PlanError> {
5333    let mut order_by_exprs = vec![];
5334    let mut col_orders = vec![];
5335    {
5336        for (i, obe) in order_by.iter().enumerate() {
5337            // Unlike `SELECT ... ORDER BY` clauses, function `ORDER BY` clauses
5338            // do not support ordinal references in PostgreSQL. So we use
5339            // `plan_expr` directly rather than `plan_order_by_or_distinct_expr`.
5340            let expr = plan_expr(ecx, &obe.expr)?.type_as_any(ecx)?;
5341            order_by_exprs.push(expr);
5342            col_orders.push(resolve_desc_and_nulls_last(obe, i));
5343        }
5344    }
5345    Ok((order_by_exprs, col_orders))
5346}
5347
5348/// Returns a human-readable rendering of `name`, falling back to a debug
5349/// dump of the `ResolvedItemName` if humanization fails. Used to construct
5350/// user-facing error messages on already-failing paths, where we'd rather
5351/// surface the raw resolved name than emit a useless `<unknown>` placeholder.
5352fn humanize_or_debug(scx: &StatementContext, name: &ResolvedItemName) -> String {
5353    scx.humanize_resolved_name(name)
5354        .map(|n| n.to_string())
5355        .unwrap_or_else(|_| format!("<error when trying to humanize `{name:?}`>"))
5356}
5357
5358/// Common part of the planning of windowed and non-windowed aggregation functions.
5359fn plan_aggregate_common(
5360    ecx: &ExprContext,
5361    Function::<Aug> {
5362        name,
5363        args,
5364        filter,
5365        over: _,
5366        distinct,
5367    }: &Function<Aug>,
5368) -> Result<AggregateExpr, PlanError> {
5369    // Normal aggregate functions, like `sum`, expect as input a single expression
5370    // which yields the datum to aggregate. Order sensitive aggregate functions,
5371    // like `jsonb_agg`, are special, and instead expect a Record whose first
5372    // element yields the datum to aggregate and whose successive elements yield
5373    // keys to order by. This expectation is hard coded within the implementation
5374    // of each of the order-sensitive aggregates. The specification of how many
5375    // order by keys to consider, and in what order, is passed via the `order_by`
5376    // field on the `AggregateFunc` variant.
5377
5378    // While all aggregate functions support the ORDER BY syntax, it's a no-op for
5379    // most, so explicitly drop it if the function doesn't care about order. This
5380    // prevents the projection into Record below from triggering on unsupported
5381    // functions.
5382
5383    let impls = match resolve_func(ecx, name, args)? {
5384        Func::Aggregate(impls) => impls,
5385        _ => bail_internal!("plan_aggregate_common called on non-aggregate function"),
5386    };
5387
5388    // We follow PostgreSQL's rule here for mapping `count(*)` into the
5389    // generalized function selection framework. The rule is simple: the user
5390    // must type `count(*)`, but the function selection framework sees an empty
5391    // parameter list, as if the user had typed `count()`. But if the user types
5392    // `count()` directly, that is an error. Like PostgreSQL, we apply these
5393    // rules to all aggregates, not just `count`, since we may one day support
5394    // user-defined aggregates, including user-defined aggregates that take no
5395    // parameters.
5396    let (args, order_by) = match &args {
5397        FunctionArgs::Star => (vec![], vec![]),
5398        FunctionArgs::Args { args, order_by } => {
5399            if args.is_empty() {
5400                sql_bail!(
5401                    "{}(*) must be used to call a parameterless aggregate function",
5402                    humanize_or_debug(ecx.qcx.scx, name)
5403                );
5404            }
5405            let args = plan_exprs(ecx, args)?;
5406            (args, order_by.clone())
5407        }
5408    };
5409
5410    let (order_by_exprs, col_orders) = plan_function_order_by(ecx, &order_by)?;
5411
5412    let (mut expr, func) = func::select_impl(ecx, FuncSpec::Func(name), impls, args, col_orders)?;
5413    if let Some(filter) = &filter {
5414        // If a filter is present, as in
5415        //
5416        //     <agg>(<expr>) FILTER (WHERE <cond>)
5417        //
5418        // we plan it by essentially rewriting the expression to
5419        //
5420        //     <agg>(CASE WHEN <cond> THEN <expr> ELSE <identity>)
5421        //
5422        // where <identity> is the identity input for <agg>.
5423        let cond =
5424            plan_expr(&ecx.with_name("FILTER"), filter)?.type_as(ecx, &SqlScalarType::Bool)?;
5425        let expr_typ = ecx.scalar_type(&expr);
5426        expr = HirScalarExpr::if_then_else(
5427            cond,
5428            expr,
5429            HirScalarExpr::literal(func.identity_datum(), expr_typ),
5430        );
5431    }
5432
5433    let mut seen_outer = false;
5434    let mut seen_inner = false;
5435    #[allow(deprecated)]
5436    expr.visit_columns(0, &mut |depth, col| {
5437        if depth == 0 && col.level == 0 {
5438            seen_inner = true;
5439        } else if col.level > depth {
5440            seen_outer = true;
5441        }
5442    });
5443    if seen_outer && !seen_inner {
5444        bail_unsupported!(
5445            3720,
5446            "aggregate functions that refer exclusively to outer columns"
5447        );
5448    }
5449
5450    // If a function supports ORDER BY (even if there was no ORDER BY specified),
5451    // map the needed expressions into the aggregate datum.
5452    if func.is_order_sensitive() {
5453        let field_names = iter::repeat(ColumnName::from(""))
5454            .take(1 + order_by_exprs.len())
5455            .collect();
5456        let mut exprs = vec![expr];
5457        exprs.extend(order_by_exprs);
5458        expr = HirScalarExpr::call_variadic(RecordCreate { field_names }, exprs);
5459    }
5460
5461    Ok(AggregateExpr {
5462        func,
5463        expr: Box::new(expr),
5464        distinct: *distinct,
5465    })
5466}
5467
5468fn plan_identifier(ecx: &ExprContext, names: &[Ident]) -> Result<HirScalarExpr, PlanError> {
5469    let mut names = names.to_vec();
5470    // The parser guarantees that an `Expr::Identifier` is constructed with a
5471    // non-empty list of name parts, so an empty list here is an internal bug.
5472    let Some(last) = names.pop() else {
5473        bail_internal!("empty identifier");
5474    };
5475    let col_name = normalize::column_name(last);
5476
5477    // If the name is qualified, it must refer to a column in a table.
5478    if !names.is_empty() {
5479        let table_name = normalize::unresolved_item_name(UnresolvedItemName(names))?;
5480        let (i, i_name) = ecx.scope.resolve_table_column(
5481            &ecx.qcx.outer_scopes,
5482            &table_name,
5483            &col_name,
5484            &mut ecx.qcx.name_manager.borrow_mut(),
5485        )?;
5486        return Ok(HirScalarExpr::named_column(i, i_name));
5487    }
5488
5489    // If the name is unqualified, first check if it refers to a column. Track any similar names
5490    // that might exist for a better error message.
5491    let similar_names = match ecx.scope.resolve_column(
5492        &ecx.qcx.outer_scopes,
5493        &col_name,
5494        &mut ecx.qcx.name_manager.borrow_mut(),
5495    ) {
5496        Ok((i, i_name)) => {
5497            return Ok(HirScalarExpr::named_column(i, i_name));
5498        }
5499        Err(PlanError::UnknownColumn { similar, .. }) => similar,
5500        Err(e) => return Err(e),
5501    };
5502
5503    // The name doesn't refer to a column. Check if it is a whole-row reference
5504    // to a table.
5505    let items = ecx.scope.items_from_table(
5506        &ecx.qcx.outer_scopes,
5507        &PartialItemName {
5508            database: None,
5509            schema: None,
5510            item: col_name.as_str().to_owned(),
5511        },
5512    )?;
5513    match items.as_slice() {
5514        // The name doesn't refer to a table either. Return an error.
5515        [] => Err(PlanError::UnknownColumn {
5516            table: None,
5517            column: col_name,
5518            similar: similar_names,
5519        }),
5520        // The name refers to a table that is the result of a function that
5521        // returned a single column. Per PostgreSQL, this is a special case
5522        // that returns the value directly.
5523        // See: https://github.com/postgres/postgres/blob/22592e10b/src/backend/parser/parse_expr.c#L2519-L2524
5524        [(column, item)] if item.from_single_column_function => Ok(HirScalarExpr::named_column(
5525            *column,
5526            ecx.qcx.name_manager.borrow_mut().intern_scope_item(item),
5527        )),
5528        // The name refers to a normal table. Return a record containing all the
5529        // columns of the table.
5530        _ => {
5531            let mut has_exists_column = None;
5532            let (exprs, field_names): (Vec<_>, Vec<_>) = items
5533                .into_iter()
5534                .filter_map(|(column, item)| {
5535                    if item.is_exists_column_for_a_table_function_that_was_in_the_target_list {
5536                        has_exists_column = Some(column);
5537                        None
5538                    } else {
5539                        let expr = HirScalarExpr::named_column(
5540                            column,
5541                            ecx.qcx.name_manager.borrow_mut().intern_scope_item(item),
5542                        );
5543                        let name = item.column_name.clone();
5544                        Some((expr, name))
5545                    }
5546                })
5547                .unzip();
5548            // For the special case of a table function with a single column, the single column is instead not wrapped.
5549            let expr = if exprs.len() == 1 && has_exists_column.is_some() {
5550                exprs.into_element()
5551            } else {
5552                HirScalarExpr::call_variadic(RecordCreate { field_names }, exprs)
5553            };
5554            if let Some(has_exists_column) = has_exists_column {
5555                Ok(HirScalarExpr::if_then_else(
5556                    HirScalarExpr::unnamed_column(has_exists_column)
5557                        .call_unary(UnaryFunc::IsNull(mz_expr::func::IsNull)),
5558                    HirScalarExpr::literal_null(ecx.scalar_type(&expr)),
5559                    expr,
5560                ))
5561            } else {
5562                Ok(expr)
5563            }
5564        }
5565    }
5566}
5567
5568fn plan_op(
5569    ecx: &ExprContext,
5570    op: &str,
5571    expr1: &Expr<Aug>,
5572    expr2: Option<&Expr<Aug>>,
5573) -> Result<HirScalarExpr, PlanError> {
5574    let impls = func::resolve_op(op)?;
5575    let args = match expr2 {
5576        None => plan_exprs(ecx, &[expr1])?,
5577        Some(expr2) => plan_exprs(ecx, &[expr1, expr2])?,
5578    };
5579    func::select_impl(ecx, FuncSpec::Op(op), impls, args, vec![])
5580}
5581
5582fn plan_function<'a>(
5583    ecx: &ExprContext,
5584    f @ Function {
5585        name,
5586        args,
5587        filter,
5588        over,
5589        distinct,
5590    }: &'a Function<Aug>,
5591) -> Result<HirScalarExpr, PlanError> {
5592    let impls = match resolve_func(ecx, name, args)? {
5593        Func::Table(_) => {
5594            sql_bail!(
5595                "table functions are not allowed in {} (function {})",
5596                ecx.name,
5597                name
5598            );
5599        }
5600        Func::Scalar(impls) => {
5601            if over.is_some() {
5602                sql_bail!(
5603                    "OVER clause not allowed on {name}. The OVER clause can only be used with window functions (including aggregations)."
5604                );
5605            }
5606            impls
5607        }
5608        Func::ScalarWindow(impls) => {
5609            let (
5610                ignore_nulls,
5611                order_by_exprs,
5612                col_orders,
5613                _window_frame,
5614                partition_by,
5615                scalar_args,
5616            ) = plan_window_function_non_aggr(ecx, f)?;
5617
5618            // All scalar window functions have 0 parameters. Let's print a nice error msg if the
5619            // user gave some args. (The below `func::select_impl` would fail anyway, but the error
5620            // msg there is less informative.)
5621            if !scalar_args.is_empty() {
5622                if let ResolvedItemName::Item {
5623                    full_name: FullItemName { item, .. },
5624                    ..
5625                } = name
5626                {
5627                    sql_bail!(
5628                        "function {} has 0 parameters, but was called with {}",
5629                        item,
5630                        scalar_args.len()
5631                    );
5632                }
5633            }
5634
5635            // Note: the window frame doesn't affect scalar window funcs, but, strangely, we should
5636            // accept a window frame here without an error msg. (Postgres also does this.)
5637            // TODO: maybe we should give a notice
5638
5639            let func = func::select_impl(ecx, FuncSpec::Func(name), impls, scalar_args, vec![])?;
5640
5641            if ignore_nulls {
5642                // If we ever add a scalar window function that supports ignore, then don't forget
5643                // to also update HIR EXPLAIN.
5644                bail_unsupported!(IGNORE_NULLS_ERROR_MSG);
5645            }
5646
5647            return Ok(HirScalarExpr::windowing(WindowExpr {
5648                func: WindowExprType::Scalar(ScalarWindowExpr {
5649                    func,
5650                    order_by: col_orders,
5651                }),
5652                partition_by,
5653                order_by: order_by_exprs,
5654            }));
5655        }
5656        Func::ValueWindow(impls) => {
5657            let window_plan = plan_window_function_non_aggr(ecx, f)?;
5658            let (ignore_nulls, order_by_exprs, col_orders, window_frame, partition_by, win_args) =
5659                window_plan;
5660
5661            let (args_encoded, func) =
5662                func::select_impl(ecx, FuncSpec::Func(name), impls, win_args, vec![])?;
5663
5664            if ignore_nulls {
5665                match func {
5666                    ValueWindowFunc::Lag | ValueWindowFunc::Lead => {}
5667                    _ => bail_unsupported!(IGNORE_NULLS_ERROR_MSG),
5668                }
5669            }
5670
5671            return Ok(HirScalarExpr::windowing(WindowExpr {
5672                func: WindowExprType::Value(ValueWindowExpr {
5673                    func,
5674                    args: Box::new(args_encoded),
5675                    order_by: col_orders,
5676                    window_frame,
5677                    ignore_nulls, // (RESPECT NULLS is the default)
5678                }),
5679                partition_by,
5680                order_by: order_by_exprs,
5681            }));
5682        }
5683        Func::Aggregate(_) => {
5684            if f.over.is_none() {
5685                // Not a window aggregate. Something is wrong.
5686                if ecx.allow_aggregates {
5687                    // Should already have been caught by `scope.resolve_expr` in `plan_expr_inner`
5688                    // (after having been planned earlier in `Step 5` of `plan_select_from_where`).
5689                    sql_bail!(
5690                        "Internal error: encountered unplanned non-windowed aggregate function: {:?}",
5691                        name,
5692                    );
5693                } else {
5694                    // scope.resolve_expr didn't catch it because we have not yet planned it,
5695                    // because it was in an unsupported context.
5696                    sql_bail!(
5697                        "aggregate functions are not allowed in {} (function {})",
5698                        ecx.name,
5699                        name
5700                    );
5701                }
5702            } else {
5703                let (ignore_nulls, order_by_exprs, col_orders, window_frame, partition_by) =
5704                    plan_window_function_common(ecx, &f.name, &f.over)?;
5705
5706                // https://github.com/MaterializeInc/database-issues/issues/6720
5707                match (&window_frame.start_bound, &window_frame.end_bound) {
5708                    (
5709                        mz_expr::WindowFrameBound::UnboundedPreceding,
5710                        mz_expr::WindowFrameBound::OffsetPreceding(..),
5711                    )
5712                    | (
5713                        mz_expr::WindowFrameBound::UnboundedPreceding,
5714                        mz_expr::WindowFrameBound::OffsetFollowing(..),
5715                    )
5716                    | (
5717                        mz_expr::WindowFrameBound::OffsetPreceding(..),
5718                        mz_expr::WindowFrameBound::UnboundedFollowing,
5719                    )
5720                    | (
5721                        mz_expr::WindowFrameBound::OffsetFollowing(..),
5722                        mz_expr::WindowFrameBound::UnboundedFollowing,
5723                    ) => bail_unsupported!("mixed unbounded - offset frames"),
5724                    (_, _) => {} // other cases are ok
5725                }
5726
5727                if ignore_nulls {
5728                    // https://github.com/MaterializeInc/database-issues/issues/6722
5729                    // If we ever add support for ignore_nulls for a window aggregate, then don't
5730                    // forget to also update HIR EXPLAIN.
5731                    bail_unsupported!(IGNORE_NULLS_ERROR_MSG);
5732                }
5733
5734                let aggregate_expr = plan_aggregate_common(ecx, f)?;
5735
5736                if aggregate_expr.distinct {
5737                    // https://github.com/MaterializeInc/database-issues/issues/6626
5738                    bail_unsupported!("DISTINCT in window aggregates");
5739                }
5740
5741                return Ok(HirScalarExpr::windowing(WindowExpr {
5742                    func: WindowExprType::Aggregate(AggregateWindowExpr {
5743                        aggregate_expr,
5744                        order_by: col_orders,
5745                        window_frame,
5746                    }),
5747                    partition_by,
5748                    order_by: order_by_exprs,
5749                }));
5750            }
5751        }
5752    };
5753
5754    if over.is_some() {
5755        bail_internal!("OVER clause should have been handled by the window function path above");
5756    }
5757
5758    if *distinct {
5759        sql_bail!(
5760            "DISTINCT specified, but {} is not an aggregate function",
5761            humanize_or_debug(ecx.qcx.scx, name)
5762        );
5763    }
5764    if filter.is_some() {
5765        sql_bail!(
5766            "FILTER specified, but {} is not an aggregate function",
5767            humanize_or_debug(ecx.qcx.scx, name)
5768        );
5769    }
5770
5771    let scalar_args = match &args {
5772        FunctionArgs::Star => {
5773            sql_bail!(
5774                "* argument is invalid with non-aggregate function {}",
5775                humanize_or_debug(ecx.qcx.scx, name)
5776            )
5777        }
5778        FunctionArgs::Args { args, order_by } => {
5779            if !order_by.is_empty() {
5780                sql_bail!(
5781                    "ORDER BY specified, but {} is not an aggregate function",
5782                    humanize_or_debug(ecx.qcx.scx, name)
5783                );
5784            }
5785            plan_exprs(ecx, args)?
5786        }
5787    };
5788
5789    func::select_impl(ecx, FuncSpec::Func(name), impls, scalar_args, vec![])
5790}
5791
5792pub const IGNORE_NULLS_ERROR_MSG: &str =
5793    "IGNORE NULLS and RESPECT NULLS options for functions other than LAG and LEAD";
5794
5795/// Resolves the name to a set of function implementations.
5796///
5797/// If the name does not specify a known built-in function, returns an error.
5798pub fn resolve_func(
5799    ecx: &ExprContext,
5800    name: &ResolvedItemName,
5801    args: &mz_sql_parser::ast::FunctionArgs<Aug>,
5802) -> Result<&'static Func, PlanError> {
5803    if let Ok(i) = ecx.qcx.scx.get_item_by_resolved_name(name) {
5804        if let Ok(f) = i.func() {
5805            return Ok(f);
5806        }
5807    }
5808
5809    // Couldn't resolve function with this name, so generate verbose error
5810    // message.
5811    let cexprs = match args {
5812        mz_sql_parser::ast::FunctionArgs::Star => vec![],
5813        mz_sql_parser::ast::FunctionArgs::Args { args, order_by } => {
5814            if !order_by.is_empty() {
5815                sql_bail!(
5816                    "ORDER BY specified, but {} is not an aggregate function",
5817                    name
5818                );
5819            }
5820            plan_exprs(ecx, args)?
5821        }
5822    };
5823
5824    let arg_types: Vec<_> = cexprs
5825        .into_iter()
5826        .map(|ty| match ecx.scalar_type(&ty) {
5827            CoercibleScalarType::Coerced(ty) => ecx.humanize_sql_scalar_type(&ty, false),
5828            CoercibleScalarType::Record(_) => "record".to_string(),
5829            CoercibleScalarType::Uncoerced => "unknown".to_string(),
5830        })
5831        .collect();
5832
5833    Err(PlanError::UnknownFunction {
5834        name: name.to_string(),
5835        arg_types,
5836    })
5837}
5838
5839fn plan_is_expr<'a>(
5840    ecx: &ExprContext,
5841    expr: &'a Expr<Aug>,
5842    construct: &IsExprConstruct<Aug>,
5843    not: bool,
5844) -> Result<HirScalarExpr, PlanError> {
5845    let expr_hir = plan_expr(ecx, expr)?;
5846
5847    let mut result = match construct {
5848        IsExprConstruct::Null => {
5849            // PostgreSQL can plan `NULL IS NULL` but not `$1 IS NULL`. This is
5850            // at odds with our type coercion rules, which treat `NULL` literals
5851            // and unconstrained parameters identically. Providing a type hint
5852            // of string means we wind up supporting both.
5853            expr_hir.type_as_any(ecx)?.call_is_null()
5854        }
5855        IsExprConstruct::Unknown => expr_hir.type_as(ecx, &SqlScalarType::Bool)?.call_is_null(),
5856        IsExprConstruct::True => expr_hir
5857            .type_as(ecx, &SqlScalarType::Bool)?
5858            .call_unary(UnaryFunc::IsTrue(expr_func::IsTrue)),
5859        IsExprConstruct::False => expr_hir
5860            .type_as(ecx, &SqlScalarType::Bool)?
5861            .call_unary(UnaryFunc::IsFalse(expr_func::IsFalse)),
5862        IsExprConstruct::DistinctFrom(expr2) => {
5863            // There are three cases:
5864            // 1. Both terms are non-null, in which case the result should be `a != b`.
5865            // 2. Exactly one term is null, in which case the result should be true.
5866            // 3. Both terms are null, in which case the result should be false.
5867            //
5868            // (a != b OR a IS NULL OR b IS NULL) AND (a IS NOT NULL OR b IS NOT NULL)
5869
5870            // We'll need `expr != expr2`, but don't just construct this HIR directly. Instead,
5871            // construct an AST expression for `expr != expr2` and plan it to get proper type
5872            // checking, implicit casts, etc. (This seems to be also what Postgres does.)
5873            let ne_ast = expr.clone().not_equals(expr2.as_ref().clone());
5874            let ne_hir = plan_expr(ecx, &ne_ast)?.type_as_any(ecx)?;
5875
5876            let expr1_hir = expr_hir.type_as_any(ecx)?;
5877            let expr2_hir = plan_expr(ecx, expr2)?.type_as_any(ecx)?;
5878
5879            let term1 = HirScalarExpr::variadic_or(vec![
5880                ne_hir,
5881                expr1_hir.clone().call_is_null(),
5882                expr2_hir.clone().call_is_null(),
5883            ]);
5884            let term2 = HirScalarExpr::variadic_or(vec![
5885                expr1_hir.call_is_null().not(),
5886                expr2_hir.call_is_null().not(),
5887            ]);
5888            term1.and(term2)
5889        }
5890    };
5891    if not {
5892        result = result.not();
5893    }
5894    Ok(result)
5895}
5896
5897fn plan_case<'a>(
5898    ecx: &ExprContext,
5899    operand: &'a Option<Box<Expr<Aug>>>,
5900    conditions: &'a [Expr<Aug>],
5901    results: &'a [Expr<Aug>],
5902    else_result: &'a Option<Box<Expr<Aug>>>,
5903) -> Result<HirScalarExpr, PlanError> {
5904    let mut cond_exprs = Vec::new();
5905    let mut result_exprs = Vec::new();
5906    for (c, r) in conditions.iter().zip_eq(results) {
5907        let c = match operand {
5908            Some(operand) => operand.clone().equals(c.clone()),
5909            None => c.clone(),
5910        };
5911        let cexpr = plan_expr(ecx, &c)?.type_as(ecx, &SqlScalarType::Bool)?;
5912        cond_exprs.push(cexpr);
5913        result_exprs.push(r);
5914    }
5915    result_exprs.push(match else_result {
5916        Some(else_result) => else_result,
5917        None => &Expr::Value(Value::Null),
5918    });
5919    let mut result_exprs = coerce_homogeneous_exprs(
5920        &ecx.with_name("CASE"),
5921        plan_exprs(ecx, &result_exprs)?,
5922        None,
5923    )?;
5924    let mut expr = result_exprs.pop().unwrap();
5925    assert_eq!(cond_exprs.len(), result_exprs.len());
5926    for (cexpr, rexpr) in cond_exprs
5927        .into_iter()
5928        .rev()
5929        .zip_eq(result_exprs.into_iter().rev())
5930    {
5931        expr = HirScalarExpr::if_then_else(cexpr, rexpr, expr);
5932    }
5933    Ok(expr)
5934}
5935
5936fn plan_literal<'a>(l: &'a Value) -> Result<CoercibleScalarExpr, PlanError> {
5937    let (datum, scalar_type) = match l {
5938        Value::Number(s) => {
5939            let d = strconv::parse_numeric(s.as_str())?;
5940            if !s.contains(&['E', '.'][..]) {
5941                // Maybe representable as an int?
5942                if let Ok(n) = d.0.try_into() {
5943                    (Datum::Int32(n), SqlScalarType::Int32)
5944                } else if let Ok(n) = d.0.try_into() {
5945                    (Datum::Int64(n), SqlScalarType::Int64)
5946                } else {
5947                    (
5948                        Datum::Numeric(d),
5949                        SqlScalarType::Numeric { max_scale: None },
5950                    )
5951                }
5952            } else {
5953                (
5954                    Datum::Numeric(d),
5955                    SqlScalarType::Numeric { max_scale: None },
5956                )
5957            }
5958        }
5959        Value::HexString(_) => bail_unsupported!("hex string literals"),
5960        Value::Boolean(b) => match b {
5961            false => (Datum::False, SqlScalarType::Bool),
5962            true => (Datum::True, SqlScalarType::Bool),
5963        },
5964        Value::Interval(i) => {
5965            let i = literal::plan_interval(i)?;
5966            (Datum::Interval(i), SqlScalarType::Interval)
5967        }
5968        Value::String(s) => return Ok(CoercibleScalarExpr::LiteralString(s.clone())),
5969        Value::Null => return Ok(CoercibleScalarExpr::LiteralNull),
5970    };
5971    let expr = HirScalarExpr::literal(datum, scalar_type);
5972    Ok(expr.into())
5973}
5974
5975/// The common part of the planning of non-aggregate window functions, i.e.,
5976/// scalar window functions and value window functions.
5977fn plan_window_function_non_aggr<'a>(
5978    ecx: &ExprContext,
5979    Function {
5980        name,
5981        args,
5982        filter,
5983        over,
5984        distinct,
5985    }: &'a Function<Aug>,
5986) -> Result<
5987    (
5988        bool,
5989        Vec<HirScalarExpr>,
5990        Vec<ColumnOrder>,
5991        mz_expr::WindowFrame,
5992        Vec<HirScalarExpr>,
5993        Vec<CoercibleScalarExpr>,
5994    ),
5995    PlanError,
5996> {
5997    let (ignore_nulls, order_by_exprs, col_orders, window_frame, partition) =
5998        plan_window_function_common(ecx, name, over)?;
5999
6000    if *distinct {
6001        sql_bail!(
6002            "DISTINCT specified, but {} is not an aggregate function",
6003            name
6004        );
6005    }
6006
6007    if filter.is_some() {
6008        bail_unsupported!("FILTER in non-aggregate window functions");
6009    }
6010
6011    let scalar_args = match &args {
6012        FunctionArgs::Star => {
6013            sql_bail!("* argument is invalid with non-aggregate function {}", name)
6014        }
6015        FunctionArgs::Args { args, order_by } => {
6016            if !order_by.is_empty() {
6017                sql_bail!(
6018                    "ORDER BY specified, but {} is not an aggregate function",
6019                    name
6020                );
6021            }
6022            plan_exprs(ecx, args)?
6023        }
6024    };
6025
6026    Ok((
6027        ignore_nulls,
6028        order_by_exprs,
6029        col_orders,
6030        window_frame,
6031        partition,
6032        scalar_args,
6033    ))
6034}
6035
6036/// The common part of the planning of all window functions.
6037fn plan_window_function_common(
6038    ecx: &ExprContext,
6039    name: &<Aug as AstInfo>::ItemName,
6040    over: &Option<WindowSpec<Aug>>,
6041) -> Result<
6042    (
6043        bool,
6044        Vec<HirScalarExpr>,
6045        Vec<ColumnOrder>,
6046        mz_expr::WindowFrame,
6047        Vec<HirScalarExpr>,
6048    ),
6049    PlanError,
6050> {
6051    if !ecx.allow_windows {
6052        sql_bail!(
6053            "window functions are not allowed in {} (function {})",
6054            ecx.name,
6055            name
6056        );
6057    }
6058
6059    let window_spec = match over.as_ref() {
6060        Some(over) => over,
6061        None => sql_bail!("window function {} requires an OVER clause", name),
6062    };
6063    if window_spec.ignore_nulls && window_spec.respect_nulls {
6064        sql_bail!("Both IGNORE NULLS and RESPECT NULLS were given.");
6065    }
6066    let window_frame = match window_spec.window_frame.as_ref() {
6067        Some(frame) => plan_window_frame(frame)?,
6068        None => mz_expr::WindowFrame::default(),
6069    };
6070    let mut partition = Vec::new();
6071    for expr in &window_spec.partition_by {
6072        partition.push(plan_expr(ecx, expr)?.type_as_any(ecx)?);
6073    }
6074
6075    let (order_by_exprs, col_orders) = plan_function_order_by(ecx, &window_spec.order_by)?;
6076
6077    Ok((
6078        window_spec.ignore_nulls,
6079        order_by_exprs,
6080        col_orders,
6081        window_frame,
6082        partition,
6083    ))
6084}
6085
6086fn plan_window_frame(
6087    WindowFrame {
6088        units,
6089        start_bound,
6090        end_bound,
6091    }: &WindowFrame,
6092) -> Result<mz_expr::WindowFrame, PlanError> {
6093    use mz_expr::WindowFrameBound::*;
6094    let units = window_frame_unit_ast_to_expr(units)?;
6095    let start_bound = window_frame_bound_ast_to_expr(start_bound);
6096    let end_bound = end_bound
6097        .as_ref()
6098        .map(window_frame_bound_ast_to_expr)
6099        .unwrap_or(CurrentRow);
6100
6101    // Validate bounds according to Postgres rules
6102    match (&start_bound, &end_bound) {
6103        // Start bound can't be UNBOUNDED FOLLOWING
6104        (UnboundedFollowing, _) => {
6105            sql_bail!("frame start cannot be UNBOUNDED FOLLOWING")
6106        }
6107        // End bound can't be UNBOUNDED PRECEDING
6108        (_, UnboundedPreceding) => {
6109            sql_bail!("frame end cannot be UNBOUNDED PRECEDING")
6110        }
6111        // Start bound should come before end bound in the list of bound definitions
6112        (CurrentRow, OffsetPreceding(_)) => {
6113            sql_bail!("frame starting from current row cannot have preceding rows")
6114        }
6115        (OffsetFollowing(_), OffsetPreceding(_) | CurrentRow) => {
6116            sql_bail!("frame starting from following row cannot have preceding rows")
6117        }
6118        // The above rules are adopted from Postgres.
6119        // The following rules are Materialize-specific.
6120        (OffsetPreceding(o1), OffsetFollowing(o2)) => {
6121            // Note that the only hard limit is that partition size + offset should fit in i64, so
6122            // in theory, we could support much larger offsets than this. But for our current
6123            // performance, even 1000000 is quite big.
6124            if *o1 > 1000000 || *o2 > 1000000 {
6125                sql_bail!("Window frame offsets greater than 1000000 are currently not supported")
6126            }
6127        }
6128        (OffsetPreceding(o1), OffsetPreceding(o2)) => {
6129            if *o1 > 1000000 || *o2 > 1000000 {
6130                sql_bail!("Window frame offsets greater than 1000000 are currently not supported")
6131            }
6132        }
6133        (OffsetFollowing(o1), OffsetFollowing(o2)) => {
6134            if *o1 > 1000000 || *o2 > 1000000 {
6135                sql_bail!("Window frame offsets greater than 1000000 are currently not supported")
6136            }
6137        }
6138        (OffsetPreceding(o), CurrentRow) => {
6139            if *o > 1000000 {
6140                sql_bail!("Window frame offsets greater than 1000000 are currently not supported")
6141            }
6142        }
6143        (CurrentRow, OffsetFollowing(o)) => {
6144            if *o > 1000000 {
6145                sql_bail!("Window frame offsets greater than 1000000 are currently not supported")
6146            }
6147        }
6148        // Other bounds are valid
6149        (_, _) => (),
6150    }
6151
6152    // RANGE is only supported in the default frame
6153    // https://github.com/MaterializeInc/database-issues/issues/6585
6154    if units == mz_expr::WindowFrameUnits::Range
6155        && (start_bound != UnboundedPreceding || end_bound != CurrentRow)
6156    {
6157        bail_unsupported!("RANGE in non-default window frames")
6158    }
6159
6160    let frame = mz_expr::WindowFrame {
6161        units,
6162        start_bound,
6163        end_bound,
6164    };
6165    Ok(frame)
6166}
6167
6168fn window_frame_unit_ast_to_expr(
6169    unit: &WindowFrameUnits,
6170) -> Result<mz_expr::WindowFrameUnits, PlanError> {
6171    match unit {
6172        WindowFrameUnits::Rows => Ok(mz_expr::WindowFrameUnits::Rows),
6173        WindowFrameUnits::Range => Ok(mz_expr::WindowFrameUnits::Range),
6174        WindowFrameUnits::Groups => bail_unsupported!("GROUPS in window frames"),
6175    }
6176}
6177
6178fn window_frame_bound_ast_to_expr(bound: &WindowFrameBound) -> mz_expr::WindowFrameBound {
6179    match bound {
6180        WindowFrameBound::CurrentRow => mz_expr::WindowFrameBound::CurrentRow,
6181        WindowFrameBound::Preceding(None) => mz_expr::WindowFrameBound::UnboundedPreceding,
6182        WindowFrameBound::Preceding(Some(offset)) => {
6183            mz_expr::WindowFrameBound::OffsetPreceding(*offset)
6184        }
6185        WindowFrameBound::Following(None) => mz_expr::WindowFrameBound::UnboundedFollowing,
6186        WindowFrameBound::Following(Some(offset)) => {
6187            mz_expr::WindowFrameBound::OffsetFollowing(*offset)
6188        }
6189    }
6190}
6191
6192pub fn scalar_type_from_sql(
6193    scx: &StatementContext,
6194    data_type: &ResolvedDataType,
6195) -> Result<SqlScalarType, PlanError> {
6196    match data_type {
6197        ResolvedDataType::AnonymousList(elem_type) => {
6198            let elem_type = scalar_type_from_sql(scx, elem_type)?;
6199            if matches!(elem_type, SqlScalarType::Char { .. }) {
6200                bail_unsupported!("char list");
6201            }
6202            Ok(SqlScalarType::List {
6203                element_type: Box::new(elem_type),
6204                custom_id: None,
6205            })
6206        }
6207        ResolvedDataType::AnonymousMap {
6208            key_type,
6209            value_type,
6210        } => {
6211            match scalar_type_from_sql(scx, key_type)? {
6212                SqlScalarType::String => {}
6213                other => sql_bail!(
6214                    "map key type must be {}, got {}",
6215                    scx.humanize_sql_scalar_type(&SqlScalarType::String, false),
6216                    scx.humanize_sql_scalar_type(&other, false)
6217                ),
6218            }
6219            Ok(SqlScalarType::Map {
6220                value_type: Box::new(scalar_type_from_sql(scx, value_type)?),
6221                custom_id: None,
6222            })
6223        }
6224        ResolvedDataType::Named { id, modifiers, .. } => {
6225            scalar_type_from_catalog(scx.catalog, *id, modifiers)
6226        }
6227        ResolvedDataType::Error => bail_internal!("should have been caught in name resolution"),
6228    }
6229}
6230
6231/// Maximum nesting depth of a custom type. A deeper type is rejected rather than
6232/// recursed into, so a long `CREATE TYPE` chain (`l0 <- l1 <- ... <- lN`) cannot
6233/// overflow the stack while resolving.
6234const MAX_TYPE_NESTING_DEPTH: usize = 128;
6235
6236/// Maximum number of sub-type resolutions performed while resolving a single
6237/// custom type. A record whose fields reference the same sub-type produces a
6238/// type tree that is exponential in its depth (fields hold owned copies, not
6239/// shared references), so bound the total work to reject such a type before it
6240/// exhausts memory / CPU rather than after.
6241const MAX_TYPE_RESOLUTION_NODES: usize = 100_000;
6242
6243pub fn scalar_type_from_catalog(
6244    catalog: &dyn SessionCatalog,
6245    id: CatalogItemId,
6246    modifiers: &[i64],
6247) -> Result<SqlScalarType, PlanError> {
6248    let (depth_limit, mut budget) = type_resolution_limits(catalog);
6249    scalar_type_from_catalog_inner(catalog, id, modifiers, 0, depth_limit, &mut budget)
6250}
6251
6252/// The `(nesting-depth, resolution-node)` limits to apply while resolving one
6253/// root custom type. See [`MAX_TYPE_NESTING_DEPTH`] and
6254/// [`MAX_TYPE_RESOLUTION_NODES`] for what each guards against.
6255///
6256/// The limits are lifted while re-planning a persisted catalog item. The
6257/// `unsafe_enable_unbounded_custom_type_resolution` flag signals this, and
6258/// `SystemVars::enable_for_item_parsing` force-enables it during bootstrap. A
6259/// type that an earlier version already accepted resolves to a finite tree, so
6260/// its persisted `create_sql` must keep re-planning after these limits were
6261/// introduced. Rejecting such a grandfathered type during rehydration would
6262/// turn a graceful planning error into a fatal bootstrap panic. A later
6263/// resolution in a normal user session still applies the limits and returns the
6264/// graceful error.
6265fn type_resolution_limits(catalog: &dyn SessionCatalog) -> (usize, usize) {
6266    if catalog
6267        .system_vars()
6268        .unsafe_enable_unbounded_custom_type_resolution()
6269    {
6270        (usize::MAX, usize::MAX)
6271    } else {
6272        (MAX_TYPE_NESTING_DEPTH, MAX_TYPE_RESOLUTION_NODES)
6273    }
6274}
6275
6276/// Bounds the total resolution work while resolving one root custom type into a
6277/// `SqlScalarType`. A single budget must span the entire root type: a record's
6278/// fields, and any containers nested within them, all draw from one shared pool.
6279/// A type that is small field-by-field but enormous in aggregate is therefore
6280/// still rejected. Resetting the budget per field would let a wide record of
6281/// individually-cheap fields resolve into an unbounded type tree and exhaust
6282/// memory, which is the denial of service this bound exists to prevent.
6283///
6284/// Use this when resolving a type that is being assembled from its parts (for
6285/// example at `CREATE TYPE` time, before the root exists in the catalog) so that
6286/// creation-time validation rejects exactly the types a later direct
6287/// [`scalar_type_from_catalog`] call would reject.
6288pub struct TypeResolutionBudget {
6289    /// Sub-type resolutions remaining before the root type is rejected as too
6290    /// complex.
6291    remaining: usize,
6292    /// Nesting depth past which the root type is rejected. Shared by every
6293    /// child so the whole root type is bounded consistently.
6294    depth_limit: usize,
6295}
6296
6297impl TypeResolutionBudget {
6298    /// Creates a budget for resolving one root type, charging the root node
6299    /// itself against the budget. Children resolved through
6300    /// [`TypeResolutionBudget::resolve_child`] begin at nesting depth one,
6301    /// mirroring a direct [`scalar_type_from_catalog`] call. The limits are
6302    /// relaxed for grandfathered persisted items, see [`type_resolution_limits`].
6303    pub fn for_root(catalog: &dyn SessionCatalog) -> TypeResolutionBudget {
6304        let (depth_limit, budget) = type_resolution_limits(catalog);
6305        TypeResolutionBudget {
6306            // The root type counts as one node.
6307            remaining: budget.saturating_sub(1),
6308            depth_limit,
6309        }
6310    }
6311
6312    /// Resolves a type referenced directly by the root (a record field, list
6313    /// element, or map value) into a `SqlScalarType`, drawing from this shared
6314    /// budget so that all such children of one root are bounded together.
6315    pub fn resolve_child(
6316        &mut self,
6317        catalog: &dyn SessionCatalog,
6318        id: CatalogItemId,
6319        modifiers: &[i64],
6320    ) -> Result<SqlScalarType, PlanError> {
6321        scalar_type_from_catalog_inner(
6322            catalog,
6323            id,
6324            modifiers,
6325            1,
6326            self.depth_limit,
6327            &mut self.remaining,
6328        )
6329    }
6330}
6331
6332fn scalar_type_from_catalog_inner(
6333    catalog: &dyn SessionCatalog,
6334    id: CatalogItemId,
6335    modifiers: &[i64],
6336    depth: usize,
6337    depth_limit: usize,
6338    budget: &mut usize,
6339) -> Result<SqlScalarType, PlanError> {
6340    if depth > depth_limit {
6341        sql_bail!("custom type nesting depth exceeds limit of {}", depth_limit);
6342    }
6343    *budget = match budget.checked_sub(1) {
6344        Some(remaining) => remaining,
6345        None => sql_bail!("custom type is too complex to resolve"),
6346    };
6347    let entry = catalog.get_item(&id);
6348    let type_details = match entry.type_details() {
6349        Some(type_details) => type_details,
6350        None => {
6351            // Resolution should never produce a `ResolvedDataType::Named` with
6352            // an ID of a non-type, but we error gracefully just in case.
6353            sql_bail!(
6354                "internal error: {} does not refer to a type",
6355                catalog.resolve_full_name(entry.name()).to_string().quoted()
6356            );
6357        }
6358    };
6359    match &type_details.typ {
6360        CatalogType::Numeric => {
6361            let mut modifiers = modifiers.iter().fuse();
6362            let precision = match modifiers.next() {
6363                Some(p) if *p < 1 || *p > i64::from(NUMERIC_DATUM_MAX_PRECISION) => {
6364                    sql_bail!(
6365                        "precision for type numeric must be between 1 and {}",
6366                        NUMERIC_DATUM_MAX_PRECISION,
6367                    );
6368                }
6369                Some(p) => Some(*p),
6370                None => None,
6371            };
6372            let scale = match modifiers.next() {
6373                Some(scale) => {
6374                    if let Some(precision) = precision {
6375                        if *scale > precision {
6376                            sql_bail!(
6377                                "scale for type numeric must be between 0 and precision {}",
6378                                precision
6379                            );
6380                        }
6381                    }
6382                    Some(NumericMaxScale::try_from(*scale)?)
6383                }
6384                None => None,
6385            };
6386            if modifiers.next().is_some() {
6387                sql_bail!("type numeric supports at most two type modifiers");
6388            }
6389            Ok(SqlScalarType::Numeric { max_scale: scale })
6390        }
6391        CatalogType::Char => {
6392            let mut modifiers = modifiers.iter().fuse();
6393            let length = match modifiers.next() {
6394                Some(l) => Some(CharLength::try_from(*l)?),
6395                None => Some(CharLength::ONE),
6396            };
6397            if modifiers.next().is_some() {
6398                sql_bail!("type character supports at most one type modifier");
6399            }
6400            Ok(SqlScalarType::Char { length })
6401        }
6402        CatalogType::VarChar => {
6403            let mut modifiers = modifiers.iter().fuse();
6404            let length = match modifiers.next() {
6405                Some(l) => Some(VarCharMaxLength::try_from(*l)?),
6406                None => None,
6407            };
6408            if modifiers.next().is_some() {
6409                sql_bail!("type character varying supports at most one type modifier");
6410            }
6411            Ok(SqlScalarType::VarChar { max_length: length })
6412        }
6413        CatalogType::Timestamp => {
6414            let mut modifiers = modifiers.iter().fuse();
6415            let precision = match modifiers.next() {
6416                Some(p) => Some(TimestampPrecision::try_from(*p)?),
6417                None => None,
6418            };
6419            if modifiers.next().is_some() {
6420                sql_bail!("type timestamp supports at most one type modifier");
6421            }
6422            Ok(SqlScalarType::Timestamp { precision })
6423        }
6424        CatalogType::TimestampTz => {
6425            let mut modifiers = modifiers.iter().fuse();
6426            let precision = match modifiers.next() {
6427                Some(p) => Some(TimestampPrecision::try_from(*p)?),
6428                None => None,
6429            };
6430            if modifiers.next().is_some() {
6431                sql_bail!("type timestamp with time zone supports at most one type modifier");
6432            }
6433            Ok(SqlScalarType::TimestampTz { precision })
6434        }
6435        t => {
6436            if !modifiers.is_empty() {
6437                sql_bail!(
6438                    "{} does not support type modifiers",
6439                    catalog.resolve_full_name(entry.name()).to_string()
6440                );
6441            }
6442            match t {
6443                CatalogType::Array {
6444                    element_reference: element_id,
6445                } => Ok(SqlScalarType::Array(Box::new(
6446                    scalar_type_from_catalog_inner(
6447                        catalog,
6448                        *element_id,
6449                        modifiers,
6450                        depth + 1,
6451                        depth_limit,
6452                        budget,
6453                    )?,
6454                ))),
6455                CatalogType::List {
6456                    element_reference: element_id,
6457                    element_modifiers,
6458                } => Ok(SqlScalarType::List {
6459                    element_type: Box::new(scalar_type_from_catalog_inner(
6460                        catalog,
6461                        *element_id,
6462                        element_modifiers,
6463                        depth + 1,
6464                        depth_limit,
6465                        budget,
6466                    )?),
6467                    custom_id: Some(id),
6468                }),
6469                CatalogType::Map {
6470                    key_reference: _,
6471                    key_modifiers: _,
6472                    value_reference: value_id,
6473                    value_modifiers,
6474                } => Ok(SqlScalarType::Map {
6475                    value_type: Box::new(scalar_type_from_catalog_inner(
6476                        catalog,
6477                        *value_id,
6478                        value_modifiers,
6479                        depth + 1,
6480                        depth_limit,
6481                        budget,
6482                    )?),
6483                    custom_id: Some(id),
6484                }),
6485                CatalogType::Range {
6486                    element_reference: element_id,
6487                } => Ok(SqlScalarType::Range {
6488                    element_type: Box::new(scalar_type_from_catalog_inner(
6489                        catalog,
6490                        *element_id,
6491                        &[],
6492                        depth + 1,
6493                        depth_limit,
6494                        budget,
6495                    )?),
6496                }),
6497                CatalogType::Record { fields } => {
6498                    let scalars: Box<[(ColumnName, SqlColumnType)]> = fields
6499                        .iter()
6500                        .map(|f| {
6501                            let scalar_type = scalar_type_from_catalog_inner(
6502                                catalog,
6503                                f.type_reference,
6504                                &f.type_modifiers,
6505                                depth + 1,
6506                                depth_limit,
6507                                budget,
6508                            )?;
6509                            Ok((
6510                                f.name.clone(),
6511                                SqlColumnType {
6512                                    scalar_type,
6513                                    nullable: true,
6514                                },
6515                            ))
6516                        })
6517                        .collect::<Result<Box<_>, PlanError>>()?;
6518                    Ok(SqlScalarType::Record {
6519                        fields: scalars,
6520                        custom_id: Some(id),
6521                    })
6522                }
6523                CatalogType::AclItem => Ok(SqlScalarType::AclItem),
6524                CatalogType::Bool => Ok(SqlScalarType::Bool),
6525                CatalogType::Bytes => Ok(SqlScalarType::Bytes),
6526                CatalogType::Date => Ok(SqlScalarType::Date),
6527                CatalogType::Float32 => Ok(SqlScalarType::Float32),
6528                CatalogType::Float64 => Ok(SqlScalarType::Float64),
6529                CatalogType::Int16 => Ok(SqlScalarType::Int16),
6530                CatalogType::Int32 => Ok(SqlScalarType::Int32),
6531                CatalogType::Int64 => Ok(SqlScalarType::Int64),
6532                CatalogType::UInt16 => Ok(SqlScalarType::UInt16),
6533                CatalogType::UInt32 => Ok(SqlScalarType::UInt32),
6534                CatalogType::UInt64 => Ok(SqlScalarType::UInt64),
6535                CatalogType::MzTimestamp => Ok(SqlScalarType::MzTimestamp),
6536                CatalogType::Interval => Ok(SqlScalarType::Interval),
6537                CatalogType::Jsonb => Ok(SqlScalarType::Jsonb),
6538                CatalogType::Oid => Ok(SqlScalarType::Oid),
6539                CatalogType::PgLegacyChar => Ok(SqlScalarType::PgLegacyChar),
6540                CatalogType::PgLegacyName => Ok(SqlScalarType::PgLegacyName),
6541                CatalogType::Pseudo => {
6542                    sql_bail!(
6543                        "cannot reference pseudo type {}",
6544                        catalog.resolve_full_name(entry.name()).to_string()
6545                    )
6546                }
6547                CatalogType::RegClass => Ok(SqlScalarType::RegClass),
6548                CatalogType::RegProc => Ok(SqlScalarType::RegProc),
6549                CatalogType::RegType => Ok(SqlScalarType::RegType),
6550                CatalogType::String => Ok(SqlScalarType::String),
6551                CatalogType::Time => Ok(SqlScalarType::Time),
6552                CatalogType::Uuid => Ok(SqlScalarType::Uuid),
6553                CatalogType::Int2Vector => Ok(SqlScalarType::Int2Vector),
6554                CatalogType::MzAclItem => Ok(SqlScalarType::MzAclItem),
6555                CatalogType::Numeric => unreachable!("handled above"),
6556                CatalogType::Char => unreachable!("handled above"),
6557                CatalogType::VarChar => unreachable!("handled above"),
6558                CatalogType::Timestamp => unreachable!("handled above"),
6559                CatalogType::TimestampTz => unreachable!("handled above"),
6560            }
6561        }
6562    }
6563}
6564
6565/// This is used to collect aggregates and table functions from within an `Expr`.
6566/// See the explanation of aggregate handling at the top of the file for more details.
6567struct AggregateTableFuncVisitor<'a> {
6568    scx: &'a StatementContext<'a>,
6569    aggs: Vec<Function<Aug>>,
6570    within_aggregate: bool,
6571    tables: BTreeMap<Function<Aug>, String>,
6572    table_disallowed_context: Vec<&'static str>,
6573    in_select_item: bool,
6574    id_gen: IdGen,
6575    err: Option<PlanError>,
6576}
6577
6578impl<'a> AggregateTableFuncVisitor<'a> {
6579    fn new(scx: &'a StatementContext<'a>) -> AggregateTableFuncVisitor<'a> {
6580        AggregateTableFuncVisitor {
6581            scx,
6582            aggs: Vec::new(),
6583            within_aggregate: false,
6584            tables: BTreeMap::new(),
6585            table_disallowed_context: Vec::new(),
6586            in_select_item: false,
6587            id_gen: Default::default(),
6588            err: None,
6589        }
6590    }
6591
6592    fn into_result(
6593        self,
6594    ) -> Result<(Vec<Function<Aug>>, BTreeMap<Function<Aug>, String>), PlanError> {
6595        match self.err {
6596            Some(err) => Err(err),
6597            None => {
6598                // Dedup while preserving the order. We don't care what the order is, but it
6599                // has to be reproducible so that EXPLAIN PLAN tests work.
6600                let mut seen = BTreeSet::new();
6601                let aggs = self
6602                    .aggs
6603                    .into_iter()
6604                    .filter(move |agg| seen.insert(agg.clone()))
6605                    .collect();
6606                Ok((aggs, self.tables))
6607            }
6608        }
6609    }
6610}
6611
6612impl<'a> VisitMut<'_, Aug> for AggregateTableFuncVisitor<'a> {
6613    fn visit_function_mut(&mut self, func: &mut Function<Aug>) {
6614        let item = match self.scx.get_item_by_resolved_name(&func.name) {
6615            Ok(i) => i,
6616            // Catching missing functions later in planning improves error messages.
6617            Err(_) => return,
6618        };
6619
6620        match item.func() {
6621            // We don't want to collect window aggregations, because these will be handled not by
6622            // plan_aggregate, but by plan_function.
6623            Ok(Func::Aggregate { .. }) if func.over.is_none() => {
6624                if self.within_aggregate {
6625                    self.err = Some(sql_err!("nested aggregate functions are not allowed",));
6626                    return;
6627                }
6628                self.aggs.push(func.clone());
6629                let Function {
6630                    name: _,
6631                    args,
6632                    filter,
6633                    over: _,
6634                    distinct: _,
6635                } = func;
6636                if let Some(filter) = filter {
6637                    self.visit_expr_mut(filter);
6638                }
6639                let old_within_aggregate = self.within_aggregate;
6640                self.within_aggregate = true;
6641                self.table_disallowed_context
6642                    .push("aggregate function calls");
6643
6644                self.visit_function_args_mut(args);
6645
6646                self.within_aggregate = old_within_aggregate;
6647                self.table_disallowed_context.pop();
6648            }
6649            Ok(Func::Table { .. }) => {
6650                self.table_disallowed_context.push("other table functions");
6651                visit_mut::visit_function_mut(self, func);
6652                self.table_disallowed_context.pop();
6653            }
6654            _ => visit_mut::visit_function_mut(self, func),
6655        }
6656    }
6657
6658    fn visit_query_mut(&mut self, _query: &mut Query<Aug>) {
6659        // Don't go into subqueries.
6660    }
6661
6662    fn visit_expr_mut(&mut self, expr: &mut Expr<Aug>) {
6663        let (disallowed_context, func) = match expr {
6664            Expr::Case { .. } => (Some("CASE"), None),
6665            Expr::HomogenizingFunction {
6666                function: HomogenizingFunction::Coalesce,
6667                ..
6668            } => (Some("COALESCE"), None),
6669            Expr::Function(func) if self.in_select_item => {
6670                // If we're in a SELECT list, replace table functions with a uuid identifier
6671                // and save the table func so it can be planned elsewhere.
6672                let mut table_func = None;
6673                if let Ok(item) = self.scx.get_item_by_resolved_name(&func.name) {
6674                    if let Ok(Func::Table { .. }) = item.func() {
6675                        if let Some(context) = self.table_disallowed_context.last() {
6676                            self.err = Some(sql_err!(
6677                                "table functions are not allowed in {} (function {})",
6678                                context,
6679                                func.name
6680                            ));
6681                            return;
6682                        }
6683                        table_func = Some(func.clone());
6684                    }
6685                }
6686                // Since we will descend into the table func below, don't add its own disallow
6687                // context here, instead use visit_function to set that.
6688                (None, table_func)
6689            }
6690            _ => (None, None),
6691        };
6692        if let Some(func) = func {
6693            // Since we are trading out expr, we need to visit the table func here.
6694            visit_mut::visit_expr_mut(self, expr);
6695            // Don't attempt to replace table functions with unsupported syntax.
6696            if let Function {
6697                name: _,
6698                args: _,
6699                filter: None,
6700                over: None,
6701                distinct: false,
6702            } = &func
6703            {
6704                // Identical table functions can be de-duplicated.
6705                let unique_id = self.id_gen.allocate_id();
6706                let id = self
6707                    .tables
6708                    .entry(func)
6709                    .or_insert_with(|| format!("table_func_{unique_id}"));
6710                // We know this is okay because id is is 11 characters + <=20 characters, which is
6711                // less than our max length.
6712                *expr = Expr::Identifier(vec![Ident::new_unchecked(id.clone())]);
6713            }
6714        }
6715        if let Some(context) = disallowed_context {
6716            self.table_disallowed_context.push(context);
6717        }
6718
6719        visit_mut::visit_expr_mut(self, expr);
6720
6721        if disallowed_context.is_some() {
6722            self.table_disallowed_context.pop();
6723        }
6724    }
6725
6726    fn visit_select_item_mut(&mut self, si: &mut SelectItem<Aug>) {
6727        let old = self.in_select_item;
6728        self.in_select_item = true;
6729        visit_mut::visit_select_item_mut(self, si);
6730        self.in_select_item = old;
6731    }
6732}
6733
6734#[derive(Default)]
6735struct WindowFuncCollector {
6736    window_funcs: Vec<Expr<Aug>>,
6737}
6738
6739impl WindowFuncCollector {
6740    fn into_result(self) -> Vec<Expr<Aug>> {
6741        // Dedup while preserving the order.
6742        let mut seen = BTreeSet::new();
6743        let window_funcs_dedupped = self
6744            .window_funcs
6745            .into_iter()
6746            .filter(move |expr| seen.insert(expr.clone()))
6747            // Reverse the order, so that in case of a nested window function call, the
6748            // inner one is evaluated first.
6749            .rev()
6750            .collect();
6751        window_funcs_dedupped
6752    }
6753}
6754
6755impl Visit<'_, Aug> for WindowFuncCollector {
6756    fn visit_expr(&mut self, expr: &Expr<Aug>) {
6757        match expr {
6758            Expr::Function(func) => {
6759                if func.over.is_some() {
6760                    self.window_funcs.push(expr.clone());
6761                }
6762            }
6763            _ => (),
6764        }
6765        visit::visit_expr(self, expr);
6766    }
6767
6768    fn visit_query(&mut self, _query: &Query<Aug>) {
6769        // Don't go into subqueries. Those will be handled by their own `plan_query`.
6770    }
6771}
6772
6773/// Specifies how long a query will live.
6774#[derive(Debug, Eq, PartialEq, Copy, Clone)]
6775pub enum QueryLifetime {
6776    /// The query's (or the expression's) result will be computed at one point in time.
6777    OneShot,
6778    /// The query (or expression) is used in a dataflow that maintains an index.
6779    Index,
6780    /// The query (or expression) is used in a dataflow that maintains a materialized view.
6781    MaterializedView,
6782    /// The query (or expression) is used in a dataflow that maintains a SUBSCRIBE.
6783    Subscribe,
6784    /// The query (or expression) is part of a (non-materialized) view.
6785    View,
6786    /// The expression is part of a source definition.
6787    Source,
6788}
6789
6790impl QueryLifetime {
6791    /// (This used to impact whether the query is allowed to reason about the time at which it is
6792    /// running, e.g., by calling the `now()` function. Nowadays, this is decided by a different
6793    /// mechanism, see `ExprPrepStyle`.)
6794    pub fn is_one_shot(&self) -> bool {
6795        let result = match self {
6796            QueryLifetime::OneShot => true,
6797            QueryLifetime::Index => false,
6798            QueryLifetime::MaterializedView => false,
6799            QueryLifetime::Subscribe => false,
6800            QueryLifetime::View => false,
6801            QueryLifetime::Source => false,
6802        };
6803        assert_eq!(!result, self.is_maintained());
6804        result
6805    }
6806
6807    /// Maintained dataflows can't have a finishing applied directly. Therefore, the finishing is
6808    /// turned into a `TopK`.
6809    pub fn is_maintained(&self) -> bool {
6810        match self {
6811            QueryLifetime::OneShot => false,
6812            QueryLifetime::Index => true,
6813            QueryLifetime::MaterializedView => true,
6814            QueryLifetime::Subscribe => true,
6815            QueryLifetime::View => true,
6816            QueryLifetime::Source => true,
6817        }
6818    }
6819
6820    /// Most maintained dataflows don't allow SHOW commands currently. However, SUBSCRIBE does.
6821    pub fn allow_show(&self) -> bool {
6822        match self {
6823            QueryLifetime::OneShot => true,
6824            QueryLifetime::Index => false,
6825            QueryLifetime::MaterializedView => false,
6826            QueryLifetime::Subscribe => true, // SUBSCRIBE allows SHOW commands!
6827            QueryLifetime::View => false,
6828            QueryLifetime::Source => false,
6829        }
6830    }
6831}
6832
6833/// Description of a CTE sufficient for query planning.
6834#[derive(Debug, Clone)]
6835pub struct CteDesc {
6836    pub name: String,
6837    pub desc: RelationDesc,
6838}
6839
6840/// The state required when planning a `Query`.
6841#[derive(Debug, Clone)]
6842pub struct QueryContext<'a> {
6843    /// The context for the containing `Statement`.
6844    pub scx: &'a StatementContext<'a>,
6845    /// The lifetime that the planned query will have.
6846    pub lifetime: QueryLifetime,
6847    /// The scopes of the outer relation expression.
6848    pub outer_scopes: Vec<Scope>,
6849    /// The type of the outer relation expressions.
6850    pub outer_relation_types: Vec<SqlRelationType>,
6851    /// CTEs for this query, mapping their assigned LocalIds to their definition.
6852    pub ctes: BTreeMap<LocalId, CteDesc>,
6853    /// A name manager, for interning column names that will be stored in HIR and MIR.
6854    pub name_manager: Rc<RefCell<NameManager>>,
6855    pub recursion_guard: RecursionGuard,
6856}
6857
6858impl CheckedRecursion for QueryContext<'_> {
6859    fn recursion_guard(&self) -> &RecursionGuard {
6860        &self.recursion_guard
6861    }
6862}
6863
6864impl<'a> QueryContext<'a> {
6865    pub fn root(scx: &'a StatementContext, lifetime: QueryLifetime) -> QueryContext<'a> {
6866        QueryContext {
6867            scx,
6868            lifetime,
6869            outer_scopes: vec![],
6870            outer_relation_types: vec![],
6871            ctes: BTreeMap::new(),
6872            name_manager: Rc::new(RefCell::new(NameManager::new())),
6873            recursion_guard: RecursionGuard::with_limit(1024), // chosen arbitrarily
6874        }
6875    }
6876
6877    fn relation_type(&self, expr: &HirRelationExpr) -> SqlRelationType {
6878        expr.typ(&self.outer_relation_types, &self.scx.param_types.borrow())
6879    }
6880
6881    /// Generate a new `QueryContext` appropriate to be used in subqueries of
6882    /// `self`.
6883    fn derived_context(&self, scope: Scope, relation_type: SqlRelationType) -> QueryContext<'a> {
6884        let ctes = self.ctes.clone();
6885        let outer_scopes = iter::once(scope).chain(self.outer_scopes.clone()).collect();
6886        let outer_relation_types = iter::once(relation_type)
6887            .chain(self.outer_relation_types.clone())
6888            .collect();
6889        // These shenanigans are simpler than adding `&mut NameManager` arguments everywhere.
6890        let name_manager = Rc::clone(&self.name_manager);
6891
6892        QueryContext {
6893            scx: self.scx,
6894            lifetime: self.lifetime,
6895            outer_scopes,
6896            outer_relation_types,
6897            ctes,
6898            name_manager,
6899            recursion_guard: self.recursion_guard.clone(),
6900        }
6901    }
6902
6903    /// Derives a `QueryContext` for a scope that contains no columns.
6904    fn empty_derived_context(&self) -> QueryContext<'a> {
6905        let scope = Scope::empty();
6906        let ty = SqlRelationType::empty();
6907        self.derived_context(scope, ty)
6908    }
6909
6910    /// Resolves `object` to a table expr, i.e. creating a `Get` or inlining a
6911    /// CTE.
6912    pub fn resolve_table_name(
6913        &self,
6914        object: ResolvedItemName,
6915    ) -> Result<(HirRelationExpr, Scope), PlanError> {
6916        match object {
6917            ResolvedItemName::Item {
6918                id,
6919                full_name,
6920                version,
6921                ..
6922            } => {
6923                let item = self.scx.get_item(&id).at_version(version);
6924                let desc = match item.relation_desc() {
6925                    Some(desc) => desc.clone(),
6926                    None => {
6927                        return Err(PlanError::InvalidDependency {
6928                            name: full_name.to_string(),
6929                            item_type: item.item_type().to_string(),
6930                        });
6931                    }
6932                };
6933                let expr = HirRelationExpr::Get {
6934                    id: Id::Global(item.global_id()),
6935                    typ: desc.typ().clone(),
6936                };
6937
6938                let name = full_name.into();
6939                let scope = Scope::from_source(Some(name), desc.iter_names().cloned());
6940
6941                Ok((expr, scope))
6942            }
6943            ResolvedItemName::Cte { id, name } => {
6944                let name = name.into();
6945                let cte = self.ctes.get(&id).unwrap();
6946                let expr = HirRelationExpr::Get {
6947                    id: Id::Local(id),
6948                    typ: cte.desc.typ().clone(),
6949                };
6950
6951                let scope = Scope::from_source(Some(name), cte.desc.iter_names());
6952
6953                Ok((expr, scope))
6954            }
6955            ResolvedItemName::Error => bail_internal!("should have been caught in name resolution"),
6956        }
6957    }
6958
6959    /// The returned String is more detailed when the `postgres_compat` flag is not set. However,
6960    /// the flag should be set in, e.g., the implementation of the `pg_typeof` function.
6961    pub fn humanize_sql_scalar_type(&self, typ: &SqlScalarType, postgres_compat: bool) -> String {
6962        self.scx.humanize_sql_scalar_type(typ, postgres_compat)
6963    }
6964}
6965
6966/// A bundle of unrelated things that we need for planning `Expr`s.
6967#[derive(Debug, Clone)]
6968pub struct ExprContext<'a> {
6969    pub qcx: &'a QueryContext<'a>,
6970    /// The name of this kind of expression eg "WHERE clause". Used only for error messages.
6971    pub name: &'a str,
6972    /// The context for the `Query` that contains this `Expr`.
6973    /// The current scope.
6974    pub scope: &'a Scope,
6975    /// The type of the current relation expression upon which this scalar
6976    /// expression will be evaluated.
6977    pub relation_type: &'a SqlRelationType,
6978    /// Are aggregate functions allowed in this context
6979    pub allow_aggregates: bool,
6980    /// Are subqueries allowed in this context
6981    pub allow_subqueries: bool,
6982    /// Are parameters allowed in this context.
6983    pub allow_parameters: bool,
6984    /// Are window functions allowed in this context
6985    pub allow_windows: bool,
6986}
6987
6988impl CheckedRecursion for ExprContext<'_> {
6989    fn recursion_guard(&self) -> &RecursionGuard {
6990        &self.qcx.recursion_guard
6991    }
6992}
6993
6994impl<'a> ExprContext<'a> {
6995    pub fn catalog(&self) -> &dyn SessionCatalog {
6996        self.qcx.scx.catalog
6997    }
6998
6999    pub fn with_name(&self, name: &'a str) -> ExprContext<'a> {
7000        let mut ecx = self.clone();
7001        ecx.name = name;
7002        ecx
7003    }
7004
7005    pub fn column_type<E>(&self, expr: &E) -> E::Type
7006    where
7007        E: AbstractExpr,
7008    {
7009        expr.typ(
7010            &self.qcx.outer_relation_types,
7011            self.relation_type,
7012            &self.qcx.scx.param_types.borrow(),
7013        )
7014    }
7015
7016    pub fn scalar_type<E>(&self, expr: &E) -> <E::Type as AbstractColumnType>::AbstractScalarType
7017    where
7018        E: AbstractExpr,
7019    {
7020        self.column_type(expr).scalar_type()
7021    }
7022
7023    fn derived_query_context(&self) -> QueryContext<'_> {
7024        let mut scope = self.scope.clone();
7025        scope.lateral_barrier = true;
7026        self.qcx.derived_context(scope, self.relation_type.clone())
7027    }
7028
7029    pub fn require_feature_flag(&self, flag: &'static FeatureFlag) -> Result<(), PlanError> {
7030        self.qcx.scx.require_feature_flag(flag)
7031    }
7032
7033    pub fn param_types(&self) -> &RefCell<BTreeMap<usize, SqlScalarType>> {
7034        &self.qcx.scx.param_types
7035    }
7036
7037    /// The returned String is more detailed when the `postgres_compat` flag is not set. However,
7038    /// the flag should be set in, e.g., the implementation of the `pg_typeof` function.
7039    pub fn humanize_sql_scalar_type(&self, typ: &SqlScalarType, postgres_compat: bool) -> String {
7040        self.qcx.scx.humanize_sql_scalar_type(typ, postgres_compat)
7041    }
7042
7043    pub fn intern(&self, item: &ScopeItem) -> Arc<str> {
7044        self.qcx.name_manager.borrow_mut().intern_scope_item(item)
7045    }
7046}
7047
7048/// Manages column names, doing lightweight string internment.
7049///
7050/// Names are stored in `HirScalarExpr` and `MirScalarExpr` using
7051/// `Option<Arc<str>>`; we use the `NameManager` when lowering from SQL to HIR
7052/// to ensure maximal sharing.
7053#[derive(Debug, Clone)]
7054pub struct NameManager(BTreeSet<Arc<str>>);
7055
7056impl NameManager {
7057    /// Creates a new `NameManager`, with no interned names
7058    pub fn new() -> Self {
7059        Self(BTreeSet::new())
7060    }
7061
7062    /// Interns a string, returning a reference-counted pointer to the interned
7063    /// string.
7064    fn intern<S: AsRef<str>>(&mut self, s: S) -> Arc<str> {
7065        let s = s.as_ref();
7066        if let Some(interned) = self.0.get(s) {
7067            Arc::clone(interned)
7068        } else {
7069            let interned: Arc<str> = Arc::from(s);
7070            self.0.insert(Arc::clone(&interned));
7071            interned
7072        }
7073    }
7074
7075    /// Interns a string representing a reference to a `ScopeItem`, returning a
7076    /// reference-counted pointer to the interned string.
7077    pub fn intern_scope_item(&mut self, item: &ScopeItem) -> Arc<str> {
7078        // TODO(mgree): extracting the table name from `item` leads to an issue with the catalog
7079        //
7080        // After an `ALTER ... RENAME` on a table, the catalog will have out-of-date
7081        // name information. Note that as of 2025-04-09, we don't support column
7082        // renames.
7083        //
7084        // A few bad alternatives:
7085        //
7086        // (1) Store it but don't write it down. This fails because the expression
7087        //     cache will erase our names on restart.
7088        // (2) When `ALTER ... RENAME` is run, re-optimize all downstream objects to
7089        //     get the right names. But the world now and the world when we made
7090        //     those objects may be different.
7091        // (3) Just don't write down the table name. Nothing fails... for now.
7092
7093        self.intern(item.column_name.as_str())
7094    }
7095}
7096
7097#[cfg(test)]
7098mod test {
7099    use super::*;
7100
7101    /// Ensure that `NameManager`'s string interning works as expected.
7102    ///
7103    /// In particular, structurally but not referentially identical strings should
7104    /// be interned to the same `Arc`ed pointer.
7105    #[mz_ore::test]
7106    pub fn test_name_manager_string_interning() {
7107        let mut nm = NameManager::new();
7108
7109        let orig_hi = "hi";
7110        let hi = nm.intern(orig_hi);
7111        let hello = nm.intern("hello");
7112
7113        assert_ne!(hi.as_ptr(), hello.as_ptr());
7114
7115        // this static string is _likely_ the same as `orig_hi``
7116        let hi2 = nm.intern("hi");
7117        assert_eq!(hi.as_ptr(), hi2.as_ptr());
7118
7119        // generate a "hi" string that doesn't get optimized to the same static string
7120        let s = format!(
7121            "{}{}",
7122            hi.chars().nth(0).unwrap(),
7123            hi2.chars().nth(1).unwrap()
7124        );
7125        // make sure that we're testing with a fresh string!
7126        assert_ne!(orig_hi.as_ptr(), s.as_ptr());
7127
7128        let hi3 = nm.intern(s);
7129        assert_eq!(hi.as_ptr(), hi3.as_ptr());
7130    }
7131}