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