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