Skip to main content

mz_expr/
relation.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#![warn(missing_docs)]
11
12use std::cell::RefCell;
13use std::cmp::{Ordering, max};
14use std::collections::{BTreeMap, BTreeSet};
15use std::fmt;
16use std::fmt::{Display, Formatter};
17use std::hash::{DefaultHasher, Hash, Hasher};
18use std::num::NonZeroU64;
19use std::time::Instant;
20
21use bytesize::ByteSize;
22use columnation::{Columnation, CopyRegion};
23use itertools::Itertools;
24use mz_ore::cast::{CastFrom, CastInto};
25use mz_ore::collections::CollectionExt;
26use mz_ore::id_gen::IdGen;
27use mz_ore::metrics::Histogram;
28use mz_ore::num::NonNeg;
29use mz_ore::soft_assert_no_log;
30use mz_ore::stack::RecursionLimitError;
31use mz_ore::str::Indent;
32use mz_repr::adt::numeric::NumericMaxScale;
33use mz_repr::explain::text::text_string_at;
34use mz_repr::explain::{
35    DummyHumanizer, ExplainConfig, ExprHumanizer, IndexUsageType, PlanRenderingContext,
36};
37use mz_repr::{
38    ColumnName, Datum, DatumVec, Diff, GlobalId, IntoRowIterator, ReprColumnType, ReprRelationType,
39    ReprScalarType, Row, RowIterator, RowRef, SqlColumnType, SqlRelationType, SqlScalarType,
40};
41use serde::{Deserialize, Serialize};
42
43use crate::Id::Local;
44use crate::explain::{HumanizedExpr, HumanizerMode};
45use crate::relation::func::{AggregateFunc, LagLeadType, TableFunc};
46use crate::row::{RowCollection, RowCollectionIter};
47use crate::scalar::columns::Columns;
48use crate::scalar::func::variadic::{
49    JsonbBuildArray, JsonbBuildObject, ListCreate, ListIndex, MapBuild, RecordCreate,
50};
51use crate::visit::{Visit, VisitChildren};
52use crate::{
53    EvalError, FilterCharacteristics, Id, LocalId, MirScalarExpr, UnaryFunc, func as scalar_func,
54};
55
56pub mod canonicalize;
57pub mod func;
58pub mod join_input_mapper;
59
60/// A recursion limit to be used for stack-safe traversals of [`MirRelationExpr`] trees.
61///
62/// The recursion limit must be large enough to accommodate for the linear representation
63/// of some pathological but frequently occurring query fragments.
64///
65/// For example, in MIR we could have long chains of
66/// - (1) `Let` bindings,
67/// - (2) `CallBinary` calls with associative functions such as `+`
68///
69/// Until we fix those, we need to stick with the larger recursion limit.
70pub const RECURSION_LIMIT: usize = 2048;
71
72/// A trait for types that describe how to build a collection.
73pub trait CollectionPlan {
74    /// Collects the set of global identifiers from dataflows referenced in Get.
75    fn depends_on_into(&self, out: &mut BTreeSet<GlobalId>);
76
77    /// Returns the set of global identifiers from dataflows referenced in Get.
78    ///
79    /// See [`CollectionPlan::depends_on_into`] to reuse an existing `BTreeSet`.
80    fn depends_on(&self) -> BTreeSet<GlobalId> {
81        let mut out = BTreeSet::new();
82        self.depends_on_into(&mut out);
83        out
84    }
85}
86
87/// An abstract syntax tree which defines a collection.
88///
89/// The AST is meant to reflect the capabilities of the `differential_dataflow::Collection` type,
90/// written generically enough to avoid run-time compilation work.
91///
92/// `derived_hash_with_manual_eq` was complaining for the wrong reason: This lint exists because
93/// it's bad when `Eq` doesn't agree with `Hash`, which is often quite likely if one of them is
94/// implemented manually. However, our manual implementation of `Eq` _will_ agree with the derived
95/// one. This is because the reason for the manual implementation is not to change the semantics
96/// from the derived one, but to avoid stack overflows.
97#[allow(clippy::derived_hash_with_manual_eq)]
98#[derive(Clone, Debug, Ord, PartialOrd, Serialize, Deserialize, Hash)]
99pub enum MirRelationExpr {
100    /// A constant relation containing specified rows.
101    ///
102    /// The runtime memory footprint of this operator is zero.
103    ///
104    /// When you would like to pattern match on this, consider using `MirRelationExpr::as_const`
105    /// instead, which looks behind `ArrangeBy`s. You might want this matching behavior because
106    /// constant folding doesn't remove `ArrangeBy`s.
107    Constant {
108        /// Rows of the constant collection and their multiplicities.
109        rows: Result<Vec<(Row, Diff)>, EvalError>,
110        /// Schema of the collection.
111        typ: ReprRelationType,
112    },
113    /// Get an existing dataflow.
114    ///
115    /// The runtime memory footprint of this operator is zero.
116    Get {
117        /// The identifier for the collection to load.
118        id: Id,
119        /// Schema of the collection.
120        typ: ReprRelationType,
121        /// If this is a global Get, this will indicate whether we are going to read from Persist or
122        /// from an index, or from a different object in `objects_to_build`. If it's an index, then
123        /// how downstream dataflow operations will use this index is also recorded. This is filled
124        /// by `prune_and_annotate_dataflow_index_imports`. Note that this is not used by the
125        /// lowering to LIR, but is used only by EXPLAIN.
126        access_strategy: AccessStrategy,
127    },
128    /// Introduce a temporary dataflow.
129    ///
130    /// The runtime memory footprint of this operator is zero.
131    Let {
132        /// The identifier to be used in `Get` variants to retrieve `value`.
133        id: LocalId,
134        /// The collection to be bound to `id`.
135        value: Box<MirRelationExpr>,
136        /// The result of the `Let`, evaluated with `id` bound to `value`.
137        body: Box<MirRelationExpr>,
138    },
139    /// Introduce mutually recursive bindings.
140    ///
141    /// Each `LocalId` is immediately bound to an initially empty  collection
142    /// with the type of its corresponding `MirRelationExpr`. Repeatedly, each
143    /// binding is evaluated using the current contents of each other binding,
144    /// and is refreshed to contain the new evaluation. This process continues
145    /// through all bindings, and repeats as long as changes continue to occur.
146    ///
147    /// The resulting value of the expression is `body` evaluated once in the
148    /// context of the final iterates.
149    ///
150    /// A zero-binding instance can be replaced by `body`.
151    /// A single-binding instance is equivalent to `MirRelationExpr::Let`.
152    ///
153    /// The runtime memory footprint of this operator is zero.
154    LetRec {
155        /// The identifiers to be used in `Get` variants to retrieve each `value`.
156        ids: Vec<LocalId>,
157        /// The collections to be bound to each `id`.
158        values: Vec<MirRelationExpr>,
159        /// Maximum number of iterations, after which we should artificially force a fixpoint.
160        /// (Whether we error or just stop is configured by `LetRecLimit::return_at_limit`.)
161        /// The per-`LetRec` limit that the user specified is initially copied to each binding to
162        /// accommodate slicing and merging of `LetRec`s in MIR transforms (e.g., `NormalizeLets`).
163        limits: Vec<Option<LetRecLimit>>,
164        /// The result of the `Let`, evaluated with `id` bound to `value`.
165        body: Box<MirRelationExpr>,
166    },
167    /// Project out some columns from a dataflow
168    ///
169    /// The runtime memory footprint of this operator is zero.
170    Project {
171        /// The source collection.
172        input: Box<MirRelationExpr>,
173        /// Indices of columns to retain.
174        outputs: Vec<usize>,
175    },
176    /// Append new columns to a dataflow
177    ///
178    /// The runtime memory footprint of this operator is zero.
179    Map {
180        /// The source collection.
181        input: Box<MirRelationExpr>,
182        /// Expressions which determine values to append to each row.
183        /// An expression may refer to columns in `input` or
184        /// expressions defined earlier in the vector
185        scalars: Vec<MirScalarExpr>,
186    },
187    /// Like Map, but yields zero-or-more output rows per input row
188    ///
189    /// The runtime memory footprint of this operator is zero.
190    FlatMap {
191        /// The source collection
192        input: Box<MirRelationExpr>,
193        /// The table func to apply
194        func: TableFunc,
195        /// The argument to the table func
196        exprs: Vec<MirScalarExpr>,
197    },
198    /// Keep rows from a dataflow where all the predicates are true
199    ///
200    /// The runtime memory footprint of this operator is zero.
201    Filter {
202        /// The source collection.
203        input: Box<MirRelationExpr>,
204        /// Predicates, each of which must be true.
205        predicates: Vec<MirScalarExpr>,
206    },
207    /// Join several collections, where some columns must be equal.
208    ///
209    /// For further details consult the documentation for [`MirRelationExpr::join`].
210    ///
211    /// The runtime memory footprint of this operator can be proportional to
212    /// the sizes of all inputs and the size of all joins of prefixes.
213    /// This may be reduced due to arrangements available at rendering time.
214    Join {
215        /// A sequence of input relations.
216        inputs: Vec<MirRelationExpr>,
217        /// A sequence of equivalence classes of expressions on the cross product of inputs.
218        ///
219        /// Each equivalence class is a list of scalar expressions, where for each class the
220        /// intended interpretation is that all evaluated expressions should be equal.
221        ///
222        /// Each scalar expression is to be evaluated over the cross-product of all records
223        /// from all inputs. In many cases this may just be column selection from specific
224        /// inputs, but more general cases exist (e.g. complex functions of multiple columns
225        /// from multiple inputs, or just constant literals).
226        equivalences: Vec<Vec<MirScalarExpr>>,
227        /// Join implementation information.
228        #[serde(default)]
229        implementation: JoinImplementation,
230    },
231    /// Group a dataflow by some columns and aggregate over each group
232    ///
233    /// The runtime memory footprint of this operator is at most proportional to the
234    /// number of distinct records in the input and output. The actual requirements
235    /// can be less: the number of distinct inputs to each aggregate, summed across
236    /// each aggregate, plus the output size. For more details consult the code that
237    /// builds the associated dataflow.
238    Reduce {
239        /// The source collection.
240        input: Box<MirRelationExpr>,
241        /// Column indices used to form groups.
242        group_key: Vec<MirScalarExpr>,
243        /// Expressions which determine values to append to each row, after the group keys.
244        aggregates: Vec<AggregateExpr>,
245        /// True iff the input is known to monotonically increase (only addition of records).
246        #[serde(default)]
247        monotonic: bool,
248        /// User hint: expected number of values per group key. Used to optimize physical rendering.
249        #[serde(default)]
250        expected_group_size: Option<u64>,
251    },
252    /// Groups and orders within each group, limiting output.
253    ///
254    /// The runtime memory footprint of this operator is proportional to its input and output.
255    TopK {
256        /// The source collection.
257        input: Box<MirRelationExpr>,
258        /// Column indices used to form groups.
259        group_key: Vec<usize>,
260        /// Column indices used to order rows within groups.
261        order_key: Vec<ColumnOrder>,
262        /// Number of records to retain
263        #[serde(default)]
264        limit: Option<MirScalarExpr>,
265        /// Number of records to skip
266        #[serde(default)]
267        offset: usize,
268        /// True iff the input is known to monotonically increase (only addition of records).
269        #[serde(default)]
270        monotonic: bool,
271        /// User-supplied hint: how many rows will have the same group key.
272        #[serde(default)]
273        expected_group_size: Option<u64>,
274    },
275    /// Return a dataflow where the row counts are negated
276    ///
277    /// The runtime memory footprint of this operator is zero.
278    Negate {
279        /// The source collection.
280        input: Box<MirRelationExpr>,
281    },
282    /// Keep rows from a dataflow where the row counts are positive
283    ///
284    /// The runtime memory footprint of this operator is proportional to its input and output.
285    Threshold {
286        /// The source collection.
287        input: Box<MirRelationExpr>,
288    },
289    /// Adds the frequencies of elements in contained sets.
290    ///
291    /// The runtime memory footprint of this operator is zero.
292    Union {
293        /// A source collection.
294        base: Box<MirRelationExpr>,
295        /// Source collections to union.
296        inputs: Vec<MirRelationExpr>,
297    },
298    /// Technically a no-op. Used to render an index. Will be used to optimize queries
299    /// on finer grain. Each `keys` item represents a different index that should be
300    /// produced from the `keys`.
301    ///
302    /// The runtime memory footprint of this operator is proportional to its input.
303    ArrangeBy {
304        /// The source collection
305        input: Box<MirRelationExpr>,
306        /// Columns to arrange `input` by, in order of decreasing primacy
307        keys: Vec<Vec<MirScalarExpr>>,
308    },
309}
310
311impl PartialEq for MirRelationExpr {
312    fn eq(&self, other: &Self) -> bool {
313        // Capture the result and test it wrt `Ord` implementation in test environments.
314        let result = structured_diff::MreDiff::new(self, other).next().is_none();
315        mz_ore::soft_assert_eq_no_log!(result, self.cmp(other) == Ordering::Equal);
316        result
317    }
318}
319impl Eq for MirRelationExpr {}
320
321impl MirRelationExpr {
322    /// Reports the schema of the relation.
323    ///
324    /// This is the SQL-type parallel of [`Self::typ`]; it is merely
325    /// a wrapper around it, returning a [`SqlRelationType`] instead of
326    /// a [`ReprRelationType`].
327    pub fn sql_typ(&self) -> SqlRelationType {
328        let repr_typ = self.typ();
329        SqlRelationType::from_repr(&repr_typ)
330    }
331
332    /// Reports the repr schema of the relation.
333    ///
334    /// This method determines the type through recursive traversal of the
335    /// relation expression, drawing from the types of base collections.
336    /// As such, this is not an especially cheap method, and should be used
337    /// judiciously.
338    ///
339    /// The relation type is computed incrementally with a recursive post-order
340    /// traversal, that accumulates the input types for the relations yet to be
341    /// visited in `type_stack`.
342    pub fn typ(&self) -> ReprRelationType {
343        let mut type_stack = Vec::new();
344        self.visit_pre_post(
345            &mut |e: &MirRelationExpr| -> Option<Vec<&MirRelationExpr>> {
346                match &e {
347                    MirRelationExpr::Let { body, .. } => Some(vec![&*body]),
348                    MirRelationExpr::LetRec { body, .. } => Some(vec![&*body]),
349                    _ => None,
350                }
351            },
352            &mut |e: &MirRelationExpr| {
353                match e {
354                    MirRelationExpr::Let { .. } => {
355                        let body_typ = type_stack.pop().unwrap();
356                        // Insert a dummy relation type for the value, since `typ_with_input_types`
357                        // won't look at it, but expects the relation type of the body to be second.
358                        type_stack.push(ReprRelationType::empty());
359                        type_stack.push(body_typ);
360                    }
361                    MirRelationExpr::LetRec { values, .. } => {
362                        let body_typ = type_stack.pop().unwrap();
363                        type_stack.extend(
364                            std::iter::repeat(ReprRelationType::empty()).take(values.len()),
365                        );
366                        // Insert dummy relation types for the values, since `typ_with_input_types`
367                        // won't look at them, but expects the relation type of the body to be last.
368                        type_stack.push(body_typ);
369                    }
370                    _ => {}
371                }
372                let num_inputs = e.num_inputs();
373                let relation_type =
374                    e.typ_with_input_types(&type_stack[type_stack.len() - num_inputs..]);
375                type_stack.truncate(type_stack.len() - num_inputs);
376                type_stack.push(relation_type);
377            },
378        );
379        assert_eq!(type_stack.len(), 1);
380        type_stack.pop().unwrap()
381    }
382
383    /// Reports the repr schema of the relation given the repr schema of the input relations.
384    pub fn typ_with_input_types(&self, input_types: &[ReprRelationType]) -> ReprRelationType {
385        let column_types = self.col_with_input_cols(input_types.iter().map(|i| &i.column_types));
386        let unique_keys = self.keys_with_input_keys(
387            input_types.iter().map(|i| i.arity()),
388            input_types.iter().map(|i| &i.keys),
389        );
390        ReprRelationType::new(column_types).with_keys(unique_keys)
391    }
392
393    /// Reports the column types of the relation given the column types of the
394    /// input relations.
395    ///
396    /// This method delegates to `try_col_with_input_cols`, panicking if an `Err`
397    /// variant is returned.
398    pub fn col_with_input_cols<'a, I>(&self, input_types: I) -> Vec<ReprColumnType>
399    where
400        I: Iterator<Item = &'a Vec<ReprColumnType>>,
401    {
402        match self.try_col_with_input_cols(input_types) {
403            Ok(col_types) => col_types,
404            Err(err) => panic!("{err}"),
405        }
406    }
407
408    /// Reports the column types of the relation given the column types of the input relations.
409    ///
410    /// `input_types` is required to contain the column types for the input relations of
411    /// the current relation in the same order as they are visited by `try_visit_children`
412    /// method, even though not all may be used for computing the schema of the
413    /// current relation. For example, `Let` expects two input types, one for the
414    /// value relation and one for the body, in that order, but only the one for the
415    /// body is used to determine the type of the `Let` relation.
416    ///
417    /// It is meant to be used during post-order traversals to compute column types
418    /// incrementally.
419    pub fn try_col_with_input_cols<'a, I>(
420        &self,
421        mut input_types: I,
422    ) -> Result<Vec<ReprColumnType>, String>
423    where
424        I: Iterator<Item = &'a Vec<ReprColumnType>>,
425    {
426        use MirRelationExpr::*;
427
428        let col_types = match self {
429            Constant { rows, typ } => {
430                let mut col_types = typ.column_types.clone();
431                let mut seen_null = vec![false; typ.arity()];
432                if let Ok(rows) = rows {
433                    for (row, _diff) in rows {
434                        for (datum, i) in row.iter().zip_eq(0..typ.arity()) {
435                            if datum.is_null() {
436                                seen_null[i] = true;
437                            }
438                        }
439                    }
440                }
441                for (&seen_null, i) in seen_null.iter().zip_eq(0..typ.arity()) {
442                    if !seen_null {
443                        col_types[i].nullable = false;
444                    } else {
445                        assert!(col_types[i].nullable);
446                    }
447                }
448                col_types
449            }
450            Get { typ, .. } => typ.column_types.clone(),
451            Project { outputs, .. } => {
452                let input = input_types.next().unwrap();
453                outputs.iter().map(|&i| input[i].clone()).collect()
454            }
455            Map { scalars, .. } => {
456                let mut result = input_types.next().unwrap().clone();
457                for scalar in scalars.iter() {
458                    result.push(scalar.typ(&result))
459                }
460                result
461            }
462            FlatMap { func, .. } => {
463                let mut result = input_types.next().unwrap().clone();
464                result.extend(
465                    func.output_sql_type()
466                        .column_types
467                        .iter()
468                        .map(ReprColumnType::from),
469                );
470                result
471            }
472            Filter { predicates, .. } => {
473                let mut result = input_types.next().unwrap().clone();
474
475                // Set as nonnull any columns where null values would cause
476                // any predicate to evaluate to null.
477                for column in non_nullable_columns(predicates) {
478                    result[column].nullable = false;
479                }
480                result
481            }
482            Join { equivalences, .. } => {
483                // Concatenate input column types
484                let mut types = input_types.flat_map(|cols| cols.to_owned()).collect_vec();
485                // In an equivalence class, if any column is non-null, then make all non-null
486                for equivalence in equivalences {
487                    let col_inds = equivalence
488                        .iter()
489                        .filter_map(|expr| match expr {
490                            MirScalarExpr::Column(col, _name) => Some(*col),
491                            _ => None,
492                        })
493                        .collect_vec();
494                    if col_inds.iter().any(|i| !types.get(*i).unwrap().nullable) {
495                        for i in col_inds {
496                            types.get_mut(i).unwrap().nullable = false;
497                        }
498                    }
499                }
500                types
501            }
502            Reduce {
503                group_key,
504                aggregates,
505                ..
506            } => {
507                let input = input_types.next().unwrap();
508                group_key
509                    .iter()
510                    .map(|e| e.typ(input))
511                    .chain(aggregates.iter().map(|agg| agg.typ(input)))
512                    .collect()
513            }
514            TopK { .. } | Negate { .. } | Threshold { .. } | ArrangeBy { .. } => {
515                input_types.next().unwrap().clone()
516            }
517            Let { .. } => {
518                // skip over the input types for `value`.
519                input_types.nth(1).unwrap().clone()
520            }
521            LetRec { values, .. } => {
522                // skip over the input types for `values`.
523                input_types.nth(values.len()).unwrap().clone()
524            }
525            Union { .. } => {
526                let mut result = input_types.next().unwrap().clone();
527                for input_col_types in input_types {
528                    for (base_col, col) in result.iter_mut().zip_eq(input_col_types) {
529                        *base_col = base_col
530                            .union(col)
531                            .map_err(|e| format!("{}\nin plan:\n{}", e, self.pretty()))?;
532                    }
533                }
534                result
535            }
536        };
537
538        Ok(col_types)
539    }
540
541    /// Reports the unique keys of the relation given the arities and the unique
542    /// keys of the input relations.
543    ///
544    /// `input_arities` and `input_keys` are required to contain the
545    /// corresponding info for the input relations of
546    /// the current relation in the same order as they are visited by `try_visit_children`
547    /// method, even though not all may be used for computing the schema of the
548    /// current relation. For example, `Let` expects two input types, one for the
549    /// value relation and one for the body, in that order, but only the one for the
550    /// body is used to determine the type of the `Let` relation.
551    ///
552    /// It is meant to be used during post-order traversals to compute unique keys
553    /// incrementally.
554    pub fn keys_with_input_keys<'a, I, J>(
555        &self,
556        mut input_arities: I,
557        mut input_keys: J,
558    ) -> Vec<Vec<usize>>
559    where
560        I: Iterator<Item = usize>,
561        J: Iterator<Item = &'a Vec<Vec<usize>>>,
562    {
563        use MirRelationExpr::*;
564
565        let mut keys = match self {
566            Constant {
567                rows: Ok(rows),
568                typ,
569            } => {
570                let n_cols = typ.arity();
571                // If the `i`th entry is `Some`, then we have not yet observed non-uniqueness in the `i`th column.
572                let mut unique_values_per_col = vec![Some(BTreeSet::<Datum>::default()); n_cols];
573                for (row, diff) in rows {
574                    for (i, datum) in row.iter().enumerate() {
575                        if datum != Datum::Dummy {
576                            if let Some(unique_vals) = &mut unique_values_per_col[i] {
577                                let is_dupe = *diff != Diff::ONE || !unique_vals.insert(datum);
578                                if is_dupe {
579                                    unique_values_per_col[i] = None;
580                                }
581                            }
582                        }
583                    }
584                }
585                if rows.len() == 0 || (rows.len() == 1 && rows[0].1 == Diff::ONE) {
586                    vec![vec![]]
587                } else {
588                    // XXX - Multi-column keys are not detected.
589                    typ.keys
590                        .iter()
591                        .cloned()
592                        .chain(
593                            unique_values_per_col
594                                .into_iter()
595                                .enumerate()
596                                .filter(|(_idx, unique_vals)| unique_vals.is_some())
597                                .map(|(idx, _)| vec![idx]),
598                        )
599                        .collect()
600                }
601            }
602            Constant { rows: Err(_), typ } | Get { typ, .. } => typ.keys.clone(),
603            Threshold { .. } | ArrangeBy { .. } => input_keys.next().unwrap().clone(),
604            Let { .. } => {
605                // skip over the unique keys for value
606                input_keys.nth(1).unwrap().clone()
607            }
608            LetRec { values, .. } => {
609                // skip over the unique keys for value
610                input_keys.nth(values.len()).unwrap().clone()
611            }
612            Project { outputs, .. } => {
613                let input = input_keys.next().unwrap();
614                input
615                    .iter()
616                    .filter_map(|key_set| {
617                        if key_set.iter().all(|k| outputs.contains(k)) {
618                            Some(
619                                key_set
620                                    .iter()
621                                    .map(|c| outputs.iter().position(|o| o == c).unwrap())
622                                    .collect(),
623                            )
624                        } else {
625                            None
626                        }
627                    })
628                    .collect()
629            }
630            Map { scalars, .. } => {
631                let mut remappings = Vec::new();
632                let arity = input_arities.next().unwrap();
633                for (column, scalar) in scalars.iter().enumerate() {
634                    // assess whether the scalar preserves uniqueness,
635                    // and could participate in a key!
636
637                    fn uniqueness(expr: &MirScalarExpr) -> Option<usize> {
638                        match expr {
639                            MirScalarExpr::CallUnary { func, expr } => {
640                                if func.preserves_uniqueness() {
641                                    uniqueness(expr)
642                                } else {
643                                    None
644                                }
645                            }
646                            MirScalarExpr::Column(c, _name) => Some(*c),
647                            _ => None,
648                        }
649                    }
650
651                    if let Some(c) = uniqueness(scalar) {
652                        remappings.push((c, column + arity));
653                    }
654                }
655
656                let mut result = input_keys.next().unwrap().clone();
657                let mut new_keys = Vec::new();
658                // Any column in `remappings` could be replaced in a key
659                // by the corresponding c. This could lead to combinatorial
660                // explosion using our current representation, so we wont
661                // do that. Instead, we'll handle the case of one remapping.
662                if remappings.len() == 1 {
663                    let (old, new) = remappings.pop().unwrap();
664                    for key in &result {
665                        if key.contains(&old) {
666                            let mut new_key: Vec<usize> =
667                                key.iter().cloned().filter(|k| k != &old).collect();
668                            new_key.push(new);
669                            new_key.sort_unstable();
670                            new_keys.push(new_key);
671                        }
672                    }
673                    result.append(&mut new_keys);
674                }
675                result
676            }
677            FlatMap { .. } => {
678                // FlatMap can add duplicate rows, so input keys are no longer
679                // valid
680                vec![]
681            }
682            Negate { .. } => {
683                // Although negate may have distinct records for each key,
684                // the multiplicity is -1 rather than 1. This breaks many
685                // of the optimization uses of "keys".
686                vec![]
687            }
688            Filter { predicates, .. } => {
689                // A filter inherits the keys of its input unless the filters
690                // have reduced the input to a single row, in which case the
691                // keys of the input are `()`.
692                let mut input = input_keys.next().unwrap().clone();
693
694                if !input.is_empty() {
695                    // Track columns equated to literals, which we can prune.
696                    let mut cols_equal_to_literal = BTreeSet::new();
697
698                    // Perform union find on `col1 = col2` to establish
699                    // connected components of equated columns. Absent any
700                    // equalities, this will be `0 .. #c` (where #c is the
701                    // greatest column referenced by a predicate), but each
702                    // equality will orient the root of the greater to the root
703                    // of the lesser.
704                    let mut union_find = Vec::new();
705
706                    for expr in predicates.iter() {
707                        if let MirScalarExpr::CallBinary {
708                            func: crate::BinaryFunc::Eq(_),
709                            expr1,
710                            expr2,
711                        } = expr
712                        {
713                            if let MirScalarExpr::Column(c, _name) = &**expr1 {
714                                if expr2.is_literal_ok() {
715                                    cols_equal_to_literal.insert(c);
716                                }
717                            }
718                            if let MirScalarExpr::Column(c, _name) = &**expr2 {
719                                if expr1.is_literal_ok() {
720                                    cols_equal_to_literal.insert(c);
721                                }
722                            }
723                            // Perform union-find to equate columns.
724                            if let (Some(c1), Some(c2)) = (expr1.as_column(), expr2.as_column()) {
725                                if c1 != c2 {
726                                    // Ensure union_find has entries up to
727                                    // max(c1, c2) by filling up missing
728                                    // positions with identity mappings.
729                                    while union_find.len() <= std::cmp::max(c1, c2) {
730                                        union_find.push(union_find.len());
731                                    }
732                                    let mut r1 = c1; // Find the representative column of [c1].
733                                    while r1 != union_find[r1] {
734                                        assert!(union_find[r1] < r1);
735                                        r1 = union_find[r1];
736                                    }
737                                    let mut r2 = c2; // Find the representative column of [c2].
738                                    while r2 != union_find[r2] {
739                                        assert!(union_find[r2] < r2);
740                                        r2 = union_find[r2];
741                                    }
742                                    // Union [c1] and [c2] by pointing the
743                                    // larger to the smaller representative (we
744                                    // update the remaining equivalence class
745                                    // members only once after this for-loop).
746                                    union_find[std::cmp::max(r1, r2)] = std::cmp::min(r1, r2);
747                                }
748                            }
749                        }
750                    }
751
752                    // Complete union-find by pointing each element at its representative column.
753                    for i in 0..union_find.len() {
754                        // Iteration not required, as each prior already references the right column.
755                        union_find[i] = union_find[union_find[i]];
756                    }
757
758                    // Remove columns bound to literals, and remap columns equated to earlier columns.
759                    // We will re-expand remapped columns in a moment, but this avoids exponential work.
760                    for key_set in &mut input {
761                        key_set.retain(|k| !cols_equal_to_literal.contains(&k));
762                        for col in key_set.iter_mut() {
763                            if let Some(equiv) = union_find.get(*col) {
764                                *col = *equiv;
765                            }
766                        }
767                        key_set.sort();
768                        key_set.dedup();
769                    }
770                    input.sort();
771                    input.dedup();
772
773                    // Expand out each key to each of its equivalent forms.
774                    // Each instance of `col` can be replaced by any equivalent column.
775                    // This has the potential to result in exponentially sized number of unique keys,
776                    // and in the future we should probably maintain unique keys modulo equivalence.
777
778                    // First, compute an inverse map from each representative
779                    // column `sub` to all other equivalent columns `col`.
780                    let mut subs = Vec::new();
781                    for (col, sub) in union_find.iter().enumerate() {
782                        if *sub != col {
783                            assert!(*sub < col);
784                            while subs.len() <= *sub {
785                                subs.push(Vec::new());
786                            }
787                            subs[*sub].push(col);
788                        }
789                    }
790                    // For each column, substitute for it in each occurrence.
791                    let mut to_add = Vec::new();
792                    for (col, subs) in subs.iter().enumerate() {
793                        if !subs.is_empty() {
794                            for key_set in input.iter() {
795                                if key_set.contains(&col) {
796                                    let mut to_extend = key_set.clone();
797                                    to_extend.retain(|c| c != &col);
798                                    for sub in subs {
799                                        to_extend.push(*sub);
800                                        to_add.push(to_extend.clone());
801                                        to_extend.pop();
802                                    }
803                                }
804                            }
805                        }
806                        // No deduplication, as we cannot introduce duplicates.
807                        input.append(&mut to_add);
808                    }
809                    for key_set in input.iter_mut() {
810                        key_set.sort();
811                        key_set.dedup();
812                    }
813                }
814                input
815            }
816            Join { equivalences, .. } => {
817                // It is important the `new_from_input_arities` constructor is
818                // used. Otherwise, Materialize may potentially end up in an
819                // infinite loop.
820                let input_mapper = crate::JoinInputMapper::new_from_input_arities(input_arities);
821
822                input_mapper.global_keys(input_keys, equivalences)
823            }
824            Reduce { group_key, .. } => {
825                // The group key should form a key, but we might already have
826                // keys that are subsets of the group key, and should retain
827                // those instead, if so.
828                let mut result = Vec::new();
829                for key_set in input_keys.next().unwrap() {
830                    if key_set
831                        .iter()
832                        .all(|k| group_key.contains(&MirScalarExpr::column(*k)))
833                    {
834                        result.push(
835                            key_set
836                                .iter()
837                                .map(|i| {
838                                    group_key
839                                        .iter()
840                                        .position(|k| k == &MirScalarExpr::column(*i))
841                                        .unwrap()
842                                })
843                                .collect::<Vec<_>>(),
844                        );
845                    }
846                }
847                if result.is_empty() {
848                    result.push((0..group_key.len()).collect());
849                }
850                result
851            }
852            TopK {
853                group_key, limit, ..
854            } => {
855                // If `limit` is `Some(1)` then the group key will become
856                // a unique key, as there will be only one record with that key.
857                let mut result = input_keys.next().unwrap().clone();
858                if limit.as_ref().and_then(|x| x.as_literal_int64()) == Some(1) {
859                    result.push(group_key.clone())
860                }
861                result
862            }
863            Union { base, inputs } => {
864                // Generally, unions do not have any unique keys, because
865                // each input might duplicate some. However, there is at
866                // least one idiomatic structure that does preserve keys,
867                // which results from SQL aggregations that must populate
868                // absent records with default values. In that pattern,
869                // the union of one GET with its negation, which has first
870                // been subjected to a projection and map, we can remove
871                // their influence on the key structure.
872                //
873                // If there are A, B, each with a unique `key` such that
874                // we are looking at
875                //
876                //     A.proj(set_containing_key) + (B - A.proj(key)).map(stuff)
877                //
878                // Then we can report `key` as a unique key.
879                //
880                // TODO: make unique key structure an optimization analysis
881                // rather than part of the type information.
882                // TODO: perhaps ensure that (above) A.proj(key) is a
883                // subset of B, as otherwise there are negative records
884                // and who knows what is true (not expected, but again
885                // who knows what the query plan might look like).
886
887                let arity = input_arities.next().unwrap();
888                let (base_projection, base_with_project_stripped) =
889                    if let MirRelationExpr::Project { input, outputs } = &**base {
890                        (outputs.clone(), &**input)
891                    } else {
892                        // A input without a project is equivalent to an input
893                        // with the project being all columns in the input in order.
894                        ((0..arity).collect::<Vec<_>>(), &**base)
895                    };
896                let mut result = Vec::new();
897                if let MirRelationExpr::Get {
898                    id: first_id,
899                    typ: _,
900                    ..
901                } = base_with_project_stripped
902                {
903                    if inputs.len() == 1 {
904                        if let MirRelationExpr::Map { input, .. } = &inputs[0] {
905                            if let MirRelationExpr::Union { base, inputs } = &**input {
906                                if inputs.len() == 1 {
907                                    if let Some((input, outputs)) = base.is_negated_project() {
908                                        if let MirRelationExpr::Get {
909                                            id: second_id,
910                                            typ: _,
911                                            ..
912                                        } = input
913                                        {
914                                            if first_id == second_id {
915                                                result.extend(
916                                                    input_keys
917                                                        .next()
918                                                        .unwrap()
919                                                        .into_iter()
920                                                        .filter(|key| {
921                                                            key.iter().all(|c| {
922                                                                outputs.get(*c) == Some(c)
923                                                                    && base_projection.get(*c)
924                                                                        == Some(c)
925                                                            })
926                                                        })
927                                                        .cloned(),
928                                                );
929                                            }
930                                        }
931                                    }
932                                }
933                            }
934                        }
935                    }
936                }
937                // Important: do not inherit keys of either input, as not unique.
938                result
939            }
940        };
941        keys.sort();
942        keys.dedup();
943        keys
944    }
945
946    /// The number of columns in the relation.
947    ///
948    /// This number is determined from the type, which is determined recursively
949    /// at non-trivial cost.
950    ///
951    /// The arity is computed incrementally with a recursive post-order
952    /// traversal, that accumulates the arities for the relations yet to be
953    /// visited in `arity_stack`.
954    pub fn arity(&self) -> usize {
955        let mut arity_stack = Vec::new();
956        self.visit_pre_post(
957            &mut |e: &MirRelationExpr| -> Option<Vec<&MirRelationExpr>> {
958                match &e {
959                    MirRelationExpr::Let { body, .. } => {
960                        // Do not traverse the value sub-graph, since it's not relevant for
961                        // determining the arity of Let operators.
962                        Some(vec![&*body])
963                    }
964                    MirRelationExpr::LetRec { body, .. } => {
965                        // Do not traverse the value sub-graph, since it's not relevant for
966                        // determining the arity of Let operators.
967                        Some(vec![&*body])
968                    }
969                    MirRelationExpr::Project { .. } | MirRelationExpr::Reduce { .. } => {
970                        // No further traversal is required; these operators know their arity.
971                        Some(Vec::new())
972                    }
973                    _ => None,
974                }
975            },
976            &mut |e: &MirRelationExpr| {
977                match &e {
978                    MirRelationExpr::Let { .. } => {
979                        let body_arity = arity_stack.pop().unwrap();
980                        arity_stack.push(0);
981                        arity_stack.push(body_arity);
982                    }
983                    MirRelationExpr::LetRec { values, .. } => {
984                        let body_arity = arity_stack.pop().unwrap();
985                        arity_stack.extend(std::iter::repeat(0).take(values.len()));
986                        arity_stack.push(body_arity);
987                    }
988                    MirRelationExpr::Project { .. } | MirRelationExpr::Reduce { .. } => {
989                        arity_stack.push(0);
990                    }
991                    _ => {}
992                }
993                let num_inputs = e.num_inputs();
994                let input_arities = arity_stack.drain(arity_stack.len() - num_inputs..);
995                let arity = e.arity_with_input_arities(input_arities);
996                arity_stack.push(arity);
997            },
998        );
999        assert_eq!(arity_stack.len(), 1);
1000        arity_stack.pop().unwrap()
1001    }
1002
1003    /// Reports the arity of the relation given the schema of the input relations.
1004    ///
1005    /// `input_arities` is required to contain the arities for the input relations of
1006    /// the current relation in the same order as they are visited by `try_visit_children`
1007    /// method, even though not all may be used for computing the schema of the
1008    /// current relation. For example, `Let` expects two input types, one for the
1009    /// value relation and one for the body, in that order, but only the one for the
1010    /// body is used to determine the type of the `Let` relation.
1011    ///
1012    /// It is meant to be used during post-order traversals to compute arities
1013    /// incrementally.
1014    pub fn arity_with_input_arities<I>(&self, mut input_arities: I) -> usize
1015    where
1016        I: Iterator<Item = usize>,
1017    {
1018        use MirRelationExpr::*;
1019
1020        match self {
1021            Constant { rows: _, typ } => typ.arity(),
1022            Get { typ, .. } => typ.arity(),
1023            Let { .. } => {
1024                input_arities.next();
1025                input_arities.next().unwrap()
1026            }
1027            LetRec { values, .. } => {
1028                for _ in 0..values.len() {
1029                    input_arities.next();
1030                }
1031                input_arities.next().unwrap()
1032            }
1033            Project { outputs, .. } => outputs.len(),
1034            Map { scalars, .. } => input_arities.next().unwrap() + scalars.len(),
1035            FlatMap { func, .. } => input_arities.next().unwrap() + func.output_arity(),
1036            Join { .. } => input_arities.sum(),
1037            Reduce {
1038                input: _,
1039                group_key,
1040                aggregates,
1041                ..
1042            } => group_key.len() + aggregates.len(),
1043            Filter { .. }
1044            | TopK { .. }
1045            | Negate { .. }
1046            | Threshold { .. }
1047            | Union { .. }
1048            | ArrangeBy { .. } => input_arities.next().unwrap(),
1049        }
1050    }
1051
1052    /// The number of child relations this relation has.
1053    pub fn num_inputs(&self) -> usize {
1054        let mut count = 0;
1055
1056        self.visit_children(|_| count += 1);
1057
1058        count
1059    }
1060
1061    /// Constructs a constant collection from specific rows and schema, where
1062    /// each row will have a multiplicity of one.
1063    pub fn constant(rows: Vec<Vec<Datum>>, typ: ReprRelationType) -> Self {
1064        let rows = rows.into_iter().map(|row| (row, Diff::ONE)).collect();
1065        MirRelationExpr::constant_diff(rows, typ)
1066    }
1067
1068    /// Constructs a constant collection from specific rows and schema, where
1069    /// each row can have an arbitrary multiplicity.
1070    pub fn constant_diff(rows: Vec<(Vec<Datum>, Diff)>, typ: ReprRelationType) -> Self {
1071        for (row, _diff) in &rows {
1072            for (datum, column_typ) in row.iter().zip_eq(typ.column_types.iter()) {
1073                assert!(
1074                    datum.is_instance_of(column_typ),
1075                    "Expected datum of type {:?}, got value {:?}",
1076                    column_typ,
1077                    datum
1078                );
1079            }
1080        }
1081        let rows = Ok(rows
1082            .into_iter()
1083            .map(move |(row, diff)| (Row::pack_slice(&row), diff))
1084            .collect());
1085        MirRelationExpr::Constant { rows, typ }
1086    }
1087
1088    /// If self is a constant, return the value and the type, otherwise `None`.
1089    /// Looks behind `ArrangeBy`s.
1090    pub fn as_const(&self) -> Option<(&Result<Vec<(Row, Diff)>, EvalError>, &ReprRelationType)> {
1091        match self {
1092            MirRelationExpr::Constant { rows, typ } => Some((rows, typ)),
1093            MirRelationExpr::ArrangeBy { input, .. } => input.as_const(),
1094            _ => None,
1095        }
1096    }
1097
1098    /// If self is a constant, mutably return the value and the type, otherwise `None`.
1099    /// Looks behind `ArrangeBy`s.
1100    pub fn as_const_mut(
1101        &mut self,
1102    ) -> Option<(
1103        &mut Result<Vec<(Row, Diff)>, EvalError>,
1104        &mut ReprRelationType,
1105    )> {
1106        match self {
1107            MirRelationExpr::Constant { rows, typ } => Some((rows, typ)),
1108            MirRelationExpr::ArrangeBy { input, .. } => input.as_const_mut(),
1109            _ => None,
1110        }
1111    }
1112
1113    /// If self is a constant error, return the error, otherwise `None`.
1114    /// Looks behind `ArrangeBy`s.
1115    pub fn as_const_err(&self) -> Option<&EvalError> {
1116        match self {
1117            MirRelationExpr::Constant { rows: Err(e), .. } => Some(e),
1118            MirRelationExpr::ArrangeBy { input, .. } => input.as_const_err(),
1119            _ => None,
1120        }
1121    }
1122
1123    /// Checks if `self` is the single element collection with no columns.
1124    pub fn is_constant_singleton(&self) -> bool {
1125        if let Some((Ok(rows), typ)) = self.as_const() {
1126            rows.len() == 1 && typ.column_types.len() == 0 && rows[0].1 == Diff::ONE
1127        } else {
1128            false
1129        }
1130    }
1131
1132    /// Constructs the expression for getting a local collection.
1133    pub fn local_get(id: LocalId, typ: ReprRelationType) -> Self {
1134        MirRelationExpr::Get {
1135            id: Id::Local(id),
1136            typ,
1137            access_strategy: AccessStrategy::UnknownOrLocal,
1138        }
1139    }
1140
1141    /// Constructs the expression for getting a global collection
1142    pub fn global_get(id: GlobalId, typ: ReprRelationType) -> Self {
1143        MirRelationExpr::Get {
1144            id: Id::Global(id),
1145            typ,
1146            access_strategy: AccessStrategy::UnknownOrLocal,
1147        }
1148    }
1149
1150    /// Retains only the columns specified by `output`.
1151    pub fn project(mut self, mut outputs: Vec<usize>) -> Self {
1152        if let MirRelationExpr::Project {
1153            outputs: columns, ..
1154        } = &mut self
1155        {
1156            // Update `outputs` to reference base columns of `input`.
1157            for column in outputs.iter_mut() {
1158                *column = columns[*column];
1159            }
1160            *columns = outputs;
1161            self
1162        } else {
1163            MirRelationExpr::Project {
1164                input: Box::new(self),
1165                outputs,
1166            }
1167        }
1168    }
1169
1170    /// Append to each row the results of applying elements of `scalar`.
1171    pub fn map(mut self, scalars: Vec<MirScalarExpr>) -> Self {
1172        if let MirRelationExpr::Map { scalars: s, .. } = &mut self {
1173            s.extend(scalars);
1174            self
1175        } else if !scalars.is_empty() {
1176            MirRelationExpr::Map {
1177                input: Box::new(self),
1178                scalars,
1179            }
1180        } else {
1181            self
1182        }
1183    }
1184
1185    /// Append to each row a single `scalar`.
1186    pub fn map_one(self, scalar: MirScalarExpr) -> Self {
1187        self.map(vec![scalar])
1188    }
1189
1190    /// Like `map`, but yields zero-or-more output rows per input row
1191    pub fn flat_map(self, func: TableFunc, exprs: Vec<MirScalarExpr>) -> Self {
1192        MirRelationExpr::FlatMap {
1193            input: Box::new(self),
1194            func,
1195            exprs,
1196        }
1197    }
1198
1199    /// Retain only the rows satisfying each of several predicates.
1200    pub fn filter<I>(mut self, predicates: I) -> Self
1201    where
1202        I: IntoIterator<Item = MirScalarExpr>,
1203    {
1204        // Extract existing predicates
1205        let mut new_predicates = if let MirRelationExpr::Filter { input, predicates } = self {
1206            self = *input;
1207            predicates
1208        } else {
1209            Vec::new()
1210        };
1211        // Normalize collection of predicates.
1212        new_predicates.extend(predicates);
1213        new_predicates.retain(|p| !p.is_literal_true());
1214        new_predicates.sort();
1215        new_predicates.dedup();
1216        // Introduce a `Filter` only if we have predicates.
1217        if !new_predicates.is_empty() {
1218            self = MirRelationExpr::Filter {
1219                input: Box::new(self),
1220                predicates: new_predicates,
1221            };
1222        }
1223
1224        self
1225    }
1226
1227    /// Form the Cartesian outer-product of rows in both inputs.
1228    pub fn product(mut self, right: Self) -> Self {
1229        if right.is_constant_singleton() {
1230            self
1231        } else if self.is_constant_singleton() {
1232            right
1233        } else if let MirRelationExpr::Join { inputs, .. } = &mut self {
1234            inputs.push(right);
1235            self
1236        } else {
1237            MirRelationExpr::join(vec![self, right], vec![])
1238        }
1239    }
1240
1241    /// Performs a relational equijoin among the input collections.
1242    ///
1243    /// The sequence `inputs` each describe different input collections, and the sequence `variables` describes
1244    /// equality constraints that some of their columns must satisfy. Each element in `variable` describes a set
1245    /// of pairs  `(input_index, column_index)` where every value described by that set must be equal.
1246    ///
1247    /// For example, the pair `(input, column)` indexes into `inputs[input][column]`, extracting the `input`th
1248    /// input collection and for each row examining its `column`th column.
1249    ///
1250    /// # Example
1251    ///
1252    /// ```rust
1253    /// use mz_repr::{Datum, SqlColumnType, ReprRelationType, ReprScalarType};
1254    /// use mz_expr::MirRelationExpr;
1255    ///
1256    /// // A common schema for each input.
1257    /// let schema = ReprRelationType::new(vec![
1258    ///     ReprScalarType::Int32.nullable(false),
1259    ///     ReprScalarType::Int32.nullable(false),
1260    /// ]);
1261    ///
1262    /// // the specific data are not important here.
1263    /// let data = vec![Datum::Int32(0), Datum::Int32(1)];
1264    ///
1265    /// // Three collections that could have been different.
1266    /// let input0 = MirRelationExpr::constant(vec![data.clone()], schema.clone());
1267    /// let input1 = MirRelationExpr::constant(vec![data.clone()], schema.clone());
1268    /// let input2 = MirRelationExpr::constant(vec![data.clone()], schema.clone());
1269    ///
1270    /// // Join the three relations looking for triangles, like so.
1271    /// //
1272    /// //     Output(A,B,C) := Input0(A,B), Input1(B,C), Input2(A,C)
1273    /// let joined = MirRelationExpr::join(
1274    ///     vec![input0, input1, input2],
1275    ///     vec![
1276    ///         vec![(0,0), (2,0)], // fields A of inputs 0 and 2.
1277    ///         vec![(0,1), (1,0)], // fields B of inputs 0 and 1.
1278    ///         vec![(1,1), (2,1)], // fields C of inputs 1 and 2.
1279    ///     ],
1280    /// );
1281    ///
1282    /// // Technically the above produces `Output(A,B,B,C,A,C)` because the columns are concatenated.
1283    /// // A projection resolves this and produces the correct output.
1284    /// let result = joined.project(vec![0, 1, 3]);
1285    /// ```
1286    pub fn join(inputs: Vec<MirRelationExpr>, variables: Vec<Vec<(usize, usize)>>) -> Self {
1287        let input_mapper = join_input_mapper::JoinInputMapper::new(&inputs);
1288
1289        let equivalences = variables
1290            .into_iter()
1291            .map(|vs| {
1292                vs.into_iter()
1293                    .map(|(r, c)| input_mapper.map_expr_to_global(MirScalarExpr::column(c), r))
1294                    .collect::<Vec<_>>()
1295            })
1296            .collect::<Vec<_>>();
1297
1298        Self::join_scalars(inputs, equivalences)
1299    }
1300
1301    /// Constructs a join operator from inputs and required-equal scalar expressions.
1302    pub fn join_scalars(
1303        mut inputs: Vec<MirRelationExpr>,
1304        equivalences: Vec<Vec<MirScalarExpr>>,
1305    ) -> Self {
1306        // Remove all constant inputs that are the identity for join.
1307        // They neither introduce nor modify any column references.
1308        inputs.retain(|i| !i.is_constant_singleton());
1309        MirRelationExpr::Join {
1310            inputs,
1311            equivalences,
1312            implementation: JoinImplementation::Unimplemented,
1313        }
1314    }
1315
1316    /// Perform a key-wise reduction / aggregation.
1317    ///
1318    /// The `group_key` argument indicates columns in the input collection that should
1319    /// be grouped, and `aggregates` lists aggregation functions each of which produces
1320    /// one output column in addition to the keys.
1321    pub fn reduce(
1322        self,
1323        group_key: Vec<usize>,
1324        aggregates: Vec<AggregateExpr>,
1325        expected_group_size: Option<u64>,
1326    ) -> Self {
1327        MirRelationExpr::Reduce {
1328            input: Box::new(self),
1329            group_key: group_key.into_iter().map(MirScalarExpr::column).collect(),
1330            aggregates,
1331            monotonic: false,
1332            expected_group_size,
1333        }
1334    }
1335
1336    /// Perform a key-wise reduction order by and limit.
1337    ///
1338    /// The `group_key` argument indicates columns in the input collection that should
1339    /// be grouped, the `order_key` argument indicates columns that should be further
1340    /// used to order records within groups, and the `limit` argument constrains the
1341    /// total number of records that should be produced in each group.
1342    pub fn top_k(
1343        self,
1344        group_key: Vec<usize>,
1345        order_key: Vec<ColumnOrder>,
1346        limit: Option<MirScalarExpr>,
1347        offset: usize,
1348        expected_group_size: Option<u64>,
1349    ) -> Self {
1350        MirRelationExpr::TopK {
1351            input: Box::new(self),
1352            group_key,
1353            order_key,
1354            limit,
1355            offset,
1356            expected_group_size,
1357            monotonic: false,
1358        }
1359    }
1360
1361    /// Negates the occurrences of each row.
1362    pub fn negate(self) -> Self {
1363        if let MirRelationExpr::Negate { input } = self {
1364            *input
1365        } else {
1366            MirRelationExpr::Negate {
1367                input: Box::new(self),
1368            }
1369        }
1370    }
1371
1372    /// Removes all but the first occurrence of each row.
1373    pub fn distinct(self) -> Self {
1374        let arity = self.arity();
1375        self.distinct_by((0..arity).collect())
1376    }
1377
1378    /// Removes all but the first occurrence of each key. Columns not included
1379    /// in the `group_key` are discarded.
1380    pub fn distinct_by(self, group_key: Vec<usize>) -> Self {
1381        self.reduce(group_key, vec![], None)
1382    }
1383
1384    /// Discards rows with a negative frequency.
1385    pub fn threshold(self) -> Self {
1386        if let MirRelationExpr::Threshold { .. } = &self {
1387            self
1388        } else {
1389            MirRelationExpr::Threshold {
1390                input: Box::new(self),
1391            }
1392        }
1393    }
1394
1395    /// Unions together any number inputs.
1396    ///
1397    /// If `inputs` is empty, then an empty relation of type `typ` is
1398    /// constructed.
1399    pub fn union_many(mut inputs: Vec<Self>, typ: ReprRelationType) -> Self {
1400        // Deconstruct `inputs` as `Union`s and reconstitute.
1401        let mut flat_inputs = Vec::with_capacity(inputs.len());
1402        for input in inputs {
1403            if let MirRelationExpr::Union { base, inputs } = input {
1404                flat_inputs.push(*base);
1405                flat_inputs.extend(inputs);
1406            } else {
1407                flat_inputs.push(input);
1408            }
1409        }
1410        inputs = flat_inputs;
1411        if inputs.len() == 0 {
1412            MirRelationExpr::Constant {
1413                rows: Ok(vec![]),
1414                typ,
1415            }
1416        } else if inputs.len() == 1 {
1417            inputs.into_element()
1418        } else {
1419            MirRelationExpr::Union {
1420                base: Box::new(inputs.remove(0)),
1421                inputs,
1422            }
1423        }
1424    }
1425
1426    /// Produces one collection where each row is present with the sum of its frequencies in each input.
1427    pub fn union(self, other: Self) -> Self {
1428        // Deconstruct `self` and `other` as `Union`s and reconstitute.
1429        let mut flat_inputs = Vec::with_capacity(2);
1430        if let MirRelationExpr::Union { base, inputs } = self {
1431            flat_inputs.push(*base);
1432            flat_inputs.extend(inputs);
1433        } else {
1434            flat_inputs.push(self);
1435        }
1436        if let MirRelationExpr::Union { base, inputs } = other {
1437            flat_inputs.push(*base);
1438            flat_inputs.extend(inputs);
1439        } else {
1440            flat_inputs.push(other);
1441        }
1442
1443        MirRelationExpr::Union {
1444            base: Box::new(flat_inputs.remove(0)),
1445            inputs: flat_inputs,
1446        }
1447    }
1448
1449    /// Arranges the collection by the specified columns
1450    pub fn arrange_by(self, keys: &[Vec<MirScalarExpr>]) -> Self {
1451        MirRelationExpr::ArrangeBy {
1452            input: Box::new(self),
1453            keys: keys.to_owned(),
1454        }
1455    }
1456
1457    /// Indicates if this is a constant empty collection.
1458    ///
1459    /// A false value does not mean the collection is known to be non-empty,
1460    /// only that we cannot currently determine that it is statically empty.
1461    pub fn is_empty(&self) -> bool {
1462        if let Some((Ok(rows), ..)) = self.as_const() {
1463            rows.is_empty()
1464        } else {
1465            false
1466        }
1467    }
1468
1469    /// If the expression is a negated project, return the input and the projection.
1470    pub fn is_negated_project(&self) -> Option<(&MirRelationExpr, &[usize])> {
1471        if let MirRelationExpr::Negate { input } = self {
1472            if let MirRelationExpr::Project { input, outputs } = &**input {
1473                return Some((&**input, outputs));
1474            }
1475        }
1476        if let MirRelationExpr::Project { input, outputs } = self {
1477            if let MirRelationExpr::Negate { input } = &**input {
1478                return Some((&**input, outputs));
1479            }
1480        }
1481        None
1482    }
1483
1484    /// Pretty-print this [MirRelationExpr] to a string.
1485    pub fn pretty(&self) -> String {
1486        let config = ExplainConfig::default();
1487        self.debug_explain(&config, None)
1488    }
1489
1490    /// Pretty-print this [MirRelationExpr] to a string using a custom
1491    /// [ExplainConfig] and an optionally provided [ExprHumanizer].
1492    /// This is intended for debugging and tests, not users.
1493    pub fn debug_explain(
1494        &self,
1495        config: &ExplainConfig,
1496        humanizer: Option<&dyn ExprHumanizer>,
1497    ) -> String {
1498        text_string_at(self, || PlanRenderingContext {
1499            indent: Indent::default(),
1500            humanizer: humanizer.unwrap_or(&DummyHumanizer),
1501            annotations: BTreeMap::default(),
1502            config,
1503            ambiguous_ids: BTreeSet::default(),
1504        })
1505    }
1506
1507    /// Take ownership of `self`, leaving an empty `MirRelationExpr::Constant` with the optionally
1508    /// given scalar types. The given scalar types should be `base_eq` with the types that `typ()`
1509    /// would find. Keys and nullability are ignored in the given `SqlRelationType`, and instead we set
1510    /// the best possible key and nullability, since we are making an empty collection.
1511    ///
1512    /// If `typ` is not given, then this calls `.typ()` (which is possibly expensive) to determine
1513    /// the correct type.
1514    pub fn take_safely(&mut self, typ: Option<ReprRelationType>) -> MirRelationExpr {
1515        if let Some(typ) = &typ {
1516            let self_typ = self.typ();
1517            soft_assert_no_log!(
1518                self_typ
1519                    .column_types
1520                    .iter()
1521                    .zip_eq(typ.column_types.iter())
1522                    .all(|(t1, t2)| t1.scalar_type == t2.scalar_type)
1523            );
1524        }
1525        let mut typ = typ.unwrap_or_else(|| self.typ());
1526        typ.keys = vec![vec![]];
1527        for ct in typ.column_types.iter_mut() {
1528            ct.nullable = false;
1529        }
1530        std::mem::replace(
1531            self,
1532            MirRelationExpr::Constant {
1533                rows: Ok(vec![]),
1534                typ,
1535            },
1536        )
1537    }
1538
1539    /// Take ownership of `self`, leaving an empty `MirRelationExpr::Constant` with the given scalar
1540    /// types. Nullability is ignored in the given `SqlColumnType`s, and instead we set the best
1541    /// possible nullability, since we are making an empty collection.
1542    pub fn take_safely_with_sql_col_types(&mut self, typ: Vec<SqlColumnType>) -> MirRelationExpr {
1543        self.take_safely(Some(ReprRelationType::from(&SqlRelationType::new(typ))))
1544    }
1545
1546    /// Like [`Self::take_safely_with_col_types`], but accepts `Vec<ReprColumnType>`.
1547    ///
1548    /// This is the preferred entry point for optimizer transforms, where repr
1549    /// types are the native currency. Internally converts to [`SqlColumnType`]
1550    /// and delegates to [`Self::take_safely_with_col_types`].
1551    pub fn take_safely_with_col_types(&mut self, typ: Vec<ReprColumnType>) -> MirRelationExpr {
1552        self.take_safely(Some(ReprRelationType::new(typ)))
1553    }
1554
1555    /// Take ownership of `self`, leaving an empty `MirRelationExpr::Constant` with an **incorrect** type.
1556    ///
1557    /// This should only be used if `self` is about to be dropped or otherwise overwritten.
1558    pub fn take_dangerous(&mut self) -> MirRelationExpr {
1559        let empty = MirRelationExpr::Constant {
1560            rows: Ok(vec![]),
1561            typ: ReprRelationType::new(Vec::new()),
1562        };
1563        std::mem::replace(self, empty)
1564    }
1565
1566    /// Replaces `self` with some logic applied to `self`.
1567    pub fn replace_using<F>(&mut self, logic: F)
1568    where
1569        F: FnOnce(MirRelationExpr) -> MirRelationExpr,
1570    {
1571        let empty = MirRelationExpr::Constant {
1572            rows: Ok(vec![]),
1573            typ: ReprRelationType::new(Vec::new()),
1574        };
1575        let expr = std::mem::replace(self, empty);
1576        *self = logic(expr);
1577    }
1578
1579    /// Store `self` in a `Let` and pass the corresponding `Get` to `body`.
1580    pub fn let_in<Body, E>(self, id_gen: &mut IdGen, body: Body) -> Result<MirRelationExpr, E>
1581    where
1582        Body: FnOnce(&mut IdGen, MirRelationExpr) -> Result<MirRelationExpr, E>,
1583    {
1584        let (get, id) = self.let_get(id_gen);
1585        let body = body(id_gen, get)?;
1586        Ok(self.let_bind(id, body))
1587    }
1588
1589    /// Returns a `Get` that stands for `self`, and the local id that `Get` refers to. A plan may
1590    /// contain the `Get` only inside a `Let` that binds that id to `self`, which
1591    /// [`MirRelationExpr::let_bind`] builds. If `self` is already a `Get`, returns a copy of
1592    /// `self` and `None`, and no `Let` is needed.
1593    ///
1594    /// The type of a new `Get` has no unique keys. (Keys carried in the `Get` would compound in
1595    /// the types of nested filters with equality predicates over such `Get`s, exponentially in the
1596    /// nesting depth, and `NormalizeLets` derives them from the `Let`'s value anyway.)
1597    ///
1598    /// Usually used together with `let_bind` through [`MirRelationExpr::let_in`].
1599    pub fn let_get(&self, id_gen: &mut IdGen) -> (MirRelationExpr, Option<LocalId>) {
1600        if let MirRelationExpr::Get { .. } = self {
1601            (self.clone(), None)
1602        } else {
1603            let id = LocalId::new(id_gen.allocate_id());
1604            let get = MirRelationExpr::Get {
1605                id: Id::Local(id),
1606                // TODO(SQL-427): Once key inference no longer enumerates key variants, this can
1607                // carry `self.typ()` with its keys.
1608                typ: ReprRelationType::new(self.typ().column_types),
1609                access_strategy: AccessStrategy::UnknownOrLocal,
1610            };
1611            (get, Some(id))
1612        }
1613    }
1614
1615    /// Wraps `body` in a `Let` that binds `self` to `id`, where `id` is what
1616    /// [`MirRelationExpr::let_get`] returned for `self`, and `self` has not changed since. For
1617    /// `None`, returns `body`.
1618    pub fn let_bind(self, id: Option<LocalId>, body: MirRelationExpr) -> MirRelationExpr {
1619        match id {
1620            None => body,
1621            Some(id) => MirRelationExpr::Let {
1622                id,
1623                value: Box::new(self),
1624                body: Box::new(body),
1625            },
1626        }
1627    }
1628
1629    /// Return every row in `self` that does not have a matching row in the first columns of `keys_and_values`, using `default` to fill in the remaining columns
1630    /// (If `default` is a row of nulls, this is the 'outer' part of LEFT OUTER JOIN)
1631    pub fn anti_lookup<E>(
1632        self,
1633        id_gen: &mut IdGen,
1634        keys_and_values: MirRelationExpr,
1635        default: Vec<(Datum, ReprScalarType)>,
1636    ) -> Result<MirRelationExpr, E> {
1637        let (data, column_types): (Vec<_>, Vec<_>) = default
1638            .into_iter()
1639            .map(|(datum, scalar_type)| {
1640                (
1641                    datum,
1642                    ReprColumnType {
1643                        scalar_type,
1644                        nullable: datum.is_null(),
1645                    },
1646                )
1647            })
1648            .unzip();
1649        assert_eq!(keys_and_values.arity() - self.arity(), data.len());
1650        self.let_in(id_gen, |_id_gen, get_keys| {
1651            let get_keys_arity = get_keys.arity();
1652            Ok(MirRelationExpr::join(
1653                vec![
1654                    // all the missing keys (with count 1)
1655                    keys_and_values
1656                        .distinct_by((0..get_keys_arity).collect())
1657                        .negate()
1658                        .union(get_keys.clone().distinct()),
1659                    // join with keys to get the correct counts
1660                    get_keys.clone(),
1661                ],
1662                (0..get_keys_arity).map(|i| vec![(0, i), (1, i)]).collect(),
1663            )
1664            // get rid of the extra copies of columns from keys
1665            .project((0..get_keys_arity).collect())
1666            // This join is logically equivalent to
1667            // `.map(<default_expr>)`, but using a join allows for
1668            // potential predicate pushdown and elision in the
1669            // optimizer.
1670            .product(MirRelationExpr::constant(
1671                vec![data],
1672                ReprRelationType::new(column_types),
1673            )))
1674        })
1675    }
1676
1677    /// Return:
1678    /// * every row in keys_and_values
1679    /// * every row in `self` that does not have a matching row in the first columns of
1680    ///   `keys_and_values`, using `default` to fill in the remaining columns
1681    /// (This is LEFT OUTER JOIN if:
1682    /// 1) `default` is a row of null
1683    /// 2) matching rows in `keys_and_values` and `self` have the same multiplicity.)
1684    pub fn lookup<E>(
1685        self,
1686        id_gen: &mut IdGen,
1687        keys_and_values: MirRelationExpr,
1688        default: Vec<(Datum<'static>, ReprScalarType)>,
1689    ) -> Result<MirRelationExpr, E> {
1690        keys_and_values.let_in(id_gen, |id_gen, get_keys_and_values| {
1691            Ok(get_keys_and_values.clone().union(self.anti_lookup(
1692                id_gen,
1693                get_keys_and_values,
1694                default,
1695            )?))
1696        })
1697    }
1698
1699    /// True iff the expression contains a `NullaryFunc::MzLogicalTimestamp`.
1700    pub fn contains_temporal(&self) -> bool {
1701        let mut contains = false;
1702        self.visit_scalars(&mut |e| contains = contains || e.contains_temporal());
1703        contains
1704    }
1705
1706    /// Fallible visitor for the [`MirScalarExpr`]s directly owned by this relation expression.
1707    ///
1708    /// The `f` visitor should not recursively descend into owned [`MirRelationExpr`]s.
1709    pub fn try_visit_scalars_mut1<F, E>(&mut self, f: &mut F) -> Result<(), E>
1710    where
1711        F: FnMut(&mut MirScalarExpr) -> Result<(), E>,
1712    {
1713        use MirRelationExpr::*;
1714        match self {
1715            Map { scalars, .. } => {
1716                for s in scalars {
1717                    f(s)?;
1718                }
1719            }
1720            Filter { predicates, .. } => {
1721                for p in predicates {
1722                    f(p)?;
1723                }
1724            }
1725            FlatMap { exprs, .. } => {
1726                for expr in exprs {
1727                    f(expr)?;
1728                }
1729            }
1730            Join {
1731                inputs: _,
1732                equivalences,
1733                implementation,
1734            } => {
1735                for equivalence in equivalences {
1736                    for expr in equivalence {
1737                        f(expr)?;
1738                    }
1739                }
1740                match implementation {
1741                    JoinImplementation::Differential((_, start_key, _), order) => {
1742                        if let Some(start_key) = start_key {
1743                            for k in start_key {
1744                                f(k)?;
1745                            }
1746                        }
1747                        for (_, lookup_key, _) in order {
1748                            for k in lookup_key {
1749                                f(k)?;
1750                            }
1751                        }
1752                    }
1753                    JoinImplementation::DeltaQuery(paths) => {
1754                        for path in paths {
1755                            for (_, lookup_key, _) in path {
1756                                for k in lookup_key {
1757                                    f(k)?;
1758                                }
1759                            }
1760                        }
1761                    }
1762                    JoinImplementation::IndexedFilter(_coll_id, _idx_id, index_key, _) => {
1763                        for k in index_key {
1764                            f(k)?;
1765                        }
1766                    }
1767                    JoinImplementation::Unimplemented => {} // No scalar exprs
1768                }
1769            }
1770            ArrangeBy { keys, .. } => {
1771                for key in keys {
1772                    for s in key {
1773                        f(s)?;
1774                    }
1775                }
1776            }
1777            Reduce {
1778                group_key,
1779                aggregates,
1780                ..
1781            } => {
1782                for s in group_key {
1783                    f(s)?;
1784                }
1785                for agg in aggregates {
1786                    f(&mut agg.expr)?;
1787                }
1788            }
1789            TopK { limit, .. } => {
1790                if let Some(s) = limit {
1791                    f(s)?;
1792                }
1793            }
1794            Constant { .. }
1795            | Get { .. }
1796            | Let { .. }
1797            | LetRec { .. }
1798            | Project { .. }
1799            | Negate { .. }
1800            | Threshold { .. }
1801            | Union { .. } => (),
1802        }
1803        Ok(())
1804    }
1805
1806    /// Fallible mutable visitor for the [`MirScalarExpr`]s in the [`MirRelationExpr`] subtree
1807    /// rooted at `self`.
1808    ///
1809    /// Note that this does not recurse into [`MirRelationExpr`] subtrees within [`MirScalarExpr`]
1810    /// nodes.
1811    pub fn try_visit_scalars_mut<F, E>(&mut self, f: &mut F) -> Result<(), E>
1812    where
1813        F: FnMut(&mut MirScalarExpr) -> Result<(), E>,
1814    {
1815        self.try_visit_mut_post(&mut |expr| expr.try_visit_scalars_mut1(f))
1816    }
1817
1818    /// Infallible mutable visitor for the [`MirScalarExpr`]s in the [`MirRelationExpr`] subtree
1819    /// rooted at `self`.
1820    ///
1821    /// Note that this does not recurse into [`MirRelationExpr`] subtrees within [`MirScalarExpr`]
1822    /// nodes.
1823    pub fn visit_scalars_mut<F>(&mut self, f: &mut F)
1824    where
1825        F: FnMut(&mut MirScalarExpr),
1826    {
1827        self.try_visit_scalars_mut(&mut |s| {
1828            f(s);
1829            Ok::<_, RecursionLimitError>(())
1830        })
1831        .expect("Unexpected error in `visit_scalars_mut` call");
1832    }
1833
1834    /// Fallible visitor for the [`MirScalarExpr`]s directly owned by this relation expression.
1835    ///
1836    /// The `f` visitor should not recursively descend into owned [`MirRelationExpr`]s.
1837    pub fn try_visit_scalars_1<F, E>(&self, f: &mut F) -> Result<(), E>
1838    where
1839        F: FnMut(&MirScalarExpr) -> Result<(), E>,
1840    {
1841        use MirRelationExpr::*;
1842        match self {
1843            Map { scalars, .. } => {
1844                for s in scalars {
1845                    f(s)?;
1846                }
1847            }
1848            Filter { predicates, .. } => {
1849                for p in predicates {
1850                    f(p)?;
1851                }
1852            }
1853            FlatMap { exprs, .. } => {
1854                for expr in exprs {
1855                    f(expr)?;
1856                }
1857            }
1858            Join {
1859                inputs: _,
1860                equivalences,
1861                implementation,
1862            } => {
1863                for equivalence in equivalences {
1864                    for expr in equivalence {
1865                        f(expr)?;
1866                    }
1867                }
1868                match implementation {
1869                    JoinImplementation::Differential((_, start_key, _), order) => {
1870                        if let Some(start_key) = start_key {
1871                            for k in start_key {
1872                                f(k)?;
1873                            }
1874                        }
1875                        for (_, lookup_key, _) in order {
1876                            for k in lookup_key {
1877                                f(k)?;
1878                            }
1879                        }
1880                    }
1881                    JoinImplementation::DeltaQuery(paths) => {
1882                        for path in paths {
1883                            for (_, lookup_key, _) in path {
1884                                for k in lookup_key {
1885                                    f(k)?;
1886                                }
1887                            }
1888                        }
1889                    }
1890                    JoinImplementation::IndexedFilter(_coll_id, _idx_id, index_key, _) => {
1891                        for k in index_key {
1892                            f(k)?;
1893                        }
1894                    }
1895                    JoinImplementation::Unimplemented => {} // No scalar exprs
1896                }
1897            }
1898            ArrangeBy { keys, .. } => {
1899                for key in keys {
1900                    for s in key {
1901                        f(s)?;
1902                    }
1903                }
1904            }
1905            Reduce {
1906                group_key,
1907                aggregates,
1908                ..
1909            } => {
1910                for s in group_key {
1911                    f(s)?;
1912                }
1913                for agg in aggregates {
1914                    f(&agg.expr)?;
1915                }
1916            }
1917            TopK { limit, .. } => {
1918                if let Some(s) = limit {
1919                    f(s)?;
1920                }
1921            }
1922            Constant { .. }
1923            | Get { .. }
1924            | Let { .. }
1925            | LetRec { .. }
1926            | Project { .. }
1927            | Negate { .. }
1928            | Threshold { .. }
1929            | Union { .. } => (),
1930        }
1931        Ok(())
1932    }
1933
1934    /// Fallible immutable visitor for the [`MirScalarExpr`]s in the [`MirRelationExpr`] subtree
1935    /// rooted at `self`.
1936    ///
1937    /// Note that this does not recurse into [`MirRelationExpr`] subtrees within [`MirScalarExpr`]
1938    /// nodes.
1939    pub fn try_visit_scalars<F, E>(&self, f: &mut F) -> Result<(), E>
1940    where
1941        F: FnMut(&MirScalarExpr) -> Result<(), E>,
1942    {
1943        self.try_visit_post(&mut |expr| expr.try_visit_scalars_1(f))
1944    }
1945
1946    /// Infallible immutable visitor for the [`MirScalarExpr`]s in the [`MirRelationExpr`] subtree
1947    /// rooted at `self`.
1948    ///
1949    /// Note that this does not recurse into [`MirRelationExpr`] subtrees within [`MirScalarExpr`]
1950    /// nodes.
1951    pub fn visit_scalars<F>(&self, f: &mut F)
1952    where
1953        F: FnMut(&MirScalarExpr),
1954    {
1955        self.try_visit_scalars(&mut |s| {
1956            f(s);
1957            Ok::<_, RecursionLimitError>(())
1958        })
1959        .expect("Unexpected error in `visit_scalars` call");
1960    }
1961
1962    /// Clears the contents of `self` even if it's so deep that simply dropping it would cause a
1963    /// stack overflow in `drop_in_place`.
1964    ///
1965    /// Leaves `self` in an unusable state, so this should only be used if `self` is about to be
1966    /// dropped or otherwise overwritten.
1967    pub fn destroy_carefully(&mut self) {
1968        let mut todo = vec![self.take_dangerous()];
1969        while let Some(mut expr) = todo.pop() {
1970            for child in expr.children_mut() {
1971                todo.push(child.take_dangerous());
1972            }
1973        }
1974    }
1975
1976    /// Computes the size (total number of nodes) and maximum depth of a MirRelationExpr for
1977    /// debug printing purposes.
1978    pub fn debug_size_and_depth(&self) -> (usize, usize) {
1979        let mut size = 0;
1980        let mut max_depth = 0;
1981        let mut todo = vec![(self, 1)];
1982        while let Some((expr, depth)) = todo.pop() {
1983            size += 1;
1984            max_depth = max(max_depth, depth);
1985            todo.extend(expr.children().map(|c| (c, depth + 1)));
1986        }
1987        (size, max_depth)
1988    }
1989
1990    /// The MirRelationExpr is considered potentially expensive if and only if
1991    /// at least one of the following conditions is true:
1992    ///
1993    ///  - It contains at least one MirScalarExpr with a function call.
1994    ///  - It contains at least one FlatMap or a Reduce operator.
1995    ///  - We run into a RecursionLimitError while analyzing the expression.
1996    ///
1997    /// !!!WARNING!!!: this method has an HirRelationExpr counterpart. The two
1998    /// should be kept in sync w.r.t. HIR ⇒ MIR lowering!
1999    pub fn could_run_expensive_function(&self) -> bool {
2000        let mut result = false;
2001        use MirRelationExpr::*;
2002        use MirScalarExpr::*;
2003        if let Err(_) = self.try_visit_scalars::<_, RecursionLimitError>(&mut |scalar| {
2004            result |= match scalar {
2005                Column(_, _) | Literal(_, _) | CallUnmaterializable(_) | If { .. } => false,
2006                // Function calls are considered expensive
2007                CallUnary { .. } | CallBinary { .. } | CallVariadic { .. } => true,
2008            };
2009            Ok(())
2010        }) {
2011            // Conservatively set `true` if on RecursionLimitError.
2012            result = true;
2013        }
2014        self.visit_pre(|e: &MirRelationExpr| {
2015            // FlatMap has a table function; Reduce has an aggregate function.
2016            // Other constructs use MirScalarExpr to run a function
2017            result |= matches!(e, FlatMap { .. } | Reduce { .. });
2018        });
2019        result
2020    }
2021
2022    /// Hash to an u64 using Rust's default Hasher. (Which is a somewhat slower, but better Hasher
2023    /// than what `Hashable::hashed` would give us.)
2024    pub fn hash_to_u64(&self) -> u64 {
2025        let mut h = DefaultHasher::new();
2026        self.hash(&mut h);
2027        h.finish()
2028    }
2029}
2030
2031// `LetRec` helpers
2032impl MirRelationExpr {
2033    /// True when `expr` contains a `LetRec` AST node.
2034    pub fn is_recursive(self: &MirRelationExpr) -> bool {
2035        let mut worklist = vec![self];
2036        while let Some(expr) = worklist.pop() {
2037            if let MirRelationExpr::LetRec { .. } = expr {
2038                return true;
2039            }
2040            worklist.extend(expr.children());
2041        }
2042        false
2043    }
2044
2045    /// Return the number of sub-expressions in the tree (including self).
2046    pub fn size(&self) -> usize {
2047        let mut size = 0;
2048        self.visit_pre(|_| size += 1);
2049        size
2050    }
2051
2052    /// Given the ids and values of a LetRec, it computes the subset of ids that are used across
2053    /// iterations. These are those ids that have a reference before they are defined, when reading
2054    /// all the bindings in order.
2055    ///
2056    /// For example:
2057    /// ```SQL
2058    /// WITH MUTUALLY RECURSIVE
2059    ///     x(...) AS f(z),
2060    ///     y(...) AS g(x),
2061    ///     z(...) AS h(y)
2062    /// ...;
2063    /// ```
2064    /// Here, only `z` is returned, because `x` and `y` are referenced only within the same
2065    /// iteration.
2066    ///
2067    /// Note that if a binding references itself, that is also returned.
2068    pub fn recursive_ids(ids: &[LocalId], values: &[MirRelationExpr]) -> BTreeSet<LocalId> {
2069        let mut used_across_iterations = BTreeSet::new();
2070        let mut defined = BTreeSet::new();
2071        for (binding_id, value) in itertools::zip_eq(ids.iter(), values.iter()) {
2072            value.visit_pre(|expr| {
2073                if let MirRelationExpr::Get {
2074                    id: Local(get_id), ..
2075                } = expr
2076                {
2077                    // If we haven't seen a definition for it yet, then this will refer
2078                    // to the previous iteration.
2079                    // The `ids.contains` part of the condition is needed to exclude
2080                    // those ids that are not really in this LetRec, but either an inner
2081                    // or outer one.
2082                    if !defined.contains(get_id) && ids.contains(get_id) {
2083                        used_across_iterations.insert(*get_id);
2084                    }
2085                }
2086            });
2087            defined.insert(*binding_id);
2088        }
2089        used_across_iterations
2090    }
2091
2092    /// Replaces `LetRec` nodes with a stack of `Let` nodes.
2093    ///
2094    /// In each `Let` binding, uses of `Get` in `value` that are not at strictly greater
2095    /// identifiers are rewritten to be the constant collection.
2096    /// This makes the computation perform exactly "one" iteration.
2097    ///
2098    /// This was used only temporarily while developing `LetRec`.
2099    pub fn make_nonrecursive(self: &mut MirRelationExpr) {
2100        let mut deadlist = BTreeSet::new();
2101        let mut worklist = vec![self];
2102        while let Some(expr) = worklist.pop() {
2103            if let MirRelationExpr::LetRec {
2104                ids,
2105                values,
2106                limits: _,
2107                body,
2108            } = expr
2109            {
2110                let ids_values = values
2111                    .drain(..)
2112                    .zip_eq(ids)
2113                    .map(|(value, id)| (*id, value))
2114                    .collect::<Vec<_>>();
2115                *expr = body.take_dangerous();
2116                for (id, mut value) in ids_values.into_iter().rev() {
2117                    // Remove references to potentially recursive identifiers.
2118                    deadlist.insert(id);
2119                    value.visit_pre_mut(|e| {
2120                        if let MirRelationExpr::Get {
2121                            id: crate::Id::Local(id),
2122                            typ,
2123                            ..
2124                        } = e
2125                        {
2126                            let typ = typ.clone();
2127                            if deadlist.contains(id) {
2128                                e.take_safely(Some(typ));
2129                            }
2130                        }
2131                    });
2132                    *expr = MirRelationExpr::Let {
2133                        id,
2134                        value: Box::new(value),
2135                        body: Box::new(expr.take_dangerous()),
2136                    };
2137                }
2138                worklist.push(expr);
2139            } else {
2140                worklist.extend(expr.children_mut().rev());
2141            }
2142        }
2143    }
2144
2145    /// For each Id `id'` referenced in `expr`, if it is larger or equal than `id`, then record in
2146    /// `expire_whens` that when `id'` is redefined, then we should expire the information that
2147    /// we are holding about `id`. Call `do_expirations` with `expire_whens` at each Id
2148    /// redefinition.
2149    ///
2150    /// IMPORTANT: Relies on the numbering of Ids to be what `renumber_bindings` gives.
2151    pub fn collect_expirations(
2152        id: LocalId,
2153        expr: &MirRelationExpr,
2154        expire_whens: &mut BTreeMap<LocalId, Vec<LocalId>>,
2155    ) {
2156        expr.visit_pre(|e| {
2157            if let MirRelationExpr::Get {
2158                id: Id::Local(referenced_id),
2159                ..
2160            } = e
2161            {
2162                // The following check needs `renumber_bindings` to have run recently
2163                if referenced_id >= &id {
2164                    expire_whens
2165                        .entry(*referenced_id)
2166                        .or_insert_with(Vec::new)
2167                        .push(id);
2168                }
2169            }
2170        });
2171    }
2172
2173    /// Call this function when `id` is redefined. It modifies `id_infos` by removing information
2174    /// about such Ids whose information depended on the earlier definition of `id`, according to
2175    /// `expire_whens`. Also modifies `expire_whens`: it removes the currently processed entry.
2176    pub fn do_expirations<I>(
2177        redefined_id: LocalId,
2178        expire_whens: &mut BTreeMap<LocalId, Vec<LocalId>>,
2179        id_infos: &mut BTreeMap<LocalId, I>,
2180    ) -> Vec<(LocalId, I)> {
2181        let mut expired_infos = Vec::new();
2182        if let Some(expirations) = expire_whens.remove(&redefined_id) {
2183            for expired_id in expirations.into_iter() {
2184                if let Some(offer) = id_infos.remove(&expired_id) {
2185                    expired_infos.push((expired_id, offer));
2186                }
2187            }
2188        }
2189        expired_infos
2190    }
2191}
2192/// Augment non-nullability of columns, by observing either
2193/// 1. Predicates that explicitly test for null values, and
2194/// 2. Columns that if null would make a predicate be null.
2195pub fn non_nullable_columns(predicates: &[MirScalarExpr]) -> BTreeSet<usize> {
2196    let mut nonnull_required_columns = BTreeSet::new();
2197    for predicate in predicates {
2198        // Add any columns that being null would force the predicate to be null.
2199        // Should that happen, the row would be discarded.
2200        predicate.non_null_requirements(&mut nonnull_required_columns);
2201
2202        /*
2203        Test for explicit checks that a column is non-null.
2204
2205        This analysis is ad hoc, and will miss things:
2206
2207        materialize=> create table a(x int, y int);
2208        CREATE TABLE
2209        materialize=> explain with(types) select x from a where (y=x and y is not null) or x is not null;
2210        Optimized Plan
2211        --------------------------------------------------------------------------------------------------------
2212        Explained Query:                                                                                      +
2213        Project (#0) // { types: "(integer?)" }                                                             +
2214        Filter ((#0) IS NOT NULL OR ((#1) IS NOT NULL AND (#0 = #1))) // { types: "(integer?, integer?)" }+
2215        Get materialize.public.a // { types: "(integer?, integer?)" }                                   +
2216                                                                                  +
2217        Source materialize.public.a                                                                           +
2218        filter=(((#0) IS NOT NULL OR ((#1) IS NOT NULL AND (#0 = #1))))                                     +
2219
2220        (1 row)
2221        */
2222
2223        if let MirScalarExpr::CallUnary {
2224            func: UnaryFunc::Not(scalar_func::Not),
2225            expr,
2226        } = predicate
2227        {
2228            if let MirScalarExpr::CallUnary {
2229                func: UnaryFunc::IsNull(scalar_func::IsNull),
2230                expr,
2231            } = &**expr
2232            {
2233                if let MirScalarExpr::Column(c, _name) = &**expr {
2234                    nonnull_required_columns.insert(*c);
2235                }
2236            }
2237        }
2238    }
2239
2240    nonnull_required_columns
2241}
2242
2243impl CollectionPlan for MirRelationExpr {
2244    /// Collects the global collections that this MIR expression directly depends on, i.e., that it
2245    /// has a `Get` for. (It does _not_ traverse view definitions transitively.)
2246    ///
2247    /// !!!WARNING!!!: this method has an HirRelationExpr counterpart. The two
2248    /// should be kept in sync w.r.t. HIR ⇒ MIR lowering!
2249    fn depends_on_into(&self, out: &mut BTreeSet<GlobalId>) {
2250        if let MirRelationExpr::Get {
2251            id: Id::Global(id), ..
2252        } = self
2253        {
2254            out.insert(*id);
2255        }
2256        self.visit_children(|expr| expr.depends_on_into(out))
2257    }
2258}
2259
2260impl MirRelationExpr {
2261    /// Iterates through references to child expressions.
2262    pub fn children(&self) -> impl DoubleEndedIterator<Item = &Self> {
2263        let mut first = None;
2264        let mut second = None;
2265        let mut rest = None;
2266        let mut last = None;
2267
2268        use MirRelationExpr::*;
2269        match self {
2270            Constant { .. } | Get { .. } => (),
2271            Let { value, body, .. } => {
2272                first = Some(&**value);
2273                second = Some(&**body);
2274            }
2275            LetRec { values, body, .. } => {
2276                rest = Some(values);
2277                last = Some(&**body);
2278            }
2279            Project { input, .. }
2280            | Map { input, .. }
2281            | FlatMap { input, .. }
2282            | Filter { input, .. }
2283            | Reduce { input, .. }
2284            | TopK { input, .. }
2285            | Negate { input }
2286            | Threshold { input }
2287            | ArrangeBy { input, .. } => {
2288                first = Some(&**input);
2289            }
2290            Join { inputs, .. } => {
2291                rest = Some(inputs);
2292            }
2293            Union { base, inputs } => {
2294                first = Some(&**base);
2295                rest = Some(inputs);
2296            }
2297        }
2298
2299        first
2300            .into_iter()
2301            .chain(second)
2302            .chain(rest.into_iter().flatten())
2303            .chain(last)
2304    }
2305
2306    /// Iterates through mutable references to child expressions.
2307    pub fn children_mut(&mut self) -> impl DoubleEndedIterator<Item = &mut Self> {
2308        let mut first = None;
2309        let mut second = None;
2310        let mut rest = None;
2311        let mut last = None;
2312
2313        use MirRelationExpr::*;
2314        match self {
2315            Constant { .. } | Get { .. } => (),
2316            Let { value, body, .. } => {
2317                first = Some(&mut **value);
2318                second = Some(&mut **body);
2319            }
2320            LetRec { values, body, .. } => {
2321                rest = Some(values);
2322                last = Some(&mut **body);
2323            }
2324            Project { input, .. }
2325            | Map { input, .. }
2326            | FlatMap { input, .. }
2327            | Filter { input, .. }
2328            | Reduce { input, .. }
2329            | TopK { input, .. }
2330            | Negate { input }
2331            | Threshold { input }
2332            | ArrangeBy { input, .. } => {
2333                first = Some(&mut **input);
2334            }
2335            Join { inputs, .. } => {
2336                rest = Some(inputs);
2337            }
2338            Union { base, inputs } => {
2339                first = Some(&mut **base);
2340                rest = Some(inputs);
2341            }
2342        }
2343
2344        first
2345            .into_iter()
2346            .chain(second)
2347            .chain(rest.into_iter().flatten())
2348            .chain(last)
2349    }
2350
2351    /// Iterative pre-order visitor.
2352    pub fn visit_pre<'a, F: FnMut(&'a Self)>(&'a self, mut f: F) {
2353        let mut worklist = vec![self];
2354        while let Some(expr) = worklist.pop() {
2355            f(expr);
2356            worklist.extend(expr.children().rev());
2357        }
2358    }
2359
2360    /// Iterative pre-order visitor.
2361    pub fn visit_pre_mut<F: FnMut(&mut Self)>(&mut self, mut f: F) {
2362        let mut worklist = vec![self];
2363        while let Some(expr) = worklist.pop() {
2364            f(expr);
2365            worklist.extend(expr.children_mut().rev());
2366        }
2367    }
2368
2369    /// Return a vector of references to the subtrees of this expression
2370    /// in post-visit order (the last element is `&self`).
2371    pub fn post_order_vec(&self) -> Vec<&Self> {
2372        let mut stack = vec![self];
2373        let mut result = vec![];
2374        while let Some(expr) = stack.pop() {
2375            result.push(expr);
2376            stack.extend(expr.children());
2377        }
2378        result.reverse();
2379        result
2380    }
2381}
2382
2383impl VisitChildren<Self> for MirRelationExpr {
2384    fn visit_children<F>(&self, mut f: F)
2385    where
2386        F: FnMut(&Self),
2387    {
2388        for child in self.children() {
2389            f(child)
2390        }
2391    }
2392
2393    fn visit_mut_children<F>(&mut self, mut f: F)
2394    where
2395        F: FnMut(&mut Self),
2396    {
2397        for child in self.children_mut() {
2398            f(child)
2399        }
2400    }
2401
2402    fn try_visit_children<F, E>(&self, mut f: F) -> Result<(), E>
2403    where
2404        F: FnMut(&Self) -> Result<(), E>,
2405    {
2406        for child in self.children() {
2407            f(child)?
2408        }
2409        Ok(())
2410    }
2411
2412    fn try_visit_mut_children<F, E>(&mut self, mut f: F) -> Result<(), E>
2413    where
2414        F: FnMut(&mut Self) -> Result<(), E>,
2415    {
2416        for child in self.children_mut() {
2417            f(child)?
2418        }
2419        Ok(())
2420    }
2421
2422    fn children<'a>(&'a self) -> impl DoubleEndedIterator<Item = &'a MirRelationExpr>
2423    where
2424        Self: 'a,
2425    {
2426        self.children()
2427    }
2428
2429    fn children_mut<'a>(&'a mut self) -> impl DoubleEndedIterator<Item = &'a mut MirRelationExpr>
2430    where
2431        Self: 'a,
2432    {
2433        self.children_mut()
2434    }
2435}
2436
2437/// Specification for an ordering by a column.
2438#[derive(
2439    Debug,
2440    Clone,
2441    Copy,
2442    Eq,
2443    PartialEq,
2444    Ord,
2445    PartialOrd,
2446    Serialize,
2447    Deserialize,
2448    Hash
2449)]
2450pub struct ColumnOrder {
2451    /// The column index.
2452    pub column: usize,
2453    /// Whether to sort in descending order.
2454    #[serde(default)]
2455    pub desc: bool,
2456    /// Whether to sort nulls last.
2457    #[serde(default)]
2458    pub nulls_last: bool,
2459}
2460
2461impl Columnation for ColumnOrder {
2462    type InnerRegion = CopyRegion<Self>;
2463}
2464
2465impl<'a, M> fmt::Display for HumanizedExpr<'a, ColumnOrder, M>
2466where
2467    M: HumanizerMode,
2468{
2469    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
2470        // If you modify this, then please also attend to Display for ColumnOrderWithExpr!
2471        write!(
2472            f,
2473            "{} {} {}",
2474            self.child(&self.expr.column),
2475            if self.expr.desc { "desc" } else { "asc" },
2476            if self.expr.nulls_last {
2477                "nulls_last"
2478            } else {
2479                "nulls_first"
2480            },
2481        )
2482    }
2483}
2484
2485/// Describes an aggregation expression.
2486#[derive(
2487    Clone,
2488    Debug,
2489    Eq,
2490    PartialEq,
2491    Ord,
2492    PartialOrd,
2493    Serialize,
2494    Deserialize,
2495    Hash
2496)]
2497pub struct AggregateExpr {
2498    /// Names the aggregation function.
2499    pub func: AggregateFunc,
2500    /// An expression which extracts from each row the input to `func`.
2501    pub expr: MirScalarExpr,
2502    /// Should the aggregation be applied only to distinct results in each group.
2503    #[serde(default)]
2504    pub distinct: bool,
2505}
2506
2507impl AggregateExpr {
2508    /// Computes the type of this `AggregateExpr`.
2509    pub fn sql_typ(&self, column_types: &[SqlColumnType]) -> SqlColumnType {
2510        self.func.output_sql_type(self.expr.sql_typ(column_types))
2511    }
2512
2513    /// Computes the type of this `AggregateExpr`.
2514    pub fn typ(&self, column_types: &[ReprColumnType]) -> ReprColumnType {
2515        self.func.output_type(self.expr.typ(column_types))
2516    }
2517
2518    /// Returns whether the expression has a constant result.
2519    pub fn is_constant(&self) -> bool {
2520        match self.func {
2521            AggregateFunc::MaxNumeric
2522            | AggregateFunc::MaxInt16
2523            | AggregateFunc::MaxInt32
2524            | AggregateFunc::MaxInt64
2525            | AggregateFunc::MaxUInt16
2526            | AggregateFunc::MaxUInt32
2527            | AggregateFunc::MaxUInt64
2528            | AggregateFunc::MaxMzTimestamp
2529            | AggregateFunc::MaxFloat32
2530            | AggregateFunc::MaxFloat64
2531            | AggregateFunc::MaxBool
2532            | AggregateFunc::MaxString
2533            | AggregateFunc::MaxDate
2534            | AggregateFunc::MaxTimestamp
2535            | AggregateFunc::MaxTimestampTz
2536            | AggregateFunc::MaxInterval
2537            | AggregateFunc::MaxTime
2538            | AggregateFunc::MaxArray
2539            | AggregateFunc::MaxList
2540            | AggregateFunc::MinNumeric
2541            | AggregateFunc::MinInt16
2542            | AggregateFunc::MinInt32
2543            | AggregateFunc::MinInt64
2544            | AggregateFunc::MinUInt16
2545            | AggregateFunc::MinUInt32
2546            | AggregateFunc::MinUInt64
2547            | AggregateFunc::MinMzTimestamp
2548            | AggregateFunc::MinFloat32
2549            | AggregateFunc::MinFloat64
2550            | AggregateFunc::MinBool
2551            | AggregateFunc::MinString
2552            | AggregateFunc::MinDate
2553            | AggregateFunc::MinTimestamp
2554            | AggregateFunc::MinTimestampTz
2555            | AggregateFunc::MinInterval
2556            | AggregateFunc::MinTime
2557            | AggregateFunc::MinArray
2558            | AggregateFunc::MinList
2559            | AggregateFunc::Any
2560            | AggregateFunc::All
2561            | AggregateFunc::Dummy => self.expr.is_literal(),
2562            AggregateFunc::Count => self.expr.is_literal_null(),
2563            AggregateFunc::SumInt16
2564            | AggregateFunc::SumInt32
2565            | AggregateFunc::SumInt64
2566            | AggregateFunc::SumUInt16
2567            | AggregateFunc::SumUInt32
2568            | AggregateFunc::SumUInt64
2569            | AggregateFunc::SumFloat32
2570            | AggregateFunc::SumFloat64
2571            | AggregateFunc::SumNumeric
2572            | AggregateFunc::SumInterval
2573            | AggregateFunc::JsonbAgg { .. }
2574            | AggregateFunc::JsonbObjectAgg { .. }
2575            | AggregateFunc::MapAgg { .. }
2576            | AggregateFunc::ArrayConcat { .. }
2577            | AggregateFunc::ListConcat { .. }
2578            | AggregateFunc::StringAgg { .. }
2579            | AggregateFunc::RowNumber { .. }
2580            | AggregateFunc::Rank { .. }
2581            | AggregateFunc::DenseRank { .. }
2582            | AggregateFunc::LagLead { .. }
2583            | AggregateFunc::FirstValue { .. }
2584            | AggregateFunc::LastValue { .. }
2585            | AggregateFunc::FusedValueWindowFunc { .. }
2586            | AggregateFunc::WindowAggregate { .. }
2587            | AggregateFunc::FusedWindowAggregate { .. } => self.expr.is_literal_err(),
2588        }
2589    }
2590
2591    /// Returns an expression that computes `self` on a group that has exactly one row.
2592    /// Instead of performing a `Reduce` with `self`, one can perform a `Map` with the expression
2593    /// returned by `on_unique`, which is cheaper. (See `ReduceElision`.)
2594    pub fn on_unique(&self, input_type: &[ReprColumnType]) -> MirScalarExpr {
2595        match &self.func {
2596            // Count is one if non-null, and zero if null.
2597            AggregateFunc::Count => self
2598                .expr
2599                .clone()
2600                .call_unary(UnaryFunc::IsNull(crate::func::IsNull))
2601                .if_then_else(
2602                    MirScalarExpr::literal_ok(Datum::Int64(0), ReprScalarType::Int64),
2603                    MirScalarExpr::literal_ok(Datum::Int64(1), ReprScalarType::Int64),
2604                ),
2605
2606            // SumInt16 takes Int16s as input, but outputs Int64s.
2607            AggregateFunc::SumInt16 => self
2608                .expr
2609                .clone()
2610                .call_unary(UnaryFunc::CastInt16ToInt64(scalar_func::CastInt16ToInt64)),
2611
2612            // SumInt32 takes Int32s as input, but outputs Int64s.
2613            AggregateFunc::SumInt32 => self
2614                .expr
2615                .clone()
2616                .call_unary(UnaryFunc::CastInt32ToInt64(scalar_func::CastInt32ToInt64)),
2617
2618            // SumInt64 takes Int64s as input, but outputs numerics.
2619            AggregateFunc::SumInt64 => self.expr.clone().call_unary(UnaryFunc::CastInt64ToNumeric(
2620                scalar_func::CastInt64ToNumeric(Some(NumericMaxScale::ZERO)),
2621            )),
2622
2623            // SumUInt16 takes UInt16s as input, but outputs UInt64s.
2624            AggregateFunc::SumUInt16 => self.expr.clone().call_unary(
2625                UnaryFunc::CastUint16ToUint64(scalar_func::CastUint16ToUint64),
2626            ),
2627
2628            // SumUInt32 takes UInt32s as input, but outputs UInt64s.
2629            AggregateFunc::SumUInt32 => self.expr.clone().call_unary(
2630                UnaryFunc::CastUint32ToUint64(scalar_func::CastUint32ToUint64),
2631            ),
2632
2633            // SumUInt64 takes UInt64s as input, but outputs numerics.
2634            AggregateFunc::SumUInt64 => {
2635                self.expr.clone().call_unary(UnaryFunc::CastUint64ToNumeric(
2636                    scalar_func::CastUint64ToNumeric(Some(NumericMaxScale::ZERO)),
2637                ))
2638            }
2639
2640            // JsonbAgg takes _anything_ as input, but must output a Jsonb array.
2641            AggregateFunc::JsonbAgg { .. } => MirScalarExpr::call_variadic(
2642                JsonbBuildArray,
2643                vec![
2644                    self.expr
2645                        .clone()
2646                        .call_unary(UnaryFunc::RecordGet(scalar_func::RecordGet(0))),
2647                ],
2648            ),
2649
2650            // JsonbAgg takes _anything_ as input, but must output a Jsonb object.
2651            AggregateFunc::JsonbObjectAgg { .. } => {
2652                let record = self
2653                    .expr
2654                    .clone()
2655                    .call_unary(UnaryFunc::RecordGet(scalar_func::RecordGet(0)));
2656                MirScalarExpr::call_variadic(
2657                    JsonbBuildObject,
2658                    (0..2)
2659                        .map(|i| {
2660                            record
2661                                .clone()
2662                                .call_unary(UnaryFunc::RecordGet(scalar_func::RecordGet(i)))
2663                        })
2664                        .collect(),
2665                )
2666            }
2667
2668            AggregateFunc::MapAgg { value_type, .. } => {
2669                let record = self
2670                    .expr
2671                    .clone()
2672                    .call_unary(UnaryFunc::RecordGet(scalar_func::RecordGet(0)));
2673                MirScalarExpr::call_variadic(
2674                    MapBuild {
2675                        value_type: value_type.clone(),
2676                    },
2677                    (0..2)
2678                        .map(|i| {
2679                            record
2680                                .clone()
2681                                .call_unary(UnaryFunc::RecordGet(scalar_func::RecordGet(i)))
2682                        })
2683                        .collect(),
2684                )
2685            }
2686
2687            // StringAgg takes nested records of strings and outputs a string
2688            AggregateFunc::StringAgg { .. } => self
2689                .expr
2690                .clone()
2691                .call_unary(UnaryFunc::RecordGet(scalar_func::RecordGet(0)))
2692                .call_unary(UnaryFunc::RecordGet(scalar_func::RecordGet(0))),
2693
2694            // ListConcat and ArrayConcat take a single level of records and output a list containing exactly 1 element
2695            AggregateFunc::ListConcat { .. } | AggregateFunc::ArrayConcat { .. } => self
2696                .expr
2697                .clone()
2698                .call_unary(UnaryFunc::RecordGet(scalar_func::RecordGet(0))),
2699
2700            // RowNumber, Rank, DenseRank take a list of records and output a list containing exactly 1 element
2701            AggregateFunc::RowNumber { .. } => {
2702                self.on_unique_ranking_window_funcs(input_type, "?row_number?")
2703            }
2704            AggregateFunc::Rank { .. } => self.on_unique_ranking_window_funcs(input_type, "?rank?"),
2705            AggregateFunc::DenseRank { .. } => {
2706                self.on_unique_ranking_window_funcs(input_type, "?dense_rank?")
2707            }
2708
2709            // The input type for LagLead is ((OriginalRow, (InputValue, Offset, Default)), OrderByExprs...)
2710            AggregateFunc::LagLead { lag_lead, .. } => {
2711                let tuple = self
2712                    .expr
2713                    .clone()
2714                    .call_unary(UnaryFunc::RecordGet(scalar_func::RecordGet(0)));
2715
2716                // Get the overall return type
2717                let return_type_with_orig_row = self
2718                    .typ(input_type)
2719                    .scalar_type
2720                    .unwrap_list_element_type()
2721                    .clone();
2722                let lag_lead_return_type =
2723                    return_type_with_orig_row.unwrap_record_element_type()[0].clone();
2724
2725                // Extract the original row
2726                let original_row = tuple
2727                    .clone()
2728                    .call_unary(UnaryFunc::RecordGet(scalar_func::RecordGet(0)));
2729
2730                // Extract the encoded args
2731                let encoded_args =
2732                    tuple.call_unary(UnaryFunc::RecordGet(scalar_func::RecordGet(1)));
2733
2734                let (result_expr, column_name) =
2735                    Self::on_unique_lag_lead(lag_lead, encoded_args, lag_lead_return_type.clone());
2736
2737                MirScalarExpr::call_variadic(
2738                    ListCreate {
2739                        elem_type: SqlScalarType::from_repr(&return_type_with_orig_row),
2740                    },
2741                    vec![MirScalarExpr::call_variadic(
2742                        RecordCreate {
2743                            field_names: vec![column_name, ColumnName::from("?record?")],
2744                        },
2745                        vec![result_expr, original_row],
2746                    )],
2747                )
2748            }
2749
2750            // The input type for FirstValue is ((OriginalRow, InputValue), OrderByExprs...)
2751            AggregateFunc::FirstValue { window_frame, .. } => {
2752                let tuple = self
2753                    .expr
2754                    .clone()
2755                    .call_unary(UnaryFunc::RecordGet(scalar_func::RecordGet(0)));
2756
2757                // Get the overall return type
2758                let return_type_with_orig_row = self
2759                    .typ(input_type)
2760                    .scalar_type
2761                    .unwrap_list_element_type()
2762                    .clone();
2763                let first_value_return_type =
2764                    return_type_with_orig_row.unwrap_record_element_type()[0].clone();
2765
2766                // Extract the original row
2767                let original_row = tuple
2768                    .clone()
2769                    .call_unary(UnaryFunc::RecordGet(scalar_func::RecordGet(0)));
2770
2771                // Extract the input value
2772                let arg = tuple.call_unary(UnaryFunc::RecordGet(scalar_func::RecordGet(1)));
2773
2774                let (result_expr, column_name) = Self::on_unique_first_value_last_value(
2775                    window_frame,
2776                    arg,
2777                    first_value_return_type,
2778                );
2779
2780                MirScalarExpr::call_variadic(
2781                    ListCreate {
2782                        elem_type: SqlScalarType::from_repr(&return_type_with_orig_row),
2783                    },
2784                    vec![MirScalarExpr::call_variadic(
2785                        RecordCreate {
2786                            field_names: vec![column_name, ColumnName::from("?record?")],
2787                        },
2788                        vec![result_expr, original_row],
2789                    )],
2790                )
2791            }
2792
2793            // The input type for LastValue is ((OriginalRow, InputValue), OrderByExprs...)
2794            AggregateFunc::LastValue { window_frame, .. } => {
2795                let tuple = self
2796                    .expr
2797                    .clone()
2798                    .call_unary(UnaryFunc::RecordGet(scalar_func::RecordGet(0)));
2799
2800                // Get the overall return type
2801                let return_type_with_orig_row = self
2802                    .typ(input_type)
2803                    .scalar_type
2804                    .unwrap_list_element_type()
2805                    .clone();
2806                let last_value_return_type =
2807                    return_type_with_orig_row.unwrap_record_element_type()[0].clone();
2808
2809                // Extract the original row
2810                let original_row = tuple
2811                    .clone()
2812                    .call_unary(UnaryFunc::RecordGet(scalar_func::RecordGet(0)));
2813
2814                // Extract the input value
2815                let arg = tuple.call_unary(UnaryFunc::RecordGet(scalar_func::RecordGet(1)));
2816
2817                let (result_expr, column_name) = Self::on_unique_first_value_last_value(
2818                    window_frame,
2819                    arg,
2820                    last_value_return_type,
2821                );
2822
2823                MirScalarExpr::call_variadic(
2824                    ListCreate {
2825                        elem_type: SqlScalarType::from_repr(&return_type_with_orig_row),
2826                    },
2827                    vec![MirScalarExpr::call_variadic(
2828                        RecordCreate {
2829                            field_names: vec![column_name, ColumnName::from("?record?")],
2830                        },
2831                        vec![result_expr, original_row],
2832                    )],
2833                )
2834            }
2835
2836            // The input type for window aggs is ((OriginalRow, InputValue), OrderByExprs...)
2837            // See an example MIR in `window_func_applied_to`.
2838            AggregateFunc::WindowAggregate {
2839                wrapped_aggregate,
2840                window_frame,
2841                order_by: _,
2842            } => {
2843                // TODO: deduplicate code between the various window function cases.
2844
2845                let tuple = self
2846                    .expr
2847                    .clone()
2848                    .call_unary(UnaryFunc::RecordGet(scalar_func::RecordGet(0)));
2849
2850                // Get the overall return type
2851                let return_type = self
2852                    .typ(input_type)
2853                    .scalar_type
2854                    .unwrap_list_element_type()
2855                    .clone();
2856                let window_agg_return_type = return_type.unwrap_record_element_type()[0].clone();
2857
2858                // Extract the original row
2859                let original_row = tuple
2860                    .clone()
2861                    .call_unary(UnaryFunc::RecordGet(scalar_func::RecordGet(0)));
2862
2863                // Extract the input value
2864                let arg_expr = tuple.call_unary(UnaryFunc::RecordGet(scalar_func::RecordGet(1)));
2865
2866                let (result, column_name) = Self::on_unique_window_agg(
2867                    window_frame,
2868                    arg_expr,
2869                    input_type,
2870                    window_agg_return_type,
2871                    wrapped_aggregate,
2872                );
2873
2874                MirScalarExpr::call_variadic(
2875                    ListCreate {
2876                        elem_type: SqlScalarType::from_repr(&return_type),
2877                    },
2878                    vec![MirScalarExpr::call_variadic(
2879                        RecordCreate {
2880                            field_names: vec![column_name, ColumnName::from("?record?")],
2881                        },
2882                        vec![result, original_row],
2883                    )],
2884                )
2885            }
2886
2887            // The input type is ((OriginalRow, (Arg1, Arg2, ...)), OrderByExprs...)
2888            AggregateFunc::FusedWindowAggregate {
2889                wrapped_aggregates,
2890                order_by: _,
2891                window_frame,
2892            } => {
2893                // Throw away OrderByExprs
2894                let tuple = self
2895                    .expr
2896                    .clone()
2897                    .call_unary(UnaryFunc::RecordGet(scalar_func::RecordGet(0)));
2898
2899                // Extract the original row
2900                let original_row = tuple
2901                    .clone()
2902                    .call_unary(UnaryFunc::RecordGet(scalar_func::RecordGet(0)));
2903
2904                // Extract the args of the fused call
2905                let all_args = tuple.call_unary(UnaryFunc::RecordGet(scalar_func::RecordGet(1)));
2906
2907                let return_type_with_orig_row = self
2908                    .typ(input_type)
2909                    .scalar_type
2910                    .unwrap_list_element_type()
2911                    .clone();
2912
2913                let all_func_return_types =
2914                    return_type_with_orig_row.unwrap_record_element_type()[0].clone();
2915                let mut func_result_exprs = Vec::new();
2916                let mut col_names = Vec::new();
2917                for (idx, wrapped_aggr) in wrapped_aggregates.iter().enumerate() {
2918                    let arg = all_args
2919                        .clone()
2920                        .call_unary(UnaryFunc::RecordGet(scalar_func::RecordGet(idx)));
2921                    let return_type =
2922                        all_func_return_types.unwrap_record_element_type()[idx].clone();
2923                    let (result, column_name) = Self::on_unique_window_agg(
2924                        window_frame,
2925                        arg,
2926                        input_type,
2927                        return_type,
2928                        wrapped_aggr,
2929                    );
2930                    func_result_exprs.push(result);
2931                    col_names.push(column_name);
2932                }
2933
2934                MirScalarExpr::call_variadic(
2935                    ListCreate {
2936                        elem_type: SqlScalarType::from_repr(&return_type_with_orig_row),
2937                    },
2938                    vec![MirScalarExpr::call_variadic(
2939                        RecordCreate {
2940                            field_names: vec![
2941                                ColumnName::from("?fused_window_aggr?"),
2942                                ColumnName::from("?record?"),
2943                            ],
2944                        },
2945                        vec![
2946                            MirScalarExpr::call_variadic(
2947                                RecordCreate {
2948                                    field_names: col_names,
2949                                },
2950                                func_result_exprs,
2951                            ),
2952                            original_row,
2953                        ],
2954                    )],
2955                )
2956            }
2957
2958            // The input type is ((OriginalRow, (Args1, Args2, ...)), OrderByExprs...)
2959            AggregateFunc::FusedValueWindowFunc {
2960                funcs,
2961                order_by: outer_order_by,
2962            } => {
2963                // Throw away OrderByExprs
2964                let tuple = self
2965                    .expr
2966                    .clone()
2967                    .call_unary(UnaryFunc::RecordGet(scalar_func::RecordGet(0)));
2968
2969                // Extract the original row
2970                let original_row = tuple
2971                    .clone()
2972                    .call_unary(UnaryFunc::RecordGet(scalar_func::RecordGet(0)));
2973
2974                // Extract the encoded args of the fused call
2975                let all_encoded_args =
2976                    tuple.call_unary(UnaryFunc::RecordGet(scalar_func::RecordGet(1)));
2977
2978                let return_type_with_orig_row = self
2979                    .typ(input_type)
2980                    .scalar_type
2981                    .unwrap_list_element_type()
2982                    .clone();
2983
2984                let all_func_return_types =
2985                    return_type_with_orig_row.unwrap_record_element_type()[0].clone();
2986                let mut func_result_exprs = Vec::new();
2987                let mut col_names = Vec::new();
2988                for (idx, func) in funcs.iter().enumerate() {
2989                    let args_for_func = all_encoded_args
2990                        .clone()
2991                        .call_unary(UnaryFunc::RecordGet(scalar_func::RecordGet(idx)));
2992                    let return_type_for_func =
2993                        all_func_return_types.unwrap_record_element_type()[idx].clone();
2994                    let (result, column_name) = match func {
2995                        AggregateFunc::LagLead {
2996                            lag_lead,
2997                            order_by,
2998                            ignore_nulls: _,
2999                        } => {
3000                            assert_eq!(order_by, outer_order_by);
3001                            Self::on_unique_lag_lead(lag_lead, args_for_func, return_type_for_func)
3002                        }
3003                        AggregateFunc::FirstValue {
3004                            window_frame,
3005                            order_by,
3006                        } => {
3007                            assert_eq!(order_by, outer_order_by);
3008                            Self::on_unique_first_value_last_value(
3009                                window_frame,
3010                                args_for_func,
3011                                return_type_for_func,
3012                            )
3013                        }
3014                        AggregateFunc::LastValue {
3015                            window_frame,
3016                            order_by,
3017                        } => {
3018                            assert_eq!(order_by, outer_order_by);
3019                            Self::on_unique_first_value_last_value(
3020                                window_frame,
3021                                args_for_func,
3022                                return_type_for_func,
3023                            )
3024                        }
3025                        _ => panic!("unknown function in FusedValueWindowFunc"),
3026                    };
3027                    func_result_exprs.push(result);
3028                    col_names.push(column_name);
3029                }
3030
3031                MirScalarExpr::call_variadic(
3032                    ListCreate {
3033                        elem_type: SqlScalarType::from_repr(&return_type_with_orig_row),
3034                    },
3035                    vec![MirScalarExpr::call_variadic(
3036                        RecordCreate {
3037                            field_names: vec![
3038                                ColumnName::from("?fused_value_window_func?"),
3039                                ColumnName::from("?record?"),
3040                            ],
3041                        },
3042                        vec![
3043                            MirScalarExpr::call_variadic(
3044                                RecordCreate {
3045                                    field_names: col_names,
3046                                },
3047                                func_result_exprs,
3048                            ),
3049                            original_row,
3050                        ],
3051                    )],
3052                )
3053            }
3054
3055            // All other variants should return the argument to the aggregation.
3056            AggregateFunc::MaxNumeric
3057            | AggregateFunc::MaxInt16
3058            | AggregateFunc::MaxInt32
3059            | AggregateFunc::MaxInt64
3060            | AggregateFunc::MaxUInt16
3061            | AggregateFunc::MaxUInt32
3062            | AggregateFunc::MaxUInt64
3063            | AggregateFunc::MaxMzTimestamp
3064            | AggregateFunc::MaxFloat32
3065            | AggregateFunc::MaxFloat64
3066            | AggregateFunc::MaxBool
3067            | AggregateFunc::MaxString
3068            | AggregateFunc::MaxDate
3069            | AggregateFunc::MaxTimestamp
3070            | AggregateFunc::MaxTimestampTz
3071            | AggregateFunc::MaxInterval
3072            | AggregateFunc::MaxTime
3073            | AggregateFunc::MaxArray
3074            | AggregateFunc::MaxList
3075            | AggregateFunc::MinNumeric
3076            | AggregateFunc::MinInt16
3077            | AggregateFunc::MinInt32
3078            | AggregateFunc::MinInt64
3079            | AggregateFunc::MinUInt16
3080            | AggregateFunc::MinUInt32
3081            | AggregateFunc::MinUInt64
3082            | AggregateFunc::MinMzTimestamp
3083            | AggregateFunc::MinFloat32
3084            | AggregateFunc::MinFloat64
3085            | AggregateFunc::MinBool
3086            | AggregateFunc::MinString
3087            | AggregateFunc::MinDate
3088            | AggregateFunc::MinTimestamp
3089            | AggregateFunc::MinTimestampTz
3090            | AggregateFunc::MinInterval
3091            | AggregateFunc::MinTime
3092            | AggregateFunc::MinArray
3093            | AggregateFunc::MinList
3094            | AggregateFunc::SumFloat32
3095            | AggregateFunc::SumFloat64
3096            | AggregateFunc::SumNumeric
3097            | AggregateFunc::SumInterval
3098            | AggregateFunc::Any
3099            | AggregateFunc::All
3100            | AggregateFunc::Dummy => self.expr.clone(),
3101        }
3102    }
3103
3104    /// `on_unique` for ROW_NUMBER, RANK, DENSE_RANK
3105    fn on_unique_ranking_window_funcs(
3106        &self,
3107        input_type: &[ReprColumnType],
3108        col_name: &str,
3109    ) -> MirScalarExpr {
3110        let sql_input_type: Vec<SqlColumnType> =
3111            input_type.iter().map(SqlColumnType::from_repr).collect();
3112        let list = self
3113            .expr
3114            .clone()
3115            // extract the list within the record
3116            .call_unary(UnaryFunc::RecordGet(scalar_func::RecordGet(0)));
3117
3118        // extract the expression within the list
3119        let record = MirScalarExpr::call_variadic(
3120            ListIndex,
3121            vec![
3122                list,
3123                MirScalarExpr::literal_ok(Datum::Int64(1), ReprScalarType::Int64),
3124            ],
3125        );
3126
3127        MirScalarExpr::call_variadic(
3128            ListCreate {
3129                elem_type: self
3130                    .sql_typ(&sql_input_type)
3131                    .scalar_type
3132                    .unwrap_list_element_type()
3133                    .clone(),
3134            },
3135            vec![MirScalarExpr::call_variadic(
3136                RecordCreate {
3137                    field_names: vec![ColumnName::from(col_name), ColumnName::from("?record?")],
3138                },
3139                vec![
3140                    MirScalarExpr::literal_ok(Datum::Int64(1), ReprScalarType::Int64),
3141                    record,
3142                ],
3143            )],
3144        )
3145    }
3146
3147    /// `on_unique` for `lag` and `lead`
3148    fn on_unique_lag_lead(
3149        lag_lead: &LagLeadType,
3150        encoded_args: MirScalarExpr,
3151        return_type: ReprScalarType,
3152    ) -> (MirScalarExpr, ColumnName) {
3153        let expr = encoded_args
3154            .clone()
3155            .call_unary(UnaryFunc::RecordGet(scalar_func::RecordGet(0)));
3156        let offset = encoded_args
3157            .clone()
3158            .call_unary(UnaryFunc::RecordGet(scalar_func::RecordGet(1)));
3159        let default_value =
3160            encoded_args.call_unary(UnaryFunc::RecordGet(scalar_func::RecordGet(2)));
3161
3162        // In this case, the window always has only one element, so if the offset is not null and
3163        // not zero, the default value should be returned instead.
3164        let value = offset
3165            .clone()
3166            .call_binary(
3167                MirScalarExpr::literal_ok(Datum::Int32(0), ReprScalarType::Int32),
3168                crate::func::Eq,
3169            )
3170            .if_then_else(expr, default_value);
3171        let result_expr = offset
3172            .call_unary(UnaryFunc::IsNull(crate::func::IsNull))
3173            .if_then_else(MirScalarExpr::literal_null(return_type), value);
3174
3175        let column_name = ColumnName::from(match lag_lead {
3176            LagLeadType::Lag => "?lag?",
3177            LagLeadType::Lead => "?lead?",
3178        });
3179
3180        (result_expr, column_name)
3181    }
3182
3183    /// `on_unique` for `first_value` and `last_value`
3184    fn on_unique_first_value_last_value(
3185        window_frame: &WindowFrame,
3186        arg: MirScalarExpr,
3187        return_type: ReprScalarType,
3188    ) -> (MirScalarExpr, ColumnName) {
3189        // If the window frame includes the current (single) row, return its value, null otherwise
3190        let result_expr = if window_frame.includes_current_row() {
3191            arg
3192        } else {
3193            MirScalarExpr::literal_null(return_type)
3194        };
3195        (result_expr, ColumnName::from("?first_value?"))
3196    }
3197
3198    /// `on_unique` for window aggregations
3199    fn on_unique_window_agg(
3200        window_frame: &WindowFrame,
3201        arg_expr: MirScalarExpr,
3202        input_type: &[ReprColumnType],
3203        return_type: ReprScalarType,
3204        wrapped_aggr: &AggregateFunc,
3205    ) -> (MirScalarExpr, ColumnName) {
3206        // If the window frame includes the current (single) row, evaluate the wrapped aggregate on
3207        // that row. Otherwise, return the default value for the aggregate.
3208        let result_expr = if window_frame.includes_current_row() {
3209            AggregateExpr {
3210                func: wrapped_aggr.clone(),
3211                expr: arg_expr,
3212                distinct: false, // We have just one input element; DISTINCT doesn't matter.
3213            }
3214            .on_unique(input_type)
3215        } else {
3216            MirScalarExpr::literal_ok(wrapped_aggr.default(), return_type)
3217        };
3218        (result_expr, ColumnName::from("?window_agg?"))
3219    }
3220
3221    /// Returns whether the expression is COUNT(*) or not.  Note that
3222    /// when we define the count builtin in sql::func, we convert
3223    /// COUNT(*) to COUNT(true), making it indistinguishable from
3224    /// literal COUNT(true), but we prefer to consider this as the
3225    /// former.
3226    ///
3227    /// (HIR has the same `is_count_asterisk`.)
3228    pub fn is_count_asterisk(&self) -> bool {
3229        self.func == AggregateFunc::Count && self.expr.is_literal_true() && !self.distinct
3230    }
3231}
3232
3233/// Describe a join implementation in dataflow.
3234#[derive(
3235    Clone,
3236    Debug,
3237    Eq,
3238    PartialEq,
3239    Ord,
3240    PartialOrd,
3241    Serialize,
3242    Deserialize,
3243    Hash
3244)]
3245pub enum JoinImplementation {
3246    /// Perform a sequence of binary differential dataflow joins.
3247    ///
3248    /// The first argument indicates
3249    /// 1) the index of the starting collection,
3250    /// 2) if it should be arranged, the keys to arrange it by, and
3251    /// 3) the characteristics of the starting collection (for EXPLAINing).
3252    /// The sequence that follows lists other relation indexes, and the key for
3253    /// the arrangement we should use when joining it in.
3254    /// The JoinInputCharacteristics are for EXPLAINing the characteristics that
3255    /// were used for join ordering.
3256    ///
3257    /// Each collection index should occur exactly once, either as the starting collection
3258    /// or somewhere in the list.
3259    Differential(
3260        (
3261            usize,
3262            Option<Vec<MirScalarExpr>>,
3263            Option<JoinInputCharacteristics>,
3264        ),
3265        Vec<(usize, Vec<MirScalarExpr>, Option<JoinInputCharacteristics>)>,
3266    ),
3267    /// Perform independent delta query dataflows for each input.
3268    ///
3269    /// The argument is a sequence of plans, for the input collections in order.
3270    /// Each plan starts from the corresponding index, and then in sequence joins
3271    /// against collections identified by index and with the specified arrangement key.
3272    /// The JoinInputCharacteristics are for EXPLAINing the characteristics that were
3273    /// used for join ordering.
3274    DeltaQuery(Vec<Vec<(usize, Vec<MirScalarExpr>, Option<JoinInputCharacteristics>)>>),
3275    /// Join a user-created index with a constant collection to speed up the evaluation of a
3276    /// predicate such as `(f1 = 3 AND f2 = 5) OR (f1 = 7 AND f2 = 9)`.
3277    /// This gets translated to a Differential join during MIR -> LIR lowering, but we still want
3278    /// to represent it in MIR, because the fast path detection wants to match on this.
3279    ///
3280    /// Consists of (`<coll_id>`, `<index_id>`, `<index_key>`, `<constants>`)
3281    IndexedFilter(GlobalId, GlobalId, Vec<MirScalarExpr>, Vec<Row>),
3282    /// No implementation yet selected.
3283    Unimplemented,
3284}
3285
3286impl Default for JoinImplementation {
3287    fn default() -> Self {
3288        JoinImplementation::Unimplemented
3289    }
3290}
3291
3292impl JoinImplementation {
3293    /// Returns `true` iff the value is not [`JoinImplementation::Unimplemented`].
3294    pub fn is_implemented(&self) -> bool {
3295        match self {
3296            Self::Unimplemented => false,
3297            _ => true,
3298        }
3299    }
3300
3301    /// Returns an optional implementation name if the value is not [`JoinImplementation::Unimplemented`].
3302    pub fn name(&self) -> Option<&'static str> {
3303        match self {
3304            Self::Differential(..) => Some("differential"),
3305            Self::DeltaQuery(..) => Some("delta"),
3306            Self::IndexedFilter(..) => Some("indexed_filter"),
3307            Self::Unimplemented => None,
3308        }
3309    }
3310}
3311
3312/// Characteristics of a join order candidate collection.
3313///
3314/// A candidate is described by a collection and a key, and may have various liabilities.
3315/// Primarily, the candidate may risk substantial inflation of records, which is something
3316/// that concerns us greatly. Additionally, the candidate may be unarranged, and we would
3317/// prefer candidates that do not require additional memory. Finally, we prefer lower id
3318/// collections in the interest of consistent tie-breaking. For more characteristics, see
3319/// comments on individual fields.
3320///
3321/// This has more than one version. `new` instantiates the appropriate version based on a
3322/// feature flag.
3323#[derive(
3324    Eq,
3325    PartialEq,
3326    Ord,
3327    PartialOrd,
3328    Debug,
3329    Clone,
3330    Serialize,
3331    Deserialize,
3332    Hash
3333)]
3334pub enum JoinInputCharacteristics {
3335    /// Old version, with `enable_join_prioritize_arranged` turned off.
3336    V1(JoinInputCharacteristicsV1),
3337    /// Newer version, with `enable_join_prioritize_arranged` turned on.
3338    V2(JoinInputCharacteristicsV2),
3339}
3340
3341impl JoinInputCharacteristics {
3342    /// Creates a new instance with the given characteristics.
3343    pub fn new(
3344        unique_key: bool,
3345        key_length: usize,
3346        arranged: bool,
3347        cardinality: Option<usize>,
3348        filters: FilterCharacteristics,
3349        input: usize,
3350        enable_join_prioritize_arranged: bool,
3351    ) -> Self {
3352        if enable_join_prioritize_arranged {
3353            Self::V2(JoinInputCharacteristicsV2::new(
3354                unique_key,
3355                key_length,
3356                arranged,
3357                cardinality,
3358                filters,
3359                input,
3360            ))
3361        } else {
3362            Self::V1(JoinInputCharacteristicsV1::new(
3363                unique_key,
3364                key_length,
3365                arranged,
3366                cardinality,
3367                filters,
3368                input,
3369            ))
3370        }
3371    }
3372
3373    /// Turns the instance into a String to be printed in EXPLAIN.
3374    pub fn explain(&self) -> String {
3375        match self {
3376            Self::V1(jic) => jic.explain(),
3377            Self::V2(jic) => jic.explain(),
3378        }
3379    }
3380
3381    /// Whether the join input described by `self` is arranged.
3382    pub fn arranged(&self) -> bool {
3383        match self {
3384            Self::V1(jic) => jic.arranged,
3385            Self::V2(jic) => jic.arranged,
3386        }
3387    }
3388
3389    /// Returns the `FilterCharacteristics` for the join input described by `self`.
3390    pub fn filters(&mut self) -> &mut FilterCharacteristics {
3391        match self {
3392            Self::V1(jic) => &mut jic.filters,
3393            Self::V2(jic) => &mut jic.filters,
3394        }
3395    }
3396}
3397
3398/// Newer version of `JoinInputCharacteristics`, with `enable_join_prioritize_arranged` turned on.
3399#[derive(
3400    Eq,
3401    PartialEq,
3402    Ord,
3403    PartialOrd,
3404    Debug,
3405    Clone,
3406    Serialize,
3407    Deserialize,
3408    Hash
3409)]
3410pub struct JoinInputCharacteristicsV2 {
3411    /// An excellent indication that record count will not increase.
3412    pub unique_key: bool,
3413    /// Cross joins are bad.
3414    /// (`key_length > 0` also implies that it is not a cross join. However, we need to note cross
3415    /// joins in a separate field, because not being a cross join is more important than `arranged`,
3416    /// but otherwise `key_length` is less important than `arranged`.)
3417    pub not_cross: bool,
3418    /// Indicates that there will be no additional in-memory footprint.
3419    pub arranged: bool,
3420    /// A weaker signal that record count will not increase.
3421    pub key_length: usize,
3422    /// Estimated cardinality (lower is better)
3423    pub cardinality: Option<std::cmp::Reverse<usize>>,
3424    /// Characteristics of the filter that is applied at this input.
3425    pub filters: FilterCharacteristics,
3426    /// We want to prefer input earlier in the input list, for stability of ordering.
3427    pub input: std::cmp::Reverse<usize>,
3428}
3429
3430impl JoinInputCharacteristicsV2 {
3431    /// Creates a new instance with the given characteristics.
3432    pub fn new(
3433        unique_key: bool,
3434        key_length: usize,
3435        arranged: bool,
3436        cardinality: Option<usize>,
3437        filters: FilterCharacteristics,
3438        input: usize,
3439    ) -> Self {
3440        Self {
3441            unique_key,
3442            not_cross: key_length > 0,
3443            arranged,
3444            key_length,
3445            cardinality: cardinality.map(std::cmp::Reverse),
3446            filters,
3447            input: std::cmp::Reverse(input),
3448        }
3449    }
3450
3451    /// Turns the instance into a String to be printed in EXPLAIN.
3452    pub fn explain(&self) -> String {
3453        let mut e = "".to_owned();
3454        if self.unique_key {
3455            e.push_str("U");
3456        }
3457        // Don't need to print `not_cross`, because that is visible in the printed key.
3458        // if !self.not_cross {
3459        //     e.push_str("C");
3460        // }
3461        for _ in 0..self.key_length {
3462            e.push_str("K");
3463        }
3464        if self.arranged {
3465            e.push_str("A");
3466        }
3467        if let Some(std::cmp::Reverse(cardinality)) = self.cardinality {
3468            e.push_str(&format!("|{cardinality}|"));
3469        }
3470        e.push_str(&self.filters.explain());
3471        e
3472    }
3473}
3474
3475/// Old version of `JoinInputCharacteristics`, with `enable_join_prioritize_arranged` turned off.
3476#[derive(
3477    Eq,
3478    PartialEq,
3479    Ord,
3480    PartialOrd,
3481    Debug,
3482    Clone,
3483    Serialize,
3484    Deserialize,
3485    Hash
3486)]
3487pub struct JoinInputCharacteristicsV1 {
3488    /// An excellent indication that record count will not increase.
3489    pub unique_key: bool,
3490    /// A weaker signal that record count will not increase.
3491    pub key_length: usize,
3492    /// Indicates that there will be no additional in-memory footprint.
3493    pub arranged: bool,
3494    /// Estimated cardinality (lower is better)
3495    pub cardinality: Option<std::cmp::Reverse<usize>>,
3496    /// Characteristics of the filter that is applied at this input.
3497    pub filters: FilterCharacteristics,
3498    /// We want to prefer input earlier in the input list, for stability of ordering.
3499    pub input: std::cmp::Reverse<usize>,
3500}
3501
3502impl JoinInputCharacteristicsV1 {
3503    /// Creates a new instance with the given characteristics.
3504    pub fn new(
3505        unique_key: bool,
3506        key_length: usize,
3507        arranged: bool,
3508        cardinality: Option<usize>,
3509        filters: FilterCharacteristics,
3510        input: usize,
3511    ) -> Self {
3512        Self {
3513            unique_key,
3514            key_length,
3515            arranged,
3516            cardinality: cardinality.map(std::cmp::Reverse),
3517            filters,
3518            input: std::cmp::Reverse(input),
3519        }
3520    }
3521
3522    /// Turns the instance into a String to be printed in EXPLAIN.
3523    pub fn explain(&self) -> String {
3524        let mut e = "".to_owned();
3525        if self.unique_key {
3526            e.push_str("U");
3527        }
3528        for _ in 0..self.key_length {
3529            e.push_str("K");
3530        }
3531        if self.arranged {
3532            e.push_str("A");
3533        }
3534        if let Some(std::cmp::Reverse(cardinality)) = self.cardinality {
3535            e.push_str(&format!("|{cardinality}|"));
3536        }
3537        e.push_str(&self.filters.explain());
3538        e
3539    }
3540}
3541
3542/// Instructions for finishing the result of a query.
3543///
3544/// The primary reason for the existence of this structure and attendant code
3545/// is that SQL's ORDER BY requires sorting rows (as already implied by the
3546/// keywords), whereas much of the rest of SQL is defined in terms of unordered
3547/// multisets. But as it turns out, the same idea can be used to optimize
3548/// trivial peeks.
3549///
3550/// The generic parameters are for accommodating prepared statement parameters in
3551/// `limit` and `offset`: the planner can hold these fields as HirScalarExpr long enough to call
3552/// `bind_parameters` on them.
3553#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
3554pub struct RowSetFinishing<L = NonNeg<i64>, O = usize> {
3555    /// Order rows by the given columns.
3556    pub order_by: Vec<ColumnOrder>,
3557    /// Include only as many rows (after offset).
3558    pub limit: Option<L>,
3559    /// Omit as many rows.
3560    pub offset: O,
3561    /// Include only given columns.
3562    pub project: Vec<usize>,
3563}
3564
3565impl<L> RowSetFinishing<L> {
3566    /// Returns a trivial finishing, i.e., that does nothing to the result set.
3567    pub fn trivial(arity: usize) -> RowSetFinishing<L> {
3568        RowSetFinishing {
3569            order_by: Vec::new(),
3570            limit: None,
3571            offset: 0,
3572            project: (0..arity).collect(),
3573        }
3574    }
3575    /// True if the finishing does nothing to any result set.
3576    pub fn is_trivial(&self, arity: usize) -> bool {
3577        self.limit.is_none()
3578            && self.order_by.is_empty()
3579            && self.offset == 0
3580            && self.project.iter().copied().eq(0..arity)
3581    }
3582    /// True if the finishing does not require an ORDER BY.
3583    ///
3584    /// LIMIT and OFFSET without an ORDER BY _are_ streamable: without an
3585    /// explicit ordering we will skip an arbitrary bag of elements and return
3586    /// the first arbitrary elements in the remaining bag. The result semantics
3587    /// are still correct but maybe surprising for some users.
3588    pub fn is_streamable(&self, arity: usize) -> bool {
3589        self.order_by.is_empty() && self.project.iter().copied().eq(0..arity)
3590    }
3591}
3592
3593impl RowSetFinishing<NonNeg<i64>, usize> {
3594    /// The number of rows needed from before the finishing to evaluate the finishing:
3595    /// offset + limit.
3596    ///
3597    /// If it returns None, then we need all the rows.
3598    pub fn num_rows_needed(&self) -> Option<usize> {
3599        self.limit
3600            .as_ref()
3601            .map(|l| usize::cast_from(u64::from(l.clone())) + self.offset)
3602    }
3603}
3604
3605impl RowSetFinishing {
3606    /// Applies finishing actions to a [`RowCollection`], and reports the total
3607    /// time it took to run.
3608    ///
3609    /// Returns a [`RowCollectionIter`] that contains all of the response data, as
3610    /// well as the size of the response in bytes.
3611    pub fn finish(
3612        &self,
3613        rows: RowCollection,
3614        max_result_size: u64,
3615        max_returned_query_size: Option<u64>,
3616        duration_histogram: &Histogram,
3617    ) -> Result<(RowCollectionIter, usize), String> {
3618        let now = Instant::now();
3619        let result = self.finish_inner(rows, max_result_size, max_returned_query_size);
3620        let duration = now.elapsed();
3621        duration_histogram.observe(duration.as_secs_f64());
3622
3623        result
3624    }
3625
3626    /// Implementation for [`RowSetFinishing::finish`].
3627    fn finish_inner(
3628        &self,
3629        rows: RowCollection,
3630        max_result_size: u64,
3631        max_returned_query_size: Option<u64>,
3632    ) -> Result<(RowCollectionIter, usize), String> {
3633        // Bail if the already-materialized collection is larger than the cap.
3634        // `byte_len` includes the per-row offset metadata, and finishing does
3635        // not allocate any additional per-entry structure.
3636        if rows.byte_len() > usize::cast_from(max_result_size) {
3637            let max_bytes = ByteSize::b(max_result_size);
3638            return Err(format!("result exceeds max size of {max_bytes}",));
3639        }
3640
3641        let sorted_view = rows;
3642        let mut iter = sorted_view
3643            .into_row_iter()
3644            .apply_offset(self.offset)
3645            .with_projection(self.project.clone());
3646
3647        if let Some(limit) = self.limit {
3648            let limit = u64::from(limit);
3649            let limit = usize::cast_from(limit);
3650            iter = iter.with_limit(limit);
3651        };
3652
3653        // TODO(parkmycar): Re-think how we can calculate the total response size without
3654        // having to iterate through the entire collection of Rows, while still
3655        // respecting the LIMIT, OFFSET, and projections.
3656        //
3657        // Note: It feels a bit bad always calculating the response size, but we almost
3658        // always need it to either check the `max_returned_query_size`, or for reporting
3659        // in the query history.
3660        let response_size: usize = iter.clone().map(|row| row.data().len()).sum();
3661
3662        // Bail if we would end up returning more data to the client than they can support.
3663        if let Some(max) = max_returned_query_size {
3664            if response_size > usize::cast_from(max) {
3665                let max_bytes = ByteSize::b(max);
3666                return Err(format!("result exceeds max size of {max_bytes}"));
3667            }
3668        }
3669
3670        Ok((iter, response_size))
3671    }
3672}
3673
3674/// A [RowSetFinishing] that can be repeatedly applied to batches of updates (in
3675/// a [RowCollection]) and keeps track of the remaining limit, offset, and cap
3676/// on query result size.
3677#[derive(Debug)]
3678pub struct RowSetFinishingIncremental {
3679    /// Include only as many rows (after offset).
3680    pub remaining_limit: Option<usize>,
3681    /// Omit as many rows.
3682    pub remaining_offset: usize,
3683    /// The maximum allowed result size, as requested by the client.
3684    pub max_returned_query_size: Option<u64>,
3685    /// Tracks our remaining allowed budget for result size.
3686    pub remaining_max_returned_query_size: Option<u64>,
3687    /// Include only given columns.
3688    pub project: Vec<usize>,
3689}
3690
3691impl RowSetFinishingIncremental {
3692    /// Turns the given [RowSetFinishing] into a [RowSetFinishingIncremental].
3693    /// Can only be used when [is_streamable](RowSetFinishing::is_streamable) is
3694    /// `true`.
3695    ///
3696    /// # Panics
3697    ///
3698    /// Panics if the result is not streamable, that is it has an ORDER BY.
3699    pub fn new(
3700        offset: usize,
3701        limit: Option<NonNeg<i64>>,
3702        project: Vec<usize>,
3703        max_returned_query_size: Option<u64>,
3704    ) -> Self {
3705        let limit = limit.map(|l| {
3706            let l = u64::from(l);
3707            let l = usize::cast_from(l);
3708            l
3709        });
3710
3711        RowSetFinishingIncremental {
3712            remaining_limit: limit,
3713            remaining_offset: offset,
3714            max_returned_query_size,
3715            remaining_max_returned_query_size: max_returned_query_size,
3716            project,
3717        }
3718    }
3719
3720    /// Applies finishing actions to the given [`RowCollection`], and reports
3721    /// the total time it took to run.
3722    ///
3723    /// Returns a [`RowCollectionIter`] that contains all of the response
3724    /// data.
3725    pub fn finish_incremental(
3726        &mut self,
3727        rows: RowCollection,
3728        max_result_size: u64,
3729        duration_histogram: &Histogram,
3730    ) -> Result<RowCollectionIter, String> {
3731        let now = Instant::now();
3732        let result = self.finish_incremental_inner(rows, max_result_size);
3733        let duration = now.elapsed();
3734        duration_histogram.observe(duration.as_secs_f64());
3735
3736        result
3737    }
3738
3739    fn finish_incremental_inner(
3740        &mut self,
3741        rows: RowCollection,
3742        max_result_size: u64,
3743    ) -> Result<RowCollectionIter, String> {
3744        // Bail if the already-materialized collection is larger than the cap.
3745        // `byte_len` includes the per-row offset metadata, and finishing does
3746        // not allocate any additional per-entry structure.
3747        if rows.byte_len() > usize::cast_from(max_result_size) {
3748            let max_bytes = ByteSize::b(max_result_size);
3749            return Err(format!("total result exceeds max size of {max_bytes}",));
3750        }
3751
3752        let batch_num_rows = rows.count();
3753
3754        let sorted_view = rows;
3755        let mut iter = sorted_view
3756            .into_row_iter()
3757            .apply_offset(self.remaining_offset)
3758            .with_projection(self.project.clone());
3759
3760        if let Some(limit) = self.remaining_limit {
3761            iter = iter.with_limit(limit);
3762        };
3763
3764        self.remaining_offset = self.remaining_offset.saturating_sub(batch_num_rows);
3765        if let Some(remaining_limit) = self.remaining_limit.as_mut() {
3766            *remaining_limit -= iter.count();
3767        }
3768
3769        // TODO(parkmycar): Re-think how we can calculate the total response size without
3770        // having to iterate through the entire collection of Rows, while still
3771        // respecting the LIMIT, OFFSET, and projections.
3772        //
3773        // Note: It feels a bit bad always calculating the response size, but we almost
3774        // always need it to either check the `max_returned_query_size`, or for reporting
3775        // in the query history.
3776        let response_size: usize = iter.clone().map(|row| row.data().len()).sum();
3777
3778        // Bail if we would end up returning more data to the client than they can support.
3779        if let Some(max) = &mut self.remaining_max_returned_query_size {
3780            if let Some(remaining) = max.checked_sub(response_size.cast_into()) {
3781                *max = remaining;
3782            } else {
3783                let max_bytes = ByteSize::b(self.max_returned_query_size.expect("known to exist"));
3784                return Err(format!("total result exceeds max size of {max_bytes}"));
3785            }
3786        }
3787
3788        Ok(iter)
3789    }
3790}
3791
3792/// Compares two rows columnwise, using [compare_columns].
3793///
3794/// Compared to the naive implementation, this allows sharing some memory and implements some
3795/// optimizations that avoid unnecessary row unpacking.
3796#[derive(Debug, Clone)]
3797pub struct RowComparator<O: AsRef<[ColumnOrder]> = Vec<ColumnOrder>> {
3798    order: O,
3799    /// Invariant: all column references in the order are less than this limit.
3800    /// This allows for partial unpacking of rows.
3801    limit: usize,
3802    left_vec: RefCell<DatumVec>,
3803    right_vec: RefCell<DatumVec>,
3804}
3805
3806impl<O: AsRef<[ColumnOrder]>> RowComparator<O> {
3807    /// Create a new row comparator from the given column ordering.
3808    pub fn new(order: O) -> Self {
3809        let limit = order
3810            .as_ref()
3811            .iter()
3812            .map(|o| o.column + 1)
3813            .max()
3814            .unwrap_or(0);
3815        Self {
3816            order,
3817            limit,
3818            left_vec: Default::default(),
3819            right_vec: Default::default(),
3820        }
3821    }
3822
3823    /// Compare two (references to) rows.
3824    pub fn compare_rows(
3825        &self,
3826        left_row: &RowRef,
3827        right_row: &RowRef,
3828        tiebreaker: impl Fn() -> Ordering,
3829    ) -> Ordering {
3830        let order = if self.limit == 0 {
3831            Ordering::Equal
3832        } else {
3833            // These borrows should never fail, since this struct is non-sync and this function
3834            // is non-recursive.
3835            let mut left_ref = self.left_vec.borrow_mut();
3836            let mut right_ref = self.right_vec.borrow_mut();
3837            let left_cols = left_ref.borrow_with_limit(left_row, self.limit);
3838            let right_cols = right_ref.borrow_with_limit(right_row, self.limit);
3839            compare_columns(self.order.as_ref(), &left_cols, &right_cols, || {
3840                Ordering::Equal
3841            })
3842        };
3843        // Tiebreak without the vecs borrowed, in case that recursively invokes this function.
3844        order.then_with(tiebreaker)
3845    }
3846}
3847
3848/// Compare `left` and `right` using `order`. If that doesn't produce a strict
3849/// ordering, call `tiebreaker`.
3850pub fn compare_columns<F>(
3851    order: &[ColumnOrder],
3852    left: &[Datum],
3853    right: &[Datum],
3854    tiebreaker: F,
3855) -> Ordering
3856where
3857    F: Fn() -> Ordering,
3858{
3859    for order in order {
3860        let cmp = match (&left[order.column], &right[order.column]) {
3861            (Datum::Null, Datum::Null) => Ordering::Equal,
3862            (Datum::Null, _) => {
3863                if order.nulls_last {
3864                    Ordering::Greater
3865                } else {
3866                    Ordering::Less
3867                }
3868            }
3869            (_, Datum::Null) => {
3870                if order.nulls_last {
3871                    Ordering::Less
3872                } else {
3873                    Ordering::Greater
3874                }
3875            }
3876            (lval, rval) => {
3877                if order.desc {
3878                    rval.cmp(lval)
3879                } else {
3880                    lval.cmp(rval)
3881                }
3882            }
3883        };
3884        if cmp != Ordering::Equal {
3885            return cmp;
3886        }
3887    }
3888    tiebreaker()
3889}
3890
3891/// Describe a window frame, e.g. `RANGE UNBOUNDED PRECEDING` or
3892/// `ROWS BETWEEN 5 PRECEDING AND CURRENT ROW`.
3893///
3894/// Window frames define a subset of the partition , and only a subset of
3895/// window functions make use of the window frame.
3896#[derive(
3897    Debug,
3898    Clone,
3899    Eq,
3900    PartialEq,
3901    Ord,
3902    PartialOrd,
3903    Serialize,
3904    Deserialize,
3905    Hash
3906)]
3907pub struct WindowFrame {
3908    /// ROWS, RANGE or GROUPS
3909    pub units: WindowFrameUnits,
3910    /// Where the frame starts
3911    pub start_bound: WindowFrameBound,
3912    /// Where the frame ends
3913    pub end_bound: WindowFrameBound,
3914}
3915
3916impl Display for WindowFrame {
3917    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
3918        write!(
3919            f,
3920            "{} between {} and {}",
3921            self.units, self.start_bound, self.end_bound
3922        )
3923    }
3924}
3925
3926impl WindowFrame {
3927    /// Return the default window frame used when one is not explicitly defined
3928    pub fn default() -> Self {
3929        WindowFrame {
3930            units: WindowFrameUnits::Range,
3931            start_bound: WindowFrameBound::UnboundedPreceding,
3932            end_bound: WindowFrameBound::CurrentRow,
3933        }
3934    }
3935
3936    fn includes_current_row(&self) -> bool {
3937        use WindowFrameBound::*;
3938        match self.start_bound {
3939            UnboundedPreceding => match self.end_bound {
3940                UnboundedPreceding => false,
3941                OffsetPreceding(0) => true,
3942                OffsetPreceding(_) => false,
3943                CurrentRow => true,
3944                OffsetFollowing(_) => true,
3945                UnboundedFollowing => true,
3946            },
3947            OffsetPreceding(0) => match self.end_bound {
3948                UnboundedPreceding => unreachable!(),
3949                OffsetPreceding(0) => true,
3950                // Any nonzero offsets here will create an empty window
3951                OffsetPreceding(_) => false,
3952                CurrentRow => true,
3953                OffsetFollowing(_) => true,
3954                UnboundedFollowing => true,
3955            },
3956            OffsetPreceding(_) => match self.end_bound {
3957                UnboundedPreceding => unreachable!(),
3958                // Window ends at the current row
3959                OffsetPreceding(0) => true,
3960                OffsetPreceding(_) => false,
3961                CurrentRow => true,
3962                OffsetFollowing(_) => true,
3963                UnboundedFollowing => true,
3964            },
3965            CurrentRow => true,
3966            OffsetFollowing(0) => match self.end_bound {
3967                UnboundedPreceding => unreachable!(),
3968                OffsetPreceding(_) => unreachable!(),
3969                CurrentRow => unreachable!(),
3970                OffsetFollowing(_) => true,
3971                UnboundedFollowing => true,
3972            },
3973            OffsetFollowing(_) => match self.end_bound {
3974                UnboundedPreceding => unreachable!(),
3975                OffsetPreceding(_) => unreachable!(),
3976                CurrentRow => unreachable!(),
3977                OffsetFollowing(_) => false,
3978                UnboundedFollowing => false,
3979            },
3980            UnboundedFollowing => false,
3981        }
3982    }
3983}
3984
3985/// Describe how frame bounds are interpreted
3986#[derive(
3987    Debug,
3988    Clone,
3989    Eq,
3990    PartialEq,
3991    Ord,
3992    PartialOrd,
3993    Serialize,
3994    Deserialize,
3995    Hash
3996)]
3997pub enum WindowFrameUnits {
3998    /// Each row is treated as the unit of work for bounds
3999    Rows,
4000    /// Each peer group is treated as the unit of work for bounds,
4001    /// and offset-based bounds use the value of the ORDER BY expression
4002    Range,
4003    /// Each peer group is treated as the unit of work for bounds.
4004    /// Groups is currently not supported, and it is rejected during planning.
4005    Groups,
4006}
4007
4008impl Display for WindowFrameUnits {
4009    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
4010        match self {
4011            WindowFrameUnits::Rows => write!(f, "rows"),
4012            WindowFrameUnits::Range => write!(f, "range"),
4013            WindowFrameUnits::Groups => write!(f, "groups"),
4014        }
4015    }
4016}
4017
4018/// Specifies [WindowFrame]'s `start_bound` and `end_bound`
4019///
4020/// The order between frame bounds is significant, as Postgres enforces
4021/// some restrictions there.
4022#[derive(
4023    Debug,
4024    Clone,
4025    Serialize,
4026    Deserialize,
4027    PartialEq,
4028    Eq,
4029    Hash,
4030    PartialOrd,
4031    Ord
4032)]
4033pub enum WindowFrameBound {
4034    /// `UNBOUNDED PRECEDING`
4035    UnboundedPreceding,
4036    /// `<N> PRECEDING`
4037    OffsetPreceding(u64),
4038    /// `CURRENT ROW`
4039    CurrentRow,
4040    /// `<N> FOLLOWING`
4041    OffsetFollowing(u64),
4042    /// `UNBOUNDED FOLLOWING`.
4043    UnboundedFollowing,
4044}
4045
4046impl Display for WindowFrameBound {
4047    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
4048        match self {
4049            WindowFrameBound::UnboundedPreceding => write!(f, "unbounded preceding"),
4050            WindowFrameBound::OffsetPreceding(offset) => write!(f, "{} preceding", offset),
4051            WindowFrameBound::CurrentRow => write!(f, "current row"),
4052            WindowFrameBound::OffsetFollowing(offset) => write!(f, "{} following", offset),
4053            WindowFrameBound::UnboundedFollowing => write!(f, "unbounded following"),
4054        }
4055    }
4056}
4057
4058/// Maximum iterations for a LetRec.
4059#[derive(
4060    Debug,
4061    Clone,
4062    Copy,
4063    PartialEq,
4064    Eq,
4065    PartialOrd,
4066    Ord,
4067    Hash,
4068    Serialize,
4069    Deserialize
4070)]
4071pub struct LetRecLimit {
4072    /// Maximum number of iterations to evaluate.
4073    pub max_iters: NonZeroU64,
4074    /// Whether to throw an error when reaching the above limit.
4075    /// If true, we simply use the current contents of each Id as the final result.
4076    pub return_at_limit: bool,
4077}
4078
4079impl LetRecLimit {
4080    /// Compute the smallest limit from a Vec of `LetRecLimit`s.
4081    pub fn min_max_iter(limits: &Vec<Option<LetRecLimit>>) -> Option<u64> {
4082        limits
4083            .iter()
4084            .filter_map(|l| l.as_ref().map(|l| l.max_iters.get()))
4085            .min()
4086    }
4087
4088    /// The default value of `LetRecLimit::return_at_limit` when using the RECURSION LIMIT option of
4089    /// WMR without ERROR AT or RETURN AT.
4090    pub const RETURN_AT_LIMIT_DEFAULT: bool = false;
4091}
4092
4093impl Display for LetRecLimit {
4094    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
4095        write!(f, "[recursion_limit={}", self.max_iters)?;
4096        if self.return_at_limit != LetRecLimit::RETURN_AT_LIMIT_DEFAULT {
4097            write!(f, ", return_at_limit")?;
4098        }
4099        write!(f, "]")
4100    }
4101}
4102
4103/// For a global Get, this indicates whether we are going to read from Persist or from an index.
4104/// (See comment in MirRelationExpr::Get.)
4105#[derive(
4106    Clone,
4107    Debug,
4108    Eq,
4109    PartialEq,
4110    Ord,
4111    PartialOrd,
4112    Serialize,
4113    Deserialize,
4114    Hash
4115)]
4116pub enum AccessStrategy {
4117    /// It's either a local Get (a CTE), or unknown at the time.
4118    /// `prune_and_annotate_dataflow_index_imports` decides it for global Gets, and thus switches to
4119    /// one of the other variants.
4120    UnknownOrLocal,
4121    /// The Get will read from Persist.
4122    Persist,
4123    /// The Get will read from an index or indexes: (index id, how the index will be used).
4124    Index(Vec<(GlobalId, IndexUsageType)>),
4125    /// The Get will read a collection that is computed by the same dataflow, but in a different
4126    /// `BuildDesc` in `objects_to_build`.
4127    SameDataflow,
4128}
4129
4130#[cfg(test)]
4131mod tests {
4132    use std::num::NonZeroUsize;
4133
4134    use mz_repr::explain::text::text_string_at;
4135
4136    use crate::explain::HumanizedExplain;
4137
4138    use super::*;
4139
4140    #[mz_ore::test]
4141    fn test_row_set_finishing_as_text() {
4142        let finishing = RowSetFinishing {
4143            order_by: vec![ColumnOrder {
4144                column: 4,
4145                desc: true,
4146                nulls_last: true,
4147            }],
4148            limit: Some(NonNeg::try_from(7).unwrap()),
4149            offset: Default::default(),
4150            project: vec![1, 3, 4, 5],
4151        };
4152
4153        let mode = HumanizedExplain::new(false);
4154        let expr = mode.expr(&finishing, None);
4155
4156        let act = text_string_at(&expr, mz_ore::str::Indent::default);
4157
4158        let exp = {
4159            use mz_ore::fmt::FormatBuffer;
4160            let mut s = String::new();
4161            write!(&mut s, "Finish");
4162            write!(&mut s, " order_by=[#4 desc nulls_last]");
4163            write!(&mut s, " limit=7");
4164            write!(&mut s, " output=[#1, #3..=#5]");
4165            writeln!(&mut s, "");
4166            s
4167        };
4168
4169        assert_eq!(act, exp);
4170    }
4171
4172    #[mz_ore::test]
4173    fn test_row_set_finishing_incremental_max_returned_query_size() {
4174        let row = Row::pack_slice(&[Datum::String("hello")]);
4175        let row_size = u64::cast_from(row.data().len());
4176        let diff = NonZeroUsize::new(1).unwrap();
4177        let batch = RowCollection::new(vec![(row, diff)], &[]);
4178
4179        // Set max_returned_query_size to hold exactly 2 batches worth of rows.
4180        let mut finishing = RowSetFinishingIncremental::new(0, None, vec![0], Some(row_size * 2));
4181
4182        let max_result_size = u64::MAX;
4183
4184        let r = finishing.finish_incremental_inner(batch.clone(), max_result_size);
4185        assert!(r.is_ok());
4186        assert_eq!(finishing.remaining_max_returned_query_size, Some(row_size));
4187
4188        let r = finishing.finish_incremental_inner(batch.clone(), max_result_size);
4189        assert!(r.is_ok());
4190        assert_eq!(finishing.remaining_max_returned_query_size, Some(0));
4191
4192        let r = finishing.finish_incremental_inner(batch, max_result_size);
4193        assert!(r.unwrap_err().contains("total result exceeds max size"));
4194    }
4195}
4196
4197/// An iterator over AST structures, which calls out nodes in difference.
4198///
4199/// The iterators visit two ASTs in tandem, continuing as long as the AST node data matches,
4200/// and yielding an output pair as soon as the AST nodes do not match. Their intent is to call
4201/// attention to the moments in the ASTs where they differ, and incidentally a stack-free way
4202/// to compare two ASTs.
4203mod structured_diff {
4204
4205    use super::MirRelationExpr;
4206    use itertools::Itertools;
4207
4208    ///  An iterator over structured differences between two `MirRelationExpr` instances.
4209    pub struct MreDiff<'a> {
4210        /// Pairs of expressions that must still be compared.
4211        todo: Vec<(&'a MirRelationExpr, &'a MirRelationExpr)>,
4212    }
4213
4214    impl<'a> MreDiff<'a> {
4215        /// Create a new `MirRelationExpr` structured difference.
4216        pub fn new(expr1: &'a MirRelationExpr, expr2: &'a MirRelationExpr) -> Self {
4217            MreDiff {
4218                todo: vec![(expr1, expr2)],
4219            }
4220        }
4221    }
4222
4223    impl<'a> Iterator for MreDiff<'a> {
4224        // Pairs of expressions that do not match.
4225        type Item = (&'a MirRelationExpr, &'a MirRelationExpr);
4226
4227        fn next(&mut self) -> Option<Self::Item> {
4228            while let Some((expr1, expr2)) = self.todo.pop() {
4229                match (expr1, expr2) {
4230                    (
4231                        MirRelationExpr::Constant {
4232                            rows: rows1,
4233                            typ: typ1,
4234                        },
4235                        MirRelationExpr::Constant {
4236                            rows: rows2,
4237                            typ: typ2,
4238                        },
4239                    ) => {
4240                        if rows1 != rows2 || typ1 != typ2 {
4241                            return Some((expr1, expr2));
4242                        }
4243                    }
4244                    (
4245                        MirRelationExpr::Get {
4246                            id: id1,
4247                            typ: typ1,
4248                            access_strategy: as1,
4249                        },
4250                        MirRelationExpr::Get {
4251                            id: id2,
4252                            typ: typ2,
4253                            access_strategy: as2,
4254                        },
4255                    ) => {
4256                        if id1 != id2 || typ1 != typ2 || as1 != as2 {
4257                            return Some((expr1, expr2));
4258                        }
4259                    }
4260                    (
4261                        MirRelationExpr::Let {
4262                            id: id1,
4263                            body: body1,
4264                            value: value1,
4265                        },
4266                        MirRelationExpr::Let {
4267                            id: id2,
4268                            body: body2,
4269                            value: value2,
4270                        },
4271                    ) => {
4272                        if id1 != id2 {
4273                            return Some((expr1, expr2));
4274                        } else {
4275                            self.todo.push((body1, body2));
4276                            self.todo.push((value1, value2));
4277                        }
4278                    }
4279                    (
4280                        MirRelationExpr::LetRec {
4281                            ids: ids1,
4282                            body: body1,
4283                            values: values1,
4284                            limits: limits1,
4285                        },
4286                        MirRelationExpr::LetRec {
4287                            ids: ids2,
4288                            body: body2,
4289                            values: values2,
4290                            limits: limits2,
4291                        },
4292                    ) => {
4293                        if ids1 != ids2 || values1.len() != values2.len() || limits1 != limits2 {
4294                            return Some((expr1, expr2));
4295                        } else {
4296                            self.todo.push((body1, body2));
4297                            self.todo.extend(values1.iter().zip_eq(values2.iter()));
4298                        }
4299                    }
4300                    (
4301                        MirRelationExpr::Project {
4302                            outputs: outputs1,
4303                            input: input1,
4304                        },
4305                        MirRelationExpr::Project {
4306                            outputs: outputs2,
4307                            input: input2,
4308                        },
4309                    ) => {
4310                        if outputs1 != outputs2 {
4311                            return Some((expr1, expr2));
4312                        } else {
4313                            self.todo.push((input1, input2));
4314                        }
4315                    }
4316                    (
4317                        MirRelationExpr::Map {
4318                            scalars: scalars1,
4319                            input: input1,
4320                        },
4321                        MirRelationExpr::Map {
4322                            scalars: scalars2,
4323                            input: input2,
4324                        },
4325                    ) => {
4326                        if scalars1 != scalars2 {
4327                            return Some((expr1, expr2));
4328                        } else {
4329                            self.todo.push((input1, input2));
4330                        }
4331                    }
4332                    (
4333                        MirRelationExpr::Filter {
4334                            predicates: predicates1,
4335                            input: input1,
4336                        },
4337                        MirRelationExpr::Filter {
4338                            predicates: predicates2,
4339                            input: input2,
4340                        },
4341                    ) => {
4342                        if predicates1 != predicates2 {
4343                            return Some((expr1, expr2));
4344                        } else {
4345                            self.todo.push((input1, input2));
4346                        }
4347                    }
4348                    (
4349                        MirRelationExpr::FlatMap {
4350                            input: input1,
4351                            func: func1,
4352                            exprs: exprs1,
4353                        },
4354                        MirRelationExpr::FlatMap {
4355                            input: input2,
4356                            func: func2,
4357                            exprs: exprs2,
4358                        },
4359                    ) => {
4360                        if func1 != func2 || exprs1 != exprs2 {
4361                            return Some((expr1, expr2));
4362                        } else {
4363                            self.todo.push((input1, input2));
4364                        }
4365                    }
4366                    (
4367                        MirRelationExpr::Join {
4368                            inputs: inputs1,
4369                            equivalences: eq1,
4370                            implementation: impl1,
4371                        },
4372                        MirRelationExpr::Join {
4373                            inputs: inputs2,
4374                            equivalences: eq2,
4375                            implementation: impl2,
4376                        },
4377                    ) => {
4378                        if inputs1.len() != inputs2.len() || eq1 != eq2 || impl1 != impl2 {
4379                            return Some((expr1, expr2));
4380                        } else {
4381                            self.todo.extend(inputs1.iter().zip_eq(inputs2.iter()));
4382                        }
4383                    }
4384                    (
4385                        MirRelationExpr::Reduce {
4386                            aggregates: aggregates1,
4387                            input: inputs1,
4388                            group_key: gk1,
4389                            monotonic: m1,
4390                            expected_group_size: egs1,
4391                        },
4392                        MirRelationExpr::Reduce {
4393                            aggregates: aggregates2,
4394                            input: inputs2,
4395                            group_key: gk2,
4396                            monotonic: m2,
4397                            expected_group_size: egs2,
4398                        },
4399                    ) => {
4400                        if aggregates1 != aggregates2 || gk1 != gk2 || m1 != m2 || egs1 != egs2 {
4401                            return Some((expr1, expr2));
4402                        } else {
4403                            self.todo.push((inputs1, inputs2));
4404                        }
4405                    }
4406                    (
4407                        MirRelationExpr::TopK {
4408                            group_key: gk1,
4409                            order_key: order1,
4410                            input: input1,
4411                            limit: l1,
4412                            offset: o1,
4413                            monotonic: m1,
4414                            expected_group_size: egs1,
4415                        },
4416                        MirRelationExpr::TopK {
4417                            group_key: gk2,
4418                            order_key: order2,
4419                            input: input2,
4420                            limit: l2,
4421                            offset: o2,
4422                            monotonic: m2,
4423                            expected_group_size: egs2,
4424                        },
4425                    ) => {
4426                        if order1 != order2
4427                            || gk1 != gk2
4428                            || l1 != l2
4429                            || o1 != o2
4430                            || m1 != m2
4431                            || egs1 != egs2
4432                        {
4433                            return Some((expr1, expr2));
4434                        } else {
4435                            self.todo.push((input1, input2));
4436                        }
4437                    }
4438                    (
4439                        MirRelationExpr::Negate { input: input1 },
4440                        MirRelationExpr::Negate { input: input2 },
4441                    ) => {
4442                        self.todo.push((input1, input2));
4443                    }
4444                    (
4445                        MirRelationExpr::Threshold { input: input1 },
4446                        MirRelationExpr::Threshold { input: input2 },
4447                    ) => {
4448                        self.todo.push((input1, input2));
4449                    }
4450                    (
4451                        MirRelationExpr::Union {
4452                            base: base1,
4453                            inputs: inputs1,
4454                        },
4455                        MirRelationExpr::Union {
4456                            base: base2,
4457                            inputs: inputs2,
4458                        },
4459                    ) => {
4460                        if inputs1.len() != inputs2.len() {
4461                            return Some((expr1, expr2));
4462                        } else {
4463                            self.todo.push((base1, base2));
4464                            self.todo.extend(inputs1.iter().zip_eq(inputs2.iter()));
4465                        }
4466                    }
4467                    (
4468                        MirRelationExpr::ArrangeBy {
4469                            keys: keys1,
4470                            input: input1,
4471                        },
4472                        MirRelationExpr::ArrangeBy {
4473                            keys: keys2,
4474                            input: input2,
4475                        },
4476                    ) => {
4477                        if keys1 != keys2 {
4478                            return Some((expr1, expr2));
4479                        } else {
4480                            self.todo.push((input1, input2));
4481                        }
4482                    }
4483                    _ => {
4484                        return Some((expr1, expr2));
4485                    }
4486                }
4487            }
4488            None
4489        }
4490    }
4491}