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}